pinn/CLAUDE.md

CLAUDE.md

This file provides guidance to Claude Code (claude.ai/code) when working with code in this repository.

Project Overview

Heat Equation PINN — A Physics-Informed Neural Network that solves the 1D time-varying heat equation with physical boundary conditions:

∂T/∂t = α ∂²T/∂x²    x ∈ [0, L],  t ∈ [0, T_END]

• IC: T(x, 0) = T0 (uniform)
• BC x=0: Neumann — heat flux step: −k ∂T/∂x = Q(t) where Q = Q_VAL if t ≥ T_STEP else 0
• BC x=L: Robin — convection: −k ∂T/∂x = h (T − T_AMB)

No experimental data is needed. A fdm/ module provides a reference numerical solution (FTCS explicit scheme) used for evaluation and visualization comparison.

All physical and numerical parameters live in config.py.

Running

Always activate the virtual environment first:

source .venv/bin/activate

PINN:

python app.py          # Train / Evaluate (L2 vs FDM) / Visualize

FDM reference solver:

python fdm/app.py      # Solve / Heatmap / Animation / Time-series

Saved artifacts (git-ignored): models/best_heat_pinn_model.pth, HTML plots in animations/ and animations/fdm/.

To retrain from scratch: rm models/best_heat_pinn_model.pth before running option 1.

Dependencies

requirements.txt exists. Key packages: torch, numpy, plotly. No pandas or scikit-learn needed.

GPU is auto-detected (cuda → mps → cpu) in engine.py:_get_device().

Architecture

config.py          ← all physical + numerical parameters (edit here to change the problem)
app.py             ← PINN CLI menu
model.py           ← HeatPINN + heat_pinn_loss()
engine.py          ← data sampling, Adam+L-BFGS training, evaluation vs FDM, visualization call
visualizer.py      ← PINN vs FDM: heatmap, animated T(x), time-series at fixed points
fdm/
  solver.py        ← FTCS explicit scheme, ghost-cell Neumann, explicit Robin
  visualizer.py    ← same 3 plot types for FDM-only output
  app.py           ← FDM CLI menu

Neural Network (model.py)

HeatPINN: 5-layer fully connected, input (x, t) → output T.

Output scaling — the network predicts a dimensionless perturbation; the forward() applies:

T = T_AMB + (Q_VAL · L / K) · net(x, t)

This keeps net outputs in [0, 1] range and ensures gradients ∂T/∂x are O(1) for the network to learn. Do not remove this scaling.

heat_pinn_loss() normalizes all three loss terms to O(1) using T_char = Q_VAL·L/K and grad_char = (Q_VAL/K)². Changing physical parameters in config.py does not require re-tuning loss weights.

Training (engine.py)

prepare_data() samples collocation points with deliberate clustering: extra points near x=0 (steep Neumann gradient) and around t=T_STEP (flux step discontinuity). Increasing N_f / N_bc here is the first lever to pull if accuracy is low.

train_model() runs Adam first, then L-BFGS fine-tuning. L-BFGS uses a closure that captures loss components in _last dict (avoids calling heat_pinn_loss outside an active grad context).

evaluate_model() runs the FDM solver and downsamples its (NX, NT) output to the PINN prediction grid (100, 100) for L2 comparison.

FDM Solver (fdm/solver.py)

Returns (T_matrix[NX, NT], x_vals, t_vals). Uses:

• Ghost cell for Neumann: T_ghost = T[1] + 2·dx·Q(t)/k
• Explicit Robin at x=L: T[N] = (T[N−1] + dx·h/k·T_amb) / (1 + dx·h/k) — uses T_cur[-2], not T_new[-2]
• CFL check at startup (warns, does not crash)

Loss Scaling Notes

If you change Q_VAL, K, H_CONV, or L in config.py, the normalization in heat_pinn_loss() adjusts automatically. If losses diverge, check that T_char = Q_VAL·L/K is not near zero.