Merge pull request #42 from dgergel/add_cmip6_model_cleanup_service

add new leapday removal service and cmip6 cleanup service

Author: brews

dodola/cli.py

@@ -51,6 +51,27 @@ def dodola_cli(debug):
     logging.basicConfig(level=loglevel)
 
 
+@dodola_cli.command(help="Clean up and standardize GCM")
+@click.argument("x", required=True)
+@click.argument("out", required=True)
+@click.option(
+    "--drop-leapdays/--no-drop-leapdays",
+    default=True,
+    help="Whether to remove leap days",
+)
+def cleancmip6(x, out, drop_leapdays):
+    """Clean and standardize CMIP6 GCM to 'out'. If drop-leapdays option is set, remove leap days"""
+    services.clean_cmip6(x, out, drop_leapdays, storage=_authenticate_storage())
+
+
+@dodola_cli.command(help="Remove leap days and update calendar")
+@click.argument("x", required=True)
+@click.argument("out", required=True)
+def removeleapdays(x, out):
+    """ Remove leap days and update calendar attribute"""
+    services.remove_leapdays(x, out, storage=_authenticate_storage())
+
+
 @dodola_cli.command(help="Bias-correct GCM on observations")
 @click.argument("x", required=True)
 @click.argument("xtrain", required=True)

dodola/core.py

@@ -6,6 +6,7 @@
 from skdownscale.pointwise_models import PointWiseDownscaler, BcsdTemperature
 import xarray as xr
 from xclim import sdba
+from xclim.core.calendar import convert_calendar
 import xesmf as xe
 
 # Break this down into a submodule(s) if needed.
@@ -117,3 +118,58 @@ def xesmf_regrid(x, domain, method, weights_path=None):
         filename=weights_path,
     )
     return regridder(x)
+
+
+def standardize_gcm(ds, leapday_removal=True):
+    """
+
+    Parameters
+    ----------
+    x : xr.Dataset
+    leapday_removal : bool, optional
+
+    Returns
+    -------
+    xr.Dataset
+    """
+    dims_to_drop = []
+    if "height" in ds.dims:
+        dims_to_drop.append("height")
+    if "member_id" in ds.dims:
+        dims_to_drop.append("member_id")
+    if "time_bnds" in ds.dims:
+        dims_to_drop.append("time_bnds")
+
+    if "member_id" in ds.dims:
+        ds_cleaned = ds.isel(member_id=0).drop(dims_to_drop)
+    else:
+        ds_cleaned = ds.isel.drop(dims_to_drop)
+
+    if leapday_removal:
+        # if calendar is just integers, xclim cannot understand it
+        if ds.time.dtype == "int64":
+            ds_cleaned["time"] = xr.decode_cf(ds_cleaned).time
+        # remove leap days and update calendar
+        ds_noleap = xclim_remove_leapdays(ds_cleaned)
+
+        # rechunk, otherwise chunks are different sizes
+        ds_out = ds_noleap.chunk(730, len(ds.lat), len(ds.lon))
+    else:
+        ds_out = ds_cleaned
+
+    return ds_out
+
+
+def xclim_remove_leapdays(ds):
+    """
+
+    Parameters
+    ----------
+    ds : xr.Dataset
+
+    Returns
+    -------
+    xr.Dataset
+    """
+    ds_noleap = convert_calendar(ds, target="noleap")
+    return ds_noleap

dodola/services.py

@@ -5,7 +5,13 @@
 import os
 from tempfile import TemporaryDirectory
 from rechunker import rechunk as rechunker_rechunk
-from dodola.core import apply_bias_correction, build_xesmf_weights_file, xesmf_regrid
+from dodola.core import (
+    apply_bias_correction,
+    build_xesmf_weights_file,
+    xesmf_regrid,
+    standardize_gcm,
+    xclim_remove_leapdays,
+)
 
 logger = logging.getLogger(__name__)
 
@@ -158,6 +164,44 @@ def regrid(x, out, method, storage, domain_file, weights_path=None):
     storage.write(out, regridded_ds)
 
 
