Skip to content

Commit

Permalink
modified examples
Browse files Browse the repository at this point in the history
  • Loading branch information
@schwemro
schwemro committed Dec 6, 2023
1 parent 1b034d9 commit 0216045
Show file tree
Hide file tree
Showing 8 changed files with 75 additions and 69 deletions.
20 changes: 11 additions & 9 deletions examples/catchment_scale/eberbaechle/make_figures.py
View file
Original file line number Diff line number Diff line change
Expand Up @@ -36,12 +36,12 @@

nrows = 53
ncols = 80
x1 = 76
y1 = 43
t_dry = 300
x1 = 12
y1 = 7
t_dry = 1017
t_wet = 468
t_drywet = 338
t_wetdry = 480
t_wetdry = 896

base_path = Path(__file__).parent
# directory of figures
Expand Down Expand Up @@ -108,7 +108,7 @@
df.to_csv(base_path_figs / f"theta_et_perc.csv", sep=";", index=True)

# load transport simulations
states_tm_file = base_path / "svat_oxygen18_distributed" / "output" / "SVAT18O.nc"
states_tm_file = base_path / "svat_oxygen18_distributed" / "output" / "SVAT18O.average.nc"
ds_tm = xr.open_dataset(states_tm_file, engine="h5netcdf", decode_times=False)
# assign date
days_tm = ds_tm["Time"].values
Expand Down Expand Up @@ -403,20 +403,22 @@
axes[2, 2].set_ylim(
0,
)

vals1 = onp.where(ds_tm["tt10_q_ss"].isel(x=x1, y=y1).values >= 1000, onp.nan, ds_tm["tt10_q_ss"].isel(x=x1, y=y1).values)
vals2 = onp.where(ds_tm["tt90_q_ss"].isel(x=x1, y=y1).values >= 1000, onp.nan, ds_tm["tt90_q_ss"].isel(x=x1, y=y1).values)
vals3 = onp.where(ds_tm["tt50_q_ss"].isel(x=x1, y=y1).values >= 1000, onp.nan, ds_tm["tt50_q_ss"].isel(x=x1, y=y1).values)
axes[3, 2].axvline(date_hm[t_dry], color="red", alpha=0.5)
axes[3, 2].axvline(date_hm[t_wetdry], color="red", alpha=0.5)
axes[3, 2].axvline(date_hm[t_drywet], color="red", alpha=0.5)
axes[3, 2].axvline(date_hm[t_wet], color="red", alpha=0.5)
axes[3, 2].fill_between(
date_tm,
ds_tm["tt10_q_ss"].isel(x=x1, y=y1).values,
ds_tm["tt90_q_ss"].isel(x=x1, y=y1).values,
vals1,
vals2,
color="grey",
edgecolor=None,
alpha=0.2,
)
axes[3, 2].plot(date_tm, ds_tm["tt50_q_ss"].isel(x=x1, y=y1).values, "-", color="grey", lw=1)
axes[3, 2].plot(date_tm, vals3, "-", color="grey", lw=1)
axes[3, 2].set_xlim(date_tm[0], date_tm[-1])
axes[3, 2].set_ylabel(r"age [days]")
axes[3, 2].set_xlabel(r"Time [year-month]")
Expand Down
8 changes: 4 additions & 4 deletions examples/catchment_scale/eberbaechle/svat_oxygen18_distributed/merge_output.py
View file
Original file line number Diff line number Diff line change
Expand Up @@ -51,12 +51,12 @@
if not f.dimensions:
f.dimensions = dict_dim
v = f.create_variable("x", ("x",), float, compression="gzip", compression_opts=1)
v.attrs["long_name"] = "model run"
v.attrs["units"] = ""
v.attrs["long_name"] = "x"
v.attrs["units"] = "m"
v[:] = onp.arange(dict_dim["x"])
v = f.create_variable("y", ("y",), float, compression="gzip", compression_opts=1)
v.attrs["long_name"] = ""
v.attrs["units"] = ""
v.attrs["long_name"] = "y"
v.attrs["units"] = "m"
v[:] = onp.arange(dict_dim["y"])
v = f.create_variable("ages", ("ages",), float, compression="gzip", compression_opts=1)
v.attrs["long_name"] = "Water ages"
Expand Down
37 changes: 19 additions & 18 deletions examples/catchment_scale/eberbaechle/svat_oxygen18_distributed/svat_oxygen18.py
View file
Original file line number Diff line number Diff line change
Expand Up @@ -185,21 +185,21 @@ def set_initial_conditions_setup(self, state):

