Split ESMF weight generation into dedicated steps for ISMIP7 forcing

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

@@ -1,3 +1,6 @@
+from compass.landice.tests.ismip7_forcing.atmosphere.build_mapping_file import (  # noqa: E501
+    BuildMappingFile,
+)
 from compass.landice.tests.ismip7_forcing.atmosphere.process_runoff import (
     ProcessRunoff,
 )
@@ -40,6 +43,7 @@ def __init__(self, test_group):
         subdir = name
         super().__init__(test_group=test_group, name=name, subdir=subdir)
 
+        self.add_step(BuildMappingFile(test_case=self))
         self.add_step(ProcessSmb(test_case=self))
         self.add_step(ProcessTemperature(test_case=self))
         self.add_step(ProcessSmbGradient(test_case=self))

compass/landice/tests/ismip7_forcing/atmosphere/build_mapping_file.py

@@ -0,0 +1,147 @@
+import os
+import shutil
+
+from compass.landice.ismip7.ice_sheet_params import get_params
+from compass.landice.ismip7.mapping import build_mapping_file
+from compass.step import Step
+
+
+class BuildMappingFile(Step):
+    """
+    A step for building the ESMF mapping file (regridding weights) from
+    the ISMIP7 atmosphere grid to the MALI mesh. This step runs
+    ESMF_RegridWeightGen with the full esmf_ntasks allocation, so it can
+    be run on a separate node allocation from the processing steps (which
+    run on a single node).
+
+    The mapping file can be reused across scenarios, since it depends only on
+    the ice sheet, mesh name, and remapping method. If mapping_files_path is
+    provided in the config, this step will symlink any existing weight file
+    from that directory and skip the build.
+    """
+
+    def __init__(self, test_case):
+        """
+        Create the step
+
+        Parameters
+        ----------
+        test_case : compass.landice.tests.ismip7_forcing.atmosphere.Atmosphere
+            The test case this step belongs to
+        """
+        super().__init__(test_case=test_case, name="build_mapping_file")
+
+    def setup(self):
+        """
+        Set up this step of the test case
+        """
+        config = self.config
+        section = config["ismip7"]
+        base_path_mali = section.get("base_path_mali")
+        mali_mesh_file = section.get("mali_mesh_file")
+
+        self.add_input_file(filename=mali_mesh_file,
+                            target=os.path.join(base_path_mali,
+                                                mali_mesh_file))
+
+        # Request the full esmf_ntasks allocation for this step
+        self.ntasks = section.getint("esmf_ntasks")
+        self.min_tasks = self.ntasks
+
+    def run(self):
+        """
+        Run this step of the test case
+        """
+        logger = self.logger
+        config = self.config
+        params = get_params(config)
+
+        section = config["ismip7"]
+        base_path_ismip7 = section.get("base_path_ismip7")
+        mali_mesh_name = section.get("mali_mesh_name")
+        mali_mesh_file = section.get("mali_mesh_file")
+        ice_sheet = section.get("ice_sheet")
+        output_base_path = section.get("output_base_path")
+
+        # Check if user provided a path to reuse existing weight files
+        mapping_files_path = section.get("mapping_files_path")
+
+        section = config["ismip7_atmosphere"]
+        method_remap = section.get("method_remap")
+
+        # Construct the canonical mapping file name
+        mapping_file = (f"map_ismip7_{ice_sheet}_atm_to_"
+                        f"{mali_mesh_name}_{method_remap}.nc")
+
+        # If mapping_files_path is provided, symlink any existing file
+        if mapping_files_path != "NotAvailable":
+            source_file = os.path.join(mapping_files_path, mapping_file)
+            if os.path.exists(source_file):
+                logger.info(f"Symlinking existing mapping file from "
+                            f"{mapping_files_path}")
+                if os.path.exists(mapping_file):
+                    os.remove(mapping_file)
+                os.symlink(source_file, mapping_file)
+
+        # Use any atmosphere file as the grid template for building the
+        # mapping file. We'll use the first year's acabf (SMB) file.
+        ismip7_section = config["ismip7"]
+        base_path_ismip7 = ismip7_section.get("base_path_ismip7")
+        mali_mesh_name = ismip7_section.get("mali_mesh_name")
+        mali_mesh_file = ismip7_section.get("mali_mesh_file")
+        ice_sheet = ismip7_section.get("ice_sheet")
+        output_base_path = ismip7_section.get("output_base_path")
+        mapping_files_path = ismip7_section.get("mapping_files_path")
+
+        atmosphere_section = config["ismip7_atmosphere"]
+        method_remap = atmosphere_section.get("method_remap")
+
+        prefix = params["prefix"]
+        resolution = params["atm_resolution"]
+        version = params["atm_version"]
+
+        model = (ismip7_section.get("model")
+                 if params["atm_model"] is None
+                 else params["atm_model"])
+
+        scenario = ismip7_section.get("scenario")
+
+        input_path = os.path.join(base_path_ismip7, "acabf", version)
+        # Use a simple pattern that will match the first available file
+        file_pattern = (f"acabf_{prefix}_{model}_{scenario}_"
+                        f"SDBN1-{resolution}_{version}_*.nc")
+        ismip7_grid_file = os.path.join(input_path, file_pattern)
+
+        import glob
+        grid_files = sorted(glob.glob(ismip7_grid_file))
+        if not grid_files:
+            raise FileNotFoundError(
+                f"No atmosphere grid file found matching pattern:\n"
+                f"  {ismip7_grid_file}")
+        ismip7_grid_file = grid_files[0]
+        logger.info(f"Using grid template: "
+                    f"{os.path.basename(ismip7_grid_file)}")
+
+        # Build the mapping file (build_mapping_file will skip if it already
+        # exists, e.g., from the symlink above or a previous failed run)
+        build_mapping_file(config, logger, ismip7_grid_file, mapping_file,
+                           mali_mesh_file=mali_mesh_file,
+                           method_remap=method_remap)
+
+        # Copy the mapping file to output_base_path/mapping_files/ for reuse
+        # in future runs
+        mapping_files_dir = os.path.join(output_base_path, "mapping_files")
+        if not os.path.exists(mapping_files_dir):
+            os.makedirs(mapping_files_dir)
+
+        dst = os.path.join(mapping_files_dir, mapping_file)
+        # Only copy if it's a real file (not a symlink we just created)
+        if not os.path.islink(mapping_file):
+            logger.info(f"Copying mapping file to {mapping_files_dir} "
+                        f"for reuse")
+            shutil.copy(mapping_file, dst)
+        else:
+            logger.info(f"Mapping file is a symlink; not copying to "
+                        f"{mapping_files_dir}")
+
+        logger.info("Done building atmosphere mapping file.")

