Issue 476: Add coordinate name standardization for lat/lon aliases

weather_mv/loader_pipeline/sinks.py

@@ -42,6 +42,12 @@
 TIF_TRANSFORM_CRS_TO = "EPSG:4326"
 # A constant for all the things in the coords key set that aren't the level name.
 DEFAULT_COORD_KEYS = frozenset(('latitude', 'time', 'step', 'valid_time', 'longitude', 'number'))
+# Common aliases for the canonical 'latitude' / 'longitude' coordinate names (matched case-insensitively).
+LATITUDE_ALIASES = frozenset(('lat',))
+LONGITUDE_ALIASES = frozenset(('lon', 'long', 'lng'))
+# CF-convention units that identify geospatial coordinates (e.g. for datasets using 'x' / 'y' names).
+CF_LATITUDE_UNITS = frozenset(('degrees_north', 'degree_north', 'degrees_n', 'degree_n'))
+CF_LONGITUDE_UNITS = frozenset(('degrees_east', 'degree_east', 'degrees_e', 'degree_e'))
 DEFAULT_TIME_ORDER_LIST = ['%Y', '%m', '%d', '%H', '%M', '%S']
 # For uploading / downloading retry logic.
 INITIAL_DELAY = 1.0  # Initial delay in seconds.
@@ -146,6 +152,52 @@ def rearrange_time_list(order_list: t.List, time_list: t.List) -> t.List:
     return datetime.datetime(*time_list)
 
 
