Merge pull request #249 from xylar/support-omega-time-coord

Support the Omega `time` coordinate

polaris/mpas/__init__.py

@@ -1,2 +1,2 @@
 from polaris.mpas.area import area_for_field
-from polaris.mpas.time import time_index_from_xtime
+from polaris.mpas.time import time_index_from_xtime, time_since_start

polaris/mpas/time.py

@@ -23,6 +23,28 @@ def time_index_from_xtime(xtime, dt_target, start_xtime=None):
     time_index : int
         Index in xtime that is closest to dt_target
     """
+    dt = time_since_start(xtime, start_xtime)
+    time_index = np.argmin(np.abs(np.subtract(dt, dt_target)))
+    return time_index
+
+
+def time_since_start(xtime, start_xtime=None):
+    """
+    Determine the time elapsed since the start of the simulation
+
+    Parameters
+    ----------
+    xtime : numpy.ndarray of numpy.char
+        Times in the dataset
+
+    start_xtime : str, optional
+        The start time, the first entry in ``xtime`` by default
+
+    Returns
+    -------
+    dt : numpy.ndarray
+        The elapsed time in seconds corresponding to each entry in xtime
+    """
     if start_xtime is None:
         start_xtime = xtime[0].decode()
 
@@ -32,5 +54,4 @@ def time_index_from_xtime(xtime, dt_target, start_xtime=None):
         t = datetime.datetime.strptime(xt.decode(),
                                        '%Y-%m-%d_%H:%M:%S')
         dt[idx] = (t - t0).total_seconds()
-    time_index = np.argmin(np.abs(np.subtract(dt, dt_target)))
-    return time_index
+    return dt

polaris/ocean/convergence/analysis.py

@@ -4,7 +4,7 @@
 import numpy as np
 import pandas as pd
 
-from polaris.mpas import area_for_field, time_index_from_xtime
+from polaris.mpas import area_for_field, time_since_start
 from polaris.ocean.convergence import (
     get_resolution_for_task,
     get_timestep_for_task,
@@ -395,11 +395,20 @@ def get_output_field(self, refinement_factor, field_name, time, zidx=None):
         field_mpas : xarray.DataArray
             model output field
         """
-        ds_out = self.open_model_dataset(f'output_r{refinement_factor:02g}.nc')
+        config = self.config
+        ds_out = self.open_model_dataset(f'output_r{refinement_factor:02g}.nc',
+                                         decode_times=False)
 
-        tidx = time_index_from_xtime(ds_out.xtime.values, time)
-        ds_out = ds_out.isel(Time=tidx)
+        model = config.get('ocean', 'model')
+        if model == 'mpas-o':
+            dt = time_since_start(ds_out.xtime.values)
+        else:
+            # time is seconds since the start of the simulation in Omega
+            dt = ds_out.Time.values
 
+        tidx = np.argmin(np.abs(dt - time))
+
+        ds_out = ds_out.isel(Time=tidx)
         field_mpas = ds_out[field_name]
         if zidx is not None:
             field_mpas = field_mpas.isel(nVertLevels=zidx)

polaris/ocean/tasks/manufactured_solution/viz.py

@@ -135,6 +135,8 @@ def run(self):
         section = config['manufactured_solution']
         eta0 = section.getfloat('ssh_amplitude')
 
+        model = config.get('ocean', 'model')
+
         use_mplstyle()
         fig, axes = plt.subplots(nrows=nres, ncols=3, figsize=(12, 2 * nres))
         rmse = []
@@ -146,16 +148,26 @@ def run(self):
             ds_init = self.open_model_dataset(
                 f'init_r{refinement_factor:02g}.nc')
             ds = self.open_model_dataset(
-                f'output_r{refinement_factor:02g}.nc')
+                f'output_r{refinement_factor:02g}.nc', decode_times=False)
             exact = ExactSolution(config, ds_init)
 
-            t0 = datetime.datetime.strptime(ds.xtime.values[0].decode(),
-                                            '%Y-%m-%d_%H:%M:%S')
-            tf = datetime.datetime.strptime(ds.xtime.values[-1].decode(),
-                                            '%Y-%m-%d_%H:%M:%S')
-            t = (tf - t0).total_seconds()
-            ssh_model = ds.ssh.values[-1, :]
-            rmse.append(np.sqrt(np.mean((ssh_model - exact.ssh(t).values)**2)))
+            if model == 'mpas-o':
+                t0 = datetime.datetime.strptime(ds.xtime.values[0].decode(),
+                                                '%Y-%m-%d_%H:%M:%S')
+                tf = datetime.datetime.strptime(ds.xtime.values[-1].decode(),
+                                                '%Y-%m-%d_%H:%M:%S')
+                t = (tf - t0).total_seconds()
+
+            else:
+                # time is seconds since the start of the simulation in Omega
+                t = ds.Time[-1].values
+
+            ssh_model = ds.ssh.isel(Time=-1)
+            if 'nVertLevels' in ssh_model.dims:
+                # Omega v0 uses stacked shallow water where ssh has nVertLevels
+                ssh_model = ssh_model.isel(nVertLevels=0)
+            rmse.append(np.sqrt(np.mean(
+                (ssh_model.values - exact.ssh(t).values)**2)))
 
             # Comparison plots
             ds['ssh_exact'] = exact.ssh(t)