Take advantage of xarray's decode_times for Omega's time variable conversion

polaris/mpas/time.py

@@ -28,7 +28,7 @@ def time_index_from_xtime(xtime, dt_target, start_xtime=None):
     return time_index
 
 
-def time_since_start(xtime, start_xtime=None):
+def time_since_start(xtime, start_xtime='0001-01-01_01:00:00'):
     """
     Determine the time elapsed since the start of the simulation
 
@@ -48,9 +48,14 @@ def time_since_start(xtime, start_xtime=None):
     if start_xtime is None:
         start_xtime = xtime[0].decode()
 
-    t0 = datetime.datetime.strptime(start_xtime, '%Y-%m-%d_%H:%M:%S')
+    try:
+        time_format = '%Y-%m-%d_%H:%M:%S.%f'
+        t0 = datetime.datetime.strptime(start_xtime, time_format)
+    except ValueError:
+        time_format = '%Y-%m-%d_%H:%M:%S'
+        t0 = datetime.datetime.strptime(start_xtime, time_format)
     dt = np.zeros((len(xtime),))
     for idx, xt in enumerate(xtime):
-        t = datetime.datetime.strptime(xt.decode(), '%Y-%m-%d_%H:%M:%S')
+        t = datetime.datetime.strptime(xt.decode(), time_format)
         dt[idx] = (t - t0).total_seconds()
     return dt

polaris/ocean/model/time.py

@@ -1,9 +1,11 @@
 import time
-from datetime import datetime
 
+import cftime
 import numpy as np
 import pandas as pd
 
+from polaris.mpas.time import time_since_start
+
 
 def get_days_since_start(ds):
     """
@@ -13,23 +15,23 @@ def get_days_since_start(ds):
     if 'daysSinceStartOfSim' in ds.keys():
         t_arr = ds.daysSinceStartOfSim.values.astype(float)
     elif 'xtime' in ds.keys():
-        timestamps = []
-        for time_str in ds.xtime.values.astype(str):
-            try:
-                timestamp = datetime.strptime(time_str, '%Y-%m-%d_%H:%M:%S.%f')
-            except ValueError:
-                timestamp = datetime.strptime(time_str, '%Y-%m-%d_%H:%M:%S')
-            timestamps.append(timestamp)
         # Calculate seconds since the first timestamp
-        seconds_since_start = [
-            (ts - timestamps[0]).total_seconds() for ts in timestamps
-        ]
+        seconds_since_start = time_since_start(ds.xtime.values)
+        # Convert to days
         t_arr = np.array(seconds_since_start, dtype=float) / 86400.0
     elif 'Time' in ds.keys():
-        t_vals = ds['Time'].values
-        t_pd = pd.to_datetime(t_vals)
-        t_arr = 1.0e9 * (t_pd - t_pd[0]) / np.timedelta64(1, 's')
-        t_arr = t_arr.astype(float) / 86400.0
+        # This option works if decode_times=True when loading xr.Dataset
+        if 'Time' in ds['Time'].coords:
+            t_arr = cftime.date2num(
+                ds['Time'].values,
+                units=ds['Time'].Units.replace('seconds', 'days'),
+                calendar=ds['Time'].dt.calendar,
+                has_year_zero=True,
+            )
+        else:
+            t_pd = pd.to_datetime(ds['Time'].values)
+            t_arr = 1.0e9 * (t_pd - t_pd[0]) / np.timedelta64(1, 's')
+            t_arr = t_arr.astype(float) / 86400.0
     else:
         raise ValueError('Could not find a time variable in dataset')
     return t_arr

polaris/tasks/ocean/seamount/__init__.py

@@ -13,6 +13,7 @@ def add_seamount_tasks(component):
     indir = 'planar/seamount'
     config_filename = 'seamount.cfg'
     config = PolarisConfigParser(filepath=f'{indir}/{config_filename}')
+    config.add_from_package('polaris.ocean.eos', 'linear.cfg')
     config.add_from_package('polaris.tasks.ocean.seamount', config_filename)
 
     init_step = Init(component=component, name='init', indir=indir)

polaris/tasks/ocean/seamount/viz.py

@@ -1,16 +1,15 @@
 import cmocean  # noqa: F401
 import matplotlib.pyplot as plt
 import numpy as np
-import pandas as pd
 import xarray as xr
 from mpas_tools.ocean.viz.transect import compute_transect, plot_transect
 
-from polaris import Step
 from polaris.mpas import cell_mask_to_edge_mask
+from polaris.ocean.model import OceanIOStep, get_days_since_start
 from polaris.viz import plot_horiz_field
 
 
-class Viz(Step):
+class Viz(OceanIOStep):
     """
     A step for plotting the results of the default seamount forward step
     """
@@ -40,9 +39,11 @@ def run(self):
         """
         Run this step of the task
         """
-        ds_mesh = xr.load_dataset('mesh.nc')
-        ds_init = xr.load_dataset('init.nc')
-        ds = xr.load_dataset('output.nc')
+        ds_mesh = self.open_model_dataset('mesh.nc')
+        ds_init = self.open_model_dataset('init.nc', self.config)
+        ds = self.open_model_dataset(
+            'output.nc', self.config, decode_times=True
+        )
 
         x_min = ds_mesh.xVertex.min().values
         x_max = ds_mesh.xVertex.max().values
@@ -115,10 +116,8 @@ def run(self):
         # Plot the time series of max velocity
         plt.figure(figsize=[12, 6], dpi=100)
         umax = np.amax(ds.normalVelocity[:, :, 0].values, axis=1)
-        t = ds.daysSinceStartOfSim.values
-        time = pd.to_timedelta(t)
-        days_float = time / pd.Timedelta(days=1)
-        plt.plot(days_float, umax, 'k-o', label='max(normalVelocity)')
+        t_days = get_days_since_start(ds)
+        plt.plot(t_days, umax, 'k-o', label='max(normalVelocity)')
         plt.xlabel('Time (days)')
         plt.ylabel('Maximum Velocity (m/s)')
         plt.legend()