[Paper Review] $ξ$-torch: differentiable scientific computing library

Muhammad Kasim, S. M. Vinko|arXiv (Cornell University)|Oct 5, 2020

Model Reduction and Neural Networks14 references4 citations

TL;DR

This paper introduces $\xi$-torch, a PyTorch-based library that provides differentiable functionals for scientific computing, such as root finders, ODE solvers, and quadrature methods, with analytically derived gradients for improved numerical stability and memory efficiency. It enables end-to-end differentiable physics-informed machine learning by supporting first- and higher-order derivatives, demonstrated in mirror design and molecular dynamics optimization with rapid convergence.

ABSTRACT

Physics-informed learning has shown to have a better generalization than learning without physical priors. However, training physics-informed deep neural networks requires some aspect of physical simulations to be written in a differentiable manner. Unfortunately, some operations and functionals commonly used in physical simulations are scattered, hard to integrate, and lack higher order derivatives which are needed in physical simulations. In this work, we present $ξ$-torch, a library of differentiable functionals for scientific simulations. Example functionals are a root finder and an initial value problem solver, among others. The gradient of functionals in $ξ$-torch are written based on their analytical expression to improve numerical stability and reduce memory requirements. $ξ$-torch also provides second and higher order derivatives of the functionals which are rarely available in existing packages. We show two applications of this library in optimizing parameters in physics simulations. The library and all test cases in this work can be found at https://github.com/xitorch/xitorch/ and the documentation at https://xitorch.readthedocs.io.

Motivation & Objective

• To address the lack of integrated, differentiable scientific functionals in existing libraries for physics-informed deep learning.
• To improve numerical stability and reduce memory usage in scientific simulations by using analytical derivatives instead of automatic differentiation on forward algorithms.
• To enable higher-order derivatives (e.g., second-order) for advanced scientific simulations, such as vibrational frequency calculations.
• To facilitate the integration of deep learning with physics simulations by providing a unified, modular library with pure function and linear operator abstractions.
• To support optimization of physical system parameters by enabling gradient flow through complex simulation components like ODE solvers and root finders.

Proposed method

• Implementing functionals such as root finders, ODE solvers (solve_ivp), and quadrature methods (mcquad, SQuad) with analytically derived gradients to enhance numerical stability and reduce memory consumption.
• Using pure functions as inputs to ensure correct higher-order derivative computation, with a mechanism to convert stateful object methods into pure functions by tracking state variables affecting outputs.
• Introducing an EditableModule base class to manually specify parameter names affecting method outputs, enabling accurate gradient computation in complex classes.
• Providing a LinearOperator abstraction for sparse linear algebra operations to avoid expensive matrix construction in high-dimensional problems.
• Leveraging PyTorch’s automatic differentiation engine with analytical derivatives to compute second and higher-order derivatives via gradgradcheck validation.
• Designing a modular library structure with dedicated modules for optimization, integration, interpolation, and linear algebra, inspired by SciPy but tailored for differentiable scientific computing.

Experimental results

Research questions

• RQ1 Can differentiable scientific functionals be implemented with analytical gradients to improve numerical stability and reduce memory usage in scientific simulations?
• RQ2 How can higher-order derivatives (e.g., second-order) be efficiently and correctly computed for complex functionals like ODE solvers and root finders?
• RQ3 To what extent can a unified library of differentiable functionals accelerate physics-informed machine learning in scientific applications?
• RQ4 How can stateful object methods be transformed into pure functions to enable correct gradient computation in differentiable scientific computing?
• RQ5 Can $\xi$-torch enable effective parameter optimization in physical systems, such as mirror design and molecular dynamics, through end-to-end differentiable simulation?

Key findings

• The $\xi$-torch library successfully computes first- and higher-order derivatives of scientific functionals using analytical expressions, improving numerical stability and reducing memory usage compared to backpropagation through algorithmic implementations.
• The library achieves correct second-order gradients, validated via PyTorch’s gradgradcheck, demonstrating robustness in complex differentiable functionals.
• In the mirror design case, the root finder functional enabled optimization of a neural network-based mirror surface in 200 iterations, achieving accurate focusing of light rays at the target point.
• In molecular dynamics, the solve_ivp functional allowed optimization of initial velocities to form the letters 'OX' at t=1, with convergence achieved in 7500 iterations using Adam.
• The library’s modular design and support for pure functions and editable state variables enable reliable gradient computation in complex physical simulation pipelines.
• The availability of differentiable functionals like Monte Carlo quadrature and fixed-samples quadrature expands the scope of differentiable scientific computing beyond standard ODE and root-finding tasks.