added runoff tests

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",
+]

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",

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()

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

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)

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)