cs.LGOct 16, 2026

Torch-DPhys: A Torch-Native, MOOSE-Compatible Framework for Differentiable Physics Simulation and Hybrid Machine Learning Integration

Authors: Akzhol Almukhametov, Yang Liu

Organizations: Department of Nuclear Engineering, Texas A&M University

Read PDF
Ask AI
Cite

Copy citation

DOINot available for this paper

Abstract

We present Torch-DPhys, a Torch-native framework implemented in PyTorch for differentiable physics simulation and hybrid physics-machine learning. It implements selected thermal-hydraulic modeling concepts and discretizations from MOOSE's Thermal Hydraulics Module and the System Analysis Module (SAM), with a preliminary extension to continuum mechanics. Native tensor residuals support CPU and GPU execution, compilation of selected numerical kernels, and sparse direct linear solves. Sensitivities are obtained by differentiating executed numerical operations or through discrete adjoints of converged residual equations. Verification combines analytical solutions, reference-code comparisons, and derivative checks. Demonstrations cover plane-strain elasticity, one- and two-loop thermal-hydraulic transients, and an axisymmetric molten-salt fast reactor. For the reactor, comparison with SAM-Lite gives a maximum normalized field discrepancy of 1.13 x 10^-5 over 200 adaptive time steps, demonstrating agreement between the tested discrete implementations. Selected transient sensitivities are assessed using tangent-linear calculations and finite differences on fixed time schedules. Two applications illustrate the framework's connection between simulation and optimization. A residual-adjoint parameter-identification calculation recovers nominal pump head and heating power with relative errors below 0.06% from noiseless synthetic MOOSE observations. A separate three-loop plant demonstration uses solver-in-the-loop training of a predictive controller by differentiating trajectory-level objectives through the thermal-hydraulic simulation. These examples illustrate a foundation for gradient-based inference, control, and coupled physics-machine learning models in nuclear engineering.

Figures & tables

