Support Omega time in mfd_sln viz step

E3SM-Project · Dec 4, 2024 · a6e420e · a6e420e
1 parent 5444602
commit a6e420e
Showing 1 changed file with 20 additions and 8 deletions.
diff --git a/polaris/ocean/tasks/manufactured_solution/viz.py b/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 as 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)