TorchSpin: Differentiable Electron Paramagnetic Resonance Simulation in PyTorch

Alec H. Follmer

physics.chem-phOct 12, 2026

TorchSpin: Differentiable Electron Paramagnetic Resonance Simulation in PyTorch

Authors: Alec H. Follmer

Organizations: Department of Chemistry, University of California, Davis, Davis, California 95616, United States · Linac Coherent Light Source, SLAC National Accelerator Laboratory, Menlo Park, California 94025, United States

Abstract

Interpreting electron paramagnetic resonance (EPR) spectra commonly requires repeated forward simulation and optimization of spin Hamiltonian parameters. Established software provides mature simulation and fitting methods, but these calculations are not generally available as differentiable components within machine-learning workflows. We present TorchSpin, an open-source PyTorch framework that ports the organization and numerical conventions of EasySpin into a tensor-based computational environment. Representative calculations spanning continuous-wave EPR, ENDOR, pulse EPR, spin dynamics, trajectory simulation, and magnetometry closely reproduced corresponding EasySpin results. In matched benchmarks, TorchSpin CPU execution was faster than EasySpin for six of nine workloads although performance depended strongly on the calculation. We therefore focused on the capabilities enabled by automatic differentiation in PyTorch. In a Cu(II) powder-spectrum demonstration, feedback propagated through TorchSpin improved the quality of a supervised convolutional neural network (CNN) without using parameter labels during refinement. Models trained from randomly initialized CNNs using only noisy spectra and TorchSpin reconstruction loss achieved median similarities of $\geq$ 0.99. CNNs trained on TorchSpin-generated libraries also transferred without adaptation to independently generated EasySpin spectra with performance comparable to EasySpin-trained controls. Once trained, the CNN produced initial parameter estimates in approximately 1 ms per spectrum while a batch of 512 was processed in approximately 1.3 ms with GPU acceleration. TorchSpin therefore provides both a tested Python implementation of established EPR calculations and a differentiable physical layer for rapid learned inference, spectral refinement, and parameter-label-free training.

Figures & tables

Figure 1: Numerical agreement between TorchSpin and EasySpin across representative EPR and ENDOR simulations. TorchSpin (red) and independently generated EasySpin (black) spectra are compared for ( A ) isotropic fast-motion cw EPR using garlic , ( B ) slow-motion cw EPR using chili , ( C ) solid-state cw EPR using pepper , and ( D ) solid-state ENDOR using salt. Blue traces show the residuals (TorchSpin $-$ EasySpin). Each EasySpin spectrum was normalized to its maximum absolute amplitude.

Figure 2: Performance of representative EasySpin and TorchSpin simulation workloads on a common computing platform. ( A ) Median wall time for EasySpin on the CPU, TorchSpin using the selected CPU thread count, and TorchSpin in CUDA mode. TorchSpin CPU calculations were evaluated using 1, 8, and 32 threads; red labels identify the selected configuration, defined as the smallest thread count whose timing interquartile range (IQR) overlapped that of the configuration with the lowest median time. Hatched blue bars denote hybrid CUDA execution, in which verified GPU kernels are combined with CPU-side operations. N/A indicates that the corresponding chili or cardamom implementation does not currently provide a CUDA execution path. ( B ) Speedup relative to EasySpin, calculated as the EasySpin wall time divided by the corresponding TorchSpin wall time. Values above unity indicate faster TorchSpin execution, and the dashed line marks equal performance. Points show ratios of median times; error bars are conservative IQR-propagated bounds calculated from the EasySpin and TorchSpin measurements. Each condition comprised one untimed warm-up followed by five timed replicates on a local workstation; bars and points report medians, and error bars span the IQR.

