Compressed active-cell solver

GPU-only. See Compressed active-cell mesh for the concepts.

geoswe.CompressedSolver is resolved lazily, because geoswe.compressed_solver imports CuPy at module load: the name exists on a CPU-only install and raises ImportError naming the gpu extra only when you touch it. For the same reason it is deliberately left out of geoswe.__all__, so that from geoswe import * keeps working without a GPU.

class geoswe.CompressedSolver

Bases: object

Shallow-water solver on the compressed active-cell mesh (GPU, float32).

Build the problem as a dense Solver2D (mesh, bed, initial state, Manning roughness), mark the cells to keep with set_inside_mask, and convert:

cs = CompressedSolver.from_dense(s)      # packs the active cells, frees ``s``
cs.set_rain(rain)                        # a RainfallForcing; optional
cs.run(t_end=3600.0)
h = cs.depth()                           # NumPy (nx, ny), zero outside the mask

Only the active cells and a two-cell halo around them are stored and updated. The halo keeps the state it had in s (dry, for a dry start), so water that reaches the edge of the active set leaves it, as with the dense "fall" boundary. The other set_* methods attach the forcings of the coastal applications (stage ring, sponge, infiltration, drains, cross-sections); they take the bundles that geoswe.runlib.driver builds. For a domain too large to build densely, see save_cache and run_cached.

classmethod from_dense(s, *, ngh=None, dx=None, cfl=None, h_min=None, g=None, m_cls_xp=None, m_tab_xp=None, x0=None, y0=None, crs_wkt='', nx_glob=None, ny_glob=None, comm=None, dims=None, i0_glob=0, j0_glob=0, Nx_loc=None, Ny_loc=None, cfl_no_sigma=False, cfl_linf=False, nx_orig=None, ny_orig=None, gauge_every_s=360.0, say=<built-in function print>)

Pack the dense Solver2D s into the flat layout, free its fields, and build the MPI halo.

On one GPU every argument but s can be left out: the grid, CFL number, depth floor, gravity and Manning roughness are read from s, and a solver with no inside mask keeps every cell. Under MPI the global grid size and the rank’s placement must be passed, as the run driver does. cfl_linf=True uses the dense solver’s velocity norm, max(|u|, |v|), in the time step, so that a compressed run takes the dense run’s steps. Rainfall and the other forcings are not copied from s; attach them with the set_* methods.

m_cls_xp (the Manning class of every cell) is on the padded grid, (nx + 2*ngh, ny + 2*ngh), the shape of s.q[0], not the interior (nx, ny); m_tab_xp is the class table it indexes. The packing gather is a CuPy fancy index, which wraps an out-of-range index instead of raising, so an interior-shaped field would be gathered from the wrong cells; the shape is therefore checked (geoswe.compressed_mesh.CompressedMesh2D.pack()).

set_ring(ring)

set_sponge(sponge)

Attach the open-boundary sponge the run driver builds: {"keep": k, "amb": a}, two fields on the padded grid (the shape of s.q[0], not the interior (nx, ny)). Each step relaxes the band toward the open-ocean state, h <- keep*h + amb with hu, hv <- keep*hu, keep*hv, so keep = 1 and amb = 0 leave a cell alone. The driver sets keep = 1 - alpha over the band and amb = alpha times the ambient still-water depth, which pulls the band toward that depth and absorbs the outgoing wave. None for no sponge.

set_rain(rain)

Attach rainfall: a RainfallForcing (uniform, or one (nx, ny) frame per time), or the bundle the run driver builds for a gridded product on its native grid (native_rate_dev, lookup_dev, t_s).

add_inflow(ghost_idx, nbr_idx, normal, ds, t_series, q_series)

Add one discharge (hydrograph) inlet; call once per river.

ghost_idx are flat indices of the inlet’s GHOST cells and nbr_idx their active neighbours (the pairing _build_ghost_bc produces – pass the subset that is the river mouth). Each inlet keeps its own normal, width and hydrograph.

set_drain(drain)

set_infil(infil)

set_ga_drain(b)

set_clamp(b)

set_cross_sections(b)

enable_max_depth(on=True)

property n_active

Number of active cells, the ones that are updated.

property n_stored

Number of stored cells, the active cells plus their two-cell halo.

depth()

Water depth as a NumPy array on the unpadded grid, zero outside the active set.

Under MPI this is the calling rank’s block.

save_cache(cache_dir)

