Remove hardcoded cpus_per_task=128 in landice mesh_gen cases

compass/landice/mesh.py

@@ -22,6 +22,8 @@
 from scipy.interpolate import interpn
 from scipy.ndimage import distance_transform_edt
 
+from compass.step import Step
+
 
 def mpas_flood_fill(seed_mask, grow_mask, cellsOnCell, nEdgesOnCell,
                     grow_iters=sys.maxsize):
@@ -864,6 +866,26 @@ def build_mali_mesh(self, cell_width, x1, y1, geom_points,
     check_call(args, logger=logger)
 
 
+class LandiceMeshStep(Step):
+    """
+    A base class for land-ice mesh-generation steps that should use all
+    available cores for serial post-processing steps.
+    """
+
+    def constrain_resources(self, available_resources):
+        """
+        Update ``cpus_per_task`` to use all available cores
+
+        Parameters
+        ----------
+        available_resources : dict
+            A dictionary containing available resources (cores, tasks, nodes
+            and cores_per_node)
+        """
+        super().constrain_resources(available_resources)
+        self.cpus_per_task = available_resources['cores']
+
+
 def make_region_masks(self, mesh_filename, mask_filename,
                       cores, tags, component='landice', all_tags=True):
     """

compass/landice/tests/antarctica/mesh.py

@@ -2,20 +2,11 @@
 import xarray as xr
 from mpas_tools.logging import check_call
 
-from compass.landice.mesh import (
-    add_bedmachine_thk_to_ais_gridded_data,
-    build_cell_width,
-    build_mali_mesh,
-    get_optional_interp_datasets,
-    make_region_masks,
-    preprocess_ais_data,
-    run_optional_interpolation,
-)
+from compass.landice import mesh as landice_mesh
 from compass.model import make_graph_file
-from compass.step import Step
 
 
-class Mesh(Step):
+class Mesh(landice_mesh.LandiceMeshStep):
     """
     A step for creating a mesh and initial condition for Antarctica test cases
 
@@ -35,8 +26,7 @@ def __init__(self, test_case):
             The test case this step belongs to
 
         """
-        super().__init__(test_case=test_case, name='mesh', cpus_per_task=128,
-                         min_cpus_per_task=1)
+        super().__init__(test_case=test_case, name='mesh')
 
         self.mesh_filename = 'Antarctica.nc'
         self.add_output_file(filename='graph.info')
@@ -64,8 +54,8 @@ def run(self):
         parallel_executable = config.get('parallel', 'parallel_executable')
         nProcs = section_ais.get('nProcs')
         src_proj = section_ais.get("src_proj")
-        bedmachine_dataset, measures_dataset = get_optional_interp_datasets(
-            section_ais, logger)
+        bedmachine_dataset, measures_dataset = \
+            landice_mesh.get_optional_interp_datasets(section_ais, logger)
 
         section_name = 'mesh'
 
@@ -74,7 +64,7 @@ def run(self):
 
         if bedmachine_dataset is not None:
             bm_updated_gridded_dataset = (
-                add_bedmachine_thk_to_ais_gridded_data(
+                landice_mesh.add_bedmachine_thk_to_ais_gridded_data(
                     self,
                     source_gridded_dataset,
                     bedmachine_dataset))
@@ -86,13 +76,13 @@ def run(self):
         ds.close()
         logger.info('calling build_cell_width')
         cell_width, x1, y1, geom_points, geom_edges, floodFillMask = \
-            build_cell_width(
+            landice_mesh.build_cell_width(
                 self, section_name=section_name,
                 gridded_dataset=bm_updated_gridded_dataset,
                 flood_fill_start=[nx // 2, ny // 2])
 
         # Now build the base mesh and perform the standard interpolation
-        build_mali_mesh(
+        landice_mesh.build_mali_mesh(
             self, cell_width, x1, y1, geom_points, geom_edges,
             mesh_name=self.mesh_filename, section_name=section_name,
             gridded_dataset=bm_updated_gridded_dataset,
@@ -115,7 +105,7 @@ def run(self):
         # Preprocess the gridded AIS source datasets to work
         # with the rest of the workflow
         logger.info('calling preprocess_ais_data')
-        preprocessed_gridded_dataset = preprocess_ais_data(
+        preprocessed_gridded_dataset = landice_mesh.preprocess_ais_data(
             self, bm_updated_gridded_dataset, floodFillMask)
 
         # interpolate fields from *preprocessed* composite dataset
@@ -140,7 +130,7 @@ def run(self):
         interpolate_data = section_ais.getboolean(
             'interpolate_data', fallback=False)
         if interpolate_data:
-            run_optional_interpolation(
+            landice_mesh.run_optional_interpolation(
                 self, self.mesh_filename, src_proj,
                 parallel_executable, nProcs,
                 bedmachine_dataset=bedmachine_dataset,
@@ -153,14 +143,16 @@ def run(self):
 
         # create a region mask
         mask_filename = f'{self.mesh_filename[:-3]}_imbie_regionMasks.nc'
-        make_region_masks(self, self.mesh_filename, mask_filename,
-                          self.cpus_per_task,
-                          tags=['EastAntarcticaIMBIE',
-                                'WestAntarcticaIMBIE',
-                                'AntarcticPeninsulaIMBIE'],
-                          all_tags=False)
+        landice_mesh.make_region_masks(
+            self, self.mesh_filename, mask_filename,
+            self.cpus_per_task,
+            tags=['EastAntarcticaIMBIE',
+                  'WestAntarcticaIMBIE',
+                  'AntarcticPeninsulaIMBIE'],
+            all_tags=False)
 
         mask_filename = f'{self.mesh_filename[:-3]}_ismip6_regionMasks.nc'
-        make_region_masks(self, self.mesh_filename, mask_filename,
-                          self.cpus_per_task,
-                          tags=['ISMIP6_Basin'])
+        landice_mesh.make_region_masks(
+            self, self.mesh_filename, mask_filename,
+            self.cpus_per_task,
+            tags=['ISMIP6_Basin'])

compass/landice/tests/crane/mesh.py

@@ -1,15 +1,8 @@
-from compass.landice.mesh import (
-    build_cell_width,
-    build_mali_mesh,
-    get_mesh_config_bounding_box,
-    get_optional_interp_datasets,
-    run_optional_interpolation,
-)
+from compass.landice import mesh as landice_mesh
 from compass.model import make_graph_file
-from compass.step import Step
 
 
-class Mesh(Step):
+class Mesh(landice_mesh.LandiceMeshStep):
     """
     A step for creating a mesh and initial condition for thwaites test cases
     """
@@ -21,8 +14,7 @@ def __init__(self, test_case):
         test_case : compass.TestCase
             The test case this step belongs to
         """
-        super().__init__(test_case=test_case, name='mesh', cpus_per_task=128,
-                         min_cpus_per_task=1)
+        super().__init__(test_case=test_case, name='mesh')
 
         self.add_output_file(filename='graph.info')
         self.add_output_file(filename='Crane.nc')
@@ -49,16 +41,16 @@ def run(self):
         section_name = 'mesh'
         section = config[section_name]
         src_proj = section.get('src_proj')
-        bedmachine_dataset, measures_dataset = get_optional_interp_datasets(
-            section, logger)
+        bedmachine_dataset, measures_dataset = \
+            landice_mesh.get_optional_interp_datasets(section, logger)
 
         logger.info('calling build_cell_width')
         cell_width, x1, y1, geom_points, geom_edges, floodMask = \
-            build_cell_width(
+            landice_mesh.build_cell_width(
                 self, section_name=section_name,
                 gridded_dataset='antarctica_8km_2024_01_29.nc')
 
-        build_mali_mesh(
+        landice_mesh.build_mali_mesh(
             self, cell_width, x1, y1, geom_points, geom_edges,
             mesh_name=mesh_name, section_name=section_name,
             gridded_dataset='antarctica_1km_2024_01_29_AP.nc',
@@ -71,9 +63,10 @@ def run(self):
         interpolate_data = section.getboolean(
             'interpolate_data', fallback=False)
         if interpolate_data:
-            run_optional_interpolation(
+            landice_mesh.run_optional_interpolation(
                 self, mesh_name, src_proj, parallel_executable, nProcs,
-                subset_bounds=get_mesh_config_bounding_box(section),
+                subset_bounds=landice_mesh.get_mesh_config_bounding_box(
+                    section),
                 bedmachine_dataset=bedmachine_dataset,
                 measures_dataset=measures_dataset)
 

compass/landice/tests/greenland/mesh.py

@@ -1,19 +1,11 @@
 import numpy as np
 import xarray as xr
 
-from compass.landice.mesh import (
-    build_cell_width,
-    build_mali_mesh,
-    get_mesh_config_bounding_box,
-    get_optional_interp_datasets,
-    make_region_masks,
-    run_optional_interpolation,
-)
+from compass.landice import mesh as landice_mesh
 from compass.model import make_graph_file
-from compass.step import Step
 
 
-class Mesh(Step):
+class Mesh(landice_mesh.LandiceMeshStep):
     """
     A step for creating a mesh and initial condition for greenland test cases
 
@@ -32,8 +24,7 @@ def __init__(self, test_case):
             The test case this step belongs to
 
         """
-        super().__init__(test_case=test_case, name='mesh', cpus_per_task=128,
-                         min_cpus_per_task=1)
+        super().__init__(test_case=test_case, name='mesh')
 
         # output files
         self.mesh_filename = 'GIS.nc'
@@ -64,6 +55,7 @@ def setup(self):
                             target=geojson_filename,
                             database=None)
 
+
     def run(self):
         """
         Run this step of the test case
@@ -78,8 +70,8 @@ def run(self):
         src_proj = section_gis.get("src_proj")
         geojson_filename = section_gis.get('geojson_filename')
 
-        bedmachine_dataset, measures_dataset = get_optional_interp_datasets(
-            section_gis, logger)
+        bedmachine_dataset, measures_dataset = \
+            landice_mesh.get_optional_interp_datasets(section_gis, logger)
 
         if bedmachine_dataset is not None:
             ds_bm = xr.open_dataset(bedmachine_dataset)
@@ -89,7 +81,7 @@ def run(self):
                 float(ds_bm.y1.min()),
                 float(ds_bm.y1.max())]
             ds_bm.close()
-            bounding_box = get_mesh_config_bounding_box(
+            bounding_box = landice_mesh.get_mesh_config_bounding_box(
                 section_gis, default_bounds=default_bounds)
         else:
             bounding_box = None
@@ -101,13 +93,13 @@ def run(self):
 
         logger.info('calling build_cell_width')
         cell_width, x1, y1, geom_points, geom_edges, floodMask = \
-            build_cell_width(
+            landice_mesh.build_cell_width(
                 self, section_name=section_name,
                 gridded_dataset=source_gridded_dataset_2km,
                 flood_fill_start=[100, 700])
 
         # Now build the base mesh and perform the standard interpolation
-        build_mali_mesh(
+        landice_mesh.build_mali_mesh(
             self, cell_width, x1, y1, geom_points, geom_edges,
             mesh_name=self.mesh_filename, section_name=section_name,
             gridded_dataset=source_gridded_dataset_1km,
@@ -120,7 +112,7 @@ def run(self):
         interpolate_data = section_gis.getboolean(
             'interpolate_data', fallback=False)
         if interpolate_data:
-            run_optional_interpolation(
+            landice_mesh.run_optional_interpolation(
                 self, self.mesh_filename, src_proj,
                 parallel_executable, nProcs,
                 bedmachine_dataset=bedmachine_dataset,
@@ -133,23 +125,25 @@ def run(self):
 
         # create region masks
         mask_filename = f'{self.mesh_filename[:-3]}_ismip6_regionMasks.nc'
-        make_region_masks(self, self.mesh_filename, mask_filename,
-                          self.cpus_per_task,
-                          tags=["Greenland", "ISMIP6", "Shelf"],
-                          component='ocean')
+        landice_mesh.make_region_masks(
+            self, self.mesh_filename, mask_filename,
+            self.cpus_per_task,
+            tags=["Greenland", "ISMIP6", "Shelf"],
+            component='ocean')
 
         mask_filename = f'{self.mesh_filename[:-3]}_zwally_regionMasks.nc'
-        make_region_masks(self, self.mesh_filename, mask_filename,
-                          self.cpus_per_task,
-                          tags=['eastCentralGreenland',
-                                'northEastGreenland',
-                                'northGreenland',
-                                'northWestGreenland',
-                                'southEastGreenland',
-                                'southGreenland',
-                                'southWestGreenland',
-                                'westCentralGreenland'],
-                          all_tags=False)
+        landice_mesh.make_region_masks(
+            self, self.mesh_filename, mask_filename,
+            self.cpus_per_task,
+            tags=['eastCentralGreenland',
+                  'northEastGreenland',
+                  'northGreenland',
+                  'northWestGreenland',
+                  'southEastGreenland',
+                  'southGreenland',
+                  'southWestGreenland',
+                  'westCentralGreenland'],
+            all_tags=False)
 
         # Do some final validation of the mesh
         ds = xr.open_dataset(self.mesh_filename)

compass/landice/tests/humboldt/mesh.py

@@ -1,15 +1,8 @@
-from compass.landice.mesh import (
-    build_cell_width,
-    build_mali_mesh,
-    get_mesh_config_bounding_box,
-    get_optional_interp_datasets,
-    run_optional_interpolation,
-)
+from compass.landice import mesh as landice_mesh
 from compass.model import make_graph_file
-from compass.step import Step
 
 
-class Mesh(Step):
+class Mesh(landice_mesh.LandiceMeshStep):
     """
     A step for creating a mesh and initial condition for humboldt test cases
 
@@ -30,8 +23,7 @@ def __init__(self, test_case):
         mesh_type : str
             The resolution or mesh type of the test case
         """
-        super().__init__(test_case=test_case, name='mesh', cpus_per_task=128,
-                         min_cpus_per_task=1)
+        super().__init__(test_case=test_case, name='mesh')
 
         self.add_output_file(filename='graph.info')
         self.add_output_file(filename='Humboldt.nc')
@@ -59,16 +51,16 @@ def run(self):
         section = config[section_name]
         mesh_name = 'Humboldt.nc'
         src_proj = section.get('src_proj')
-        bedmachine_dataset, measures_dataset = get_optional_interp_datasets(
-            section, logger)
+        bedmachine_dataset, measures_dataset = \
+            landice_mesh.get_optional_interp_datasets(section, logger)
 
         logger.info('calling build_cell_width')
         cell_width, x1, y1, geom_points, geom_edges, floodMask = \
-            build_cell_width(
+            landice_mesh.build_cell_width(
                 self, section_name=section_name,
                 gridded_dataset='greenland_2km_2024_01_29.epsg3413.nc')
 
-        build_mali_mesh(
+        landice_mesh.build_mali_mesh(
             self, cell_width, x1, y1, geom_points, geom_edges,
             mesh_name=mesh_name, section_name=section_name,
             gridded_dataset='humboldt_1km_2024_01_29.epsg3413.icesheetonly.nc',
@@ -81,9 +73,10 @@ def run(self):
         interpolate_data = section.getboolean(
             'interpolate_data', fallback=False)
         if interpolate_data:
-            run_optional_interpolation(
+            landice_mesh.run_optional_interpolation(
                 self, mesh_name, src_proj, parallel_executable, nProcs,
-                subset_bounds=get_mesh_config_bounding_box(section),
+                subset_bounds=landice_mesh.get_mesh_config_bounding_box(
+                    section),
                 bedmachine_dataset=bedmachine_dataset,
                 measures_dataset=measures_dataset)