Switch soma to use ModelStep

This involves breaking the `initial_state` step into 4 steps:
* convert base mesh
* compute initial condition (including land mask) on base mesh
* cull mesh
* compute initial condition on culled mesh

compass/ocean/tests/soma/base_mesh.py

@@ -0,0 +1,56 @@
+import xarray
+
+from mpas_tools.io import write_netcdf
+from mpas_tools.mesh.conversion import convert
+
+from compass.step import Step
+
+
+class BaseMesh(Step):
+    """
+    A step for converting a base mesh to MPAS format for SOMA test cases
+
+    Attributes
+    ----------
+    resolution : str
+        The resolution of the test case
+    """
+
+    def __init__(self, test_case, resolution):
+        """
+        Create the step
+
+        Parameters
+        ----------
+        test_case : compass.TestCase
+            The test case this step belongs to
+
+        resolution : str
+            The resolution of the test case
+        """
+        self.resolution = resolution
+
+        super().__init__(test_case=test_case, name='base_mesh')
+
+        mesh_filenames = {'32km': 'SOMA_32km_grid.161202.nc',
+                          '16km': 'SOMA_16km_grid.161202.nc',
+                          '8km': 'SOMA_8km_grid.161202.nc',
+                          '4km': 'SOMA_4km_grid.161202.nc'}
+        if resolution not in mesh_filenames:
+            raise ValueError(f'Unexpected SOMA resolution: {resolution}')
+
+        self.add_input_file(filename='base_mesh.nc',
+                            target=mesh_filenames[resolution],
+                            database='mesh_database')
+
+        for file in ['mesh.nc', 'base_graph.info']:
+            self.add_output_file(filename=file)
+
+    def run(self):
+        """
+        Convert the base mesh to MPAS format
+        """
+        ds_mesh = convert(xarray.open_dataset('base_mesh.nc'),
+                          graphInfoFileName='base_graph.info',
+                          logger=self.logger)
+        write_netcdf(ds_mesh, 'mesh.nc')

compass/ocean/tests/soma/culled_mesh.py

@@ -0,0 +1,38 @@
+import xarray
+
+from mpas_tools.io import write_netcdf
+from mpas_tools.mesh.conversion import cull
+
+from compass.step import Step
+
+
+class CulledMesh(Step):
+    """
+    A step for culling the base mesh for SOMA test cases
+    """
+
+    def __init__(self, test_case):
+        """
+        Create the step
+
+        Parameters
+        ----------
+        test_case : compass.TestCase
+            The test case this step belongs to
+        """
+        super().__init__(test_case=test_case, name='culled_mesh')
+
+        self.add_input_file(filename='masked_initial_state.nc',
+                            target='../init_on_base_mesh/initial_state.nc')
+
+        for file in ['culled_mesh.nc', 'culled_graph.info']:
+            self.add_output_file(filename=file)
+
+    def run(self):
+        """
+        Run this step of the test case
+        """
+        ds_mesh = cull(xarray.open_dataset('masked_initial_state.nc'),
+                       graphInfoFileName='culled_graph.info',
+                       logger=self.logger)
+        write_netcdf(ds_mesh, 'culled_mesh.nc')

compass/ocean/tests/soma/forward.py

@@ -1,11 +1,10 @@
 import numpy as np
 
-from compass.step import Step
-from compass.model import partition, run_model
+from compass.model import ModelStep
 from compass.ocean.particles import build_particle_simple
 
 
-class Forward(Step):
+class Forward(ModelStep):
     """
     A step for performing forward MPAS-Ocean runs as part of SOMA test cases.
 
@@ -78,7 +77,8 @@ def __init__(self, test_case, resolution, with_particles,
 
         super().__init__(test_case=test_case, name='forward', subdir=None,
                          ntasks=res_params['cores'],
-                         min_tasks=res_params['min_tasks'])
+                         min_tasks=res_params['min_tasks'],
+                         openmp_threads=1)
         # make sure output is double precision
         self.add_streams_file('compass.ocean.streams', 'streams.output')
 
@@ -121,13 +121,13 @@ def __init__(self, test_case, resolution, with_particles,
         self.add_namelist_options(options=options)
 
         self.add_input_file(filename='mesh.nc',
-                            target='../initial_state/culled_mesh.nc')
+                            target='../culled_mesh/culled_mesh.nc')
+        self.add_input_file(filename='graph.info',
+                            target='../culled_mesh/culled_graph.info')
         self.add_input_file(filename='init.nc',
                             target='../initial_state/initial_state.nc')
         self.add_input_file(filename='forcing.nc',
                             target='../initial_state/forcing.nc')
-        self.add_input_file(filename='graph.info',
-                            target='../initial_state/graph.info')
 
         self.add_model_as_input()
 
@@ -137,21 +137,24 @@ def __init__(self, test_case, resolution, with_particles,
             self.add_output_file(
                 filename='analysis_members/lagrPartTrack.0001-01-01_00.00.00.nc')
 
-    def run(self):
+    def runtime_setup(self):
         """
