Add support for reading, defining and outputting CF-compliant projection info

.github/workflows/python-package-pip.yml

@@ -26,6 +26,10 @@ jobs:
           python -m pip install . "zarr${{ matrix.zarr-version }}"
           python -m pip install pytest
 
+      - name: Run doctests
+        run: |
+          python -m pytest --doctest-modules mllam_data_prep/
+
       - name: Run tests
         env:
           AWS_ACCESS_KEY_ID: ${{ secrets.AWS_ACCESS_KEY_ID }}

CHANGELOG.md

@@ -49,6 +49,8 @@ interface and addresses bugs around optional dependencies for
 
 ### Added
 
+- add support for reading, setting (via `projections["__common__"]` section in config) and outputting CF-Conventions complient projection information. ![\#38](https://github.com/mllam/mllam-data-prep/pull/38)
+- add `doctest-modules` to CI. ![\#38](https://github.com/mllam/mllam-data-prep/pull/38)
 - add optional output path argument to parser. ![\#26](https://github.com/mllam/mllam-data-prep/pull/26)
 
 ### Changed

example.danra.yaml

@@ -42,6 +42,10 @@ inputs:
         altitude:
           values: [100, ]
           units: m
+    projections:
+      danra_projection:
+        dims: [x, y]
+        crs_wkt: 'PROJCRS["DMI HARMONIE DANRA lambert projection", BASEGEOGCRS["DMI HARMONIE DANRA lambert CRS", DATUM["DMI HARMONIE DANRA lambert datum", ELLIPSOID["Sphere", 6367470, 0, LENGTHUNIT["metre", 1, ID["EPSG", 9001]]]], PRIMEM["Greenwich", 0, ANGLEUNIT["degree", 0.0174532925199433, ID["EPSG", 8901]]], ID["EPSG",4035]], CONVERSION["Lambert Conic Conformal (2SP)", METHOD["Lambert Conic Conformal (2SP)", ID["EPSG", 9802]], PARAMETER["Latitude of false origin", 56.7, ANGLEUNIT["degree", 0.0174532925199433, ID["EPSG", 8821]]], PARAMETER["Longitude of false origin", 25, ANGLEUNIT["degree", 0.0174532925199433, ID["EPSG", 8822]]], PARAMETER["Latitude of 1st standard parallel", 56.7, ANGLEUNIT["degree", 0.0174532925199433, ID["EPSG", 8823]]], PARAMETER["Latitude of 2nd standard parallel", 56.7, ANGLEUNIT["degree", 0.0174532925199433, ID["EPSG", 8824]]], PARAMETER["Easting at false origin", 0, LENGTHUNIT["metre", 1, ID["EPSG", 8826]]], PARAMETER["Northing at false origin", 0, LENGTHUNIT["metre", 1, ID["EPSG", 8827]]]], CS[Cartesian, 2], AXIS["(E)", east, ORDER[1], LENGTHUNIT["metre", 1, ID["EPSG", 9001]]], AXIS["(N)", north, ORDER[2], LENGTHUNIT["metre", 1, ID["EPSG", 9001]]], USAGE[AREA["Denmark and surrounding regions"], BBOX[47, -3, 65, 25], SCOPE["Danra reanalysis projection"]]]'
     dim_mapping:
       time:
         method: rename
@@ -61,6 +65,10 @@ inputs:
     variables:
       # use surface incoming shortwave radiation as forcing
       - swavr0m
+    projections:
+      danra_projection:
+        dims: [x, y]
+        crs_wkt: 'PROJCRS["DMI HARMONIE DANRA lambert projection", BASEGEOGCRS["DMI HARMONIE DANRA lambert CRS", DATUM["DMI HARMONIE DANRA lambert datum", ELLIPSOID["Sphere", 6367470, 0, LENGTHUNIT["metre", 1, ID["EPSG", 9001]]]], PRIMEM["Greenwich", 0, ANGLEUNIT["degree", 0.0174532925199433, ID["EPSG", 8901]]], ID["EPSG",4035]], CONVERSION["Lambert Conic Conformal (2SP)", METHOD["Lambert Conic Conformal (2SP)", ID["EPSG", 9802]], PARAMETER["Latitude of false origin", 56.7, ANGLEUNIT["degree", 0.0174532925199433, ID["EPSG", 8821]]], PARAMETER["Longitude of false origin", 25, ANGLEUNIT["degree", 0.0174532925199433, ID["EPSG", 8822]]], PARAMETER["Latitude of 1st standard parallel", 56.7, ANGLEUNIT["degree", 0.0174532925199433, ID["EPSG", 8823]]], PARAMETER["Latitude of 2nd standard parallel", 56.7, ANGLEUNIT["degree", 0.0174532925199433, ID["EPSG", 8824]]], PARAMETER["Easting at false origin", 0, LENGTHUNIT["metre", 1, ID["EPSG", 8826]]], PARAMETER["Northing at false origin", 0, LENGTHUNIT["metre", 1, ID["EPSG", 8827]]]], CS[Cartesian, 2], AXIS["(E)", east, ORDER[1], LENGTHUNIT["metre", 1, ID["EPSG", 9001]]], AXIS["(N)", north, ORDER[2], LENGTHUNIT["metre", 1, ID["EPSG", 9001]]], USAGE[AREA["Denmark and surrounding regions"], BBOX[47, -3, 65, 25], SCOPE["Danra reanalysis projection"]]]'
     derived_variables:
       # derive variables to be used as forcings
       toa_radiation:
@@ -96,6 +104,10 @@ inputs:
     dims: [x, y]
     variables:
       - lsm
+    projections:
+      danra_projection:
+        dims: [x, y]
+        crs_wkt: 'PROJCRS["DMI HARMONIE DANRA lambert projection", BASEGEOGCRS["DMI HARMONIE DANRA lambert CRS", DATUM["DMI HARMONIE DANRA lambert datum", ELLIPSOID["Sphere", 6367470, 0, LENGTHUNIT["metre", 1, ID["EPSG", 9001]]]], PRIMEM["Greenwich", 0, ANGLEUNIT["degree", 0.0174532925199433, ID["EPSG", 8901]]], ID["EPSG",4035]], CONVERSION["Lambert Conic Conformal (2SP)", METHOD["Lambert Conic Conformal (2SP)", ID["EPSG", 9802]], PARAMETER["Latitude of false origin", 56.7, ANGLEUNIT["degree", 0.0174532925199433, ID["EPSG", 8821]]], PARAMETER["Longitude of false origin", 25, ANGLEUNIT["degree", 0.0174532925199433, ID["EPSG", 8822]]], PARAMETER["Latitude of 1st standard parallel", 56.7, ANGLEUNIT["degree", 0.0174532925199433, ID["EPSG", 8823]]], PARAMETER["Latitude of 2nd standard parallel", 56.7, ANGLEUNIT["degree", 0.0174532925199433, ID["EPSG", 8824]]], PARAMETER["Easting at false origin", 0, LENGTHUNIT["metre", 1, ID["EPSG", 8826]]], PARAMETER["Northing at false origin", 0, LENGTHUNIT["metre", 1, ID["EPSG", 8827]]]], CS[Cartesian, 2], AXIS["(E)", east, ORDER[1], LENGTHUNIT["metre", 1, ID["EPSG", 9001]]], AXIS["(N)", north, ORDER[2], LENGTHUNIT["metre", 1, ID["EPSG", 9001]]], USAGE[AREA["Denmark and surrounding regions"], BBOX[47, -3, 65, 25], SCOPE["Danra reanalysis projection"]]]'
     dim_mapping:
       grid_index:
         method: stack

mllam_data_prep/config.py

@@ -180,6 +180,27 @@ class DimMapping:
     name_format: str = field(default=None)
 
 
+@dataclass
+class Projection:
+    """
+    Definition of a single coordinate reference system (CRS) projection that
+    variables in the dataset refer to.
+
+    Attributes
+    ----------
+    dims : List[str]
+        The dimensions of the projection, usually ["x", "y"].
+    crs_wkt : str
+        The well-known text (WKT) representation of the CRS. For compatibility with cartopy the WKT string
+        needs to contain a BBOX parameter, e.g. `crs_wkt = "PROJCS[...], BBOX[...]"`.
+        Registered CRSs at e.g. EPSG can be used as well, e.g. `crs_wkt = "EPSG:4326"`. These often will resolve to
+        a WKT string that includes the BBOX parameter (see https://epsg.io/4326).
+    """
+
+    dims: List[str]
+    crs_wkt: str
+
+
 @dataclass
 class InputDataset:
     """
@@ -233,6 +254,7 @@ class InputDataset:
     target_output_variable: str
     variables: Optional[Union[List[str], Dict[str, Dict[str, ValueSelection]]]] = None
     derived_variables: Optional[Dict[str, DerivedVariable]] = None
+    projections: Optional[Dict[str, Projection]] = None
     attributes: Optional[Dict[str, Any]] = field(default_factory=dict)
 
 

mllam_data_prep/create_dataset.py

@@ -4,6 +4,7 @@
 from pathlib import Path
 
 import numpy as np
+import pyproj
 import xarray as xr
 import zarr
 from loguru import logger
@@ -15,6 +16,13 @@
 from .ops.derive_variable import derive_variable
 from .ops.loading import load_input_dataset
 from .ops.mapping import map_dims_and_variables
+from .ops.projection import (
+    ProjectionCompatibilityWarning,
+    ProjectionInconsistencyWarning,
+    check_projection_compatibility,
+    get_projection_crs,
+    validate_projection_consistency,
+)
 from .ops.selection import select_by_kwargs
 from .ops.statistics import calc_stats
 from .ops.subsetting import extract_variable
@@ -137,6 +145,7 @@ def create_dataset(config: Config):
     chunking_config = config.output.chunking
 
     dataarrays_by_target = defaultdict(list)
+    projections = {}
 
     for dataset_name, input_config in config.inputs.items():
         path = input_config.path
@@ -145,6 +154,7 @@ def create_dataset(config: Config):
         target_output_var = input_config.target_output_variable
         expected_input_attributes = input_config.attributes
         expected_input_var_dims = input_config.dims
+        dataset_projections = input_config.projections or {}
 
         output_dims = output_config.variables[target_output_var]
 
@@ -192,6 +202,40 @@ def create_dataset(config: Config):
             dataset_name=dataset_name,
         )
 
+        # get the projection information from the dataset and update it with the projection
+        # information given in the input config
+        # TODO: read the projection information from the config
+        try:
+            projection_crs = get_projection_crs(ds)
+        except ValueError:
+            projection_crs = None
+        if projection_crs is None and dataset_projections:
+            logger.warning(
+                f"Projection information not found in dataset {dataset_name}, using projection information from config."
+            )
+            projection_crs = {
+                crs_var: {"crs_wkt": p.crs_wkt}
+                for crs_var, p in dataset_projections.items()
+            }
+        elif projection_crs is None and not dataset_projections:
+            logger.warning(
+                f"No projection information found neither in dataset {dataset_name}, nor in config."
+            )
+        elif projection_crs is not None and dataset_projections:
+            logger.warning(
+                f"Projection information given in {dataset_name} is overwritten by that given in config."
+            )
+            projection_crs = {
+                crs_var: {"crs_wkt": p.crs_wkt}
+                for crs_var, p in dataset_projections.items()
+            }
+        if projection_crs is not None:
+            projections.update(projection_crs)
+
+        # TODO: remove this once grid_mapping is set in config file
+        for var in ds.data_vars:
+            ds[var].attrs["grid_mapping"] = list(projection_crs.keys())[0]
+
         dim_mapping = input_config.dim_mapping
 
         # check that there is an entry for each arch dimension
@@ -233,6 +277,24 @@ def create_dataset(config: Config):
 
         dataarrays_by_target[target_output_var].append(da_target)
 
+    # validate projections across input datasets
+    if projections:
+        try:
+            validate_projection_consistency(projections.values())
+        except ProjectionInconsistencyWarning as e:
+            logger.warning(f"Projection information might be ambiguous: {e}")
+        projection = pyproj.CRS.from_cf(list(projections.values())[0])
+        try:
+            check_projection_compatibility(list(projections.values())[0])
+        except ProjectionCompatibilityWarning:
+            logger.warning(
+                "WKT string should include a BBOX definition to be compatible with e.g. cartopy."
+            )
+
+        # TODO: generalize the retrieval of x and y coords
+        # coords = (dataarrays_by_target[target_output_var][0]['x'], dataarrays_by_target[target_output_var][0]['y'])
+        # lon, lat = get_latitude_longitude_from_projection(projection, coords)
+
     ds = _merge_dataarrays_by_target(dataarrays_by_target=dataarrays_by_target)
 
     # need to drop the encoding so that we can write to zarr with new chunksizes
@@ -278,6 +340,14 @@ def create_dataset(config: Config):
         )
         ds["splits"] = da_splits
 
+    # add the projection information to the dataset
+    if projections:
+        ds[projections.keys()] = xr.DataArray(0, attrs=projection.to_cf()).astype(
+            "int16"
+        )
+        # for var in ["VARIABLES_WITH_PROKJECTION"]:
+        #     ds[var].attrs["grid_mapping"] = "__common__"
+
     ds.attrs = {}
     ds.attrs["schema_version"] = config.schema_version
     ds.attrs["dataset_version"] = config.dataset_version