TorchSpin：在 PyTorch 中进行可微分电子顺磁共振模拟

arXiv:2610.18924v1 公告类型：交叉发布 摘要：解读电子顺磁共振（EPR）谱通常需要反复进行正向模拟和自旋哈密顿量参数优化。现有软件提供了成熟的模拟与拟合方法，但这些计算通常无法作为可微分组件用于机器学习工作流。我们提出 TorchSpin，这是一个开源 PyTorch 框架，将 EasySpin 的组织方式和数值计算约定移植到基于张量的计算环境中。涵盖连续波 EPR、电子核双共振（ENDOR）、脉冲 EPR、自旋动力学、轨迹模拟和磁测量的代表性计算均高度复现了相应的 EasySpin 结果。在条件匹配的基准测试中，TorchSpin 在九项工作负载中的六项上 CPU 执行速度快于 EasySpin，但性能高度依赖于具体计算。因此，我们重点关注 PyTorch 自动微分所带来的能力。在一项 Cu(II) 粉末谱演示中，经由 TorchSpin 传播的反馈提升了监督式卷积神经网络（CNN）的质量，且在优化过程中未使用参数标签。从随机初始化的 CNN 出发，仅使用含噪声谱和 TorchSpin 重建损失进行训练的模型，其相似度中位数达到 $\geq$ 0.99。在 TorchSpin 生成的谱库上训练的 CNN 也能无需适配直接迁移至独立生成的 EasySpin 谱，性能与在 EasySpin 数据上训练的对照模型相当。训练完成后，CNN 生成初始参数估计的耗时约为每条谱 1 ms，而借助 GPU 加速，处理包含 512 条谱的一个批次仅需约 1.3 ms。因此，TorchSpin 不仅提供了经过测试的成熟 EPR 计算方法的 Python 实现，还提供了一个可微分物理层，用于快速的学习式推断、谱优化以及无需参数标签的训练。