Include Omega fields needed for p* coordinate in initial conditions

polaris/ocean/model/mpaso_to_omega.yaml

@@ -27,6 +27,7 @@ variables:
   maxLevelCell: MaxLayerCell
   # currently hard-coded in horizontal mesh
   # bottomDepth: BottomDepth
+  vertCoordMovementWeights: VertCoordMovementWeights
 
   # tracers
   temperature: Temperature

polaris/ocean/ocean.cfg

@@ -51,3 +51,6 @@ min_vert_levels = 1
 
 # Minimum thickness of each layer for z-star coordinate
 min_layer_thickness = 0
+
+# The number of iterations to use when computing pseudothickness for TEOS-10
+pseudothickness_iter_count = 6

polaris/ocean/vertical/diagnostics.py

@@ -46,20 +46,31 @@ def geom_thickness_from_ds(ds, config):
         )
 
 
-def pseudothickness_from_ds(ds, config):
+def pseudothickness_from_ds(
+    ds, config, src_var_name='layerThickness', iter_count=None
+):
     """
     Compute pseudothickness from temperature and salinity in dataset.
 
     Parameters
     ----------
     ds : xarray.Dataset
         An ocean dataset containing 'temperature', 'salinity',
-        'layerThickness', and 'ssh'
+        src_var_name, and 'ssh'
 
     config : polaris.config.PolarisConfigParser
         Configuration options for the test case, including
         'vertical_grid:surface_pressure'
 
+    src_var_name : str, optional
+        The name of the variable in ds to use as the source for the
+        geometric thickness, by default 'layerThickness'.
+
+    iter_count : int, optional
+        The number of iterations to use when computing pressure and
+        specific volume, by default ``pseudothickness_iter_count`` for teos-10
+        and 1 for linear or constant EOS.
+
     Returns
     -------
     pseudothickness : xarray.DataArray or None
@@ -78,14 +89,25 @@ def pseudothickness_from_ds(ds, config):
         )
         return None, None
 
+    if iter_count is None:
+        eos_type = config.get('ocean', 'eos_type')
+        if eos_type in ['constant', 'linear']:
+            iter_count = 1
+        elif eos_type == 'teos-10':
+            iter_count = config.getint(
+                'vertical_grid', 'pseudothickness_iter_count'
+            )
+        else:
+            raise ValueError(f'Unsupported equation of state type: {eos_type}')
+
     surface_pressure = config.getfloat('vertical_grid', 'surface_pressure')
     p_interface, _, spec_vol = pressure_and_spec_vol_from_state_at_geom_height(
         config,
-        ds.layerThickness,
+        ds[src_var_name],
         ds.temperature,
         ds.salinity,
         surface_pressure * xr.ones_like(ds.ssh),
-        iter_count=1,
+        iter_count=iter_count,
     )
 
     pseudothickness = pseudothickness_from_pressure(p_interface)

polaris/tasks/ocean/__init__.py

@@ -172,19 +172,22 @@ def write_model_dataset(self, ds, filename, config=None):
         filename : str
             The path for the NetCDF file to write
         """
-        if (
-            self.model == 'omega'
-            and 'layerThickness' in ds.keys()
-            and 'PseudoThickness' not in ds.keys()
-            and config is not None
-        ):
-            pseudothickness, spec_vol = pseudothickness_from_ds(
-                ds, config=config
-            )
-            if pseudothickness is not None and spec_vol is not None:
-                ds['PseudoThickness'] = pseudothickness
-                ds['SpecVol'] = spec_vol
-                ds['layerThickness'] = ds['PseudoThickness'].copy()
+        if self.model == 'omega' and config is not None:
+            # fields to be converted from geometric to pseudo thickness
+            mpas_to_omega_vars = {
+                'layerThickness': 'PseudoThickness',
+                'restingThickness': 'RefPseudoThickness',
+            }
+            for mpas_var, omega_var in mpas_to_omega_vars.items():
+                if mpas_var in ds.keys() and omega_var not in ds.keys():
+                    pseudothickness, spec_vol = pseudothickness_from_ds(
+                        ds, config=config, src_var_name=mpas_var
+                    )
+                    if pseudothickness is not None and spec_vol is not None:
+                        ds[omega_var] = pseudothickness
+                        if mpas_var == 'layerThickness':
+                            ds['SpecVol'] = spec_vol
+                            ds['layerThickness'] = ds['PseudoThickness'].copy()
 
         # After map_to_native_model_vars, the dataset contains
         # LayerThickness and PseudoThickness which are both