vs.S_snow = update(
vs.S_snow,
at[2:-2, 2:-2, : vs.taup1],
at[2:-2, 2:-2, :vs.taup1],
self._read_var_from_nc("S_snow", self._base_path, "SVAT.nc")[:, :, vs.itt, npx.newaxis],
)
vs.S_rz = update(
vs.S_rz,
at[2:-2, 2:-2, : vs.taup1],
at[2:-2, 2:-2, :vs.taup1],
self._read_var_from_nc("S_rz", self._base_path, "SVAT.nc")[:, :, vs.itt, npx.newaxis],
)
vs.S_ss = update(
vs.S_ss,
at[2:-2, 2:-2, : vs.taup1],
at[2:-2, 2:-2, :vs.taup1],
self._read_var_from_nc("S_ss", self._base_path, "SVAT.nc")[:, :, vs.itt, npx.newaxis],
)
vs.S_s = update(
vs.S_s, at[2:-2, 2:-2, : vs.taup1], vs.S_rz[2:-2, 2:-2, : vs.taup1] + vs.S_ss[2:-2, 2:-2, : vs.taup1]
vs.S_s, at[2:-2, 2:-2, :vs.taup1], vs.S_rz[2:-2, 2:-2, :vs.taup1] + vs.S_ss[2:-2, 2:-2, :vs.taup1]
)
vs.S_rz_init = update(vs.S_rz_init, at[2:-2, 2:-2], vs.S_rz[2:-2, 2:-2, 0])
vs.S_ss_init = update(vs.S_ss_init, at[2:-2, 2:-2], vs.S_ss[2:-2, 2:-2, 0])
Expand All @@ -212,14 +212,14 @@ def set_initial_conditions(self, state):
arr0 = allocate(state.dimensions, ("x", "y"))
vs.sa_rz = update(
vs.sa_rz,
at[2:-2, 2:-2, : vs.taup1, 1:],
at[2:-2, 2:-2, :vs.taup1, 1:],
npx.diff(npx.linspace(arr0[2:-2, 2:-2], vs.S_rz[2:-2, 2:-2, vs.tau], settings.ages, axis=-1), axis=-1)[
:, :, npx.newaxis, :
],
)
vs.sa_ss = update(
vs.sa_ss,
at[2:-2, 2:-2, : vs.taup1, 1:],
at[2:-2, 2:-2, :vs.taup1, 1:],
npx.diff(npx.linspace(arr0[2:-2, 2:-2], vs.S_ss[2:-2, 2:-2, vs.tau], settings.ages, axis=-1), axis=-1)[
:, :, npx.newaxis, :
],
Expand Down Expand Up @@ -249,37 +249,38 @@ def set_initial_conditions(self, state):
)

