Make isomip plotting more flexible

MPAS-Dev · Oct 24, 2024 · 824648b · 824648b
1 parent a28a73a
commit 824648b
Showing 1 changed file with 5 additions and 1 deletion.
diff --git a/compass/ocean/tests/isomip_plus/viz/plot.py b/compass/ocean/tests/isomip_plus/viz/plot.py
@@ -239,6 +239,8 @@ def __init__(self, inFolder, streamfunctionFolder, outFolder, expt,
         self.sectionY = sectionY
         self.showProgress = showProgress
 
+        if 'Time' in dsMesh.dims:
+            dsMesh = dsMesh.isel(Time=0)
         self.dsMesh = dsMesh
         self.ds = ds
 
@@ -508,7 +510,9 @@ def plot_horiz_series(self, da, nameInTitle, prefix, oceanDomain,
             The time indices at which to plot. If not provided, set to all.
         """
 
-        nTime = self.ds.sizes['Time']
+        if 'Time' not in da.dims:
+            da = da.expand_dims(dim='Time', axis=0)
+        nTime = da.sizes['Time']
         if self.showProgress:
             widgets = ['plotting {}: '.format(nameInTitle),
                        progressbar.Percentage(), ' ',