pplot: make contourf more robust

ec-jrc · Oct 1, 2023 · c67b866 · c67b866
1 parent e65b211
commit c67b866
Showing 1 changed file with 12 additions and 1 deletion.
diff --git a/pyposeidon/utils/pplot.py b/pyposeidon/utils/pplot.py
@@ -372,7 +372,18 @@ def contourf(self, ax=None, it=None, **kwargs):
                 pass
 
         var = kwargs.get("var", "depth")
-        z = kwargs.get("z", self._obj[var].data[it, :].flatten())
+
+        if len(self._obj[var].shape) == 1:
+            zv = self._obj[var].data.flatten()
+        elif len(self._obj[var].shape) == 2:
+            if t_var in self._obj[var].coords:
+                zv = self._obj[var][it, :].data.flatten()
+            else:
+                raise Exception(f"{t_var} not in {var} dims")
+        else:
+            raise Exception(f"{var} dimension is larger than 2, please subset")
+
+        z = kwargs.get("z", zv)
 
         nv = kwargs.get("nv", 10)