geoswe 1.0.0__py3-none-any.whl

geoswe/__init__.py

@@ -0,0 +1,102 @@
+"""GeoSWE: Geophysical Shallow-Water Engine.
+
+A finite-volume solver for the 2D nonlinear shallow-water equations, built for
+flood modeling from county to continental scale. It runs on NVIDIA and AMD GPUs
+through CuPy, across GPUs with mpi4py, and on the CPU with NumPy.
+
+Highlights
+----------
+* HLLC and local Lax-Friedrichs numerical fluxes.
+* Well-balanced schemes: the Xia (2017) surface-reconstruction method (SRM)
+  and Audusse hydrostatic reconstruction; exact lake-at-rest preservation.
+* Implicit Manning friction, rainfall, wetting and drying, and prescribed
+  water levels (tide, surge).
+* A compressed active-cell mesh (GPU) that stores and updates only the cells
+  of the modeled region, reaching continental scale on one node (Florida at
+  10 m, the conterminous United States at 30 m).
+
+Quick start
+-----------
+>>> import numpy as np
+>>> from geoswe import Mesh2D, Config, Solver2D
+>>> nx, ny = 200, 160
+>>> mesh = Mesh2D(nx=nx, ny=ny, dx=1.0, dy=1.0)
+>>> cfg = Config()        # first-order HLLC + SRM, forward Euler, CFL 0.5
+>>> q0 = np.zeros((3, nx, ny)); q0[0] = 1.0             # h, hu, hv
+>>> q0[0, 80:120, 60:100] = 2.0                         # a column of water
+>>> s = Solver2D(mesh, cfg, q0, np.zeros((nx, ny)))
+>>> _ = s.run(t_end=1.0)
+>>> h = s.depth()                                       # NumPy array (nx, ny)
+
+GeoSWE uses the GPU when CuPy and a CUDA device are present and NumPy
+otherwise. Set the ``GEOSWE_BACKEND`` environment variable to ``"numpy"`` or
+``"cupy"`` before importing to choose.
+"""
+from __future__ import annotations
+
+# --- always importable: pure-NumPy-capable core ---------------------------
+from .backend import (
+    xp,
+    set_backend,
+    get_backend,
+    gpu_platform,
+    to_host,
+    to_device,
+    sync,
+    USING_CUPY,
+)
+from .swe import G, H_MIN
+from .mesh import Mesh1D, Mesh2D
+from .solver import Config, Solver1D, Solver2D
+from .forcing import RainfallForcing, StageBoundary
+
+__version__ = "1.0.0"
+
+__all__ = [
+    # backend control
+    "xp",
+    "set_backend",
+    "get_backend",
+    "gpu_platform",
+    "to_host",
+    "to_device",
+    "sync",
+    "USING_CUPY",
+    # physical constants / defaults
+    "G",
+    "H_MIN",
+    # meshes
+    "Mesh1D",
+    "Mesh2D",
+    # configuration + solvers
+    "Config",
+    "Solver1D",
+    "Solver2D",
+    # forcings
+    "RainfallForcing",
+    "StageBoundary",
+    # NOTE: the GPU-only CompressedSolver is exposed lazily via __getattr__ but
+    # deliberately NOT listed in __all__, so `from geoswe import *` works on
+    # CPU-only installs.
+    "__version__",
+]
+
+
+def __getattr__(name):
+    """Lazily expose the GPU-only compressed solver.
+
+    ``geoswe.compressed_solver`` hard-imports CuPy at module load, so importing
+    it eagerly would break ``import geoswe`` on a CPU-only machine. Resolve it
+    only when the caller actually asks for ``geoswe.CompressedSolver``.
+    """
+    if name == "CompressedSolver":
+        try:
+            from .compressed_solver import CompressedSolver
+        except ImportError as e:  # clear message on CPU-only installs
+            raise ImportError(
+                "geoswe.CompressedSolver needs CuPy and SciPy (the gpu extra): "
+                "pip install 'cupy-cuda12x[ctk]' scipy, or on an AMD GPU (the gpu-rocm "
+                "extra): pip install 'cupy-rocm-7-0' scipy"
+            ) from e
+        return CompressedSolver
+    raise AttributeError(f"module {__name__!r} has no attribute {name!r}")

