NUCDIFF: A Differentiable Nuclide Depletion Solver for Sensitivity Analysis of Long-Term Decay Heat

With the renewed focus on nuclear power for energy creation in the public and private sectors, renewed interest has emerged on the practical application of various industry codes for the effective management of spent nuclear fuel. Of the standard nuclide depletion codes known, such codes treat the matrix-exponential solver in regard to the Bateman equations as a black box, where the solver is implemented through compiled routines that expose only the forward map and not the intermediate quantities required for reverse-mode automatic differentiation.As a result, sensitivity analyses with respect to nuclear data parameters require many forward solves to an almost extraordinary amount, one per parameter for finite-difference gradients, resulting in extraordinary computational complexity. To improve computational efficiency, we present NUCDIFF, a JAX-based implementation of the 16th-order Chebyshev rational approximation method (CRAM-16) that exposes derivatives of any forward observable with respect to half-lives, branching ratios, initial inventory, and time step through reverse-mode automatic differentiation, using a custom Jacobian–vector product derived from the resolvent identity.For application, we find that on a stylized 132-nuclide pressurized water reactor spent fuel inventory, a single reverse-mode pass computes the sensitivity vector by a ∼ 360× speedup that grows linearly with the network, compared to an equivalent central difference sweep. Forward passes agree to 1% for handbook decay heat values. We also identify the dominant contributors at each timescale, and display autodiff/finite-difference agreement to −10−6 relative. The code, tests, and data used to produce every figure in this paper are available in the supplementary information.