Update algorithm for core counts for MPAS partition files

compass/ocean/suites/pr.txt

@@ -16,7 +16,6 @@ ocean/global_ocean/QU240/PHC/decomp_test
 ocean/global_ocean/QU240/PHC/threads_test
 ocean/global_ocean/QU240/PHC/analysis_test
 ocean/global_ocean/QU240/PHC/dynamic_adjustment
-ocean/global_ocean/QU240/PHC/files_for_e3sm
 
 ocean/global_ocean/QU240/PHC/RK4/performance_test
 ocean/global_ocean/QU240/PHC/RK4/restart_test

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

@@ -1,7 +1,7 @@
 import numpy as np
 
 
-def get_core_list(ncells, max_cells_per_core=6000, min_cells_per_core=100):
+def get_core_list(ncells, max_cells_per_core=30000, min_cells_per_core=2):
     """
     Get a fairly exhaustive list of core counts to partition a given number of
     cells into
@@ -24,17 +24,36 @@ def get_core_list(ncells, max_cells_per_core=6000, min_cells_per_core=100):
     cores : numpy.ndarray
         Likely numbers of cores to run with
     """
-    min_graph_size = int(ncells / max_cells_per_core)
+    min_graph_size = max(2, int(ncells / max_cells_per_core))
     max_graph_size = int(ncells / min_cells_per_core)
-    n_power2 = 2**np.arange(1, 21)
-    n_multiples12 = 12 * np.arange(1, 9)
 
-    cores = n_power2
-    for power10 in range(3):
-        cores = np.concatenate([cores, 10**power10 * n_multiples12])
-
-    mask = np.logical_and(cores >= min_graph_size,
-                          cores <= max_graph_size)
-    cores = cores[mask]
-
-    return cores
+    cores = []
+    for candidate in range(min_graph_size, max_graph_size):
+        factors = _prime_factors(candidate)
+        twos = np.count_nonzero(factors == 2)
+        fives = np.count_nonzero(factors == 5)
+        gt_five = np.count_nonzero(factors > 5)
+        big_factor = factors.max()
+        if twos > 0 and fives <= twos and gt_five <= 1 and big_factor <= 7:
+            cores.append(candidate)
+        # small odd multiples of 3 and a few that correspond to divisors of the
+        # ne30 (30x30x6=5400) size
+        elif candidate in [3, 9, 15, 21, 225, 675, 1350]:
+            cores.append(candidate)
+
+    return np.array(cores)
+
+
+# https://stackoverflow.com/a/22808285
+def _prime_factors(n):
+    i = 2
+    factors = []
+    while i * i <= n:
+        if n % i:
+            i += 1
+        else:
+            n //= i
+            factors.append(i)
+    if n > 1:
+        factors.append(n)
+    return np.array(factors)

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

@@ -78,7 +78,13 @@ def run(self):
             check_call(args, logger)
 
         # create link in assembled files directory
+
+        inputdata_dir = os.path.join(self.ocean_inputdata_dir, 'partitions')
+        try:
+            os.makedirs(inputdata_dir)
+        except FileExistsError:
+            pass
         files = glob('mpas-o.graph.info.*')
         for file in files:
             symlink(os.path.abspath(file),
-                    f'{self.ocean_inputdata_dir}/{file}')
+                    f'{inputdata_dir}/{file}')

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

