145 lines
3.6 KiB
Python
145 lines
3.6 KiB
Python
# 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}]")
|