Learning Neural Differential Algebraic Equations via Operator Splitting

Abstract

Differential algebraic equations (DAEs) describe the temporal evolution of systems that obey both differential and algebraic constraints. Of particular interest are systems that contain implicit relationships between their components, such as conservation laws. Here, we present an Operator Splitting (OS) numerical integration scheme for learning unknown components of DAEs from time-series data. In this work, we show that the proposed OS-based time-stepping scheme is suitable for relevant system-theoretic data-driven modeling tasks. Presented examples include (i) the inverse problem of tank-manifold dynamics and (ii) discrepancy modeling of a network of pumps, tanks, and pipes. Our experiments demonstrate the proposed method's robustness to noise and extrapolation ability to (i) learn the behaviors of the system components and their interaction physics and (ii) disambiguate between data trends and mechanistic relationships contained in the system.

Problem Formulation

Differential-Algebraic Equations (DAEs) model dynamical systems governed by both differential dynamics and algebraic constraints—common in physical domains with conservation laws or implicit relationships. Traditional neural ODE models fail to capture these algebraic structures, motivating the need for neural architectures that can learn both explicit and implicit dynamics. The goal is to develop a data-driven framework for learning unknown or partially known DAE components from time-series data while maintaining consistency with the system’s inherent physical constraints.
We formulate the problem as learning two mappings $ f_\theta $ and $ g_\phi $ that represent the differential and algebraic components of the DAE system in the following parameter estimation problem:

$$ \begin{aligned} &\underset{\theta_f}{\text{minimize}} && \int_{t_0}^{t_N} \left( \|x(t) - \hat{x}(t)\|_2^2 + \|y(t) - \hat{y}(t)\|_2^2 \right) dt \\\\ &\text{subject to} && \dot{x}(t) = f(x(t), y(t), u(t); \theta_f), \\\\ & && 0 = g(x(t), y(t), u(t)), \\\\ & && x(t_0) = x_0, \quad y(t_0) = y_0, \end{aligned} $$

where $ x(t) \in \mathbb{R}^{n_x} $ are the differential states, $ y(t) \in \mathbb{R}^{n_z} $ are algebraic states, $ u(t) \in \mathbb{R}^{n_u} $ are known inputs, and $ \theta_f $ are the learnable parameters.

Method

We propose a modular neural DAE learning framework that decomposes the DAE system into two components: a neural surrogate for the algebraic solver, and a physics-constrained (gray-box) neural ordinary differential equation (ODE) model. The algebraic states are updated using a trainable neural network, while the differential states evolve via an ODE solver applied to a structured ODE.
The proposed operator splitting strategy is inspired by fractional-step numerical integration and is defined as:

$$ \begin{aligned} y^{(t+\Delta t)} & = h\left( x^{(t)}, y^{(t)}, u^{(t)}; \theta_h \right), \\\\ x^{(t+\Delta t)} &= \text{ODESolve}\left(f, \{ x^{(t)}, y^{(t+\Delta t)}, u^{(t)} \}; \theta_f \right), \end{aligned} $$

where $ h: \mathbb{R}^{n_x} \times \mathbb{R}^{n_y} \times \mathbb{R}^{n_u} \rightarrow \mathbb{R}^{n_y} $ is a neural surrogate for the algebraic update, $ f: \mathbb{R}^{n_x} \times \mathbb{R}^{n_y} \times \mathbb{R}^{n_u} \rightarrow \mathbb{R}^{n_x} $ is a neural ODE model, and $\text{ODESolve}$ denotes a numerical ODE solver (e.g., explicit Runge-Kutta).
This architecture enables end-to-end training of neural DAEs with a constraint-aware loss function:

$$ \mathcal{L}(\Theta) = \lambda_1 \mathcal{L}_{\text{residual}}(X, \hat{X}; \Theta) + \lambda_2 \mathcal{L}_{\text{constraints}}(X, \hat{X}; \Theta), $$

where $ \mathcal{L}_{\text{residual}} $ measures the error in predicted states, and $ \mathcal{L}_{\text{constraints}} $ penalizes violations of the algebraic constraints. Here, $ X $ are the model-predicted trajectories, and $ \hat{X} $ are the corresponding observations from the training dataset.

Examples

Our Neural DAE framework and all our examples are implemented using the Neuromancer library.

Let's consider tank-manifold-pump system governed by differential-algebraic equations. The goal is to learn the nonlinear area-height profile of a tank, and the implicit flow splitting rule at a manifold junction, only from observed time series of tank heights and input flow rates.

Try this example interactively in Google Colab

In this example, the full DAE system is given by:

$$ \begin{aligned} \frac{dx_1}{dt} &= \frac{y_1}{3}, \\ \frac{dx_2}{dt} &= \frac{y_2}{\sqrt{x_2} + 0.1}, \\ 0 &= u - y_1 - y_2, \\ 0 &= x_1 - x_2, \end{aligned} $$

where $x_1, x_2$ are the tank heights and $y_1, y_2$ are the outlet flows. The model enforces mass conservation and equal pressure heads through algebraic constraints, while learning $\phi_2(x_2) = \sqrt{x_2} + 0.1$ using a neural network.

Figure: Schematic of tank-manifold-pump system.

The learned model accurately reconstructs the tank heights and flow rates under both training and extrapolation conditions.

Figure: (a) Tank heights and (b) flow rates during training trajectory with constant inflow through the manifold.

Figure: Extrapolation of learned area-height relationship and resulting flow trajectories under unseen varying inflow through the manifold.

The model demonstrates the ability to learn interpretable component dynamics while satisfying algebraic constraints and to generalize beyond the training regime.

Conclusions and Related Work

We presented a data-driven framework for learning Differential-Algebraic Equation (DAE) systems using a novel operator splitting strategy. By alternating updates between a neural surrogate for algebraic states and a physics-informed ODE solver for differential dynamics, the method enables consistent modeling of constrained systems directly from time-series data. Our experiments show robustness to noise and extrapolation, with the ability to learn interpretable physical behaviors and interactions among system components.

Future research directions include:

Benchmarking against baseline methods such as neural ODEs and PINNs.
Adapting the framework for stiff or discontinuous dynamics via implicit solvers.
Extending to higher-index DAEs common in optimal control and system design.
Handling of partial observability for large-scale systems.

Related work on neural DAEs:
Machine Learning with Hard Constraints: Neural DAEs as a General Formalism
A Simultaneous Approach for Training Neural Differential-Algebraic Systems of Equations

Acknowledgment

This research was supported by the U.S. Department of Energy through the Building Technologies Office under the Advancing Market-Ready Building Energy Management by Cost-Effective Differentiable Predictive Control projects. Pacific Northwest National Laboratory (PNNL) is a multi-program national laboratory operated for the U.S. Department of Energy (DOE) by Battelle Memorial Institute under Contract No. DE-AC05-76RL0-1830.

This work was also supported by the Ralph O’Connor Sustainable Energy Institute (ROSEI) at Johns Hopkins University.

BibTeX

@inproceedings{koch2025, title={Learning Neural Differential Algebraic Equations via Operator Splitting}, author={James Koch and Madelyn Shapiro and Himanshu Sharma and Draguna Vrabie and Jan Drgona}, year={2025}, eprint={2403.12938}, archivePrefix={arXiv}, primaryClass={cs.LG}, url={https://arxiv.org/abs/2403.12938}, volume={}, number={}, booktitle={Conference on Decision and Control (CDC)}, }