Write the flat mesh, state, bed, roughness and forcing bundles to cache_dir (per-rank r## subdirectories under MPI) for later run_cached loads.

The ring, sponge, rain and karst-drain bundles are persisted. Green-Ampt infiltration, the stage clamp and the cross-section gauges are not, and run_cached cannot apply them: a replay of a run that used them omits those terms, which changes the trajectory. This warns for the two that do (Green-Ampt, the clamp); the cross sections only sample. Forcings set after this call, and the set_drain/set_infil bundles read at run() time, are not in the file either.

run(t_end=None, *, out_dir=None, frame_every_s=0.0, checkpoint_every_s=0.0, ckpt_dir=None, resume=False, max_wall_s=0.0, stop_at_epoch=0.0, say=<built-in function print>, dt_max=None, dt_min=None)

Time-step the compressed mesh from t = 0 to t_end.

Writes depth frames every frame_every_s of simulated time into out_dir (plus gauges, cross-sections and the running depth maximum when enabled), checkpoints every checkpoint_every_s and at max_wall_s/stop_at_epoch deadlines, and resumes bit-identically from ckpt_dir when resume is set. With no out_dir nothing is written; read the result with depth(). dt_max caps the time step. On one rank, rain also keeps the step below the CFL step of the film it lays down, as in geoswe.Solver2D.run(), so rain on a dry bed is not deposited minutes at a time. dt_min raises when the step collapses below it rather than grinding on to the job’s wall clock (default GEOSWE_DT_MIN, 0 = off). say=None gives a silent run.

The frames are not GeoTIFFs. Each is one compressed .npz per rank, <out_dir>/frames_parallel/depth_<index:05d>_t<seconds:07d>_r<rank:02d>.npz, holding a single array h: that rank’s interior subdomain of the depth field, float16 by default (a 0.008 m step at a depth of 10 m), or float32 with SWE_FRAME_FP32=1. frames_parallel/manifest.json, written beside them, carries the grid, the CRS and each rank’s (i0, j0, nx, ny) offsets, which is the only way to place the shards back on the global grid; no reader for the format ships with the library. enable_max_depth writes GeoTIFFs instead.

Cache replay

run_cached runs a cache written by CompressedSolver.save_cache() without rebuilding the mesh. It is not re-exported at the top level: import it as from geoswe.compressed_solver import run_cached.

geoswe.compressed_solver.run_cached(cache_dir, *, inflows=None, t_end, frame_every_s, out_dir, cfl=0.5, h_min=1e-06, g=9.81, comm=None, checkpoint_every_s=0.0, ckpt_dir=None, resume=False, max_wall_s=0.0, stop_at_epoch=0.0, say=<built-in function print>, bench=False, dt_max=None, dt_min=None)

Load the flat cache straight to GPU (no dense domain) and run the loop. Under MPI (comm.size>1) each rank loads cache_dir/r<rank>/ and rebuilds its halo.

dt_max caps the time step, in seconds. Unlike CompressedSolver.run this path applies no rain cap of its own: the published cached benchmark replays with rain, and the cap would change its step schedule. Without one, rain on a dry bed takes the dry-partition CFL step, 478.9 s at dx = 3 m, and lays that whole interval of rain down in one go. Pass dt_max="rain" for the same film bound CompressedSolver.run applies (27.4 s on that grid at 40 mm/h), or a number for a cap of your own. That bound is read from the rain table in the cache, so the GEOSWE_RAIN_NPZ and GEOSWE_RAIN_UNIFORM_MMHR overrides, which the loop applies afterwards, do not move it; pass the cap as a number when you replay another deck through them. dt_min raises when the step collapses below it, instead of grinding on to the job’s wall clock; it defaults to GEOSWE_DT_MIN (0 = off). say=None gives a silent run, as in CompressedSolver.run.

The compressed mesh

The flat active-cell layout the solver is packed into. Reading it is a measurement, not part of running a case: memory_summary sizes the indirection tables, and pack/unpack move a padded (nx + 2*ngh, ny + 2*ngh) array in and out of the flat (N_active,) one.

class geoswe.compressed_mesh.CompressedMesh2D(mesh, inside_mask)

Bases: object

Wraps a regular Cartesian Mesh2D + an inside_mask into a flat compressed-cell layout with neighbor indirection.

pack(arr_2d)

Gather the active subset of a padded (nxp, nyp) or (…, nxp, nyp) array into a flat (N_active,) or (…, N_active) array.

Uses fancy indexing. ij_active lives on the ACTIVE backend (a CuPy device array under the GPU backend), so arr_2d must be on the same backend – pack(numpy_array) raises under CuPy; cp.asarray the input first.

The last two axes must be the PADDED extent. CuPy’s fancy indexing wraps an out-of-range index per axis rather than raising (NumPy raises IndexError, and the compressed solver is CuPy-only), so a field on another extent is gathered from the wrong cells and comes back plausible and finite: measured on a 24x20 mesh handed a (22, 18) field, 61 of 328 entries came from a wrapped index, the worst 2218 off. Hence the shape check, the same hazard set_manning_table names.

unpack(arr_flat, fill=0.0, out=None)

Scatter a flat (N_active,) array back to (nxp, nyp) with fill at outside/ghost cells. If out is given, write into it in place.

The mirror of pack’s check, for the same reason: a CuPy scatter wraps an out-of-range index too (measured: index 5 into a 4-long axis writes row 1), so a wrongly shaped out would be filled at the wrong cells with no error at all.

property memory_summary

Approximate memory cost of the compressed-mesh tables in MB.