Add soil temperature and soil moisture to ERA5 data, STEMMUS_SCOPE recipe

pyproject.toml

@@ -131,10 +131,11 @@ markers = [
 ignore_missing_imports = true
 disallow_untyped_defs = true
 python_version = "3.10"
+exclude = "tests"
 
 [tool.ruff]
 line-length = 88
-exclude = ["docs", "build"]
+exclude = ["docs", "build", "tests"]
 target-version = "py310"
 
 [tool.ruff.lint]

src/zampy/__init__.py

@@ -1,4 +1,5 @@
 """zampy."""
+
 from zampy import datasets
 
 

src/zampy/cli.py

@@ -1,6 +1,8 @@
 """Implements CLI interface for Zampy."""
+
 from pathlib import Path
 import click
+import dask.distributed
 from zampy.recipe import RecipeManager
 
 
@@ -14,4 +16,5 @@ def run_recipe(recipe: Path) -> None:
 
 
 if __name__ == "__main__":
+    dask.distributed.Client()
     run_recipe()

src/zampy/conventions/ALMA.json

@@ -84,5 +84,13 @@
     "land_cover": {
         "variable": "land_cover",
         "units": ""
+    },
+    "soil_temperature": {
+        "variable": "SoilTemp",
+        "units": "kelvin"
+    },
+    "soil_moisture": {
+        "variable": "SoilMoist",
+        "units": "kilogram/meter**3"
     }
 }

src/zampy/datasets/__init__.py

@@ -1,4 +1,5 @@
 """Datasets implementations."""
+
 from zampy.datasets import dataset_protocol
 from zampy.datasets import validation
 from zampy.datasets.catalog import DATASETS

src/zampy/datasets/catalog.py

@@ -1,4 +1,5 @@
 """Catalog of datasets."""
+
 from zampy.datasets import dataset_protocol
 from zampy.datasets.cams import CAMS
 from zampy.datasets.era5 import ERA5

src/zampy/datasets/cds_utils.py

@@ -1,5 +1,6 @@
 """CDS utilities used by ECMWF datasets."""
 
+from copy import copy
 from pathlib import Path
 import cdsapi
 import numpy as np
@@ -44,6 +45,21 @@
     "21:00", "22:00", "23:00",
 ]  # fmt: skip
 