+@log_service
+def clean_cmip6(x, out, leapday_removal, storage):
+    """Cleans and standardizes CMIP6 GCM
+
+    Parameters
+    ----------
+    x : str
+        Storage URL to input xr.Dataset that will be cleaned.
+    out : str
+        Storage URL to write cleaned GCM output to.
+    leapday_removal : bool
+        Whether or not to remove leap days.
+    storage : dodola.repository._ZarrRepo
+        Storage abstraction for data IO.
+    """
+    ds = storage.read(x)
+    cleaned_ds = standardize_gcm(ds, leapday_removal)
+    storage.write(out, cleaned_ds)
+
+
+@log_service
+def remove_leapdays(x, out, storage):
+    """Removes leap days and updates calendar attribute
+
+    Parameters
+    ----------
+    x : str
+        Storage URL to input xr.Dataset that will be regridded.
+    out : str
+        Storage URL to write regridded output to.
+    storage : dodola.repository._ZarrRepo
+        Storage abstraction for data IO.
+    """
+    ds = storage.read(x)
+    noleap_ds = xclim_remove_leapdays(ds)
+    storage.write(out, noleap_ds)
+
+
 @log_service
 def disaggregate(x, weights, out, repo):
     """This is just an example. Please replace or delete."""

dodola/tests/test_services.py

@@ -6,7 +6,15 @@
 import xarray as xr
 from xesmf.data import wave_smooth
 from xesmf.util import grid_global
-from dodola.services import bias_correct, build_weights, rechunk, regrid
+from xclim.core.calendar import convert_calendar
+from dodola.services import (
+    bias_correct,
+    build_weights,
+    rechunk,
+    regrid,
+    remove_leapdays,
+    clean_cmip6,
+)
 from dodola.repository import memory_repository
 
 
@@ -32,6 +40,39 @@ def _datafactory(x, start_time="1950-01-01"):
     return out
 
 
+def _gcmfactory(x, start_time="1950-01-01"):
+    """Populate xr.Dataset with synthetic GCM data for testing
+    that includes extra dimensions and leap days to be removed.
+    """
+    start_time = str(start_time)
+    if x.ndim != 1:
+        raise ValueError("'x' needs dim of one")
+
+    time = xr.cftime_range(
+        start=start_time, freq="D", periods=len(x), calendar="standard"
+    )
+
+    out = xr.Dataset(
+        {
+            "fakevariable": (
+                ["time", "lon", "lat", "member_id"],
+                x[:, np.newaxis, np.newaxis, np.newaxis],
+            )
+        },
+        coords={
+            "index": time,
+            "time": time,
+            "lon": (["lon"], [1.0]),
+            "lat": (["lat"], [1.0]),
+            "member_id": (["member_id"], [1.0]),
+            "height": (["height"], [1.0]),
+            "time_bnds": (["time_bnds"], [1.0]),
+        },
+    )
+    # out['time'] = out['time'].assign_attrs({'calendar': 'standard'})
+    return out
+
+
 @pytest.fixture
 def domain_file(request):
     """ Creates a fake domain Dataset for testing"""
@@ -299,3 +340,50 @@ def test_regrid_weights_integration(domain_file, tmpdir):
     )
     actual_shape = fakestorage.read("an/output/path.zarr")["fakevariable"].shape
     assert actual_shape == expected_shape
+
+
+def test_clean_cmip6():
+    """ Tests that cmip6 cleanup removes extra dimensions on dataset """
+    # Setup input data
+    n = 1500  # need over four years of daily data
+    ts = np.sin(np.linspace(-10 * np.pi, 10 * np.pi, n)) * 0.5
+    ds_gcm = _gcmfactory(ts, start_time="1950-01-01")
+
+    fakestorage = memory_repository(
+        {
+            "an/input/path.zarr": ds_gcm,
+        }
+    )
+
+    clean_cmip6(
+        "an/input/path.zarr",
+        "an/output/path.zarr",
+        storage=fakestorage,
+        leapday_removal=True,
+    )
+    ds_cleaned = fakestorage.read("an/output/path.zarr")
+
+    assert "height" not in ds_cleaned.dims
+    assert "member_id" not in ds_cleaned.dims
+    assert "time_bnds" not in ds_cleaned.dims
+
+
+def test_remove_leapdays():
+    """ Test that leapday removal service removes leap days """
+    # Setup input data
+    n = 1500  # need over four years of daily data
+    ts = np.sin(np.linspace(-10 * np.pi, 10 * np.pi, n)) * 0.5
+    ds_leap = _gcmfactory(ts, start_time="1950-01-01")
+
+    fakestorage = memory_repository(
+        {
+            "an/input/path.zarr": ds_leap,
+        }
+    )
+
+    remove_leapdays("an/input/path.zarr", "an/output/path.zarr", storage=fakestorage)
+    ds_noleap = fakestorage.read("an/output/path.zarr")
+    ds_leapyear = ds_noleap.loc[dict(time=slice("1952-01-01", "1952-12-31"))]
+
+    # check to be sure that leap days have been removed
+    assert len(ds_leapyear.time) == 365