Figure 3: Overview of the CNN training and evaluation workflows used to test TorchSpin. ( A ) A CNN was first trained by supervised regression against the seven generating parameters of an EasySpin spectral library. During subsequent differentiable refinement, the CNN mapped noisy EasySpin spectra to the principal values of the g and 63 Cu hyperfine tensors and a Gaussian linewidth. TorchSpin reconstructed the predicted spectrum, and a spectral loss against the corresponding clean EasySpin target was backpropagated through the simulator to update the CNN. ( B ) In the parameter-label-free experiment, independently initialized CNNs were trained without supervised pretraining, parameter labels, or clean spectral targets. The same fixed noisy EasySpin spectrum served as both the CNN input and the reconstruction target, and TorchSpin supplied the spectral gradients used to update the network. ( C ) For the cross-engine library comparison, supervised CNNs were trained on independent labeled libraries generated with EasySpin or TorchSpin. The two conditions used the same architecture and training protocol, with identical initial weights within each paired initialization seed. The resulting checkpoints were frozen before prospective evaluation. ( D,E ) CNNs trained on the TorchSpin ( D ) or EasySpin ( E ) library were evaluated on the same held-out noisy EasySpin spectra. Predictions from both conditions were reconstructed using the same TorchSpin forward model and compared with the corresponding clean EasySpin spectra. Black arrows denote input and forward-computation paths; red arrows denote spectral-error propagation or comparison.

Figure 4: Differentiable spectral feedback and simulation-trained CNN inversion of Cu(II) powder EPR spectra. ( A ) Representative held-out spectrum showing the noisy input, clean EasySpin target, and TorchSpin reconstructions obtained from the supervised CNN before and after differentiable spectral refinement. Residuals are shown relative to the clean target. ( B ) Paired comparison across all 512 test spectra. Refinement improved 373 spectra (72.9%), increasing the median cosine similarity from 0.977 to 0.985. The dashed diagonal denotes equal performance before and after refinement. ( C ) Empirical cumulative distributions of spectral agreement for the supervised control and three independently initialized, parameter-label-free CNNs trained from random initialization using TorchSpin spectral feedback. All models were evaluated on the same 512-spectrum EasySpin test library. The fraction exceeding the 0.950 threshold marked by the dashed line was 65.2% for the supervised control and 98.6% for each label-free model. ( D ) Mean absolute error (MAE) in the predicted g x , g y , g z , A x , A y , A z , and Gaussian linewidth ( lw ), expressed as a percentage of the full prespecified sampling span for each parameter. Label-free bars report the mean across three independently initialized models, and error bars indicate the sample standard deviation. ( E ) Cross-engine comparison of supervised CNNs trained on independent labeled EasySpin or TorchSpin spectral libraries using paired initial weights and identical training protocols. The frozen models were evaluated on the same held-out EasySpin spectra, and all predicted parameters were reconstructed using TorchSpin. Color identifies the training-library simulator, line style distinguishes the three paired initialization seeds, and the gray curve shows TorchSpin reconstructions generated from the true sampled parameters. The dashed vertical line marks a cosine similarity of 0.950. Spectral agreement in all panels was evaluated against the corresponding noise-free EasySpin spectrum.

Figure 5: Interactive Jupyter interface for conventional cw EPR fitting with TorchSpin. Parameter sliders and fixed-parameter controls allow adjustment of the spin system parameters. Additional controls select the optimization method and fitting limits. The plots display the input spectrum, starting simulation, best fit, residual, and best RMSD as a function of model evaluations. The interface supports repeated adjustment and refitting within the notebook.

Figure S1: Comparison of trajectory-based cw EPR simulations generated by TorchSpin and EasySpin. TorchSpin (red) and EasySpin (black) spectra were calculated using cardamom with the same spin-system and experimental parameters. Blue shows the residual (TorchSpin $-$ EasySpin). Because the calculations use independently sampled rotational trajectories, the residual includes finite-sampling variation. Each TorchSpin spectrum was multiplied by a single amplitude factor to account for differences in absolute-intensity convention.

Figure S2: Comparison of time-domain two-pulse ESEEM simulations generated by TorchSpin and EasySpin. The real components of the TorchSpin (red) and EasySpin (black) echo signals were calculated using saffron with equivalent spin-system, pulse-sequence, and orientation-grid parameters. Blue shows the residuals (TorchSpin $-$ EasySpin). A single amplitude factor was applied to the complete TorchSpin trace to account for different absolute-intensity conventions.

