Produce a CMIP7 historical nitrogen forcing ancillary file for ACCESS-ESM1.6

CMIP7/esm1p6/atmosphere/cmip7_HI.py

@@ -1,8 +1,13 @@
 import calendar
 from pathlib import Path
 
+import iris
 import mule
-from cmip7_ancil_constants import ANCIL_TODAY, UM_VERSION
+from cmip7_ancil_constants import (
+    ANCIL_TODAY,
+    MONTHS_IN_A_YEAR,
+    UM_VERSION,
+)
 from cmip7_PI import CMIP7_PI_YEAR
 
 CMIP7_HI_BEG_YEAR = CMIP7_PI_YEAR
@@ -41,3 +46,39 @@ def fix_esm15_hi_ancil_date(ifile, ofile):
         field.lbdatd = calendar.monthrange(field.lbyr, field.lbmon)[1]
 
     ff.to_file(ofile)
+
+
+def extend_hi_years(cube):
+    """
+    Extend a cube representing a monthly time series by duplicating
+    and adjusting the first and last years.
+    Based on Crown copyright code from ozone_cmip6_ancillary_for_suite.py
+    by Steven Hardiman of the UK Met Office.
+    """
+    time_coord = cube.coord("time")
+    time_points = time_coord.points
+    if (months := len(time_points)) <= MONTHS_IN_A_YEAR:
+        raise ValueError(
+            f"Cannot extend a cube containing {months} months of data. "
+            f"More than {MONTHS_IN_A_YEAR} months are required."
+        )
+
+    # Duplicate the first year.
+    length_one_year = time_points[MONTHS_IN_A_YEAR] - time_points[0]
+    beg_year = cube[:MONTHS_IN_A_YEAR].copy()
+    beg_year_tc = beg_year.coord("time")
+    beg_year_tc.points = beg_year_tc.points - length_one_year
+    if time_coord.has_bounds():
+        beg_year_tc.bounds = beg_year_tc.bounds - length_one_year
+
+    # Duplicate the last year.
+    length_one_year = time_points[-1] - time_points[-1 - MONTHS_IN_A_YEAR]
+    end_year = cube[-MONTHS_IN_A_YEAR:].copy()
+    end_year_tc = end_year.coord("time")
+    end_year_tc.points = end_year_tc.points + length_one_year
+    if time_coord.has_bounds():
+        end_year_tc.bounds = end_year_tc.bounds + length_one_year
+
+    # Return a cube with extended years.
+    cubelist = iris.cube.CubeList((beg_year, cube, end_year))
+    return cubelist.concatenate_cube()

CMIP7/esm1p6/atmosphere/nitrogen/cmip7_HI_nitrogen_generate.py

@@ -1,6 +1,15 @@
 from argparse import ArgumentParser
+from pathlib import Path
 
 from cmip7_ancil_argparse import common_parser
+from cmip7_ancil_constants import ANCIL_TODAY
+from cmip7_HI import extend_hi_years
+from nitrogen.cmip7_nitrogen import (
+    cmip7_nitrogen_dirpath,
+    load_cmip7_nitrogen,
+    regrid_cmip7_nitrogen,
+    save_cmip7_nitrogen,
+)
 
 
 def parse_args():
@@ -11,10 +20,40 @@ def parse_args():
             "Generate input files from CMIP7 historical nitrogen forcings"
         ),
     )
-    parser.add_argument("--dataset-date-range-list")
+    parser.add_argument("--dataset-date-range")
     parser.add_argument("--save-filename")
     return parser.parse_args()
 
 
+def cmip7_hi_nitrogen_filepath(args, species):
+    dirpath = cmip7_nitrogen_dirpath(args, "mon", species)
+    filename = (
+        f"{species}_input4MIPs_surfaceFluxes_CMIP_"
+        f"{args.dataset_version}_gn_"
+        f"{args.dataset_date_range}.nc"
+    )
+    return dirpath / filename
+
+
+def esm_hi_nitrogen_save_dirpath(args):
+    return (
+        Path(args.ancil_target_dirname)
+        / "modern"
+        / "historical"
+        / "atmosphere"
+        / "land"
+        / "biogeochemistry"
+        / args.esm_grid_rel_dirname
+        / ANCIL_TODAY
+    )
+
+
 if __name__ == "__main__":
     args = parse_args()
+
+    # Load the CMIP7 datasets
+    nitrogen_cube = load_cmip7_nitrogen(args, cmip7_hi_nitrogen_filepath)
+    # Regrid to match the ESM1.5 mask and extend the time series
+    esm_cube = extend_hi_years(regrid_cmip7_nitrogen(args, nitrogen_cube))
+    # Save the ancillary
+    save_cmip7_nitrogen(args, esm_cube, esm_hi_nitrogen_save_dirpath)

CMIP7/esm1p6/atmosphere/nitrogen/cmip7_PI_nitrogen_generate.py

@@ -1,20 +1,13 @@
 from argparse import ArgumentParser
 from pathlib import Path
 
-import iris
 from cmip7_ancil_argparse import common_parser
-from cmip7_ancil_common import (
-    INTERPOLATION_SCHEME,
-    esm_grid_mask_cube,
-    fix_coords,
-    save_ancil,
-)
 from cmip7_ancil_constants import ANCIL_TODAY
-from iris.util import equalise_attributes
 from nitrogen.cmip7_nitrogen import (
-    NITROGEN_SPECIES,
-    NITROGEN_STASH_ITEM,
     cmip7_nitrogen_dirpath,
+    load_cmip7_nitrogen,
+    regrid_cmip7_nitrogen,
+    save_cmip7_nitrogen,
 )
 
 
