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

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

    Jun 13, 2026Abdeladhim TahimiAutomatic DifferentiationBackpropagation

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

  2. Mosaic: A Benchmark Suite for Differentiable Physics Solvers

    Jun 26, 2026Andrin Rehmann, Heiko Zimmermann, Dion HäfnerAutomatic DifferentiationPDE Solving

    Pasteur Labs, Brooklyn, NY, USA

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

    May 20, 2026Haozhe Jia, Pengyu Yin, Wenshuo Chen +6Physics-Informed Diffusion ModelsPhysics-Informed Generative Modeling

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