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

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

@@ -24,6 +24,10 @@ jobs:
           python -m pip install .
           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 }}

example.danra.yaml

@@ -42,6 +42,26 @@ inputs:
         altitude:
           values: [100, ]
           units: m
+    projections:
+      crs:
+        dims: [x, y]
+        attributes:
+          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]]],CONVERSION["unknown",METHOD["Lambert Conic Conformal (2SP)",ID["EPSG",9802]],PARAMETER["Latitude of false origin",0.0567,ANGLEUNIT["degree",0.0174532925199433],ID["EPSG",8821]],PARAMETER["Longitude of false origin",0.025,ANGLEUNIT["degree",0.0174532925199433],ID["EPSG",8822]],PARAMETER["Latitude of 1st standard parallel",0.0567,ANGLEUNIT["degree",0.0174532925199433],ID["EPSG",8823]],PARAMETER["Latitude of 2nd standard parallel",0.0567,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]]]]'
+          semi_major_axis: 6367470.0
+          semi_minor_axis: 6367470.0
+          inverse_flattening: 0.0
+          reference_ellipsoid_name: 'Sphere'
+          longitude_of_prime_meridian: 0.0
+          prime_meridian_name: 'Greenwich'
+          geographic_crs_name: 'DMI HARMONIE DANRA lambert CRS'
+          horizontal_datum_name: 'DMI HARMONIE DANRA lambert datum'
+          projected_crs_name: 'DMI HARMONIE DANRA lambert projection'
+          grid_mapping_name: 'lambert_conformal_conic'
+          standard_parallel: [0.0567, 0.0567]
+          latitude_of_projection_origin: 0.0567
+          longitude_of_central_meridian: 0.025
+          false_easting: 0.0
+          false_northing: 0.0
     dim_mapping:
       time:
         method: rename

mllam_data_prep/config.py

@@ -122,6 +122,26 @@ 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"].
+    attributes : Dict[str, Any]
+        The attributes that define the projection. The attributes need to adhere to the CF-conventions
+        and contain either a `crs_wkt` attribute or individual attributes that are specific to the used projection.
+        If both are given, then the `crs_wkt` attribute will be used to create the projection object.
+    """
+
+    dims: List[str]
+    attributes: Dict[str, Any]
+
+
 @dataclass
 class InputDataset:
     """
@@ -167,10 +187,15 @@ class InputDataset:
 
     path: str
     dims: List[str]
-    variables: Union[List[str], Dict[str, Dict[str, ValueSelection]]]
+    variables: Union[
+        List[str],
+        Dict[str, Dict[str, ValueSelection]],
+        Dict[str, Dict[str, Dict[str, str]]],
+    ]
     dim_mapping: Dict[str, DimMapping]
     target_output_variable: str
     attributes: Dict[str, Any] = None
+    projections: Optional[Dict[str, Projection]] = None
 
 
 @dataclass

mllam_data_prep/create_dataset.py

@@ -4,6 +4,7 @@
 from pathlib import Path
 
 import numpy as np
+import pyproj
 import xarray as xr
 from loguru import logger
 from numcodecs import Blosc
@@ -12,6 +13,11 @@
 from .config import Config, InvalidConfigException
 from .ops.loading import load_and_subset_dataset
 from .ops.mapping import map_dims_and_variables
+from .ops.projection import (
+    ProjectionInconsistencyWarning,
+    get_projection_crs,
+    validate_projection_consistency,
+)
 from .ops.selection import select_by_kwargs
 from .ops.statistics import calc_stats
 
@@ -106,13 +112,15 @@ def create_dataset(config: Config):
     output_coord_ranges = output_config.coord_ranges
 
     dataarrays_by_target = defaultdict(list)
+    projections = []
 
     for dataset_name, input_config in config.inputs.items():
         path = input_config.path
         variables = input_config.variables
         target_output_var = input_config.target_output_variable
         expected_input_attributes = input_config.attributes or {}
         expected_input_var_dims = input_config.dims
+        dataset_projections = input_config.projections or {}
 
         output_dims = output_config.variables[target_output_var]
 
@@ -127,6 +135,10 @@ def create_dataset(config: Config):
             dataset_name=dataset_name,
         )
 
+        # TODO: remove this once grid_mapping is set in config file
+        for var in ds.data_vars:
+            ds[var].attrs["grid_mapping"] = "crs"
+
         dim_mapping = input_config.dim_mapping
 
         # check that there is an entry for each arch dimension
@@ -158,6 +170,24 @@ def create_dataset(config: Config):
 
         da_target.attrs["source_dataset"] = 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
+        projection_crs = get_projection_crs(ds)
+        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: p.attributes 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."
+            )
+        if projection_crs is not None:
+            projections.extend(projection_crs.values())
+
         # only need to do selection for the coordinates that the input dataset actually has
         if output_coord_ranges is not None:
             selection_kwargs = {}
@@ -168,6 +198,18 @@ 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)
+        except ProjectionInconsistencyWarning as e:
+            logger.warning(f"Projection information might be ambiguous: {e}")
+        projection = pyproj.CRS.from_cf(projections[0])
+
+        # 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
@@ -214,6 +256,12 @@ def create_dataset(config: Config):
         )
         ds["splits"] = da_splits
 
+    # add the projection information to the dataset
+    if projections:
+        ds["crs"] = xr.DataArray(0, attrs=projection.to_cf()).astype("int16")
+        # for var in ["VARIABLES_WITH_PROKJECTION"]:
+        #     ds[var].attrs["grid_mapping"] = "crs"
+
     ds.attrs = {}
     ds.attrs["schema_version"] = config.schema_version
     ds.attrs["dataset_version"] = config.dataset_version