Merge pull request #123 from xylar/clean-up-rpe-analysis

Clean up RPE analysis step and framework

polaris/ocean/rpe.py

@@ -5,38 +5,42 @@
 from mpas_tools.cime.constants import constants
 
 
-def compute_rpe(initial_state_file_name, output_files):
+def compute_rpe(mesh_filename, initial_state_filename, output_filenames):
     """
     Computes the reference (resting) potential energy for the whole domain
 
     Parameters
     ----------
-    initial_state_file_name : str
+    mesh_filename : str
+        Name of the netCDF file containing the MPAS horizontal mesh variables
+
+    initial_state_filename : str
         Name of the netCDF file containing the initial state
 
-    output_files : list
+    output_filenames : list
         List of netCDF files containing output of forward steps
 
     Returns
     -------
     rpe : numpy.ndarray
         the reference potential energy of size ``Time`` x ``len(output_files)``
     """
-    num_files = len(output_files)
+    num_files = len(output_filenames)
     if num_files == 0:
         raise ValueError('Must provide at least one output filename')
 
     gravity = constants['SHR_CONST_G']
 
-    dsInit = xr.open_dataset(initial_state_file_name)
-    nVertLevels = dsInit.sizes['nVertLevels']
+    ds_mesh = xr.open_dataset(mesh_filename)
+    ds_init = xr.open_dataset(initial_state_filename)
+    nVertLevels = ds_init.sizes['nVertLevels']
 
-    xEdge = dsInit.xEdge
-    yEdge = dsInit.yEdge
-    areaCell = dsInit.areaCell
-    minLevelCell = dsInit.minLevelCell - 1
-    maxLevelCell = dsInit.maxLevelCell - 1
-    bottomDepth = dsInit.bottomDepth
+    xEdge = ds_mesh.xEdge
+    yEdge = ds_mesh.yEdge
+    areaCell = ds_mesh.areaCell
+    minLevelCell = ds_init.minLevelCell - 1
+    maxLevelCell = ds_init.maxLevelCell - 1
+    bottomDepth = ds_init.bottomDepth
 
     areaCellMatrix = np.tile(areaCell, (nVertLevels, 1)).transpose()
     bottomMax = np.max(bottomDepth.values)
@@ -54,13 +58,13 @@ def compute_rpe(initial_state_file_name, output_files):
                                vert_index <= maxLevelCell)
     cell_mask = np.swapaxes(cell_mask, 0, 1)
 
-    with xr.open_dataset(output_files[0]) as ds:
+    with xr.open_dataset(output_filenames[0]) as ds:
         nt = ds.sizes['Time']
         xtime = ds.xtime.values
 
     rpe = np.ones((num_files, nt))
 
-    for file_index, out_filename in enumerate(output_files):
+    for file_index, out_filename in enumerate(output_filenames):
 
         ds = xr.open_dataset(out_filename)
 

polaris/ocean/tasks/baroclinic_channel/baroclinic_channel.cfg

@@ -37,9 +37,6 @@ btr_dt_per_km = 1.5
 # 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
 
@@ -63,3 +60,17 @@ salinity = 35.0
 
 # Coriolis parameter for entire domain.
 coriolis_parameter = -1.2e-4
+
+
+# config options for the baroclinic channel RPE tasks
+[baroclinic_channel_rpe]
+
+# Viscosity values to test for rpe test case
+viscosities = 1.0, 5.0, 10.0, 20.0, 200.0
+
+# plot time (days)
+plot_time = 20.0
+
+# min and max temperature range
+min_temp = 11.8
+max_temp = 13.0

polaris/ocean/tasks/baroclinic_channel/forward.py

@@ -72,10 +72,6 @@ def __init__(self, component, resolution, name='forward', subdir=None,
                          indir=indir, ntasks=ntasks, min_tasks=min_tasks,
                          openmp_threads=openmp_threads)
 
-        if nu is not None:
-            # update the viscosity to the requested value
-            self.add_model_config_options(options=dict(config_mom_del2=nu))
-
         # make sure output is double precision
         self.add_yaml_file('polaris.ocean.config', 'output.yaml')
 