geoswe/backend.py

@@ -0,0 +1,335 @@
+"""Backend abstraction: numpy on CPU, cupy on GPU.
+
+Usage:
+    from .backend import xp, to_host, sync, USING_CUPY
+    a = xp.zeros((100, 100))   # numpy or cupy array depending on backend
+
+The reliable way to choose the backend is the ``GEOSWE_BACKEND`` environment
+variable ("cupy" or "numpy"), set before ``import geoswe``; several modules
+specialise for the backend at import time. ``set_backend()`` works only while
+no solver module has been imported yet; afterwards it raises.
+
+Once set, all modules that import `xp` from this file will use the selected backend.
+
+CuPy drives NVIDIA GPUs through CUDA and AMD GPUs through ROCm/HIP, and
+``gpu_platform()`` says which one is in use. The kernels are written in CUDA C;
+build them with ``raw_kernel()`` / ``raw_module()`` / ``elementwise_kernel()`` rather
+than the CuPy classes so the ROCm compiler receives a source and options it accepts.
+"""
+from __future__ import annotations
+
+import os
+import re
+import warnings
+
+USING_CUPY = False
+GPU_PLATFORM = None   # "cuda" (NVIDIA), "hip" (AMD ROCm) or None (NumPy backend)
+_xp_module = None
+
+
+# CuPy ships one wheel per CUDA major (cupy-cuda12x, cupy-cuda13x, ...) and every one
+# installs the same `cupy` package, so two of them overwrite each other's files. The
+# usual way in: install the default CUDA 12 build, then add the CUDA 13 one on top.
+# The ROCm builds (cupy-rocm-7-0 on PyPI, amd-cupy on AMD's index) collide the same way.
+_CUPY_BUILD = re.compile(r"(amd-)?cupy(-cuda\d+x|-rocm-[\d-]+)?")
+
+
+def _installed_cupy_builds():
+    """Names of the installed CuPy distributions (normally zero or one)."""
+    from importlib import metadata
+    names = set()
+    for dist in metadata.distributions():
+        name = (dist.metadata.get("Name") or "").lower().replace("_", "-")
+        if _CUPY_BUILD.fullmatch(name):
+            names.add(name)
+    return sorted(names)
+
+
+def _check_single_cupy():
+    """Refuse the GPU backend when more than one CuPy build is installed."""
+    if os.environ.get("GEOSWE_ALLOW_MULTIPLE_CUPY") == "1":
+        return
+    builds = _installed_cupy_builds()
+    if len(builds) > 1:
+        if any("rocm" in b or b.startswith("amd-") for b in builds):
+            keep = ("pip install 'cupy-cuda12x[ctk]' for an NVIDIA GPU or 'cupy-rocm-7-0' for an "
+                    "AMD one (the builds the gpu and gpu-rocm extras install). ")
+        else:
+            keep = ("pip install 'cupy-cuda12x[ctk]'. That build, which the gpu extra installs, "
+                    "runs on both CUDA 12 and CUDA 13 drivers. ")
+        raise RuntimeError(
+            f"several CuPy builds are installed ({', '.join(builds)}). They overwrite each "
+            f"other in the same `cupy` package, and kernels then fail in confusing ways. "
+            f"Keep one: pip uninstall -y {' '.join(builds)} && {keep}"
+            f"(GEOSWE_ALLOW_MULTIPLE_CUPY=1 skips this check.)")
+
+
+def _platform_of(cp):
+    """Which GPU stack this CuPy build targets: "hip" for AMD ROCm, else "cuda"."""
+    try:
+        return "hip" if cp.cuda.runtime.is_hip else "cuda"
+    except AttributeError:   # CuPy too old to have a ROCm build
+        return "cuda"
+
+
+def _cuda_device_count(cp):
+    """Number of visible GPUs (CuPy counts AMD devices through the same call); 0 when the
+    driver or every device is missing."""
+    try:
+        return int(cp.cuda.runtime.getDeviceCount())
+    except Exception:          # CUDARuntimeError: no driver, no device
+        return 0
+
+
+def _no_device(cp):
+    """What to say when this CuPy build sees no device."""
+    if _platform_of(cp) == "hip":
+        return "no AMD GPU is visible (check rocm-smi and HIP_VISIBLE_DEVICES)"
+    return "no CUDA device is visible (check nvidia-smi and CUDA_VISIBLE_DEVICES)"
+
+
+def _select_default():
+    global _xp_module, USING_CUPY, GPU_PLATFORM
+    # GEOSWE_BACKEND is the documented name. SWELL_BACKEND and SWE_IGR_BACKEND are the
+    # research tree's names, kept so that both trees respond to the same job scripts.
+    # First one set wins; with none set the GPU is used when it is there.
+    asked = (os.environ.get("GEOSWE_BACKEND")
+             or os.environ.get("SWELL_BACKEND")
+             or os.environ.get("SWE_IGR_BACKEND"))
+    backend = (asked or "cupy").lower()
+    if backend not in ("cupy", "numpy"):
+        raise ValueError(f"GEOSWE_BACKEND must be 'cupy' (GPU) or 'numpy' (CPU), got {backend!r}")
+    if backend == "cupy":
+        _check_single_cupy()   # raise here, not inside the NumPy fallback below
+        try:
+            import cupy as cp  # type: ignore
+        except ImportError:
+            cp = None
+            if asked:          # the GPU was asked for by name: say why it is not used
+                warnings.warn("GEOSWE_BACKEND=cupy, but CuPy is not installed; using the NumPy "
+                              "(CPU) backend. The gpu extra installs it: "
+                              "pip install 'cupy-cuda12x[ctk]' (on an AMD GPU, the gpu-rocm "
+                              "extra: pip install 'cupy-rocm-7-0')",
+                              stacklevel=3)
+        if cp is not None:
+            if _cuda_device_count(cp) > 0:
+                _xp_module = cp
+                USING_CUPY = True
+                GPU_PLATFORM = _platform_of(cp)
+                return
+            if asked:
+                raise RuntimeError(
+                    f"GEOSWE_BACKEND=cupy, but {_no_device(cp)}. "
+                    f"Set GEOSWE_BACKEND=numpy to run on the CPU.")
+            warnings.warn(f"CuPy is installed but {_no_device(cp)}; using the NumPy "
+                          f"(CPU) backend. Set GEOSWE_BACKEND=numpy to make this choice explicit.",
+                          stacklevel=3)
+    import numpy as np
+    _xp_module = np
+    USING_CUPY = False
+    GPU_PLATFORM = None
+
+
+_select_default()
+
+
+def set_backend(name: str):
+    """Switch the global backend to 'numpy' or 'cupy'.
+
+    Must be called before any solver module is imported: ``geoswe.solver`` (and
+    the CUDA kernels behind it) freeze the backend choice at import time, so a
+    late switch would silently run the wrong code path. Prefer setting the
+    ``GEOSWE_BACKEND`` environment variable before ``import geoswe``.
+    """
+    global _xp_module, USING_CUPY, GPU_PLATFORM
+    name = str(name).lower()  # normalize
+    if name not in ("cupy", "numpy"):
+        raise ValueError(f"backend must be 'cupy' or 'numpy', got {name!r}")
+    # Fail loud when the switch cannot take effect: solver.py freezes
+    # USING_CUPY/fused-kernel dispatch at its own import.
+    import sys
+    if "geoswe.solver" in sys.modules and name != get_backend():
+        raise RuntimeError(
+            f"set_backend({name!r}) called after geoswe.solver was imported with the "
+            f"{get_backend()!r} backend; the switch cannot take effect. Set the "
+            f"GEOSWE_BACKEND environment variable before importing geoswe instead."
+        )
+    if name == "cupy":
+        _check_single_cupy()
+        import cupy as cp  # type: ignore
+        if _cuda_device_count(cp) < 1:
+            raise RuntimeError(f"set_backend('cupy'): {_no_device(cp)}")
+        _xp_module = cp
+        USING_CUPY = True
+        GPU_PLATFORM = _platform_of(cp)
+    else:
+        import numpy as np
+        _xp_module = np
+        USING_CUPY = False
+        GPU_PLATFORM = None
+
+
+def get_backend() -> str:
+    """Return ``"cupy"`` when the GPU backend is active, else ``"numpy"``."""
+    return "cupy" if USING_CUPY else "numpy"
+
+
+def gpu_platform():
+    """GPU stack behind the CuPy backend: ``"cuda"`` (NVIDIA), ``"hip"`` (AMD ROCm),
+    or ``None`` on the NumPy backend."""
+    return GPU_PLATFORM
+
+
+def nvidia_compute_capability() -> str:
+    """Compute capability of the current NVIDIA device as CuPy reports it (``"90"`` on
+    an H100), or ``""`` when the device is not an NVIDIA one.
+
+    Use this to recognise an NVIDIA architecture, not ``cp.cuda.Device().compute_capability``:
+    under ROCm that attribute carries the AMD ISA number, so an MI200-series card
+    (gfx90a) reports ``"90"`` too.
+    """
+    if GPU_PLATFORM != "cuda":
+        return ""
+    try:
+        import cupy as cp  # type: ignore
+        return str(cp.cuda.Device().compute_capability)
+    except Exception:
+        return ""
+
+
+# Declared ahead of every kernel under ROCm. CUDA guarantees that __fmul_rn/__fadd_rn
+# (and the double forms) are never contracted into a fused multiply-add; the kernels use
+# them where two code paths must round identically. HIP declares them as plain `*` and
+# `+`, which clang then fuses after inlining, so they are redefined here with contraction
+# switched off.
+_HIP_COMPAT = r"""
+static __device__ __forceinline__ float _geoswe_fmul_rn(float a, float b) {
+#pragma clang fp contract(off)
+    return a * b;
+}
+static __device__ __forceinline__ float _geoswe_fadd_rn(float a, float b) {
+#pragma clang fp contract(off)
+    return a + b;
+}
+static __device__ __forceinline__ double _geoswe_dmul_rn(double a, double b) {
+#pragma clang fp contract(off)
+    return a * b;
+}
+static __device__ __forceinline__ double _geoswe_dadd_rn(double a, double b) {
+#pragma clang fp contract(off)
+    return a + b;
+}
+#define __fmul_rn(a, b) _geoswe_fmul_rn((a), (b))
+#define __fadd_rn(a, b) _geoswe_fadd_rn((a), (b))
+#define __dmul_rn(a, b) _geoswe_dmul_rn((a), (b))
+#define __dadd_rn(a, b) _geoswe_dadd_rn((a), (b))
+"""
+# `#line 1` keeps compiler diagnostics on the kernel's own line numbers.
+_HIP_PREAMBLE = _HIP_COMPAT + "#line 1\n"
+
+# Compiler options that only NVRTC understands; the ROCm compiler rejects them.
+_NVRTC_ONLY_OPTIONS = ("-maxrregcount", "--maxrregcount")
+_warned_options = set()
+
+
+def _hip_options(options):
+    """`options` without the NVRTC-only ones, warning once for each one dropped."""
+    kept = []
+    for opt in options:
+        if str(opt).startswith(_NVRTC_ONLY_OPTIONS):
+            if opt not in _warned_options:
+                _warned_options.add(opt)
+                warnings.warn(f"kernel option {opt!r} is specific to NVIDIA's compiler and "
+                              f"is ignored on AMD GPUs (ROCm)", RuntimeWarning, stacklevel=3)
+            continue
+        kept.append(opt)
+    return tuple(kept)
+
+
+# How the ROCm compiler may contract `a*b + c` into a fused multiply-add, chosen by
+# GEOSWE_HIP_FP_CONTRACT. clang's own default for HIP ("fast") fuses across statements at
+# the optimizer's discretion, so the same expression can round differently in two kernels:
+# the fused and the split time step then differ in the last bit (measured on an MI200,
+# tests/test_gpu_storage_fused.py). "off" never fuses and "on" fuses only inside a single
+# source expression, so both are fixed by the source text. "off" is the default: every
+# operation rounds once, in source order, so two kernels agree whenever their statements
+# do, and the dense and compressed solvers agree bit for bit. "on" is about 3 % faster
+# (MI250X, 147 M cells) but leaves 1 ulp between those two (test_gpu_compressed_equiv.py).
+_HIP_FP_CONTRACT = {"off": "-ffp-contract=off", "on": "-ffp-contract=on",
+                    "fast": "-ffp-contract=fast-honor-pragmas"}
+_HIP_FP_CONTRACT_DEFAULT = "off"
+
+
+def _hip_fp_contract():
+    mode = os.environ.get("GEOSWE_HIP_FP_CONTRACT", _HIP_FP_CONTRACT_DEFAULT).lower()
+    if mode not in _HIP_FP_CONTRACT:
+        raise ValueError(f"GEOSWE_HIP_FP_CONTRACT must be one of {sorted(_HIP_FP_CONTRACT)}, "
+                         f"got {mode!r}")
+    return (_HIP_FP_CONTRACT[mode],)
+
+
+def raw_kernel(code, name, options=(), **kwargs):
+    """``cupy.RawKernel`` for a CUDA C kernel, on either GPU vendor.
+
+    On NVIDIA the source and the options reach CuPy unchanged. Under ROCm the source
+    gets the compatibility preamble, NVRTC-only options are dropped, and the
+    floating-point contraction mode is set (``GEOSWE_HIP_FP_CONTRACT``).
+    """
+    import cupy as cp  # type: ignore
+    if GPU_PLATFORM == "hip":
+        code, options = _HIP_PREAMBLE + code, _hip_options(options) + _hip_fp_contract()
+    return cp.RawKernel(code, name, options=tuple(options), **kwargs)
+
+
+def raw_module(code, options=(), **kwargs):
+    """``cupy.RawModule`` counterpart of :func:`raw_kernel`."""
+    import cupy as cp  # type: ignore
+    if GPU_PLATFORM == "hip":
+        code, options = _HIP_PREAMBLE + code, _hip_options(options) + _hip_fp_contract()
+    return cp.RawModule(code=code, options=tuple(options), **kwargs)
+
+
+def elementwise_kernel(in_params, out_params, operation, name="kernel", **kwargs):
+    """``cupy.ElementwiseKernel`` counterpart of :func:`raw_kernel`.
+
+    The split time step is made of these, so under ROCm they must be compiled with the
+    same contraction mode as the raw kernels they are compared against.
+    """
+    import cupy as cp  # type: ignore
+    if GPU_PLATFORM == "hip":
+        kwargs["options"] = _hip_options(kwargs.get("options", ())) + _hip_fp_contract()
+        kwargs["preamble"] = _HIP_COMPAT + kwargs.get("preamble", "")
+    return cp.ElementwiseKernel(in_params, out_params, operation, name, **kwargs)
+
+
+class _XPProxy:
+    """Proxy that forwards attribute access to the currently-selected backend module."""
+
+    def __getattr__(self, name):
+        return getattr(_xp_module, name)
+
+
+xp = _XPProxy()
+
+
+def to_host(arr):
+    """Return a NumPy array regardless of backend."""
+    if USING_CUPY and hasattr(arr, "get"):
+        return arr.get()
+    import numpy as np
+    return np.asarray(arr)
+
+
+def to_device(arr):
+    """Move a host array to the current backend (GPU if CuPy)."""
+    if USING_CUPY:
+        import cupy as cp  # type: ignore
+        return cp.asarray(arr)
+    return arr
+
+
+def sync():
+    """Block until all GPU operations have completed (no-op on CPU)."""
+    if USING_CUPY:
+        import cupy as cp  # type: ignore
+        cp.cuda.Stream.null.synchronize()

