[WIP] Nonhydrostatic free surface: CG solver support for IBG and stretched-horizontal grids

[xkykai]

After meeting and having discussions with @shriyafruitwala on #4740 I think it is straightforward to extend the nonhydrostatic free surface to other grids and setups as well, utilizing the CG solver to either deal with immersed boundaries or stretched grids in non-z directions.

Summary

NonhydrostaticModel supports an implicit free surface via a Robin (mixed) boundary condition on the 3D pressure Poisson equation. The Robin BC modifies the rigid-lid operator at k = Nz:

• Diagonal correction: subtract Azᶜᶜᶠ(i,j,Nz+1) / den * ϕ[i,j,Nz] (eliminates the null space, so no gauge condition is needed)
• Inhomogeneous RHS term: subtract Azᶜᶜᶠ(i,j,Nz+1) * g * Δt * η★ / den where η★ = η + Δt * w̃ and den = g*Δt² + Δzᶠ/2

Previously this was handled only by FourierTridiagonalPoissonSolver with InhomogeneousFormulation(ZDirection()), which bakes both terms into the direct solve. This PR extends implicit free surface support to grids and configurations where that direct approach is not applicable.

Bug fixed: incorrect dispatch for XZRegularRG/YZRegularRG + free surface

The old dispatch Union{GridWithFourierTridiagonalSolver, XYZRegularRG} + free_surface applied InhomogeneousFormulation(ZDirection()) to all GridWithFourierTridiagonalSolver types, which includes XZRegularRG (stretched y) and YZRegularRG (stretched x). This is invalid — those grids have a non-uniform horizontal direction that cannot be handled by FFT in the z-tridiagonal formulation.

New cases covered

Grid	Free surface	Strategy
XYZRegularRG or XYRegularRG	non-null	Unchanged: FourierTridiagonalPoissonSolver(InhomogeneousFormulation(ZDirection()))
XZRegularRG (stretched y)	non-null	New: ConjugateGradientPoissonSolver with FreeSurfaceLaplacian + FT (HomogNeumann y-dir) preconditioner
YZRegularRG (stretched x)	non-null	New: ConjugateGradientPoissonSolver with FreeSurfaceLaplacian + FT (HomogNeumann x-dir) preconditioner
IBGWithFFT (any FFT-compatible underlying)	non-null	New: ConjugateGradientPoissonSolver with FreeSurfaceLaplacian + fft_free_surface_preconditioner(underlying)

For IBG cases, fft_free_surface_preconditioner selects the best available FT preconditioner: InhomogeneousFormulation(ZDirection()) when the underlying grid has both x and y uniform (so the preconditioner directly encodes the Robin BC), or HomogeneousNeumann FT otherwise.

Key additions

FreeSurfaceLaplacian (callable struct, Solvers/conjugate_gradient_poisson_solver.jl)

CG linear operator for the free-surface pressure equation. For each CG iteration it:

1. Fills halos then overrides the top ghost with the Neumann value ϕ[Nz+1] = ϕ[Nz], preventing the MixedBoundaryCondition on p_Δt (set at model construction from the previous time step) from polluting the operator
2. Applies V∇² via the existing _symmetric_laplacian_operator! kernel
3. Adds the Robin diagonal correction -Az/den * ϕ[Nz]

solve_for_pressure! split (CG dispatch, solve_for_pressure.jl)

ConjugateGradientPoissonSolver now has separate dispatches for ::Nothing (rigid lid, unchanged) and free_surface (new). The free-surface dispatch:

• Builds the V∇²-scaled RHS from divergence, then subtracts the Robin inhomogeneous term Az * g * Δt * η★ / den at k = Nz
• Updates the FT preconditioner's diagonal at k = Nz via update_fourier_tridiagonal_solver! (guarded against HomogeneousNeumann FT preconditioners and non-FT preconditioners)
• Calls solve!(pressure, cg_solver, rhs, free_surface, Δt), propagating (free_surface, Δt) through to FreeSurfaceLaplacian

Dispatch table redesign to avoid Julia method ambiguity

XYRegularRG <: XZRegularRG (a grid with both x and y regular is a subtype of "x regular"). This means mixing union dispatches on the grid type with per-type dispatches on ::Nothing vs free_surface creates cross-dimension ambiguities that Julia cannot resolve. The fix is to use per-type dispatches throughout — one method per concrete grid type for each (grid_type, free_surface_type) combination. The same treatment was applied to the Distributed solver dispatches.

Test plan

• All 16 dispatch tests pass (all grid type × free surface combinations)
• Full NonhydrostaticModel construction with ImplicitFreeSurface() on XYZRegularRG (regression)
• Full NonhydrostaticModel construction with ImplicitFreeSurface() on IBGWithFFT → uses ConjugateGradientPoissonSolver with FreeSurfaceLaplacian
• End-to-end time stepping with XZRegularRG + ImplicitFreeSurface (blocked by pre-existing materialize_free_surface limitation: the hydrostatic FFTImplicitFreeSurfaceSolver construction inside ImplicitFreeSurface fails for non-uniform-in-horizontal grids before our new pressure solver dispatch is reached)

🤖 Generated with Claude Code