Dissipative hamiltonian neural network for exact mass conservation

[jakobgradl]

The proposed changes include the implementation of two mass conservation models and two mass-conserving SSA models as dissipative hamiltonian neural networks.

The implementation is based on a Helmholtz decomposition of the 2-dimensional mass flux vector field Hv̄ into a divergent, curl-free component d and a rotational, divergence-free component r. The two components are modelled based on two scalar potential fields D and R that replace variables u, v, and a in the network output. These variables are instead computed from the potential fields like this:

Hv̄ = d + r = grad(D) + J grad(R) = [D_x, D_y] + [-R_y, R_x]
with
div(d) = div(Hv̄) and curl(d) = 0
curl(r) = curl(Hv̄) and div(r) = 0

u = (D_x - R_y) / H
v = (D_y + R_x)/ H
a = D_xx + D_yy
with _x and _xx indicating the first and second derivative in the indicated direction, and with u,v being the depth-integrated velocity components and a being the mass flux divergence.

Constructing the mass flux and velocity field in this way exactly satisfies the 2-dimensional continuity equation. Evaluating the residual of the continuity equation in the physics loss of the network is thus no longer necessary. Rather, the equations are entirely implemented as data misfit functions, currently for misfit with observed velocity and dHdt.

References:
Bente et al. 2025, http://arxiv.org/abs/2510.06286
Sosanya and Greydanus 2022, http://arxiv.org/abs/2201.10085
Schoder et al. 2020, https://pmc.ncbi.nlm.nih.gov/articles/PMC7648005/
Bathia et al. 2013, https://ieeexplore.ieee.org/abstract/document/6365629

[enigne]

Hi @jakobgradl, thank you so much for the implementation, it looks pretty good. Before we merge it to the main branch, we need to resolve a few issues:

1. CI for Jax is complaining about how the second derivative is handled in pinnicle/physics/stressbalance_MC.py:133, an example can be found in pinnicle/physics/stressbalance.py:124

2. CI for paddle is complaining about "The Op divide_double_grad doesn't have any grad op", which is usually related to the change here

3. After reviewing the code, I realized not all of them are covered by unit test, could you add a few more test to cover every line of code you contributed. An easy way is to use codecov and see the coverage.