Direct Computation of Viscosity from Differentiable Atomistic Simulations

Rayan Chatterjee, Srikanth Sastry, N M Anoop Krishnan

Keywords: differentiable simulation, viscosity, jax-md

TL;DR: We propose a method to compute viscosity using differentiable simulation

Abstract: Shear viscosity calculation from molecular dynamics simulations demands long equilibration times and extensive statistical averaging to achieve convergence. We address this challenge by presenting a direct automatic-differentiation method to compute shear viscosity from differentiable molecular dynamics simulations. Our approach differentiates microscopic shear stress with respect to applied shear rate, but crucially identifies a characteristic timescale $\tau_\alpha$ that defines a stable window for reliable gradient computation. We demonstrate that $\tau_\alpha$ marks the onset of chaotic divergence in stress dynamics and corresponds to the timescale where stress autocorrelations decay by $>90\%$, providing both theoretical justification and physical insight. Through systematic validation on Weeks-Chandler-Andersen systems across multiple realizations, our method yields viscosity estimates ($1.92 \pm 0.38$) that agree with Green-Kubo predictions ($2.24 \pm 0.24$) within statistical uncertainty, while circumventing the noise accumulation inherent in long-time correlation approaches. The identified stability window concept establishes a general framework for extracting transport properties from differentiable simulations of chaotic systems before gradient explosion occurs, with promising applications to thermal conductivity and diffusion coefficients. This work provides a principled solution to gradient instability in differentiable physics, enabling reliable parameter optimization and property prediction in complex molecular systems.

Submission Track: Paper Track (Short Paper)

Submission Category: AI-Guided Design

Institution Location: New Delhi, India

AI4Mat Journal Track: Yes

AI4Mat RLSF: Yes

Submission Number: 126