Skip to content

Commit

Permalink
modified oneD_distributed tutorial
Browse files Browse the repository at this point in the history
  • Loading branch information
@schwemro
schwemro committed Jul 13, 2023
1 parent 664cab7 commit 8a8b011
Show file tree
Hide file tree
Showing 4 changed files with 111 additions and 117 deletions.
64 changes: 19 additions & 45 deletions examples/hillslope_scale/oneD_distributed_tutorial/oneD.py
View file
Original file line number Diff line number Diff line change
Expand Up @@ -62,6 +62,7 @@ def set_settings(self, state):
settings.nx, settings.ny = self._config["nx"], self._config["ny"]
# derive total number of time steps from forcing
settings.runlen = self._get_runlen(self._input_dir, "forcing.nc")
settings.nitt_forc = len(self._read_var_from_nc("Time", self._input_dir, 'forcing.nc'))

# spatial discretization (in meters)
settings.dx = self._config["dx"]
Expand Down Expand Up @@ -245,67 +246,40 @@ def set_boundary_conditions_setup(self, state):
def set_boundary_conditions(self, state):
pass

@roger_routine
def set_forcing_setup(self, state):
pass

@roger_routine(
dist_safe=False,
local_variables=[
"time",
"itt_day",
"itt_forc",
"prec_day",
"ta_day",
"pet_day",
"year",
"month",
"doy",
"PREC",
"TA",
"PET",
],
)
def set_forcing_setup(self, state):
vs = state.variables

vs.PREC = update(vs.PREC, at[:], self._read_var_from_nc("PREC", self._input_dir, 'forcing.nc')[0, 0, :])
vs.TA = update(vs.TA, at[:], self._read_var_from_nc("TA", self._input_dir, 'forcing.nc')[0, 0, :])
vs.PET = update(vs.PET, at[:], self._read_var_from_nc("PET", self._input_dir, 'forcing.nc')[0, 0, :])

@roger_routine
def set_forcing(self, state):
vs = state.variables

condt = vs.time % (24 * 60 * 60) == 0
condt = (vs.time % (24 * 60 * 60) == 0)
if condt:
vs.itt_day = 0
vs.year = update(
vs.year, at[1], self._read_var_from_nc("YEAR", self._input_dir, "forcing.nc")[vs.itt_forc]
)
vs.month = update(
vs.month, at[1], self._read_var_from_nc("MONTH", self._input_dir, "forcing.nc")[vs.itt_forc]
)
vs.doy = update(
vs.doy, at[1], self._read_var_from_nc("DOY", self._input_dir, "forcing.nc")[vs.itt_forc]
)
vs.prec_day = update(
vs.prec_day,
at[2:-2, 2:-2, :],
self._read_var_from_nc("PREC", self._input_dir, "forcing.nc")[
:, :, vs.itt_forc : vs.itt_forc + 6 * 24
],
)
vs.ta_day = update(
vs.ta_day,
at[2:-2, 2:-2, :],
self._read_var_from_nc("TA", self._input_dir, "forcing.nc")[
:, :, vs.itt_forc : vs.itt_forc + 6 * 24
],
)
vs.pet_day = update(
vs.pet_day,
at[2:-2, 2:-2, :],
self._read_var_from_nc("PET", self._input_dir, "forcing.nc")[
:, :, vs.itt_forc : vs.itt_forc + 6 * 24
],
)
vs.year = update(vs.year, at[1], self._read_var_from_nc("YEAR", self._input_dir, 'forcing.nc')[vs.itt_forc])
vs.month = update(vs.month, at[1], self._read_var_from_nc("MONTH", self._input_dir, 'forcing.nc')[vs.itt_forc])
vs.doy = update(vs.doy, at[1], self._read_var_from_nc("DOY", self._input_dir, 'forcing.nc')[vs.itt_forc])
vs.prec_day = update(vs.prec_day, at[:, :, :], vs.PREC[npx.newaxis, npx.newaxis, vs.itt_forc:vs.itt_forc+6*24])
vs.ta_day = update(vs.ta_day, at[:, :, :], vs.TA[npx.newaxis, npx.newaxis, vs.itt_forc:vs.itt_forc+6*24])
vs.pet_day = update(vs.pet_day, at[:, :, :], vs.PET[npx.newaxis, npx.newaxis, vs.itt_forc:vs.itt_forc+6*24])
vs.itt_forc = vs.itt_forc + 6 * 24

@roger_routine
def set_diagnostics(self, state, base_path=tmp_dir):
diagnostics = state.diagnostics

# variables written to output files
diagnostics["rate"].output_variables = self._config["OUTPUT_RATE"]
# values are aggregated to daily
diagnostics["rate"].output_frequency = 24 * 60 * 60 # in seconds
Expand Down
121 changes: 65 additions & 56 deletions examples/plot_scale/freiburg_altheim_kupferzell/post_processing.py
View file
Original file line number Diff line number Diff line change
Expand Up @@ -1122,62 +1122,71 @@
# fig.savefig(file, dpi=250)
# plt.close('all')

