Switch global ocean mesh culling and files for E3SM to use python masks

compass/ocean/tests/global_ocean/files_for_e3sm/diagnostics_files.py

@@ -6,9 +6,9 @@
     MpasMeshDescriptor, Remapper
 from geometric_features import GeometricFeatures
 from geometric_features.aggregation import get_aggregator_by_name
-from mpas_tools.mesh.conversion import mask
+from mpas_tools.logging import check_call
+from mpas_tools.ocean.moc import add_moc_southern_boundary_transects
 from mpas_tools.io import write_netcdf
-from mpas_tools.ocean.moc import make_moc_basins_and_transects
 
 from compass.io import symlink
 from compass.step import Step
@@ -63,7 +63,6 @@ def run(self):
         logger = self.logger
 
         restart_filename = 'restart.nc'
-
         with xarray.open_dataset(restart_filename) as ds:
             mesh_short_name = ds.attrs['MPAS_Mesh_Short_Name']
 
@@ -76,27 +75,30 @@ def run(self):
                 os.makedirs(directory)
             except OSError:
                 pass
-
-        _make_moc_masks(mesh_short_name, logger)
+        _make_moc_masks(mesh_short_name, logger, cores)
 
         gf = GeometricFeatures()
-
         region_groups = ['Antarctic Regions', 'Arctic Ocean Regions',
                          'Arctic Sea Ice Regions', 'Ocean Basins',
-                         'Ocean Subbasins', 'ISMIP6 Regions',
-                         'Transport Transects']
+                         'Ocean Subbasins', 'ISMIP6 Regions']
 
         if with_ice_shelf_cavities:
             region_groups.append('Ice Shelves')
 
         for region_group in region_groups:
             function, prefix, date = get_aggregator_by_name(region_group)
-
             suffix = '{}{}'.format(prefix, date)
-
             fcMask = function(gf)
             _make_region_masks(mesh_short_name, suffix=suffix, fcMask=fcMask,
-                               logger=logger)
+                               logger=logger, cores=cores)
+
+        transect_groups = ['Transport Transects']
+        for transect_group in transect_groups:
+            function, prefix, date = get_aggregator_by_name(transect_group)
+            suffix = '{}{}'.format(prefix, date)
+            fcMask = function(gf)
+            _make_transect_masks(mesh_short_name, suffix=suffix, fcMask=fcMask,
+                                 logger=logger, cores=cores)
 
         _make_analysis_lat_lon_map(config, mesh_short_name, cores, logger)
         _make_analysis_polar_map(config, mesh_short_name,
@@ -115,19 +117,47 @@ def run(self):
                     '{}/{}'.format(output_dir, filename))
 
 
-def _make_region_masks(mesh_name, suffix, fcMask, logger):
+def _make_region_masks(mesh_name, suffix, fcMask, logger, cores):
     mesh_filename = 'restart.nc'
 
     geojson_filename = '{}.geojson'.format(suffix)
     mask_filename = '{}_{}.nc'.format(mesh_name, suffix)
 
     fcMask.to_geojson(geojson_filename)
 
-    dsMesh = xarray.open_dataset(mesh_filename)
+    args = ['compute_mpas_region_masks',
+            '-m', mesh_filename,
+            '-g', geojson_filename,
+            '-o', mask_filename,
+            '-t', 'cell',
+            '--process_count', '{}'.format(cores)]
+    check_call(args, logger=logger)
+
+    # make links in output directory
+    output_dir = '../assembled_files/diagnostics/mpas_analysis/' \
+                 'region_masks'
+    symlink('../../../../diagnostics_files/{}'.format(mask_filename),
+            '{}/{}'.format(output_dir, mask_filename))
+
 
-    dsMask = mask(dsMesh, fcMask=fcMask, logger=logger)
+def _make_transect_masks(mesh_name, suffix, fcMask, logger, cores,
+                         subdivision_threshold=10e3):
+    mesh_filename = 'restart.nc'
+
+    geojson_filename = '{}.geojson'.format(suffix)
+    mask_filename = '{}_{}.nc'.format(mesh_name, suffix)
 