+SPLIT_VARIABLES = {
+    "soil_temperature": (
+        "soil_temperature_level_1",
+        "soil_temperature_level_2",
+        "soil_temperature_level_3",
+        "soil_temperature_level_4",
+    ),
+    "soil_moisture": (
+        "volumetric_soil_water_layer_1",
+        "volumetric_soil_water_layer_2",
+        "volumetric_soil_water_layer_3",
+        "volumetric_soil_water_layer_4",
+    ),
+}
+
 
 def cds_request(
     dataset: str,
@@ -226,6 +242,12 @@ def retrieve_era5(
     # create list of year/month pairs
     year_month_pairs = time_bounds_to_year_month(time_bounds)
 
+    variables = copy(variables)  # Prevent original input from being modified in-place
+    for split_var in SPLIT_VARIABLES:
+        if split_var in variables:
+            variables.remove(split_var)
+            variables.extend(SPLIT_VARIABLES[split_var])
+
     for (year, month), variable in product(
         year_month_pairs, variables, position=0, leave=True
     ):
@@ -354,7 +376,8 @@ def convert_to_zampy(
         print(f"File '{ncfile.name}' already exists, skipping...")
     else:
         ds = parse_nc_file(file)
-
+        # Rename the vswl data:
+        ncfile = Path(str(ncfile).replace("volumetric_soil_water", "soil_moisture"))
         ds.to_netcdf(path=ncfile)
 
 
@@ -373,6 +396,28 @@ def convert_to_zampy(
     "co2": "co2_concentration",
 }
 
+VAR_REFERENCE_MULTI_LAYER = {
+    "stl1": "soil_temperature",
+    "stl2": "soil_temperature",
+    "stl3": "soil_temperature",
+    "stl4": "soil_temperature",
+    "swvl1": "soil_moisture",
+    "swvl2": "soil_moisture",
+    "swvl3": "soil_moisture",
+    "swvl4": "soil_moisture",
+}
+
+LAYER_BOUNDS = {
+    "stl1": [[0.0, 7.0]],
+    "stl2": [[7.0, 28.0]],
+    "stl3": [[28.0, 100.0]],
+    "stl4": [[100.0, 289.0]],
+    "swvl1": [[0.0, 7.0]],
+    "swvl2": [[7.0, 28.0]],
+    "swvl3": [[28.0, 100.0]],
+    "swvl4": [[100.0, 289.0]],
+}
+
 WATER_DENSITY = 997.0  # kg/m3
 
 
@@ -416,6 +461,33 @@ def parse_nc_file(file: Path) -> xr.Dataset:
                 variable_name
             ].desc
 
+        if variable in VAR_REFERENCE_MULTI_LAYER:
+            if (  # Soil temperature/moisture routine
+                str(variable).startswith("stl") or str(variable).startswith("swvl")
+            ):
+                if str(variable).startswith("swvl"):
+                    varname = "soil_moisture"
+                    standard_name = "moisture_content_of_soil_layer"
+                    ds[variable] *= WATER_DENSITY
+                    ds[variable].attrs.update({"units": "kg m**-3"})
+                else:
+                    varname = "soil_temperature"
+                    standard_name = "temperature_in_ground"
+
+                da = ds[variable]
+                name = str(da.name)
+                da = da.expand_dims({"depth": [np.mean(LAYER_BOUNDS[name])]})
+                da = da.rename(varname)
+                da.attrs.update(
+                    {
+                        "long_name": varname.replace("_", " "),
+                        "standard_name": standard_name,
+                    }
+                )
+
+                ds = da.to_dataset()
+                ds["depth_bounds"] = (("depth", "nv"), LAYER_BOUNDS[name])
+
     # TODO: add dataset attributes.
 
     return ds

src/zampy/datasets/converter.py

@@ -1,4 +1,5 @@
 """Dataset formatter for different conventions."""
+
 import json
 import warnings
 from pathlib import Path

src/zampy/datasets/dataset_protocol.py

@@ -1,4 +1,5 @@
 """Outline of the dataset protocol."""
+
 import json
 import shutil
 from dataclasses import dataclass

src/zampy/datasets/era5.py

@@ -47,15 +47,30 @@ class ERA5Land(ECMWFDataset):  # noqa: D101
     raw_variables = [
         Variable(name="t2m", unit=unit_registry.kelvin),
         Variable(name="d2m", unit=unit_registry.kelvin),
+        Variable(name="st", unit=unit_registry.kelvin),
+        Variable(name="swvl", unit=unit_registry.fraction),
     ]
 
     # variable names used in cdsapi downloading request
     cds_var_names = {
         "air_temperature": "2m_temperature",
         "dewpoint_temperature": "2m_dewpoint_temperature",
+        "soil_temperature_level_1": "soil_temperature_level_1",  # Note: split variables
+        "soil_temperature_level_2": "soil_temperature_level_2",
+        "soil_temperature_level_3": "soil_temperature_level_3",
+        "soil_temperature_level_4": "soil_temperature_level_4",
+        "volumetric_soil_water_layer_1": "volumetric_soil_water_layer_1",
+        "volumetric_soil_water_layer_2": "volumetric_soil_water_layer_2",
+        "volumetric_soil_water_layer_3": "volumetric_soil_water_layer_3",
+        "volumetric_soil_water_layer_4": "volumetric_soil_water_layer_4",
     }
 
-    variable_names = list(cds_var_names.keys())
+    variable_names = [
+        "air_temperature",
+        "dewpoint_temperature",
+        "soil_temperature",
+        "soil_moisture",
+    ]
 
     variables = [VARIABLE_REFERENCE_LOOKUP[var] for var in variable_names]
 

src/zampy/datasets/eth_canopy_height.py

@@ -1,4 +1,5 @@
 """ETH canopy height dataset."""
+
 import gzip
 from pathlib import Path
 import numpy as np
@@ -269,10 +270,17 @@ def parse_tiff_file(file: Path, sd_file: bool = False) -> xr.Dataset:
     da = da.isel(band=0)  # get rid of band dim
     da = da.drop_vars(["band", "spatial_ref"])  # drop unnecessary coords
     ds = da.to_dataset()
-    ds = ds.assign_coords(  # halfway in the year
-        {"time": np.datetime64("2020-07-01").astype("datetime64[ns]")}
+    ds = xr.concat(  # Cover entirety of 2020
+        (
+            ds.assign_coords(
+                {"time": np.datetime64("2020-01-01").astype("datetime64[ns]")}
+            ),
+            ds.assign_coords(
+                {"time": np.datetime64("2021-01-01").astype("datetime64[ns]")}
+            ),
+        ),
+        dim="time",
     )
-    ds = ds.expand_dims("time")
     ds = ds.rename(
         {
             "band_data": "height_of_vegetation_standard_deviation"

src/zampy/datasets/prism_dem.py

@@ -1,4 +1,5 @@
 """Prism DEM dataset."""
+
 import gzip
 import tarfile
 from pathlib import Path

src/zampy/datasets/utils.py

@@ -1,4 +1,5 @@
 """Shared utilities from datasets."""
+
 import urllib.request
 from pathlib import Path
 import requests

src/zampy/datasets/validation.py

@@ -1,4 +1,5 @@
 """Checks for user input validation."""
+
 from pathlib import Path
 from zampy.datasets.dataset_protocol import Dataset
 from zampy.datasets.dataset_protocol import SpatialBounds

src/zampy/recipe.py

@@ -1,13 +1,17 @@
-""""All functionality to read and execute Zampy recipes."""
+"""All functionality to read and execute Zampy recipes."""
+
 from pathlib import Path
 from typing import Any
 import numpy as np
+import pandas as pd
+import xarray as xr
 import yaml
 from zampy.datasets import DATASETS
 from zampy.datasets import converter
 from zampy.datasets.dataset_protocol import Dataset
 from zampy.datasets.dataset_protocol import SpatialBounds
 from zampy.datasets.dataset_protocol import TimeBounds
+from zampy.datasets.validation import validate_download_request
 
 
 def recipe_loader(recipe_path: Path) -> dict:
@@ -31,8 +35,8 @@ def recipe_loader(recipe_path: Path) -> dict:
         for key in ["convention", "frequency", "resolution"]
     ):
         msg = (
-            "One of the following items are missing from the recipe:\n"
-            "name, download, convert."
+            "One of the following 'convert' items are missing from the recipe:\n"
+            "convention, frequency, resolution."
         )
         raise ValueError(msg)
 
@@ -92,9 +96,22 @@ def __init__(self, recipe_path: Path) -> None:
 
     def run(self) -> None:
         """Run the full recipe."""
+        # First validate all inputs (before downloading, processing...)
         for dataset_name in self.datasets:
             _dataset = DATASETS[dataset_name.lower()]
             dataset: Dataset = _dataset()
+
+            validate_download_request(
+                dataset,
+                self.download_dir,
+                self.timebounds,
+                self.spatialbounds,
+                self.datasets[dataset_name]["variables"],
+            )
+
+        for dataset_name in self.datasets:
+            _dataset = DATASETS[dataset_name.lower()]
+            dataset = _dataset()
             variables: list[str] = self.datasets[dataset_name]["variables"]
 
             # Download datset
@@ -118,7 +135,18 @@ def run(self) -> None:
             ds = converter.convert(ds, dataset, convention=self.convention)
 
             if "time" in ds.dims:  # Dataset with only DEM (e.g.) has no time dim.
-                ds = ds.resample(time=self.frequency).mean()
+                freq = xr.infer_freq(ds["time"])
+                if freq is None:  # fallback:
+                    freq = (
+                        ds["time"].isel(time=1).to_numpy()
+                        - ds["time"].isel(time=0).to_numpy()
+                    )
+                data_freq = pd.to_timedelta(pd.tseries.frequencies.to_offset(freq))
+
+                if data_freq < pd.Timedelta(self.frequency):
+                    ds = ds.resample(time=self.frequency).mean()
+                elif data_freq > pd.Timedelta(self.frequency):
+                    ds = ds.resample(time=self.frequency).interpolate("nearest")
 
             comp = dict(zlib=True, complevel=5)
             encoding = {var: comp for var in ds.data_vars}
@@ -127,6 +155,7 @@ def run(self) -> None:
             # e.g. "era5_2010-2020.nc"
             fname = f"{dataset_name.lower()}_{time_start}-{time_end}.nc"
             ds.to_netcdf(path=self.data_dir / fname, encoding=encoding)
+            del ds
 
         print(
             "Finished running the recipe. Output data can be found at:\n"

src/zampy/reference/variables.py

@@ -1,4 +1,5 @@
 """Variable reference for Zampy."""
+
 from pint import UnitRegistry
 from zampy.datasets.dataset_protocol import Variable
 
@@ -13,6 +14,7 @@ def unit_registration() -> UnitRegistry:
     unit_registry.define("degree_east = degree = degree_E = degreeE")
     unit_registry.define("watt_per_square_meter = watt/meter**2")
     unit_registry.define("joule_per_square_meter = joule/meter**2")
+    unit_registry.define("kilogram_per_square_meter = kilogram/(meter**2)")
     unit_registry.define(
         "kilogram_per_square_meter_second = kilogram/(meter**2*second)"
     )
@@ -56,6 +58,8 @@ def unit_registration() -> UnitRegistry:
     Variable("elevation", unit=unit_registry.meter),
     Variable("leaf_area_index", unit=unit_registry.fraction),
     Variable("land_cover", unit=unit_registry.dimensionless),
+    Variable("soil_temperature", unit=unit_registry.kelvin),
+    Variable("soil_moisture", unit=unit_registry.kilogram_per_square_meter),
 )
 
 VARIABLE_REFERENCE_LOOKUP = {var.name: var for var in VARIABLE_REFERENCE}

tests/test_data/fapar-lai/generate_fake_data.py

@@ -1,4 +1,5 @@
 """Generate fake data for the fapar-lai tests."""
+
 import zipfile
 from pathlib import Path
 import numpy as np

tests/test_datasets/__init__.py

@@ -1,4 +1,5 @@
 """This module contains all tests for datasets included in zampy."""
+
 from pathlib import Path