if settings.enable_oxygen18:
vs.C_iso_snow = update(vs.C_iso_snow, at[2:-2, 2:-2, : vs.taup1], npx.nan)
vs.C_iso_rz = update(vs.C_iso_rz, at[2:-2, 2:-2, : vs.taup1], self._config["d18O_RZ"])
vs.C_iso_ss = update(vs.C_iso_ss, at[2:-2, 2:-2, : vs.taup1], self._config["d18O_SS"])
vs.C_snow = update(vs.C_snow, at[2:-2, 2:-2, :vs.taup1], npx.nan)
vs.C_iso_snow = update(vs.C_iso_snow, at[2:-2, 2:-2, :vs.taup1], npx.nan)
vs.C_iso_rz = update(vs.C_iso_rz, at[2:-2, 2:-2, :vs.taup1], self._config["d18O_RZ"])
vs.C_iso_ss = update(vs.C_iso_ss, at[2:-2, 2:-2, :vs.taup1], self._config["d18O_SS"])
vs.C_rz = update(
vs.C_rz,
at[2:-2, 2:-2, : vs.taup1],
at[2:-2, 2:-2, :vs.taup1],
delta_to_conc(state, vs.C_iso_rz[2:-2, 2:-2, vs.tau, npx.newaxis]),
)
vs.msa_rz = update(
vs.msa_rz,
at[2:-2, 2:-2, : vs.taup1, :],
vs.C_rz[2:-2, 2:-2, : vs.taup1, npx.newaxis],
at[2:-2, 2:-2, :vs.taup1, :],
vs.C_rz[2:-2, 2:-2, :vs.taup1, npx.newaxis],
)
vs.msa_rz = update(
vs.msa_rz,
at[2:-2, 2:-2, : vs.taup1, 0],
at[2:-2, 2:-2, :vs.taup1, 0],
0,
)
vs.C_ss = update(
vs.C_ss,
at[2:-2, 2:-2, : vs.taup1],
at[2:-2, 2:-2, :vs.taup1],
delta_to_conc(state, vs.C_iso_ss[2:-2, 2:-2, vs.tau, npx.newaxis]),
)
vs.msa_ss = update(
vs.msa_ss,
at[2:-2, 2:-2, : vs.taup1, :],
vs.C_ss[2:-2, 2:-2, : vs.taup1, npx.newaxis],
at[2:-2, 2:-2, :vs.taup1, :],
vs.C_ss[2:-2, 2:-2, :vs.taup1, npx.newaxis],
)
vs.msa_ss = update(
vs.msa_ss,
at[2:-2, 2:-2, : vs.taup1, 0],
at[2:-2, 2:-2, :vs.taup1, 0],
0,
)
vs.msa_s = update(
Expand All @@ -296,7 +297,7 @@ def set_initial_conditions(self, state):
)
vs.msa_s = update(
vs.msa_s,
at[2:-2, 2:-2, : vs.taup1, 0],
at[2:-2, 2:-2, :vs.taup1, 0],
0,
)
vs.C_s = update(
Expand Down
1 change: 1 addition & 0 deletions examples/plot_scale/rietholzbach/svat_oxygen18/svat_oxygen18.py
View file
Original file line number Diff line number Diff line change
Expand Up @@ -281,6 +281,7 @@ def set_initial_conditions(self, state):
)

if settings.enable_oxygen18:
vs.C_snow = update(vs.C_snow, at[2:-2, 2:-2, :vs.taup1], npx.nan)
vs.C_iso_snow = update(vs.C_iso_snow, at[2:-2, 2:-2, : vs.taup1], npx.nan)
vs.C_iso_rz = update(vs.C_iso_rz, at[2:-2, 2:-2, : vs.taup1], -10.5)
vs.C_iso_ss = update(vs.C_iso_ss, at[2:-2, 2:-2, : vs.taup1], -10.5)
Expand Down
37 changes: 19 additions & 18 deletions examples/plot_scale/rietholzbach/svat_oxygen18_monte_carlo/svat_oxygen18.py
View file
Original file line number Diff line number Diff line change
Expand Up @@ -524,27 +524,27 @@ def set_initial_conditions_setup(self, state):