-        Run this step of the test case
+        Write particles initial condition if requested
         """
-        ntasks = self.ntasks
-        partition(ntasks, self.config, self.logger)
+        # we need to partition the graph file first in the
+        # ModelStep.runtime_setup() method, since we need that info for
+        # particles
+        super().runtime_setup()
 
         if self.with_particles:
+            ntasks = self.ntasks
             section = self.config['soma']
             min_den = section.getfloat('min_particle_density')
             max_den = section.getfloat('max_particle_density')
             nsurf = section.getint('surface_count')
+            # todo: This might be better a separate step rather than runtime
+            #       setup depending on how time consuming it is
             build_particle_simple(
                 f_grid='mesh.nc', f_name='particles.nc',
                 f_decomp=f'graph.info.part.{ntasks}',
                 buoySurf=np.linspace(min_den, max_den, nsurf))
-
-        run_model(self, partition_graph=False)

compass/ocean/tests/soma/initial_state.py

@@ -1,15 +1,10 @@
-import xarray
+from compass.model import ModelStep
 
-from mpas_tools.io import write_netcdf
-from mpas_tools.mesh.conversion import convert, cull
 
-from compass.step import Step
-from compass.model import run_model
-
-
-class InitialState(Step):
+class InitialState(ModelStep):
     """
-    A step for creating a mesh and initial condition for SOMA test cases
+    A step for creating a mesh and initial condition, either with or without
+    a mask for culling land cells, for SOMA test cases
 
     Attributes
     ----------
@@ -18,7 +13,7 @@ class InitialState(Step):
     """
 
     def __init__(self, test_case, resolution, with_surface_restoring,
-                 three_layer):
+                 three_layer, mark_land):
         """
         Create the step
 