Figure 1: Torch-DPhys workflow. Solid arrows show forward evaluation and parameter updates. Dashed arrows show the two differentiation routes: unrolled differentiation through executed solver operations and discrete adjoints of the converged residual equations.
Figure 2: Coarse mesh of the axisymmetric MSFR computational domain, with enlarged detail. The axial direction is vertical and the symmetry axis is on the left.
QuantityConfiguration
Fluid propertiesConstant ρ=4306.71\rho=4306.71 kg/m 3 , cp=1593.94c_{p}=1593.94 J/(kg K), and μ=0.0253665\mu=0.0253665 Pa s. Molecular conductivity varies linearly from 0.9952 at 800 K to 1.0288 W/(m K) at 1200 K.
Mixing length0.07 m for r<1.25r<1.25 m, decreasing linearly to 0.018 m at r=1.75r=1.75 m, and 0.018 m beyond. Turbulent Prandtl number Prt=0.9Pr_{t}=0.9 .
InletRecorded normalized Nek5000 velocity profile, multiplied by −1.25-1.25 m/s and the coarse-boundary flux correction 1.02705297222. Radial velocity is zero and temperature is 900 K.
Outlet, walls and axisOutlet pressure 161 kPa. Walls are no-slip and adiabatic. The symmetry axis has zero radial velocity and zero normal gradients of axial velocity and temperature.
Initial fieldsBoth velocity components 10−510^{-5} m/s, pressure 110 kPa and temperature 900 K, subject to the prescribed boundary values.
Body force and heatingGravity is +9.81+9.81 m/s 2 in the mesh axial direction. The fixed volumetric heat source is given by Eq. ( 5 ).
Table 1: Physical settings of the MSFR initialization benchmark. Coordinates xx and rr denote the axial and radial directions.
Figure 3: Plane-strain cantilever verification. (a) Deformed mesh colored by displacement magnitude (unmagnified). (b) Upper-boundary vertical displacement profile compared against archived MOOSE results.
Figure 4: Schematics of the one-dimensional thermal-hydraulic verification cases: (a) Case 03 open helium chain, (b) Case 04 closed loop with constant pump head, and (c) Case 05 coupled primary–secondary system with PID control. Arrows indicate nominal flow direction and dots mark output sensor locations.
Figure 5: Comparison between Torch-DPhys and MOOSE for the 03_upper_loop.i configuration: (a) core outlet fluid temperature, (b) heat transfer rate from the solid core to the coolant, (c) cooler outlet temperature, and (d) core inlet and outlet pressures relative to 11 MPa.
Figure 6: Signed differences (Torch-DPhys minus MOOSE) corresponding to the four observables in Figure 5 .
Figure 7: Comparisons between Torch-DPhys and MOOSE for the closed primary loop ( 04_loop.i with constant pump head): (a) core outlet mass flow and (b) core outlet temperature; and for the coupled two-loop system ( 05_secondary_side.i ): (c) primary and secondary heat-exchanger outlet temperatures and (d) PID-controlled primary-loop mass flow rate.
Figure 8: MSFR forward verification over the 200-step adaptive transient. (a,b) Reference SAM-Lite temperature and velocity magnitude fields at t=43.7t=43.7 s. (c) Nodal absolute temperature difference ∣TTorch−TSAM∣|T^{\mathrm{Torch}}-T^{\mathrm{SAM}}| at t=43.7t=43.7 s in units of 10−1310^{-13} K (gray indicates exact bitwise zero). (d) Evolution of the normalized field discrepancies ev(k)e_{v}(k) defined in Eq. ( 8 ) over all 200 steps.
FieldUnitMaximum absolute differenceStep
Temperature TTK1.94×10−51.94\times 10^{-5}31
Axial velocity uzu_{z}m/s1.02×10−51.02\times 10^{-5}3
Radial velocity uru_{r}m/s2.28×10−52.28\times 10^{-5}3
Pressure ppPa2.85×10−22.85\times 10^{-2}4
Table 2: Maximum absolute nodal discrepancies between Torch-DPhys and SAM-Lite across all 1,399 vertices and steps 1–200 of the MSFR initialization transient.
Figure 9: Relative discrepancy between the MSFR discrete residual adjoint and forward tangent-linear sensitivities across 39 parameter directions on trajectory windows of 1 to 40 accepted steps. Triangles in the lower band indicate exact bitwise zero discrepancy; “+ retry” denotes the 10-step window containing a rejected step attempt.
ItemSetting
Workload200 adaptive forward steps to t=43.7148t=43.7148 s at baseline tolerances (266 Newton iterations), followed by a 200-step discrete residual adjoint.
Hardware and precisionIntel Core Ultra 7 265K CPU, NVIDIA RTX 2000 Ada GPU (16 GB), float64 arithmetic under Windows Subsystem for Linux 2 (WSL2).
Software stackPython 3.12.3, PyTorch 2.13.0+cu129, Triton 3.7.1, CUDA 12.9, cuDSS 0.8.0.10, NumPy 2.1.2, SciPy 1.16.1, Intel oneAPI Math Kernel Library (MKL) 2024.2.
Thread and rank settingsTorch-DPhys: 4, 8, or 16 Torch/OpenMP/MKL threads (with MKL_CBWR=AUTO,STRICT and 1 OpenBLAS thread). SAM-Lite: 1 rank with 1 or 4 MOOSE threads, or 4, 8, and 16 MPI ranks with 1 thread per rank.
Measurement protocolMedians and min–max ranges over three warm repetitions for Torch-DPhys; single runs for cold compilation and SAM-Lite. Forward timings measure Transient.execute() (Torch-DPhys) and cumulative PicardSolve (SAM-Lite).
Table 3: Hardware and software configuration for the 200-step MSFR performance benchmark.
Figure 10: Wall-clock execution costs on the 200-step MSFR benchmark. (a) Forward and discrete-adjoint runtimes across 4-thread CPU and GPU configurations in Torch-DPhys (showing warm medians, min–max ranges, and cold first-call totals) alongside 4-rank MPI SAM-Lite. (b) Forward runtime versus configured CPU threads (Torch-DPhys) or MPI ranks (SAM-Lite).
Figure 11: Per-step component runtime breakdown (in ms per accepted step) for (a) forward stepping and (b) the discrete adjoint prior to auxiliary reuse. (c) Compiled GPU adjoint breakdown before and after enabling auxiliary-variable reuse (hatched), with diamonds indicating the independently measured uninstrumented step totals.
Observation source / Initial guessInitial HH (m)Initial QQ (W)L-BFGS-B updatesRecovered (H,Q)(H,Q) (m, W)
Same solver / Start A0.0661513008(0.094500, 2000.00)
Same solver / Start B0.1228527007(0.094500, 2000.00)
MOOSE / Start A0.0661513008(0.094516, 1998.87)
Table 4: Parameter identification results for the closed thermal-hydraulic loop ( Hnom=0.0945H^{\mathrm{nom}}=0.0945 m, Qnom=2000Q^{\mathrm{nom}}=2000 W).
Figure 12: Optimization trajectories for same-solver parameter recovery in the 04_loop.i thermal-hydraulic loop: (a) normalized parameter estimates θi/θinom\theta_{i}/\theta_{i}^{\mathrm{nom}} and (b) objective value JJ versus L-BFGS-B update. Solid and dashed curves denote the discrete residual adjoint and unrolled Newton differentiation, respectively.
Differentiation routeBackward time (s)Peak allocation (MiB)
Converged-residual discrete adjoint26.9 [26.1, 27.0]24.7
Unrolled Newton AD (with step recomputation)268 [263, 274]201.0
Table 5: Backward-pass runtime (median and [min, max] over three fresh processes) and peak live GPU memory allocation for a 122-step trajectory at Start A.
Figure 13: CAD representation of the three-loop thermal-fluid facility modeled in the differentiable predictive control demonstration (HX: heat exchanger).
Figure 14: Two-level hybrid physics–ML training architecture for differentiable predictive control: (a) inner loop backpropagating trajectory losses through the mechanistic plant solver to update GRU policy weights, and (b) outer Bayesian optimization loop tuning objective weights and temporal discount factors.
Figure 15: Closed-loop DPC tracking performance across four power-demand profiles. Top row: delivered versus demanded thermal power. Bottom row: peak heater-rod temperature relative to the 97 ∘C97\,^{\circ}\mathrm{C} soft penalty barrier and 100 ∘C100\,^{\circ}\mathrm{C} operating limit. Solid curves and shaded bands indicate the mean and ±1\pm 1 standard deviation across the four selected policies.
Appendix figures & tables7 assets