-    write_netcdf(dsMask, mask_filename)
+    fcMask.to_geojson(geojson_filename)
+
+    args = ['compute_mpas_transect_masks',
+            '-m', mesh_filename,
+            '-g', geojson_filename,
+            '-o', mask_filename,
+            '-t', 'edge',
+            '-s', '{}'.format(subdivision_threshold),
+            '--process_count', '{}'.format(cores),
+            '--add_edge_sign']
+    check_call(args, logger=logger)
 
     # make links in output directory
     output_dir = '../assembled_files/diagnostics/mpas_analysis/' \
@@ -195,22 +225,40 @@ def _make_mapping_file(mesh_name, outGridName, inDescriptor, outDescriptor,
                                 esmf_parallel_exec=parallel_executable)
 
 
-def _make_moc_masks(mesh_short_name, logger):
+def _make_moc_masks(mesh_short_name, logger, cores):
     gf = GeometricFeatures()
 
     mesh_filename = 'restart.nc'
 
-    mask_filename = '{}_moc_masks.nc'.format(mesh_short_name)
-    mask_and_transect_filename = '{}_moc_masks_and_transects.nc'.format(
-        mesh_short_name)
+    function, prefix, date = get_aggregator_by_name('MOC Basins')
+    fcMask = function(gf)
+
+    suffix = '{}{}'.format(prefix, date)
+
+    geojson_filename = '{}.geojson'.format(suffix)
+    mask_filename = '{}_{}.nc'.format(mesh_short_name, suffix)
+
+    fcMask.to_geojson(geojson_filename)
+
+    args = ['compute_mpas_region_masks',
+            '-m', mesh_filename,
+            '-g', geojson_filename,
+            '-o', mask_filename,
+            '-t', 'cell',
+            '--process_count', '{}'.format(cores)]
+    check_call(args, logger=logger)
+
+    mask_and_transect_filename = '{}_mocBasinsAndTransects{}.nc'.format(
+        mesh_short_name, date)
+
+    dsMesh = xarray.open_dataset(mesh_filename)
+    dsMask = xarray.open_dataset(mask_filename)
 
-    geojson_filename = 'moc_basins.geojson'
+    dsMasksAndTransects = add_moc_southern_boundary_transects(
+        dsMask, dsMesh, logger=logger)
 
-    make_moc_basins_and_transects(gf, mesh_filename,
-                                  mask_and_transect_filename,
-                                  geojson_filename=geojson_filename,
-                                  mask_filename=mask_filename,
-                                  logger=logger)
+    write_netcdf(dsMasksAndTransects, mask_and_transect_filename,
+                 char_dim_name='StrLen')
 
     # make links in output directories (both inputdata and diagnostics)
     output_dir = '../assembled_files/inputdata/ocn/mpas-o/{}'.format(

compass/ocean/tests/global_ocean/global_ocean.cfg

@@ -5,7 +5,7 @@
 
 ## config options related to the mesh step
 # number of cores to use
-mesh_cores = 1
+mesh_cores = 18
 # minimum of cores, below which the step fails
 mesh_min_cores = 1
 # maximum memory usage allowed (in MB)

compass/ocean/tests/global_ocean/mesh/cull.py

@@ -2,17 +2,19 @@
 
 from geometric_features import GeometricFeatures, FeatureCollection, \
     read_feature_collection
-from mpas_tools.mesh import conversion
+from mpas_tools.mesh.conversion import cull
+from mpas_tools.mesh.mask import compute_mpas_flood_fill_mask
 from mpas_tools.io import write_netcdf
 from mpas_tools.ocean.coastline_alteration import widen_transect_edge_masks, \
     add_critical_land_blockages, add_land_locked_cells_to_mask
 from mpas_tools.viz.paraview_extractor import extract_vtk
-from mpas_tools.logging import LoggingContext
+from mpas_tools.logging import LoggingContext, check_call
 
 
 def cull_mesh(with_cavities=False, with_critical_passages=False,
               custom_critical_passages=None, custom_land_blockages=None,
-              preserve_floodplain=False, logger=None, use_progress_bar=True):
+              preserve_floodplain=False, logger=None, use_progress_bar=True,
+              process_count=1):
     """
     First step of initializing the global ocean:
 
@@ -69,17 +71,22 @@ def cull_mesh(with_cavities=False, with_critical_passages=False,
 
     use_progress_bar : bool, optional
         Whether to display progress bars (problematic in logging to a file)
+
+    process_count : int, optional
+        The number of cores to use to create masks (``None`` to use all
+        available cores)
     """
     with LoggingContext(name=__name__, logger=logger) as logger:
         _cull_mesh_with_logging(
             logger, with_cavities, with_critical_passages,
             custom_critical_passages, custom_land_blockages,
-            preserve_floodplain, use_progress_bar)
+            preserve_floodplain, use_progress_bar, process_count)
 
 
 def _cull_mesh_with_logging(logger, with_cavities, with_critical_passages,
                             custom_critical_passages, custom_land_blockages,
-                            preserve_floodplain, use_progress_bar):
+                            preserve_floodplain, use_progress_bar,
+                            process_count):
     """ Cull the mesh once the logger is defined for sure """
 
     # required for compatibility with MPAS
