Merge pull request #410 from hollyhan/add_ismip6_forcing_process_RACMO

Add ismip6 forcing testgroup

compass/landice/__init__.py

@@ -8,6 +8,7 @@
 from compass.landice.tests.greenland import Greenland
 from compass.landice.tests.humboldt import Humboldt
 from compass.landice.tests.hydro_radial import HydroRadial
+from compass.landice.tests.ismip6_forcing import Ismip6Forcing
 from compass.landice.tests.kangerlussuaq import Kangerlussuaq
 from compass.landice.tests.koge_bugt_s import KogeBugtS
 from compass.landice.tests.mismipplus import MISMIPplus
@@ -34,6 +35,7 @@ def __init__(self):
         self.add_test_group(Greenland(mpas_core=self))
         self.add_test_group(Humboldt(mpas_core=self))
         self.add_test_group(HydroRadial(mpas_core=self))
+        self.add_test_group(Ismip6Forcing(mpas_core=self))
         self.add_test_group(Kangerlussuaq(mpas_core=self))
         self.add_test_group(KogeBugtS(mpas_core=self))
         self.add_test_group(MISMIPplus(mpas_core=self))

compass/landice/tests/ismip6_forcing/__init__.py

@@ -0,0 +1,26 @@
+from compass.testgroup import TestGroup
+from compass.landice.tests.ismip6_forcing.atmosphere import Atmosphere
+from compass.landice.tests.ismip6_forcing.ocean_thermal import OceanThermal
+from compass.landice.tests.ismip6_forcing.ocean_basal import OceanBasal
+
+
+class Ismip6Forcing(TestGroup):
+    """
+    A test group for processing the ISMIP6 ocean and atmosphere forcing data
+    """
+
+    def __init__(self, mpas_core):
+        """
+        Create the test group
+
+        Parameters
+        ----------
+        mpas_core : compass.landice.Landice
+            the MPAS core that this test group belongs to
+        """
+        super().__init__(mpas_core=mpas_core, name="ismip6_forcing")
+
+        self.add_test_case(Atmosphere(test_group=self))
+        self.add_test_case(OceanBasal(test_group=self))
+        self.add_test_case(OceanThermal(test_group=self, process_obs=True))
+        self.add_test_case(OceanThermal(test_group=self, process_obs=False))

compass/landice/tests/ismip6_forcing/atmosphere/__init__.py

@@ -0,0 +1,42 @@
+from compass.testcase import TestCase
+from compass.landice.tests.ismip6_forcing.atmosphere.process_smb \
+    import ProcessSMB
+from compass.landice.tests.ismip6_forcing.atmosphere.process_smb_racmo \
+    import ProcessSmbRacmo
+from compass.landice.tests.ismip6_forcing.configure import configure as \
+    configure_testgroup
+
+
+class Atmosphere(TestCase):
+    """
+    A test case for processing the ISMIP6 atmosphere forcing data.
+    The test case builds a mapping file for interpolation between the
+    ISMIP6 8km polarstereo grid and MALI mesh, regrids the forcing data
+    and rename the ISMIP6 variables to corresponding MALI variables.
+    """
+
+    def __init__(self, test_group):
+        """
+        Create the test case
+
+        Parameters
+        ----------
+        test_group : compass.landice.tests.ismip6_forcing.Ismip6Forcing
+            The test group that this test case belongs to
+        """
+        name = "atmosphere"
+        subdir = name
+        super().__init__(test_group=test_group, name=name, subdir=subdir)
+
+        step = ProcessSmbRacmo(test_case=self)
+        self.add_step(step)
+        step = ProcessSMB(test_case=self)
+        self.add_step(step)
+
+    def configure(self):
+        """
+        Configures test case
+        """
+
+        configure_testgroup(config=self.config,
+                            check_model_options=True)

compass/landice/tests/ismip6_forcing/atmosphere/create_mapfile_smb.py