for var_sim in vars_sim:
for delta in deltas:
for soil_depth in soil_depths:
cond = _soil_depths[soil_depth]
for future in ["nf", "ff"]:
fig, axes = plt.subplots(
len(locations), len(land_cover_scenarios), sharex="col", sharey="row", figsize=(6, 4.5)
)
for i, location in enumerate(locations):
for j, land_cover_scenario in enumerate(land_cover_scenarios):
values = dict_glm[location][land_cover_scenario]["CCCma-CanESM2_CCLM4-8-17"][var_sim][delta][
soil_depth
][future]["params"][1:]
df_params_canesm = pd.DataFrame(index=range(len(values)), columns=["value", "Parameter"])
df_params_canesm.loc[:, "value"] = values
df_params_canesm.loc[:, "Parameter"] = [
r"$\theta_{pwp}$",
r"$\theta_{ufc}$",
r"$\theta_{ac}$",
r"$k_s$",
]
df_params_canesm.loc[:, "Climate model"] = "CCCma-CanESM2_CCLM4-8-17"

values = dict_glm[location][land_cover_scenario]["MPI-M-MPI-ESM-LR_RCA4"][var_sim][delta][
soil_depth
][future]["params"][1:]
df_params_mpiesm = pd.DataFrame(index=range(len(values)), columns=["value", "Parameter"])
df_params_mpiesm.loc[:, "value"] = values
df_params_mpiesm.loc[:, "Parameter"] = [
r"$\theta_{pwp}$",
r"$\theta_{ufc}$",
r"$\theta_{ac}$",
r"$k_s$",
]
df_params_mpiesm.loc[:, "Climate model"] = "MPI-M-MPI-ESM-LR_RCA4"

df_params = pd.concat([df_params_canesm, df_params_mpiesm], ignore_index=True)
sns.catplot(
x="Parameter",
y="value",
hue="Climate model",
palette=["red", "blue"],
data=df_params,
ax=axes[i, j],
kind="bar",
)
axes[i, j].legend([], [], frameon=False)
axes[0, j].set_title(f"{Land_cover_scenarios[j]}")
axes[i, j].set_ylabel("")
axes[i, j].set_xlabel("")
axes[i, j].tick_params(axis="x", rotation=33)
axes[i, 0].set_ylabel(f"{Locations[i]}\n{_lab[delta]} {_lab[var_sim]} [%]")
fig.tight_layout()
file = base_path_figs / f"{var_sim}_{delta}_{soil_depth}_{future}_barplot.png"
fig.savefig(file, dpi=250)
plt.close("all")
import pickle
# Store data (serialize)
with open(base_path_figs / 'glm_results.pkl', 'wb') as handle:
pickle.dump(dict_glm, handle, protocol=pickle.HIGHEST_PROTOCOL)

# # Load data (deserialize)
# with open(base_path_figs / 'glm_results.pkl', 'rb') as handle:
# dict_glm = pickle.load(handle)

# for var_sim in vars_sim:
# for delta in deltas:
# for soil_depth in soil_depths:
# cond = _soil_depths[soil_depth]
# for future in ["nf", "ff"]:
# fig, axes = plt.subplots(
# len(locations), len(land_cover_scenarios), sharex="col", sharey="row", figsize=(6, 4.5)
# )
# for i, location in enumerate(locations):
# for j, land_cover_scenario in enumerate(land_cover_scenarios):
# values = dict_glm[location][land_cover_scenario]["CCCma-CanESM2_CCLM4-8-17"][var_sim][delta][
# soil_depth
# ][future]["params"][1:]
# df_params_canesm = pd.DataFrame(index=range(len(values)), columns=["value", "Parameter"])
# df_params_canesm.loc[:, "value"] = values
# df_params_canesm.loc[:, "Parameter"] = [
# r"$\theta_{pwp}$",
# r"$\theta_{ufc}$",
# r"$\theta_{ac}$",
# r"$k_s$",
# ]
# df_params_canesm.loc[:, "Climate model"] = "CCCma-CanESM2_CCLM4-8-17"

# values = dict_glm[location][land_cover_scenario]["MPI-M-MPI-ESM-LR_RCA4"][var_sim][delta][
# soil_depth
# ][future]["params"][1:]
# df_params_mpiesm = pd.DataFrame(index=range(len(values)), columns=["value", "Parameter"])
# df_params_mpiesm.loc[:, "value"] = values
# df_params_mpiesm.loc[:, "Parameter"] = [
# r"$\theta_{pwp}$",
# r"$\theta_{ufc}$",
# r"$\theta_{ac}$",
# r"$k_s$",
# ]
# df_params_mpiesm.loc[:, "Climate model"] = "MPI-M-MPI-ESM-LR_RCA4"

