diff --git a/ilamb3/analysis/__init__.py b/ilamb3/analysis/__init__.py index 8d04021..33b7482 100644 --- a/ilamb3/analysis/__init__.py +++ b/ilamb3/analysis/__init__.py @@ -3,5 +3,11 @@ from ilamb3.analysis.bias import bias_analysis from ilamb3.analysis.nbp import nbp_analysis from ilamb3.analysis.relationship import relationship_analysis +from ilamb3.analysis.runoff_sensitivity import runoff_sensitivity_analysis -__all__ = ["bias_analysis", "relationship_analysis", "nbp_analysis"] +__all__ = [ + "bias_analysis", + "nbp_analysis", + "relationship_analysis", + "runoff_sensitivity_analysis", +] diff --git a/ilamb3/analysis/runoff_sensitivity.py b/ilamb3/analysis/runoff_sensitivity.py index f3c3919..5200493 100644 --- a/ilamb3/analysis/runoff_sensitivity.py +++ b/ilamb3/analysis/runoff_sensitivity.py @@ -131,7 +131,7 @@ def compute_runoff_sensitivity( return df -class runoff_sensitivty_analysis(ILAMBAnalysis): +class runoff_sensitivity_analysis(ILAMBAnalysis): """ Runoff sensitivity to temperature and precipitation per river basin. @@ -320,7 +320,7 @@ def plots( # Initialize the analysis ref_file = Path("ref.parquet") - analysis = runoff_sensitivty_analysis( + analysis = runoff_sensitivity_analysis( basin_source="mrb_basins.nc", sensitivity_frame=ref_file if ref_file.is_file() else None, mrro_source="mrro | LORA", diff --git a/ilamb3/post.py b/ilamb3/post.py index a469d96..2c5cd02 100644 --- a/ilamb3/post.py +++ b/ilamb3/post.py @@ -30,7 +30,7 @@ def get_plots(ds_ref: xr.Dataset, dsd_com: dict[str, xr.Dataset]) -> list[str]: """ # Which plots do we find in both? plots = set(ds_ref) - plots.union(*[set(ds) for _, ds in dsd_com.items()]) + plots = plots.union(*[set(ds) for _, ds in dsd_com.items()]) return list(plots) @@ -71,12 +71,14 @@ def _append_finite(values, da): # numpydoc ignore=GL08 plots = get_plots(ds_ref, dsd_com) limits = {key: [] for key in plots} - limits = {key: _append_finite(limits[key], ds_ref[key]) for key in ds_ref} - limits = { - key: _append_finite(limits[key], com[key]) - for _, com in dsd_com.items() - for key in com - } + limits.update({key: _append_finite(limits[key], ds_ref[key]) for key in ds_ref}) + limits.update( + { + key: _append_finite(limits[key], com[key]) + for _, com in dsd_com.items() + for key in com + } + ) limits = { key: np.quantile(np.hstack(arr), [outlier_fraction, 1 - outlier_fraction]) for key, arr in limits.items() diff --git a/ilamb3/tests/test_dataset.py b/ilamb3/tests/test_dataset.py index a3332f8..cb80c90 100644 --- a/ilamb3/tests/test_dataset.py +++ b/ilamb3/tests/test_dataset.py @@ -5,10 +5,10 @@ from ilamb3 import dataset as dset -def generate_test_dset(seed: int = 1): +def generate_test_dset(seed: int = 1, shift: bool = False): rs = np.random.RandomState(seed) lat = [-67.5, -22.5, 22.5, 67.5] - lon = [-135.0, -45.0, 45.0, 135.0] + lon = [-135.0 + 360 * shift, -45.0 + 360 * shift, 45.0, 135.0] time = pd.date_range(start="2000-01-15", periods=5, freq="30D") ds = xr.Dataset( data_vars={ @@ -179,3 +179,9 @@ def test_is_spatial_or_site(): ds = generate_test_site_dset() assert not dset.is_spatial(ds) assert dset.is_site(ds) + + +def test_shift_lon(): + ds = generate_test_dset(shift=True) + ds = dset.shift_lon(ds) + assert ds["lon"].min() < -120 diff --git a/ilamb3/tests/test_post.py b/ilamb3/tests/test_post.py new file mode 100644 index 0000000..1123305 --- /dev/null +++ b/ilamb3/tests/test_post.py @@ -0,0 +1,63 @@ +import cartopy.crs as ccrs +import numpy as np +import xarray as xr + +from ilamb3 import post + + +def generate_test_dset(seed: int = 1, shift: bool = False): + rs = np.random.RandomState(seed) + lat = [-67.5, -22.5, 22.5, 67.5] + lon = [-135.0 + 360 * shift, -45.0 + 360 * shift, 45.0, 135.0] + ds = xr.Dataset( + data_vars={ + "da": xr.DataArray( + rs.rand(len(lat), len(lon)) * 1e-8, + coords=[lat, lon], + dims=["lat", "lon"], + ), + } + ) + ds["da"].attrs["units"] = "kg m-2 s-1" + return ds + + +def generate_test_site_dset(seed: int = 1): + rs = np.random.RandomState(seed) + lat = xr.DataArray(data=[-67.5, -22.5, 22.5, 67.5], dims=["site"]) + lon = xr.DataArray(data=[-135.0, -45.0, 45.0, 135.0], dims=["site"]) + ds = xr.Dataset( + data_vars={ + "da": xr.DataArray( + rs.rand(lat.size) * 1e-8, + dims=["time", "site"], + ), + } + ) + ds["da"] = ds["da"].assign_coords({"lat": lat, "lon": lon}) + ds["da"].attrs["units"] = "kg m-2 s-1" + ds["da"].attrs["coordinates"] = "lat lon" + return ds + + +def test_get_plot_limits(): + limits = post.get_plot_limits( + generate_test_dset(seed=1).rename_vars({"da": "a"}), + { + "m": generate_test_dset(seed=1).rename_vars({"da": "b"}), + "n": generate_test_dset(seed=1).rename_vars({"da": "c"}), + }, + ) + assert not (set(limits) - set(["a", "b", "c"])) + + +def test_extents(): + ds = generate_test_dset() + ext = post._plot_extents(ds["da"]) + assert np.allclose(ext, [-135.0, 135.0, -67.5, 67.5]) + + +def test_proj(): + ds = generate_test_dset() + proj = post._plot_projection(post._plot_extents(ds["da"])) + assert isinstance(proj[0], ccrs.PlateCarree) diff --git a/ilamb3/tests/test_runoff.py b/ilamb3/tests/test_runoff.py new file mode 100644 index 0000000..1b45e2a --- /dev/null +++ b/ilamb3/tests/test_runoff.py @@ -0,0 +1,27 @@ +import numpy as np +import xarray as xr + +from ilamb3.analysis.runoff_sensitivity import compute_runoff_sensitivity +from ilamb3.regions import Regions +from ilamb3.tests.test_compare import generate_test_dset + + +def test_runoff_sensitivity(): + ilamb_regions = Regions() + ilamb_regions.add_latlon_bounds("test", "test", [-80, 80], [-170, 170]) + ds = xr.Dataset( + { + "pr": 1e3 * generate_test_dset(seed=1, ntime=240, nlat=10, nlon=20)["da"], + "tas": 273 + + 1e9 * generate_test_dset(seed=2, ntime=240, nlat=10, nlon=20)["da"], + "mrro": generate_test_dset(seed=3, ntime=240, nlat=10, nlon=20)["da"], + } + ) + ds["mrro"] = ( + ds["mrro"].mean() + + 1e2 * (ds["pr"] - ds["pr"].mean()) + + 1e-3 * (ds["tas"] - ds["tas"].mean()) + ) + df = compute_runoff_sensitivity(ds, ["test"]) + assert np.allclose(df.loc["test"]["tas Sensitivity"], 28006.878361) + assert np.allclose(df.loc["test"]["pr Sensitivity"], 139.850275)