Merge pull request #99 from altheaden/make-cell-count-heuristic

Make cell count heuristic

docs/developers_guide/ocean/framework.md

@@ -44,11 +44,7 @@ respectively) are set automatically if `ntasks` and `min_tasks` have not
 already been set explicitly.  In such cases, a subclass of `OceanModelStep`
 must override the
 {py:meth}`polaris.ocean.model.OceanModelStep.compute_cell_count()` method
-to compute the number of cells in the mesh.  Since it is typically not possible
-to read the cell count from a file during setup, this method may need to have
-a heuristic way of approximating the number of cells during setup (i.e. when
-the `at_setup` parameter is `True`.  Then, it can return the exact number of 
-cells at runtime (i.e. `at_setup == False`).
+to approximate the number of cells in the mesh, using a simple heuristic.
 
 The algorithm for determining the resources is:
 

docs/developers_guide/ocean/test_groups/baroclinic_channel.md

@@ -53,7 +53,7 @@ constructor, the associate namelist option (`config_mom_del2`) will be given
 this value. Namelist and streams files are updated in
 {py:meth}`polaris.ocean.tests.baroclinic_channel.forward.Forward.dynamic_model_config()`
 with time steps determined algorithmically based on config options.  The
-number of cells is computed from config options in
+number of cells is approximated from config options in
 {py:meth}`polaris.ocean.tests.baroclinic_channel.forward.Forward.compute_cell_count()`
 so that this can be used to constrain the number of MPI tasks that tests
 have as their target and minimum (if the resources are not explicitly

docs/developers_guide/ocean/test_groups/inertial_gravity_wave.md

@@ -50,7 +50,7 @@ defines a step for running MPAS-Ocean from the initial condition produced in the
 `init` step.  Namelist and streams files are updated in
 {py:meth}`polaris.ocean.tests.inertial_gravity_wave.forward.Forward.dynamic_model_config()`
 with time steps determined algorithmically based on config options.  The number
-of cells is computed from config options in
+of cells is approximated from config options in
 {py:meth}`polaris.ocean.tests.inertial_gravity_wave.forward.Forward.compute_cell_count()`
 so that this can be used to constrain the number of MPI tasks that tests have as
 their target and minimum (if the resources are not explicitly prescribed).  For

docs/developers_guide/ocean/test_groups/manufactured_solution.md

@@ -50,7 +50,7 @@ defines a step for running MPAS-Ocean from the initial condition produced in
 the `init` step.  Namelist and streams files are updated in
 {py:meth}`polaris.ocean.tests.manufactured_solution.forward.Forward.dynamic_model_config()`
 with time steps determined algorithmically based on config options.  The
-number of cells is computed from config options in
+number of cells is approximated from config options in
 {py:meth}`polaris.ocean.tests.manufactured_solution.forward.Forward.compute_cell_count()`
 so that this can be used to constrain the number of MPI tasks that tests
 have as their target and minimum (if the resources are not explicitly

docs/tutorials/dev_add_test_group.md

@@ -1180,7 +1180,7 @@ you might try to stress-test your own test case.
 
 ## Adding the forward step
 
-Now that we know that the first step steems to be working, we're ready to add 
+Now that we know that the first step seems to be working, we're ready to add 
 another. We will add a `forward` step for running the MPAS-Ocean model forward
 in time from the initial condition created in `init`.  `forward`
 will be a little more complicated than `init` as we get started.
@@ -2022,16 +2022,15 @@ dt_per_km = 30
 btr_dt_per_km = 1.5
 ```
 
-Unlike in `compute_cell_count()`, we don't do anything differently in
-`dynamic_model_config()` here whether it's run at setup or at runtime.  The
-reason we run it twice is to update the model config options in case the user
-modified `dt_per_km` or `btr_dt_per_km` in the config file in the work 
-directory before running the step.
+We don't do anything differently in `dynamic_model_config()` here whether it's 
+run at setup or at runtime.  The reason we run it twice is to update the model 
+config options in case the user modified `dt_per_km` or `btr_dt_per_km` in the 
+config file in the work directory before running the step.
 
 ### Computing the cell count
 
 In the ocean component, we have infrastructure for determining good values
-for `ntask` and `min_tasks` (the reasonable range of MPI tasks that a forward
+for `ntasks` and `min_tasks` (the reasonable range of MPI tasks that a forward
 model step should use).  Using this infrastructure requires overriding the
 {py:meth}`polaris.ocean.model.OceanModelStep.compute_cell_count()` method:
 
@@ -2045,7 +2044,7 @@ class Forward(OceanModelStep):
 
     ...
 
-    def compute_cell_count(self, at_setup):
+    def compute_cell_count(self):
         """
         Compute the approximate number of cells in the mesh, used to constrain
         resources
@@ -2061,32 +2060,21 @@ class Forward(OceanModelStep):
         cell_count : int or None
             The approximate number of cells in the mesh
         """
-        if at_setup:
-            # no file to read from, so we'll compute it based on config options
-            section = self.config['yet_another_channel']
-            lx = section.getfloat('lx')
-            ly = section.getfloat('ly')
-            nx, ny = compute_planar_hex_nx_ny(lx, ly, self.resolution)
-            cell_count = nx * ny
-        else:
-            # get nCells from the input file
-            with xr.open_dataset('initial_state.nc') as ds:
-                cell_count = ds.sizes['nCells']
+        section = self.config['yet_another_channel']
+        lx = section.getfloat('lx')
+        ly = section.getfloat('ly')
+        nx, ny = compute_planar_hex_nx_ny(lx, ly, self.resolution)
+        cell_count = nx * ny
         return cell_count
 ```
-This method gets called once at setup (with `at_setup=True`) and once at
-runtime (with `at_setup=False`).  We do this because we need to compute or 
-estimate the size of the mesh during setup so we have a good guess at the
+We need to estimate the size of the mesh so we have a good guess at the 
 resources it will need when we add it to a test suite and make a job script for
-running it.  Then, at runtime, we may need to revise that size, either because
-it was an estimate only or because it was based on config options that a user
-may have changed after setup.  Here, we use 
+running it.  Here, we use 
 {py:func}`polaris.mesh.planar.compute_planar_hex_nx_ny()` to get `nx` and `ny`
 (and thus the total cell count) during setup because we have no other way to
-get them (the initial condition has not yet been created, for example).  But
-at runtime, rather than recomputing them, we look at the initial condition
-file that must exist (otherwise the step already would have failed) to see
-what the actual number of cells ended up being.
+get them.  When using task parallelism, we must use this approximation at 
+runtime, because we cannot rely on any tests being completed to use as a basis 
+for computation. 
 
 `cell_count` is used in `OceanModelStep` to compute `ntasks` and `min_tasks`
 by also using 2 config options:

polaris/ocean/model/ocean_model_step.py

@@ -102,7 +102,7 @@ def setup(self):
         self.dynamic_ntasks = (self.ntasks is None and self.min_tasks is None)
 
         if self.dynamic_ntasks:
-            self._update_ntasks(at_setup=True)
+            self._update_ntasks()
 
         super().setup()
 
@@ -111,20 +111,14 @@ def constrain_resources(self, available_cores):
         Update the number of MPI tasks to use based on the estimated mesh size
         """
         if self.dynamic_ntasks:
-            self._update_ntasks(at_setup=False)
+            self._update_ntasks()
         super().constrain_resources(available_cores)
 
-    def compute_cell_count(self, at_setup):
+    def compute_cell_count(self):
         """
         Compute the approximate number of cells in the mesh, used to constrain
         resources
 
-        Parameters
-        ----------
-        at_setup : bool
-            Whether this method is being run during setup of the step, as
-            opposed to at runtime
-
         Returns
         -------
         cell_count : int or None
@@ -193,13 +187,13 @@ def add_streams_file(self, package, streams, template_replacements=None):
         raise ValueError('Input streams files are not supported in '
                          'OceanModelStep')
 
-    def _update_ntasks(self, at_setup):
+    def _update_ntasks(self):
         """
         Update ``ntasks`` and ``min_tasks`` for the step based on the estimated
         mesh size
         """
         config = self.config
-        cell_count = self.compute_cell_count(at_setup)
+        cell_count = self.compute_cell_count()
         if cell_count is None:
             raise ValueError('ntasks and min_tasks were not set explicitly '
                              'but they also cannot be computed because '

polaris/ocean/tests/baroclinic_channel/forward.py

@@ -1,7 +1,5 @@
 import time
 
-import xarray as xr
-
 from polaris.mesh.planar import compute_planar_hex_nx_ny
 from polaris.ocean.model import OceanModelStep
 
@@ -95,33 +93,21 @@ def __init__(self, test_case, resolution, name='forward', subdir=None,
         self.dt = None
         self.btr_dt = None
 
-    def compute_cell_count(self, at_setup):
+    def compute_cell_count(self):
         """
         Compute the approximate number of cells in the mesh, used to constrain
         resources
 
-        Parameters
-        ----------
-        at_setup : bool
-            Whether this method is being run during setup of the step, as
-            opposed to at runtime
-
         Returns
         -------
         cell_count : int or None
             The approximate number of cells in the mesh
         """
-        if at_setup:
-            # no file to read from, so we'll compute it based on config options
-            section = self.config['baroclinic_channel']
-            lx = section.getfloat('lx')
-            ly = section.getfloat('ly')
-            nx, ny = compute_planar_hex_nx_ny(lx, ly, self.resolution)
-            cell_count = nx * ny
-        else:
-            # get nCells from the input file
-            with xr.open_dataset('initial_state.nc') as ds:
-                cell_count = ds.sizes['nCells']
+        section = self.config['baroclinic_channel']
+        lx = section.getfloat('lx')
+        ly = section.getfloat('ly')
+        nx, ny = compute_planar_hex_nx_ny(lx, ly, self.resolution)
+        cell_count = nx * ny
         return cell_count
 
     def dynamic_model_config(self, at_setup):

polaris/ocean/tests/global_convergence/cosine_bell/forward.py

@@ -1,7 +1,5 @@
 import time
 
-import xarray as xr
-
 from polaris.ocean.model import OceanModelStep
 
 
@@ -55,30 +53,18 @@ def __init__(self, test_case, resolution, mesh_name):
 
         self.add_output_file(filename='output.nc')
 
-    def compute_cell_count(self, at_setup):
+    def compute_cell_count(self):
         """
         Compute the approximate number of cells in the mesh, used to constrain
         resources
 
-        Parameters
-        ----------
-        at_setup : bool
-            Whether this method is being run during setup of the step, as
-            opposed to at runtime
-
         Returns
         -------
         cell_count : int or None
             The approximate number of cells in the mesh
         """
-        if at_setup:
-            # use a heuristic based on QU30 (65275 cells) and QU240 (10383
-            # cells)
-            cell_count = 6e8 / self.resolution**2
-        else:
-            # get nCells from the input file
-            with xr.open_dataset('init.nc') as ds:
-                cell_count = ds.sizes['nCells']
+        # use a heuristic based on QU30 (65275 cells) and QU240 (10383 cells)
+        cell_count = 6e8 / self.resolution**2
         return cell_count
 
     def dynamic_model_config(self, at_setup):

polaris/ocean/tests/inertial_gravity_wave/forward.py

@@ -1,7 +1,6 @@
 import time
 
 import numpy as np
-import xarray as xr
 
 from polaris.mesh.planar import compute_planar_hex_nx_ny
 from polaris.ocean.model import OceanModelStep
@@ -61,33 +60,21 @@ def __init__(self, test_case, resolution,
         self.add_yaml_file('polaris.ocean.tests.inertial_gravity_wave',
                            'forward.yaml')
 
-    def compute_cell_count(self, at_setup):
+    def compute_cell_count(self):
         """
         Compute the approximate number of cells in the mesh, used to constrain
         resources
 
-        Parameters
-        ----------
-        at_setup : bool
-            Whether this method is being run during setup of the step, as
-            opposed to at runtime
-
         Returns
         -------
         cell_count : int or None
             The approximate number of cells in the mesh
         """
-        if at_setup:
-            # no file to read from, so we'll compute it based on config options
-            section = self.config['inertial_gravity_wave']
-            lx = section.getfloat('lx')
-            ly = np.sqrt(3.0) / 2.0 * lx
-            nx, ny = compute_planar_hex_nx_ny(lx, ly, self.resolution)
-            cell_count = nx * ny
-        else:
-            # get nCells from the input file
-            with xr.open_dataset('initial_state.nc') as ds:
-                cell_count = ds.sizes['nCells']
+        section = self.config['inertial_gravity_wave']
+        lx = section.getfloat('lx')
+        ly = np.sqrt(3.0) / 2.0 * lx
+        nx, ny = compute_planar_hex_nx_ny(lx, ly, self.resolution)
+        cell_count = nx * ny
         return cell_count
 
     def dynamic_model_config(self, at_setup):

polaris/ocean/tests/manufactured_solution/forward.py

@@ -1,7 +1,6 @@
 import time
 
 import numpy as np
-import xarray as xr
 
 from polaris.mesh.planar import compute_planar_hex_nx_ny
 from polaris.ocean.model import OceanModelStep
@@ -64,33 +63,22 @@ def __init__(self, test_case, resolution,
         self.add_yaml_file('polaris.ocean.tests.manufactured_solution',
                            'forward.yaml')
 
-    def compute_cell_count(self, at_setup):
+    def compute_cell_count(self):
         """
         Compute the approximate number of cells in the mesh, used to constrain
         resources
 
-        Parameters
-        ----------
-        at_setup : bool
-            Whether this method is being run during setup of the step, as
-            opposed to at runtime
-
         Returns
         -------
         cell_count : int or None
             The approximate number of cells in the mesh
         """
-        if at_setup:
-            # no file to read from, so we'll compute it based on config options
-            section = self.config['manufactured_solution']
-            lx = section.getfloat('lx')
-            ly = np.sqrt(3.0) / 2.0 * lx
-            nx, ny = compute_planar_hex_nx_ny(lx, ly, self.resolution)
-            cell_count = nx * ny
-        else:
-            # get nCells from the input file
-            with xr.open_dataset('initial_state.nc') as ds:
-                cell_count = ds.sizes['nCells']
+        # no file to read from, so we'll compute it based on config options
+        section = self.config['manufactured_solution']
+        lx = section.getfloat('lx')
+        ly = np.sqrt(3.0) / 2.0 * lx
+        nx, ny = compute_planar_hex_nx_ny(lx, ly, self.resolution)
+        cell_count = nx * ny
         return cell_count
 
     def dynamic_model_config(self, at_setup):