@@ -41,38 +34,6 @@ def cmip7_pi_nitrogen_filepath(args, species):
     return dirpath / filename
 
 
-def load_cmip7_pi_nitrogen_species(args, species):
-    species_filepath = cmip7_pi_nitrogen_filepath(args, species)
-    cube = iris.load_cube(species_filepath)
-    return cube
-
-
-def load_cmip7_pi_nitrogen(args):
-    # Load all of the PI nitrogen datasets into a CubeList
-    nitrogen_cubes = iris.cube.CubeList()
-    for species in NITROGEN_SPECIES:
-        species_cube = load_cmip7_pi_nitrogen_species(args, species)
-        nitrogen_cubes.append(species_cube)
-    # Remove all attributes that differ between cubes
-    equalise_attributes(nitrogen_cubes)
-    # Add the cubes together
-    cube_tot = nitrogen_cubes[0]
-    for cube in nitrogen_cubes[1:]:
-        cube_tot += cube
-    # Change the units from kg m-2 s-1 to g m-2 day-1
-    cube_tot.convert_units("g m-2 day-1")
-    return cube_tot
-
-
-def regrid_cmip7_pi_nitrogen(args, cube):
-    # Make the coordinates comaptible with the ESM1.5 grid mask
-    fix_coords(args, cube)
-    # Regrid using the ESM1.5 grid mask
-    esm_cube = cube.regrid(esm_grid_mask_cube(args), INTERPOLATION_SCHEME)
-    esm_cube.data = esm_cube.data.filled(0.0)
-    return esm_cube
-
-
 def esm_pi_nitrogen_save_dirpath(args):
     return (
         Path(args.ancil_target_dirname)
@@ -86,22 +47,12 @@ def esm_pi_nitrogen_save_dirpath(args):
     )
 
 
-def save_cmip7_pi_nitrogen(args, cube):
-    # Add STASH metadata
-    cube.attributes["STASH"] = iris.fileformats.pp.STASH(
-        model=1, section=0, item=NITROGEN_STASH_ITEM
-    )
-    # Save as an ancillary file
-    save_dirpath = esm_pi_nitrogen_save_dirpath(args)
-    save_ancil(cube, save_dirpath, args.save_filename)
-
-
 if __name__ == "__main__":
     args = parse_args()
 
     # Load the CMIP7 datasets
-    nitrogen_cube = load_cmip7_pi_nitrogen(args)
+    nitrogen_cube = load_cmip7_nitrogen(args, cmip7_pi_nitrogen_filepath)
     # Regrid to match the ESM1.5 mask
-    esm_cube = regrid_cmip7_pi_nitrogen(args, nitrogen_cube)
+    esm_cube = regrid_cmip7_nitrogen(args, nitrogen_cube)
     # Save the ancillary
-    save_cmip7_pi_nitrogen(args, esm_cube)
+    save_cmip7_nitrogen(args, esm_cube, esm_pi_nitrogen_save_dirpath)

CMIP7/esm1p6/atmosphere/nitrogen/cmip7_nitrogen.py

@@ -1,5 +1,14 @@
 from pathlib import Path
 
+import iris
+from cmip7_ancil_common import (
+    INTERPOLATION_SCHEME,
+    esm_grid_mask_cube,
+    fix_coords,
+    save_ancil,
+)
+from iris.util import equalise_attributes
+
 NITROGEN_SPECIES = ("drynhx", "drynoy", "wetnhx", "wetnoy")
 """
 The STASH code to use is m01s00i884 NITROGEN DEPOSITION as per
@@ -19,3 +28,40 @@ def cmip7_nitrogen_dirpath(args, period, species):
         / "gn"
         / args.dataset_vdate
     )
+
+
+def load_cmip7_nitrogen(args, load_filepath_fn):
+    # Load all of the PI nitrogen datasets into a CubeList
+    nitrogen_cubes = iris.cube.CubeList()
+    for species in NITROGEN_SPECIES:
+        species_filepath = load_filepath_fn(args, species)
+        species_cube = iris.load_cube(species_filepath)
+        nitrogen_cubes.append(species_cube)
+    # Remove all attributes that differ between cubes
+    equalise_attributes(nitrogen_cubes)
+    # Add the cubes together
+    cube_tot = nitrogen_cubes[0]
+    for cube in nitrogen_cubes[1:]:
+        cube_tot += cube
+    # Change the units from kg m-2 s-1 to g m-2 day-1
+    cube_tot.convert_units("g m-2 day-1")
+    return cube_tot
+
+
+def regrid_cmip7_nitrogen(args, cube):
+    # Make the coordinates comaptible with the ESM1.5 grid mask
+    fix_coords(args, cube)
+    # Regrid using the ESM1.5 grid mask
+    esm_cube = cube.regrid(esm_grid_mask_cube(args), INTERPOLATION_SCHEME)
+    esm_cube.data = esm_cube.data.filled(0.0)
+    return esm_cube
+
+
+def save_cmip7_nitrogen(args, cube, save_dirpath_fn):
+    # Add STASH metadata
+    cube.attributes["STASH"] = iris.fileformats.pp.STASH(
+        model=1, section=0, item=NITROGEN_STASH_ITEM
+    )
+    # Save as an ancillary file
+    save_dirpath = save_dirpath_fn(args)
+    save_ancil(cube, save_dirpath, args.save_filename)