القائمة
Physics-informed neural networks, neural ODEs, and data-driven force-field learning using PyTorch, DeepXDE, and torchdiffeq.
يُحدَّد Prompt الذي يراجع المصدر أولًا بشكل افتراضي. يمكنك التبديل إلى أمر مباشر أو تنزيل نسخة محلية.
اقرأ SKILL.md وأي ملفات مرافقة يعرضها SkillsMP قبل أن تقرر التثبيت.
التثبيت باستخدام Codex أو Claude انسخ هذا Prompt والصقه في Codex أو Claude أو مساعد آخر ليراجع صفحة Skill ويثبّتها لك.
يتجاوز الأمر المباشر Prompt المخصّص للمراجعة. افحص المصدر قبل تشغيله.
npx skills add https://github.com/xjtulyc/awesome-rosetta-skills --skill pytorch-physicsيبقى الأمر في سطر واحد. مرّر أفقيًا لمراجعته كاملًا قبل النسخ.
تفضّل نسخة محلية؟ نزّل الملفات المتاحة حاليًا لدى SkillsMP.
استنادا إلى تصنيف SOC المهني
عرض SKILL.md
| name | pytorch-physics |
| description | Physics-informed neural networks, neural ODEs, and data-driven force-field learning using PyTorch, DeepXDE, and torchdiffeq. |
| tags | ["physics","pytorch","pinn","neural-ode","scientific-computing","deep-learning"] |
| version | 1.0.0 |
| authors | [{"name":"Rosetta Skills Contributors","github":"@xjtulyc"}] |
| license | MIT |
| platforms | ["claude-code","codex","gemini-cli","cursor"] |
| dependencies | ["torch>=2.0","deepxde>=1.9","torchdiffeq>=0.2","numpy>=1.24","matplotlib>=3.7"] |
| last_updated | 2026-03-17 |
| status | stable |
This skill covers the full workflow for applying deep learning to physical systems: Physics-Informed Neural Networks (PINNs) that embed differential equations as loss terms, Neural Ordinary Differential Equations (Neural ODEs) for continuous-time dynamics, and data-driven force-field learning for molecular or mechanical simulations.
Use this skill when you need to:
Do not use this skill for purely statistical forecasting where no physics is known, or when a classical numerical solver is already fast enough and no generalization is required.
A PINN approximates the solution u(x, t) of a PDE with a neural network f_θ. The training loss has three components:
L = L_data + λ_r * L_residual + λ_bc * L_boundary
Because PyTorch tracks the full computation graph, ∂f/∂x and ∂²f/∂x² are computed
exactly with torch.autograd.grad, not finite differences.
A Neural ODE replaces the discrete update rule of a ResNet with a continuous-time ODE:
dz/dt = f_θ(z, t)
The hidden state z(t) is integrated by a black-box ODE solver (e.g. Dormand-Prince RK45).
Gradients flow back through the solver via the adjoint method, keeping memory O(1).
torchdiffeq provides odeint and odeint_adjoint for this purpose.
A force field maps atomic positions {r_i} to potential energy E and forces F = -∇E. Training a neural network to predict E from a descriptor and back-propagating to get F ensures energy conservation by construction. This is the foundation of NequIP, SchNet, etc.
| Operation | Code |
|---|---|
| First derivative | torch.autograd.grad(u, x, create_graph=True) |
| Second derivative | nested autograd.grad calls |
| Jacobian | torch.func.jacrev (functorch) |
| Hessian | torch.func.hessian |
pip install torch>=2.0 deepxde>=1.9 torchdiffeq>=0.2 numpy>=1.24 matplotlib>=3.7
For GPU support (CUDA 12):
pip install torch --index-url https://download.pytorch.org/whl/cu121
import torch
import deepxde as dde
import torchdiffeq
print("PyTorch:", torch.__version__)
print("DeepXDE:", dde.__version__)
print("CUDA available:", torch.cuda.is_available())
# Optional: weights & biases logging
export WANDB_API_KEY="<paste-your-key>"
# Optional: custom data directory
export PHYSICS_DATA_DIR="/data/physics_datasets"
import os
wandb_key = os.getenv("WANDB_API_KEY", "")
data_dir = os.getenv("PHYSICS_DATA_DIR", "./data")
# DeepXDE supports multiple backends; set to pytorch
export DDE_BACKEND=pytorch
The problem: ∂u/∂t = α ∂²u/∂x², x ∈ [0, 1], t ∈ [0, 1] with u(x, 0) = sin(πx), u(0, t) = u(1, t) = 0. Analytical solution: u(x, t) = exp(-α π² t) sin(πx).
import numpy as np
import torch
import torch.nn as nn
import matplotlib.pyplot as plt
# ── Hyper-parameters ────────────────────────────────────────────────────────
ALPHA = 0.01 # thermal diffusivity
N_COLLOC = 5000 # interior collocation points
N_BC = 200 # boundary / IC points
N_EPOCHS = 10000
LR = 1e-3
LAMBDA_R = 1.0 # residual weight
LAMBDA_BC = 10.0 # boundary weight
DEVICE = "cuda" if torch.cuda.is_available() else "cpu"
# ── Network ──────────────────────────────────────────────────────────────────
class PINN(nn.Module):
def __init__(self, layers=(2, 64, 64, 64, 1)):
super().__init__()
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, x: torch.Tensor, t: torch.Tensor) -> torch.Tensor:
xt = torch.cat([x, t], dim=1)
return self.net(xt)
def pde_residual(model, x, t, alpha=ALPHA):
"""Return PDE residual: u_t - alpha * u_xx."""
x.requires_grad_(True)
t.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]
return u_t - alpha * u_xx
def sample_collocation(n):
x = torch.rand(n, 1, device=DEVICE)
t = torch.rand(n, 1, device=DEVICE)
return x, t
def sample_boundary(n):
"""Sample IC (t=0) and Dirichlet BC (x=0, x=1)."""
n3 = n // 3
# Initial condition
x_ic = torch.rand(n3, 1, device=DEVICE)
t_ic = torch.zeros(n3, 1, device=DEVICE)
u_ic = torch.sin(torch.pi * x_ic)
# Left BC
x_l = torch.zeros(n3, 1, device=DEVICE)
t_l = torch.rand(n3, 1, device=DEVICE)
u_l = torch.zeros(n3, 1, device=DEVICE)
# Right BC
x_r = torch.ones(n3, 1, device=DEVICE)
t_r = torch.rand(n3, 1, device=DEVICE)
u_r = torch.zeros(n3, 1, device=DEVICE)
xb = torch.cat([x_ic, x_l, x_r])
tb = torch.cat([t_ic, t_l, t_r])
ub = torch.cat([u_ic, u_l, u_r])
return xb, tb, ub
# ── Training loop ────────────────────────────────────────────────────────────
model = PINN().to(DEVICE)
optimizer = torch.optim.Adam(model.parameters(), lr=LR)
scheduler = torch.optim.lr_scheduler.CosineAnnealingLR(optimizer, N_EPOCHS)
losses = []
for epoch in range(N_EPOCHS):
optimizer.zero_grad()
# Residual loss
xc, tc = sample_collocation(N_COLLOC)
res = pde_residual(model, xc, tc)
loss_r = LAMBDA_R * (res ** 2).mean()
# Boundary loss
xb, tb, ub = sample_boundary(N_BC)
u_pred_b = model(xb, tb)
loss_bc = LAMBDA_BC * ((u_pred_b - ub) ** 2).mean()
loss = loss_r + loss_bc
loss.backward()
optimizer.step()
scheduler.step()
if epoch % 500 == 0:
losses.append(loss.item())
print(f"Epoch {epoch:5d} | loss={loss.item():.4e} "
f"res={loss_r.item():.4e} bc={loss_bc.item():.4e}")
# ── Evaluation ───────────────────────────────────────────────────────────────
x_test = torch.linspace(0, 1, 100, device=DEVICE).unsqueeze(1)
t_test = torch.full_like(x_test, 0.5)
with torch.no_grad():
u_pred = model(x_test, t_test).cpu().numpy().ravel()
u_exact = np.exp(-ALPHA * np.pi**2 * 0.5) * np.sin(np.pi * x_test.cpu().numpy().ravel())
plt.figure(figsize=(8, 4))
plt.plot(x_test.cpu().numpy(), u_exact, label="Exact", lw=2)
plt.plot(x_test.cpu().numpy(), u_pred, "--", label="PINN", lw=2)
plt.xlabel("x"); plt.ylabel("u(x, 0.5)")
plt.title("Heat equation — PINN vs exact at t=0.5")
plt.legend(); plt.tight_layout(); plt.savefig("heat_pinn.png", dpi=150)
print("L2 error:", np.linalg.norm(u_pred - u_exact) / np.linalg.norm(u_exact))
Fit a Neural ODE to noisy observations of a damped harmonic oscillator and then integrate forward in time to forecast unseen states.
import torch
import torch.nn as nn
from torchdiffeq import odeint_adjoint as odeint
import numpy as np
import matplotlib.pyplot as plt
DEVICE = "cuda" if torch.cuda.is_available() else "cpu"
# ── Generate ground-truth data ───────────────────────────────────────────────
def true_dynamics(t, y):
"""Damped oscillator: y = [x, v], dy/dt = [v, -2ζω v - ω² x]."""
omega, zeta = 2.0, 0.1
x, v = y[..., 0:1], y[..., 1:2]
dxdt = v
dvdt = -2 * zeta * omega * v - omega**2 * x
return torch.cat([dxdt, dvdt], dim=-1)
t_train = torch.linspace(0, 5, 60, device=DEVICE)
y0_true = torch.tensor([[1.0, 0.0]], device=DEVICE)
with torch.no_grad():
y_true = odeint(true_dynamics, y0_true, t_train, method="rk4") # (T, 1, 2)
noise = 0.05 * torch.randn_like(y_true)
y_obs = (y_true + noise).squeeze(1) # (T, 2)
# ── Neural ODE model ─────────────────────────────────────────────────────────
class ODEFunc(nn.Module):
def __init__(self, state_dim=2, hidden=64):
super().__init__()
self.net = nn.Sequential(
nn.Linear(state_dim, hidden), nn.SiLU(),
nn.Linear(hidden, hidden), nn.SiLU(),
nn.Linear(hidden, state_dim),
)
def forward(self, t, y):
return self.net(y)
class NeuralODE(nn.Module):
def __init__(self):
super().__init__()
self.func = ODEFunc()
def forward(self, y0, t_span):
return odeint(self.func, y0, t_span, method="dopri5",
rtol=1e-4, atol=1e-6)
# ── Training ─────────────────────────────────────────────────────────────────
node = NeuralODE().to(DEVICE)
optimizer = torch.optim.Adam(node.parameters(), lr=5e-3)
y0_est = y_obs[0:1] # use first observation as initial condition
for epoch in range(2000):
optimizer.zero_grad()
y_pred = node(y0_est, t_train) # (T, 1, 2)
loss = ((y_pred.squeeze(1) - y_obs) ** 2).mean()
loss.backward()
optimizer.step()
if epoch % 200 == 0:
print(f"Epoch {epoch:4d} | MSE={loss.item():.5f}")
# ── Forecasting ──────────────────────────────────────────────────────────────
t_fore = torch.linspace(0, 10, 120, device=DEVICE)
with torch.no_grad():
y_fore = node(y0_est, t_fore).squeeze(1).cpu().numpy()
y_true_np = y_true.squeeze(1).cpu().numpy()
t_np = t_train.cpu().numpy()
fig, axes = plt.subplots(1, 2, figsize=(12, 4))
for i, label in enumerate(["position x", "velocity v"]):
axes[i].scatter(t_np, y_obs[:, i].cpu().numpy(), s=10, label="observations")
axes[i].plot(t_fore.cpu().numpy(), y_fore[:, i], label="Neural ODE forecast")
axes[i].axvline(5, ls="--", color="gray", label="forecast boundary")
axes[i].set_xlabel("time"); axes[i].set_ylabel(label); axes[i].legend()
plt.tight_layout(); plt.savefig("neural_ode_oscillator.png", dpi=150)
Train a simple equivariant-inspired force network where forces are derived from a learned potential energy to ensure conservation.
import torch
import torch.nn as nn
import numpy as np
import matplotlib.pyplot as plt
DEVICE = "cuda" if torch.cuda.is_available() else "cpu"
# ── Synthetic dataset: 2-body Lennard-Jones ──────────────────────────────────
def lj_potential(r, eps=1.0, sigma=1.0):
"""Lennard-Jones 6-12 potential energy."""
sr6 = (sigma / r) ** 6
return 4 * eps * (sr6**2 - sr6)
def generate_lj_data(n=2000):
r = torch.FloatTensor(n, 1).uniform_(0.9, 3.0)
E = lj_potential(r)
F = -torch.autograd.functional.jacobian(
lambda x: lj_potential(x).sum(), r
).squeeze() # dE/dr; shape (n,)
# pack as (n, 1) for convenient batching
return r, E, F.unsqueeze(1)
r_data, E_data, F_data = generate_lj_data()
r_data = r_data.to(DEVICE)
E_data = E_data.to(DEVICE)
F_data = F_data.to(DEVICE)
# ── Potential network ─────────────────────────────────────────────────────────
class PotentialNet(nn.Module):
"""Predict scalar E(r); forces are F = -dE/dr computed at inference."""
def __init__(self, hidden=128):
super().__init__()
self.net = nn.Sequential(
nn.Linear(1, hidden), nn.SiLU(),
nn.Linear(hidden, hidden), nn.SiLU(),
nn.Linear(hidden, hidden), nn.SiLU(),
nn.Linear(hidden, 1),
)
def forward(self, r):
return self.net(r)
def energy_and_force(self, r):
r = r.requires_grad_(True)
E = self.net(r)
F = -torch.autograd.grad(E.sum(), r, create_graph=True)[0]
return E, F
net = PotentialNet().to(DEVICE)
optimizer = torch.optim.Adam(net.parameters(), lr=1e-3)
LAMBDA_F = 0.9 # weight for force matching
BATCH = 256
dataset = torch.utils.data.TensorDataset(r_data, E_data, F_data)
loader = torch.utils.data.DataLoader(dataset, batch_size=BATCH, shuffle=True)
for epoch in range(200):
epoch_loss = 0.0
for r_b, E_b, F_b in loader:
optimizer.zero_grad()
E_pred, F_pred = net.energy_and_force(r_b)
loss_E = ((E_pred - E_b) ** 2).mean()
loss_F = ((F_pred - F_b) ** 2).mean()
loss = (1 - LAMBDA_F) * loss_E + LAMBDA_F * loss_F
loss.backward()
optimizer.step()
epoch_loss += loss.item()
if epoch % 20 == 0:
print(f"Epoch {epoch:3d} | loss={epoch_loss/len(loader):.5f}")
# ── Evaluation ────────────────────────────────────────────────────────────────
r_test = torch.linspace(0.95, 3.0, 200, device=DEVICE).unsqueeze(1)
with torch.no_grad():
# Temporarily enable grad for force calculation
pass
r_test.requires_grad_(True)
E_test, F_test = net.energy_and_force(r_test)
r_np = r_test.detach().cpu().numpy().ravel()
E_np = E_test.detach().cpu().numpy().ravel()
F_np = F_test.detach().cpu().numpy().ravel()
E_true = lj_potential(r_test.detach()).cpu().numpy().ravel()
fig, (ax1, ax2) = plt.subplots(1, 2, figsize=(12, 4))
ax1.plot(r_np, E_true, label="LJ exact"); ax1.plot(r_np, E_np, "--", label="NN")
ax1.set_ylim(-2, 5); ax1.set_xlabel("r"); ax1.set_ylabel("E"); ax1.legend()
ax2.plot(r_np, F_np, label="NN force")
ax2.set_xlabel("r"); ax2.set_ylabel("F = -dE/dr"); ax2.legend()
plt.tight_layout(); plt.savefig("force_field.png", dpi=150)
Given noisy measurements, identify the unknown parameter α in u_t = α u_xx.
import torch
import torch.nn as nn
DEVICE = "cuda" if torch.cuda.is_available() else "cpu"
class InversePINN(nn.Module):
def __init__(self):
super().__init__()
self.net = nn.Sequential(
nn.Linear(2, 64), nn.Tanh(),
nn.Linear(64, 64), nn.Tanh(),
nn.Linear(64, 1),
)
# alpha is a learnable parameter
self.log_alpha = nn.Parameter(torch.tensor(0.0))
@property
def alpha(self):
return torch.exp(self.log_alpha)
def forward(self, x, t):
return self.net(torch.cat([x, t], 1))
def residual(self, x, t):
x = x.requires_grad_(True)
t = t.requires_grad_(True)
u = self(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]
return u_t - self.alpha * u_xx
def train_inverse(alpha_true=0.02, n_data=200, n_colloc=3000, epochs=8000):
model = InversePINN().to(DEVICE)
optimizer = torch.optim.Adam(model.parameters(), lr=1e-3)
# Generate synthetic noisy data from analytical solution
import numpy as np
x_d = torch.rand(n_data, 1, device=DEVICE)
t_d = torch.rand(n_data, 1, device=DEVICE)
u_d = (torch.exp(torch.tensor(-alpha_true) * np.pi**2 * t_d)
* torch.sin(np.pi * x_d)
+ 0.01 * torch.randn(n_data, 1, device=DEVICE))
for ep in range(epochs):
optimizer.zero_grad()
# Data loss
u_pred = model(x_d, t_d)
loss_data = ((u_pred - u_d) ** 2).mean()
# PDE residual loss
xc = torch.rand(n_colloc, 1, device=DEVICE)
tc = torch.rand(n_colloc, 1, device=DEVICE)
res = model.residual(xc, tc)
loss_res = (res ** 2).mean()
loss = loss_data + loss_res
loss.backward()
optimizer.step()
if ep % 1000 == 0:
print(f"Epoch {ep:5d} | loss={loss.item():.4e} "
f"alpha_est={model.alpha.item():.5f} (true={alpha_true})")
return model
trained = train_inverse()
print("Final alpha estimate:", trained.alpha.item())
DeepXDE provides a high-level API for common PDE problem types.
import deepxde as dde
import numpy as np
import os
# Set backend
os.environ.setdefault("DDE_BACKEND", "pytorch")
# Burgers' equation: u_t + u * u_x = nu * u_xx
nu = 0.01 / np.pi
def burgers_pde(x, y):
"""x: (N,2) [x_coord, t]; y: (N,1) [u]"""
dy_x = dde.grad.jacobian(y, x, i=0, j=0)
dy_t = dde.grad.jacobian(y, x, i=0, j=1)
dy_xx = dde.grad.hessian(y, x, i=0, j=0)
return dy_t + y * dy_x - nu * dy_xx
geom = dde.geometry.Interval(-1, 1)
timedomain = dde.geometry.TimeDomain(0, 0.99)
geomtime = dde.geometry.GeometryXTime(geom, timedomain)
bc = dde.icbc.DirichletBC(geomtime, lambda x: 0, lambda x, on_b: on_b)
ic = dde.icbc.IC(geomtime, lambda x: -np.sin(np.pi * x[:, 0:1]),
lambda x, on_i: on_i)
data = dde.data.TimePDE(
geomtime, burgers_pde, [bc, ic],
num_domain=2500, num_boundary=100, num_initial=200,
)
net = dde.nn.FNN([2] + [64] * 4 + [1], "tanh", "Glorot normal")
model = dde.Model(data, net)
model.compile("adam", lr=1e-3)
losshistory, train_state = model.train(iterations=15000)
dde.saveplot(losshistory, train_state, issave=False, isplot=True)
Conserve energy exactly by parameterizing the Hamiltonian H(q, p).
import torch
import torch.nn as nn
from torchdiffeq import odeint
DEVICE = "cuda" if torch.cuda.is_available() else "cpu"
class HamiltonianNet(nn.Module):
"""Learn H(q, p) and derive canonical equations from it."""
def __init__(self, dim=1, hidden=128):
super().__init__()
self.H = nn.Sequential(
nn.Linear(2 * dim, hidden), nn.Tanh(),
nn.Linear(hidden, hidden), nn.Tanh(),
nn.Linear(hidden, 1),
)
self.dim = dim
def hamiltonian(self, qp):
return self.H(qp)
def forward(self, t, qp):
"""Returns dq/dt = dH/dp, dp/dt = -dH/dq."""
qp = qp.requires_grad_(True)
H_val = self.H(qp).sum()
grad = torch.autograd.grad(H_val, qp, create_graph=True)[0]
dqdt = grad[..., self.dim:] # dH/dp
dpdt = -grad[..., :self.dim] # -dH/dq
return torch.cat([dqdt, dpdt], dim=-1)
def train_hnn(n_traj=50, t_span=torch.linspace(0, 2, 40)):
"""Train on simple pendulum trajectories."""
model = HamiltonianNet(dim=1).to(DEVICE)
optimizer = torch.optim.Adam(model.parameters(), lr=1e-3)
t_span = t_span.to(DEVICE)
# Generate ground truth pendulum: H = 0.5 p^2 - cos(q)
def pendulum(t, qp):
q, p = qp[..., 0:1], qp[..., 1:2]
return torch.cat([p, -torch.sin(q)], dim=-1)
for step in range(3000):
# Random initial conditions
q0 = torch.FloatTensor(1, 1).uniform_(-2, 2).to(DEVICE)
p0 = torch.FloatTensor(1, 1).uniform_(-2, 2).to(DEVICE)
qp0 = torch.cat([q0, p0], dim=-1)
with torch.no_grad():
qp_true = odeint(pendulum, qp0, t_span, method="rk4")
qp_pred = odeint(model, qp0, t_span, method="rk4")
loss = ((qp_pred - qp_true) ** 2).mean()
optimizer.zero_grad(); loss.backward(); optimizer.step()
if step % 300 == 0:
print(f"Step {step:4d} | loss={loss.item():.5f}")
return model
Symptom: loss.backward() raises a warning and gradients are zero.
Cause: Using odeint (non-adjoint) with too many time steps can exhaust memory
or detach the graph.
Fix: Switch to odeint_adjoint which uses the adjoint method and does not keep
the full trajectory in memory.
from torchdiffeq import odeint_adjoint as odeint
# use exactly the same call signature as odeint
Symptom: loss_r oscillates and does not decrease after several thousand epochs.
Cause 1: Learning rate too high — the residual term involves second derivatives which amplify gradient magnitudes.
Fix: Lower LR to 1e-4 and add gradient clipping.
torch.nn.utils.clip_grad_norm_(model.parameters(), max_norm=1.0)
Cause 2: Collocation points are not covering the domain evenly.
Fix: Use Latin-hypercube sampling instead of pure random.
from scipy.stats.qmc import LatinHypercube
sampler = LatinHypercube(d=2)
pts = torch.tensor(sampler.random(n=N_COLLOC), dtype=torch.float32, device=DEVICE)
xc, tc = pts[:, 0:1], pts[:, 1:2]
Symptom: ImportError or wrong backend active.
Fix:
export DDE_BACKEND=pytorch
python -c "import deepxde; print(deepxde.backend.backend_name)"
Symptom: Loss becomes NaN after a few iterations when using create_graph=True.
Cause: Exploding gradients through the double backward pass.
Fix: Add gradient clipping and ensure input distances are bounded away from zero.
r = torch.clamp(r, min=0.8) # avoid singularity at r → 0
torch.nn.utils.clip_grad_norm_(net.parameters(), 5.0)
Solve the steady incompressible Navier-Stokes equations inside a unit square cavity where the top wall moves at velocity U=1. Re = 100.
import torch
import torch.nn as nn
import numpy as np
import matplotlib.pyplot as plt
DEVICE = "cuda" if torch.cuda.is_available() else "cpu"
RE = 100.0
class NSNet(nn.Module):
"""Predict (u, v, p) from (x, y)."""
def __init__(self, hidden=128, depth=6):
super().__init__()
layers = [nn.Linear(2, hidden), nn.Tanh()]
for _ in range(depth - 1):
layers += [nn.Linear(hidden, hidden), nn.Tanh()]
layers.append(nn.Linear(hidden, 3))
self.net = nn.Sequential(*layers)
def forward(self, xy):
return self.net(xy)
def ns_residuals(model, xy):
xy = xy.requires_grad_(True)
uvp = model(xy)
u, v, p = uvp[:, 0:1], uvp[:, 1:2], uvp[:, 2:3]
def grad1(f, x):
return torch.autograd.grad(f, x, torch.ones_like(f), create_graph=True)[0]
uvp_x = grad1(uvp, xy)[:, 0:1], grad1(uvp, xy)[:, 1:2], grad1(uvp, xy)[:, 2:3]
# recompute per-component
u_xy = torch.autograd.grad(u, xy, torch.ones_like(u), create_graph=True)[0]
v_xy = torch.autograd.grad(v, xy, torch.ones_like(v), create_graph=True)[0]
p_xy = torch.autograd.grad(p, xy, torch.ones_like(p), create_graph=True)[0]
ux, uy = u_xy[:, 0:1], u_xy[:, 1:2]
vx, vy = v_xy[:, 0:1], v_xy[:, 1:2]
px, py = p_xy[:, 0:1], p_xy[:, 1:2]
uxx = torch.autograd.grad(ux, xy, torch.ones_like(ux), create_graph=True)[0][:, 0:1]
uyy = torch.autograd.grad(uy, xy, torch.ones_like(uy), create_graph=True)[0][:, 1:2]
vxx = torch.autograd.grad(vx, xy, torch.ones_like(vx), create_graph=True)[0][:, 0:1]
vyy = torch.autograd.grad(vy, xy, torch.ones_like(vy), create_graph=True)[0][:, 1:2]
# Continuity
cont = ux + vy
# Momentum x
mom_x = u * ux + v * uy + px - (uxx + uyy) / RE
# Momentum y
mom_y = u * vx + v * vy + py - (vxx + vyy) / RE
return cont, mom_x, mom_y
def sample_boundary_ns(n):
"""Return (xy, u_target, v_target) for all four walls."""
n4 = n // 4
# Bottom y=0: u=v=0
bot = torch.cat([torch.rand(n4, 1), torch.zeros(n4, 1)], 1).to(DEVICE)
# Top y=1: u=1, v=0
top = torch.cat([torch.rand(n4, 1), torch.ones(n4, 1)], 1).to(DEVICE)
# Left x=0: u=v=0
left = torch.cat([torch.zeros(n4, 1), torch.rand(n4, 1)], 1).to(DEVICE)
# Right x=1: u=v=0
right = torch.cat([torch.ones(n4, 1), torch.rand(n4, 1)], 1).to(DEVICE)
xy_bc = torch.cat([bot, top, left, right])
u_bc = torch.zeros(4 * n4, 1, device=DEVICE)
u_bc[n4:2*n4] = 1.0 # top wall
v_bc = torch.zeros(4 * n4, 1, device=DEVICE)
return xy_bc, u_bc, v_bc
model = NSNet().to(DEVICE)
optimizer = torch.optim.Adam(model.parameters(), lr=1e-3)
for epoch in range(5000):
optimizer.zero_grad()
# Interior residuals
xy_int = torch.rand(2000, 2, device=DEVICE)
cont, mom_x, mom_y = ns_residuals(model, xy_int)
loss_pde = (cont**2 + mom_x**2 + mom_y**2).mean()
# Boundary
xy_bc, u_bc, v_bc = sample_boundary_ns(400)
uvp_bc = model(xy_bc)
loss_bc = ((uvp_bc[:, 0:1] - u_bc)**2 + (uvp_bc[:, 1:2] - v_bc)**2).mean()
loss = loss_pde + 10 * loss_bc
loss.backward()
optimizer.step()
if epoch % 500 == 0:
print(f"Epoch {epoch:4d} | total={loss.item():.4e}")
# Visualise u-velocity field
nx = 50
x_ = torch.linspace(0, 1, nx); y_ = torch.linspace(0, 1, nx)
X, Y = torch.meshgrid(x_, y_, indexing="ij")
xy_grid = torch.stack([X.ravel(), Y.ravel()], 1).to(DEVICE)
with torch.no_grad():
uvp_grid = model(xy_grid).cpu().numpy()
U = uvp_grid[:, 0].reshape(nx, nx)
plt.figure(figsize=(6, 5))
plt.contourf(X.numpy(), Y.numpy(), U, levels=30, cmap="RdBu_r")
plt.colorbar(label="u velocity")
plt.title(f"Lid-driven cavity Re={RE} — PINN u-field")
plt.tight_layout(); plt.savefig("ns_cavity.png", dpi=150)
Encode a noisy trajectory with an RNN, decode the latent ODE, and reconstruct the full trajectory including missing segments.
import torch
import torch.nn as nn
from torchdiffeq import odeint_adjoint as odeint
import numpy as np
import matplotlib.pyplot as plt
DEVICE = "cuda" if torch.cuda.is_available() else "cpu"
LATENT_DIM = 8
OBS_DIM = 2
class LatentODEFunc(nn.Module):
def __init__(self, latent_dim=LATENT_DIM, hidden=64):
super().__init__()
self.net = nn.Sequential(
nn.Linear(latent_dim, hidden), nn.ELU(),
nn.Linear(hidden, hidden), nn.ELU(),
nn.Linear(hidden, latent_dim),
)
def forward(self, t, z):
return self.net(z)
class Encoder(nn.Module):
"""GRU encoder: maps a variable-length observation sequence to z0."""
def __init__(self, obs_dim=OBS_DIM, hidden=32, latent_dim=LATENT_DIM):
super().__init__()
self.gru = nn.GRU(obs_dim + 1, hidden, batch_first=True)
self.out = nn.Linear(hidden, latent_dim * 2) # mean + log-var
def forward(self, obs, times):
# Concatenate time deltas to observations
dt = torch.diff(times, prepend=times[:, :1], dim=1).unsqueeze(-1)
inp = torch.cat([obs, dt], dim=-1)
_, h = self.gru(inp)
params = self.out(h.squeeze(0))
mu, log_var = params.chunk(2, dim=-1)
return mu, log_var
class Decoder(nn.Module):
def __init__(self, latent_dim=LATENT_DIM, obs_dim=OBS_DIM, hidden=64):
super().__init__()
self.net = nn.Sequential(
nn.Linear(latent_dim, hidden), nn.ReLU(),
nn.Linear(hidden, obs_dim),
)
def forward(self, z):
return self.net(z)
class LatentODE(nn.Module):
def __init__(self):
super().__init__()
self.encoder = Encoder()
self.ode_func = LatentODEFunc()
self.decoder = Decoder()
def forward(self, obs, t_obs, t_pred):
mu, log_var = self.encoder(obs, t_obs)
std = torch.exp(0.5 * log_var)
eps = torch.randn_like(std)
z0 = mu + eps * std # reparameterization
# Integrate latent ODE from t=0 to all prediction times
all_times = torch.cat([t_pred[0:1], t_pred]).unique(sorted=True)
z_traj = odeint(self.ode_func, z0, all_times, method="dopri5")
x_pred = self.decoder(z_traj)
# KL divergence
kl = -0.5 * (1 + log_var - mu**2 - log_var.exp()).sum(dim=-1).mean()
return x_pred, kl
# Quick sanity-check forward pass
model = LatentODE().to(DEVICE)
T, B = 30, 4
obs_dummy = torch.randn(B, T, OBS_DIM, device=DEVICE)
t_obs_dummy = torch.linspace(0, 3, T, device=DEVICE).unsqueeze(0).expand(B, -1)
t_pred_dummy = torch.linspace(0, 5, 50, device=DEVICE)
x_pred, kl = model(obs_dummy, t_obs_dummy, t_pred_dummy)
print("x_pred shape:", x_pred.shape) # (50, B, OBS_DIM)
print("KL:", kl.item())