Extract forecast reference time

test/test_um2netcdf.py

@@ -1,6 +1,9 @@
 import unittest.mock as mock
+import warnings
 from dataclasses import dataclass
 from collections import namedtuple
+
+import cf_units
 from iris.exceptions import CoordinateNotFoundError
 import operator
 
@@ -79,6 +82,7 @@ class DummyStash:
     item: int
 
 
+# TODO: would making this a dataclass provide any benefit?
 class DummyCube:
     """
     Imitation iris Cube for unit testing.
@@ -116,6 +120,20 @@ def coord(self, _name):
             msg = f"{self.__class__}[{self.var_name}]: lacks coord for '{_name}'"
             raise CoordinateNotFoundError(msg)
 
+    def remove_coord(self, coord):
+        del self._coordinates[coord]
+
+    # Auxiliary methods
+    # These methods DO NOT exist in the iris cube API, these are helper methods
+    # to configure DummyCubes for testing.
+    #
+    # All methods should see an `aux_` prefix
+
+    def aux_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 {}
+
 
 # NB: these cube fixtures have been chosen to mimic cubes for testing key parts
 # of the process() workflow. Some cubes require pressure level masking with the
@@ -1107,3 +1125,109 @@ 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_points():
+    """Use for cube.coord('forecast_reference_time').points attribute."""
+    return [-16383336.]
+
+
+# FIXME: unit.calendar needs updating as per:
+#  https://github.com/ACCESS-NRI/um2nc-standalone/pull/118#issuecomment-2404034473
+@pytest.fixture
+def forecast_ref_time_coord(forecast_ref_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 == cf_units.CALENDAR_STANDARD
+
+    return iris.coords.DimCoord(forecast_ref_time_points,
+                                standard_name=um2nc.FORECAST_REFERENCE_TIME,
+                                units=unit)
+
+
+# FIXME: unit.calendar needs updating as per:
+#  https://github.com/ACCESS-NRI/um2nc-standalone/pull/118#issuecomment-2404034473
+@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=cf_units.CALENDAR_GREGORIAN)
+    assert unit.calendar == cf_units.CALENDAR_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):
+    # verify fix_forecast_ref_time() exits early if the coord is missing
+    with pytest.raises(iris.exceptions.CoordinateNotFoundError):
+        forecast_cube.coord(um2nc.FORECAST_REFERENCE_TIME)
+
+    # TODO: fails to assert the cube is unmodified
+    # TODO: is this test even needed?
+    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):
+    # verify fix_forecast_ref_time() exits early if the coord is missing
+    forecast_cube.aux_update_coords([forecast_ref_time_coord])
+
+    with pytest.raises(iris.exceptions.CoordinateNotFoundError):
+        forecast_cube.coord(um2nc.TIME)
+
+    # TODO: fails to assert the cube is unmodified
+    # TODO: is this test even needed?
+    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):
+    # TODO: executes part of the ref time fix code
+    # TODO: needs to assert the changes!
+    forecast_period = iris.coords.DimCoord([372.0],
+                                           standard_name=um2nc.FORECAST_PERIOD)
+
+    forecast_cube.aux_update_coords([forecast_ref_time_coord,
+                                     time_coord,
+                                     forecast_period])
+
+    assert um2nc.fix_forecast_reference_time(forecast_cube) is None
+
+    # TODO: add assertions here
+    warnings.warn("test_fix_forecast_reference_time_standard asserts nothing")
+
+
+@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

@@ -78,6 +78,11 @@
 SIGMA_THETA = "sigma_theta"
 SIGMA_RHO = "sigma_rho"
 
+# forecast reference time constants
+FORECAST_REFERENCE_TIME = "forecast_reference_time"
+FORECAST_PERIOD = "forecast_period"
+TIME = "time"
+
 
 class PostProcessingError(Exception):
     """Generic class for um2nc specific errors."""
@@ -110,34 +115,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
@@ -346,33 +323,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
@@ -414,6 +365,90 @@ def cubewrite(cube, sman, compression, use64bit, verbose):
         sman.write(cube, zlib=True, complevel=compression, fill_value=fill_value)
 
 
+# TODO: this review https://github.com/ACCESS-NRI/um2nc-standalone/pull/118
+#     indicates these time functions can be simplified. Suggestion: modularise
+#     functionality to simplify logic:
+#   * Extract origin time shift to a function?
+#   * Extract hours to days conversion to a function?
+def fix_forecast_reference_time(cube):
+    # If reference date is before 1600 use proleptic gregorian
+    # calendar and change units from hours to days
+
+    # TODO: constrain the exception handler to the first 2 coord lookups?
+    # TODO: detect dump files & exit early (look up all coords early)
+    try:
+        reftime = cube.coord(FORECAST_REFERENCE_TIME)
+        time = cube.coord(TIME)
+        refdate = reftime.units.num2date(reftime.points[0])
+
+        # TODO: replace with `if` check to prevent assert vanishing in optimised Python mode
+        assert time.units.origin == 'hours since 1970-01-01 00:00:00'
+
+        # TODO: add reference year as a configurable arg?
+        # TODO: determine if the start year should be changed from 1600
+        #  see https://github.com/ACCESS-NRI/um2nc-standalone/pull/118/files#r1792886613
+        if time.units.calendar == cf_units.CALENDAR_PROLEPTIC_GREGORIAN and refdate.year < 1600:
+            convert_proleptic(time)
+        else:
+            if time.units.calendar == cf_units.CALENDAR_GREGORIAN:
+                new_calendar = cf_units.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.
+
+        # TODO: remove_coord() calls the coord() lookup, which raises
+        #       CoordinateNotFoundError if the forecast period is missing, this
+        #       ties remove_coords() to the exception handler below
+        #
+        # TODO: if removing the forecast period fails, forecast reference time is
+        #       NOT removed. What is the desired behaviour?
+        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 arg to pre-emptively avoid a clash with builtin time module
+# TODO: refactor to add a calendar arg here, see
+#       https://github.com/ACCESS-NRI/um2nc-standalone/pull/118/files#r1794321677
+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=cf_units.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
+
+    # TODO: refactor manual looping with pythonic iteration
+    #
+    # TODO: limit the looping as per this review suggestion:
+    #  https://github.com/ACCESS-NRI/um2nc-standalone/pull/118/files#r1794315128
+    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.