Add ability to derive variables and add selected derived forcings

CHANGELOG.md

@@ -11,6 +11,7 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0
 
 ### Added
 
+- add ability to derive variables from input datasets [\#34](https://github.com/mllam/mllam-data-prep/pull/34), @ealerskans
 - add github PR template to guide development process on github [\#44](https://github.com/mllam/mllam-data-prep/pull/44), @leifdenby
 
 ## [v0.5.0](https://github.com/mllam/mllam-data-prep/releases/tag/v0.5.0)

README.md

@@ -187,6 +187,32 @@ inputs:
         name_format: "{var_name}"
     target_output_variable: forcing
 
+  danra_derived_forcings:
+    path: https://mllam-test-data.s3.eu-north-1.amazonaws.com/single_levels.zarr
+    dims: [time, x, y]
+    derived_variables:
+      toa_radiation:
+        kwargs:
+          time: time
+          lat: lat
+          lon: lon
+        function: mllam_data_prep.derived_variables.calculate_toa_radiation
+      hour_of_day:
+        kwargs:
+          time: time
+        function: mllam_data_prep.derived_variables.calculate_hour_of_day
+    dim_mapping:
+      time:
+        method: rename
+        dim: time
+      grid_index:
+        method: stack
+        dims: [x, y]
+      forcing_feature:
+        method: stack_variables_by_var_name
+        name_format: "{var_name}"
+    target_output_variable: forcing
+
   danra_lsm:
     path: https://mllam-test-data.s3.eu-north-1.amazonaws.com/lsm.zarr
     dims: [x, y]
@@ -286,15 +312,40 @@ inputs:
       grid_index:
         method: stack
         dims: [x, y]
-    target_architecture_variable: state
+    target_output_variable: state
 
   danra_surface:
     path: https://mllam-test-data.s3.eu-north-1.amazonaws.com/single_levels.zarr
     dims: [time, x, y]
     variables:
-      # shouldn't really be using sea-surface pressure as "forcing", but don't
-      # have radiation varibles in danra yet
-      - pres_seasurface
+      # use surface incoming shortwave radiation as forcing
+      - swavr0m
+    dim_mapping:
+      time:
+        method: rename
+        dim: time
+      grid_index:
+        method: stack
+        dims: [x, y]
+      forcing_feature:
+        method: stack_variables_by_var_name
+        name_format: "{var_name}"
+    target_output_variable: forcing
+
+  danra_derived_forcings:
+    path: https://mllam-test-data.s3.eu-north-1.amazonaws.com/single_levels.zarr
+    dims: [time, x, y]
+    derived_variables:
+      toa_radiation:
+        kwargs:
+          time: time
+          lat: lat
+          lon: lon
+        function: mllam_data_prep.derived_variables.calculate_toa_radiation
+      hour_of_day:
+        kwargs:
+          time: time
+        function: mllam_data_prep.derived_variables.calculate_hour_of_day
     dim_mapping:
       time:
         method: rename
@@ -305,7 +356,7 @@ inputs:
       forcing_feature:
         method: stack_variables_by_var_name
         name_format: "{var_name}"
-    target_architecture_variable: forcing
+    target_output_variable: forcing
 
   ...
 ```
@@ -315,11 +366,15 @@ The `inputs` section defines the source datasets to extract data from. Each sour
 - `path`: the path to the source dataset. This can be a local path or a URL to e.g. a zarr dataset or netCDF file, anything that can be read by `xarray.open_dataset(...)`.
 - `dims`: the dimensions that the source dataset is expected to have. This is used to check that the source dataset has the expected dimensions and also makes it clearer in the config file what the dimensions of the source dataset are.
 - `variables`: selects which variables to extract from the source dataset. This may either be a list of variable names, or a dictionary where each key is the variable name and the value defines a dictionary of coordinates to do selection on. When doing selection you may also optionally define the units of the variable to check that the units of the variable match the units of the variable in the model architecture.
-- `target_architecture_variable`: the variable in the model architecture that the source dataset should be mapped to.
+- `target_output_variable`: the variable in the model architecture that the source dataset should be mapped to.
 - `dim_mapping`: defines how the dimensions of the source dataset should be mapped to the dimensions of the model architecture. This is done by defining a method to apply to each dimension. The methods are:
   - `rename`: simply rename the dimension to the new name
   - `stack`: stack the listed dimension to create the dimension in the output
   - `stack_variables_by_var_name`: stack the dimension into the new dimension, and also stack the variable name into the new variable name. This is useful when you have multiple variables with the same dimensions that you want to stack into a single variable.
+- `derived_variables`: defines the variables to be derived from the variables available in the source dataset. This should be a dictionary where each key is the variable to be derived and the value defines a dictionary with the following additional information.
+  - `function`: the function to be used to derive a variable. This should be a string and may either be the full namespace of the function (e.g. `mllam_data_prep.derived_variables.calculate_toa_radiation`) or in case the function is included in the `mllam_data_prep.derived_variables` module it is enough with the function name only.
+  - `kwargs`: arguments for the function used to derive a variable. This is a dictionary where each key is the variable name to select from the source dataset and each value is the named argument to `function`.
+  - `attributes`: section where users can specify attributes (e.g. `units` and `long_name`) as a dictionary (not included in the example config file), where the keys are the attribute names and the values are strings. If using a function defined in `mllam_data_prep.derived_variables` this section is optional as the attributes should already be defined. In this case, adding the attributes to the config file will overwrite the already-defined ones. If using an external function, where the attributes `units` and `long_name` are not set, this section is a requirement.
 
 
 ### Config schema versioning

example.danra.yaml

@@ -73,6 +73,32 @@ inputs:
         name_format: "{var_name}"
     target_output_variable: forcing
 
+  danra_derived_forcings:
+    path: https://mllam-test-data.s3.eu-north-1.amazonaws.com/single_levels.zarr
+    dims: [time, x, y]
+    derived_variables:
+      toa_radiation:
+        kwargs:
+          time: time
+          lat: lat
+          lon: lon
+        function: mllam_data_prep.derived_variables.calculate_toa_radiation
+      hour_of_day:
+        kwargs:
+          time: time
+        function: mllam_data_prep.derived_variables.calculate_hour_of_day
+    dim_mapping:
+      time:
+        method: rename
+        dim: time
+      grid_index:
+        method: stack
+        dims: [x, y]
+      forcing_feature:
+        method: stack_variables_by_var_name
+        name_format: "{var_name}"
+    target_output_variable: forcing
+
   danra_lsm:
     path: https://mllam-test-data.s3.eu-north-1.amazonaws.com/lsm.zarr
     dims: [x, y]

mllam_data_prep/config.py

@@ -52,6 +52,24 @@ class ValueSelection:
     units: str = None
 
 
+@dataclass
+class DerivedVariable:
+    """
+    Defines a derived variables, where the kwargs (variables required
+    for the calculation) and the function (for calculating the variable)
+    are specified.
+
+    Attributes:
+        kwargs: Variables required for calculating the derived variable.
+        function: Function used to calculate the derived variable.
+        attributes: Attributes (e.g. `units` and `long_name`) for the derived variable.
+    """
+
+    kwargs: Dict[str, str]
+    function: str
+    attributes: Optional[Dict[str, str]] = field(default_factory=dict)
+
+
 @dataclass
 class DimMapping:
     """
@@ -120,7 +138,8 @@ class InputDataset:
         1) the path to the dataset,
         2) the expected dimensions of the dataset,
         3) the variables to select from the dataset (and optionally subsection
-           along the coordinates for each variable) and finally
+           along the coordinates for each variable) or the variables to derive
+           from the dataset, and finally
         4) the method by which the dimensions and variables of the dataset are
            mapped to one of the output variables (this includes stacking of all
            the selected variables into a new single variable along a new coordinate,
@@ -134,11 +153,6 @@ class InputDataset:
     dims: List[str]
         List of the expected dimensions of the dataset. E.g. `["time", "x", "y"]`.
         These will be checked to ensure consistency of the dataset being read.
-    variables: Union[List[str], Dict[str, Dict[str, ValueSelection]]]
-        List of the variables to select from the dataset. E.g. `["temperature", "precipitation"]`
-        or a dictionary where the keys are the variable names and the values are dictionaries
-        defining the selection for each variable. E.g. `{"temperature": levels: {"values": [1000, 950, 900]}}`
-        would select the "temperature" variable and only the levels 1000, 950, and 900.
     dim_mapping: Dict[str, DimMapping]
         Mapping of the variables and dimensions in the input dataset to the dimensions of the
         output variable (`target_output_variable`). The key is the name of the output dimension to map to
@@ -151,14 +165,28 @@ class InputDataset:
         (e.g. two datasets that coincide in space and time will only differ in the feature dimension,
         so the two will be combined by concatenating along the feature dimension).
         If a single shared coordinate cannot be found then an exception will be raised.
+    variables: Union[List[str], Dict[str, Dict[str, ValueSelection]]]
+        List of the variables to select from the dataset. E.g. `["temperature", "precipitation"]`
+        or a dictionary where the keys are the variable names and the values are dictionaries
+        defining the selection for each variable. E.g. `{"temperature": levels: {"values": [1000, 950, 900]}}`
+        would select the "temperature" variable and only the levels 1000, 950, and 900.
+    derived_variables: Dict[str, DerivedVariable]
+        Dictionary of variables to derive from the dataset, where the keys are the variable names and
+        the values are dictionaries defining the necessary function and kwargs. E.g.
+        `{"toa_radiation": {"kwargs": {"time": "time", "lat": "lat", "lon": "lon"}, "function": "calculate_toa_radiation"}}`
+        would derive the "toa_radiation" variable using the `calculate_toa_radiation` function, which
+        takes `time`, `lat` and `lon` as arguments.
+    attributes: Dict[str, Any]
+        Optional dictionary with dataset attributes.
     """
 
     path: str
     dims: List[str]
-    variables: Union[List[str], Dict[str, Dict[str, ValueSelection]]]
     dim_mapping: Dict[str, DimMapping]
     target_output_variable: str
-    attributes: Dict[str, Any] = None
+    variables: Optional[Union[List[str], Dict[str, Dict[str, ValueSelection]]]] = None
+    derived_variables: Optional[Dict[str, DerivedVariable]] = None
+    attributes: Optional[Dict[str, Any]] = field(default_factory=dict)
 
 
 @dataclass
@@ -258,7 +286,7 @@ class Output:
 
     variables: Dict[str, List[str]]
     coord_ranges: Dict[str, Range] = None
-    chunking: Dict[str, int] = None
+    chunking: Dict[str, int] = field(default_factory=dict)
     splitting: Splitting = None
 
 

mllam_data_prep/create_dataset.py

@@ -10,7 +10,8 @@
 
 from . import __version__
 from .config import Config, InvalidConfigException
-from .ops.loading import load_and_subset_dataset
+from .derived_variables import derive_variables
+from .ops.loading import load_dataset, subset_dataset
 from .ops.mapping import map_dims_and_variables
 from .ops.selection import select_by_kwargs
 from .ops.statistics import calc_stats
@@ -30,11 +31,14 @@ def _check_dataset_attributes(ds, expected_attributes, dataset_name):
 
     # check for attributes having the wrong value
     incorrect_attributes = {
-        k: v for k, v in expected_attributes.items() if ds.attrs[k] != v
+        key: val for key, val in expected_attributes.items() if ds.attrs[key] != val
     }
     if len(incorrect_attributes) > 0:
         s_list = "\n".join(
-            [f"{k}: {v} != {ds.attrs[k]}" for k, v in incorrect_attributes.items()]
+            [
+                f"{key}: {val} != {ds.attrs[key]}"
+                for key, val in incorrect_attributes.items()
+            ]
         )
         raise ValueError(
             f"Dataset {dataset_name} has the following incorrect attributes: {s_list}"
@@ -120,23 +124,51 @@ def create_dataset(config: Config):
 
     output_config = config.output
     output_coord_ranges = output_config.coord_ranges
+    chunking_config = config.output.chunking
 
     dataarrays_by_target = defaultdict(list)
 
     for dataset_name, input_config in config.inputs.items():
         path = input_config.path
         variables = input_config.variables
+        derived_variables = input_config.derived_variables
         target_output_var = input_config.target_output_variable
-        expected_input_attributes = input_config.attributes or {}
+        expected_input_attributes = input_config.attributes
         expected_input_var_dims = input_config.dims
 
         output_dims = output_config.variables[target_output_var]
 
         logger.info(f"Loading dataset {dataset_name} from {path}")
         try:
-            ds = load_and_subset_dataset(fp=path, variables=variables)
+            ds_source = load_dataset(fp=path)
         except Exception as ex:
             raise Exception(f"Error loading dataset {dataset_name} from {path}") from ex
+
+        if variables:
+            logger.info(f"Subsetting dataset {dataset_name}")
+            try:
+                ds = subset_dataset(
+                    ds=ds_source, variables=variables, chunking=chunking_config
+                )
+            except Exception as ex:
+                raise Exception(
+                    f"Error subsetting dataset {dataset_name} from {path}"
+                ) from ex
+
+        if derived_variables:
+            logger.info(f"Deriving variables from {dataset_name}")
+            try:
+                ds = derive_variables(
+                    ds=ds_source,
+                    derived_variables=derived_variables,
+                    chunking=chunking_config,
+                )
+            except Exception as ex:
+                raise Exception(
+                    f"Error deriving variables '{', '.join(list(derived_variables.keys()))}'"
+                    f" from dataset {dataset_name} from {path}"
+                ) from ex
+
         _check_dataset_attributes(
             ds=ds,
             expected_attributes=expected_input_attributes,
@@ -191,9 +223,8 @@ def create_dataset(config: Config):
 
     # default to making a single chunk for each dimension if chunksize is not specified
     # in the config
-    chunking_config = config.output.chunking or {}
     logger.info(f"Chunking dataset with {chunking_config}")
-    chunks = {d: chunking_config.get(d, int(ds[d].count())) for d in ds.dims}
+    chunks = {dim: chunking_config.get(dim, int(ds[dim].count())) for dim in ds.dims}
     ds = ds.chunk(chunks)
 
     splitting = config.output.splitting