@@ -36,6 +31,8 @@ def __init__(self, test_case, resolution, with_surface_restoring,
         three_layer : bool
             Whether to use only 3 vertical layers and no continental shelf
 
+        mark_land : bool
+            Whether to mark land cells for culling
         """
         self.resolution = resolution
 
@@ -55,22 +52,15 @@ def __init__(self, test_case, resolution, with_surface_restoring,
 
         res_params = res_params[resolution]
 
-        super().__init__(test_case=test_case, name='initial_state',
-                         ntasks=res_params['cores'],
-                         min_tasks=res_params['min_tasks'])
-
-        mesh_filenames = {'32km': 'SOMA_32km_grid.161202.nc',
-                          '16km': 'SOMA_16km_grid.161202.nc',
-                          '8km': 'SOMA_8km_grid.161202.nc',
-                          '4km': 'SOMA_4km_grid.161202.nc'}
-        if resolution not in mesh_filenames:
-            raise ValueError(f'Unexpected SOMA resolution: {resolution}')
-
-        self.add_input_file(filename='base_mesh.nc',
-                            target=mesh_filenames[resolution],
-                            database='mesh_database')
+        if mark_land:
+            name = 'init_on_base_mesh'
+        else:
+            name = 'initial_state'
 
-        self.add_model_as_input()
+        super().__init__(test_case=test_case, name=name,
+                         ntasks=res_params['cores'],
+                         min_tasks=res_params['min_tasks'],
+                         openmp_threads=1)
 
         package = 'compass.ocean.tests.soma'
 
@@ -86,46 +76,40 @@ def __init__(self, test_case, resolution, with_surface_restoring,
             options['config_soma_vert_levels'] = '60'
             options['config_vertical_grid'] = "'60layerPHC'"
 
-        self.add_namelist_file(package, 'namelist.init', mode='init',
-                               out_name='namelist_mark_land.ocean')
-
-        mark_land_options = dict(
-            config_write_cull_cell_mask='.true.',
-            config_block_decomp_file_prefix="'base_graph.info.part.'",
-            config_proc_decomp_file_prefix="'base_graph.info.part.'")
-
-        mark_land_options.update(options)
-
-        self.add_namelist_options(options=mark_land_options, mode='init',
-                                  out_name='namelist_mark_land.ocean')
-
-        self.add_streams_file(
-            package, 'streams.init', mode='init',
-            template_replacements={'mesh_filename': 'mesh.nc',
-                                   'init_filename': 'masked_initial_state.nc',
-                                   'forcing_filename': 'masked_forcing.nc'},
-            out_name='streams_mark_land.ocean')
+        if mark_land:
+            options['config_write_cull_cell_mask'] = '.true.'
+        else:
+            options['config_write_cull_cell_mask'] = '.false.'
 
         self.add_namelist_file(package, 'namelist.init', mode='init',
                                out_name='namelist.ocean')
-        options['config_write_cull_cell_mask'] = '.false.'
+
         self.add_namelist_options(options=options, mode='init',
                                   out_name='namelist.ocean')
 
         self.add_streams_file(
             package, 'streams.init', mode='init',
-            template_replacements={'mesh_filename': 'culled_mesh.nc',
-                                   'init_filename': 'initial_state.nc',
-                                   'forcing_filename': 'forcing.nc'},
             out_name='streams.ocean')
 
-        for file in ['initial_state.nc', 'forcing.nc', 'graph.info']:
+        if mark_land:
+            self.add_input_file(filename='mesh.nc',
+                                target='../base_mesh/mesh.nc')
+            self.add_input_file(filename='graph.info',
+                                target='../base_mesh/base_graph.info')
+        else:
+            self.add_input_file(filename='mesh.nc',
+                                target='../culled_mesh/culled_mesh.nc')
+            self.add_input_file(filename='graph.info',
+                                target='../culled_mesh/culled_graph.info')
+
+        for file in ['initial_state.nc', 'forcing.nc']:
             self.add_output_file(filename=file)
 
-    def run(self):
+    def runtime_setup(self):
         """
-        Run this step of the test case
+        Set namelist options from config options
         """
+        super().runtime_setup()
 
         config = self.config
         section = config['soma']
@@ -142,21 +126,5 @@ def run(self):
             config_soma_shelf_depth=section.get('shelf_depth'),
             config_soma_bottom_depth=section.get('bottom_depth'))
 
-        for out_name in ['namelist_mark_land.ocean', 'namelist.ocean']:
-            self.update_namelist_at_runtime(options=options, out_name=out_name)
-        ds_mesh = convert(xarray.open_dataset('base_mesh.nc'),
-                          graphInfoFileName='base_graph.info',
-                          logger=self.logger)
-        write_netcdf(ds_mesh, 'mesh.nc')
-
-        run_model(self, namelist='namelist_mark_land.ocean',
-                  streams='streams_mark_land.ocean',
-                  graph_file='base_graph.info')
-
-        ds_mesh = cull(xarray.open_dataset('masked_initial_state.nc'),
-                       graphInfoFileName='graph.info',
-                       logger=self.logger)
-        write_netcdf(ds_mesh, 'culled_mesh.nc')
-
-        run_model(self, namelist='namelist.ocean', streams='streams.ocean',
-                  graph_file='graph.info')
+        self.update_namelist_at_runtime(options=options,
+                                        out_name='namelist.ocean')

compass/ocean/tests/soma/soma_test_case.py

@@ -1,4 +1,6 @@
 from compass.testcase import TestCase
+from compass.ocean.tests.soma.base_mesh import BaseMesh
+from compass.ocean.tests.soma.culled_mesh import CulledMesh
 from compass.ocean.tests.soma.initial_state import InitialState
 from compass.ocean.tests.soma.forward import Forward
 from compass.ocean.tests.soma.analysis import Analysis
@@ -89,10 +91,22 @@ def __init__(self, test_group, resolution, with_particles,
 
         super().__init__(test_group=test_group, name=name, subdir=subdir)
 
+        self.add_step(BaseMesh(
+            test_case=self, resolution=resolution))
+
         self.add_step(InitialState(
             test_case=self, resolution=resolution,
             with_surface_restoring=with_surface_restoring,
-            three_layer=three_layer))
+            three_layer=three_layer, mark_land=True))
+
+        self.add_step(CulledMesh(
+            test_case=self))
+
+        self.add_step(InitialState(
+            test_case=self, resolution=resolution,
+            with_surface_restoring=with_surface_restoring,
+            three_layer=three_layer, mark_land=False))
+
         options = dict()
         if with_surface_restoring:
             options['config_use_activeTracers_surface_restoring'] = '.true.'

compass/ocean/tests/soma/streams.init

@@ -2,18 +2,18 @@
 
 <immutable_stream name="mesh"
                   type="input"
-                  filename_template="{{ mesh_filename }}"
+                  filename_template="mesh.nc"
                   input_interval="initial_only" />
 
 <immutable_stream name="input_init"
-        filename_template="{{ mesh_filename }}">
+        filename_template="mesh.nc">
 </immutable_stream>
 
 <stream name="output_init"
         type="output"
         output_interval="0000_00:00:01"
         clobber_mode="truncate"
-        filename_template="{{ init_filename }}">
+        filename_template="initial_state.nc">
 
     <stream name="input_init"/>
     <var_struct name="tracers"/>
@@ -38,7 +38,7 @@
         type="output"
         output_interval="0000_00:00:01"
         clobber_mode="truncate"
-        filename_template="{{ forcing_filename }}">
+        filename_template="forcing.nc">
 
     <var_array name="activeTracersPistonVelocity"/>
     <var_array name="activeTracersSurfaceRestoringValue"/>