+def _coordinate_renames(ds: xr.Dataset) -> t.Dict[str, str]:
+    """Returns a mapping of coordinate names in 'ds' to the canonical 'latitude' / 'longitude' names.
+
+    A coordinate is identified as a latitude (or longitude) coordinate if either:
+    * its name is a known alias, e.g. 'lat', 'Lon' (matched case-insensitively), or
+    * its CF-convention metadata identifies it, i.e. a 'standard_name' attribute of
+      'latitude' / 'longitude' or a 'units' attribute like 'degrees_north' / 'degrees_east'.
+      This also covers datasets that name their geospatial coordinates 'x' / 'y'.
+
+    Coordinates already using canonical names are left untouched.
+    """
+
+    def _find_coordinate(canonical: str, aliases: t.FrozenSet[str], cf_units: t.FrozenSet[str]) -> t.Optional[str]:
+        for coord in ds.coords:
+            name = str(coord)
+            if name.lower() in aliases or name.lower() == canonical:
+                return name
+        for coord in ds.coords:
+            attrs = ds[coord].attrs
+            standard_name = str(attrs.get('standard_name', '')).strip().lower()
+            units = str(attrs.get('units', '')).strip().lower()
+            if standard_name == canonical or units in cf_units:
+                return str(coord)
+        return None
+
+    renames = {}
+    for canonical, aliases, cf_units in [('latitude', LATITUDE_ALIASES, CF_LATITUDE_UNITS),
+                                         ('longitude', LONGITUDE_ALIASES, CF_LONGITUDE_UNITS)]:
+        if canonical in ds.coords:
+            continue
+        found = _find_coordinate(canonical, aliases, cf_units)
+        if found is not None:
+            renames[found] = canonical
+    return renames
+
+
+def standardize_coordinate_names(ds: xr.Dataset) -> xr.Dataset:
+    """Renames latitude / longitude coordinate aliases (e.g. 'lat' / 'lon') to their canonical names.
+    """
+    renames = _coordinate_renames(ds)
+    if renames:
+        logger.info(f'Renaming coordinates to their canonical names: {renames}.')
+        ds = ds.rename(renames)
+    return ds
+
+
 def _preprocess_tif(
     ds: xr.Dataset,
     tif_metadata_for_start_time: str,

weather_mv/loader_pipeline/sinks_test.py

@@ -22,7 +22,7 @@
 import xarray as xr
 
 import weather_mv
-from .sinks import match_datetime, open_dataset
+from .sinks import match_datetime, open_dataset, standardize_coordinate_names
 
 
 class TestDataBase(unittest.TestCase):
@@ -126,6 +126,87 @@ def test_group_common_hypercubes(self):
             self.assertEqual(isinstance(ds, list), True)
 
 
+class StandardizeCoordinateNamesTest(unittest.TestCase):
+
+    def _dataset_with_coords(self, lat_name: str, lon_name: str,
+                             lat_attrs: dict = None, lon_attrs: dict = None) -> xr.Dataset:
+        return xr.Dataset(
+            {'temperature': ((lat_name, lon_name), np.zeros((3, 2)))},
+            coords={
+                lat_name: ((lat_name,), np.array([0.0, 1.0, 2.0]), lat_attrs or {}),
+                lon_name: ((lon_name,), np.array([10.0, 11.0]), lon_attrs or {}),
+            })
+
+    def assertCanonicalCoords(self, ds: xr.Dataset):
+        self.assertIn('latitude', ds.coords)
+        self.assertIn('longitude', ds.coords)
+
+    def test_renames_lat_lon(self):
+        ds = standardize_coordinate_names(self._dataset_with_coords('lat', 'lon'))
+        self.assertCanonicalCoords(ds)
+
+    def test_renames_long_alias(self):
+        ds = standardize_coordinate_names(self._dataset_with_coords('lat', 'long'))
+        self.assertCanonicalCoords(ds)
+
+    def test_renames_case_insensitively(self):
+        ds = standardize_coordinate_names(self._dataset_with_coords('Lat', 'LON'))
+        self.assertCanonicalCoords(ds)
+
+    def test_renames_capitalized_canonical_names(self):
+        ds = standardize_coordinate_names(self._dataset_with_coords('Latitude', 'Longitude'))
+        self.assertCanonicalCoords(ds)
+
+    def test_canonical_names_left_untouched(self):
+        original = self._dataset_with_coords('latitude', 'longitude')
+        ds = standardize_coordinate_names(original)
+        self.assertCanonicalCoords(ds)
+        xr.testing.assert_identical(original, ds)
+
+    def test_renames_x_y_with_cf_units(self):
+        ds = standardize_coordinate_names(self._dataset_with_coords(
+            'y', 'x', lat_attrs={'units': 'degrees_north'}, lon_attrs={'units': 'degrees_east'}))
+        self.assertCanonicalCoords(ds)
+
+    def test_renames_x_y_with_cf_standard_names(self):
+        ds = standardize_coordinate_names(self._dataset_with_coords(
+            'y', 'x', lat_attrs={'standard_name': 'latitude'}, lon_attrs={'standard_name': 'longitude'}))
+        self.assertCanonicalCoords(ds)
+
+    def test_renames_cf_metadata_case_insensitively_and_strips_whitespace(self):
+        ds = standardize_coordinate_names(self._dataset_with_coords(
+            'y', 'x',
+            lat_attrs={'standard_name': ' Latitude '},
+            lon_attrs={'units': ' Degrees_East '}))
+        self.assertCanonicalCoords(ds)
+
+    def test_x_y_without_cf_metadata_left_untouched(self):
+        # Bare 'x' / 'y' coordinates may be projected (non-degree) values; without
+        # CF metadata identifying them, they must not be relabeled.
+        ds = standardize_coordinate_names(self._dataset_with_coords('y', 'x'))
+        self.assertNotIn('latitude', ds.coords)
+        self.assertNotIn('longitude', ds.coords)
+
+    def test_alias_ignored_when_canonical_name_present(self):
+        ds = xr.Dataset(
+            {'temperature': (('latitude', 'longitude'), np.zeros((3, 2)))},
+            coords={
+                'latitude': np.array([0.0, 1.0, 2.0]),
+                'longitude': np.array([10.0, 11.0]),
+                'lat': (('latitude',), np.array([5.0, 6.0, 7.0])),
+            })
+        ds = standardize_coordinate_names(ds)
+        self.assertCanonicalCoords(ds)
+        self.assertIn('lat', ds.coords)
+
+    def test_dataset_without_spatial_coords_left_untouched(self):
+        original = xr.Dataset(
+            {'temperature': (('time',), np.zeros(3))},
+            coords={'time': np.array([0, 1, 2])})
+        ds = standardize_coordinate_names(original)
+        xr.testing.assert_identical(original, ds)
+
+
 class DatetimeTest(unittest.TestCase):
 
     def test_datetime_regex_string(self):