Merge pull request #516 from cbegeman/cleanup-ocean-rpe-analysis

Cleanup ocean RPE tests

compass/ocean/rpe.py

@@ -9,18 +9,19 @@ def compute_rpe(initial_state_file_name='initial_state.nc',
     Computes the reference (resting) potential energy for the whole domain
     Parameters
     ----------
-    initial_state_file_name : string
+    initial_state_file_name : str
         Name of the netCDF file containing the initial state
 
-    output_file_prefix : string
+    output_file_prefix : str
         Prefix for the netCDF file containing output of forward step
 
     num_files : int
         Number of files on which to compute rpe
 
     Returns
     -------
-    rpe : numpy array of size (num_timesteps) x (num_files)
+    rpe : numpy.ndarray
+          the reference potential energy of size (num_timesteps) x (num_files)
     """
     gravity = 9.80616
 

compass/ocean/tests/baroclinic_channel/baroclinic_channel.cfg

@@ -26,6 +26,9 @@ min_pc_fraction = 0.1
 # or fractions. False means fractions.
 use_distances = False
 
+# Viscosity values to test for rpe test case
+viscosities = 1, 5, 10, 20, 200
+
 # Temperature of the surface in the northern half of the domain.
 surface_temperature = 13.1
 
@@ -48,4 +51,4 @@ gradient_width_dist = 40e3
 salinity = 35.0
 
 # Coriolis parameter for entire domain.
-coriolis_parameter = -1.2e-4
+coriolis_parameter = -1.2e-4

compass/ocean/tests/baroclinic_channel/rpe_test/__init__.py

@@ -33,8 +33,16 @@ def __init__(self, test_group, resolution):
         subdir = f'{resolution}/{name}'
         super().__init__(test_group=test_group, name=name,
                          subdir=subdir)
+        self.resolution = resolution
 
-        nus = [1, 5, 10, 20, 200]
+
+    def configure(self):
+        """
+        Modify the configuration options for this test case.
+        """
+        resolution = self.resolution
+        config = self.config
+        baroclinic_channel.configure(resolution, config)
 
         res_params = {'1km': {'ntasks': 144, 'min_tasks': 36},
                       '4km': {'ntasks': 36, 'min_tasks': 8},
@@ -47,13 +55,12 @@ def __init__(self, test_group, resolution):
 
         params = res_params[resolution]
 
-        self.resolution = resolution
-
         self.add_step(
             InitialState(test_case=self, resolution=resolution))
 
+        nus = config.getlist('baroclinic_channel', 'viscosities', dtype=float)
         for index, nu in enumerate(nus):
-            name = 'rpe_test_{}_nu_{}'.format(index + 1, nu)
+            name = f'rpe_test_{index + 1}_nu_{int(nu)}'
             step = Forward(
                 test_case=self, name=name, subdir=name,
                 ntasks=params['ntasks'], min_tasks=params['min_tasks'],
@@ -69,11 +76,3 @@ def __init__(self, test_group, resolution):
 
         self.add_step(
             Analysis(test_case=self, resolution=resolution, nus=nus))
-
-    def configure(self):
-        """
-        Modify the configuration options for this test case.
-        """
-        baroclinic_channel.configure(self.resolution, self.config)
-
-    # no run() is needed because we're doing the default: running all steps

compass/ocean/tests/baroclinic_channel/rpe_test/analysis.py

@@ -45,11 +45,11 @@ def __init__(self, test_case, resolution, nus):
 
         for index, nu in enumerate(nus):
             self.add_input_file(
-                filename='output_{}.nc'.format(index+1),
-                target='../rpe_test_{}_nu_{}/output.nc'.format(index+1, nu))
+                filename=f'output_{index+1}.nc',
+                target=f'../rpe_test_{index+1}_nu_{int(nu)}/output.nc')
 
         self.add_output_file(
-            filename='sections_baroclinic_channel_{}.png'.format(resolution))
+            filename=f'sections_baroclinic_channel_{resolution}.png')
         self.add_output_file(filename='rpe_t.png')
 
     def run(self):
@@ -81,7 +81,8 @@ def _plot(nx, ny, filename, nus, rpe):
     nus : list of float
         The viscosity values
 
-    rpe : float, dim len(nu) x len(time)
+    rpe : numpy.ndarray
+        The reference potential energy with size len(nu) x len(time)
     """
 
     plt.switch_backend('Agg')
@@ -98,7 +99,7 @@ def _plot(nx, ny, filename, nus, rpe):
     for i in range(num_files):
         rpe_norm = np.divide((rpe[i, :]-rpe[i, 0]), rpe[i, 0])
         plt.plot(times, rpe_norm,
-                 label="$\\nu_h=${}".format(nus[i]))
+                 label=f"$\\nu_h=${nus[i]}")
     plt.xlabel('Time, days')
     plt.ylabel('RPE-RPE(0)/RPE(0)')
     plt.legend()
@@ -109,7 +110,7 @@ def _plot(nx, ny, filename, nus, rpe):
         2.1 * num_files, 5.0), constrained_layout=True)
 
     for iCol in range(num_files):
-        ds = xarray.open_dataset('output_{}.nc'.format(iCol + 1))
+        ds = xarray.open_dataset(f'output_{iCol + 1}.nc')
         times = ds.daysSinceStartOfSim.values
         times = np.divide(times.tolist(), nanosecondsPerDay)
         tidx = np.argmin(np.abs(times-time))
@@ -132,8 +133,8 @@ def _plot(nx, ny, filename, nus, rpe):
             cmap='cmo.thermal',
             vmin=11.8,
             vmax=13.0)
