Merge pull request #506 from xylar/simplified-seaice-partition-tool

Add `simple_seaice_partition` tool

conda_package/docs/seaice/partition.rst

@@ -149,6 +149,55 @@ line, in a file. You can optionally save a NetCDF file with partition
 information ``partition_diag.nc``, which will contain a ``partition_{nProcs}``
 field for each number of processors requested.
 
+A simplified tool, primarily intended for use on LCRC machines Anvil and
+Chrysalis, has only a few arguments:
+
+.. code-block:: none
+
+    $ simple_seaice_partitions --help
+    usage: simple_seaice_partitions [-h] -m MESHFILENAME -p OUTPUTPREFIX -n
+                                    [NPROCSARRAY ...] [-d DATADIR]
+
+    Create sea-ice partitions on LCRC.
+
+    options:
+      -h, --help            show this help message and exit
+      -m MESHFILENAME, --mesh MESHFILENAME
+                            MPAS-Seaice mesh file.
+      -p OUTPUTPREFIX, --prefix OUTPUTPREFIX
+                            prefix for output partition filenames.
+      -n [NPROCSARRAY ...], --nprocs [NPROCSARRAY ...]
+                            list of the number of processors to create partition
+                            for.
+      -d DATADIR, --datadir DATADIR
+                            Directory with seaice_QU60km_polar.nc and
+                            icePresent_QU60km_polar.nc.
+
+The mesh file is any file that contains the MPAS-Seaice mesh.  Some meshes
+are available in `inputdata/share/meshes/mpas/sea-ice` and also each
+MPAS-Seaice initial condition in `inputdata/ice/mpas-seaice/<mesh_name>`
+contains the MPAS mesh.  Which specific initial condition you choose should
+not matter because the mesh should be identical.
+
+The output prefix can be an absolute or relative path prefix for the graph
+partition file to be created.  Typically, this will be something like
+``partitions/mpas-seaice.graph.info.230313``.  It should end in a date that
+matches other existing partition files (i.e. it can't typically be today's
+date or E3SM won't find the new partition file) and should not contain the
+``.part.<task_count>`` suffix.
+
+You can provide several task counts with ``-n`` for efficiency.  There is a
+significant overhead in calling the tool multiple times for different task
+counts.
+
+Here is an example:
+
+.. code-block:: bash
+
+    cd /lcrc/group/e3sm/public_html/inputdata/ice/mpas-seaice/WC14to60E2r3
+    simple_seaice_partitions -m seaice.WC14to60E2r3.210210.nc -p partitions/mpas-seaice.graph.info.230313 -n 468
+
+
 Graph partition function
 ------------------------
 

conda_package/mpas_tools/seaice/partition.py

@@ -6,6 +6,7 @@
 import subprocess
 import argparse
 import shutil
+import glob
 
 from .regrid import regrid_to_other_mesh
 from .mask import extend_seaice_mask
@@ -322,6 +323,121 @@ def create_partitions():
                               cullEquatorialRegion=False)
 
 
