The semi-implicit free-surface solve

The semi-implicit core couples surface elevation and depth-averaged transport through an elliptic equation. Its forward solve uses preconditioned conjugate gradients (PCG), with a discrete sine transform preconditioner in the published configuration. Its backward pass differentiates the solved system implicitly. (Manuscript, Section 2.1)

From a surface-pressure response to a linear system

Let $\eta$ be the unknown surface elevation, $D$ the depth entering the operator, ${\mathrm{\nabla }}_h$ the horizontal gradient, and $\kappa$ the positive coefficient determined by the time discretization. With $b$ collecting known forcing, history, and boundary contributions, the operator has the form

$$A(D)\eta =b,A(D)\eta =\eta -\kappa {\mathrm{\nabla }}_h\cdot (D{\mathrm{\nabla }}_h\eta ).$$

The C-grid implementation evaluates face fluxes and their divergence with the grid metrics and masks. This connects adjacent surface cells without explicitly assembling a dense matrix. Exact coefficients and boundary handling are defined in implicit_barotropic.py.

Forward solve and preconditioning

PCG repeatedly applies the discrete operator and a preconditioner to reduce the residual. The preconditioner accelerates convergence; the physical operator still defines the solved equation. The portable semi-implicit configurations select the registered Triton backend and compiled vertical operators. The published coarse physical step is 600 s. (Core configuration, portable examples)

Differentiate the solution

Let $\mathcal{L}$ be a scalar loss, $\lambda$ an adjoint vector, and $A^{\mathrm{\top }}$ the transpose of the discrete linear operator under the chosen algebraic representation. Given the loss derivative with respect to the solved elevation, the adjoint system is

$$A^{\mathrm{\top }}\lambda =\frac{\partial \mathcal{L}}{\partial \eta }.$$

For variations $\mathrm{d}b$ and $\mathrm{d}A$, the corresponding loss variation is

$$\mathrm{d}\mathcal{L}={\lambda }^{\mathrm{\top }}(\mathrm{d}b-\mathrm{d}A\eta ).$$

This relation propagates sensitivity through the solved equation. The production backward operator includes its geometry and boundary conventions; it is implemented in autograd.py. The sea-surface adjoint checks are in tests/test_fixed_core_runtime.py.

Why Midtrain needs this derivative

An earlier correction changes the state entering the next physical interval. The free-surface solve is one part of that physical response. Midtrain uses its derivative to connect the later loss back to the earlier network parameters. See Staged training for the complete gradient path.