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, @mafdmi
 - add github PR template to guide development process on github [\#44](https://github.com/mllam/mllam-data-prep/pull/44), @leifdenby
 - add support for zarr 3.0.0 and above [\#51](https://github.com/mllam/mllam-data-prep/pull/51), @kashif
 

README.md

@@ -112,7 +112,7 @@ ds = mdp.create_dataset(config=config)
 A full example configuration file is given in [example.danra.yaml](example.danra.yaml), and reproduced here for completeness:
 
 ```yaml
-schema_version: v0.5.0
+schema_version: v0.6.0
 dataset_version: v0.1.0
 
 output:
@@ -175,6 +175,24 @@ inputs:
     variables:
       # use surface incoming shortwave radiation as forcing
       - swavr0m
+    derived_variables:
+      # derive variables to be used as forcings
+      toa_radiation:
+        kwargs:
+          time: ds_input.time
+          lat: ds_input.lat
+          lon: ds_input.lon
+        function: mllam_data_prep.ops.derive_variable.physical_field.calculate_toa_radiation
+      hour_of_day_sin:
+        kwargs:
+          time: ds_input.time
+          component: sin
+        function: mllam_data_prep.ops.derive_variable.time_components.calculate_hour_of_day
+      hour_of_day_cos:
+        kwargs:
+          time: ds_input.time
+          component: cos
+        function: mllam_data_prep.ops.derive_variable.time_components.calculate_hour_of_day
     dim_mapping:
       time:
         method: rename
@@ -286,15 +304,32 @@ 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
+    derived_variables:
+      # derive variables to be used as forcings
+      toa_radiation:
+        kwargs:
+          time: ds_input.time
+          lat: ds_input.lat
+          lon: ds_input.lon
+        function: mllam_data_prep.ops.derive_variable.physical_field.calculate_toa_radiation
+      hour_of_day_sin:
+        kwargs:
+          time: ds_input.time
+          component: sin
+        function: mllam_data_prep.ops.derive_variable.time_components.calculate_hour_of_day
+      hour_of_day_cos:
+        kwargs:
+          time: ds_input.time
+          component: cos
+        function: mllam_data_prep.ops.derive_variable.time_components.calculate_hour_of_day
     dim_mapping:
       time:
         method: rename
@@ -305,7 +340,7 @@ inputs:
       forcing_feature:
         method: stack_variables_by_var_name
         name_format: "{var_name}"
-    target_architecture_variable: forcing
+    target_output_variable: forcing
 
   ...
 ```
@@ -315,11 +350,49 @@ 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 name of the variable to be derived and the value defines a dictionary with the following additional information. See also the 'Derived Variables' section for more details.
+  - `function`: the function used to derive a variable. This should be a string with the full namespace of the function, e.g. `mllam_data_prep.ops.derived_variables.physical_field.calculate_toa_radiation`.
+  - `kwargs`: arguments to `function`. This is a dictionary where each key is the named argument to `function` and each value is the input to the function. Here we distinguish between values to be extracted/selected from the input dataset and values supplied by the users themselves. Arguments with values to be extracted from the input dataset need to be prefixed with "ds_input." to distinguish them from other arguments. See the 'Derived Variables' section for more details.
+
+#### Derived Variables
+Variables that are not part of the source dataset but can be derived from variables in the source dataset can also be included. They should be defined in their own section, called `derived_variables` as illustrated in the example config above and in the example config file [example.danra.yaml](example.danra.yaml).
+
+To derive a variable, the function to be used (`function`) and the arguments to this function (`kwargs`) need to be specified, as explained above. Here we need to distinguish between arguments that should be data from the input dataset and arguments that should be supplied by the users themselves. The example below illustrates how to derive the cosine component of the cyclically encoded hour of day variable
+
+```yaml
+    derived_variables:
+      hour_of_day_cos:
+        kwargs:
+          time: ds_input.time
+          component: cos
+        function: mllam_data_prep.ops.derive_variable.time_components.calculate_hour_of_day
+        attrs:
+          units: 1
+          long_name: cos component of cyclically encoded hour of day
+```
+
+The function `mllam_data_prep.ops.derive_variable.time_components.calculate_hour_of_day` takes two arguments; `time` and `component`. The `time` argument should extract the `time` variable from the input dataset and has therefore been prefixed with "ds_input." to distinguish it from other arguments that should not  be extracted from the source dataset. The `component` argument, on the other hand, is a string (either "sin" or "cos") and decides if the returned derived variable is the sine or cosine component of the cyclically encoded hour of day.
+
+In addition, an optional section called `attrs` can be added. In this section, the user can add attributes to the derived variable, as illustrated in the example above. Note that the attributes `units` and `long_name` are **required**. This means that if the function used to derive a variable does not set these attributes they are **required** to be set in the config file. If using a function defined in `mllam_data_prep.ops.derive_variable` the `attrs` section is optional as the required attributes should already be defined. In this case, adding the `units` and `long_name` attributes to the `attrs` section of the derived variable in config file will **overwrite** the already-defined attributes in the function. It is also possible to set other attributes. This can be done by adding them under the `attrs` section in the same way as shown for `unit` and `long_name` in the example above.
+
+Currently, the following derived variables are included as part of `mllam-data-prep`:
+- `toa_radiation`:
+  - Top-of-atmosphere incoming radiation
+  - function: `mllam_data_prep.ops.derive_variable.physical_field.calculate_toa_radiation`
+  - arguments: `lat`, `lon`, `time`
+- `hour_of_day_[sin/cos]`:
+  - Sine or cosine part of cyclically encoded hour of day
+  - function: `mllam_data_prep.ops.derive_variable.time_components.calculate_hour_of_day`
+  - arguments: `time`, `component`
+- `day_of_year_[sin/cos]`:
+  - Sine or cosine part of cyclically encoded day of year
+  - function: `mllam_data_prep.ops.derive_variable.time_components.calculate_day_of_year`
+  - arguments: `time`, `component`
 
 
 ### Config schema versioning

example.danra.yaml

@@ -1,4 +1,4 @@
-schema_version: v0.5.0
+schema_version: v0.6.0
 dataset_version: v0.1.0
 
 output:
@@ -61,6 +61,24 @@ inputs:
     variables:
       # use surface incoming shortwave radiation as forcing
       - swavr0m
+    derived_variables:
+      # derive variables to be used as forcings
+      toa_radiation:
+        kwargs:
+          time: ds_input.time
+          lat: ds_input.lat
+          lon: ds_input.lon
+        function: mllam_data_prep.ops.derive_variable.physical_field.calculate_toa_radiation
+      hour_of_day_sin:
+        kwargs:
+          time: ds_input.time
+          component: sin
+        function: mllam_data_prep.ops.derive_variable.time_components.calculate_hour_of_day
+      hour_of_day_cos:
+        kwargs:
+          time: ds_input.time
+          component: cos
+        function: mllam_data_prep.ops.derive_variable.time_components.calculate_hour_of_day
     dim_mapping:
       time:
         method: rename

mllam_data_prep/config.py

@@ -9,6 +9,50 @@ class InvalidConfigException(Exception):
     pass
 
 
+def validate_config(config_inputs):
+    """
+    Validate that, in the config:
+    - either `variables` or `derived_variables` are present in the config
+    - if both `variables` and `derived_variables` are present, that they don't
+      add the same variables to the dataset
+
+    Parameters
+    ----------
+    config_inputs: Dict[str, InputDataset]
+
+    Returns
+    -------
+    """
+
+    for input_dataset_name, input_dataset in config_inputs.items():
+        if not input_dataset.variables and not input_dataset.derived_variables:
+            raise InvalidConfigException(
+                f"Input dataset '{input_dataset_name}' is missing the keys `variables` and/or"
+                " `derived_variables`. Make sure that you update the config so that the input"
+                f" dataset '{input_dataset_name}' contains at least either a `variables` or"
+                " `derived_variables` section."
+            )
+        elif input_dataset.variables and input_dataset.derived_variables:
+            # Check so that there are no overlapping variables
+            if isinstance(input_dataset.variables, list):
+                variable_vars = input_dataset.variables
+            elif isinstance(input_dataset.variables, dict):
+                variable_vars = input_dataset.variables.keys()
+            else:
+                raise TypeError(
+                    f"Expected an instance of list or dict, but got {type(input_dataset.variables)}."
+                )
+            derived_variable_vars = input_dataset.derived_variables.keys()
+            common_vars = list(set(variable_vars) & set(derived_variable_vars))
+            if len(common_vars) > 0:
+                raise InvalidConfigException(
+                    "Both `variables` and `derived_variables` include the following variables name(s):"
+                    f" '{', '.join(common_vars)}'. This is not allowed. Make sure that there"
+                    " are no overlapping variable names between `variables` and `derived_variables`,"
+                    f" either by renaming or removing '{', '.join(common_vars)}' from one of them."
+                )
+
+
 @dataclass
 class Range:
     """
@@ -52,6 +96,32 @@ class ValueSelection:
     units: str = None
 
 
+@dataclass
+class DerivedVariable:
+    """
+    Defines a derived variables, where the function (for calculating the variable) and
+    the kwargs (arguments to function) are specified. kwargs can contain both arguments
+    which should extract/select data from the input dataset, in which case they should
+    have the "ds_input." prefix to distinguish them from other argument that should not
+    be extracted from the dataset (e.g. a string to indicate if the sine or cosine
+    component should be extracted).
+
+    Optionally, attributes to the derived variable can be specified in `attrs`, e.g.
+    {"attrs": "units": "W*m**-2, "long_name": "top-of-the-atmosphere radiation"}.
+    In case a function does not return an `xr.DataArray` with the required attributes
+    (`units` and `long_name`) set, these have to be specified in `attrs`.
+
+    Attributes:
+        kwargs: Variables required for calculating the derived variable.
+        function: Function used to calculate the derived variable.
+        attrs: Attributes (e.g. `units` and `long_name`) to set for the derived variable.
+    """
+
+    kwargs: Dict[str, str]
+    function: str
+    attrs: Optional[Dict[str, str]] = field(default_factory=dict)
+
+
 @dataclass
 class DimMapping:
     """
@@ -120,7 +190,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 +205,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 +217,23 @@ 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 names variables will be given and
+        the values are `DerivedVariable` definitions that specify how to derive a variable.
     """
 
     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 +333,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
 
 
@@ -298,6 +373,9 @@ class Config(dataclass_wizard.JSONWizard, dataclass_wizard.YAMLWizard):
     dataset_version: str
     extra: Dict[str, Any] = None
 
+    def __post_init__(self):
+        validate_config(self.inputs)
+
     class _(JSONWizard.Meta):
         raise_on_unknown_json_key = True