+def simple_partitions():
+    """
+    An entry point for creating sea-ice partitions on LCRC (Anvil and
+    Chrysalis)
+    """
+
+    data_dir = '/lcrc/group/e3sm/public_html/mpas_standalonedata/' \
+               'mpas-seaice/partition'
+
+    # parsing input
+    parser = argparse.ArgumentParser(
+        description='Create sea-ice partitions on LCRC.')
+
+    parser.add_argument(
+        '-m', '--mesh', dest="meshFilename", required=True,
+        help='MPAS-Seaice mesh file.')
+    parser.add_argument(
+        '-p', '--prefix', dest="outputPrefix", required=True,
+        help='prefix for output partition filenames.')
+    parser.add_argument(
+        '-n', '--nprocs', dest="nProcsArray", nargs='*', required=True,
+        help='list of the number of processors to create partition for.',
+        type=int)
+    parser.add_argument(
+        '-d', '--datadir', dest="dataDir", required=False,
+        default=data_dir,
+        help='Directory with seaice_QU60km_polar.nc and '
+             'icePresent_QU60km_polar.nc.')
+
+    args = parser.parse_args()
+
+    meshFilenameDst = os.path.abspath(args.meshFilename)
+
+    tmpdir = 'tmp_seaice_part_dir'
+    try:
+        shutil.rmtree(tmpdir)
+    except FileNotFoundError:
+        pass
+
+    os.makedirs(tmpdir)
+
+    cwd = os.getcwd()
+
+    os.chdir(tmpdir)
+
+    # make a local link to the mesh file
+    basename = os.path.basename(meshFilenameDst)
+    command = ['ln', '-s', meshFilenameDst, basename]
+    subprocess.run(command, check=True)
+    meshFilenameDst = basename
+
+    # 1) Regrid the ice presence from the input data mesh to the grid of choice
+    print("Regrid to desired mesh...")
+    filenameOut = "icePresent_regrid.nc"
+
+    meshFilenameSrc = os.path.join(data_dir, 'seaice_QU60km_polar.nc')
+    filenameData = os.path.join(data_dir, 'icePresent_QU60km_polar.nc')
+
+    regrid_to_other_mesh(
+        meshFilenameSrc=meshFilenameSrc,
+        filenameData=filenameData,
+        meshFilenameDst=meshFilenameDst,
+        filenameOut=filenameOut,
+        generateWeights=True,
+        weightsFilename=None)
+
+    # 2) create icePresence variable
+    print("fix_regrid_output...")
+
+    inputFile = "icePresent_regrid.nc"
+    outputFile = "icePresent_regrid_modify.nc"
+    subprocess.call(["fix_regrid_output.exe", inputFile, meshFilenameDst,
+                     outputFile])
+
+    # 3) create variable icePresenceExtended
+    print("extend_seaice_mask...")
+    filenamePresence = "icePresent_regrid_modify.nc"
+    extend_seaice_mask(meshFilenameDst, filenamePresence, 0.0, False)
+
+    # 4) Make the regions file from the icePresenceExtended variable
+    print("make_regions_file...")
+    filenameIcePresent = "icePresent_regrid_modify.nc"
+    filenameOut = "regions.nc"
+    make_regions_file(filenameIcePresent=filenameIcePresent,
+                      filenameMesh=meshFilenameDst,
+                      regionType="two_region_eq",
+                      varname="icePresenceExtended",
+                      limit=0.5,
+                      filenameOut=filenameOut)
+
+    nProcsArray = args.nProcsArray
+
+    # create partitions
+    regionFilename = "regions.nc"
+    outputPrefix = os.path.join(cwd, args.outputPrefix)
+
+    gen_seaice_mesh_partition(meshFilename=meshFilenameDst,
+                              regionFilename=regionFilename,
+                              nProcsArray=nProcsArray,
+                              mpasCullerLocation=None,
+                              outputPrefix=outputPrefix,
+                              plotting=False,
+                              metis="gpmetis",
+                              cullEquatorialRegion=False)
+
+    for file in glob.glob(f'{outputPrefix}*'):
+        command = ['chmod', 'ug+rw', file]
+        subprocess.run(command, check=True)
+        command = ['chmod', 'o+r', file]
+        subprocess.run(command, check=True)
+
+    os.chdir(cwd)
+    shutil.rmtree(tmpdir)
+
+
 # ---------------------------------------------------------------------
 # Private functions
 # ---------------------------------------------------------------------

conda_package/recipe/meta.yaml

@@ -32,6 +32,7 @@ build:
     - compute_projection_region_masks = mpas_tools.mesh.mask:entry_point_compute_projection_grid_region_masks
     - prepare_seaice_partitions = mpas_tools.seaice.partition:prepare_partitions
     - create_seaice_partitions = mpas_tools.seaice.partition:create_partitions
+    - simple_seaice_partitions = mpas_tools.seaice.partition:simple_partitions
 
 requirements:
   build:

conda_package/setup.py

@@ -96,4 +96,5 @@
                      'compute_lon_lat_region_masks = mpas_tools.mesh.mask:entry_point_compute_lon_lat_region_masks',
                      'compute_projection_region_masks = mpas_tools.mesh.mask:entry_point_compute_projection_region_masks',
                      'prepare_seaice_partitions = mpas_tools.seaice.partition:prepare_partitions',
-                     'create_seaice_partitions = mpas_tools.seaice.partition:create_partitions']})
+                     'create_seaice_partitions = mpas_tools.seaice.partition:create_partitions',
+                     'simple_seaice_partitions = mpas_tools.seaice.partition:simple_partitions']})