vs.S_snow = update(
vs.S_snow,
at[2:-2, 2:-2, : vs.taup1],
at[2:-2, 2:-2, :vs.taup1],
self._read_var_from_nc("S_snow", self._base_path / "input", self._states_hm_file)[
:, :, vs.itt, npx.newaxis
],
)
vs.S_rz = update(
vs.S_rz,
at[2:-2, 2:-2, : vs.taup1],
at[2:-2, 2:-2, :vs.taup1],
self._read_var_from_nc("S_rz", self._base_path / "input", self._states_hm_file)[
:, :, vs.itt, npx.newaxis
],
)
vs.S_ss = update(
vs.S_ss,
at[2:-2, 2:-2, : vs.taup1],
at[2:-2, 2:-2, :vs.taup1],
self._read_var_from_nc("S_ss", self._base_path / "input", self._states_hm_file)[
:, :, vs.itt, npx.newaxis
],
)
vs.S_s = update(
vs.S_s, at[2:-2, 2:-2, : vs.taup1], vs.S_rz[2:-2, 2:-2, : vs.taup1] + vs.S_ss[2:-2, 2:-2, : vs.taup1]
vs.S_s, at[2:-2, 2:-2, :vs.taup1], vs.S_rz[2:-2, 2:-2, :vs.taup1] + vs.S_ss[2:-2, 2:-2, :vs.taup1]
)
vs.S_rz_init = update(vs.S_rz_init, at[2:-2, 2:-2], vs.S_rz[2:-2, 2:-2, 0])
vs.S_ss_init = update(vs.S_ss_init, at[2:-2, 2:-2], vs.S_ss[2:-2, 2:-2, 0])
Expand All @@ -558,14 +558,14 @@ def set_initial_conditions(self, state):
arr0 = allocate(state.dimensions, ("x", "y"))
vs.sa_rz = update(
vs.sa_rz,
at[2:-2, 2:-2, : vs.taup1, 1:],
at[2:-2, 2:-2, :vs.taup1, 1:],
npx.diff(npx.linspace(arr0[2:-2, 2:-2], vs.S_rz[2:-2, 2:-2, vs.tau], settings.ages, axis=-1), axis=-1)[
:, :, npx.newaxis, :
],
)
vs.sa_ss = update(
vs.sa_ss,
at[2:-2, 2:-2, : vs.taup1, 1:],
at[2:-2, 2:-2, :vs.taup1, 1:],
npx.diff(npx.linspace(arr0[2:-2, 2:-2], vs.S_ss[2:-2, 2:-2, vs.tau], settings.ages, axis=-1), axis=-1)[
:, :, npx.newaxis, :
],
Expand Down Expand Up @@ -595,37 +595,38 @@ def set_initial_conditions(self, state):
)

