diff --git a/compass/landice/mesh.py b/compass/landice/mesh.py index 0d7455bcc4..88f230290f 100644 --- a/compass/landice/mesh.py +++ b/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): """ diff --git a/compass/landice/tests/antarctica/mesh.py b/compass/landice/tests/antarctica/mesh.py index d18b80ebcc..352cbf369f 100644 --- a/compass/landice/tests/antarctica/mesh.py +++ b/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']) diff --git a/compass/landice/tests/crane/mesh.py b/compass/landice/tests/crane/mesh.py index 649b2eb309..f9e5430bfc 100644 --- a/compass/landice/tests/crane/mesh.py +++ b/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) diff --git a/compass/landice/tests/greenland/mesh.py b/compass/landice/tests/greenland/mesh.py index 029ac3444e..666d14f2e6 100644 --- a/compass/landice/tests/greenland/mesh.py +++ b/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) diff --git a/compass/landice/tests/humboldt/mesh.py b/compass/landice/tests/humboldt/mesh.py index 11e1eb639e..37bc0abfc2 100644 --- a/compass/landice/tests/humboldt/mesh.py +++ b/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) diff --git a/compass/landice/tests/isunnguata_sermia/mesh.py b/compass/landice/tests/isunnguata_sermia/mesh.py index e2c50ed1db..07ed90b001 100644 --- a/compass/landice/tests/isunnguata_sermia/mesh.py +++ b/compass/landice/tests/isunnguata_sermia/mesh.py @@ -1,10 +1,4 @@ -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 @@ -59,16 +53,16 @@ def run(self): section = config[section_name] mesh_name = 'Isunnguata_Sermia.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='greenland_1km_2024_01_29.epsg3413.icesheetonly.nc', # noqa @@ -80,9 +74,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) diff --git a/compass/landice/tests/kangerlussuaq/mesh.py b/compass/landice/tests/kangerlussuaq/mesh.py index 94bac4d727..d54aaacef8 100644 --- a/compass/landice/tests/kangerlussuaq/mesh.py +++ b/compass/landice/tests/kangerlussuaq/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 kangerlussuaq test cases @@ -31,8 +24,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='Kangerlussuaq.nc') @@ -60,16 +52,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='greenland_8km_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='greenland_1km_2024_01_29.epsg3413.icesheetonly.nc', # noqa @@ -82,9 +74,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) diff --git a/compass/landice/tests/koge_bugt_s/mesh.py b/compass/landice/tests/koge_bugt_s/mesh.py index 61c4f8f940..93f36bee56 100644 --- a/compass/landice/tests/koge_bugt_s/mesh.py +++ b/compass/landice/tests/koge_bugt_s/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 koge_bugt_s test cases @@ -31,8 +24,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='Koge_Bugt_S.nc') @@ -60,16 +52,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='greenland_8km_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='greenland_1km_2024_01_29.epsg3413.icesheetonly.nc', # noqa @@ -82,9 +74,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) diff --git a/compass/landice/tests/thwaites/mesh.py b/compass/landice/tests/thwaites/mesh.py index 0c1b6a439b..6563873fdb 100644 --- a/compass/landice/tests/thwaites/mesh.py +++ b/compass/landice/tests/thwaites/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='Thwaites.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_ASE.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)