@@ -0,0 +1,204 @@
+import os
+import numpy as np
+import shutil
+import netCDF4
+import xarray as xr
+from mpas_tools.scrip.from_mpas import scrip_from_mpas
+from mpas_tools.logging import check_call
+from pyremap.descriptor import interp_extrap_corners_2d
+
+
+def build_mapping_file(config, cores, logger, ismip6_grid_file,
+                       mapping_file, mali_mesh_file=None, method_remap=None):
+    """
+    Build a mapping file if it does not exist.
+    Mapping file is then used to remap the ismip6 source file in polarstero
+    coordinate to unstructured mali mesh
+
+    Parameters
+    ----------
+    config : compass.config.CompassConfigParser
+        Configuration options for a ismip6 forcing test case
+
+    cores : int
+        the number of cores for the ESMF_RegridWeightGen
+
+    logger : logging.Logger
+        A logger for output from the step
+
+    ismip6_grid_file : str
+        ismip6 grid file
+
+    mapping_file : str
+        weights for interpolation from ismip6_grid_file to mali_mesh_file
+
+    mali_mesh_file : str, optional
+        The MALI mesh file is used if mapping file does not exist
+
+    method_remap : str, optional
+        Remapping method used in building a mapping file
+    """
+
+    if os.path.exists(mapping_file):
+        logger.info(f"Mapping file exists. Not building a new one.")
+        return
+
+    # create the scrip files if mapping file does not exist
+    logger.info(f"Mapping file does not exist. Building one based on the"
+                f" input/output meshes")
+    logger.info(f"Creating temporary scrip files for source and "
+                f"destination grids...")
+
+    if mali_mesh_file is None:
+        raise ValueError("Mapping file does not exist. To build one, Mali "
+                         "mesh file with '-f' should be provided. "
+                         "Type --help for info")
+
+    # name temporary scrip files that will be used to build mapping file
+    source_grid_scripfile = "temp_source_scrip.nc"
+    mali_scripfile = "temp_mali_scrip.nc"
+    # this is the projection of ismip6 data for Antarctica
+    ismip6_projection = "ais-bedmap2"
+
+    # create a scripfile for the atmosphere forcing data
+    create_atm_scrip(ismip6_grid_file, source_grid_scripfile)
+
+    # create a MALI mesh scripfile
+    # make sure the mali mesh file uses the longitude convention of [0 2pi]
+    # make changes on a duplicated file to avoid making changes to the
+    # original mesh file
+
+    mali_mesh_copy = f"{mali_mesh_file}_copy"
+    shutil.copy(mali_mesh_file, f"{mali_mesh_file}_copy")
+
+    args = ["set_lat_lon_fields_in_planar_grid.py",
+            "--file", mali_mesh_copy,
+            "--proj", ismip6_projection]
+
+    check_call(args, logger=logger)
+
+    # create a MALI mesh scripfile if mapping file does not exist
+    scrip_from_mpas(mali_mesh_copy, mali_scripfile)
+
+    # create a mapping file using ESMF weight gen
+    print(f"Creating a mapping file. Mapping method used: {method_remap}")
+
+    if method_remap is None:
+        raise ValueError("Desired remapping option should be provided with "
+                         "--method. Available options are 'bilinear',"
+                         "'neareststod', 'conserve'.")
+
+    parallel_executable = config.get('parallel', 'parallel_executable')
+    # split the parallel executable into constituents in case it includes flags
+    args = parallel_executable.split(' ')
+    args.extend(["-n", f"{cores}",
+                 "ESMF_RegridWeightGen",
+                 "-s", source_grid_scripfile,
+                 "-d", mali_scripfile,
+                 "-w", mapping_file,
+                 "-m", method_remap,
+                 "-i", "-64bit_offset",
+                 "--dst_regional", "--src_regional"])
+
+    check_call(args, logger=logger)
+
+    # remove the temporary scripfiles once the mapping file is generated
+    logger.info(f"Removing the temporary mesh and scripfiles...")
+    os.remove(source_grid_scripfile)
+    os.remove(mali_scripfile)
+    os.remove(mali_mesh_copy)
+
+
+def create_atm_scrip(source_grid_file, source_grid_scripfile):
+    """
+    create a scripfile for the ismip6 atmospheric forcing data.
+    Note: the atmospheric forcing data do not have 'x' and 'y' coordinates and
+    only have dimensions of them. This function uses 'lat' and 'lon'
+    coordinates to generate a scripfile.
+
+    Parameters
+    ----------
+    source_grid_file : str
+        input smb grid file
+
+    source_grid_scripfile : str
+        output scrip file of the input smb data
+    """
+
+    ds = xr.open_dataset(source_grid_file)
+    out_file = netCDF4.Dataset(source_grid_scripfile, 'w')
+
+    if "rlon" in ds and "rlat" in ds:  # this is for RACMO's rotated-pole grid
+        nx = ds.sizes["rlon"]
+        ny = ds.sizes["rlat"]
+    else:
+        nx = ds.sizes["x"]
+        ny = ds.sizes["y"]
+    units = "degrees"
+
+    grid_size = nx * ny
+
+    out_file.createDimension("grid_size", grid_size)
+    out_file.createDimension("grid_corners", 4)
+    out_file.createDimension("grid_rank", 2)
+
+    # Variables
+    grid_center_lat = out_file.createVariable("grid_center_lat", "f8",
+                                              ("grid_size",))
+    grid_center_lat.units = units
+    grid_center_lon = out_file.createVariable("grid_center_lon", "f8",
+                                              ("grid_size",))
+    grid_center_lon.units = units
+    grid_corner_lat = out_file.createVariable("grid_corner_lat", "f8",
+                                              ("grid_size", "grid_corners"))
+    grid_corner_lat.units = units
+    grid_corner_lon = out_file.createVariable("grid_corner_lon", "f8",
+                                              ("grid_size", "grid_corners"))
+    grid_corner_lon.units = units
+    grid_imask = out_file.createVariable("grid_imask", "i4", ("grid_size",))
+    grid_imask.units = "unitless"
+    out_file.createVariable("grid_dims", "i4", ("grid_rank",))
+
+    out_file.variables["grid_center_lat"][:] = ds.lat.values.flat
+    out_file.variables["grid_center_lon"][:] = ds.lon.values.flat
+    out_file.variables["grid_dims"][:] = [nx, ny]
+    out_file.variables["grid_imask"][:] = 1
+
+    if "lat_bnds" in ds and "lon_bnds" in ds:
+        lat_corner = ds.lat_bnds
+        if "time" in lat_corner.dims:
+            lat_corner = lat_corner.isel(time=0)
+
+        lon_corner = ds.lon_bnds
+        if "time" in lon_corner.dims:
+            lon_corner = lon_corner.isel(time=0)
+
+        lat_corner = lat_corner.values
+        lon_corner = lon_corner.values
+    else:
+        # this part is used for RACMO as it does not have lat_bnds & lon_bnds
+        lat_corner = _unwrap_corners(interp_extrap_corners_2d(ds.lat.values))
+        lon_corner = _unwrap_corners(interp_extrap_corners_2d(ds.lon.values))
+
+    grid_corner_lat = lat_corner.reshape((grid_size, 4))
+    grid_corner_lon = lon_corner.reshape((grid_size, 4))
+
+    out_file.variables["grid_corner_lat"][:] = grid_corner_lat
+    out_file.variables["grid_corner_lon"][:] = grid_corner_lon
+
+    out_file.close()
+
+
+def _unwrap_corners(in_field):
+    """
+    Turn a 2D array of corners into an array of rectangular mesh elements
+    """
+    out_field = np.zeros(((in_field.shape[0] - 1) *
+                          (in_field.shape[1] - 1), 4))
+    # corners are counterclockwise
+    out_field[:, 0] = in_field[0:-1, 0:-1].flat
+    out_field[:, 1] = in_field[0:-1, 1:].flat
+    out_field[:, 2] = in_field[1:, 1:].flat
+    out_field[:, 3] = in_field[1:, 0:-1].flat
+
+    return out_field