TorchSpin: Differentiable Electron Paramagnetic Resonance Simulation in PyTorch

The gist: TorchSpin presents an open-source PyTorch framework that ports EasySpin's organization into a differentiable tensor-based environment, enabling rapid learned inference and parameter-label-free training for electron paramagnetic resonance simulations.

How it works

TorchSpin is a PyTorch port of the EasySpin simulation framework, retaining established EasySpin program names like garlic, chili, pepper, salt, saffron, spidyan, and curry to identify TorchSpin ports of the corresponding EasySpin simulation pathways N. Representative calculations spanning continuous-wave EPR, ENDOR, pulse EPR, spin dynamics, trajectory simulation and magnetometry closely reproduced corresponding EasySpin results N.

The framework is designed to connect a forward calculation directly to learned parameter inference through automatic differentiation N. This enables a second form of training where parameters predicted from a spectrum can be returned to the physical model, and disagreement between the reconstructed and measured spectra can be propagated through the simulation to update either the predicted parameters or the network itself N. This allows for spectral feedback even when the generating Hamiltonian is unknown N.

Numerical Correspondence and Validation

TorchSpin reproduces EasySpin across representative simulation regimes, with comparisons showing high fidelity across various EPR and ENDOR simulations N. For instance, the fast-motion continuous wave (cw) EPR spectrum calculated with garlic was reproduced to numerical precision, with a normalized root mean square error (NRMSE) below 10-15 N. Slow-motion cw EPR calculations for chili and solid-state pepper gave NRMSE values of 7.23 x 10-5 and 9.09 x 10-5, respectively N. The trajectory-based cardamom calculation reproduced the cw EPR envelope with a cosine similarity greater than 0.999 and an NRMSE of 0.0201 N.

Performance Benchmarking

Performance was evaluated across nine matched EasySpin and TorchSpin workloads chosen to represent distinct computational demands N. TorchSpin CPU execution was faster than EasySpin for six of the nine workloads N, with the largest difference observed for the two-nucleus slow-motion chili calculation, where TorchSpin was 7.32-fold faster N. CUDA execution was verified for five pepper workloads and the saffron HYSCORE calculation, but these were classified as hybrid because substantial numerical work remained on the CPU N.

Differentiable Spectral Feedback in Machine Learning

The paper demonstrated that differentiable spectral feedback improves learned EPR inversion N. In a proof-of-concept workflow, a CNN trained to predict principal values of the g and 63Cu hyperfine tensors and a Gaussian linewidth from noisy Cu(II) powder spectra was refined using TorchSpin N. This differentiable spectral refinement improved 373 of the 512 test reconstructions (72.9%) N, increasing the median cosine similarity from 0.977 to 0.985 N.

Furthermore, parameter-label-free training showed that three independently initialized CNNs trained without parameter labels or clean spectral targets produced a median cosine similarity to the clean EasySpin test spectra of 0.9923 N, compared with 0.970 for the supervised control N. This indicates that TorchSpin supplied sufficient spectral gradients to learn accurate reconstruction across most of the defined Cu(II) domain without parameter labels or clean-spectrum targets N.

Cross-Engine Library Transfer

TorchSpin can generate labeled spectral libraries for conventional supervised learning, and CNNs trained on TorchSpin-generated libraries transferred to independently generated EasySpin spectra with median reconstruction accuracy similar to that of the EasySpin-trained controls N. The source of the training library had little effect on the central performance of the inverse model, with a median cosine similarity to clean EasySpin test spectra averaging 0.966 for TorchSpin-trained models and 0.967 for the EasySpin-trained references N.

TorchSpin also supports conventional cw EPR fitting through an interactive Jupyter interface modeled on the workflow of EasySpin’s esfit N, allowing users to adjust initial parameter values and monitor the fitted spectrum within the notebook N. TorchSpin provides a practical route for bringing established EPR simulation into the PyTorch computational environment, where the principal advantage is that the same physical model can generate training data, evaluate predicted parameters, and propagate spectral disagreement through an optimization or learning system N.

Conclusions

TorchSpin brings representative EPR simulation methods into a differentiable PyTorch framework while preserving the organization and nomenclature familiar to EasySpin users N. The more consequential capability is that supported forward calculations can participate directly in automatic differentiation N, allowing spectral disagreement to refine a supervised inverse model and to train models from random initialization without parameter labels or clean spectral targets N. The accompanying code, validation data, benchmarks and notebooks provide a reproducible foundation for extending this approach to additional EPR methods and, importantly, for testing it against experimental spectra N.

The paper's work establishes that TorchSpin is competitive across diverse workloads but that neither CPU parallelism nor hybrid CUDA execution provides a universal speed advantage N. The final CNN required a median of about 0.976 ms to analyze one spectrum using a single CPU thread and 1.05 ms on an RTX 4090 when spectrum normalization and host-to-device transfer were included N. The practical value of the tensor implementation extends beyond forward-simulation speed, as the forward calculation can serve as a differentiable physical constraint, a decoder for parameter-label-free learning, or a generator of synthetic libraries for reusable inverse models N.

The methods section details the software implementation in Python using PyTorch for tensor operations, while retaining NumPy and SciPy for numerical procedures not yet converted to tensor-native form N. Performance benchmarking was conducted on a local workstation equipped with an AMD Ryzen Threadripper PRO 5995WX processor and two NVIDIA GeForce RTX 4090 GPUs N. The Cu(II) spectral libraries were generated independently with EasySpin pepper and contained 4,096 training, 512 validation and 512 prospective test spectra N. The CNN architecture used a one-dimensional residual CNN containing 1,297,127 trainable parameters N. The differentiable refinement objective combined first-derivative cosine loss, integrated-spectrum cosine loss and peak-weighted mean-squared error with relative weights of 1.0, 0.25 and 0.25 respectively N. The inference timing shows that a batch of 512 spectra was processed in approximately 1.3 ms with GPU acceleration N, while TorchSpin forward reconstruction required an additional 68.4 ms N. The results confirm that TorchSpin brings representative EPR simulation methods into a differentiable PyTorch framework N. The paper's organization and program names are derived directly from EasySpin, preserving a familiar framework while exposing supported calculations to tensor batching, automatic differentiation and machine-learning tools N. The accompanying code, validation data, benchmarks and notebooks provide a reproducible foundation for extending this approach to additional EPR methods and, importantly, for testing it against experimental spectra N.

1. Bold header: Differentiable spectral feedback for supervised inverse models. The system can propagate spectral disagreement through the simulation to update either the predicted parameters or the network itself, allowing for a second form of training and improving spectral reconstruction quality, as shown by an increase in the fraction of spectra reaching a cosine similarity of at least 0.950 from 76.6% to 90.2%.

2. Bold header: Parameter-label-free training without parameter supervision. The improved system can train the same inverse model from random initialization without parameter labels or clean spectral targets, achieving a median cosine similarity of 0.992 on test spectra, demonstrating that the simulator can serve as a physical decoder for parameter-label-free learning.

3. Bold header: Cross-engine library transfer capability. The system can analyze spectra generated independently by different simulators without retraining, as CNNs trained on one library transferred without adaptation to independently generated EasySpin spectra with median reconstruction accuracy similar to that of the EasySpin-trained controls.

4. Bold header: Rapid learned inference and parameter estimation. Once trained, the CNN can produce initial parameter estimates in approximately 1 ms per spectrum, enabling a rapid screening of large spectral series or providing informed starting parameters before conventional fitting is applied.

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 at least 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.