Improve export scripts and small patches

padeopsIO/budgetIO.py

@@ -218,7 +218,7 @@ def __init__(
 
         elif npz or src == "npz":  # .npz saved files
             self.associate_npz = True
-            self._init_npz()
+            self._init_npz(normalize_origin=normalize_origin)
             self.printv(f"Initialized BudgetIO at {dirname} from .npz files. ")
 
         elif npy or src == "npy":
@@ -2007,6 +2007,7 @@ def read_turb_vvel(self, tidx=None, **kwargs):
         """
         return self.read_turb_property(tidx, "vvel", **kwargs)
 
+
     def get_logfiles(self, path=None, search_str="*.o[0-9]*", id=-1):
         """
         Searches for all logfiles formatted "*.o[0-9]" (Stampede3 format)
@@ -2026,22 +2027,40 @@ def get_logfiles(self, path=None, search_str="*.o[0-9]*", id=-1):
         path = path or self.dir_name
         return tools.get_logfiles(path, search_str=search_str, id=id)
 
+
+    def query_logfile(self, search_terms, logfile=None, search_str=None, id=-1, **kwargs):
+        """
+        Queries the PadeOps output log file for text lines printed out by temporalhook.F90.
+
+        By default, the search looks for TERM followed by any arbitrary characters, then at least 1
+        character of white space followed by a number of format %e (exponential). Casts the resulting
+        string into a float.
+
+        see `io_utils.query_logfile()` for more information.
+        """
+        if logfile is None:
+            if search_str is None:
+                search_str = "*.o[0-9]*"
+            logfile = self.get_logfiles(search_str=search_str, id=id)
+        return io.query_logfile(logfile, search_terms=search_terms, **kwargs)
+
+
     def get_ustar(self, logfile=None, crop_budget=True, average=True):
         """
         Gleans ustar from the logfile.
 
         Parameters
         ----------
         logfile : path-like, optional
-            Path to logfile. If None, searches for all files ending in '.o[0-9]*'.
+            Path to logfile. If None, searches for all files ending in '*.o[0-9]*'.
             Default is None.
         crop_budget : bool, optional
             Crops time axis to budgets. Defaults to True.
         average : bool, optional
-            Time averages. Defaults to True.
+            Time averages over the budget_time_avg window. Defaults to True.
         """
         return tools.get_ustar(
-            self, search_str=logfile, crop_budget=crop_budget, average=average
+            self, logfile=logfile, crop_budget=crop_budget, average=average
         )
 
     def get_uhub(self, z_hub=0, use_fields=False, **slice_kwargs):

padeopsIO/gridslice.py

@@ -288,6 +288,49 @@ def __call__(self, ax=None, cbar=True, figsize=None, **kwargs):
             return im
 
 
+@xr.register_dataset_accessor("pcolormesh")
+@xr.register_dataarray_accessor("pcolormesh")
+class XRPcolormesh:
+    """
+    Add plotting function `pcolormesh` for 2d arrays
+
+    This is similar to `imshow` except useful for non-uniform grids. It 
+    is also similar to the native `plot` function, but plots the transpose
+    automatically and sets the apsect ratio to 1.0 by default.
+    """
+
+    def __init__(self, xarray_obj):
+        self._obj = xarray_obj
+
+    def __call__(self, ax=None, cbar=True, figsize=None, aspect=1, **kwargs):
+        if isinstance(self._obj, xr.Dataset):
+            if len(self._obj.keys()) > 1:
+                raise ValueError("Cannot plot type `Dataset` with more than 1 key")
+            else:
+                self._obj[next(iter(self._obj))].pcolormesh(ax=ax, cbar=cbar, **kwargs)
+
+        else:
+            if self._obj.ndim != 2:
+                raise AttributeError("pcolormesh() requires 2D data")
+            if ax is None:
+                _, ax = plt.subplots(figsize=figsize)
+
+            im = ax.pcolormesh(
+                *self._obj.grid.xi, self._obj.T, **kwargs
+            )
+            axes = self._obj.grid.keys()
+            ax.set_xlabel(labels[axes[0]])
+            ax.set_ylabel(labels[axes[1]])
+            ax.set_aspect(aspect)
+            if cbar:
+                if self._obj.name in labels.keys():
+                    label = labels[self._obj.name]
+                else:
+                    label = self._obj.name
+                plt.colorbar(im, ax=ax, label=label)
+            return im
+
+
 # ================= helper functions =================