@@ -87,6 +83,11 @@ def __init__(self, component, resolution, name='forward', subdir=None,
         self.add_yaml_file('polaris.ocean.tasks.baroclinic_channel',
                            'forward.yaml')
 
+        if nu is not None:
+            # update the viscosity to the requested value *after* loading
+            # forward.yaml
+            self.add_model_config_options(options=dict(config_mom_del2=nu))
+
         self.add_output_file(
             filename='output.nc',
             validate_vars=['temperature', 'salinity', 'layerThickness',

polaris/ocean/tasks/baroclinic_channel/rpe/__init__.py

@@ -68,7 +68,8 @@ def _add_rpe_and_analysis_steps(self, config=None):
         component = self.component
         resolution = self.resolution
 
-        nus = config.getlist('baroclinic_channel', 'viscosities', dtype=float)
+        nus = config.getlist('baroclinic_channel_rpe', 'viscosities',
+                             dtype=float)
         for nu in nus:
             name = f'nu_{nu:g}'
             step = Forward(

polaris/ocean/tasks/baroclinic_channel/rpe/analysis.py

@@ -5,6 +5,7 @@
 
 from polaris import Step
 from polaris.ocean.rpe import compute_rpe
+from polaris.viz import plot_horiz_field
 
 
 class Analysis(Step):
@@ -38,7 +39,11 @@ def __init__(self, component, indir, resolution, nus):
         self.nus = nus
 
         self.add_input_file(
-            filename='initial_state.nc',
+            filename='mesh.nc',
+            target='../../init/culled_mesh.nc')
+
+        self.add_input_file(
+            filename='init.nc',
             target='../../init/initial_state.nc')
 
         for nu in nus:
@@ -55,100 +60,53 @@ def run(self):
         """
         Run this step of the test case
         """
-        section = self.config['baroclinic_channel']
-        lx = section.getfloat('lx')
-        ly = section.getfloat('ly')
-        init_filename = self.inputs[0]
-        rpe = compute_rpe(initial_state_file_name=init_filename,
-                          output_files=self.inputs[1:])
-        with xr.open_dataset(init_filename) as ds_init:
-            nx = ds_init.attrs['nx']
-            ny = ds_init.attrs['ny']
-        _plot(nx, ny, lx, ly, self.outputs[0], self.nus, rpe)
-
-
-def _plot(nx, ny, lx, ly, filename, nus, rpe):
-    """
-    Plot section of the baroclinic channel 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)
-
-    lx : float
-        The size of the domain in km in the x direction
-
-    ly : int
-        The size of the domain in km in the y direction
-
-    filename : str
-        The output file name
-
-    nus : list
-        The viscosity values
-
-    rpe : numpy.ndarray
-        The reference potential energy with size len(nu) x len(time)
-    """
-
-    plt.switch_backend('Agg')
-    num_files = len(nus)
-    time = 20
-
-    ds = xr.open_dataset(f'output_nu_{nus[0]:g}.nc', decode_times=False)
-    times = ds.daysSinceStartOfSim.values
-
-    fig = plt.figure()
-    for i in range(num_files):
-        rpe_norm = np.divide((rpe[i, :] - rpe[i, 0]), rpe[i, 0])
-        plt.plot(times, rpe_norm,
-                 label=f"$\\nu_h=${nus[i]}")
-    plt.xlabel('Time, days')
-    plt.ylabel('RPE-RPE(0)/RPE(0)')
-    plt.legend()
-    plt.savefig('rpe_t.png')
-    plt.close(fig)
-
-    fig, axs = plt.subplots(1, num_files, figsize=(
-        2.1 * num_files, 5.0), constrained_layout=True)
-
-    # ***NOTE***: This is a quick-and-dirty plotting technique for regular
-    # planar hex meshes that we do not recommend adopting in other tasks
-    for iCol, nu in enumerate(nus):
-        ds = xr.open_dataset(f'output_nu_{nu:g}.nc', decode_times=False)
+        mesh_filename = 'mesh.nc'
+        init_filename = 'init.nc'
+        output_filename = self.outputs[0]
+        nus = self.nus
+        section = self.config['baroclinic_channel_rpe']
+
+        rpe = compute_rpe(mesh_filename=mesh_filename,
+                          initial_state_filename=init_filename,
+                          output_filenames=self.inputs[2:])
+
+        plt.switch_backend('Agg')
+        sim_count = len(nus)
+        time = section.getfloat('plot_time')
+        min_temp = section.getfloat('min_temp')
+        max_temp = section.getfloat('max_temp')
+
+        ds = xr.open_dataset(f'output_nu_{nus[0]:g}.nc', decode_times=False)
         times = ds.daysSinceStartOfSim.values
-        time_index = np.argmin(np.abs(times - time))
-        var = ds.temperature.values
-        var1 = np.reshape(var[time_index, :, 0], [ny, nx])
-        # flip in y-dir
-        var = np.flipud(var1)
-
-        # Every other row in y needs to average two neighbors in x on
-        # planar hex mesh
-        var_avg = var
-        for j in range(0, ny, 2):
-            for i in range(1, nx - 2):
-                var_avg[j, i] = (var[j, i + 1] + var[j, i]) / 2.0
-
-        ax = axs[iCol]
-        dis = ax.imshow(
-            var_avg,
-            extent=[0, lx, 0, ly],
-            cmap='cmo.thermal',
-            vmin=11.8,
-            vmax=13.0)
-        ax.set_title(f'day {times[time_index]}, '
-                     f'$\\nu_h=${nus[iCol]}')
-        ax.set_xticks(np.linspace(0, lx, 5))
-        ax.set_yticks(np.arange(0, ly, 11))
-
-        ax.set_xlabel('x, km')
-        if iCol == 0:
-            ax.set_ylabel('y, km')
-        if iCol == num_files - 1:
-            fig.colorbar(dis, ax=axs[num_files - 1], aspect=40)
-    plt.savefig(filename)
+
+        fig = plt.figure()
+        for i in range(sim_count):
+            rpe_norm = np.divide((rpe[i, :] - rpe[i, 0]), rpe[i, 0])
+            plt.plot(times, rpe_norm, label=f'$\\nu_h=${nus[i]}')
+        plt.xlabel('Time, days')
+        plt.ylabel('RPE-RPE(0)/RPE(0)')
+        plt.legend()
+        plt.savefig('rpe_t.png')
+        plt.close(fig)
+
+        ds_mesh = xr.open_dataset(mesh_filename)
+        ds_init = xr.open_dataset(init_filename)
+
+        fig, axes = plt.subplots(1, sim_count, figsize=(
+            3 * sim_count, 5.0), constrained_layout=True)
+
+        for row_index, nu in enumerate(nus):
+            ax = axes[row_index]
+            ds = xr.open_dataset(f'output_nu_{nu:g}.nc', decode_times=False)
+            ds = ds.isel(nVertLevels=0)
+            ds['maxLevelCell'] = ds_init.maxLevelCell
+            times = ds.daysSinceStartOfSim.values
+            time_index = np.argmin(np.abs(times - time))
+
+            plot_horiz_field(ds, ds_mesh, 'temperature', ax=ax,
+                             cmap='cmo.thermal', t_index=time_index,
+                             vmin=min_temp, vmax=max_temp,
+                             cmap_title='SST (C)')
+            ax.set_title(f'day {times[time_index]:g}, $\\nu_h=${nu:g}')
+
+        plt.savefig(output_filename)

polaris/ocean/tasks/cosine_bell/__init__.py

@@ -78,6 +78,8 @@ def configure(self):
         super().configure()
         config = self.config
         config.add_from_package('polaris.mesh', 'mesh.cfg')
+        config.add_from_package('polaris.ocean.tasks.cosine_bell',
+                                'cosine_bell.cfg')
 
         # set up the steps again in case a user has provided new resolutions
         self._setup_steps(config)

polaris/ocean/tasks/cosine_bell/analysis.py

@@ -9,7 +9,7 @@
 
 class Analysis(Step):
     """
-    A step for visualizing the output from the cosine bell test case
+    A step for analyzing the output from the cosine bell test case
 
     Attributes
     ----------