feat: add ess lab3

2026-02-27 17:06:05 +01:00
parent a639c34cee
commit a0b45dd36b
1 changed files with 144 additions and 0 deletions
@@ -0,0 +1,144 @@
+# Requirements:
+# conda
+# conda install fastscape -c conda-forge
+# conda install -c conda-forge "zarr<3"
+
+import numpy as np
+import xarray as xr
+import xsimlab as xs
+import matplotlib.pyplot as plt
+from fastscape.models import basic_model
+
+Lx = 100e3
+Ly = 100e3
+nx = 101
+ny = 101
+
+t_end = 5e6
+dt = 2e4
+time = np.arange(0.0, t_end + dt, dt)
+
+save_every = 10
+save_time = time[::save_every]
+
+Kf = 1e-5
+m_exp = 0.5
+n_exp = 1.0
+Kd = 0.3
+
+U_default = 1e-3
+
+BC = "fixed_value"
+
+def experiment(uplift_rate=U_default, k_coef=Kf, area_exp=m_exp, slope_exp=n_exp, diffusivity=Kd, seed=42, label="exp"):
+
+    in_vars = {
+        "grid__shape": xr.DataArray([ny, nx], dims=("shape_yx",)),
+        "grid__length": xr.DataArray([Ly, Lx], dims=("shape_yx",)),
+        "boundary__status": BC,
+        "uplift__rate": uplift_rate,
+        "diffusion__diffusivity": diffusivity,
+        "spl__k_coef": k_coef,
+        "spl__area_exp": area_exp,
+        "spl__slope_exp": slope_exp,
+        "init_topography__seed": seed,
+    }
+
+    clocks = {
+        "time": xr.IndexVariable("time", time),
+        "save": xr.IndexVariable("save", save_time),
+    }
+
+    out_vars = {
+        "topography__elevation": "save",
+        "terrain__slope": "save",
+        "drainage__area": "save",
+        "erosion__cumulative_height": "save",
+    }
+
+    setup = xs.create_setup(
+        model=basic_model,
+        clocks=clocks,
+        master_clock="time",
+        input_vars=in_vars,
+        output_vars=out_vars,
+    )
+
+    ds = setup.xsimlab.run(model=basic_model)
+
+    ds = ds.assign_attrs(
+        experiment_label=label,
+        uplift_rate=uplift_rate,
+        k_coef=k_coef,
+        area_exp=area_exp,
+        slope_exp=slope_exp,
+        diffusivity=diffusivity,
+        seed=seed,
+        Lx=Lx,
+        Ly=Ly,
+        nx=nx,
+        ny=ny,
+        dt=dt,
+        t_end=t_end,
+        boundary_condition=BC,
+    )
+
+    return ds
+
+def plot_maps(ds, title_prefix=""):
+    elev = ds["topography__elevation"]
+    slope = ds["terrain__slope"]
+
+    elev_f = elev.isel(save=-1)
+    slope_f = slope.isel(save=-1)
+
+    _, ax = plt.subplots(figsize=(6, 5))
+    im = ax.imshow(elev_f, origin="lower")
+    ax.set_title(f"{title_prefix} final elevation (m)")
+    plt.colorbar(im, ax=ax, shrink=0.85)
+    plt.show()
+
+    _, ax = plt.subplots(figsize=(6, 5))
+    im = ax.imshow(slope_f, origin="lower")
+    ax.set_title(f"{title_prefix} final slope (-)")
+    plt.colorbar(im, ax=ax, shrink=0.85)
+    plt.show()
+
+
+def plot_max_elevation(ds, title_prefix=""):
+    elev = ds["topography__elevation"]
+    t_myr = ds["save"].values / 1e6
+    zmax = elev.max(dim=("y", "x")).values
+
+    _, ax = plt.subplots(figsize=(6, 4))
+    ax.plot(t_myr, zmax)
+    ax.set_xlabel("Time (Myr)")
+    ax.set_ylabel("Max elevation (m)")
+    ax.set_title(f"{title_prefix} max elevation vs time")
+    ax.grid(True, alpha=0.3)
+    plt.show()
+
+experiments = [
+    dict(label="baseline", uplift_rate=1e-3, k_coef=1e-5, diffusivity=0.3, seed=42),
+    dict(label="high_uplift", uplift_rate=2e-3, k_coef=1e-5, diffusivity=0.3, seed=42),
+    dict(label="low_Kf", uplift_rate=1e-3, k_coef=3e-6, diffusivity=0.3, seed=42),
+    dict(label="high_Kd", uplift_rate=1e-3, k_coef=1e-5, diffusivity=0.9, seed=42),
+]
+
+results = {}
+
+for p in experiments:
+    ds = experiment(
+        uplift_rate=p["uplift_rate"],
+        k_coef=p["k_coef"],
+        diffusivity=p["diffusivity"],
+        seed=p["seed"],
+        label=p["label"],
+    )
+    results[p["label"]] = ds
+
+print("Done. Experiments:", list(results.items()))
+
+#for label, ds in results.items():
+#    plot_maps(ds, title_prefix=f"[{label}]")
+#    plot_max_elevation(ds, title_prefix=f"[{label}]")