Supplementary material from the paper’s appendix.

Appendix

ObservationsParameterAdjoint derivativeh=10−2h=10^{-2}h=10−3h=10^{-3}h=10−4h=10^{-4}
Same solverPump head−0.02324309-0.023243097.63×10−57.63\times 10^{-5}6.13×10−76.13\times 10^{-7}1.93×10−71.93\times 10^{-7}
Same solverHeating power−0.01525549-0.015255499.90×10−59.90\times 10^{-5}1.01×10−61.01\times 10^{-6}1.20×10−61.20\times 10^{-6}
MOOSEPump head−0.02326628-0.023266287.62×10−57.62\times 10^{-5}6.12×10−76.12\times 10^{-7}1.93×10−71.93\times 10^{-7}
MOOSEHeating power−0.01522336-0.015223369.92×10−59.92\times 10^{-5}1.02×10−61.02\times 10^{-6}1.20×10−61.20\times 10^{-6}
Appendix
Table 6: Comparison of residual-adjoint logarithmic derivatives against central finite differences at Initial Guess A ( zH=log⁡(0.7)z_{H}=\log(0.7) , zQ=log⁡(0.65)z_{Q}=\log(0.65) ) for perturbation step sizes h∈{10−2,10−3,10−4}h\in\{10^{-2},10^{-3},10^{-4}\} .
Figure 16: Nonlinear residual convergence histories on six matched single-step MSFR problems comparing SAM-Lite PJFNK against Torch-DPhys Newton with exact AD Jacobians and direct LU solves: (a) backward-Euler startup step and (b–f) BDF2 continuation steps.
Figure 17: Attempt sequence and nonlinear iteration counts during the tight-tolerance ( rtol=10−12r_{\mathrm{tol}}=10^{-12} , unit scaling) MSFR replay. Top: step sizes and acceptance outcomes. Bottom: nonlinear iteration counts for SAM-Lite and Torch-DPhys.
QuantityChange ( Δt=0.03125\Delta t=0.03125 s minus 0.250.25 s)Context
Mean outlet temperature+0.46+0.46 K 
Mean inlet pressure+2.2+2.2 kPa5.8%5.8\% of the inlet–outlet pressure drop
Outlet mass flow from velocity field+36+36 kg/sGlobal mass-flux imbalance rises from 1.65%1.65\% to 1.90%1.90\%
Maximum core temperature+0.8+0.8 K 
Appendix
Table 7: Endpoint solution sensitivity of the production MSFR formulation when refining a constant time step from Δt=0.25\Delta t=0.25 s to Δt=0.03125\Delta t=0.03125 s at t=43.7t=43.7 s.
QuantityComparison scopeMaximum discrepancy
Jacobian assemblyTested states within each execution modeBitwise identical
Forward statesEager trajectories (Eq. ( 8 ))1.53×10−141.53\times 10^{-14}
Adjoint gradientsEvaluated on common dense trajectories7.14×10−127.14\times 10^{-12}
Adjoint gradientsSparse forward trajectory vs. dense reference1.49×10−111.49\times 10^{-11}
Appendix
Table 8: Maximum normalized discrepancies between the sparse SuperLU and dense LU implementations across the 10-step retry and 40-step MSFR verification windows.
Trajectory window / scheduleMax. discrepancyFD countsReference change estabe_{\mathrm{stab}}
1 step, backward Euler1.15×10−121.15\times 10^{-12}11/32.10×10−122.10\times 10^{-12}
3 steps, backward Euler + BDF26.75×10−126.75\times 10^{-12}11/35.28×10−125.28\times 10^{-12}
10 steps, with step-9 retry7.02×10−127.02\times 10^{-12}11/31.40×10−111.40\times 10^{-11}
10 steps, baseline schedule1.38×10−111.38\times 10^{-11}3/01.47×10−101.47\times 10^{-10}
40 steps, baseline schedule7.37×10−107.37\times 10^{-10}3/07.93×10−117.93\times 10^{-11}
Appendix
Table 9: MSFR transient-adjoint verification across five trajectory windows. “Max. discrepancy” gives the maximum relative difference between the adjoint and forward tangent-linear recurrence across 39 directions; “FD counts” lists resolved-and-agreeing versus unresolved finite-difference directions; and “Reference change” gives estabe_{\mathrm{stab}} (Eq. ( 27 )) after applying two additional root-polishing Newton updates per step.
Evaluated quantityMaximum normalized discrepancy
Assembled residual vector8.35×10−168.35\times 10^{-16}
Residual vector–Jacobian product (VJP)4.27×10−164.27\times 10^{-16}
Element-Jacobian blocks6.61×10−166.61\times 10^{-16}
Dense LU factors and linear-system solutionsBitwise identical
Adjoint gradient (40-step window)1.14×10−131.14\times 10^{-13}
Adjoint gradient (10-step retry window)1.40×10−141.40\times 10^{-14}
Appendix
Table 10: Maximum normalized discrepancies between compiled ( torch.compile ) and eager MSFR execution on CPU ( float64 , 4 threads).
More in this topic

