PR review changes

compass/ocean/tests/baroclinic_gyre/baroclinic_gyre.cfg

@@ -5,13 +5,13 @@
 # the type of vertical grid
 grid_type = linear_dz
 
-# the linear rate of thickness increase for linear_dz
+# the linear rate of thickness (m) increase for linear_dz
 linear_dz_rate = 10.
 
 # Number of vertical levels
 vert_levels = 15
 
-# Total water column depth
+# Total water column depth in m
 bottom_depth = 1800.
 
 # The type of vertical coordinate (e.g. z-level, z-star)
@@ -32,24 +32,26 @@ lat_max = 75
 lon_min = 0
 lon_max = 60
 
-# Initial vertical temperature profile
+# Initial vertical temperature profile (C)
 initial_temp_top = 33.
 initial_temp_bot = 1.
 
 # Constant salinity value (also used in restoring)
 initial_salinity = 34.
 
-# Maximum zonal wind stress value
+# Maximum zonal wind stress value (N m-2)
 wind_stress_max = 0.1
 
-# Surface temperature restoring
+# Surface temperature restoring profile
 restoring_temp_min = 0.
 restoring_temp_max = 30.
 
 # Restoring timescale for surface temperature (in days)
 restoring_temp_timescale = 30.
 
-# config options for the mean state visualization
-[mean_state_viz]
+# config options for the post processing (moc and viz)
+[baroclinic_gyre_post]
+# latitude bin increment for the moc calculation
+dlat = 0.25
 # number of years to average over for the mean state plots
 time_averaging_length = 1

compass/ocean/tests/baroclinic_gyre/cull_mesh.py

@@ -8,7 +8,7 @@
 
 class CullMesh(Step):
     """
-    Cull a global mesh to only a signle basin
+    Cull a global mesh to only a single basin
     Attributes
     ----------
     """

compass/ocean/tests/baroclinic_gyre/forward.py