if settings.enable_oxygen18:
vs.C_iso_snow = update(vs.C_iso_snow, at[2:-2, 2:-2, : vs.taup1], npx.nan)
vs.C_iso_rz = update(vs.C_iso_rz, at[2:-2, 2:-2, : vs.taup1], -10.5)
vs.C_iso_ss = update(vs.C_iso_ss, at[2:-2, 2:-2, : vs.taup1], -10.5)
vs.C_snow = update(vs.C_snow, at[2:-2, 2:-2, :vs.taup1], npx.nan)
vs.C_iso_snow = update(vs.C_iso_snow, at[2:-2, 2:-2, :vs.taup1], npx.nan)
vs.C_iso_rz = update(vs.C_iso_rz, at[2:-2, 2:-2, :vs.taup1], -10.5)
vs.C_iso_ss = update(vs.C_iso_ss, at[2:-2, 2:-2, :vs.taup1], -10.5)
vs.C_rz = update(
vs.C_rz,
at[2:-2, 2:-2, : vs.taup1],
at[2:-2, 2:-2, :vs.taup1],
delta_to_conc(state, vs.C_iso_rz[2:-2, 2:-2, vs.tau, npx.newaxis]),
)
vs.msa_rz = update(
vs.msa_rz,
at[2:-2, 2:-2, : vs.taup1, :],
vs.C_rz[2:-2, 2:-2, : vs.taup1, npx.newaxis],
at[2:-2, 2:-2, :vs.taup1, :],
vs.C_rz[2:-2, 2:-2, :vs.taup1, npx.newaxis],
)
vs.msa_rz = update(
vs.msa_rz,
at[2:-2, 2:-2, : vs.taup1, 0],
at[2:-2, 2:-2, :vs.taup1, 0],
0,
)
vs.C_ss = update(
vs.C_ss,
at[2:-2, 2:-2, : vs.taup1],
at[2:-2, 2:-2, :vs.taup1],
delta_to_conc(state, vs.C_iso_ss[2:-2, 2:-2, vs.tau, npx.newaxis]),
)
vs.msa_ss = update(
vs.msa_ss,
at[2:-2, 2:-2, : vs.taup1, :],
vs.C_ss[2:-2, 2:-2, : vs.taup1, npx.newaxis],
at[2:-2, 2:-2, :vs.taup1, :],
vs.C_ss[2:-2, 2:-2, :vs.taup1, npx.newaxis],
)
vs.msa_ss = update(
vs.msa_ss,
at[2:-2, 2:-2, : vs.taup1, 0],
at[2:-2, 2:-2, :vs.taup1, 0],
0,
)
vs.msa_s = update(
Expand All @@ -642,7 +643,7 @@ def set_initial_conditions(self, state):
)
vs.msa_s = update(
vs.msa_s,
at[2:-2, 2:-2, : vs.taup1, 0],
at[2:-2, 2:-2, :vs.taup1, 0],
0,
)
vs.C_s = update(
Expand Down
37 changes: 19 additions & 18 deletions examples/plot_scale/rietholzbach/svat_oxygen18_sensitivity/svat_oxygen18.py
View file
Original file line number Diff line number Diff line change
Expand Up @@ -285,27 +285,27 @@ def set_initial_conditions_setup(self, state):

vs.S_snow = update(
vs.S_snow,
at[2:-2, 2:-2, : vs.taup1],
at[2:-2, 2:-2, :vs.taup1],
self._read_var_from_nc(
"S_snow", self._base_path / "input", f"states_hm_saltelli_for_{transport_model_structure}.nc"
)[:, :, vs.itt, npx.newaxis],
)
vs.S_rz = update(
vs.S_rz,
at[2:-2, 2:-2, : vs.taup1],
at[2:-2, 2:-2, :vs.taup1],
self._read_var_from_nc(
"S_rz", self._base_path / "input", f"states_hm_saltelli_for_{transport_model_structure}.nc"
)[:, :, vs.itt, npx.newaxis],
)
vs.S_ss = update(
vs.S_ss,
at[2:-2, 2:-2, : vs.taup1],
at[2:-2, 2:-2, :vs.taup1],
self._read_var_from_nc(
"S_ss", self._base_path / "input", f"states_hm_saltelli_for_{transport_model_structure}.nc"
)[:, :, vs.itt, npx.newaxis],
)
vs.S_s = update(
vs.S_s, at[2:-2, 2:-2, : vs.taup1], vs.S_rz[2:-2, 2:-2, : vs.taup1] + vs.S_ss[2:-2, 2:-2, : vs.taup1]
vs.S_s, at[2:-2, 2:-2, :vs.taup1], vs.S_rz[2:-2, 2:-2, :vs.taup1] + vs.S_ss[2:-2, 2:-2, :vs.taup1]
)
vs.S_rz_init = update(vs.S_rz_init, at[2:-2, 2:-2], vs.S_rz[2:-2, 2:-2, 0])
vs.S_ss_init = update(vs.S_ss_init, at[2:-2, 2:-2], vs.S_ss[2:-2, 2:-2, 0])
Expand All @@ -318,14 +318,14 @@ def set_initial_conditions(self, state):
arr0 = allocate(state.dimensions, ("x", "y"))
vs.sa_rz = update(
vs.sa_rz,
at[2:-2, 2:-2, : vs.taup1, 1:],
at[2:-2, 2:-2, :vs.taup1, 1:],
npx.diff(npx.linspace(arr0[2:-2, 2:-2], vs.S_rz[2:-2, 2:-2, vs.tau], settings.ages, axis=-1), axis=-1)[
:, :, npx.newaxis, :
],
)
vs.sa_ss = update(
vs.sa_ss,
at[2:-2, 2:-2, : vs.taup1, 1:],
at[2:-2, 2:-2, :vs.taup1, 1:],
npx.diff(npx.linspace(arr0[2:-2, 2:-2], vs.S_ss[2:-2, 2:-2, vs.tau], settings.ages, axis=-1), axis=-1)[
:, :, npx.newaxis, :
],
Expand Down Expand Up @@ -355,37 +355,38 @@ def set_initial_conditions(self, state):
)

if settings.enable_oxygen18:
vs.C_iso_snow = update(vs.C_iso_snow, at[2:-2, 2:-2, : vs.taup1], npx.nan)
vs.C_iso_rz = update(vs.C_iso_rz, at[2:-2, 2:-2, : vs.taup1], -10.5)
vs.C_iso_ss = update(vs.C_iso_ss, at[2:-2, 2:-2, : vs.taup1], -10.5)
vs.C_snow = update(vs.C_snow, at[2:-2, 2:-2, :vs.taup1], npx.nan)
vs.C_iso_snow = update(vs.C_iso_snow, at[2:-2, 2:-2, :vs.taup1], npx.nan)
vs.C_iso_rz = update(vs.C_iso_rz, at[2:-2, 2:-2, :vs.taup1], -10.5)
vs.C_iso_ss = update(vs.C_iso_ss, at[2:-2, 2:-2, :vs.taup1], -10.5)
vs.C_rz = update(
vs.C_rz,
at[2:-2, 2:-2, : vs.taup1],
at[2:-2, 2:-2, :vs.taup1],
delta_to_conc(state, vs.C_iso_rz[2:-2, 2:-2, vs.tau, npx.newaxis]),
)
vs.msa_rz = update(
vs.msa_rz,
at[2:-2, 2:-2, : vs.taup1, :],
vs.C_rz[2:-2, 2:-2, : vs.taup1, npx.newaxis],
at[2:-2, 2:-2, :vs.taup1, :],
vs.C_rz[2:-2, 2:-2, :vs.taup1, npx.newaxis],
)
vs.msa_rz = update(
vs.msa_rz,
at[2:-2, 2:-2, : vs.taup1, 0],
at[2:-2, 2:-2, :vs.taup1, 0],
0,
)
vs.C_ss = update(
vs.C_ss,
at[2:-2, 2:-2, : vs.taup1],
at[2:-2, 2:-2, :vs.taup1],
delta_to_conc(state, vs.C_iso_ss[2:-2, 2:-2, vs.tau, npx.newaxis]),
)
vs.msa_ss = update(
vs.msa_ss,
at[2:-2, 2:-2, : vs.taup1, :],
vs.C_ss[2:-2, 2:-2, : vs.taup1, npx.newaxis],
at[2:-2, 2:-2, :vs.taup1, :],
vs.C_ss[2:-2, 2:-2, :vs.taup1, npx.newaxis],
)
vs.msa_ss = update(
vs.msa_ss,
at[2:-2, 2:-2, : vs.taup1, 0],
at[2:-2, 2:-2, :vs.taup1, 0],
0,
)
vs.msa_s = update(
Expand All @@ -402,7 +403,7 @@ def set_initial_conditions(self, state):
)
vs.msa_s = update(
vs.msa_s,
at[2:-2, 2:-2, : vs.taup1, 0],
at[2:-2, 2:-2, :vs.taup1, 0],
0,
)
vs.C_s = update(
Expand Down
Loading

0 comments on commit 0216045

Please sign in to comment.