Differentiable Physics

Differentiable physics makes physical simulators usable for gradient-based learning, parameter estimation, design, and control.

See all 64 papers

Explore similar work

CardsList
Jun 13, 2026

Automatic Differentiation from Scratch: How PyTorch Computes Gradients in Physics-Informed Neural Networks

This paper traces, with explicit numerical values, how PyTorch's automatic differentiation (AD) engine computes gradients for Physics-Informed Neural Network (PINN) training -- a setting that requires two levels of differentiation: computing the physics derivative y^′(t)=dy^/dt\hat{y}'(t)=d\hat{y}/dt through the network, and computing parameter gradients ∇θL\nabla_θL of a loss that itself depends on y^′(t)\hat{y}'(t). Using a 1-3-3-1 multilayer perceptron and the initial value problem y′(t)+y(t)=0y'(t)+y(t)=0, y(0)=1y(0)=1, we trace the complete pipeline at every node: the computational graph built during the forward pass, the reverse-mode backward traversal that computes all 22 parameter gradients in a single pass, and the graph-on-graph mechanism by which create\_graph=True enables correct differentiation through the physics-informed residual. Every adjoint value is verified against the hand derivations of Tahimi (2026), connecting the P/QP/Q sensitivity framework to the vector--Jacobian products used by PyTorch's autograd engine.

