Add config options for tasks

MPAS-Dev · Oct 16, 2023 · 58a577b · 58a577b
1 parent 4544d42
commit 58a577b
Show file tree

Hide file tree

Showing 4 changed files with 44 additions and 4 deletions.
diff --git a/compass/ocean/tests/drying_slope/convergence/__init__.py b/compass/ocean/tests/drying_slope/convergence/__init__.py
@@ -1,3 +1,5 @@
+from numpy import ceil
+
 from compass.config import CompassConfigParser
 from compass.ocean.tests.drying_slope.analysis import Analysis
 from compass.ocean.tests.drying_slope.forward import Forward
@@ -77,6 +79,9 @@ def _setup_steps(self, config, subdir, method):
         self.steps_to_run = list()
 
         self.resolutions = resolutions
+        section = config['drying_slope']
+        ntasks_baseline = section.getint('ntasks_baseline')
+        min_tasks = section.getint('min_tasks')
 
         for resolution in self.resolutions:
 
@@ -89,10 +94,13 @@ def _setup_steps(self, config, subdir, method):
                                        name=init_name,
                                        resolution=resolution,
                                        coord_type=self.coord_type))
+            ntasks = max(min_tasks,
+                         int(ceil(ntasks_baseline / resolution**2.)))
             forward_step = Forward(test_case=self, resolution=resolution,
                                    name=f'forward_{res_name}',
                                    input_path=f'../{init_name}',
-                                   ntasks=4, openmp_threads=1,
+                                   ntasks=ntasks, min_tasks=min_tasks,
+                                   openmp_threads=1,
                                    damping_coeff=self.damping_coeffs[0],
                                    coord_type=self.coord_type)
             if method == 'ramp':

diff --git a/compass/ocean/tests/drying_slope/default/__init__.py b/compass/ocean/tests/drying_slope/default/__init__.py
@@ -1,3 +1,6 @@
+from numpy import ceil
+
+from compass.config import CompassConfigParser
 from compass.ocean.tests.drying_slope.forward import Forward
 from compass.ocean.tests.drying_slope.initial_state import InitialState
 from compass.ocean.tests.drying_slope.viz import Viz
@@ -32,6 +35,9 @@ def __init__(self, test_group, resolution, coord_type, method):
 
         coord_type : str
             The type of vertical coordinate (``sigma``, ``single_layer``)
+
+        method : str
+            The type of wetting-and-drying algorithm
         """
         name = 'default'
 
@@ -47,9 +53,17 @@ def __init__(self, test_group, resolution, coord_type, method):
         self.add_step(InitialState(test_case=self, resolution=resolution,
                                    coord_type=coord_type))
         damping_coeffs = None
+        config = CompassConfigParser()
+        config.add_from_package('compass.ocean.tests.drying_slope',
+                                'drying_slope.cfg')
+        section = config['drying_slope']
+        ntasks_baseline = section.getint('ntasks_baseline')
+        min_tasks = section.getint('min_tasks')
+        ntasks = max(min_tasks, int(ceil(ntasks_baseline / resolution**2.)))
         if coord_type == 'single_layer':
             forward_step = Forward(test_case=self, resolution=resolution,
-                                   ntasks=4, openmp_threads=1,
+                                   ntasks=ntasks, min_tasks=min_tasks,
+                                   openmp_threads=1,
                                    coord_type=coord_type)
             if method == 'ramp':
                 forward_step.add_namelist_options(
@@ -60,7 +74,8 @@ def __init__(self, test_group, resolution, coord_type, method):
             for damping_coeff in damping_coeffs:
                 forward_step = Forward(test_case=self, resolution=resolution,
                                        name=f'forward_{damping_coeff}',
-                                       ntasks=4, openmp_threads=1,
+                                       ntasks=ntasks, min_tasks=min_tasks,
+                                       openmp_threads=1,
                                        damping_coeff=damping_coeff,
                                        coord_type=coord_type)
                 if method == 'ramp':

diff --git a/compass/ocean/tests/drying_slope/drying_slope.cfg b/compass/ocean/tests/drying_slope/drying_slope.cfg
@@ -16,6 +16,12 @@ nx = 6
 # time step in s per km of horizontal resolution
 dt_per_km = 30
 
+# Number of tasks at 1km resolution
+ntasks_baseline = 4
+
+# Minimum number of tasks
+min_tasks = 1
+
 # config options for visualizing drying slope ouptut
 [drying_slope_convergence]
 

diff --git a/compass/ocean/tests/drying_slope/loglaw/__init__.py b/compass/ocean/tests/drying_slope/loglaw/__init__.py
@@ -1,3 +1,6 @@
+from numpy import ceil
+
+from compass.config import CompassConfigParser
 from compass.ocean.tests.drying_slope.forward import Forward
 from compass.ocean.tests.drying_slope.initial_state import InitialState
 from compass.ocean.tests.drying_slope.viz import Viz
@@ -46,8 +49,16 @@ def __init__(self, test_group, resolution, coord_type, method):
                          subdir=subdir)
         self.add_step(InitialState(test_case=self, coord_type=coord_type,
                                    resolution=resolution))
+        config = CompassConfigParser()
+        config.add_from_package('compass.ocean.tests.drying_slope',
+                                'drying_slope.cfg')
+        section = config['drying_slope']
+        ntasks_baseline = section.getint('ntasks_baseline')
+        min_tasks = section.getint('min_tasks')
+        ntasks = max(min_tasks, int(ceil(ntasks_baseline / resolution**2.)))
         forward_step = Forward(test_case=self, resolution=resolution,
-                               ntasks=4, openmp_threads=1,
+                               ntasks=ntasks, min_tasks=min_tasks,
+                               openmp_threads=1,
                                coord_type=coord_type)
         forward_step.add_namelist_options(
             {'config_implicit_bottom_drag_type': "'loglaw'"})