compass/landice/tests/ismip7_forcing/atmosphere/process_runoff.py

@@ -7,7 +7,6 @@
 from mpas_tools.logging import check_call
 
 from compass.landice.ismip7.ice_sheet_params import get_params
-from compass.landice.ismip7.mapping import build_mapping_file
 from compass.landice.ismip7.remap import extrapolate_source
 from compass.step import Step
 
@@ -38,11 +37,22 @@ def setup(self):
         section = config["ismip7"]
         base_path_mali = section.get("base_path_mali")
         mali_mesh_file = section.get("mali_mesh_file")
+        mali_mesh_name = section.get("mali_mesh_name")
+        ice_sheet = section.get("ice_sheet")
+
+        section_atm = config["ismip7_atmosphere"]
+        method_remap = section_atm.get("method_remap")
 
         self.add_input_file(filename=mali_mesh_file,
                             target=os.path.join(base_path_mali,
                                                 mali_mesh_file))
 
+        # Add input file for the mapping file from build_mapping_file step
+        mapping_file = (f"map_ismip7_{ice_sheet}_atm_to_"
+                        f"{mali_mesh_name}_{method_remap}.nc")
+        self.add_input_file(filename=mapping_file,
+                            target=f"../build_mapping_file/{mapping_file}")
+
     def run(self):
         """
         Run this step of the test case
@@ -54,7 +64,6 @@ def run(self):
         section = config["ismip7"]
         base_path_ismip7 = section.get("base_path_ismip7")
         mali_mesh_name = section.get("mali_mesh_name")
-        mali_mesh_file = section.get("mali_mesh_file")
         model = section.get("model")
         scenario = section.get("scenario")
         output_base_path = section.get("output_base_path")
@@ -103,17 +112,11 @@ def run(self):
         logger.info(f"Found {len(input_files)} runoff files for years "
                     f"{start_year}-{end_year}")
 
-        # Build mapping file (reuse if already created by other atm steps)
+        # Construct the mapping file name (symlinked from
+        # build_mapping_file step)
         mapping_file = (f"map_ismip7_{ice_sheet}_atm_to_"
                         f"{mali_mesh_name}_{method_remap}.nc")
 
-        if not os.path.exists(mapping_file):
-            logger.info("Building mapping file...")
-            build_mapping_file(config, logger,
-                               input_files[0], mapping_file,
-                               mali_mesh_file=mali_mesh_file,
-                               method_remap=method_remap)
-
         # Remap each year file
         remapped_files = []
         for input_file in input_files:

compass/landice/tests/ismip7_forcing/atmosphere/process_smb.py

@@ -7,7 +7,6 @@
 from mpas_tools.logging import check_call
 
 from compass.landice.ismip7.ice_sheet_params import get_params
-from compass.landice.ismip7.mapping import build_mapping_file
 from compass.landice.ismip7.remap import extrapolate_source
 from compass.step import Step
 
@@ -38,11 +37,22 @@ def setup(self):
         section = config["ismip7"]
         base_path_mali = section.get("base_path_mali")
         mali_mesh_file = section.get("mali_mesh_file")
+        mali_mesh_name = section.get("mali_mesh_name")
+        ice_sheet = section.get("ice_sheet")
+
+        section_atm = config["ismip7_atmosphere"]
+        method_remap = section_atm.get("method_remap")
 
         self.add_input_file(filename=mali_mesh_file,
                             target=os.path.join(base_path_mali,
                                                 mali_mesh_file))
 
+        # Add input file for the mapping file from build_mapping_file step
+        mapping_file = (f"map_ismip7_{ice_sheet}_atm_to_"
+                        f"{mali_mesh_name}_{method_remap}.nc")
+        self.add_input_file(filename=mapping_file,
+                            target=f"../build_mapping_file/{mapping_file}")
+
     def run(self):
         """
         Run this step of the test case