-        ax.set_title("day {}, $\\nu_h=${}".format(
-            int(times[tidx]), nus[iCol]))
+        ax.set_title(f'day {times[tidx]}, '
+                     f'$\\nu_h=${nus[iCol]}')
         ax.set_xticks(np.arange(0, 161, step=40))
         ax.set_yticks(np.arange(0, 501, step=50))
 

compass/ocean/tests/internal_wave/internal_wave.cfg

@@ -29,6 +29,9 @@ ny = 50
 # the size of grid cells (m)
 dc = 5000.0
 
+# Viscosity values to test for rpe test case
+viscosities = 0.01, 1, 15, 150
+
 # Logical flag that determines if locations of features are defined by distance
 # or fractions. False means fractions.
 use_distances = False

compass/ocean/tests/internal_wave/rpe_test/__init__.py

@@ -23,12 +23,17 @@ def __init__(self, test_group):
         name = 'rpe_test'
         super().__init__(test_group=test_group, name=name)
 
-        nus = [0.01, 1, 15, 150]
 
+    def configure(self):
+        """
+        Modify the configuration options for this test case.
+        """
+        config = self.config
         self.add_step(InitialState(test_case=self))
 
+        nus = config.getlist('internal_wave', 'viscosities', dtype=float)
         for index, nu in enumerate(nus):
-            name = 'rpe_test_{}_nu_{:g}'.format(index + 1, nu)
+            name = f'rpe_test_{index + 1}_nu_{nu:g}'
             step = Forward(
                 test_case=self, name=name, subdir=name, ntasks=4,
                 openmp_threads=1, nu=float(nu))

compass/ocean/tests/internal_wave/rpe_test/analysis.py

@@ -41,8 +41,8 @@ def __init__(self, test_case, nus):
 
         for index, nu in enumerate(nus):
             self.add_input_file(
-                filename='output_{}.nc'.format(index+1),
-                target='../rpe_test_{}_nu_{}/output.nc'.format(index+1, nu))
+                filename=f'output_{index+1}.nc',
+                target=f'../rpe_test_{index+1}_nu_{nu:g}/output.nc')
 
         self.add_output_file(
             filename='sections_internal_wave.png')
@@ -52,32 +52,24 @@ def run(self):
         """
         Run this step of the test case
         """
-        section = self.config['internal_wave']
-        nx = section.getint('nx')
-        ny = section.getint('ny')
         rpe = compute_rpe(num_files=len(self.nus))
-        _plot(nx, ny, self.outputs[0], self.nus, rpe)
+        _plot(self.outputs[0], self.nus, rpe)
 
 
-def _plot(nx, ny, filename, nus, rpe):
+def _plot(filename, nus, rpe):
     """
     Plot section of the internal wave at different viscosities
 
     Parameters
     ----------
-    nx : int
-        The number of cells in the x direction
-
-    ny : int
-        The number of cells in the y direction (before culling)
-
     filename : str
         The output file name
 
     nus : list of float
         The viscosity values
 
-    rpe : float, dim len(nu) x len(time)
+    rpe : numpy.ndarray
+        The reference potential energy with size len(nu) x len(time)
     """
 
     plt.switch_backend('Agg')
@@ -94,7 +86,7 @@ def _plot(nx, ny, filename, nus, rpe):
     for i in range(num_files):
         rpe_norm = np.divide((rpe[i, :]-rpe[i, 0]), rpe[i, 0])
         plt.plot(times, rpe_norm,
-                 label="$\\nu_h=${}".format(nus[i]))
+                 label=f"$\\nu_h=${nus[i]}")
     plt.xlabel('Time, days')
     plt.ylabel('RPE-RPE(0)/RPE(0)')
     plt.legend()
@@ -108,7 +100,7 @@ def _plot(nx, ny, filename, nus, rpe):
         4.0 * nCol, 3.7 * nRow), constrained_layout=True)
 
     for iRow in range(nRow):
-        ds = xarray.open_dataset('output_{}.nc'.format(iRow + 1))
+        ds = xarray.open_dataset(f'output_{iRow + 1}.nc')
         times = ds.daysSinceStartOfSim.values
         times = np.divide(times.tolist(), nanosecondsPerDay)
         var = ds.temperature.values
@@ -128,6 +120,6 @@ def _plot(nx, ny, filename, nus, rpe):
                 ax.set_ylabel('depth, m')
             if iCol == nCol - 1:
                 fig.colorbar(dis, ax=axs[iRow, iCol], aspect=10)
-            ax.set_title("day {}, $\\nu_h=${}".format(time[iCol], nus[iRow]))
+            ax.set_title(f"day {time[iCol]}, $\\nu_h=${nus[iRow]}")
 
     plt.savefig('sections_internal_wave.png')

docs/users_guide/ocean/test_groups/baroclinic_channel.rst

@@ -62,6 +62,9 @@ All 5 test cases share the same set of config options:
     # or fractions. False means fractions.
     use_distances = False
 
+    # Viscosity values to test for rpe test case
+    viscosities = 1, 5, 10, 20, 200
+
     # Temperature of the surface in the northern half of the domain.
     surface_temperature = 13.1
 

docs/users_guide/ocean/test_groups/internal_wave.rst

@@ -66,6 +66,9 @@ All 3 test cases share the same set of config options:
     # the size of grid cells (m)
     dc = 5000.0
 
+    # Viscosity values to test for rpe test case
+    viscosities = 0.01, 1, 15, 150
+
     # Logical flag that determines if locations of features are defined by distance
     # or fractions. False means fractions.
     use_distances = False