Extract forecast reference time

test/test_um2netcdf.py

@@ -1,6 +1,8 @@
 import unittest.mock as mock
 from dataclasses import dataclass
 from collections import namedtuple
+
+import cf_units
 from iris.exceptions import CoordinateNotFoundError
 import operator
 
@@ -406,10 +408,10 @@ def __init__(self, item_code, var_name=None, attributes=None,
         self.standard_name = None
         self.long_name = None
         self.data = None
+        self._coordinates = {}
 
-        # Mimic a coordinate dictionary keys for iris coordinate names. This
-        # ensures the access key for coord() matches the coordinate's name
-        self._coordinates = {c.name(): c for c in coords} if coords else {}
+        if coords:
+            self.update_coords(coords)
 
     def name(self):
         return self.var_name
@@ -421,6 +423,14 @@ def coord(self, name):
             msg = f"{self.__class__}: lacks coord for '{name}'"
             raise CoordinateNotFoundError(msg)
 
+    def update_coords(self, coords):
+        # Mimic a coordinate dictionary keys for iris coordinate names. This
+        # ensures the access key for coord() matches the coordinate's name
+        self._coordinates = {c.name(): c for c in coords} if coords else {}
+
+    def remove_coord(self, coord):
+        del self._coordinates[coord]
+
 
 def test_set_item_codes_avoid_overwrite():
     item_code = 1007
@@ -1177,3 +1187,89 @@ def test_convert_32_bit_with_float64(ua_plev_cube):
     ua_plev_cube.data = array
     um2nc.convert_32_bit(ua_plev_cube)
     assert ua_plev_cube.data.dtype == np.float32
+
+
+# fix forecast reference time tests
+@pytest.fixture
+def forecast_cube():
+    # NB: using a non-existent item code for fake forecast cube
+    return DummyCube(item_code=999)
+
+
+@pytest.fixture
+def time_points():
+    """Use for cube.coord('time').points attribute."""
+    return [-16382964.]
+
+
+@pytest.fixture
+def forecast_ref_time_coord(time_points):
+    # units & point data ripped from aiihca.paa1jan data file:
+    # cubes = iris.load("aiihca.paa1jan")
+    # cubes[0].long_name --> 'atmosphere_optical_thickness_due_to_sulphate_ambient_aerosol'
+    # cubes[0].coord("time").points --> array([-16382964.])
+    unit = cf_units.Unit(unit="hours since 1970-01-01 00:00:00")
+    assert unit.calendar == um2nc.STANDARD
+
+    return iris.coords.DimCoord(time_points,
+                                standard_name=um2nc.FORECAST_REFERENCE_TIME,
+                                units=unit)
+
+
+@pytest.fixture
+def time_coord(time_points):
+    # units data ripped from aiihca data file
+    unit = cf_units.Unit(unit="hours since 1970-01-01 00:00:00",
+                         calendar=um2nc.GREGORIAN)
+    assert unit.calendar == um2nc.STANDARD
+
+    return iris.coords.DimCoord(time_points,
+                                standard_name=um2nc.TIME,
+                                units=unit)
+
+
+def test_fix_forecast_reference_time_exit_on_missing_ref_time(forecast_cube):
+    with pytest.raises(iris.exceptions.CoordinateNotFoundError):
+        forecast_cube.coord(um2nc.FORECAST_REFERENCE_TIME)
+
+    assert um2nc.fix_forecast_reference_time(forecast_cube) is None
+
+
+def test_fix_forecast_reference_time_exit_on_missing_time(forecast_cube,
+                                                          forecast_ref_time_coord):
+    forecast_cube.update_coords([forecast_ref_time_coord])
+
+    with pytest.raises(iris.exceptions.CoordinateNotFoundError):
+        forecast_cube.coord(um2nc.TIME)
+
+    assert um2nc.fix_forecast_reference_time(forecast_cube) is None
+
+
+def test_fix_forecast_reference_time_standard(forecast_cube,
+                                              forecast_ref_time_coord,
+                                              time_coord):
+
+    forecast_period = iris.coords.DimCoord([372.0],
+                                           standard_name=um2nc.FORECAST_PERIOD)
+
+    forecast_cube.update_coords([forecast_ref_time_coord,
+                                 time_coord,
+                                 forecast_period])
+
+    assert um2nc.fix_forecast_reference_time(forecast_cube) is None
+
+
+@pytest.mark.skip
+def test_fix_forecast_reference_time_gregorian(forecast_cube,
+                                               forecast_ref_time_coord,
+                                               time_coord):
+    msg = "Is time.units.calendar == 'gregorian' branch & testing required?"
+    raise NotImplementedError(msg)
+
+
+@pytest.mark.skip
+def test_fix_forecast_reference_time_proleptic_gregorian(forecast_cube,
+                                                         forecast_ref_time_coord,
+                                                         time_coord):
+    msg = "Is time.units.calendar == 'proleptic_gregorian' branch & testing required?"
+    raise NotImplementedError(msg)

umpost/um2netcdf.py

@@ -70,6 +70,14 @@
 LEVEL_HEIGHT = "level_height"
 SIGMA = "sigma"
 
+# forecast reference time constants
+FORECAST_REFERENCE_TIME = "forecast_reference_time"
+TIME = "time"
+PROLEPTIC_GREGORIAN = "proleptic_gregorian"
+GREGORIAN = "gregorian"
+FORECAST_PERIOD = "forecast_period"
+STANDARD = "standard"
+
 
 class PostProcessingError(Exception):
     """Generic class for um2nc specific errors."""
@@ -102,34 +110,6 @@ def pg_calendar(self):
 PPField.calendar = pg_calendar
 
 
-# TODO: rename time to avoid clash with builtin time module
-def convert_proleptic(time):
-    # Convert units from hours to days and shift origin from 1970 to 0001
-    newunits = cf_units.Unit("days since 0001-01-01 00:00", calendar='proleptic_gregorian')
-    tvals = np.array(time.points)  # Need a copy because can't assign to time.points[i]
-    tbnds = np.array(time.bounds) if time.bounds is not None else None
-
-    for i in range(len(time.points)):
-        date = time.units.num2date(tvals[i])
-        newdate = cftime.DatetimeProlepticGregorian(date.year, date.month, date.day,
-                                                    date.hour, date.minute, date.second)
-        tvals[i] = newunits.date2num(newdate)
-
-        if tbnds is not None:  # Fields with instantaneous data don't have bounds
-            for j in range(2):
-                date = time.units.num2date(tbnds[i][j])
-                newdate = cftime.DatetimeProlepticGregorian(date.year, date.month, date.day,
-                                                            date.hour, date.minute, date.second)
-                tbnds[i][j] = newunits.date2num(newdate)
-
-    time.points = tvals
-
-    if tbnds is not None:
-        time.bounds = tbnds
-
-    time.units = newunits
-
-
 def fix_lat_coord_name(lat_coordinate, grid_type, dlat):
     """
     Add a 'var_name' attribute to a latitude coordinate object
@@ -338,33 +318,7 @@ def cubewrite(cube, sman, compression, use64bit, verbose):
 
     cube.attributes['missing_value'] = np.array([fill_value], cube.data.dtype)
 
-    # If reference date is before 1600 use proleptic gregorian
-    # calendar and change units from hours to days
-    try:
-        reftime = cube.coord('forecast_reference_time')
-        time = cube.coord('time')
-        refdate = reftime.units.num2date(reftime.points[0])
-        assert time.units.origin == 'hours since 1970-01-01 00:00:00'
-
-        if time.units.calendar == 'proleptic_gregorian' and refdate.year < 1600:
-            convert_proleptic(time)
-        else:
-            if time.units.calendar == 'gregorian':
-                new_calendar = 'proleptic_gregorian'
-            else:
-                new_calendar = time.units.calendar
-
-            time.units = cf_units.Unit("days since 1970-01-01 00:00", calendar=new_calendar)
-            time.points = time.points/24.
-
-            if time.bounds is not None:
-                time.bounds = time.bounds/24.
-
-        cube.remove_coord('forecast_period')
-        cube.remove_coord('forecast_reference_time')
-    except iris.exceptions.CoordinateNotFoundError:
-        # Dump files don't have forecast_reference_time
-        pass
+    fix_forecast_reference_time(cube)
 
     # Check whether any of the coordinates is a pseudo-dimension with integer values and
     # if so, reset to int32 to prevent problems with possible later conversion to netCDF3
@@ -406,6 +360,64 @@ def cubewrite(cube, sman, compression, use64bit, verbose):
         sman.write(cube, zlib=True, complevel=compression, fill_value=fill_value)
 
 
+def fix_forecast_reference_time(cube):
+    # If reference date is before 1600 use proleptic gregorian
+    # calendar and change units from hours to days
+    try:
+        reftime = cube.coord('forecast_reference_time')
+        time = cube.coord('time')
+        refdate = reftime.units.num2date(reftime.points[0])
+        assert time.units.origin == 'hours since 1970-01-01 00:00:00'
+
+        if time.units.calendar == 'proleptic_gregorian' and refdate.year < 1600:
+            convert_proleptic(time)
+        else:
+            if time.units.calendar == 'gregorian':
+                new_calendar = 'proleptic_gregorian'
+            else:
+                new_calendar = time.units.calendar
+
+            time.units = cf_units.Unit("days since 1970-01-01 00:00", calendar=new_calendar)
+            time.points = time.points / 24.
+
+            if time.bounds is not None:
+                time.bounds = time.bounds / 24.
+
+        cube.remove_coord('forecast_period')
+        cube.remove_coord('forecast_reference_time')
+    except iris.exceptions.CoordinateNotFoundError:
+        # Dump files don't have forecast_reference_time
+        pass
+
+
+# TODO: rename time to avoid clash with builtin time module
+def convert_proleptic(time):
+    # Convert units from hours to days and shift origin from 1970 to 0001
+    newunits = cf_units.Unit("days since 0001-01-01 00:00", calendar='proleptic_gregorian')
+    tvals = np.array(time.points)  # Need a copy because can't assign to time.points[i]
+    tbnds = np.array(time.bounds) if time.bounds is not None else None
+
+    for i in range(len(time.points)):
+        date = time.units.num2date(tvals[i])
+        newdate = cftime.DatetimeProlepticGregorian(date.year, date.month, date.day,
+                                                    date.hour, date.minute, date.second)
+        tvals[i] = newunits.date2num(newdate)
+
+        if tbnds is not None:  # Fields with instantaneous data don't have bounds
+            for j in range(2):
+                date = time.units.num2date(tbnds[i][j])
+                newdate = cftime.DatetimeProlepticGregorian(date.year, date.month, date.day,
+                                                            date.hour, date.minute, date.second)
+                tbnds[i][j] = newunits.date2num(newdate)
+
+    time.points = tvals
+
+    if tbnds is not None:
+        time.bounds = tbnds
+
+    time.units = newunits
+
+
 def fix_cell_methods(cell_methods):
     """
     Removes misleading 'hour' from interval naming, leaving other names intact.