@@ -60,6 +60,9 @@ def run(self):
         """
         super().run()
         logger = self.logger
+        config = self.config
+        plotting = config.getboolean('files_for_e3sm',
+                                     'plot_seaice_partitions')
         creation_date = self.creation_date
 
         with xr.open_dataset('restart.nc') as ds:
@@ -92,16 +95,23 @@ def run(self):
                 '-o', '.',
                 '-p', f'mpas-seaice.graph.info.{creation_date}',
                 '-g', 'gpmetis',
-                '--plotting',
                 '-n']
+
+        if plotting:
+            args.append('--plotting')
         args = args + [f'{ncores}' for ncores in cores]
         check_call(args, logger)
 
         # create link in assembled files directory
+        inputdata_dir = os.path.join(self.seaice_inputdata_dir, 'partitions')
+        try:
+            os.makedirs(inputdata_dir)
+        except FileExistsError:
+            pass
         files = glob('mpas-seaice.graph.info.*')
         for file in files:
             symlink(os.path.abspath(file),
-                    f'{self.seaice_inputdata_dir}/{file}')
+                    f'{inputdata_dir}/{file}')
 
 
 def _make_mapping_file(in_mesh_filename, in_mesh_name, out_mesh_filename,

compass/ocean/tests/global_ocean/global_ocean.cfg

@@ -142,3 +142,6 @@ graph_filename = autodetect
 
 # whether the mesh has ice-shelf cavities
 with_ice_shelf_cavities = autodetect
+
+# whether to write out sea-ice partition info for plotting in paraview
+plot_seaice_partitions = False

docs/developers_guide/ocean/api.rst

@@ -170,8 +170,12 @@ test cases and steps
    files_for_e3sm.ocean_initial_condition.OceanInitialCondition.run
    files_for_e3sm.scrip.Scrip
    files_for_e3sm.scrip.Scrip.run
+   files_for_e3sm.seaice_graph_partition.SeaiceGraphPartition
+   files_for_e3sm.seaice_graph_partition.SeaiceGraphPartition.run
    files_for_e3sm.seaice_initial_condition.SeaiceInitialCondition
    files_for_e3sm.seaice_initial_condition.SeaiceInitialCondition.run
+   files_for_e3sm.e3sm_to_cmip_maps.E3smToCmipMaps
+   files_for_e3sm.e3sm_to_cmip_maps.E3smToCmipMaps.run
    files_for_e3sm.diagnostic_maps.DiagnosticMaps
    files_for_e3sm.diagnostic_maps.DiagnosticMaps.run
    files_for_e3sm.diagnostic_masks.DiagnosticMasks

docs/developers_guide/ocean/test_groups/global_ocean.rst

@@ -442,7 +442,7 @@ Kuroshio8to60 and Kuroshio12to60
 ++++++++++++++++++++++++++++++++
 
 The ``Kuroshio8to60`` and ``Kuroshio12to60`` mehses are designed to explore
-dynamics of the Western Boundary Current (WBC) in the North Pacific Ocean, 
+dynamics of the Western Boundary Current (WBC) in the North Pacific Ocean,
 the Kuroshio.
 
 The class
@@ -925,19 +925,19 @@ The test case is constructed with an argument ``restart_filename``. the final
 restart file produced by the :ref:`dev_ocean_global_ocean_dynamic_adjustment`
 for the given mesh.
 
-The test case is made up of 5 steps:
+The test case is made up of 8 steps:
 
 :py:class:`compass.ocean.tests.global_ocean.files_for_e3sm.ocean_initial_condition.OceanInitialCondition`
     takes out the ``xtime`` variable from the restart file, creating a symlink
-    at ``assembled_files/inputdata/ocn/mpas-o/<mesh_short_name>/<mesh_short_name>_no_xtime.nc``
+    at ``assembled_files/inputdata/ocn/mpas-o/<mesh_short_name>/mpaso.<mesh_short_name>.<datestamp>.nc``
 
 :py:class:`compass.ocean.tests.global_ocean.files_for_e3sm.ocean_graph_partition.OceanGraphPartition`
     computes graph partitions (see :ref:`dev_model`) appropriate for a wide
-    range of core counts between ``min_graph_size = int(nCells / 6000)`` and
-    ``max_graph_size = int(nCells / 100)``.  Possible processor counts are
-    any power of 2 or any multiple of 12, 120 and 1200 in the range.  Symlinks
-    to the graph files are placed at
-    ``assembled_files/inputdata/ocn/mpas-o/<mesh_short_name>/mpas-o.graph.info.<core_count>``
+    range of core counts between ``min_graph_size = int(nCells / 30000)`` and
+    ``max_graph_size = int(nCells / 2)``.  About 400 different processor counts
+    are produced for each mesh (keeping only counts with small prime factors).
+    Symlinks to the graph files are placed at
+    ``assembled_files/inputdata/ocn/mpas-o/<mesh_short_name>/partitions/mpas-o.graph.info.<datestamp>.part.<core_count>``
 
 :py:class:`compass.ocean.tests.global_ocean.files_for_e3sm.seaice_initial_condition.SeaiceInitialCondition`
     extracts the following variables from the restart file:
@@ -959,8 +959,18 @@ The test case is made up of 5 steps:
            keep_vars.append('landIceMask')
 
     A symlink to the resulting file is placed at
-    ``assembled_files/inputdata/ocn/mpas-cice/<mesh_short_name>/seaice.<mesh_short_name>_no_xtime.nc``
+    ``assembled_files/inputdata/ocn/mpas-seaice/<mesh_short_name>/mpassi.<mesh_short_name>.<datestamp>.nc``
 
+:py:class:`compass.ocean.tests.global_ocean.files_for_e3sm.seaice_graph_partition.SeaiceGraphPartition`
+    computes graph partitions (see :ref:`dev_model`) appropriate for a wide
+    range of core counts between ``min_graph_size = int(nCells / 30000)`` and
+    ``max_graph_size = int(nCells / 2)``.  The sea-ice graph partitions
+    include cells for each processor in both polar and equatorial regions for
+    better load balancing.  See `Graph partitioning <http://mpas-dev.github.io/MPAS-Tools/stable/seaice/partition.html>`_
+    from the MPAS-Tools documentation for details.  About 400 different
+    processor counts are produced for each mesh (keeping only counts with small
+    prime factors). Symlinks to the graph files are placed at
+    ``assembled_files/inputdata/ice/mpas-seaice/<mesh_short_name>/partitions/mpas-seaice.graph.info.<datestamp>.part.<core_count>``
 
 :py:class:`compass.ocean.tests.global_ocean.files_for_e3sm.scrip.Scrip`
     generates a SCRIP file (see :ref:`global_ocean_files_for_e3sm` in the
@@ -977,6 +987,17 @@ The test case is made up of 5 steps:
     Otherwise, only one file is symlinked, and it is named
     ``ocean.<mesh_short_name>.scrip.<creation_date>.nc``
 
+
+:py:class:`compass.ocean.tests.global_ocean.files_for_e3sm.e3sm_to_cmip_maps.E3smToCmipMaps`
+    creates mapping files for
+    `e3sm_to_cmip <https://e3sm-to-cmip.readthedocs.io/en/latest/>`_.
+
+    Mapping files are created from the MPAS-Ocean and -Seaice mesh to a
+    standard 1-degree latitude-longitude grid using three methods: `aave`
+    (conservative), `mono` (monotonic) and `nco` (NCO's conservative
+    algorithm). The mapping files are symlinked in the directory
+    ``assembled_files/diagnostics/maps/``.
+
 :py:class:`compass.ocean.tests.global_ocean.files_for_e3sm.diagnostic_maps.DiagnosticMaps`
     creates mapping files for
     `MPAS-Analysis <https://mpas-dev.github.io/MPAS-Analysis/stable/>`_.