@@ -122,10 +129,16 @@ def _cull_mesh_with_logging(logger, with_cavities, with_critical_passages,
     fcLandCoverage.to_geojson('land_coverage.geojson')
 
     # Create the land mask based on the land coverage, i.e. coastline data
-    dsBaseMesh = xarray.open_dataset('base_mesh.nc')
-    dsLandMask = conversion.mask(dsBaseMesh, fcMask=fcLandCoverage,
-                                 logger=logger)
+    args = ['compute_mpas_region_masks',
+            '-m', 'base_mesh.nc',
+            '-g', 'land_coverage.geojson',
+            '-o', 'land_mask.nc',
+            '-t', 'cell',
+            '--process_count', '{}'.format(process_count)]
+    check_call(args, logger=logger)
 
+    dsBaseMesh = xarray.open_dataset('base_mesh.nc')
+    dsLandMask = xarray.open_dataset('land_mask.nc')
     dsLandMask = add_land_locked_cells_to_mask(dsLandMask, dsBaseMesh,
                                                latitude_threshold=43.0,
                                                nSweeps=20)
@@ -151,8 +164,16 @@ def _cull_mesh_with_logging(logger, with_cavities, with_critical_passages,
                 custom_land_blockages))
 
         # create masks from the transects
-        dsCritBlockMask = conversion.mask(dsBaseMesh, fcMask=fcCritBlockages,
-                                          logger=logger)
+        fcCritBlockages.to_geojson('critical_blockages.geojson')
+        args = ['compute_mpas_transect_masks',
+                '-m', 'base_mesh.nc',
+                '-g', 'critical_blockages.geojson',
+                '-o', 'critical_blockages.nc',
+                '-t', 'cell',
+                '-s', '10e3',
+                '--process_count', '{}'.format(process_count)]
+        check_call(args, logger=logger)
+        dsCritBlockMask = xarray.open_dataset('critical_blockages.nc')
 
         dsLandMask = add_critical_land_blockages(dsLandMask, dsCritBlockMask)
 
@@ -171,8 +192,16 @@ def _cull_mesh_with_logging(logger, with_cavities, with_critical_passages,
                 custom_critical_passages))
 
         # create masks from the transects
-        dsCritPassMask = conversion.mask(dsBaseMesh, fcMask=fcCritPassages,
-                                         logger=logger)
+        fcCritPassages.to_geojson('critical_passages.geojson')
+        args = ['compute_mpas_transect_masks',
+                '-m', 'base_mesh.nc',
+                '-g', 'critical_passages.geojson',
+                '-o', 'critical_passages.nc',
+                '-t', 'cell', 'edge',
+                '-s', '10e3',
+                '--process_count', '{}'.format(process_count)]
+        check_call(args, logger=logger)
+        dsCritPassMask = xarray.open_dataset('critical_passages.nc')
 
         # Alter critical passages to be at least two cells wide, to avoid sea
         # ice blockage