# df_params = pd.concat([df_params_canesm, df_params_mpiesm], ignore_index=True)
# sns.barplot(
# x="Parameter",
# y="value",
# hue="Climate model",
# palette=["red", "blue"],
# data=df_params,
# ax=axes[i, j],
# errorbar=None
# )
# axes[i, j].legend([], [], frameon=False)
# axes[0, j].set_title(f"{Land_cover_scenarios[j]}")
# axes[i, j].set_ylabel("")
# axes[i, j].set_xlabel("")
# axes[i, j].tick_params(axis="x", rotation=33)
# axes[i, 0].set_ylabel(f"{Locations[i]}\n{_lab[delta]} {_lab[var_sim]} [%]")
# fig.tight_layout()
# file = base_path_figs / f"{var_sim}_{delta}_{soil_depth}_{future}_barplot.png"
# fig.savefig(file, dpi=250)
# plt.close("all")


# norm = mpl.colors.Normalize(vmin=-1, vmax=1)
Expand Down
31 changes: 15 additions & 16 deletions examples/plot_scale/oneD_tutorial/oneD.py
View file
Original file line number Diff line number Diff line change
Expand Up @@ -60,6 +60,7 @@ def set_settings(self, state):
settings.nx, settings.ny = 1, 1
# derive total number of time steps from forcing
settings.runlen = self._get_runlen(self._input_dir, 'forcing.nc')
settings.nitt_forc = len(self._read_var_from_nc("Time", self._input_dir, 'forcing.nc'))

# spatial discretization (in meters)
settings.dx = 1
Expand Down Expand Up @@ -192,24 +193,22 @@ def set_boundary_conditions_setup(self, state):
def set_boundary_conditions(self, state):
pass

@roger_routine
def set_forcing_setup(self, state):
pass

@roger_routine(
dist_safe=False,
local_variables=[
"time",
"itt_day",
"itt_forc",
"prec_day",
"ta_day",
"pet_day",
"year",
"month",
"doy",
"PREC",
"TA",
"PET",
],
)
def set_forcing_setup(self, state):
vs = state.variables

vs.PREC = update(vs.PREC, at[:], self._read_var_from_nc("PREC", self._input_dir, 'forcing.nc')[0, 0, :])
vs.TA = update(vs.TA, at[:], self._read_var_from_nc("TA", self._input_dir, 'forcing.nc')[0, 0, :])
vs.PET = update(vs.PET, at[:], self._read_var_from_nc("PET", self._input_dir, 'forcing.nc')[0, 0, :])

@roger_routine
def set_forcing(self, state):
vs = state.variables

Expand All @@ -219,9 +218,9 @@ def set_forcing(self, state):
vs.year = update(vs.year, at[1], self._read_var_from_nc("YEAR", self._input_dir, 'forcing.nc')[vs.itt_forc])
vs.month = update(vs.month, at[1], self._read_var_from_nc("MONTH", self._input_dir, 'forcing.nc')[vs.itt_forc])
vs.doy = update(vs.doy, at[1], self._read_var_from_nc("DOY", self._input_dir, 'forcing.nc')[vs.itt_forc])
vs.prec_day = update(vs.prec_day, at[2:-2, 2:-2, :], self._read_var_from_nc("PREC", self._input_dir, 'forcing.nc')[:, :, vs.itt_forc:vs.itt_forc+6*24])
vs.ta_day = update(vs.ta_day, at[2:-2, 2:-2, :], self._read_var_from_nc("TA", self._input_dir, 'forcing.nc')[:, :, vs.itt_forc:vs.itt_forc+6*24])
vs.pet_day = update(vs.pet_day, at[2:-2, 2:-2, :], self._read_var_from_nc("PET", self._input_dir, 'forcing.nc')[:, :, vs.itt_forc:vs.itt_forc+6*24])
vs.prec_day = update(vs.prec_day, at[:, :, :], vs.PREC[npx.newaxis, npx.newaxis, vs.itt_forc:vs.itt_forc+6*24])
vs.ta_day = update(vs.ta_day, at[:, :, :], vs.TA[npx.newaxis, npx.newaxis, vs.itt_forc:vs.itt_forc+6*24])
vs.pet_day = update(vs.pet_day, at[:, :, :], vs.PET[npx.newaxis, npx.newaxis, vs.itt_forc:vs.itt_forc+6*24])
vs.itt_forc = vs.itt_forc + 6 * 24

@roger_routine
Expand Down
12 changes: 12 additions & 0 deletions roger/core/soil.py
View file
Original file line number Diff line number Diff line change
Expand Up @@ -398,6 +398,18 @@ def loop_body(i, z_root):
* vs.maskCatch[2:-2, 2:-2],
)

vs.z_root = update(
vs.z_root,
at[2:-2, 2:-2, 0],
z_root[2:-2, 2:-2],
)

vs.z_root = update(
vs.z_root,
at[2:-2, 2:-2, 1],
z_root[2:-2, 2:-2],
)

# set thickness of upper soil water storage to 20 cm for bare soils
mask_crops = npx.isin(vs.lu_id, npx.arange(500, 600, 1, dtype=npx.int32))
vs.z_root = update(
Expand Down

0 comments on commit 8a8b011

Please sign in to comment.