Figure S3: Comparison of pulse-driven spin dynamics calculated by TorchSpin and EasySpin. The real-valued electron z -magnetization trajectories generated by TorchSpin (red) and EasySpin (black) are shown during a frequency-swept pulse and the subsequent free-evolution period. Blue shows the pointwise residual (TorchSpin $-$ EasySpin). A single amplitude factor was applied to the complete TorchSpin trajectory to account for differences in signal convention.

Figure S4: Comparison of temperature-dependent effective magnetic moments calculated by TorchSpin and EasySpin. The normalized effective magnetic moments obtained with TorchSpin (red) and EasySpin (black) are shown as a function of temperature for an exchange-coupled Gd(III)–Cu(II) spin system. Blue shows the pointwise residual (TorchSpin $-$ EasySpin), magnified by 10 6 for visibility. A single amplitude factor was applied to the complete TorchSpin trace before normalization.

Figure S5: CPU thread scaling of representative TorchSpin simulation workloads. Median wall times were measured using 1, 8, and 32 CPU threads on a local workstation. Each cell reports the median execution time with its IQR in brackets and the speedup relative to the corresponding one-thread calculation. Cell color encodes the speedup on a log 2 scale: blue indicates faster execution, red indicates slower execution, and white indicates little or no change. The black outline and bold text identify the selected CPU configuration, defined as the smallest thread count whose timing interquartile range overlapped that of the configuration with the lowest median wall time. Each condition comprised one untimed warm-up followed by five timed replicates; quartiles were calculated as the linearly interpolated 25th and 75th percentiles.

Workload	EasySpin (s)	TorchSpin CPU (s)	ES/TS CPU	TorchSpin CUDA mode (s)	ES/TS CUDA
pepper: Cu + 2 N, matrix	0.292 [0.229–0.351]	0.847 (32 t)	0.344	0.674 (hybrid)	0.433
pepper: Mn(II), matrix	3.03 [3.03–3.14]	1.07 (32 t)	2.83	0.954 (hybrid)	3.18
pepper: perturbative grid	0.0818 [0.072–0.0825]	0.0616 (1 t)	1.33	0.162 (hybrid)	0.506
pepper: strain summation	0.814 [0.808–0.816]	0.994 (8 t)	0.82	1.23 (hybrid)	0.661
pepper: 20-call fit pattern	2.5 [2.5–2.5]	1.29 (8 t)	1.93	1.77 (hybrid)	1.41
chili: nitroxide + proton	1.15 [1.14–1.16]	0.157 (1 t)	7.32	N/A	N/A

Table S1: Matched-host performance measurements used in the manuscript comparison.

Model coordinate	Sampled interval	Physical transformation
g $\perp$ center	2.040–2.080	g x = g $\perp$ $-$ $\Delta$ g/2; g y = g $\perp$ + $\Delta$ g/2
g rhombicity, $\Delta$ g	0–0.030	g y $-$ g x
g z	2.18–2.30	–
A $\perp$ center	35–85 MHz	A x = A $\perp$ $-$ $\Delta$ A/2; A y = A $\perp$ + $\Delta$ A/2
A rhombicity, $\Delta$ A	0–40 MHz	A y $-$ A x
A z	450–650 MHz	–

Table S2: Bounded coordinates used to generate the Cu(II) spectral libraries.

Figure S6: Analysis of differentiable spectral refinement after supervised initialization. ( A ) Cumulative distribution of cosine similarity to the clean EasySpin test spectrum for the supervised and physics-refined CNNs. ( B ) Median cosine similarity within six predefined noise intervals; error bars show the interquartile range. The bin counts from lowest to highest noise were 82, 97, 84, 83, 79, and 87 spectra. ( C ) Training and validation histories for supervised parameter regression ( left ) and parameter-label-free spectral refinement ( right ); dotted lines mark the selected checkpoints. ( D ) Predicted and generating g and hyperfine components for the refined model.

Metric	Supervised CNN	Physics-refined CNN
Median clean-spectrum cosine	0.9768	0.9849
Minimum cosine	0.5934	0.8178
Fraction with cosine $\geq$ 0.950	76.6%	90.2%
Spectra improved	—	373/512 (72.9%)

