SpaiNN: Equivariant Message Passing for Excited-State Nonadiabatic Molecular Dynamics

Excited-state molecular dynamics simulations are crucial for understanding processes like photosynthesis, vision, and radiation damage. However, the computational complexity of quantum chemical calculations restricts their scope. Machine learning (ML) offers a solution by delivering high accuracy properties at lower computational costs. We present SpaiNN, an open-source Python software for ML-driven surface hopping nonadiabatic molecular dynamics simulations. SpaiNN combines the invariant and equivariant neural network architectures of SchNetPack with SHARC for surface hopping dynamics. Its modular design allows users to implement and adapt modules easily. We compare rotationally-invariant and equivariant representations in fitting potential energy surfaces of multiple electronic states and properties arising from the interaction of two electronic states. Simulations of the methylenimmonium cation and various alkenes demonstrate the superior performance of equivariant SpaiNN models, improving accuracy, generalization, and efficiency in both training and inference.