-RoGeR, *Runoff Generation Research*, is a process-based hydrological model that can be applied from plot to catchment scale. RoGeR is written in pure Python, which facilitates model setup and model workflows. We want to enable high-performance hydrologic modelling with a clear focus on flexibility and usability.
+RoGeR, *Runoff Generation Research*, is a process-based hydrological model that can be applied from plot to catchment scale. RoGeR is written in pure Python, which facilitates model setup and model workflows. We want to enable high-performance hydrological modelling with a clear focus on flexibility and usability.
RoGeR supports a NumPy backend for small-scale problems, and a
high-performance [JAX](https://github.com/google/jax) backend
@@ -36,16 +36,15 @@ We strongly recommend to [visit our documentation](https://roger.readthedocs.io/
## Features
(25 square meter resolved simulations -of a rectangular soil covered -by grass, click for better +of the Eberbaechle catchment, +Germany (2019-2022), click for better quality)
@@ -112,7 +111,7 @@ If you use Roger in scientific work, please consider citing [the following publi ```bibtex @article{ - title = {Roger v3.0 – a process-based hydrologic toolbox model in {Python}}, + title = {Roger v3.0.3 – a process-based hydrologic toolbox model in {Python}}, volume = {...}, doi = {https://doi.org/10.5194/gmd-2023-118}, journal = {Geosci. Model Dev.}, diff --git a/doc/_images/fluxes_theta_and_tt_rt.gif b/doc/_images/fluxes_theta_and_tt_rt.gif new file mode 100644 index 000000000..35cda789d Binary files /dev/null and b/doc/_images/fluxes_theta_and_tt_rt.gif differ diff --git a/doc/_images/fluxes_theta_and_tt_rt.mp4 b/doc/_images/fluxes_theta_and_tt_rt.mp4 new file mode 100644 index 000000000..b0c947974 Binary files /dev/null and b/doc/_images/fluxes_theta_and_tt_rt.mp4 differ diff --git a/doc/_images/theta_and_tt.gif b/doc/_images/theta_and_tt.gif deleted file mode 100644 index 3104c6769..000000000 Binary files a/doc/_images/theta_and_tt.gif and /dev/null differ diff --git a/doc/_images/theta_and_tt.mp4 b/doc/_images/theta_and_tt.mp4 deleted file mode 100644 index 387a76a00..000000000 Binary files a/doc/_images/theta_and_tt.mp4 and /dev/null differ diff --git a/examples/catchment_scale/eberbaechle/make_figures.py b/examples/catchment_scale/eberbaechle/make_figures.py index a570decb8..83fe4adb9 100644 --- a/examples/catchment_scale/eberbaechle/make_figures.py +++ b/examples/catchment_scale/eberbaechle/make_figures.py @@ -38,10 +38,10 @@ ncols = 80 x1 = 76 y1 = 43 -t_dry = 1079 -t_wet = 10 -t_drywet = 217 -t_wetdry = 234 +t_dry = 300 +t_wet = 468 +t_drywet = 338 +t_wetdry = 480 base_path = Path(__file__).parent # directory of figures @@ -66,19 +66,46 @@ params_hm_file = base_path / "svat_distributed" / "parameters.nc" ds_params_hm = xr.open_dataset(params_hm_file, engine="h5netcdf") -# lu_id = ds_params_hm["lu_id"].values -# ds_params_hm["lu_id"].isel(x=x1, y=y1) -# ds_params_hm["theta_pwp"].isel(x=x1, y=y1) -# ds_params_hm["z_soil"].isel(x=x1, y=y1) - -# import pandas as pd -# df = pd.DataFrame(index=date_hm, columns=["theta", "transp"]) -# for t in range(len(date_hm)): -# vals = onp.where(ds_hm["theta"].isel(Time=t).values <= -9999, onp.nan, ds_hm["theta"].isel(Time=t).values) -# df.iloc[t, 0] = onp.nanmean(vals) -# vals = onp.where(ds_hm["transp"].isel(Time=t).values <= -9999, onp.nan, ds_hm["transp"].isel(Time=t).values) -# df.iloc[t, 1] = onp.nanmax(vals) -# df.to_csv(base_path_figs / "theta_transp.csv", sep=";", index=True) + +import pandas as pd + +df = pd.DataFrame( + index=date_hm, + columns=[ + "theta_min", + "theta_50", + "theta_avg", + "theta_max", + "evap", + "transp", + "perc_50", + "perc_avg", + "perc_max", + "day", + "rank", + "prob", + ], +) +for t in range(len(date_hm)): + vals = onp.where(ds_hm["theta"].isel(Time=t).values <= -9999, onp.nan, ds_hm["theta"].isel(Time=t).values) + df.iloc[t, 0] = onp.nanmin(vals) + df.iloc[t, 1] = onp.nanmedian(vals) + df.iloc[t, 2] = onp.nanmean(vals) + df.iloc[t, 3] = onp.nanmax(vals) + vals = onp.where(ds_hm["evap_soil"].isel(Time=t).values <= -9999, onp.nan, ds_hm["evap_soil"].isel(Time=t).values) + df.iloc[t, 4] = onp.nanmax(vals) + vals = onp.where(ds_hm["transp"].isel(Time=t).values <= -9999, onp.nan, ds_hm["transp"].isel(Time=t).values) + df.iloc[t, 5] = onp.nanmax(vals) + vals = onp.where(ds_hm["q_ss"].isel(Time=t).values <= -9999, onp.nan, ds_hm["q_ss"].isel(Time=t).values) + df.iloc[t, 6] = onp.nanmedian(vals) + df.iloc[t, 7] = onp.nanmean(vals) + df.iloc[t, 8] = onp.nanmax(vals) +df.loc[:, "day"] = range(len(date_hm)) +df = df.sort_values(by=["theta_50"]) +df.loc[:, "rank"] = range(len(date_hm)) +df.loc[:, "prob"] = df.loc[:, "rank"] / len(date_hm) +df = df.sort_values(by=["day"]) +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" @@ -88,14 +115,14 @@ date_tm = num2date(days_tm, units="days since 2019-10-31", calendar="standard", only_use_cftime_datetimes=False) ds_tm = ds_tm.assign_coords(Time=("Time", date_tm)) -# plot spatially distributed soil moisture and median travel time of percolation at different dates +# plot spatially distributed soil moisture and average travel time of percolation at different dates fig, axes = plt.subplots(2, 2, figsize=(6, 4)) axes[0, 0].imshow( onp.where(ds_hm["theta"].isel(Time=t_dry).values <= -9999, onp.nan, ds_hm["theta"].isel(Time=t_dry).values), extent=(0, 80 * 25, 0, 53 * 25), cmap="viridis_r", - vmin=0.15, - vmax=0.3, + vmin=0.1, + vmax=0.35, ) axes[0, 0].grid(zorder=0) axes[0, 0].set_xlabel("[m]") @@ -105,20 +132,133 @@ onp.where(ds_hm["theta"].isel(Time=t_wet).values <= -9999, onp.nan, ds_hm["theta"].isel(Time=t_wet).values), extent=(0, 80 * 25, 0, 53 * 25), cmap="viridis_r", - vmin=0.15, - vmax=0.3, + vmin=0.1, + vmax=0.35, ) axes[0, 1].grid(zorder=0) axes[0, 1].set_xlabel("[m]") axes[0, 1].set_xlabel("[m]") axes[0, 1].set_title(str(ds_hm["Time"].values[t_wet]).split("T")[0]) cmap = copy.copy(mpl.colormaps.get_cmap("viridis_r")) -norm = mpl.colors.Normalize(vmin=0.15, vmax=0.3) +norm = mpl.colors.Normalize(vmin=0.1, vmax=0.35) axl1 = fig.add_axes([0.85, 0.6, 0.02, 0.3]) -cb1 = mpl.colorbar.ColorbarBase(axl1, cmap=cmap, norm=norm, orientation="vertical", ticks=[0.15, 0.2, 0.25, 0.3]) -cb1.ax.set_yticklabels(["<0.15", "0.2", "0.25", ">0.3"]) +cb1 = mpl.colorbar.ColorbarBase(axl1, cmap=cmap, norm=norm, orientation="vertical", ticks=[0.1, 0.2, 0.3, 0.35]) +cb1.ax.set_yticklabels(["<0.1", "0.2", "0.3", ">0.35"]) cb1.set_label(r"$\theta$ [-]") +axes[1, 0].imshow( + ds_tm["ttavg_q_ss"].isel(Time=t_dry).values, extent=(0, 80 * 25, 0, 53 * 25), cmap="viridis", vmin=1, vmax=200 +) +axes[1, 0].grid(zorder=0) +axes[1, 0].set_xlabel("[m]") +axes[1, 0].set_ylabel("[m]") +axes[1, 0].set_title(str(ds_tm["Time"].values[t_dry]).split("T")[0]) +axes[1, 1].imshow( + ds_tm["ttavg_q_ss"].isel(Time=t_wet).values, extent=(0, 80 * 25, 0, 53 * 25), cmap="viridis", vmin=1, vmax=200 +) +axes[1, 1].grid(zorder=0) +axes[1, 1].set_xlabel("[m]") +axes[1, 1].set_title(str(ds_tm["Time"].values[t_wet]).split("T")[0]) +cmap = copy.copy(mpl.colormaps.get_cmap("viridis")) +norm = mpl.colors.Normalize(vmin=1, vmax=200) +axl2 = fig.add_axes([0.85, 0.13, 0.02, 0.3]) +cb2 = mpl.colorbar.ColorbarBase(axl2, cmap=cmap, norm=norm, orientation="vertical", ticks=[1, 100, 200]) +cb2.ax.set_yticklabels(["1", "100", ">200"]) +cb2.ax.invert_yaxis() +cb2.set_label(r"$\overline{TT}_{PERC}$ [days]") +fig.subplots_adjust(left=0.1, bottom=0.1, top=0.94, right=0.8, wspace=0.3, hspace=0.3) +file = base_path_figs / "theta_and_ttavg.png" +fig.savefig(file, dpi=300) +file = base_path_figs / "theta_and_ttavg.pdf" +fig.savefig(file, dpi=300) + +# plot spatially distributed percolation and average travel time of percolation at different dates +fig, axes = plt.subplots(2, 2, figsize=(6, 4)) +axes[0, 0].imshow( + onp.where(ds_hm["q_ss"].isel(Time=t_dry).values <= -9999, onp.nan, ds_hm["q_ss"].isel(Time=t_dry).values), + extent=(0, 80 * 25, 0, 53 * 25), + cmap="viridis_r", + vmin=0, + vmax=50, +) +axes[0, 0].grid(zorder=0) +axes[0, 0].set_xlabel("[m]") +axes[0, 0].set_ylabel("[m]") +axes[0, 0].set_title(str(ds_hm["Time"].values[t_dry]).split("T")[0]) +axes[0, 1].imshow( + onp.where(ds_hm["q_ss"].isel(Time=t_wet).values <= -9999, onp.nan, ds_hm["q_ss"].isel(Time=t_wet).values), + extent=(0, 80 * 25, 0, 53 * 25), + cmap="viridis_r", + vmin=0, + vmax=50, +) +axes[0, 1].grid(zorder=0) +axes[0, 1].set_xlabel("[m]") +axes[0, 1].set_xlabel("[m]") +axes[0, 1].set_title(str(ds_hm["Time"].values[t_wet]).split("T")[0]) +cmap = copy.copy(mpl.colormaps.get_cmap("viridis_r")) +norm = mpl.colors.Normalize(vmin=0, vmax=50) +axl1 = fig.add_axes([0.85, 0.6, 0.02, 0.3]) +cb1 = mpl.colorbar.ColorbarBase(axl1, cmap=cmap, norm=norm, orientation="vertical", ticks=[0, 25, 50]) +cb1.ax.set_yticklabels(["0", "25", ">50"]) +cb1.set_label(r"$PERC$ [mm/day]") + +axes[1, 0].imshow( + ds_tm["ttavg_q_ss"].isel(Time=t_dry).values, extent=(0, 80 * 25, 0, 53 * 25), cmap="viridis", vmin=1, vmax=200 +) +axes[1, 0].grid(zorder=0) +axes[1, 0].set_xlabel("[m]") +axes[1, 0].set_ylabel("[m]") +axes[1, 0].set_title(str(ds_tm["Time"].values[t_dry]).split("T")[0]) +axes[1, 1].imshow( + ds_tm["ttavg_q_ss"].isel(Time=t_wet).values, extent=(0, 80 * 25, 0, 53 * 25), cmap="viridis", vmin=1, vmax=200 +) +axes[1, 1].grid(zorder=0) +axes[1, 1].set_xlabel("[m]") +axes[1, 1].set_title(str(ds_tm["Time"].values[t_wet]).split("T")[0]) +cmap = copy.copy(mpl.colormaps.get_cmap("viridis")) +norm = mpl.colors.Normalize(vmin=1, vmax=200) +axl2 = fig.add_axes([0.85, 0.13, 0.02, 0.3]) +cb2 = mpl.colorbar.ColorbarBase(axl2, cmap=cmap, norm=norm, orientation="vertical", ticks=[1, 100, 200]) +cb2.ax.set_yticklabels(["1", "100", ">200"]) +cb2.ax.invert_yaxis() +cb2.set_label(r"$\overline{TT}_{PERC}$ [days]") +fig.subplots_adjust(left=0.1, bottom=0.1, top=0.94, right=0.8, wspace=0.3, hspace=0.3) +file = base_path_figs / "perc_and_ttavg.png" +fig.savefig(file, dpi=300) +file = base_path_figs / "perc_and_ttavg.pdf" +fig.savefig(file, dpi=300) + +fig, axes = plt.subplots(2, 2, figsize=(6, 4)) +axes[0, 0].imshow( + onp.where(ds_hm["q_ss"].isel(Time=t_dry).values <= -9999, onp.nan, ds_hm["q_ss"].isel(Time=t_dry).values), + extent=(0, 80 * 25, 0, 53 * 25), + cmap="viridis_r", + vmin=0, + vmax=0.1, +) +axes[0, 0].grid(zorder=0) +axes[0, 0].set_xlabel("[m]") +axes[0, 0].set_ylabel("[m]") +axes[0, 0].set_title(str(ds_hm["Time"].values[t_dry]).split("T")[0]) +axes[0, 1].imshow( + onp.where(ds_hm["q_ss"].isel(Time=t_wet).values <= -9999, onp.nan, ds_hm["q_ss"].isel(Time=t_wet).values), + extent=(0, 80 * 25, 0, 53 * 25), + cmap="viridis_r", + vmin=0, + vmax=0.1, +) +axes[0, 1].grid(zorder=0) +axes[0, 1].set_xlabel("[m]") +axes[0, 1].set_xlabel("[m]") +axes[0, 1].set_title(str(ds_hm["Time"].values[t_wet]).split("T")[0]) +cmap = copy.copy(mpl.colormaps.get_cmap("viridis_r")) +norm = mpl.colors.Normalize(vmin=0, vmax=0.1) +axl1 = fig.add_axes([0.85, 0.6, 0.02, 0.3]) +cb1 = mpl.colorbar.ColorbarBase(axl1, cmap=cmap, norm=norm, orientation="vertical", ticks=[0, 0.05, 0.1]) +cb1.ax.set_yticklabels(["0", "0.05", ">0.1"]) +cb1.set_label(r"$PERC$ [mm/day]") + axes[1, 0].imshow( ds_tm["ttavg_q_ss"].isel(Time=t_dry).values, extent=(0, 80 * 25, 0, 53 * 25), cmap="viridis", vmin=1, vmax=1000 ) @@ -140,11 +280,12 @@ cb2.ax.invert_yaxis() cb2.set_label(r"$\overline{TT}_{PERC}$ [days]") fig.subplots_adjust(left=0.1, bottom=0.1, top=0.94, right=0.8, wspace=0.3, hspace=0.3) -file = base_path_figs / "theta_and_ttavg.png" +file = base_path_figs / "perc1_and_ttavg.png" fig.savefig(file, dpi=300) -file = base_path_figs / "theta_and_ttavg.pdf" +file = base_path_figs / "perc1_and_ttavg.pdf" fig.savefig(file, dpi=300) + # plot fluxes and isotopic signals of a single grid cell fig, axes = plt.subplots(4, 3, figsize=(6, 5)) axes[0, 0].bar(date_hm, ds_hm["prec"].isel(x=x1, y=y1).values, width=0.1, color="blue", align="edge", edgecolor="blue") @@ -162,6 +303,9 @@ axes[1, 0].plot(date_hm, ds_hm["evap_soil"].isel(x=x1, y=y1).values, "--", color="#c2e699", lw=0.8) axes[1, 0].set_xlim(date_hm[0], date_hm[-1]) axes[1, 0].set_ylabel(r"ET [mm/day]") +axes[1, 0].set_ylim( + 0, +) axes[1, 0].xaxis.set_major_formatter(mdates.DateFormatter("%y-%m")) axes[1, 0].tick_params(axis="x", labelrotation=60) @@ -181,6 +325,9 @@ axes[3, 0].axvline(date_hm[t_wet], color="red", alpha=0.5) axes[3, 0].plot(date_hm, ds_hm["q_ss"].isel(x=x1, y=y1).values, "-", color="grey", lw=1) axes[3, 0].set_xlim(date_hm[0], date_hm[-1]) +axes[3, 0].set_ylim( + 0, +) axes[3, 0].set_ylabel(r"PERC [mm/day]") axes[3, 0].xaxis.set_major_formatter(mdates.DateFormatter("%y-%m")) axes[3, 0].set_xlabel(r"Time [year-month]") @@ -227,11 +374,14 @@ edgecolor=None, alpha=0.2, ) -axes[1, 2].plot(date_tm, ds_tm["ttavg_transp"].isel(x=x1, y=y1).values, "-", color="#31a354", lw=1) +axes[1, 2].plot(date_tm, ds_tm["tt50_transp"].isel(x=x1, y=y1).values, "-", color="#31a354", lw=1) axes[1, 2].set_xlim(date_tm[0], date_tm[-1]) axes[1, 2].set_ylabel(r"age [days]") axes[1, 2].xaxis.set_major_formatter(mdates.DateFormatter("%y-%m")) axes[1, 2].tick_params(axis="x", labelrotation=60) +axes[1, 2].set_ylim( + 0, +) axes[2, 2].axvline(date_hm[t_dry], color="red", alpha=0.5) axes[2, 2].axvline(date_hm[t_wetdry], color="red", alpha=0.5) @@ -245,11 +395,14 @@ edgecolor=None, alpha=0.2, ) -axes[2, 2].plot(date_tm, ds_tm["rtavg_s"].isel(x=x1, y=y1).values, "-", color="brown", lw=1) +axes[2, 2].plot(date_tm, ds_tm["rt50_s"].isel(x=x1, y=y1).values, "-", color="brown", lw=1) axes[2, 2].set_xlim(date_tm[0], date_tm[-1]) axes[2, 2].set_ylabel(r"age [days]") axes[2, 2].xaxis.set_major_formatter(mdates.DateFormatter("%y-%m")) axes[2, 2].tick_params(axis="x", labelrotation=60) +axes[2, 2].set_ylim( + 0, +) 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) @@ -263,12 +416,13 @@ edgecolor=None, alpha=0.2, ) -axes[3, 2].plot(date_tm, ds_tm["ttavg_q_ss"].isel(x=x1, y=y1).values, "-", color="grey", lw=1) +axes[3, 2].plot(date_tm, ds_tm["tt50_q_ss"].isel(x=x1, y=y1).values, "-", 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]") axes[3, 2].xaxis.set_major_formatter(mdates.DateFormatter("%y-%m")) axes[3, 2].tick_params(axis="x", labelrotation=60) +axes[3, 2].set_ylim(0, 500) fig.subplots_adjust(left=0.1, bottom=0.13, top=0.95, right=0.98, hspace=0.6, wspace=0.42) file = base_path_figs / "ts_single_grid_cell.png" @@ -385,15 +539,15 @@ fig, axes = plt.subplots(4, 3, figsize=(6, 5)) axes[0, 0].hist( - ds_hm["evap_soil"].isel(Time=t_wetdry).values.flatten(), - color="#6baed6", + ds_hm["evap_soil"].isel(Time=t_dry).values.flatten(), + color="#fed976", bins=10, range=(0, 1), align="mid", alpha=0.5, ) axes[0, 0].hist( - ds_hm["evap_soil"].isel(Time=t_wetdry).values.flatten(), + ds_hm["evap_soil"].isel(Time=t_wet).values.flatten(), color="#08306b", bins=10, range=(0, 1), @@ -404,7 +558,7 @@ axes[0, 0].set_ylabel("# grid cells") axes[0, 0].text( 0.95, - 1.12, + 1.14, "(a)", fontsize=9, horizontalalignment="center", @@ -413,7 +567,7 @@ ) axes[1, 0].hist( - ds_hm["transp"].isel(Time=t_wetdry).values.flatten(), color="#6baed6", bins=40, range=(0, 4), align="mid", alpha=0.5 + ds_hm["transp"].isel(Time=t_dry).values.flatten(), color="#fed976", bins=40, range=(0, 4), align="mid", alpha=0.5 ) axes[1, 0].hist( ds_hm["transp"].isel(Time=t_wet).values.flatten(), color="#08306b", bins=40, range=(0, 4), align="mid", alpha=0.5 @@ -422,7 +576,7 @@ axes[1, 0].set_ylabel("# grid cells") axes[1, 0].text( 0.95, - 1.12, + 1.14, "(b)", fontsize=9, horizontalalignment="center", @@ -431,21 +585,26 @@ ) axes[2, 0].hist( - ds_hm["theta"].isel(Time=t_wetdry).values.flatten(), - color="#6baed6", - bins=20, - range=(0.2, 0.4), + ds_hm["theta"].isel(Time=t_dry).values.flatten(), + color="#fed976", + bins=30, + range=(0.15, 0.45), align="mid", alpha=0.5, ) axes[2, 0].hist( - ds_hm["theta"].isel(Time=t_wet).values.flatten(), color="#08306b", bins=20, range=(0.2, 0.4), align="mid", alpha=0.5 + ds_hm["theta"].isel(Time=t_wet).values.flatten(), + color="#08306b", + bins=30, + range=(0.15, 0.45), + align="mid", + alpha=0.5, ) axes[2, 0].set_xlabel(r"$\theta$ [-]") axes[2, 0].set_ylabel("# grid cells") axes[2, 0].text( 0.95, - 1.12, + 1.14, "(c)", fontsize=9, horizontalalignment="center", @@ -454,7 +613,7 @@ ) axes[3, 0].hist( - ds_hm["q_ss"].isel(Time=t_wetdry).values.flatten(), color="#6baed6", bins=30, range=(0, 15), align="mid", alpha=0.5 + ds_hm["q_ss"].isel(Time=t_dry).values.flatten(), color="#fed976", bins=30, range=(0, 15), align="mid", alpha=0.5 ) axes[3, 0].hist( ds_hm["q_ss"].isel(Time=t_wet).values.flatten(), color="#08306b", bins=30, range=(0, 15), align="mid", alpha=0.5 @@ -463,7 +622,7 @@ axes[3, 0].set_ylabel("# grid cells") axes[3, 0].text( 0.95, - 1.12, + 1.14, "(d)", fontsize=9, horizontalalignment="center", @@ -472,8 +631,8 @@ ) axes[0, 1].hist( - ds_tm["C_iso_evap_soil"].isel(Time=t_wetdry).values.flatten(), - color="#6baed6", + ds_tm["C_iso_evap_soil"].isel(Time=t_dry).values.flatten(), + color="#fed976", bins=24, range=(-12, -6), align="mid", @@ -490,7 +649,7 @@ axes[0, 1].set_xlabel(r"$\delta^{18}$$O_{EVAP_{soil}}$ [‰]") axes[0, 1].text( 0.95, - 1.12, + 1.14, "(e)", fontsize=9, horizontalalignment="center", @@ -499,8 +658,8 @@ ) axes[1, 1].hist( - ds_tm["C_iso_transp"].isel(Time=t_wetdry).values.flatten(), - color="#6baed6", + ds_tm["C_iso_transp"].isel(Time=t_dry).values.flatten(), + color="#fed976", bins=24, range=(-12, -6), align="mid", @@ -517,7 +676,7 @@ axes[1, 1].set_xlabel(r"$\delta^{18}$$O_{TRANSP}$ [‰]") axes[1, 1].text( 0.95, - 1.12, + 1.14, "(f)", fontsize=9, horizontalalignment="center", @@ -526,8 +685,8 @@ ) axes[2, 1].hist( - ds_tm["C_iso_s"].isel(Time=t_wetdry).values.flatten(), - color="#6baed6", + ds_tm["C_iso_s"].isel(Time=t_dry).values.flatten(), + color="#fed976", bins=24, range=(-12, -6), align="mid", @@ -544,7 +703,7 @@ axes[2, 1].set_xlabel(r"$\delta^{18}$$O_{\theta}$ [‰]") axes[2, 1].text( 0.95, - 1.12, + 1.14, "(g)", fontsize=9, horizontalalignment="center", @@ -553,8 +712,8 @@ ) axes[3, 1].hist( - ds_tm["C_iso_q_ss"].isel(Time=t_wetdry).values.flatten(), - color="#6baed6", + ds_tm["C_iso_q_ss"].isel(Time=t_dry).values.flatten(), + color="#fed976", bins=24, range=(-12, -6), align="mid", @@ -571,7 +730,7 @@ axes[3, 1].set_xlabel(r"$\delta^{18}$$O_{PERC}$ [‰]") axes[3, 1].text( 0.95, - 1.12, + 1.14, "(h)", fontsize=9, horizontalalignment="center", @@ -582,8 +741,8 @@ axes[0, 2].set_axis_off() axes[1, 2].hist( - ds_tm["ttavg_transp"].isel(Time=t_wetdry).values.flatten(), - color="#6baed6", + ds_tm["ttavg_transp"].isel(Time=t_dry).values.flatten(), + color="#fed976", bins=50, range=(0, 600), align="mid", @@ -600,10 +759,10 @@ label=r"dry condtions", ) axes[1, 2].set_xlabel(r"$\overline{TT_{TRANSP}}$ [days]") -axes[1, 2].legend(frameon=False, loc="upper left", bbox_to_anchor=(0.0, 1.6)) +axes[1, 2].legend(frameon=False, loc="upper left", bbox_to_anchor=(0.0, 1.8)) axes[1, 2].text( 0.95, - 1.12, + 1.14, "(i)", fontsize=9, horizontalalignment="center", @@ -612,8 +771,8 @@ ) axes[2, 2].hist( - ds_tm["rtavg_s"].isel(Time=t_wetdry).values.flatten(), - color="#6baed6", + ds_tm["rtavg_s"].isel(Time=t_dry).values.flatten(), + color="#fed976", bins=50, range=(0, 600), align="mid", @@ -625,7 +784,7 @@ axes[2, 2].set_xlabel(r"$\overline{RT_{\theta}}$ [days]") axes[2, 2].text( 0.95, - 1.12, + 1.14, "(j)", fontsize=9, horizontalalignment="center", @@ -634,8 +793,8 @@ ) axes[3, 2].hist( - ds_tm["ttavg_q_ss"].isel(Time=t_wetdry).values.flatten(), - color="#6baed6", + ds_tm["ttavg_q_ss"].isel(Time=t_dry).values.flatten(), + color="#fed976", bins=50, range=(0, 600), align="mid", @@ -652,7 +811,7 @@ axes[3, 2].set_xlabel(r"$\overline{TT_{PERC}}$ [days]") axes[3, 2].text( 0.95, - 1.12, + 1.14, "(k)", fontsize=9, horizontalalignment="center", @@ -672,7 +831,7 @@ ds_hm["evap_soil"].isel(Time=t_wet).values.flatten(), color="#08306b", bins=50, - range=(0, 2), + range=(0, 2.1), histtype="step", cumulative=True, lw=2.0, @@ -682,7 +841,7 @@ ds_hm["evap_soil"].isel(Time=t_drywet).values.flatten(), color="#6baed6", bins=50, - range=(0, 2), + range=(0, 2.1), histtype="step", cumulative=True, lw=1.5, @@ -691,13 +850,13 @@ ds_hm["evap_soil"].isel(Time=t_wetdry).values.flatten(), color="#feb24c", bins=50, - range=(0, 2), + range=(0, 2.1), histtype="step", cumulative=True, lw=1.0, ls="-.", ) -axes[0, 0].set_xlim(0, 2) +axes[0, 0].set_xlim(0, 2.1) axes[0, 0].set_xlabel(r"$EVAP_{soil}$ [mm/day]") axes[0, 0].set_ylabel("# grid cells") axes[0, 0].text( @@ -813,8 +972,8 @@ axes[3, 0].hist( ds_hm["q_ss"].isel(Time=t_wet).values.flatten(), color="#08306b", - bins=50, - range=(0, 10), + bins=100, + range=(0, 100), histtype="step", cumulative=True, lw=2.0, @@ -822,8 +981,8 @@ axes[3, 0].hist( ds_hm["q_ss"].isel(Time=t_drywet).values.flatten(), color="#6baed6", - bins=50, - range=(0, 10), + bins=100, + range=(0, 100), histtype="step", cumulative=True, lw=1.5, @@ -831,8 +990,8 @@ axes[3, 0].hist( ds_hm["q_ss"].isel(Time=t_wetdry).values.flatten(), color="#feb24c", - bins=50, - range=(0, 10), + bins=100, + range=(0, 100), histtype="step", cumulative=True, lw=1.0, @@ -841,14 +1000,14 @@ axes[3, 0].hist( ds_hm["q_ss"].isel(Time=t_dry).values.flatten(), color="#fed976", - bins=50, - range=(0, 10), + bins=100, + range=(0, 100), histtype="step", cumulative=True, lw=0.5, ls="--", ) -axes[3, 0].set_xlim(0, 10) +axes[3, 0].set_xlim(0, 100) axes[3, 0].set_xlabel(r"$PERC$ [mm/day]") axes[3, 0].set_ylabel("# grid cells") axes[3, 0].text( @@ -867,7 +1026,7 @@ ds_tm["C_iso_evap_soil"].isel(Time=t_wet).values.flatten(), color="#08306b", bins=50, - range=(-15, -5), + range=(-15, -4.9), histtype="step", cumulative=True, lw=2.0, @@ -877,7 +1036,7 @@ ds_tm["C_iso_evap_soil"].isel(Time=t_drywet).values.flatten(), color="#6baed6", bins=50, - range=(-15, -5), + range=(-15, -4.9), histtype="step", cumulative=True, lw=1.5, @@ -886,7 +1045,7 @@ ds_tm["C_iso_evap_soil"].isel(Time=t_wetdry).values.flatten(), color="#feb24c", bins=50, - range=(-15, -5), + range=(-15, -4.9), histtype="step", cumulative=True, lw=1.0, @@ -907,7 +1066,7 @@ ds_tm["C_iso_transp"].isel(Time=t_wet).values.flatten(), color="#08306b", bins=50, - range=(-15, -5), + range=(-15, -4.9), histtype="step", cumulative=True, lw=2.0, @@ -916,7 +1075,7 @@ ds_tm["C_iso_transp"].isel(Time=t_wetdry).values.flatten(), color="#feb24c", bins=50, - range=(-15, -5), + range=(-15, -4.9), histtype="step", cumulative=True, lw=1.5, @@ -926,7 +1085,7 @@ ds_tm["C_iso_transp"].isel(Time=t_drywet).values.flatten(), color="#6baed6", bins=50, - range=(-15, -5), + range=(-15, -4.9), histtype="step", cumulative=True, lw=1.0, @@ -935,7 +1094,7 @@ ds_tm["C_iso_transp"].isel(Time=t_dry).values.flatten(), color="#fed976", bins=50, - range=(-15, -5), + range=(-15, -4.9), histtype="step", cumulative=True, lw=0.5, @@ -956,7 +1115,7 @@ ds_tm["C_iso_s"].isel(Time=t_wet).values.flatten(), color="#08306b", bins=50, - range=(-15, -5), + range=(-15, -4.9), histtype="step", cumulative=True, lw=2.0, @@ -965,7 +1124,7 @@ ds_tm["C_iso_s"].isel(Time=t_drywet).values.flatten(), color="#6baed6", bins=50, - range=(-15, -5), + range=(-15, -4.9), histtype="step", cumulative=True, lw=1.5, @@ -974,7 +1133,7 @@ ds_tm["C_iso_s"].isel(Time=t_wetdry).values.flatten(), color="#feb24c", bins=50, - range=(-15, -5), + range=(-15, -4.9), histtype="step", cumulative=True, lw=1.0, @@ -984,7 +1143,7 @@ ds_tm["C_iso_s"].isel(Time=t_dry).values.flatten(), color="#fed976", bins=50, - range=(-15, -5), + range=(-15, -4.9), histtype="step", cumulative=True, lw=0.5, @@ -1005,7 +1164,7 @@ ds_tm["C_iso_q_ss"].isel(Time=t_wet).values.flatten(), color="#08306b", bins=50, - range=(-15, -5), + range=(-15, -4.9), histtype="step", cumulative=True, lw=2.0, @@ -1014,7 +1173,7 @@ ds_tm["C_iso_q_ss"].isel(Time=t_drywet).values.flatten(), color="#6baed6", bins=50, - range=(-15, -5), + range=(-15, -4.9), histtype="step", cumulative=True, lw=1.5, @@ -1023,7 +1182,7 @@ ds_tm["C_iso_q_ss"].isel(Time=t_wetdry).values.flatten(), color="#feb24c", bins=50, - range=(-15, -5), + range=(-15, -4.9), histtype="step", cumulative=True, lw=1.0, @@ -1033,7 +1192,7 @@ ds_tm["C_iso_q_ss"].isel(Time=t_dry).values.flatten(), color="#fed976", bins=50, - range=(-15, -5), + range=(-15, -4.9), histtype="step", cumulative=True, lw=0.5, @@ -1091,7 +1250,7 @@ lw=0.5, ls="--", ) -axes[1, 2].set_xlim(0, 1000) +axes[1, 2].set_xlim(0, 400) axes[1, 2].set_xlabel(r"$\overline{TT_{TRANSP}}$ [days]") axes[1, 2].text( 0.95, @@ -1107,7 +1266,7 @@ axes[1, 2].plot([], [], color="#6baed6", label="transition to wet condtions", lw=1.5) axes[1, 2].plot([], [], color="#08306b", label="wet condtions", lw=2.0) lines, labels = axes[1, 2].get_legend_handles_labels() -fig.legend(lines, labels, loc="upper right", frameon=False, bbox_to_anchor=(0.98, 1.005)) +fig.legend(lines, labels, loc="upper right", frameon=False, bbox_to_anchor=(1.01, 0.96)) axes[2, 2].hist( ds_tm["rtavg_s"].isel(Time=t_wet).values.flatten(), @@ -1147,7 +1306,7 @@ lw=0.5, ls="--", ) -axes[2, 2].set_xlim(0, 1000) +axes[2, 2].set_xlim(0, 400) axes[2, 2].set_xlabel(r"$\overline{RT}$ [days]") axes[2, 2].text( 0.95, @@ -1196,7 +1355,7 @@ lw=0.5, ls="--", ) -axes[3, 2].set_xlim(0, 1000) +axes[3, 2].set_xlim(0, 400) axes[3, 2].set_xlabel(r"$\overline{TT_{PERC}}$ [days]") axes[3, 2].text( 0.95, @@ -1208,12 +1367,12 @@ transform=axes[3, 2].transAxes, ) -inset = fig.add_axes([0.225, 0.15, 0.11, 0.08]) +inset = fig.add_axes([0.225, 0.14, 0.11, 0.075]) inset.hist( ds_hm["q_ss"].isel(Time=t_wet).values.flatten(), color="#08306b", - bins=40, - range=(0, 0.4), + bins=1000, + range=(0, 100), histtype="step", cumulative=True, lw=2.0, @@ -1221,8 +1380,8 @@ inset.hist( ds_hm["q_ss"].isel(Time=t_drywet).values.flatten(), color="#6baed6", - bins=40, - range=(0, 0.4), + bins=1000, + range=(0, 100), histtype="step", cumulative=True, lw=1.5, @@ -1230,8 +1389,8 @@ inset.hist( ds_hm["q_ss"].isel(Time=t_wetdry).values.flatten(), color="#feb24c", - bins=40, - range=(0, 0.4), + bins=1000, + range=(0, 100), histtype="step", cumulative=True, lw=1.0, @@ -1240,14 +1399,14 @@ inset.hist( ds_hm["q_ss"].isel(Time=t_dry).values.flatten(), color="#fed976", - bins=40, - range=(0, 0.4), + bins=1000, + range=(0, 100), histtype="step", cumulative=True, lw=0.5, ls="--", ) -inset.set_xlim(0, 0.4) +inset.set_xlim(0, 0.2) fig.subplots_adjust(left=0.1, bottom=0.1, top=0.95, right=0.98, hspace=0.6, wspace=0.25) file = base_path_figs / "cumulated_dist_states_dry_normal_wet.png" @@ -1303,7 +1462,7 @@ # # make GIF for Online-Documentation -# for t in range(1, 1000): +# for t in range(1, 1097): # fig, axes = plt.subplots(3, 3, figsize=(6, 6)) # axes[0, 0].set_axis_off() @@ -1332,16 +1491,16 @@ # cb1.ax.set_xticklabels(["0", "2.5", ">5"]) # cb1.set_label("TRANSP [mm/day]") -# axes[1, 1].imshow(onp.where(ds_hm["theta"].isel(Time=t).values <= -9999, onp.nan, ds_hm["theta"].isel(Time=t).values), extent=(0, 80*25, 0, 53*25), cmap="viridis_r", vmin=0.15, vmax=0.3) +# axes[1, 1].imshow(onp.where(ds_hm["theta"].isel(Time=t).values <= -9999, onp.nan, ds_hm["theta"].isel(Time=t).values), extent=(0, 80*25, 0, 53*25), cmap="viridis_r", vmin=0.1, vmax=0.3) # axes[1, 1].grid(zorder=0) # axes[1, 1].set_xlabel("") # axes[1, 1].set_ylabel("") # axes[1, 1].set_yticklabels([]) # cmap = copy.copy(mpl.colormaps.get_cmap("viridis_r")) -# norm = mpl.colors.Normalize(vmin=0.15, vmax=0.3) +# norm = mpl.colors.Normalize(vmin=0.1, vmax=0.4) # axl1 = fig.add_axes([0.38, 0.66, 0.2, 0.01]) -# cb1 = mpl.colorbar.ColorbarBase(axl1, cmap=cmap, norm=norm, orientation="horizontal", ticks=[0.15, 0.2, 0.25, 0.3]) -# cb1.ax.set_xticklabels(["<0.15", "0.2", "0.25", ">0.3"]) +# cb1 = mpl.colorbar.ColorbarBase(axl1, cmap=cmap, norm=norm, orientation="horizontal", ticks=[0.1, 0.2, 0.3, 0.4]) +# cb1.ax.set_xticklabels(["<0.1", "0.2", "0.3", ">0.4"]) # cb1.set_label(r"$\theta$ [-]") # axes[1, 2].imshow(onp.where(ds_hm["q_ss"].isel(Time=t).values <= -9999, onp.nan, ds_hm["q_ss"].isel(Time=t).values), extent=(0, 80*25, 0, 53*25), cmap="viridis_r", vmin=0, vmax=20) @@ -1356,39 +1515,39 @@ # cb1.ax.set_xticklabels(["0", "10", ">20"]) # cb1.set_label("PERC [mm/day]") -# axes[2, 0].imshow(ds_tm["ttavg_transp"].isel(Time=t).values, extent=(0, 80*25, 0, 53*25), cmap="viridis_r", vmin=1, vmax=1000) +# axes[2, 0].imshow(ds_tm["ttavg_transp"].isel(Time=t).values, extent=(0, 80*25, 0, 53*25), cmap="viridis_r", vmin=1, vmax=200) # axes[2, 0].grid(zorder=0) # axes[2, 0].set_xlabel("[m]") # axes[2, 0].set_ylabel("[m]") # cmap = copy.copy(mpl.colormaps.get_cmap("viridis")) -# norm = mpl.colors.Normalize(vmin=1, vmax=1000) +# norm = mpl.colors.Normalize(vmin=1, vmax=200) # axl2 = fig.add_axes([0.1, 0.38, 0.2, 0.01]) -# cb2 = mpl.colorbar.ColorbarBase(axl2, cmap=cmap, norm=norm, orientation="horizontal", ticks=[1, 500, 1000]) -# cb2.ax.set_xticklabels(["1", "500", ">1000"]) +# cb2 = mpl.colorbar.ColorbarBase(axl2, cmap=cmap, norm=norm, orientation="horizontal", ticks=[1, 100, 200]) +# cb2.ax.set_xticklabels(["1", "100", ">200"]) # cb2.ax.invert_yaxis() # cb2.set_label(r"$\overline{TT}_{TRANSP}$ [days]") -# axes[2, 1].imshow(ds_tm["rtavg_s"].isel(Time=t).values, extent=(0, 80*25, 0, 53*25), cmap="viridis_r", vmin=1, vmax=1000) +# axes[2, 1].imshow(ds_tm["rtavg_s"].isel(Time=t).values, extent=(0, 80*25, 0, 53*25), cmap="viridis_r", vmin=1, vmax=200) # axes[2, 1].grid(zorder=0) # axes[2, 1].set_xlabel("[m]") # axes[2, 1].set_yticklabels([]) # cmap = copy.copy(mpl.colormaps.get_cmap("viridis")) -# norm = mpl.colors.Normalize(vmin=1, vmax=1000) +# norm = mpl.colors.Normalize(vmin=1, vmax=200) # axl2 = fig.add_axes([0.38, 0.38, 0.2, 0.01]) -# cb2 = mpl.colorbar.ColorbarBase(axl2, cmap=cmap, norm=norm, orientation="horizontal", ticks=[1, 500, 1000]) -# cb2.ax.set_xticklabels(["1", "500", ">1000"]) +# cb2 = mpl.colorbar.ColorbarBase(axl2, cmap=cmap, norm=norm, orientation="horizontal", ticks=[1, 100, 200]) +# cb2.ax.set_xticklabels(["1", "100", ">200"]) # cb2.ax.invert_yaxis() -# cb2.set_label(r"$\overline{RT}_{SOIL}$ [days]") +# cb2.set_label(r"$\overline{RT}_{\theta}$ [days]") -# axes[2, 2].imshow(ds_tm["ttavg_q_ss"].isel(Time=t).values, extent=(0, 80*25, 0, 53*25), cmap="viridis_r", vmin=1, vmax=1000) +# axes[2, 2].imshow(ds_tm["ttavg_q_ss"].isel(Time=t).values, extent=(0, 80*25, 0, 53*25), cmap="viridis_r", vmin=1, vmax=200) # axes[2, 2].grid(zorder=0) # axes[2, 2].set_xlabel("[m]") # axes[2, 2].set_yticklabels([]) # cmap = copy.copy(mpl.colormaps.get_cmap("viridis")) -# norm = mpl.colors.Normalize(vmin=1, vmax=1000) +# norm = mpl.colors.Normalize(vmin=1, vmax=200) # axl2 = fig.add_axes([0.65, 0.38, 0.2, 0.01]) -# cb2 = mpl.colorbar.ColorbarBase(axl2, cmap=cmap, norm=norm, orientation="horizontal", ticks=[1, 500, 1000]) -# cb2.ax.set_xticklabels(["1", "500", ">1000"]) +# cb2 = mpl.colorbar.ColorbarBase(axl2, cmap=cmap, norm=norm, orientation="horizontal", ticks=[1, 100, 200]) +# cb2.ax.set_xticklabels(["1", "100", ">200"]) # cb2.ax.invert_yaxis() # cb2.set_label(r"$\overline{TT}_{PERC}$ [days]") # fig.subplots_adjust(left=0.1, bottom=0.1, top=0.94, right=0.85, wspace=0.35, hspace=0.1) @@ -1397,12 +1556,12 @@ # plt.close("all") # images_data = [] -# #load 10 images -# for t in range(1, 1000): +# #load images +# for t in range(1, 1097): # data = imageio.v2.imread(base_path_figs / f"fluxes_theta_and_tt_rt_{t}.png") # images_data.append(data) # file = base_path_figs / "fluxes_theta_and_tt_rt.gif" -# imageio.mimwrite(file, images_data, format='.gif', fps=12) +# imageio.mimwrite(file, images_data, format='.gif', fps=14) -# plt.close("all") +plt.close("all") diff --git a/examples/catchment_scale/eberbaechle/svat_distributed/svat.py b/examples/catchment_scale/eberbaechle/svat_distributed/svat.py index db2ad1406..5ed38f994 100644 --- a/examples/catchment_scale/eberbaechle/svat_distributed/svat.py +++ b/examples/catchment_scale/eberbaechle/svat_distributed/svat.py @@ -36,7 +36,7 @@ def _read_var_from_nc(self, var, path_dir, file): if var in ["lu_id"]: vals = npx.array(var_obj) vals1 = npx.where(vals <= -9999, 999, vals) - elif var in ["dmph", "dmpv", "lmpv", "z_soil", "sealing"]: + elif var in ["dmph", "dmpv", "lmpv", "z_soil", "sealing", "prec_weight", "pet_weight", "ta_offset"]: vals = npx.array(var_obj) vals1 = npx.where(vals <= -9999, 0, vals) else: @@ -130,7 +130,9 @@ def set_parameters_setup(self, state): vs = state.variables # land use ID (see README for description) - vs.lu_id = update(vs.lu_id, at[2:-2, 2:-2], self._read_var_from_nc("lu_id", self._base_path, "parameters.nc")) + vs.lu_id = update( + vs.lu_id, at[2:-2, 2:-2], self._read_var_from_nc("lu_id", self._base_path, "parameters.nc") + ) # degree of surface sealing (-) vs.sealing = update( vs.sealing, at[2:-2, 2:-2], self._read_var_from_nc("sealing", self._base_path, "parameters.nc") @@ -269,21 +271,24 @@ def set_forcing(self, state): 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.prec_weight[2:-2, 2:-2, npx.newaxis], + ] + * vs.prec_weight[2:-2, 2:-2, npx.newaxis], ) 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.ta_offset[2:-2, 2:-2, npx.newaxis], + ] + + vs.ta_offset[2:-2, 2:-2, npx.newaxis], ) 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.pet_weight[2:-2, 2:-2, npx.newaxis], + ] + * vs.pet_weight[2:-2, 2:-2, npx.newaxis], ) vs.itt_forc = vs.itt_forc + 6 * 24 diff --git a/roger/core/adaptive_time_stepping.py b/roger/core/adaptive_time_stepping.py index c1ca5d388..7ee75ab1c 100644 --- a/roger/core/adaptive_time_stepping.py +++ b/roger/core/adaptive_time_stepping.py @@ -15,13 +15,49 @@ def adaptive_time_stepping_kernel(state): vs = state.variables settings = state.settings - cond0 = npx.array((vs.prec_day[2:-2, 2:-2, :] <= 0).all() & (vs.swe[2:-2, 2:-2, vs.tau] <= 0).all() & (vs.swe_top[2:-2, 2:-2, vs.tau] <= 0).all() & (vs.ta_day[2:-2, 2:-2, :] > settings.ta_fm).all()) - cond00 = npx.array(((vs.prec_day[2:-2, 2:-2, :] > 0) & (vs.ta_day[2:-2, 2:-2, :] <= settings.ta_fm)).any() | ((vs.prec_day[2:-2, 2:-2, :] <= 0) & (vs.ta_day[2:-2, 2:-2, :] <= settings.ta_fm)).all()) - cond1 = npx.array((vs.prec_day[2:-2, 2:-2, :] > settings.hpi).any() & (vs.prec_day[2:-2, 2:-2, :] > 0).any() & (vs.ta_day[2:-2, 2:-2, :] > settings.ta_fm).any()) - cond2 = npx.array((vs.prec_day[2:-2, 2:-2, :] <= settings.hpi).all() & (vs.prec_day[2:-2, 2:-2, :] > 0).any() & (vs.ta_day[2:-2, 2:-2, :] > settings.ta_fm).any()) - cond3 = npx.array((vs.prec_day[2:-2, 2:-2, :] > settings.hpi).any() & (vs.prec_day[2:-2, 2:-2, :] > 0).any() & (((vs.swe[2:-2, 2:-2, vs.tau] > 0).any() | (vs.swe_top[2:-2, 2:-2, vs.tau] > 0).any()) & (vs.ta_day[2:-2, 2:-2, :] > settings.ta_fm).any())) - cond4 = npx.array((vs.prec_day[2:-2, 2:-2, :] <= settings.hpi).all() & (vs.prec_day[2:-2, 2:-2, :] > 0).any() & (((vs.swe[2:-2, 2:-2, vs.tau] > 0).any() | (vs.swe_top[2:-2, 2:-2, vs.tau] > 0).any()) & (vs.ta_day[2:-2, 2:-2, :] > settings.ta_fm).any())) - cond5 = npx.array((vs.prec_day[2:-2, 2:-2, :] <= 0).all() & (((vs.swe[2:-2, 2:-2, vs.tau] > 0).any() | (vs.swe_top[2:-2, 2:-2, vs.tau] > 0).any()) & (vs.ta_day[2:-2, 2:-2, :] > settings.ta_fm).any())) + cond0 = npx.array( + (vs.prec_day[2:-2, 2:-2, :] <= 0).all() + & (vs.swe[2:-2, 2:-2, vs.tau] <= 0).all() + & (vs.swe_top[2:-2, 2:-2, vs.tau] <= 0).all() + & (vs.ta_day[2:-2, 2:-2, :] > settings.ta_fm).all() + ) + cond00 = npx.array( + ((vs.prec_day[2:-2, 2:-2, :] > 0) & (vs.ta_day[2:-2, 2:-2, :] <= settings.ta_fm)).any() + | ((vs.prec_day[2:-2, 2:-2, :] <= 0) & (vs.ta_day[2:-2, 2:-2, :] <= settings.ta_fm)).all() + ) + cond1 = npx.array( + (vs.prec_day[2:-2, 2:-2, :] > settings.hpi).any() + & (vs.prec_day[2:-2, 2:-2, :] > 0).any() + & (vs.ta_day[2:-2, 2:-2, :] > settings.ta_fm).any() + ) + cond2 = npx.array( + (vs.prec_day[2:-2, 2:-2, :] <= settings.hpi).all() + & (vs.prec_day[2:-2, 2:-2, :] > 0).any() + & (vs.ta_day[2:-2, 2:-2, :] > settings.ta_fm).any() + ) + cond3 = npx.array( + (vs.prec_day[2:-2, 2:-2, :] > settings.hpi).any() + & (vs.prec_day[2:-2, 2:-2, :] > 0).any() + & ( + ((vs.swe[2:-2, 2:-2, vs.tau] > 0).any() | (vs.swe_top[2:-2, 2:-2, vs.tau] > 0).any()) + & (vs.ta_day[2:-2, 2:-2, :] > settings.ta_fm).any() + ) + ) + cond4 = npx.array( + (vs.prec_day[2:-2, 2:-2, :] <= settings.hpi).all() + & (vs.prec_day[2:-2, 2:-2, :] > 0).any() + & ( + ((vs.swe[2:-2, 2:-2, vs.tau] > 0).any() | (vs.swe_top[2:-2, 2:-2, vs.tau] > 0).any()) + & (vs.ta_day[2:-2, 2:-2, :] > settings.ta_fm).any() + ) + ) + cond5 = npx.array( + (vs.prec_day[2:-2, 2:-2, :] <= 0).all() + & ( + ((vs.swe[2:-2, 2:-2, vs.tau] > 0).any() | (vs.swe_top[2:-2, 2:-2, vs.tau] > 0).any()) + & (vs.ta_day[2:-2, 2:-2, :] > settings.ta_fm).any() + ) + ) cond_time = npx.array((vs.time % (24 * 60 * 60) == 0)) cond0_arr = allocate(state.dimensions, ("x", "y")) @@ -33,31 +69,38 @@ def adaptive_time_stepping_kernel(state): cond5_arr = allocate(state.dimensions, ("x", "y")) cond0_arr = update( cond0_arr, - at[2:-2, 2:-2], npx.where(cond0[npx.newaxis, npx.newaxis], 1, 0), + at[2:-2, 2:-2], + npx.where(cond0[npx.newaxis, npx.newaxis], 1, 0), ) cond00_arr = update( cond00_arr, - at[2:-2, 2:-2], npx.where(cond00[npx.newaxis, npx.newaxis], 1, 0), + at[2:-2, 2:-2], + npx.where(cond00[npx.newaxis, npx.newaxis], 1, 0), ) cond1_arr = update( cond1_arr, - at[2:-2, 2:-2], npx.where(cond1[npx.newaxis, npx.newaxis], 1, 0), + at[2:-2, 2:-2], + npx.where(cond1[npx.newaxis, npx.newaxis], 1, 0), ) cond2_arr = update( cond2_arr, - at[2:-2, 2:-2], npx.where(cond2[npx.newaxis, npx.newaxis], 1, 0), + at[2:-2, 2:-2], + npx.where(cond2[npx.newaxis, npx.newaxis], 1, 0), ) cond3_arr = update( cond3_arr, - at[2:-2, 2:-2], npx.where(cond3[npx.newaxis, npx.newaxis], 1, 0), + at[2:-2, 2:-2], + npx.where(cond3[npx.newaxis, npx.newaxis], 1, 0), ) cond4_arr = update( cond4_arr, - at[2:-2, 2:-2], npx.where(cond4[npx.newaxis, npx.newaxis], 1, 0), + at[2:-2, 2:-2], + npx.where(cond4[npx.newaxis, npx.newaxis], 1, 0), ) cond5_arr = update( cond5_arr, - at[2:-2, 2:-2], npx.where(cond5[npx.newaxis, npx.newaxis], 1, 0), + at[2:-2, 2:-2], + npx.where(cond5[npx.newaxis, npx.newaxis], 1, 0), ) prec_daily, ta_daily, pet_daily = calc_prec_ta_pet_daily(state) @@ -67,11 +110,15 @@ def adaptive_time_stepping_kernel(state): # no event or snowfall - daily time steps vs.prec = update( vs.prec, - at[2:-2, 2:-2, vs.tau], npx.where((cond0_arr[2:-2, 2:-2] == 1) | (cond00_arr[2:-2, 2:-2] == 1), prec_daily, vs.prec[2:-2, 2:-2, vs.tau]), + at[2:-2, 2:-2, vs.tau], + npx.where( + (cond0_arr[2:-2, 2:-2] == 1) | (cond00_arr[2:-2, 2:-2] == 1), prec_daily, vs.prec[2:-2, 2:-2, vs.tau] + ), ) vs.ta = update( vs.ta, - at[2:-2, 2:-2, vs.tau], npx.where(((cond0_arr[2:-2, 2:-2] == 1) | (cond00_arr[2:-2, 2:-2] == 1)), ta_daily, vs.ta[2:-2, 2:-2, vs.tau]), + at[2:-2, 2:-2, vs.tau], + npx.where(((cond0_arr[2:-2, 2:-2] == 1) | (cond00_arr[2:-2, 2:-2] == 1)), ta_daily, vs.ta[2:-2, 2:-2, vs.tau]), ) vs.dt_secs = npx.where((cond0 | cond00), 24 * 60 * 60, vs.dt_secs) vs.dt_secs = npx.where(cond_time, 24 * 60 * 60, 60 * 60) @@ -79,27 +126,59 @@ def adaptive_time_stepping_kernel(state): # rainfall/snow melt event - hourly time steps vs.prec = update( vs.prec, - at[2:-2, 2:-2, vs.tau], npx.where(((cond2_arr[2:-2, 2:-2] == 1) | (cond4_arr[2:-2, 2:-2] == 1) | (cond5_arr[2:-2, 2:-2] == 1)) & ((cond1_arr[2:-2, 2:-2] == 0) & (cond3_arr[2:-2, 2:-2] == 0)), prec_hourly, vs.prec[2:-2, 2:-2, vs.tau]), + at[2:-2, 2:-2, vs.tau], + npx.where( + ((cond2_arr[2:-2, 2:-2] == 1) | (cond4_arr[2:-2, 2:-2] == 1) | (cond5_arr[2:-2, 2:-2] == 1)) + & ((cond1_arr[2:-2, 2:-2] == 0) & (cond3_arr[2:-2, 2:-2] == 0)), + prec_hourly, + vs.prec[2:-2, 2:-2, vs.tau], + ), ) vs.ta = update( vs.ta, - at[2:-2, 2:-2, vs.tau], npx.where(((cond2_arr[2:-2, 2:-2] == 1) | (cond4_arr[2:-2, 2:-2] == 1) | (cond5_arr[2:-2, 2:-2] == 1)) & ((cond1_arr[2:-2, 2:-2] == 0) & (cond3_arr[2:-2, 2:-2] == 0)), ta_hourly, vs.ta[2:-2, 2:-2, vs.tau]), + at[2:-2, 2:-2, vs.tau], + npx.where( + ((cond2_arr[2:-2, 2:-2] == 1) | (cond4_arr[2:-2, 2:-2] == 1) | (cond5_arr[2:-2, 2:-2] == 1)) + & ((cond1_arr[2:-2, 2:-2] == 0) & (cond3_arr[2:-2, 2:-2] == 0)), + ta_hourly, + vs.ta[2:-2, 2:-2, vs.tau], + ), ) vs.dt_secs = npx.where((cond2 | cond4 | cond5) & ~cond1 & ~cond3, 60 * 60, vs.dt_secs) # heavy rainfall event - 10 minutes time steps vs.prec = update( vs.prec, - at[2:-2, 2:-2, vs.tau], npx.where(((cond1_arr[2:-2, 2:-2] == 1) | (cond3_arr[2:-2, 2:-2] == 1)) & ((cond2_arr[2:-2, 2:-2] == 0) & (cond4_arr[2:-2, 2:-2] == 0) & (cond5_arr[2:-2, 2:-2] == 0)), prec_10mins, vs.prec[2:-2, 2:-2, vs.tau]), + at[2:-2, 2:-2, vs.tau], + npx.where( + ((cond1_arr[2:-2, 2:-2] == 1) | (cond3_arr[2:-2, 2:-2] == 1)) + & ((cond2_arr[2:-2, 2:-2] == 0) & (cond4_arr[2:-2, 2:-2] == 0) & (cond5_arr[2:-2, 2:-2] == 0)), + prec_10mins, + vs.prec[2:-2, 2:-2, vs.tau], + ), ) vs.ta = update( vs.ta, - at[2:-2, 2:-2, vs.tau], npx.where(((cond1_arr[2:-2, 2:-2] == 1) | (cond3_arr[2:-2, 2:-2] == 1)) & ((cond2_arr[2:-2, 2:-2] == 0) & (cond4_arr[2:-2, 2:-2] == 0) & (cond5_arr[2:-2, 2:-2] == 0)), ta_10mins, vs.ta[2:-2, 2:-2, vs.tau]), + at[2:-2, 2:-2, vs.tau], + npx.where( + ((cond1_arr[2:-2, 2:-2] == 1) | (cond3_arr[2:-2, 2:-2] == 1)) + & ((cond2_arr[2:-2, 2:-2] == 0) & (cond4_arr[2:-2, 2:-2] == 0) & (cond5_arr[2:-2, 2:-2] == 0)), + ta_10mins, + vs.ta[2:-2, 2:-2, vs.tau], + ), ) - vs.dt_secs = npx.where((cond1 | cond3) & ~ cond2 & ~cond4 & ~cond5, 10 * 60, vs.dt_secs) + vs.dt_secs = npx.where((cond1 | cond3) & ~cond2 & ~cond4 & ~cond5, 10 * 60, vs.dt_secs) # determine end of event - cond_event1 = ((vs.prec[2:-2, 2:-2] > 0).any() | ((vs.swe[2:-2, 2:-2, vs.tau] > 0) | (vs.swe_top[2:-2, 2:-2, vs.tau] > 0)).any() & (vs.ta[2:-2, 2:-2] > settings.ta_fm)).any() - cond_event2 = ((vs.prec[2:-2, 2:-2] <= 0) & (vs.ta[2:-2, 2:-2] > settings.ta_fm)).all() | ((vs.prec[2:-2, 2:-2] > 0) & (vs.ta[2:-2, 2:-2] <= settings.ta_fm)).all() | ((vs.swe[2:-2, 2:-2, vs.taum1] > 0).any() & (vs.swe[2:-2, 2:-2, vs.tau] <= 0).all()) + cond_event1 = ( + (vs.prec[2:-2, 2:-2] > 0).any() + | ((vs.swe[2:-2, 2:-2, vs.tau] > 0) | (vs.swe_top[2:-2, 2:-2, vs.tau] > 0)).any() + & (vs.ta[2:-2, 2:-2] > settings.ta_fm) + ).any() + cond_event2 = ( + ((vs.prec[2:-2, 2:-2] <= 0) & (vs.ta[2:-2, 2:-2] > settings.ta_fm)).all() + | ((vs.prec[2:-2, 2:-2] > 0) & (vs.ta[2:-2, 2:-2] <= settings.ta_fm)).all() + | ((vs.swe[2:-2, 2:-2, vs.taum1] > 0).any() & (vs.swe[2:-2, 2:-2, vs.tau] <= 0).all()) + ) vs.time_event0 = npx.where(cond_event1, 0, vs.time_event0) vs.time_event0 = npx.where(cond_event2, vs.time_event0 + vs.dt_secs, vs.time_event0) @@ -109,8 +188,14 @@ def adaptive_time_stepping_kernel(state): cond7 = npx.array((vs.time_event0 <= settings.end_event) & (vs.dt_secs == 60 * 60)) cond8 = npx.array((vs.time_event0 <= settings.end_event) & (vs.dt_secs == 24 * 60 * 60)) cond9 = npx.array((vs.time_event0 > settings.end_event) & (vs.time % (60 * 60) != 0) & (vs.dt_secs == 10 * 60)) - cond10 = npx.array((vs.time_event0 > settings.end_event) & (vs.time % (60 * 60) == 0) & ((vs.dt_secs == 10 * 60) | (vs.dt_secs == 60 * 60))) - cond11 = npx.array((vs.time_event0 > settings.end_event) & (vs.time % (24 * 60 * 60) == 0) & (vs.dt_secs == 24 * 60 * 60)) + cond10 = npx.array( + (vs.time_event0 > settings.end_event) + & (vs.time % (60 * 60) == 0) + & ((vs.dt_secs == 10 * 60) | (vs.dt_secs == 60 * 60)) + ) + cond11 = npx.array( + (vs.time_event0 > settings.end_event) & (vs.time % (24 * 60 * 60) == 0) & (vs.dt_secs == 24 * 60 * 60) + ) cond6_arr = allocate(state.dimensions, ("x", "y")) cond7_arr = allocate(state.dimensions, ("x", "y")) cond8_arr = allocate(state.dimensions, ("x", "y")) @@ -119,81 +204,98 @@ def adaptive_time_stepping_kernel(state): cond11_arr = allocate(state.dimensions, ("x", "y")) cond6_arr = update( cond6_arr, - at[2:-2, 2:-2], npx.where(cond6[npx.newaxis, npx.newaxis], 1, 0), + at[2:-2, 2:-2], + npx.where(cond6[npx.newaxis, npx.newaxis], 1, 0), ) cond7_arr = update( cond7_arr, - at[2:-2, 2:-2], npx.where(cond7[npx.newaxis, npx.newaxis], 1, 0), + at[2:-2, 2:-2], + npx.where(cond7[npx.newaxis, npx.newaxis], 1, 0), ) cond8_arr = update( cond8_arr, - at[2:-2, 2:-2], npx.where(cond8[npx.newaxis, npx.newaxis], 1, 0), + at[2:-2, 2:-2], + npx.where(cond8[npx.newaxis, npx.newaxis], 1, 0), ) cond9_arr = update( cond9_arr, - at[2:-2, 2:-2], npx.where(cond9[npx.newaxis, npx.newaxis], 1, 0), + at[2:-2, 2:-2], + npx.where(cond9[npx.newaxis, npx.newaxis], 1, 0), ) cond10_arr = update( cond10_arr, - at[2:-2, 2:-2], npx.where(cond10[npx.newaxis, npx.newaxis], 1, 0), + at[2:-2, 2:-2], + npx.where(cond10[npx.newaxis, npx.newaxis], 1, 0), ) cond11_arr = update( cond11_arr, - at[2:-2, 2:-2], npx.where(cond11[npx.newaxis, npx.newaxis], 1, 0), + at[2:-2, 2:-2], + npx.where(cond11[npx.newaxis, npx.newaxis], 1, 0), ) vs.pet = update( vs.pet, - at[2:-2, 2:-2], npx.where((cond6_arr[2:-2, 2:-2] == 1), pet_10mins, vs.pet[2:-2, 2:-2]), + at[2:-2, 2:-2], + npx.where((cond6_arr[2:-2, 2:-2] == 1), pet_10mins, vs.pet[2:-2, 2:-2]), ) vs.ta = update( vs.ta, - at[2:-2, 2:-2, vs.tau], npx.where((cond6_arr[2:-2, 2:-2] == 1), ta_10mins, vs.ta[2:-2, 2:-2, vs.tau]), + at[2:-2, 2:-2, vs.tau], + npx.where((cond6_arr[2:-2, 2:-2] == 1), ta_10mins, vs.ta[2:-2, 2:-2, vs.tau]), ) vs.event_id = update( vs.event_id, - at[vs.tau], npx.where(cond6, vs.event_id_counter, vs.event_id[vs.tau]), + at[vs.tau], + npx.where(cond6, vs.event_id_counter, vs.event_id[vs.tau]), ) vs.dt = npx.where(cond6, 1 / 6, vs.dt) vs.itt_day = npx.where(cond6, vs.itt_day + 1, vs.itt_day) vs.pet = update( vs.pet, - at[2:-2, 2:-2], npx.where((cond7_arr[2:-2, 2:-2] == 1), pet_hourly, vs.pet[2:-2, 2:-2]), + at[2:-2, 2:-2], + npx.where((cond7_arr[2:-2, 2:-2] == 1), pet_hourly, vs.pet[2:-2, 2:-2]), ) vs.ta = update( vs.ta, - at[2:-2, 2:-2, vs.tau], npx.where((cond7_arr[2:-2, 2:-2] == 1), ta_hourly, vs.ta[2:-2, 2:-2, vs.tau]), + at[2:-2, 2:-2, vs.tau], + npx.where((cond7_arr[2:-2, 2:-2] == 1), ta_hourly, vs.ta[2:-2, 2:-2, vs.tau]), ) vs.event_id = update( vs.event_id, - at[vs.tau], npx.where(cond7, vs.event_id_counter, vs.event_id[vs.tau]), + at[vs.tau], + npx.where(cond7, vs.event_id_counter, vs.event_id[vs.tau]), ) vs.dt = npx.where(cond7, 1, vs.dt) vs.itt_day = npx.where(cond7, vs.itt_day + 6, vs.itt_day) vs.pet = update( vs.pet, - at[2:-2, 2:-2], npx.where((cond8_arr[2:-2, 2:-2] == 1), pet_daily, vs.pet[2:-2, 2:-2]), + at[2:-2, 2:-2], + npx.where((cond8_arr[2:-2, 2:-2] == 1), pet_daily, vs.pet[2:-2, 2:-2]), ) vs.ta = update( vs.ta, - at[2:-2, 2:-2, vs.tau], npx.where((cond8_arr[2:-2, 2:-2] == 1), ta_daily, vs.ta[2:-2, 2:-2, vs.tau]), + at[2:-2, 2:-2, vs.tau], + npx.where((cond8_arr[2:-2, 2:-2] == 1), ta_daily, vs.ta[2:-2, 2:-2, vs.tau]), ) vs.dt = npx.where(cond8, 24, vs.dt) vs.itt_day = npx.where(cond8, 0, vs.itt_day) vs.pet = update( vs.pet, - at[2:-2, 2:-2], npx.where((cond9_arr[2:-2, 2:-2] == 1), pet_10mins, vs.pet[2:-2, 2:-2]), + at[2:-2, 2:-2], + npx.where((cond9_arr[2:-2, 2:-2] == 1), pet_10mins, vs.pet[2:-2, 2:-2]), ) vs.ta = update( vs.ta, - at[2:-2, 2:-2, vs.tau], npx.where((cond9_arr[2:-2, 2:-2] == 1), ta_10mins, vs.ta[2:-2, 2:-2, vs.tau]), + at[2:-2, 2:-2, vs.tau], + npx.where((cond9_arr[2:-2, 2:-2] == 1), ta_10mins, vs.ta[2:-2, 2:-2, vs.tau]), ) vs.event_id = update( vs.event_id, - at[vs.tau], npx.where(cond9, 0, vs.event_id[vs.tau]), + at[vs.tau], + npx.where(cond9, 0, vs.event_id[vs.tau]), ) vs.dt = npx.where(cond9, 1 / 6, vs.dt) vs.dt_secs = npx.where(cond9, 10 * 60, vs.dt_secs) @@ -201,15 +303,18 @@ def adaptive_time_stepping_kernel(state): vs.pet = update( vs.pet, - at[2:-2, 2:-2], npx.where((cond10_arr[2:-2, 2:-2] == 1), pet_hourly, vs.pet[2:-2, 2:-2]), + at[2:-2, 2:-2], + npx.where((cond10_arr[2:-2, 2:-2] == 1), pet_hourly, vs.pet[2:-2, 2:-2]), ) vs.ta = update( vs.ta, - at[2:-2, 2:-2, vs.tau], npx.where((cond10_arr[2:-2, 2:-2] == 1), ta_hourly, vs.ta[2:-2, 2:-2, vs.tau]), + at[2:-2, 2:-2, vs.tau], + npx.where((cond10_arr[2:-2, 2:-2] == 1), ta_hourly, vs.ta[2:-2, 2:-2, vs.tau]), ) vs.event_id = update( vs.event_id, - at[vs.tau], npx.where(cond10, 0, vs.event_id[vs.tau]), + at[vs.tau], + npx.where(cond10, 0, vs.event_id[vs.tau]), ) vs.dt = npx.where(cond10, 1, vs.dt) vs.dt_secs = npx.where(cond10, 60 * 60, vs.dt_secs) @@ -217,37 +322,43 @@ def adaptive_time_stepping_kernel(state): vs.pet = update( vs.pet, - at[2:-2, 2:-2], npx.where((cond11_arr[2:-2, 2:-2] == 1), pet_daily, vs.pet[2:-2, 2:-2]), + at[2:-2, 2:-2], + npx.where((cond11_arr[2:-2, 2:-2] == 1), pet_daily, vs.pet[2:-2, 2:-2]), ) vs.ta = update( vs.ta, - at[2:-2, 2:-2, vs.tau], npx.where((cond11_arr[2:-2, 2:-2] == 1), ta_daily, vs.ta[2:-2, 2:-2, vs.tau]), + at[2:-2, 2:-2, vs.tau], + npx.where((cond11_arr[2:-2, 2:-2] == 1), ta_daily, vs.ta[2:-2, 2:-2, vs.tau]), ) vs.event_id = update( vs.event_id, - at[vs.tau], npx.where(cond11, 0, vs.event_id[vs.tau]), + at[vs.tau], + npx.where(cond11, 0, vs.event_id[vs.tau]), ) vs.dt = npx.where(cond11, 24, vs.dt) vs.dt_secs = npx.where(cond11, 24 * 60 * 60, vs.dt_secs) vs.itt_day = npx.where(cond11, 0, vs.itt_day) # set event id for next event - vs.event_id_counter = npx.where((vs.event_id[vs.taum1] > 0) & (vs.event_id[vs.tau] == 0), vs.event_id_counter + 1, vs.event_id_counter) + vs.event_id_counter = npx.where( + (vs.event_id[vs.taum1] > 0) & (vs.event_id[vs.tau] == 0), vs.event_id_counter + 1, vs.event_id_counter + ) # set residual PET vs.pet_res = update(vs.pet_res, at[2:-2, 2:-2], vs.pet[2:-2, 2:-2]) - return KernelOutput(prec=vs.prec, - ta=vs.ta, - pet=vs.pet, - pet_res=vs.pet_res, - dt=vs.dt, - dt_secs=vs.dt_secs, - itt_day=vs.itt_day, - event_id=vs.event_id, - time_event0=vs.time_event0, - event_id_counter=vs.event_id_counter - ) + return KernelOutput( + prec=vs.prec, + ta=vs.ta, + pet=vs.pet, + pet_res=vs.pet_res, + dt=vs.dt, + dt_secs=vs.dt_secs, + itt_day=vs.itt_day, + event_id=vs.event_id, + time_event0=vs.time_event0, + event_id_counter=vs.event_id_counter, + ) @roger_kernel @@ -264,15 +375,33 @@ def calc_prec_ta_pet_10_minutes(state): @roger_kernel def calc_prec_ta_pet_hourly(state): vs = state.variables - idx = allocate(state.dimensions, ("x", "y", 6*24)) + idx = allocate(state.dimensions, ("x", "y", 6 * 24)) idx = update( idx, - at[2:-2, 2:-2, :], npx.arange(0, 6*24)[npx.newaxis, npx.newaxis, :], + at[2:-2, 2:-2, :], + npx.arange(0, 6 * 24)[npx.newaxis, npx.newaxis, :], ) - prec = npx.sum(npx.where((idx[2:-2, 2:-2, :] >= vs.itt_day) & (idx[2:-2, 2:-2, :] < vs.itt_day+6), vs.prec_day[2:-2, 2:-2, :], 0), axis=-1) - ta = npx.nanmean(npx.where((idx[2:-2, 2:-2, :] >= vs.itt_day) & (idx[2:-2, 2:-2, :] < vs.itt_day+6), vs.ta_day[2:-2, 2:-2, :], npx.nan), axis=-1) - pet = npx.sum(npx.where((idx[2:-2, 2:-2, :] >= vs.itt_day) & (idx[2:-2, 2:-2, :] < vs.itt_day+6), vs.pet_day[2:-2, 2:-2, :], 0), axis=-1) + prec = npx.sum( + npx.where( + (idx[2:-2, 2:-2, :] >= vs.itt_day) & (idx[2:-2, 2:-2, :] < vs.itt_day + 6), vs.prec_day[2:-2, 2:-2, :], 0 + ), + axis=-1, + ) + ta = npx.nanmean( + npx.where( + (idx[2:-2, 2:-2, :] >= vs.itt_day) & (idx[2:-2, 2:-2, :] < vs.itt_day + 6), + vs.ta_day[2:-2, 2:-2, :], + npx.nan, + ), + axis=-1, + ) + pet = npx.sum( + npx.where( + (idx[2:-2, 2:-2, :] >= vs.itt_day) & (idx[2:-2, 2:-2, :] < vs.itt_day + 6), vs.pet_day[2:-2, 2:-2, :], 0 + ), + axis=-1, + ) return prec, ta, pet @@ -282,7 +411,7 @@ def calc_prec_ta_pet_daily(state): vs = state.variables prec = npx.sum(vs.prec_day, axis=-1)[2:-2, 2:-2] - ta = npx.mean(vs.ta_day[2:-2, 2:-2, 0:24*6], axis=-1) - pet = npx.sum(vs.pet_day[2:-2, 2:-2, 0:24*6], axis=-1) + ta = npx.nanmean(vs.ta_day[2:-2, 2:-2, 0 : 24 * 6], axis=-1) + pet = npx.sum(vs.pet_day[2:-2, 2:-2, 0 : 24 * 6], axis=-1) return prec, ta, pet
+ 