geoswe/bc.py

@@ -0,0 +1,235 @@
+"""Boundary conditions implemented via ghost cells.
+
+For SWE state q in 1D shape (Nvar, Nx_pad) or 2D shape (Nvar, Nx_pad, Ny_pad).
+Ghost cells are the first/last ngh cells along each spatial axis.
+"""
+from __future__ import annotations
+
+from .backend import xp as np  # backend-agnostic
+
+
+def apply_bc_1d(q, ngh: int, kind: str = "extrapolate", left=None, right=None):
+    """Apply boundary conditions to a 1D state.
+
+    kind:
+        'extrapolate': zero-gradient (open)
+        'periodic': wrap-around
+        'wall': reflective (h mirrored, hu sign-flipped)
+        'dirichlet': both ends fixed; provide `left` and `right` (each (Nvar,))
+    """
+    if kind == "periodic":
+        q[:, :ngh] = q[:, -2 * ngh : -ngh]
+        q[:, -ngh:] = q[:, ngh : 2 * ngh]
+        return
+    if kind == "extrapolate":
+        q[:, :ngh] = q[:, ngh : ngh + 1]
+        q[:, -ngh:] = q[:, -ngh - 1 : -ngh]
+        return
+    if kind == "wall":
+        # mirror with sign flip on the normal momentum
+        for j in range(ngh):
+            q[0, j] = q[0, 2 * ngh - 1 - j]
+            q[1, j] = -q[1, 2 * ngh - 1 - j]
+            q[0, -1 - j] = q[0, -2 * ngh + j]
+            q[1, -1 - j] = -q[1, -2 * ngh + j]
+        return
+    if kind == "dirichlet":
+        assert left is not None and right is not None
+        for j in range(ngh):
+            q[:, j] = left
+            q[:, -1 - j] = right
+        return
+    if kind == "fall":
+        # 'fall' = free outflow: zero depth and momentum in the ghost cells (same as apply_bc_2d).
+        q[:, :ngh] = 0.0
+        q[:, -ngh:] = 0.0
+        return
+    raise ValueError(f"unknown BC kind {kind}")
+
+
+def apply_bc_2d(q, ngh: int, kind_x: str = "extrapolate", kind_y: str = "extrapolate"):
+    """Apply boundary conditions to a 2D state. Independent BCs along x and y."""
+    # x-direction ghost cells
+    if kind_x == "periodic":
+        q[:, :ngh, :] = q[:, -2 * ngh : -ngh, :]
+        q[:, -ngh:, :] = q[:, ngh : 2 * ngh, :]
+    elif kind_x == "extrapolate":
+        q[:, :ngh, :] = q[:, ngh : ngh + 1, :]
+        q[:, -ngh:, :] = q[:, -ngh - 1 : -ngh, :]
+    elif kind_x == "wall":
+        # Vectorized mirror-with-sign-flip avoids launching 6 kernels per
+        # ghost layer on CuPy. The
+        # ::-1 indexing on q[:, ngh:2*ngh, :] produces the mirror of the first
+        # ngh interior cells in reverse order, matching the per-j formula
+        # q[..., j] = q[..., 2*ngh - 1 - j].
+        q[0, :ngh, :] =  q[0, ngh : 2 * ngh, :][::-1, :]
+        q[1, :ngh, :] = -q[1, ngh : 2 * ngh, :][::-1, :]
+        q[2, :ngh, :] =  q[2, ngh : 2 * ngh, :][::-1, :]
+        q[0, -ngh:, :] =  q[0, -2 * ngh : -ngh, :][::-1, :]
+        q[1, -ngh:, :] = -q[1, -2 * ngh : -ngh, :][::-1, :]
+        q[2, -ngh:, :] =  q[2, -2 * ngh : -ngh, :][::-1, :]
+    elif kind_x == "fall":
+        # 'fall': h=0, hU=0 in the ghost cells -- forces free outflow by
+        # creating a maximum hydraulic gradient at the boundary
+        q[:, :ngh, :] = 0.0
+        q[:, -ngh:, :] = 0.0
+    else:
+        raise ValueError(f"unknown BC kind_x {kind_x}")
+
+    # y-direction ghost cells
+    if kind_y == "periodic":
+        q[:, :, :ngh] = q[:, :, -2 * ngh : -ngh]
+        q[:, :, -ngh:] = q[:, :, ngh : 2 * ngh]
+    elif kind_y == "extrapolate":
+        q[:, :, :ngh] = q[:, :, ngh : ngh + 1]
+        q[:, :, -ngh:] = q[:, :, -ngh - 1 : -ngh]
+    elif kind_y == "wall":
+        # Vectorized mirror-with-sign-flip on y-normal momentum.
+        q[0, :, :ngh] =  q[0, :, ngh : 2 * ngh][:, ::-1]
+        q[1, :, :ngh] =  q[1, :, ngh : 2 * ngh][:, ::-1]
+        q[2, :, :ngh] = -q[2, :, ngh : 2 * ngh][:, ::-1]
+        q[0, :, -ngh:] =  q[0, :, -2 * ngh : -ngh][:, ::-1]
+        q[1, :, -ngh:] =  q[1, :, -2 * ngh : -ngh][:, ::-1]
+        q[2, :, -ngh:] = -q[2, :, -2 * ngh : -ngh][:, ::-1]
+    elif kind_y == "fall":
+        q[:, :, :ngh] = 0.0
+        q[:, :, -ngh:] = 0.0
+    else:
+        raise ValueError(f"unknown BC kind_y {kind_y}")
+
+
+def apply_bc_2d_face(q, ngh: int, face: str, kind: str):
+    """Apply a physical BC to a single ghost face of a 2D state.
+
+    face: one of 'x-', 'x+', 'y-', 'y+'
+    kind: 'extrapolate', 'wall', or 'periodic'  (periodic only makes sense
+          for x- and x+ together / y- and y+ together; calling it on a
+          single face wraps to the local subgrid's opposite interior side
+          which is rarely useful in multi-rank settings; used here only
+          when the caller has already determined the face is physical-only).
+
+    Used by the multi-GPU dispatcher: halo.exchange fills MPI faces, then
+    this fills only the remaining physical faces (avoiding the clobber
+    that ``apply_bc_2d`` with kind='periodic' would cause on an MPI face).
+    """
+    if face == "x-":
+        if kind == "extrapolate":
+            q[:, :ngh, :] = q[:, ngh : ngh + 1, :]
+        elif kind == "wall":
+            for j in range(ngh):
+                q[0, j, :] =  q[0, 2 * ngh - 1 - j, :]
+                q[1, j, :] = -q[1, 2 * ngh - 1 - j, :]
+                q[2, j, :] =  q[2, 2 * ngh - 1 - j, :]
+        elif kind == "fall":
+            q[:, :ngh, :] = 0.0
+        elif kind == "periodic":
+            q[:, :ngh, :] = q[:, -2 * ngh : -ngh, :]
+        else:
+            raise ValueError(f"unknown BC kind {kind!r} on face {face}")
+    elif face == "x+":
+        if kind == "extrapolate":
+            q[:, -ngh:, :] = q[:, -ngh - 1 : -ngh, :]
+        elif kind == "wall":
+            for j in range(ngh):
+                q[0, -1 - j, :] =  q[0, -2 * ngh + j, :]
+                q[1, -1 - j, :] = -q[1, -2 * ngh + j, :]
+                q[2, -1 - j, :] =  q[2, -2 * ngh + j, :]
+        elif kind == "fall":
+            q[:, -ngh:, :] = 0.0
+        elif kind == "periodic":
+            q[:, -ngh:, :] = q[:, ngh : 2 * ngh, :]
+        else:
+            raise ValueError(f"unknown BC kind {kind!r} on face {face}")
+    elif face == "y-":
+        if kind == "extrapolate":
+            q[:, :, :ngh] = q[:, :, ngh : ngh + 1]
+        elif kind == "wall":
+            for j in range(ngh):
+                q[0, :, j] =  q[0, :, 2 * ngh - 1 - j]
+                q[1, :, j] =  q[1, :, 2 * ngh - 1 - j]
+                q[2, :, j] = -q[2, :, 2 * ngh - 1 - j]
+        elif kind == "fall":
+            q[:, :, :ngh] = 0.0
+        elif kind == "periodic":
+            q[:, :, :ngh] = q[:, :, -2 * ngh : -ngh]
+        else:
+            raise ValueError(f"unknown BC kind {kind!r} on face {face}")
+    elif face == "y+":
+        if kind == "extrapolate":
+            q[:, :, -ngh:] = q[:, :, -ngh - 1 : -ngh]
+        elif kind == "wall":
+            for j in range(ngh):
+                q[0, :, -1 - j] =  q[0, :, -2 * ngh + j]
+                q[1, :, -1 - j] =  q[1, :, -2 * ngh + j]
+                q[2, :, -1 - j] = -q[2, :, -2 * ngh + j]
+        elif kind == "fall":
+            q[:, :, -ngh:] = 0.0
+        elif kind == "periodic":
+            q[:, :, -ngh:] = q[:, :, ngh : 2 * ngh]
+        else:
+            raise ValueError(f"unknown BC kind {kind!r} on face {face}")
+    else:
+        raise ValueError(f"unknown face {face!r}")
+
+
+def apply_inflow_discharge(q, bed, idx_i, idx_j, nx_n, ny_n, ds, Q, ngh,
+                           hsum=None, dry_frac=0.10):
+    """Impose a discharge (hydrograph) inflow through the GHOST cells of an inlet.
+
+    The stage ring of \\S3.6 prescribes a free surface and zeroes momentum; a
+    river inlet is the opposite case -- the discharge is what is known and the
+    depth has to come from the flow. The split across the cross-section follows
+    the usual depth-weighted convention, so a single hydrograph drives the
+    whole inlet:
+
+      * total discharge ``Q`` (m^3/s) is split over the inlet cells in
+        proportion to their depth, ``w_k = h_k / sum(h)``, so the deeper part of
+        a cross-section carries more of the flow;
+      * the per-cell normal unit discharge is ``Q*w_k/ds`` (m^2/s);
+      * a DRY cross-section has no depth to weight by, so the surface is seeded
+        from the bed relief, ``z_min + dry_frac*(z_max - z_min)``; where that
+        section is FLAT this degenerates to zero depth and the hydrograph would
+        never start, so it falls back to the critical depth ``(q^2/g)^(1/3)`` --
+        the minimum-energy depth carrying q, which adds no free parameter.
+
+    The values are written into the ``ngh`` GHOST cells outward of each inlet
+    cell, NOT into the interior. Writing the interior instead makes the inlet
+    depend on whatever face BC surrounds it -- a wall would mirror the imposed
+    momentum straight back and deliver nothing -- and it would also clobber the
+    solution every step. Writing the ghost lets the scheme compute the boundary
+    flux itself, so the inlet works on a wall, an open face or a mask edge
+    alike. Depth is carried out to the ghost unchanged (zero-gradient), which is
+    the "impose one condition, take the other from the flow" a subcritical inlet
+    wants.
+
+    ``nx_n``, ``ny_n`` are the INWARD unit normal and must be axis-aligned
+    (one of them +-1, the other 0). ``hsum`` takes the MPI-reduced depth sum
+    when a cross-section spans ranks; otherwise each rank would weight against
+    its own share.
+    """
+    if len(idx_i) == 0:
+        return
+    h = q[0][idx_i, idx_j]
+    z = bed[idx_i, idx_j]
+    s = float(h.sum()) if hsum is None else float(hsum)
+    if s > 0.0:
+        w = h / s
+        h_face = h                          # zero-gradient depth into the ghost
+    else:                                   # dry cross-section: see docstring
+        zmin = z.min(); zmax = z.max()
+        h_face = np.maximum(zmin + dry_frac * (zmax - zmin) - z, 0.0)
+        low = (z <= zmin)                   # put the flow in the thalweg
+        n_low = float(low.sum())
+        if float(h_face.max()) <= 0.0:      # FLAT inlet -> critical depth
+            g = 9.80665
+            q_cell = abs(Q) / (max(n_low, 1.0) * ds)
+            h_face = h_face + (q_cell * q_cell / g) ** (1.0 / 3.0)
+        w = np.where(low, 1.0 / max(n_low, 1.0), 0.0)
+    qn = (Q / ds) * w                       # normal unit discharge, m^2/s
+    si, sj = int(round(nx_n)), int(round(ny_n))
+    for k in range(1, int(ngh) + 1):        # ghost cells, outward of the inlet
+        gi = idx_i - k * si
+        gj = idx_j - k * sj
+        q[0][gi, gj] = h_face
+        q[1][gi, gj] = qn * nx_n
+        q[2][gi, gj] = qn * ny_n