Table S3: Test-set effects of parameter-label-free spectral refinement after supervised initialization.

Parameter	Supervised MAE	Refined MAE
g x	0.00125	0.00116
g y	0.00111	0.00099
g z	0.00369	0.00396
A x (MHz)	5.72	5.89
A y (MHz)	4.25	4.44
A z (MHz)	8.26	8.42

Table S4: Component-wise parameter MAE before and after parameter-label-free spectral refinement.

Figure S7: Optimization histories for the three parameter-label-free CNNs. Validation spectral loss and training spectral loss across the 24-epoch schedules. Line styles distinguish the three initialization seeds.

Figure S8: Parameter recovery and predicted-parameter correlations for training without labels. Top : generating versus predicted physical parameters for the supervised control (black) and all three label-free seeds (red); the dashed line denotes equality. Bottom : correlation matrices of the seven predicted parameters for each label-free seed.

Figure S9: Representative parameter-label-free reconstructions across the performance distribution. Best, median-ranked, and worst examples are shown for each initialization seed. Gray traces are fixed noisy observations, black traces are clean EasySpin spectra, and red traces are TorchSpin reconstructions from the label-free CNN predictions. Test indices and cosine similarities are given above each panel.

Model	Median cosine	Fraction $\geq$ 0.950	g x MAE	g y MAE	g z MAE	A x MAE (MHz)	A y MAE (MHz)	A z MAE (MHz)	lw MAE (mT)
Supervised control	0.9698	65.2%	0.00144	0.00140	0.00432	6.45	4.41	10.0	0.0855
Label-free seed 1	0.9923	98.6%	0.00118	0.000821	0.00392	6.98	5.01	9.04	0.107
Label-free seed 2	0.9921	98.6%	0.00115	0.000879	0.00454	7.01	4.82	9.81	0.105
Label-free seed 3	0.9913	98.6%	0.00126	0.000945	0.00444	6.96	5.05	9.87	0.124

Table S5: Prospective test-set performance for parameter-label-free training from random initialization.

Figure S10: Cumulative distributions of reconstruction NRMSE for CNNs trained on labeled EasySpin or TorchSpin libraries and evaluated on the same prospective EasySpin test set. Line styles identify the three paired initialization seeds. All predictions were decoded with TorchSpin.

Panel of Figure 21 — Figure S11: Ranked prospective reconstructions for paired initialization seed 1. Left : CNN trained on the TorchSpin library. Right : CNN trained on the EasySpin library. For each condition, the best, median-ranked, poor, and worst cases were selected by cosine similarity to the clean EasySpin test spectrum. Gray traces are noisy inputs, black traces are clean EasySpin references, and red traces are TorchSpin reconstructions.

Panel of Figure 22 — Figure S12: Ranked prospective reconstructions for paired initialization seed 2. Left : CNN trained on the TorchSpin library. Right : CNN trained on the EasySpin library. For each condition, the best, median-ranked, poor, and worst cases were selected by cosine similarity to the clean EasySpin test spectrum. Gray traces are noisy inputs, black traces are clean EasySpin references, and red traces are TorchSpin reconstructions.

Panel of Figure 23 — Figure S13: Ranked prospective reconstructions for paired initialization seed 3. Left : CNN trained on the TorchSpin library. Right : CNN trained on the EasySpin library. For each condition, the best, median-ranked, poor, and worst cases were selected by cosine similarity to the clean EasySpin test spectrum. Gray traces are noisy inputs, black traces are clean EasySpin references, and red traces are TorchSpin reconstructions.

Training library	Median cosine	Fraction $\geq$ 0.950	Mean NRMSE
TorchSpin	0.9656 ± 0.0008	61.59 ± 2.07%	0.3075 ± 0.0067
EasySpin	0.9667 ± 0.0004	63.93 ± 1.59%	0.3005 ± 0.0063

Table S6: Prospective cross-engine comparison averaged across three paired CNN initialization seeds. Values are mean ± sample standard deviation across three paired initialization seeds evaluated on the same 512-spectrum test set; percentage uncertainties are percentage points, not confidence intervals.

