Shock formation, viscous regularization, and coupled ODE systems for A→B→C reactions, where physics-informed training earns its keep.
The heat equation from Heat Equation is linear: if $u_1$ and $u_2$ are solutions, so is $\alpha u_1 + \beta u_2$. Linearity is what makes Fourier analysis work. You decompose the initial condition into eigenfunctions of the spatial operator, each mode evolves independently, and you superpose the results. The analytical machinery is beautiful and exact.
Real engineering systems are relentlessly nonlinear. In fluid dynamics, fluid carries its own momentum, the velocity field advects itself, producing the nonlinear convective term $(\mathbf{u}\cdot\nabla)\mathbf{u}$ in the Navier-Stokes equations. In chemical kinetics, rate laws are products of concentrations: $r = k[A][B]$ for a bimolecular reaction, or the hyperbolic Michaelis-Menten form $r = V_\mathrm{max}[S]/(K_m + [S])$ for enzyme kinetics. In combustion, temperature-dependent Arrhenius rate constants introduce exponential nonlinearity.
Classical numerical methods handle nonlinearity through linearisation (Newton iterations at each time step) or explicit time-stepping with severe stability constraints. Finite element methods can struggle with convection-dominated problems (the Péclet number problem) and require upwinding schemes or stabilisation terms. The PINN approach sidesteps mesh generation entirely and encodes the nonlinear residual directly, though it introduces its own set of challenges around spectral bias and loss landscape geometry that we examine in this article.
Burgers' equation is the benchmark nonlinear PDE for PINNs. It is the headline example in Raissi, Perdikaris & Karniadakis (2019) and has been reproduced and extended in hundreds of subsequent papers. Its importance stems from two properties: it is simple enough to analyse rigorously, and it produces a near-shock at low viscosity that genuinely tests the network's ability to represent steep gradients.
This is the 1D Navier-Stokes momentum equation with the pressure gradient suppressed and the velocity field called $u$. The left side is the material derivative $Du/Dt = \partial u/\partial t + u\,\partial u/\partial x$, the rate of change of velocity following a fluid parcel. The right side is viscous dissipation. When $\nu \to 0$ the equation becomes the inviscid Burgers equation, which develops a genuine discontinuity (shock) in finite time from smooth initial data.
The initial condition $u(x,0) = -\sin(\pi x)$ is negative for $x > 0$ and positive for $x < 0$. This means fluid is moving toward $x = 0$ from both sides, creating a compressive wave that steepens into a near-shock around $t \approx 0.4$. The viscosity $\nu = 0.01/\pi \approx 0.00318$ is small enough that the layer is very thin but just sufficient to prevent an actual discontinuity, yielding a steep but smooth gradient that the network must resolve.
The residual that must be driven to zero at every interior collocation point is:
The critical observation is that the nonlinear term $u\,\partial u/\partial x$ involves the product of the network output $u$ (from a forward pass) with the spatial gradient $\partial u/\partial x$ (from autograd). This product is automatically handled by PyTorch's autograd engine, no special treatment needed. The residual computation reads exactly like the mathematical formula:
nu = 0.01 / np.pi # ≈ 0.00318 (small viscosity → near-shock) def residual(model, x, t): u = model(x, t) u_t = torch.autograd.grad(u, t, torch.ones_like(u), create_graph=True)[0] u_x = torch.autograd.grad(u, x, torch.ones_like(u), create_graph=True)[0] u_xx= torch.autograd.grad(u_x, x, torch.ones_like(u_x), create_graph=True)[0] # The nonlinear term u*u_x is computed elementwise, autograd handles it return u_t + u * u_x - nu * u_xx def loss_fn(): r = residual(model, x_f, t_f) L_f = torch.mean(r**2) L_0 = torch.mean((model(x_0, t_0) - u_0)**2) # IC: -sin(πx) L_b = (torch.mean(model(x_bL, t_b)**2) + torch.mean(model(x_bR, t_b)**2)) # BC: u=0 at x=±1 return L_f + L_0 + L_b
The small viscosity $\nu = 0.01/\pi$ is simultaneously the most physically interesting aspect of the problem and the most numerically challenging. It means the layer thickness at the shock is $O(\nu/|\nabla u|) \approx O(0.003)$, occupying a tiny fraction of the domain. Uniform random sampling of 10,000 collocation points across $[-1,1]\times[0,1]$ places on average only $\sim$30 points inside this thin layer, far too few to strongly constrain the residual where it matters most.
The notebook addresses this with a deeper and wider network ([2, 40, 40, 40, 40, 40, 1], five hidden layers of width 40) and critically, a two-stage optimisation strategy that has become the de facto standard in the PINN literature.
Adam is run for 15,000 epochs at $10^{-3}$. Its adaptive learning rates per parameter make it robust to poorly conditioned loss landscapes and effective at escaping sharp valleys. However, Adam alone typically leaves the solution slightly diffuse around steep features, it converges to a good basin but cannot tightly minimise the residual there.
L-BFGS (Limited-memory Broyden-Fletcher-Goldfarb-Shanno) is a quasi-Newton method. It approximates the inverse Hessian from the history of gradient vectors, allowing much larger and better-directed steps than gradient descent. Starting from Adam's solution, L-BFGS can drive the loss down several additional orders of magnitude and sharpen the resolution of the steep layer.
# Stage 1: Adam robust, handles the difficult early training opt = torch.optim.Adam(model.parameters(), lr=1e-3) for ep in range(15000): opt.zero_grad(); loss = loss_fn(); loss.backward(); opt.step() if ep % 2500 == 0: print(f"[adam] epoch {ep:5d} | loss {loss.item():.3e}") # Stage 2: L-BFGS quasi-Newton, drives loss to near-machine precision opt2 = torch.optim.LBFGS( model.parameters(), max_iter=3000, tolerance_grad=1e-9, tolerance_change=1e-12, history_size=50, line_search_fn="strong_wolfe" ) def closure(): opt2.zero_grad() l = loss_fn() l.backward() return l opt2.step(closure) print("L-BFGS done | final loss:", loss_fn().item())
The network architecture uses five hidden layers of width 40 rather than the three layers of width 64 used for the heat equation. For the Burgers problem with its sharp interior layer, depth (more layers) helps more than width because it enables the network to compose multiple nonlinear transformations needed to represent the steep gradient, while keeping the total parameter count manageable for the L-BFGS inverse Hessian approximation.
Notebook 05_reaction_kinetics.ipynb addresses a different kind of problem: not a PDE over space-time, but a system of coupled ODEs representing chemical species concentrations evolving over time. The consecutive reaction $A \xrightarrow{k_1} B \xrightarrow{k_2} C$ describes a process where a reactant $A$ is converted to an intermediate $B$, which then converts to a product $C$. This is the simplest form of a reaction network and a foundational model in pharmaceutical synthesis, metabolic pathway analysis, and polymer degradation kinetics.
With initial conditions $C_A(0) = 1$, $C_B(0) = C_C(0) = 0$ and rate constants $k_1 = 1.5$, $k_2 = 0.6$ from the notebook. Mass is conserved exactly: $C_A(t) + C_B(t) + C_C(t) = 1$ for all $t$.
For $k_1 \neq k_2$, the system admits a closed-form solution derived by successive integration of the first-order linear ODEs:
With $k_1 = 1.5 > k_2 = 0.6$, species $B$ accumulates before $C_A$ is depleted because the formation rate ($k_1 C_A$) initially exceeds the consumption rate ($k_2 C_B$). The intermediate $C_B$ peaks at $t^* = \ln(k_1/k_2)/(k_1-k_2) \approx 1.02$ before decaying to zero as the reaction completes.
The reaction PINN outputs three values simultaneously: a single shared network with one input ($t$) and three outputs ($C_A, C_B, C_C$). This shared trunk, multiple head architecture is the standard approach for coupled systems. The shared hidden layers learn common temporal features (e.g., characteristic timescales, overall curvature shapes) that all three species exhibit, while the separate output dimensions specialise to each species' profile.
class KineticsPINN(nn.Module): # Single input t → three outputs [C_A, C_B, C_C] def __init__(self, width=64, depth=4): super().__init__() layers = [1] + [width]*depth + [3] seq = [] for i in range(len(layers)-1): seq.append(nn.Linear(layers[i], layers[i+1])) if i < len(layers)-2: seq.append(nn.Tanh()) self.net = nn.Sequential(*seq) def forward(self, t): return self.net(t) # shape (N, 3) def d_dt(y, t): # Differentiate each output scalar w.r.t. t return torch.autograd.grad(y, t, torch.ones_like(y), create_graph=True)[0] def loss_fn(): y = model(t_col) # shape (N, 3) A, B, C = y[:,0:1], y[:,1:2], y[:,2:3] # Each species' time derivative At = d_dt(A, t_col) Bt = d_dt(B, t_col) Ct = d_dt(C, t_col) # ODE residuals (should all be zero) rA = At + k1 * A rB = Bt - k1 * A + k2 * B rC = Ct - k2 * B L_pde = torch.mean(rA**2 + rB**2 + rC**2) L_ic = torch.mean((model(t0) - y0)**2) # IC weighted ×20, helps during early training return L_pde + 20.0 * L_ic
Several design choices in this code are worth noting. First, y[:, 0:1] preserves the batch dimension (shape (N,1)), which is required for autograd.grad to compute per-sample gradients correctly, using y[:, 0] would produce shape (N,) and silently fail. Second, the three residuals are summed before taking the mean, which effectively weights each species equally. Third, the IC loss carries a weight of 20, this is a form of manual loss weighting that ensures the initial condition is satisfied early in training, providing a reliable anchor for the time-integration structure of the ODE system.
Notice that mass conservation $C_A + C_B + C_C = 1$ is not imposed as a loss term in the notebook. After training, the constraint is evaluated as a diagnostic: if $C_A + C_B + C_C \approx 1$ everywhere, it confirms that the three ODE residuals are genuinely satisfied (since their sum implies the conservation law). This is a powerful cross-check, a network that satisfies mass conservation only because it was penalised to do so tells us less about ODE accuracy than one that satisfies it as an emergent consequence of correct residuals.
For a single run with known $(k_1, k_2)$, a classical ODE solver (e.g., SciPy's solve_ivp with an adaptive Runge-Kutta method) solves the system faster and more accurately than a PINN. Classical solvers for smooth ODEs are mature, fast, and reliable. So why use a PINN here?
The answer is not the forward problem, it is everything that comes after:
1. The inverse problem. If you have measured concentration profiles from inline spectroscopy (e.g., Raman or NIR), the PINN framework can treat $k_1$ and $k_2$ as learnable parameters and fit them to the data while simultaneously satisfying the ODE structure. A classical solver requires an outer optimisation loop (e.g., Bard 1974, nonlinear least squares) that calls the solver many times. The PINN computes everything in a single differentiable computational graph. Inverse Problems covers this in detail.
2. Data assimilation mid-trajectory. If process analytical technology (PAT) instruments provide inline concentration readings at multiple time points, these can be added directly as a data loss term $\mathcal{L}_\mathrm{data} = \sum_i (u(\hat{t}_i;\theta) - \hat{u}_i)^2$. The PINN then produces a solution that is simultaneously consistent with the ODE physics and with the measurements, a constrained regression that classical interpolation cannot perform.
3. Uncertainty and extrapolation. Neural network ensembles or Bayesian PINNs can propagate uncertainty from sparse observations through the physics constraints, yielding principled confidence intervals on quantities like time-to-peak for $C_B$ or yield of $C_C$. This is directly relevant to ICH Q10 Pharmaceutical Quality System requirements for process understanding and real-time release testing.
4. Complex nonlinear kinetics. First-order kinetics is the simplest case. Michaelis-Menten enzyme kinetics, autocatalytic reactions, and multi-step mechanisms with reversibility and inhibition produce ODE systems where analytical solutions do not exist. The PINN residual is written in exactly the same way regardless of how complex the rate law is, you simply change the expression inside residual().
The $A \to B \to C$ reaction network is not merely a pedagogical model. It describes a wide class of pharmaceutical manufacturing steps directly:
| Process | A (feed) | B (intermediate) | C (product) | Relevance of PINN |
|---|---|---|---|---|
| API synthesis step | Starting material | Reactive intermediate | Active ingredient | Maximise $C_C$ while limiting impurity $B$ accumulation |
| Hydrolysis degradation | Drug substance | Hydrolysis product 1 | Final degradant | Shelf-life prediction from accelerated stability data |
| Enzyme catalysis | Substrate | ES complex | Product | Parameter estimation of $K_m$, $V_\mathrm{max}$ from sparse data |
| Metabolic pathway | Parent drug | Active metabolite | Inactive metabolite | PK/PD modelling with inline biomarker data |
ICH Q10 emphasises the Pharmaceutical Quality System as a science- and risk-based framework for process understanding. Real-time monitoring of reaction progress using inline PAT instruments generates a stream of spectroscopic data that must be interpreted against kinetic models. A PINN provides a natural data-physics fusion architecture: the physics constrains the solution to be consistent with the rate equations even between measurement time points, while the data loss pulls the trajectory toward the observed signal.
After 10,000 Adam epochs the KineticsPINN achieves relative $L_2$ error below $10^{-3}$ on all three species evaluated on 300 uniform time points in $[0,6]$. Mass conservation is satisfied to within $\pm 0.002$ across the entire time window without being explicitly enforced as a loss term. The intermediate $C_B$ peak location (time-to-maximum yield) is recovered to within 1% of the analytical value $t^* = 1.02$.
Sensitivity to hyperparameters is greater than for the heat equation: reducing the IC weight below 5 typically leads to drift in the initial condition after $\sim$3,000 epochs, with the solution converging to a physically wrong trajectory that minimises the ODE residual but violates $C_A(0)=1$. This underscores a general principle, for initial value problems, the IC loss must be weighted strongly enough to serve as a reliable anchor, especially during the early training phase when the PDE residual gradient is large and can dominate.
The Burgers equation and the reaction kinetics system together establish the two most important extensions of the basic PINN framework: nonlinear PDE terms and multi-output ODE systems. But both problems were solved in the forward direction: given the physics and initial conditions, compute the solution.
Inverse Problems inverts this. We will take the same reaction kinetics network and remove $k_1$, $k_2$ from the code, instead treating them as learnable parameters in the model's parameter dictionary. When sparse noisy measurements of the concentration profiles are provided, the PINN will jointly optimise the network weights and the rate constants, doing physics-constrained parameter estimation in a single differentiable computation. This is the core capability that makes PINNs genuinely useful for real experimental data.
Adaptive Sampling returns to Burgers to address the adaptive sampling problem: how to automatically concentrate collocation points near the shock, and how to dynamically reweight the loss terms to counteract the spectral bias that prevents the network from resolving steep gradients.