add get_extent to PastaStore

pastas · Oct 24, 2023 · 6767be3 · 6767be3
1 parent 6cf415f
commit 6767be3
Showing 1 changed file with 17 additions and 0 deletions.
diff --git a/pastastore/store.py b/pastastore/store.py
@@ -340,6 +340,23 @@ def get_tmin_tmax(self, libname, names=None, progressbar=False):
             tmintmax.loc[n, "tmax"] = s.last_valid_index()
         return tmintmax
 
+    def get_extent(self, libname, names=None):
+        names = self.conn._parse_names(names, libname=libname)
+        if libname in ["oseries", "stresses"]:
+            df = getattr(self, libname)
+        elif libname == "models":
+            df = self.oseries
+        else:
+            raise ValueError(f"Cannot get extent for library '{libname}'.")
+
+        extent = [
+            df.loc[names, "x"].min(),
+            df.loc[names, "x"].max(),
+            df.loc[names, "y"].min(),
+            df.loc[names, "y"].max(),
+        ]
+        return extent
+
     def get_parameters(
         self,
         parameters: Optional[List[str]] = None,