@@ -185,40 +214,53 @@ def _cull_mesh_with_logging(logger, with_cavities, with_critical_passages,
         dsPreserve.append(dsBaseMesh)
 
     # cull the mesh based on the land mask
-    dsCulledMesh = conversion.cull(dsBaseMesh, dsMask=dsLandMask,
-                                   dsPreserve=dsPreserve, logger=logger)
+    dsCulledMesh = cull(dsBaseMesh, dsMask=dsLandMask,
+                        dsPreserve=dsPreserve, logger=logger)
 
     # create a mask for the flood fill seed points
-    dsSeedMask = conversion.mask(dsCulledMesh, fcSeed=fcSeed, logger=logger)
+    dsSeedMask = compute_mpas_flood_fill_mask(dsMesh=dsCulledMesh,
+                                              fcSeed=fcSeed,
+                                              logger=logger)
 
     # cull the mesh a second time using a flood fill from the seed points
-    dsCulledMesh = conversion.cull(dsCulledMesh, dsInverse=dsSeedMask,
-                                   graphInfoFileName='culled_graph.info',
-                                   logger=logger)
+    dsCulledMesh = cull(dsCulledMesh, dsInverse=dsSeedMask,
+                        graphInfoFileName='culled_graph.info', logger=logger)
     write_netcdf(dsCulledMesh, 'culled_mesh.nc', format=netcdf_format)
 
     if critical_passages:
         # make a new version of the critical passages mask on the culled mesh
-        dsCritPassMask = conversion.mask(dsCulledMesh, fcMask=fcCritPassages,
-                                         logger=logger)
-        write_netcdf(dsCritPassMask, 'critical_passages_mask_final.nc',
-                     format=netcdf_format)
+        fcCritPassages.to_geojson('critical_passages.geojson')
+        args = ['compute_mpas_transect_masks',
+                '-m', 'culled_mesh.nc',
+                '-g', 'critical_passages.geojson',
+                '-o', 'critical_passages_mask_final.nc',
+                '-t', 'cell',
+                '-s', '10e3',
+                '--process_count', '{}'.format(process_count)]
+        check_call(args, logger=logger)
 
     if with_cavities:
         fcAntarcticIce = gf.read(
             componentName='bedmachine', objectType='region',
             featureNames=['AntarcticIceCoverage'])
+
         fcAntarcticIce.to_geojson('ice_coverage.geojson')
-        dsMask = conversion.mask(dsCulledMesh, fcMask=fcAntarcticIce,
-                                 logger=logger)
+        args = ['compute_mpas_region_masks',
+                '-m', 'culled_mesh.nc',
+                '-g', 'ice_coverage.geojson',
+                '-o', 'ice_coverage.nc',
+                '-t', 'cell',
+                '--process_count', '{}'.format(process_count)]
+        check_call(args, logger=logger)
+        dsMask = xarray.open_dataset('ice_coverage.nc')
+
         landIceMask = dsMask.regionCellMasks.isel(nRegions=0)
         dsLandIceMask = xarray.Dataset()
         dsLandIceMask['landIceMask'] = landIceMask
 
         write_netcdf(dsLandIceMask, 'land_ice_mask.nc', format=netcdf_format)
 
-        dsLandIceCulledMesh = conversion.cull(dsCulledMesh, dsMask=dsMask,
-                                              logger=logger)
+        dsLandIceCulledMesh = cull(dsCulledMesh, dsMask=dsMask, logger=logger)
         write_netcdf(dsLandIceCulledMesh, 'no_ISC_culled_mesh.nc',
                      format=netcdf_format)
 

docs/developers_guide/quick_start.rst

@@ -62,9 +62,9 @@ follows:
 .. code-block:: bash
 
     conda create -n dev_compass python=3.8 affine cartopy cartopy_offlinedata \
-        cmocean "esmf=*=mpi_mpich*" ffmpeg "geometric_features=0.3.0" git \
+        cmocean "esmf=*=mpi_mpich*" ffmpeg "geometric_features=0.4.0" git \
         ipython "jigsaw=0.9.14" "jigsawpy=0.3.3" jupyter lxml matplotlib \
-        metis "mpas_tools=0.2.0" mpich nco "netcdf4=*=nompi_*" numpy \
+        metis "mpas_tools=0.5.1" mpich nco "netcdf4=*=nompi_*" numpy \
         progressbar2 pyamg "pyremap>=0.0.7,<0.1.0" rasterio requests scipy \
         xarray
 

recipe/meta.yaml

@@ -46,7 +46,7 @@ requirements:
     - cmocean
     - esmf * {{ mpi_prefix }}_*
     - ffmpeg
-    - geometric_features 0.3.0
+    - geometric_features 0.4.0
     - git
     - ipython
     - jigsaw 0.9.14
@@ -55,7 +55,7 @@ requirements:
     - lxml
     - matplotlib-base
     - metis
-    - mpas_tools 0.2.0
+    - mpas_tools 0.5.1
     - {{ mpi }}  # [mpi != 'nompi']
     - nco
     - netcdf4 * nompi_*