@@ -89,9 +89,7 @@ def __init__(self, test_case, resolution, name='forward', subdir=None,
             options[f'config_{option}'] = res_params[option]
         if long:
             # run for 3 years instead of 3 time steps
-            options['config_start_time'] = "'0001-01-01_00:00:00'"
-            options['config_stop_time'] = "'0004-01-01_00:00:00'"
-            options['config_run_duration'] = "'none'"
+            options['config_run_duration'] = "'0003-00-00_00:00:00'"
 
         self.add_namelist_options(options=options)
         self.add_input_file(filename='init.nc',

compass/ocean/tests/baroclinic_gyre/initial_state.py

@@ -49,7 +49,6 @@ def run(self):
         dsMesh = xr.open_dataset('culled_mesh.nc')
 
         ds = _write_initial_state(config, dsMesh)
-        print('bottomDepth0', ds.refBottomDepth[0])
         _write_forcing(config, ds.latCell, ds.refBottomDepth)
 
 
@@ -69,22 +68,17 @@ def _write_initial_state(config, dsMesh):
     temp_top = section.getfloat('initial_temp_top')
     temp_bot = section.getfloat('initial_temp_bot')
 
-    print(f'zmid: {ds.zMid[0,0,:].values}')
     cc = 0.049  # 0.049 value from optimization on 1800m
     aa = temp_top / np.log(cc)
     bb = (cc ** (temp_bot / temp_top) - cc)
     temperature = aa * np.log(bb * -1.0 * ds.zMid / bottom_depth + cc)
-    print(f'Tinit: {temperature[0,0,:].values}')
     val_bot = aa * np.log(bb + cc)
     val_top = aa * np.log(cc)
     print(f'analytical bottom T: {val_bot} at '
           f'depth : {bottom_depth}')
     print(f'analytical surface T: {val_top} at '
           f'depth = 0')
 
-#    temperature = (-11. * np.log(0.0414 *
-#                   (-1. * ds.zMid + 100.3)) + 48.8)
-#    temperature = temperature.transpose('Time', 'nCells', 'nVertLevels')
     print(f'bottom layer T: {temperature[0, 0, -1]} and '
           f'surface layer T: {temperature[0, 0, 0]}')
     salinity = initial_salinity * xr.ones_like(temperature)

compass/ocean/tests/baroclinic_gyre/moc.py

@@ -43,28 +43,19 @@ def run(self):
         in_dir = '../forward/output'
         out_dir = '.'
 
-        # show progress only if we're not writing to a log file
-        # show_progress = self.log_filename is None
-
         lat_min = self.config.getfloat('baroclinic_gyre', 'lat_min')
         lat_max = self.config.getfloat('baroclinic_gyre', 'lat_max')
-        dlat = 0.25  # set in config next
-        # latbins = np.arange(15.5, 75.5, 0.25)
+        dlat = self.config.getfloat('baroclinic_gyre_post', 'dlat')
         latBins = np.arange(lat_min + 2 * dlat,
                             lat_max + 2 * dlat, dlat)
         nz = self.config.getint('vertical_grid', 'vert_levels')
 
         dsMesh = xarray.open_dataset('../initial_state/initial_state.nc')
 
         ds = xarray.open_mfdataset(
-            '{}/timeSeriesStatsMonthly*.nc'.format(in_dir),
+            f'{in_dir}/timeSeriesStatsMonthly*.nc',
             concat_dim='Time', combine='nested')
 
-        # print(dsMesh)
-        # print(ds)
-        # print(latBins)
-        # print(nz)
-
         moc = self._compute_amoc(dsMesh, ds, latBins, nz)
 
         dsMOC = xarray.Dataset()
@@ -76,9 +67,8 @@ def run(self):
         dsMOC.moc.attrs['units'] = 'Sv'
         dsMOC.moc.attrs['description'] = \
             'zonally-averaged meridional overturning streamfunction'
-        outputFileName = '{}/moc.nc'.format(out_dir)
-        # if file_complete(ds, outputFileName):
-        #    return
+
+        outputFileName = f'{out_dir}/moc.nc'
         write_netcdf(dsMOC, outputFileName)
 
     def _compute_amoc(self, dsMesh, ds, latBins, nz):

compass/ocean/tests/baroclinic_gyre/viz.py

@@ -50,7 +50,8 @@ def _plot_moc(self, ds, dsMesh, out_dir):
         """
         Plot the time-mean moc state for the test case
         """
-        avg_len = self.config.getint('mean_state_viz', 'time_averaging_length')
+        avg_len = self.config.getint('baroclinic_gyre_post',
+                                     'time_averaging_length')
         moc = ds.moc[-12 * avg_len:, :, :].mean(axis=0).T.values
         latbins = ds.latBins
         plt.contourf(
@@ -60,16 +61,21 @@ def _plot_moc(self, ds, dsMesh, out_dir):
         plt.ylabel('Depth (m)')
         plt.xlabel('Latitude')
         idx = np.unravel_index(np.argmax(moc), moc.shape)
-        amoc = "max MOC = {:.1e}".format(round(np.max(moc), 1))
-        maxloc = 'at lat = {} and z = {}m'.format(
-            latbins[idx[-1]].values, int(dsMesh.refInterfaces[idx[0]].values))
-        maxval = 'max MOC = {:.1e} at lat={}-{}'.format(
-            round(np.max(moc[:, 175]), 1),
-            latbins[175 - 1].values, latbins[175].values)
+        max_moc = round(np.max(moc), 1)
+        max_moc_loc = [latbins[idx[-1]].values,
+                       int(dsMesh.refInterfaces[idx[0]].values)]
+        ref_moc_loc = 175  # to be improved
+        max_moc_ref = round(np.max(moc[:, ref_moc_loc]), 1)
+        ref_moc_lat = [latbins[175 - 1].values, latbins[175].values]
+
+        amoc = f"max MOC = {max_moc:.1e}"
+        maxloc = f'at lat = {max_moc_loc[0]} and z = {max_moc_loc[1]}m'
+        maxval = (f'max MOC = {max_moc_ref:.1e} at '
+                  f'lat={ref_moc_lat[0]}-{ref_moc_lat[1]}')
         plt.annotate(amoc + '\n' + maxloc + '\n' + maxval,
                      xy=(0.01, 0.05), xycoords='axes fraction')
         plt.colorbar()
-        plt.savefig('{}/time_avg_moc_last{}years.png'.format(out_dir, avg_len))
+        plt.savefig(f'{out_dir}/time_avg_moc_last{avg_len}years.png')
 
     def _plot_spinup(self, mon_dir, dsMesh, out_dir):
         """
@@ -78,7 +84,7 @@ def _plot_spinup(self, mon_dir, dsMesh, out_dir):
         """
 
         ds = xarray.open_mfdataset(
-            '{}/timeSeriesStatsMonthly*.nc'.format(mon_dir),
+            f'{mon_dir}/timeSeriesStatsMonthly*.nc',
             concat_dim='Time', combine='nested')
         KE = ds.timeMonthly_avg_kineticEnergyCell[:, :, :].mean(axis=1)
         T = ds.timeMonthly_avg_activeTracers_temperature[:, :, :].mean(axis=1)
@@ -89,27 +95,28 @@ def _plot_spinup(self, mon_dir, dsMesh, out_dir):
 
         fig, ax = plt.subplots(2, 1, figsize=(6, 8))
         for ll in [0, 3, 6, 10, 14]:
-            ax[0].plot(KE[:, ll], label='{}m'.format(int(midlayer[ll])))
-            ax[1].plot(T[:, ll], label='{}m'.format(int(midlayer[ll])))
+            ax[0].plot(KE[:, ll], label=f'{int(midlayer[ll])}m')
+            ax[1].plot(T[:, ll], label=f'{int(midlayer[ll])}m')
         ax[0].legend()
         ax[1].legend()
         ax[0].set_xlabel('Time (months)')
         ax[1].set_xlabel('Time (months)')
         ax[0].set_ylabel('Layer Mean Kinetic Energy ($m^2 s^{-2}$)')
         ax[1].set_ylabel(r'Layer Mean Temperature ($^{\circ}$C)')
-        plt.savefig('{}/spinup_ft.png'.format(out_dir), bbox_inches='tight')
+        plt.savefig(f'{out_dir}/spinup_ft.png', bbox_inches='tight')
 
     def _plot_mean_surface_state(self, mon_dir, dsMesh, out_dir):
 
         lon = 180. / np.pi * dsMesh.variables['lonCell'][:]
         lat = 180. / np.pi * dsMesh.variables['latCell'][:]
         ds = xarray.open_mfdataset(
-            '{}/timeSeriesStatsMonthly*.nc'.format(mon_dir),
+            f'{mon_dir}/timeSeriesStatsMonthly*.nc',
             concat_dim='Time', combine='nested')
         heatflux = (
             ds.timeMonthly_avg_activeTracersSurfaceFlux_temperatureSurfaceFlux[:, :] *  # noqa: E501
             constants['SHR_CONST_CPSW'] * constants['SHR_CONST_RHOSW'])
-        avg_len = self.config.getint('mean_state_viz', 'time_averaging_length')
+        avg_len = self.config.getint('baroclinic_gyre_post',
+                                     'time_averaging_length')
         absmax = np.max(np.abs(np.mean(heatflux[-12 * avg_len:, :].values, axis=0)))  # noqa: E501
         fig, ax = plt.subplots(1, 3, figsize=[18, 5])
         ax[0].tricontour(lon, lat, np.mean(ds.timeMonthly_avg_ssh[-12 * avg_len:, :], axis=0),  # noqa: E501
@@ -131,10 +138,10 @@ def _plot_mean_surface_state(self, mon_dir, dsMesh, out_dir):
 
         ax[0].set_title('SSH (m)')
         ax[1].set_title(r'SST ($^\circ$C)')
-        ax[2].set_title('Heat Flux (W/m^{2})')
+        ax[2].set_title('Heat Flux (W/m$^{2}$)')
 
         ax[0].set_ylabel(r'Latitude ($^\circ$)')
         for axis in ax:
             axis.set_xlabel(r'Longitude ($^\circ$)')
-        plt.savefig('{}/meansurfacestate_last{}years.png'.format(
-            out_dir, avg_len), bbox_inches='tight')
+        plt.savefig(f'{out_dir}/meansurfacestate_last{avg_len}years.png',
+                    bbox_inches='tight')

docs/developers_guide/ocean/api.rst

@@ -96,12 +96,20 @@ baroclinic_gyre
    GyreTestCase
    GyreTestCase.configure
 
+   cull_mesh.CullMesh
+   cull_mesh.CullMesh.run
+
    forward.Forward
    forward.Forward.run
 
    initial_state.InitialState
    initial_state.InitialState.run
 
+   moc.Moc
+   moc.Moc.run
+
+   viz.Viz
+   viz.Viz.run
 
 dam_break
 ~~~~~~~~~

docs/developers_guide/ocean/test_groups/baroclinic_gyre.rst

@@ -7,39 +7,58 @@ The ``baroclinic_gyre`` test group implements variants of the
 Baroclinic ocean gyre set-up from the 
 `MITgcm test case <https://mitgcm.readthedocs.io/en/latest/examples/baroclinic_gyre/baroclinic_gyre.html>`_.
 
-This test case is described in detail in the User guide (see :ref:`baroclinic_gyre`). Here,
-we describe the shared framework for this test group.
+This test case is described in detail in the User guide
+(see :ref:`baroclinic_gyre`). Here, we describe the shared framework
+for this test group.
 
 Framework
 ---------
 
 At this stage, the test case is available at 80-km and 20-km horizontal
-resolution.  By default, the 15 vertical layers vary linearly in thickness with depth, from 50m at the surface to 190m at depth (full depth: 1800m).
+resolution.  By default, the 15 vertical layers vary linearly in thickness
+with depth, from 50m at the surface to 190m at depth (full depth: 1800m).
 
-The test group includes 2 test cases, called ``performance_test`` for a short (10-day) run, and ``3_year_test`` for a 3-year simulation. Note that 3 years are insufficient to bring the standard test case to full equilibrium.  Both test cases have 2 steps,
-``initial_state``, which defines the mesh and initial conditions for the model,
-and ``forward``, which performs the time integration of the model.
+The test group includes 2 test cases, called ``performance_test`` for a short
+(3-time-step) run, and ``3_year_test`` for a 3-year simulation. Note that 
+3 years are insufficient to bring the standard test case to full equilibrium.
+Both test cases have 2 steps, ``initial_state``, which defines the mesh and 
+initial conditions for the model, and ``forward``, which performs the time 
+integration of the model.
 
 Additionally, the test group has a shared ``namelist.forward`` file with
 a few common namelist options related to horizontal
 and vertical momentum and tracer diffusion, as well as a shared
 ``streams.forward`` file that defines ``mesh``, ``input``, ``restart``, and
 ``output`` streams. 
 
+cull_mesh
+~~~~~~~~~
+
+The class :py:class:`compass.ocean.tests.baroclinic_gyre.cull_mesh.CullMesh`
+defines a step for setting up the mesh for the baroclinic gyre test case.
+
+First, a mesh appropriate for the resolution is generated using
+:py:class:`compass.mesh.QuasiUniformSphericalMeshStep`.  Then, the mesh is
+culled to keep a single ocean basin (with lat, lon bounds set in `.cfg` file).
+
 initial_state
 ~~~~~~~~~~~~~
 
 The class :py:class:`compass.ocean.tests.baroclinic_gyre.initial_state.InitialState`
 defines a step for setting up the initial state for each test case.
 
-First, a mesh appropriate for the resolution is generated using
-:py:class:`compass.mesh.QuasiUniformSphericalMeshStep`.  Then, the mesh is
-culled to keep a single ocean basin (with lat, lon bounds set in `.cfg` file).  A vertical grid is generated,
-with 15 layers of thickness that increases linearly with depth (10m increase by default), from 50m at the surface to 190m at depth (full depth: 1800m).
-Finally, the initial temperature field is initialized with a vertical profile to approximate the discrete values set in the `MITgcm test case <https://mitgcm.readthedocs.io/en/latest/examples/baroclinic_gyre/baroclinic_gyre.html>`_, uniform in horizontal space. The surface and bottom values are set in the `.cfg` file. Salinity is set to a constant value (34 psu, set in the `.cfg` file)  and initial
-velocity is set to zero. 
+A vertical grid is generated, with 15 layers of thickness that increases 
+linearly with depth (10m increase by default), from 50m at the surface to 190m 
+at depth (full depth: 1800m). Finally, the initial temperature field is 
+initialized with a vertical profile to approximate the discrete values set in 
+the `MITgcm test case <https://mitgcm.readthedocs.io/en/latest/examples/baroclinic_gyre/baroclinic_gyre.html>`_,
+uniform in horizontal space. The surface and bottom values are set in the 
+`.cfg` file. Salinity is set to a constant value (34 psu, set in the `.cfg` 
+file)  and initial velocity is set to zero. 
 
-The ``initial_state`` step also generates the forcing, defined as zonal wind stress that varies with latitude, surface temperature restoring that varies with latitutde, and writes it to `forcing.nc`.
+The ``initial_state`` step also generates the forcing, defined as zonal wind 
+stress that varies with latitude, surface temperature restoring that varies 
+with latitutde, and writes it to `forcing.nc`.
 
 forward
 ~~~~~~~
@@ -60,11 +79,47 @@ resolutions are supported.
 -----------
 
 ``ocean/baroclinic_gyre/80km/3_year_test`` performs a longer (3 year) integration
-of the model forward in time. The point is to (ultimately) compare the quasi-steady state with theroretical scaling and results from other models. Currently, only the 80-km and 20-km horizontal
-resolutions are supported. Note that 3 years is not long enough to reach steady state.
+of the model forward in time. The point is to (ultimately) compare the quasi-
+steady state with theroretical scaling and results from other models. 
+Currently, only the 80-km and 20-km horizontal resolutions are supported. 
+Note that 3 years is not long enough to reach steady state.
 For the 80km configuration, the estimated time to equilibration is roughly 50 years.
 For a detailed comparison of the mean state against theory and results from other models,
 the mean state at 100 years may be most appropriate to be aligned with the MITgcm results. 
 
+moc
+---
+
+The class :py:class:`compass.ocean.tests.baroclinic_gyre.moc.Moc`
+defines a step for computing the zonally-averaged meridional overturning
+in the baroclinic gyre test case. 
+
+It calculates the moc by summing the monthly mean vertical velocities
+in latitude bins (the bin width ``dlat`` is set in the ``.cfg`` file). 
+It outputs a netcdf file ``moc.nc`` in the local folder, with the monthly
+value of zonally-averaged meriodional overturning in Sv. 
+
+By default, this step is only called after the ``3_year_test`` step since
+the ``performance_test`` is too short to output monthly mean output.
+
+viz
+---
+
+The class :py:class:`compass.ocean.tests.baroclinic_gyre.viz.Viz`
+defines a step for visualizing output from baroclinic gyre.
+Currently, it includes 3 separate plots.
+
+First, it plots 2 timeseries to assess the test case spin-up: the top plot is 
+layer-mean kinetic energy over 5 different layers (to show dynamical spin-up),
+while the bottom plot is layer mean temperature (to show thermal adjustment).
+The values plotted are all basin-averaged monthly output. 
 
+It also plots the time-averaged final state of the simulation. The averaging 
+period is set up in the ``.cfg`` file (by default, 1 year). Plots include
+time-mean sea surface height (SSH), sea surface temperature, and surface heat
+flux (in W/m2, due to surface restoring) with contours of SSH superimposed. 
+A separate figure is generated showing the time-mean overturning streamfunction,
+based on the ``moc.nc`` file generated in the ``moc`` step described above. 
 
+By default, this step is only called after the ``3_year_test`` step since
+the ``performance_test`` is too short to output monthly mean output.