From a0b45dd36b5859cc222e7a3b5aef04f7bae999db Mon Sep 17 00:00:00 2001 From: huaqo Date: Fri, 27 Feb 2026 17:06:05 +0100 Subject: [PATCH] feat: add ess lab3 --- .../FastScape_MountainBelt_JoaquinGottlebe.py | 144 ++++++++++++++++++ 1 file changed, 144 insertions(+) create mode 100644 2026_Earth_System_Science/lab3/FastScape_MountainBelt_JoaquinGottlebe.py diff --git a/2026_Earth_System_Science/lab3/FastScape_MountainBelt_JoaquinGottlebe.py b/2026_Earth_System_Science/lab3/FastScape_MountainBelt_JoaquinGottlebe.py new file mode 100644 index 0000000..c2885f3 --- /dev/null +++ b/2026_Earth_System_Science/lab3/FastScape_MountainBelt_JoaquinGottlebe.py @@ -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}]")