@@ -54,7 +64,6 @@ def run(self):
         section = config["ismip7"]
         base_path_ismip7 = section.get("base_path_ismip7")
         mali_mesh_name = section.get("mali_mesh_name")
-        mali_mesh_file = section.get("mali_mesh_file")
         model = section.get("model")
         scenario = section.get("scenario")
         output_base_path = section.get("output_base_path")
@@ -102,18 +111,12 @@ def run(self):
         logger.info(f"Found {len(input_files)} SMB files for years "
                     f"{start_year}-{end_year}")
 
-        # Build mapping file using the first input file as the grid template
+        # Construct the mapping file name (symlinked from
+        # build_mapping_file step)
         ice_sheet = config.get("ismip7", "ice_sheet")
         mapping_file = (f"map_ismip7_{ice_sheet}_atm_to_"
                         f"{mali_mesh_name}_{method_remap}.nc")
 
-        if not os.path.exists(mapping_file):
-            logger.info("Building mapping file...")
-            build_mapping_file(config, logger,
-                               input_files[0], mapping_file,
-                               mali_mesh_file=mali_mesh_file,
-                               method_remap=method_remap)
-
         # Remap each year file
         remapped_files = []
         for input_file in input_files:

compass/landice/tests/ismip7_forcing/atmosphere/process_smb_gradient.py

@@ -7,7 +7,6 @@
 from mpas_tools.logging import check_call
 
 from compass.landice.ismip7.ice_sheet_params import get_params
-from compass.landice.ismip7.mapping import build_mapping_file
 from compass.landice.ismip7.remap import extrapolate_source
 from compass.step import Step
 
@@ -39,11 +38,22 @@ def setup(self):
         section = config["ismip7"]
         base_path_mali = section.get("base_path_mali")
         mali_mesh_file = section.get("mali_mesh_file")
+        mali_mesh_name = section.get("mali_mesh_name")
+        ice_sheet = section.get("ice_sheet")
+
+        section_atm = config["ismip7_atmosphere"]
+        method_remap = section_atm.get("method_remap")
 
         self.add_input_file(filename=mali_mesh_file,
                             target=os.path.join(base_path_mali,
                                                 mali_mesh_file))
 
+        # Add input file for the mapping file from build_mapping_file step
+        mapping_file = (f"map_ismip7_{ice_sheet}_atm_to_"
+                        f"{mali_mesh_name}_{method_remap}.nc")
+        self.add_input_file(filename=mapping_file,
+                            target=f"../build_mapping_file/{mapping_file}")
+
     def run(self):
         """
         Run this step of the test case
@@ -55,7 +65,6 @@ def run(self):
         section = config["ismip7"]
         base_path_ismip7 = section.get("base_path_ismip7")
         mali_mesh_name = section.get("mali_mesh_name")
-        mali_mesh_file = section.get("mali_mesh_file")
         model = section.get("model")
         scenario = section.get("scenario")
         output_base_path = section.get("output_base_path")
@@ -103,18 +112,12 @@ def run(self):
         logger.info(f"Found {len(input_files)} SMB gradient files for years "
                     f"{start_year}-{end_year}")
 
-        # Build mapping file (reuse if already created by process_smb)
+        # Construct the mapping file name (symlinked from
+        # build_mapping_file step)
         ice_sheet = config.get("ismip7", "ice_sheet")
         mapping_file = (f"map_ismip7_{ice_sheet}_atm_to_"
                         f"{mali_mesh_name}_{method_remap}.nc")
 
-        if not os.path.exists(mapping_file):
-            logger.info("Building mapping file...")
-            build_mapping_file(config, logger,
-                               input_files[0], mapping_file,
-                               mali_mesh_file=mali_mesh_file,
-                               method_remap=method_remap)
-
         # Remap each year file
         remapped_files = []
         for input_file in input_files: