autosim.experimental.simulations.shallow_water

autosim.experimental.simulations.shallow_water#

Shallow-water equation simulator.

class ShallowWater2D(parameters_range=None, output_names=None, return_timeseries=False, log_level='progress_bar', nx=64, ny=128, Lx=64.0, Ly=128.0, T=90.0, dt_save=0.2, skip_nt=0, cfl=0.12, g=9.81, h_mean=1.0, nu=0.0005, drag=0.002, dtype=torch.float64)[source]#

Bases: SpatioTemporalSimulator

Full 2D shallow-water simulator with prognostic \([h, u, v]\).

The solver evolves fluid height \(h\) and horizontal velocity \((u, v)\) using:

\[\begin{split}\begin{aligned} \partial_t h + \nabla\cdot(h\mathbf{u}) &= 0, \\ \partial_t u + u\partial_x u + v\partial_y u &= f v - g\partial_x h + \nu\nabla^2 u - r u, \\ \partial_t v + u\partial_x v + v\partial_y v &= -f u - g\partial_y h + \nu\nabla^2 v - r v. \end{aligned}\end{split}\]
Parameters:

  • parameters_range (dict[str, tuple[float, float]] | None) –

    Input parameter (min, max) ranges. Supported keys:

    • amp (required): initial-condition amplitude scale.

    • h_mean: mean layer depth (scalar) around which spatial variations are generated (default 1.0 if omitted).

    • drag: linear drag coefficient (default 2e-3).

    • nu: Laplacian viscosity (default 5e-4).

    If None, uses {"amp": (0.05, 0.14)} only.

  • output_names (list[str] | None) – Passed to base. Default outputs: [“h”, “u”, “v”].

  • return_timeseries (bool) – Passed to base. Default outputs: [“h”, “u”, “v”].

  • log_level (str) – Passed to base. Default outputs: [“h”, “u”, “v”].

  • nx (int) – Grid, domain, time and CFL settings.

  • ny (int) – Grid, domain, time and CFL settings.

  • Lx (float) – Grid, domain, time and CFL settings.

  • Ly (float) – Grid, domain, time and CFL settings.

  • T (float) – Grid, domain, time and CFL settings.

  • dt_save (float) – Grid, domain, time and CFL settings.

  • skip_nt (int) – Grid, domain, time and CFL settings.

  • cfl (float) – Grid, domain, time and CFL settings.

  • g (float) – Physics constants (used when not in parameters_range).

  • h_mean (float) – Physics constants (used when not in parameters_range).

  • nu (float) – Physics constants (used when not in parameters_range).

  • drag (float) – Physics constants (used when not in parameters_range).

  • dtype (torch.dtype) – torch.float32 or torch.float64.

forward_samples_spatiotemporal(n, random_seed=None, ensure_exact_n=False)[source]#

Run sampled trajectories and return [batch,time,x,y,channels] data.

Parameters:

  • n (int)

  • random_seed (int | None)

  • ensure_exact_n (bool)

Return type:

dict

simulate_swe_2d(amp, return_timeseries, nx, ny, Lx, Ly, T, dt_save, cfl, g, h_mean, nu, drag, dtype=torch.float64, skip_nt=0)[source]#

Integrate full shallow-water equations with PDEArena-style random2 ICs.

Parameters:
Return type:

Tensor