Physics-Informed Neural Networks Neural Odes
# Physics-Informed Neural Networks & Neural ODEs
## Introduction & Motivation
Physics-Informed Neural Networks (PINNs) and Neural Ordinary Differential Equations (Neural ODEs) represent two complementary approaches to embedding the mathematical structure of dynamical systems directly into neural network architectures and training objectives, rather than treating deep learning purely as black-box function approximation. PINNs, introduced by Raissi, Perdikaris, and Karniadakis in 2019, incorporate the governing differential equations of a physical system directly into the loss function, allowing a neural network to be trained to satisfy known physics (conservation laws, boundary conditions, governing PDEs) even with sparse or noisy observational data. Neural ODEs, introduced by Chen et al. in 2018, take a different but related approach: rather than defining a neural network as a fixed sequence of discrete layers, they parameterize the derivative of a hidden state with a neural network and use a differentiable ODE solver to compute the network's output, effectively treating depth as a continuous rather than discrete quantity. Both approaches are motivated by the observation that purely data-driven deep learning often requires enormous datasets to learn dynamics that are already well understood mathematically, and that ignoring known physical structure both wastes this prior knowledge and can produce predictions that violate basic physical laws such as energy conservation. These methods have found application across computational fluid dynamics, climate modeling, molecular dynamics, biomedical simulation, and any domain where the underlying system is governed by differential equations but full high-fidelity numerical simulation is too computationally expensive for the desired use case (e.g., real-time control or large-scale parameter sweeps).
## Core Concepts & Theory
The central idea of a PINN is to use automatic differentiation, the same machinery used to compute gradients for backpropagation, to compute the derivatives of the network's output with respect to its input coordinates (space and time), and then to penalize the degree to which those derivatives fail to satisfy a known governing differential equation. This means a PINN's loss function combines a data loss (mismatch between network predictions and any available observed data points) with a physics loss (residual of the governing PDE evaluated at a large set of "collocation points" sampled throughout the domain, where no data is required). The physics loss acts as a powerful regularizer, allowing PINNs to produce physically plausible solutions even in regions of the domain with no observational data whatsoever. Neural ODEs, in contrast, generalize the residual connection pattern found in ResNets: a ResNet layer update h_{t+1} = h_t + f(h_t, heta_t) can be viewed as an Euler discretization of the continuous-time ODE \frac{dh}{dt} = f(h(t), t, heta). A Neural ODE directly parameterizes this continuous dynamics function f with a neural network and computes the hidden state at any desired output time by numerically integrating the ODE, using adaptive-step solvers such as Runge-Kutta or Dormand-Prince methods, rather than a fixed number of discrete layer evaluations. Both frameworks share the conceptual thread of embedding differential-equation structure into the learning process, but PINNs typically solve a specific known PDE with unknown solution, while Neural ODEs learn an unknown dynamics function from data with a known (differentiable, continuous) integration structure.
## Mathematical Formulation
For a PINN solving a PDE of the general form \mathcal{N}[u](x, t) = 0 over domain \Omega with boundary and initial conditions, the network u_ heta(x, t) approximates the solution u(x, t), and automatic differentiation computes the required partial derivatives (e.g., \partial u_ heta / \partial t, \partial^2 u_ heta / \partial x^2) needed to evaluate the PDE residual. The total loss combines three terms:
$$ \mathcal{L}( heta) = \lambda_{data} \mathcal{L}_{data} + \lambda_{pde} \mathcal{L}_{pde} + \lambda_{bc} \mathcal{L}_{bc} $$
where the physics residual loss is evaluated at collocation points \{(x_i, t_i)\}_{i=1}^{N_c} sampled throughout the domain:
$$ \mathcal{L}_{pde} = \frac{1}{N_c}\sum_{i=1}^{N_c} \left| \mathcal{N}[u_ heta](x_i, t_i) ight|^2 $$
For example, for the 1D heat equation \frac{\partial u}{\partial t} - \alpha \frac{\partial^2 u}{\partial x^2} = 0, the residual \mathcal{N}[u_ heta] = \frac{\partial u_ heta}{\partial t} - \alpha \frac{\partial^2 u_ heta}{\partial x^2} is computed via automatic differentiation of the network output with respect to its inputs.
For a Neural ODE, the hidden state evolves according to:
$$ \frac{dh(t)}{dt} = f_ heta(h(t), t), \qquad h(t_0) = h_0 $$
and the output at time t_1 is obtained by numerical integration:
$$ h(t_1) = h(t_0) + \int_{t_0}^{t_1} f_ heta(h(t), t) \, dt = ext{ODESolve}(h(t_0), f_ heta, t_0, t_1) $$
Gradients with respect to heta are computed either by backpropagating directly through the operations of the ODE solver, or more memory-efficiently via the adjoint sensitivity method, which solves a second, augmented ODE backward in time:
$$ \frac{da(t)}{dt} = -a(t)^ op \frac{\partial f_ heta(h(t), t)}{\partial h}, \qquad a(t_1) = \frac{\partial \mathcal{L}}{\partial h(t_1)} $$
where a(t) is the adjoint state, yielding O(1) memory cost with respect to the number of solver steps, since intermediate activations from the forward solve do not need to be stored.
## Advanced Theory & Extensions
Neural Controlled Differential Equations (Neural CDEs) extend Neural ODEs to handle irregularly sampled time series by interpolating the observed input path and driving the hidden-state ODE by this continuous input path rather than by time alone, making them well suited to real-world clinical or sensor time-series data with missing or irregular observations. Latent ODEs combine Neural ODEs with variational autoencoders to model irregularly sampled sequential data probabilistically, encoding an observed sequence into a latent initial state and evolving it forward continuously, providing well-calibrated uncertainty estimates for interpolation and extrapolation. On the PINN side, variational PINNs and hp-VPINNs incorporate finite-element-inspired variational (weak-form) formulations of the governing PDE rather than the strong-form residual, which can improve convergence for problems with sharp gradients or discontinuities that are difficult for pure strong-form collocation to resolve. Fourier Neural Operators (FNOs) and DeepONets generalize beyond PINNs by learning operators that map entire functions (e.g., an initial condition or a boundary condition) to entire solution functions, rather than learning a single fixed solution for one specific instance of the PDE; this operator-learning framing allows a trained model to generalize across a family of PDE instances (e.g., different initial conditions) without retraining, which is a substantial advantage over standard PINNs, which must be retrained for each new instance of the boundary/initial conditions. Stiff Neural ODEs address the numerical challenge that occurs when a learned dynamics function exhibits widely varying time scales, requiring implicit solvers or specialized regularization to prevent training instability and prohibitively small adaptive step sizes.
## Computational Considerations
Automatic differentiation for higher-order derivatives, as required by PINNs for second- or higher-order PDEs, is computationally more expensive than first-order backpropagation because it requires nested differentiation (differentiating the gradient computation graph itself), and the cost grows with both the PDE order and the network depth. Collocation point sampling strategy significantly affects both training cost and solution accuracy: uniform random sampling is simple but can under-resolve regions with sharp solution gradients, motivating adaptive resampling strategies that concentrate collocation points in high-residual regions as training progresses. For Neural ODEs, the choice of numerical solver directly trades off computational cost against integration accuracy: fixed-step solvers like Euler or RK4 have predictable, low per-step cost but can produce inaccurate results for stiff dynamics, while adaptive-step solvers like Dormand-Prince (`dopri5`) automatically adjust step size for accuracy but incur unpredictable and sometimes very high computational cost during training when the learned dynamics happen to be stiff at a given point in training. The adjoint sensitivity method reduces memory cost from O( ext{number of solver steps}) to O(1) relative to the forward solve, but at the cost of an additional backward-time ODE solve, and can suffer from numerical inaccuracy in some settings because the backward integration is not guaranteed to retrace the exact same trajectory as the forward solve in floating-point arithmetic. Multi-GPU and distributed training of PINNs for large-scale 3D physical domains requires domain decomposition strategies analogous to classical parallel PDE solvers, since a single collocation-point-based loss over a very fine 3D-plus-time domain can require enormous batch sizes.
## Practical Implementation Strategies
Loss term weighting (\lambda_{data}, \lambda_{pde}, \lambda_{bc}) is one of the most consequential and difficult-to-tune aspects of PINN training, since poorly balanced weights can cause the optimizer to prioritize trivially satisfying one loss term (e.g., a PDE residual near zero everywhere from a near-constant solution) while ignoring boundary conditions or data fit; adaptive weighting schemes that rescale loss terms based on gradient magnitude statistics during training (e.g., NTK-based or gradient-norm-based reweighting) substantially improve training stability over fixed manual weights. Non-dimensionalizing the physical problem (rescaling spatial, temporal, and physical quantities to be of order unity) before training is standard practice and often critical, since neural networks trained with standard initialization schemes and optimizers perform poorly when inputs or targets span many orders of magnitude, as raw physical units frequently do. For Neural ODEs, starting with a fixed-step solver during early experimentation and switching to an adaptive solver only once the architecture and training procedure are validated is a practical strategy for keeping iteration speed high, since adaptive solvers can be substantially slower during the unstable early phase of training when the learned dynamics function is far from converged. Curriculum strategies that progressively extend the time horizon over which a Neural ODE is trained (rather than training on the full trajectory length from the start) help avoid the vanishing/exploding sensitivity problems that can occur when integrating an untrained, unstable dynamics function over a long time horizon. For both PINNs and Neural ODEs, validating against a known analytical or high-fidelity numerical solution on a simplified test case before scaling to the full problem of interest is essential, since silent failures (a network that trains to low loss but represents an incorrect or trivial solution) are common failure modes that are not always obvious from the loss curve alone.
## Benchmark Datasets & Evaluation
Canonical PINN benchmark problems include the Burgers' equation (a 1D nonlinear PDE combining advection and diffusion, popular because it develops sharp shock-like gradients that stress-test collocation-based solvers), the Navier-Stokes equations for incompressible flow (used to test PINNs on realistic fluid dynamics problems including flow around obstacles), and the Schrödinger equation (testing complex-valued PDE solutions). The PDEBench and PINNacle benchmark suites aggregate a broad range of PDE problems with standardized data splits and evaluation protocols, enabling systematic comparison of PINN variants, Fourier Neural Operators, and classical numerical baselines on relative L2 error against high-fidelity ground-truth simulations. For Neural ODEs, standard evaluation tasks include continuous-time extensions of classical sequence modeling benchmarks: irregularly sampled time series from the PhysioNet clinical dataset (testing Latent ODE and Neural CDE performance on real-world medical time series with missing observations), and synthetic dynamical systems benchmarks (e.g., learning a spiral ODE or a damped pendulum from noisy trajectory observations) used to validate basic correctness before scaling to more complex systems. Density estimation and generative modeling variants of Neural ODEs (Continuous Normalizing Flows) are evaluated on standard generative modeling benchmarks (CIFAR-10, ImageNet bits-per-dimension) alongside comparison to discrete normalizing flow architectures. Evaluation for both families typically reports relative L2 or L-infinity error against a trusted reference solution (either an analytical solution or a high-fidelity classical numerical solver), rather than a single aggregate task-accuracy metric as used in standard supervised learning.
## Key Challenges & Limitations
PINNs are notoriously difficult to train reliably: the multi-term loss landscape is often ill-conditioned, gradient pathologies between the data and PDE-residual terms are common, and training can converge to trivial or non-physical solutions that nonetheless achieve low overall loss, requiring careful diagnostic monitoring beyond the aggregate loss value. PINNs generally do not scale gracefully to very high-dimensional PDEs or very long time horizons, since collocation-based residual minimization over a spatiotemporal domain that grows combinatorially with dimension quickly becomes computationally infeasible, an issue sometimes called the curse of dimensionality for collocation methods. A standard PINN must be retrained from scratch for each new instance of boundary or initial conditions, which is a substantial practical limitation compared to classical numerical solvers (which can be re-run cheaply with new conditions) or operator-learning approaches like Fourier Neural Operators (which generalize across a family of conditions after one training run). Neural ODEs face their own distinct difficulties: training can be unstable when the learned dynamics become stiff, adaptive-step solvers can require an unpredictably large and slow-to-converge number of function evaluations during early training, and the adjoint method's backward-time reconstruction of the forward trajectory can accumulate numerical error in long or chaotic trajectories. Both families of methods currently lack the decades of convergence theory, error bounds, and stability guarantees that classical numerical methods (finite element, finite difference, spectral methods) provide, making them harder to trust in safety-critical engineering contexts without extensive empirical validation against classical solvers.
## Hyperparameter Tuning
The relative weighting of the data, PDE-residual, and boundary-condition loss terms in a PINN's objective is the single most sensitive hyperparameter, and adaptive reweighting methods (balancing terms based on their gradient norms with respect to shared network parameters) generally outperform manually fixed weights across a wide range of problems. Network width and depth for PINNs tend to favor moderately deep (4-8 layers), moderately wide (20-100 units) fully connected networks with smooth activation functions (tanh or sinusoidal activations rather than ReLU, since ReLU's discontinuous second derivative is poorly suited to computing the higher-order derivatives that PDE residuals require). The number and distribution of collocation points must scale with the complexity of the expected solution; problems with sharp gradients or shocks benefit from adaptive or residual-based resampling that concentrates points where the current residual is largest, rather than a fixed uniform grid chosen at the start of training. For Neural ODEs, solver tolerance settings (relative and absolute tolerance for adaptive-step solvers) directly trade off training speed against integration accuracy, with looser tolerances speeding up early training at the risk of insufficiently accurate gradients, and tightening tolerances as training progresses being a common practical schedule. The choice between the direct backpropagation-through-solver-operations approach and the adjoint sensitivity method for computing Neural ODE gradients should be guided by trajectory length and memory constraints: the adjoint method's constant memory cost becomes essential for long trajectories, while direct backpropagation is often more numerically robust and preferred for short trajectories where memory is not a binding constraint.
## Real-World Applications & Case Studies
In computational fluid dynamics, PINNs have been applied to reconstruct full velocity and pressure fields from sparse, noisy experimental measurements (e.g., particle image velocimetry data), effectively performing physics-constrained data assimilation that would be underdetermined without the governing Navier-Stokes constraints. In biomedical modeling, PINNs have been used to infer patient-specific cardiovascular parameters (e.g., arterial stiffness or blood flow characteristics) from limited, non-invasive clinical measurements by embedding known hemodynamic equations directly into the inference process. Latent ODEs and Neural CDEs have been applied to clinical time-series forecasting from electronic health records, where lab measurements and vital signs are recorded at irregular, patient-specific intervals that violate the fixed-timestep assumption of standard RNN and Transformer architectures. In materials science and molecular dynamics, Neural ODE-based approaches model continuous-time particle trajectories and energy landscapes, providing more physically consistent extrapolation than fixed-timestep discrete models when simulating over variable or adaptive time steps. Climate and weather modeling research has explored PINN and neural-operator hybrid approaches to accelerate expensive components of numerical weather prediction pipelines (e.g., learned parameterizations of subgrid-scale physical processes) while retaining consistency with known conservation laws for mass, momentum, and energy.
## Integration with Other Methods
PINNs and Neural ODEs are increasingly combined with operator-learning frameworks such as Fourier Neural Operators and DeepONets, using neural operators as fast surrogate models that can optionally be fine-tuned or corrected with PINN-style physics-residual losses to improve physical consistency of the operator's predictions on out-of-distribution inputs. Bayesian and uncertainty-quantification extensions combine PINNs with variational inference or ensemble methods to produce calibrated uncertainty estimates over the solution field, which is essential in engineering and scientific applications where knowing the confidence of a prediction is as important as the prediction itself. Neural ODEs form the continuous-time backbone of Continuous Normalizing Flows, connecting this line of work directly to the generative modeling and density estimation literature, including diffusion models, which share the conceptual framing of transforming a simple base distribution into a complex data distribution via a continuous-time process. Reinforcement learning for continuous control has explored Neural ODE-based dynamics models as differentiable, sample-efficient world models, integrating physics-informed structure into model-based RL pipelines to improve sample efficiency over black-box learned dynamics. Hybrid physics-ML pipelines increasingly use classical numerical solvers for well-understood, computationally cheap parts of a simulation while delegating expensive or poorly understood sub-components to learned PINN or neural-operator surrogates, a design pattern sometimes called "differentiable physics" that allows end-to-end gradient-based optimization across the hybrid pipeline.
## Future Research Directions
Improving the training reliability and convergence theory of PINNs remains a central open problem, with active research into loss landscape analysis, neural tangent kernel theory for PINNs, and principled automatic loss-balancing methods that reduce the current heavy reliance on empirical trial and error. Extending operator-learning approaches (FNOs, DeepONets) to handle irregular geometries, adaptive meshes, and multi-physics coupled systems (where multiple different PDEs interact, as in fluid-structure interaction) is an active area bridging classical computational science with modern deep learning. Scaling Neural ODEs and their variants to very high-dimensional state spaces (e.g., full 3D physical fields rather than low-dimensional latent states) while maintaining tractable adjoint-based training remains computationally challenging and is an active systems-and-algorithms research direction. Combining PINNs and Neural ODEs with foundation-model-style pretraining, learning general-purpose physical priors from large corpora of simulation data that can then be efficiently fine-tuned or adapted to new, specific physical systems with minimal additional data, is an emerging direction analogous to foundation models in vision and language. Finally, formal error bounds and stability guarantees comparable to those available for classical numerical methods remain an important open theoretical goal, particularly for safety-critical applications in aerospace, structural, and biomedical engineering where physics-informed neural methods are beginning to be considered as complements to, rather than replacements for, traditional certified numerical solvers.
## Summary & Key Takeaways
Physics-Informed Neural Networks and Neural ODEs both embed differential-equation structure into deep learning, but via distinct mechanisms: PINNs add a PDE-residual term to the training loss, computed via automatic differentiation at sampled collocation points, to enforce known governing physics alongside sparse observational data, while Neural ODEs parameterize a continuous-time dynamics function with a neural network and use differentiable numerical integration (with efficient adjoint-based gradients) in place of a fixed sequence of discrete layers. PINNs excel at solving specific instances of known PDEs from sparse or noisy data but must be retrained for each new boundary/initial condition, motivating operator-learning extensions like Fourier Neural Operators that generalize across families of conditions. Neural ODEs naturally handle continuous-time and irregularly sampled data, extending cleanly to Latent ODEs, Neural CDEs, and Continuous Normalizing Flows for generative modeling. Both families face real training and scalability challenges — ill-conditioned multi-term losses for PINNs, and stiff or unstable dynamics for Neural ODEs — and lack the mature convergence theory of classical numerical methods, positioning them for now as complements to, rather than wholesale replacements of, traditional numerical solvers in scientific and engineering applications.
---
## Appendix: Practical Labs
### Lab 1: PINN for the 1D Heat Equation
import torch
import torch.nn as nn
class PINN(nn.Module):
"""Simple fully connected network mapping (x, t) -> u(x, t), the
approximate solution to a PDE."""
def __init__(self, hidden_dim=32, n_layers=4):
super().__init__()
layers = [nn.Linear(2, hidden_dim), nn.Tanh()]
for _ in range(n_layers - 1):
layers += [nn.Linear(hidden_dim, hidden_dim), nn.Tanh()]
layers += [nn.Linear(hidden_dim, 1)]
self.net = nn.Sequential(*layers)
def forward(self, x, t):
return self.net(torch.cat([x, t], dim=1))
def heat_equation_residual(model, x, t, alpha=0.1):
"""Residual of the 1D heat equation: u_t - alpha * u_xx = 0,
computed via automatic differentiation."""
x = x.clone().requires_grad_(True)
t = t.clone().requires_grad_(True)
u = model(x, t)
u_t = torch.autograd.grad(u, t, grad_outputs=torch.ones_like(u), create_graph=True)[0]
u_x = torch.autograd.grad(u, x, grad_outputs=torch.ones_like(u), create_graph=True)[0]
u_xx = torch.autograd.grad(u_x, x, grad_outputs=torch.ones_like(u_x), create_graph=True)[0]
residual = u_t - alpha * u_xx
return residual
def test_pinn_heat_equation():
torch.manual_seed(0)
model = PINN(hidden_dim=16, n_layers=3)
optimizer = torch.optim.Adam(model.parameters(), lr=1e-3)
n_collocation = 200
n_epochs = 300
losses = []
for epoch in range(n_epochs):
optimizer.zero_grad()
# Collocation points sampled uniformly over the domain [0,1] x [0,1].
x_c = torch.rand(n_collocation, 1)
t_c = torch.rand(n_collocation, 1)
residual = heat_equation_residual(model, x_c, t_c)
pde_loss = (residual ** 2).mean()
# Initial condition: u(x, 0) = sin(pi * x).
x_ic = torch.rand(50, 1)
t_ic = torch.zeros(50, 1)
u_ic_pred = model(x_ic, t_ic)
u_ic_true = torch.sin(torch.pi * x_ic)
ic_loss = ((u_ic_pred - u_ic_true) ** 2).mean()
# Boundary conditions: u(0, t) = u(1, t) = 0.
t_bc = torch.rand(50, 1)
u_bc_left = model(torch.zeros(50, 1), t_bc)
u_bc_right = model(torch.ones(50, 1), t_bc)
bc_loss = (u_bc_left ** 2).mean() + (u_bc_right ** 2).mean()
loss = pde_loss + 10.0 * ic_loss + 10.0 * bc_loss
loss.backward()
optimizer.step()
losses.append(loss.item())
print(f"Initial loss: {losses[0]:.6f}")
print(f"Final loss: {losses[-1]:.6f}")
assert losses[-1] < losses[0], "PINN training loss should decrease"
print("PINN heat equation test passed.")
if __name__ == "__main__":
test_pinn_heat_equation()### Lab 2: Neural ODE with Manual Euler Integration
import torch
import torch.nn as nn
class ODEFunc(nn.Module):
"""Parameterizes the derivative dh/dt = f(h, t) with a small MLP."""
def __init__(self, hidden_dim=16):
super().__init__()
self.net = nn.Sequential(
nn.Linear(hidden_dim, hidden_dim),
nn.Tanh(),
nn.Linear(hidden_dim, hidden_dim),
)
def forward(self, h):
return self.net(h)
def euler_odesolve(func, h0, t0, t1, n_steps=20):
"""Fixed-step Euler integration of dh/dt = func(h) from t0 to t1."""
dt = (t1 - t0) / n_steps
h = h0
trajectory = [h0]
for _ in range(n_steps):
h = h + dt * func(h)
trajectory.append(h)
return h, torch.stack(trajectory)
def test_neural_ode_euler():
torch.manual_seed(0)
hidden_dim = 4
func = ODEFunc(hidden_dim)
h0 = torch.randn(1, hidden_dim)
h_final, trajectory = euler_odesolve(func, h0, t0=0.0, t1=1.0, n_steps=50)
print(f"Initial state: {h0.detach().numpy().flatten()}")
print(f"Final state: {h_final.detach().numpy().flatten()}")
print(f"Trajectory shape: {trajectory.shape}")
assert trajectory.shape == (51, 1, hidden_dim), "Trajectory should have n_steps+1 points"
assert not torch.allclose(h0, h_final), "State should evolve under nontrivial dynamics"
# Gradients should flow through the full integration for training.
loss = h_final.sum()
loss.backward()
grad_norm = sum(p.grad.norm().item() for p in func.parameters() if p.grad is not None)
assert grad_norm > 0, "Gradients should flow back through the Euler-integrated trajectory"
print(f"Total gradient norm through ODE solve: {grad_norm:.4f}")
print("Neural ODE Euler integration test passed.")
if __name__ == "__main__":
test_neural_ode_euler()### Lab 3: Adjoint Sensitivity Method (Simplified Demonstration)
import torch
import torch.nn as nn
class SimpleDynamics(nn.Module):
def __init__(self, dim=4):
super().__init__()
self.linear = nn.Linear(dim, dim, bias=False)
def forward(self, h):
return self.linear(h)
def forward_euler(func, h0, t0, t1, n_steps):
dt = (t1 - t0) / n_steps
h = h0
states = [h0]
for _ in range(n_steps):
h = h + dt * func(h)
states.append(h)
return states
def adjoint_backward_euler(func, states, dt, grad_output):
"""Manually propagate the adjoint state a(t) backward through the
recorded trajectory, computing parameter gradients without needing to
store the full autograd graph of the forward pass (illustrative, not
a production-grade adjoint implementation)."""
adjoint = grad_output.clone()
param_grads = [torch.zeros_like(p) for p in func.parameters()]
for h in reversed(states[:-1]):
h = h.detach().requires_grad_(True)
out = func(h)
# d(adjoint . out)/dh and d(adjoint . out)/dparams via autograd,
# emulating the adjoint ODE's right-hand side.
grads = torch.autograd.grad(out, [h] + list(func.parameters()),
grad_outputs=adjoint, retain_graph=False)
dh_grad, *p_grads = grads
for pg, g in zip(param_grads, p_grads):
pg += dt * g
adjoint = adjoint + dt * dh_grad # backward Euler step for the adjoint
return adjoint, param_grads
def test_adjoint_method():
torch.manual_seed(0)
dim = 4
func = SimpleDynamics(dim)
h0 = torch.randn(1, dim)
n_steps = 20
states = forward_euler(func, h0, t0=0.0, t1=1.0, n_steps=n_steps)
dt = 1.0 / n_steps
# Suppose the loss is the sum of the final state.
grad_output = torch.ones_like(states[-1])
adjoint_final, param_grads = adjoint_backward_euler(func, states, dt, grad_output)
# Cross-check against direct autograd through the full forward trajectory.
h = h0.clone().requires_grad_(True)
h_direct = h
for _ in range(n_steps):
h_direct = h_direct + dt * func(h_direct)
loss = h_direct.sum()
loss.backward()
direct_grad = func.linear.weight.grad
manual_grad = param_grads[0]
rel_error = (manual_grad - direct_grad).norm() / (direct_grad.norm() + 1e-8)
print(f"Adjoint-computed gradient norm: {manual_grad.norm().item():.4f}")
print(f"Direct autograd gradient norm: {direct_grad.norm().item():.4f}")
print(f"Relative error: {rel_error.item():.6f}")
assert rel_error.item() < 0.1, "Adjoint-based gradient should approximately match direct autograd"
print("Adjoint sensitivity method test passed.")
if __name__ == "__main__":
test_adjoint_method()### Lab 4: Learning an Unknown Dynamical System with a Neural ODE
import torch
import torch.nn as nn
class LearnedDynamics(nn.Module):
def __init__(self, dim=2, hidden_dim=32):
super().__init__()
self.net = nn.Sequential(
nn.Linear(dim, hidden_dim),
nn.Tanh(),
nn.Linear(hidden_dim, dim),
)
def forward(self, h):
return self.net(h)
def true_dynamics(h):
"""Ground-truth dynamics: a simple rotation, dh/dt = A h, with A a
skew-symmetric matrix producing circular trajectories."""
A = torch.tensor([[0.0, -1.0], [1.0, 0.0]])
return h @ A.T
def rk4_step(func, h, dt):
k1 = func(h)
k2 = func(h + 0.5 * dt * k1)
k3 = func(h + 0.5 * dt * k2)
k4 = func(h + dt * k3)
return h + (dt / 6.0) * (k1 + 2 * k2 + 2 * k3 + k4)
def rollout(func, h0, dt, n_steps):
h = h0
traj = [h0]
for _ in range(n_steps):
h = rk4_step(func, h, dt)
traj.append(h)
return torch.stack(traj)
def test_learn_dynamics():
torch.manual_seed(0)
dt, n_steps = 0.1, 30
h0 = torch.tensor([[1.0, 0.0]])
with torch.no_grad():
true_traj = rollout(true_dynamics, h0, dt, n_steps)
model = LearnedDynamics(dim=2, hidden_dim=32)
optimizer = torch.optim.Adam(model.parameters(), lr=1e-2)
losses = []
for epoch in range(200):
optimizer.zero_grad()
pred_traj = rollout(model, h0, dt, n_steps)
loss = ((pred_traj - true_traj) ** 2).mean()
loss.backward()
optimizer.step()
losses.append(loss.item())
print(f"Initial trajectory MSE: {losses[0]:.4f}")
print(f"Final trajectory MSE: {losses[-1]:.4f}")
assert losses[-1] < losses[0] * 0.5, "Neural ODE should substantially fit the true rotational dynamics"
print("Learned dynamics Neural ODE test passed.")
if __name__ == "__main__":
test_learn_dynamics()