Add functions for computing and plotting transects

deploy/default.cfg

@@ -24,7 +24,7 @@ geometric_features = 1.2.0
 jigsaw = 0.9.14
 jigsawpy = 0.3.3
 mache = 1.16.0
-mpas_tools = 0.25.0
+mpas_tools = 0.27.0
 otps = 2021.10
 parallelio = 2.6.0
 

docs/developers_guide/ocean/api.md

@@ -314,3 +314,22 @@
    vertical.zlevel.compute_z_level_resting_thickness
    vertical.zstar.init_z_star_vertical_coord
 ```
+
+### Visualization
+
+```{eval-rst}
+.. currentmodule:: polaris.ocean.viz
+
+.. autosummary::
+   :toctree: generated/
+
+   compute_transect
+   plot_transect
+   transect.horiz.find_spherical_transect_cells_and_weights
+   transect.horiz.find_planar_transect_cells_and_weights
+   transect.horiz.make_triangle_tree
+   transect.horiz.mesh_to_triangles
+   transect.vert.find_transect_levels_and_weights
+   transect.vert.interp_mpas_to_transect_cells
+   transect.vert.interp_mpas_to_transect_nodes
+

docs/developers_guide/ocean/framework.md

@@ -422,3 +422,13 @@ density, which is horizontally constant and increases with depth.
 The {py:func}`polaris.ocean.rpe.compute_rpe()` is used to compute the RPE as
 a function of time in a series of one or more output files.  The RPE is stored
 in `rpe.csv` and also returned as a numpy array for plotting and analysis.
+
+## Visualization
+
+The `polaris.ocean.viz` module provides functions for making plots that are
+specific to the ocean component.
+
+The `polaris.ocean.viz.transect` modules includes functions for computing
+({py:func}`polaris.ocean.viz.compute_transect()`) and plotting
+({py:func}`polaris.ocean.viz.plot_transect()`) transects through the ocean
+from a sequence of x-y or latitude-longitude coordinates.

polaris/ocean/tasks/baroclinic_channel/init.py

@@ -8,6 +8,7 @@
 from polaris import Step
 from polaris.mesh.planar import compute_planar_hex_nx_ny
 from polaris.ocean.vertical import init_vertical_coord
+from polaris.ocean.viz import compute_transect, plot_transect
 from polaris.viz import plot_horiz_field
 
 
@@ -162,8 +163,30 @@ def run(self):
         write_netcdf(ds, 'initial_state.nc')
 
         cell_mask = ds.maxLevelCell >= 1
-        plot_horiz_field(ds, ds_mesh, 'temperature',
-                         'initial_temperature.png', cell_mask=cell_mask)
+
         plot_horiz_field(ds, ds_mesh, 'normalVelocity',
                          'initial_normal_velocity.png', cmap='cmo.balance',
                          show_patch_edges=True, cell_mask=cell_mask)
+
+        x_mid = 0.5 * (x_min + x_max)
+
+        y = xr.DataArray(data=np.linspace(y_min, y_max, 2), dims=('nPoints',))
+        x = x_mid * xr.ones_like(y)
+
+        ds_transect = compute_transect(
+            x=x, y=y, ds_horiz_mesh=ds_mesh,
+            layer_thickness=ds.layerThickness.isel(Time=0),
+            bottom_depth=ds.bottomDepth, min_level_cell=ds.minLevelCell - 1,
+            max_level_cell=ds.maxLevelCell - 1, spherical=False)
+
+        field_name = 'temperature'
+        plot_transect(ds_transect=ds_transect,
+                      mpas_field=ds[field_name].isel(Time=0),
+                      title=f'{field_name} at x={1e-3 * x_mid:.1f} km',
+                      out_filename=f'initial_{field_name}_section.png',
+                      vmin=9., vmax=13., cmap='cmo.thermal',
+                      colorbar_label=r'$^\circ$C')
+
+        plot_horiz_field(ds, ds_mesh, 'temperature', 'initial_temperature.png',
+                         vmin=9., vmax=13., cmap='cmo.thermal',
+                         cell_mask=cell_mask, transect_x=x, transect_y=y)

polaris/ocean/tasks/baroclinic_channel/viz.py

@@ -3,6 +3,7 @@
 import xarray as xr
 
 from polaris import Step
+from polaris.ocean.viz import compute_transect, plot_transect
 from polaris.viz import plot_horiz_field
 
 
@@ -42,13 +43,44 @@ def run(self):
         ds = xr.load_dataset('output.nc')
         t_index = ds.sizes['Time'] - 1
         cell_mask = ds_init.maxLevelCell >= 1
-        plot_horiz_field(ds, ds_mesh, 'temperature',
-                         'final_temperature.png', t_index=t_index,
-                         cell_mask=cell_mask)
         max_velocity = np.max(np.abs(ds.normalVelocity.values))
         plot_horiz_field(ds, ds_mesh, 'normalVelocity',
                          'final_normalVelocity.png',
                          t_index=t_index,
                          vmin=-max_velocity, vmax=max_velocity,
                          cmap='cmo.balance', show_patch_edges=True,
                          cell_mask=cell_mask)
+
+        x_cell = ds_mesh.xCell
+        y_cell = ds_mesh.yCell
+
+        x_min = x_cell.min().values
+        x_max = x_cell.max().values
+        y_min = y_cell.min().values
+        y_max = y_cell.max().values
+
+        x_mid = 0.5 * (x_min + x_max)
+
+        y = xr.DataArray(data=np.linspace(y_min, y_max, 2), dims=('nPoints',))
+        x = x_mid * xr.ones_like(y)
+
+        ds_transect = compute_transect(
+            x=x, y=y, ds_horiz_mesh=ds_mesh,
+            layer_thickness=ds.layerThickness.isel(Time=t_index),
+            bottom_depth=ds_init.bottomDepth,
+            min_level_cell=ds_init.minLevelCell - 1,
+            max_level_cell=ds_init.maxLevelCell - 1,
+            spherical=False)
+
+        field_name = 'temperature'
+        mpas_field = ds[field_name].isel(Time=t_index)
+        plot_transect(ds_transect=ds_transect, mpas_field=mpas_field,
+                      title=f'{field_name} at x={1e-3 * x_mid:.1f} km',
+                      out_filename=f'final_{field_name}_section.png',
+                      vmin=9., vmax=13., cmap='cmo.thermal',
+                      colorbar_label=r'$^\circ$C')
+
+        plot_horiz_field(ds, ds_mesh, 'temperature', 'final_temperature.png',
+                         t_index=t_index, vmin=9., vmax=13.,
+                         cmap='cmo.thermal', cell_mask=cell_mask, transect_x=x,
+                         transect_y=y)

polaris/ocean/viz/__init__.py

@@ -0,0 +1,2 @@
+from polaris.ocean.viz.transect.plot import plot_transect
+from polaris.ocean.viz.transect.vert import compute_transect