Abdeladhim Tahimi

CECA, Universidade Federal de Alagoas (UFAL), Rio Largo, Alagoas, 57100-000, Brazil

Automatic DifferentiationBackpropagation
AskPDF
Jun 26, 2026

Mosaic: A Benchmark Suite for Differentiable Physics Solvers

Differentiable partial differential equation (PDE) solvers underpin solver-in-the-loop ML training, gradient-based optimal control, and inverse problems, yet the practical cost of obtaining correct, usable gradients from a given solver on a given problem is largely undocumented. Integration effort, computational cost, gradient accuracy, and numerical conditioning vary widely across solvers and are discoverable only by trial and error. We introduce Mosaic, an extensible benchmarking framework for differentiable PDE solvers that standardizes access to solver gradients. Each solver is packaged as a containerized component (Tesseract) exposing a uniform gradient API regardless of language or automatic differentiation (AD) strategy, enabling researchers to evaluate, compare, and build on non-trivial physical solvers. Our evaluation of 14 solvers across fluid dynamics, structural mechanics, and heat transfer demonstrates that the benchmark surfaces practically relevant differences: order-of-magnitude variation in computational cost and Jacobian conditioning, alongside structural incompatibilities that eliminate solvers from realistic tasks entirely. Despite this variation, all solvers that produce gradients converge to similar optima, indicating that the practical barriers are memory limits, numerical stability, and setup compatibility rather than gradient accuracy alone. Mosaic is open-source and available at https://github.com/pasteurlabs/mosaic.

Andrin Rehmann, Heiko Zimmermann, Dion Häfner

Pasteur Labs, Brooklyn, NY, USA

Automatic DifferentiationPDE SolvingBenchmark DesignDifferentiable ProgrammingDifferentiable PhysicsPartial Differential Equations
AskPDF
May 20, 2026

Learning to Think in Physics: Breaking Shortcut Learning in Scientific Diffusion via Representation Alignment

Physics-informed diffusion models typically enforce PDE constraints only on final outputs, leaving intermediate representations unconstrained and prone to shortcut learning under shifted boundary conditions. We introduce REPA-P, a teacher-free, architecture-agnostic framework that aligns intermediate features with physical states using first-principles residuals. REPA-P attaches lightweight 1×11{\times}1 projection heads to selected layers, decodes hidden activations into physical quantities, and applies PDE residual losses during training. These heads are discarded at inference, introducing zero overhead. Across four PDE tasks, including Darcy flow, topology optimization, electrostatic potential, and turbulent channel flow, REPA-P accelerates convergence by up to 2×2{\times}, reduces physics residuals by up to 66.4%66.4\%, and improves out-of-distribution robustness by up to 49.3%49.3\%, with consistent gains on both U-Net and Diffusion Transformer backbones. Ablations show that supervising a small set of intermediate layers captures most benefits and complements output-level physics losses. Code is available at https://github.com/Hxxxz0/REPA-P.

Haozhe Jia, Pengyu Yin, Wenshuo Chen +6

The Hong Kong University of Science and Technology (Guangzhou), Guangzhou, China · Shandong University, Jinan, China · Data61/CSIRO, Australia +5

Physics-Informed Diffusion ModelsPhysics-Informed Generative ModelingPDE Surrogate ModelingPhysics-Informed MLShortcut LearningDiffusion Models
AskPDF