Additional shallow water param variation

examples/experimental/01_01_shallow_water_equation.ipynb

@@ -70,7 +70,18 @@
     "\n",
     "Main difference: **geometry**. PDEArena production data is generated on the sphere (SpeedyWeather/Julia), while this notebook uses a periodic Cartesian beta-plane box.\n",
     "\n",
-    "## Why intersting/useful\n",
+    "## Parameters\n",
+    "\n",
+    "The simulator supports optional parameters in `parameters_range` (in addition to required `amp`):\n",
+    "\n",
+    "- **`amp`** (required): initial-condition amplitude scale.\n",
+    "- **`H`**: mean layer depth (default 1.0 if omitted).\n",
+    "- **`drag`**: linear drag coefficient (default 2e-3).\n",
+    "- **`nu`**: Laplacian viscosity (default 5e-4).\n",
+    "\n",
+    "Use any subset; omitted parameters fall back to constructor defaults. Parameter order in `parameters_range` does not matter (parsing is by name).\n",
+    "\n",
+    "## Why interesting/useful\n",
     "\n",
     "- Produces rich vortex-dominated trajectories with physically interpretable channels.\n",
     "- Closer to weather-like dynamics than simple advection-diffusion benchmarks.\n",
@@ -87,6 +98,12 @@
     "from IPython.display import HTML\n",
     "\n",
     "from autosim.experimental.simulations import ShallowWater2D\n",
+    "from autosim.experimental.simulations.shallow_water import (\n",
+    "    DEFAULT_AMP_RANGE,\n",
+    "    DEFAULT_DRAG_RANGE,\n",
+    "    DEFAULT_H_RANGE,\n",
+    "    DEFAULT_NU_RANGE,\n",
+    ")\n",
     "from autosim.utils import plot_spatiotemporal_video\n"
    ]
   },
@@ -97,6 +114,7 @@
    "outputs": [],
    "source": [
     "# Interactive preview config: keep CFL fixed for stability.\n",
+    "# Use amp only (default), or add H, drag, nu for multi-parameter studies.\n",
     "sim = ShallowWater2D(\n",
     "    return_timeseries=True,\n",
     "    nx=64,\n",
@@ -108,10 +126,15 @@
     "    cfl=0.12,\n",
     "    log_level=\"warning\",\n",
     "    dtype=torch.float32,\n",
-    "    parameters_range={\"amp\": (0.05, 0.2)},\n",
+    "    parameters_range={\n",
+    "        \"amp\": DEFAULT_AMP_RANGE,\n",
+    "        \"H\": DEFAULT_H_RANGE,\n",
+    "        \"drag\": DEFAULT_DRAG_RANGE,\n",
+    "        \"nu\": DEFAULT_NU_RANGE,\n",
+    "    },\n",
     ")\n",
     "\n",
-    "batch = sim.forward_samples_spatiotemporal(n=1)\n",
+    "batch = sim.forward_samples_spatiotemporal(n=3)\n",
     "params = batch[\"constant_scalars\"]\n",
     "outputs = batch[\"data\"]\n"
    ]
@@ -124,7 +147,10 @@
    "source": [
     "print(\"constant_scalars shape:\", params.shape)\n",
     "print(\"data shape:\", outputs.shape)\n",
-    "print(\"sample params (trajectory 0):\", {\"amp\": float(params[0, 0])})\n"
+    "sample_params = {\n",
+    "    name: float(params[0, sim.get_parameter_idx(name)]) for name in sim.param_names\n",
+    "}\n",
+    "print(\"sample params (trajectory 0):\", sample_params)\n"
    ]
   },
   {
@@ -135,8 +161,9 @@
    "source": [
     "anim = plot_spatiotemporal_video(\n",
     "    outputs[:, :100],\n",
-    "    batch_idx=0,\n",
+    "    batch_idx=2,\n",
     "    channel_names=sim.output_names,\n",
+    "    preserve_aspect=True,\n",
     ")\n",
     "\n",
     "HTML(anim.to_jshtml())"
@@ -155,6 +182,12 @@
    "metadata": {},
    "outputs": [],
    "source": [
+    "from autosim.experimental.simulations.shallow_water import (\n",
+    "    DEFAULT_AMP_RANGE,\n",
+    "    DEFAULT_DRAG_RANGE,\n",
+    "    DEFAULT_H_RANGE,\n",
+    "    DEFAULT_NU_RANGE,\n",
+    ")\n",
     "from autosim.utils import generate_output_data\n",
     "\n",
     "sim_train = ShallowWater2D(\n",
@@ -167,7 +200,12 @@
     "    T=10.0,\n",
     "    dt_save=1.0,\n",
     "    cfl=0.12,\n",
-    "    parameters_range={\"amp\": (0.05, 0.2)},\n",
+    "    parameters_range={\n",
+    "        \"amp\": DEFAULT_AMP_RANGE,\n",
+    "        \"H\": DEFAULT_H_RANGE,\n",
+    "        \"drag\": DEFAULT_DRAG_RANGE,\n",
+    "        \"nu\": DEFAULT_NU_RANGE,\n",
+    "    },\n",
     ")\n",
     "\n",
     "outputs_data = generate_output_data(sim_train)\n"
@@ -195,7 +233,7 @@
  ],
  "metadata": {
   "kernelspec": {
-   "display_name": "autosim (3.11.14)",
+   "display_name": ".venv (3.11.15)",
    "language": "python",
    "name": "python3"
   },
@@ -209,7 +247,7 @@
    "name": "python",
    "nbconvert_exporter": "python",
    "pygments_lexer": "ipython3",
-   "version": "3.11.14"
+   "version": "3.11.15"
   }
  },
  "nbformat": 4,

src/autosim/configs/simulator/shallow_water2d.yaml

@@ -10,7 +10,7 @@ dt_save: 0.2
 skip_nt: 50
 cfl: 0.12
 g: 9.81
-H: 1.0
+h_mean: 1.0
 nu: 0.0005
 drag: 0.002
 parameters_range:

src/autosim/configs/simulator/shallow_water2d_4.yaml

@@ -0,0 +1,20 @@
+_target_: autosim.experimental.simulations.ShallowWater2D
+return_timeseries: true
+log_level: warning
+nx: 64
+ny: 64
+Lx: 64.0
+Ly: 64.0
+T: 74.0
+dt_save: 0.2
+skip_nt: 50
+cfl: 0.12
+g: 9.81
+h_mean: 1.0
+nu: 0.0005
+drag: 0.002
+parameters_range:
+  amp: [0.07, 0.14]
+  h_mean: [0.7, 1.5]
+  drag: [1e-3, 4e-3]
+  nu: [2e-4, 8e-4]