Open source Differentiable ODE Solving Infrastructure

TL;DR

This paper introduces an open-source, GPU-accelerated, differentiable ODE solving infrastructure integrated into DeepChem, enabling accurate, scalable solutions for complex dynamic systems across various scientific domains.

Contribution

It presents a flexible, GPU-accelerated, fully differentiable ODE solver integrated into DeepChem, supporting multiple methods and large systems, with demonstrated high accuracy and scalability.

Findings

Achieved mean squared errors between 10^{-4} and 10^{-6}

Demonstrated scalability to systems with up to 100 compartments

Validated on diverse applications including predator-prey, pharmacokinetics, neural ODEs, and PDEs

Abstract

Ordinary Differential Equations (ODEs) are widely used in physics, chemistry, and biology to model dynamic systems, including reaction kinetics, population dynamics, and biological processes. In this work, we integrate GPU-accelerated ODE solvers into the open-source DeepChem framework, making these tools easily accessible. These solvers support multiple numerical methods and are fully differentiable, enabling easy integration into more complex differentiable programs. We demonstrate the capabilities of our implementation through experiments on Lotka-Volterra predator-prey dynamics, pharmacokinetic compartment models, neural ODEs, and solving PDEs using reaction-diffusion equations. Our solvers achieved high accuracy with mean squared errors ranging from $10^{-4}$ to $10^{-6}$ and showed scalability in solving large systems with up to 100 compartments.