Training library	g x	g y	g z	A x (MHz)	A y (MHz)	A z (MHz)	lw (mT)
TorchSpin	0.00142	0.00138	0.00481	5.65	4.56	10.9	0.0799
EasySpin	0.00141	0.00133	0.00490	5.99	4.55	10.5	0.0807

Table S7: Mean component-wise parameter MAE across three paired initialization seeds in the cross-engine training-library comparison.

Operation	Device/configuration	Median [IQR] (ms)	Per-spectrum value (ms)
Single CNN inference, resident	CPU, 1 thread	0.976 [0.974–0.985]	0.976
Host normalization + CNN prediction	CPU, 1 thread	0.994 [0.992–0.998]	0.994
Single CNN inference, resident	RTX 4090	1.019 [1.016–1.024]	1.019
Host normalization + transfer + prediction	RTX 4090	1.047 [1.045–1.053]	1.047
Batch CNN inference, n = 32	RTX 4090	1.029 [1.026–1.034]	0.0321
Batch CNN inference, n = 128	RTX 4090	1.016 [1.014–1.021]	0.00794

Table S8: CNN inference and TorchSpin reconstruction latency.

Magnetic Resonance Spectroscopy

See all 13 papers

Oct 16, 2026

TorchTPU - A Hybrid Eager/Compiled Architecture for Native PyTorch Execution on TPU

TorchTPU is an open-source collaboration between Google and Meta that introduces a native and principled PyTorch backend for Google TPUs. It enables seamless out-of-the-box model execution and TPU migration while delivering industry-competitive inference and training performance through eager execution optimizations, torch.compile integration, and native hardware acceleration.

Claudio Basile, Michael Voznesensky, Gunhyun Park +50

Neural Network Hardware Acceleration Efficient Neural Network Inference

Ask PDF

Oct 16, 2026

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

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.

Akzhol Almukhametov, Yang Liu

Department of Nuclear Engineering, Texas A&M University

Differentiable Physics Physics-Based Simulation Multiphysics Simulation

Ask PDF

Aug 10, 2023

Spintronics for image recognition: performance benchmarking via data-driven simulations

We present a demonstration of image classification using an extreme learning machine (ELM) based on a unique simulated magnetic tunnel junction (MTJ) delayed in time. As the ground state of the MTJ is a magnetic vortex, we refer to it as a vortex-based spin-torque oscillator (STVO). The dynamics of the magnetic vortex is simulated with a model called the data-driven Thiele equation approach (DD-TEA). This allows to avoid the constraints associated with repeated experimental manipulation for hyperparameters search and benchmarking. We showcase the versatility of our implementation by using it successfully for classification tasks on the MNIST, EMNIST-letters and Fashion MNIST datasets. Through simulations, we show that within an ELM with a sufficient number of parameters, the performance reached using the STVO dynamics as a source of nonlinearity is equivalent to the ones obtained with classical software activation functions such as the reLU and the sigmoid. While achieving state-of-the-art accuracy levels on the MNIST dataset, our model's performance on EMNIST-letters and Fashion MNIST is lower due to the simplicity of the network architecture and the increased complexity of the data. We expect that the DD-TEA framework will enable the exploration of deeper and more complex STVO-based architectures, ultimately leading to improved classification accuracy.

Anatole Moureaux, Chloé Chopin, Simon de Wergifosse +2

Institute of Condensed Matter and Nanosciences, Université catholique de Louvain, Place Croix du Sud 1, 1348 Louvain-la-Neuve, Belgium · Univ. Grenoble Alpes, CEA, CNRS, Grenoble INP, SPINTEC, 3800 Grenoble, France · Institute for Information and Communication Technologies, Electronics and Applied Mathematics, Université catholique de Louvain, Avenue Georges Lemaître 4, 1348 Louvain-la-Neuve, Belgium

Neuromorphic Computing Image Classification

Ask PDF

Abstract

Magnetic Resonance Spectroscopy

Explore similar work

TorchTPU - A Hybrid Eager/Compiled Architecture for Native PyTorch Execution on TPU

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

Spintronics for image recognition: performance benchmarking via data-driven simulations