Get data refactor

.gitignore

@@ -127,3 +127,4 @@ dask-worker-space/
 
 #ruff linting
 .ruff_cache
+.envrc

docs/changelog.rst

@@ -17,6 +17,7 @@ Added
 - Add support for reading model configs in ``TOML`` format. (PR #444)
 - new ``force-overwrite`` option in ``hydromt update`` CLI to force overwritting updated netcdf files. (PR #460)
 - add ``open_mfcsv`` function in ``io`` module for combining multiple CSV files into one dataset. (PR #486)
+- Adapters can now clip data that is passed through a python object the same way as through the data catalog. (PR #481)
 
 Changed
 -------
@@ -28,7 +29,6 @@ Changed
 Fixed
 -----
 - when a model component (eg maps, forcing, grid) is updated using the set_ methods, it will first be read to avoid loosing data. (PR #460)
--
 
 Deprecated
 ----------

hydromt/data_adapter/dataframe.py

@@ -196,6 +196,48 @@ def get_data(
         based on the properties of this DataFrameAdapter. For a detailed
         description see: :py:func:`~hydromt.data_catalog.DataCatalog.get_dataframe`
         """
+        kwargs = self._parse_args()
+        df = self._load_data(variables, **kwargs)
+        df = DataFrameAdapter.slice_temporal_dimension(df, time_tuple)
+        df = self._uniformize_data(df)
+        return df
+
+    def _load_data(self, variables, **kwargs):
+        df = self._read_data(**kwargs)
+        df = self._rename_vars(df, variables)
+        return df
+
+    def _uniformize_data(self, df):
+        df = self._apply_unit_conversion(df)
+        df = self._set_meta_data(df)
+        return df
+
+    @staticmethod
+    def slice_temporal_dimension(df, time_tuple):
+        """Return a sliced DataFrame.
+
+        Parameters
+        ----------
+        df : pd.DataFrame
+            the dataframe to be sliced.
+        time_tuple : tuple of str, datetime, optional
+            Start and end date of period of interest. By default the entire time period
+            of the dataset is returned.
+
+        Returns
+        -------
+        pd.DataFrame
+            Tabular data
+        """
+        if time_tuple is not None and np.dtype(df.index).type == np.datetime64:
+            logger.debug(f"DataFrame: Slicing time dime {time_tuple}")
+            df = df[df.index.slice_indexer(*time_tuple)]
+            if df.size == 0:
+                raise IndexError("DataFrame: Time slice out of range.")
+
+        return df
+
+    def _parse_args(self):
         # Extract storage_options from kwargs to instantiate fsspec object correctly
         so_kwargs = {}
         if "storage_options" in self.driver_kwargs:
@@ -209,13 +251,14 @@ def get_data(
         _ = self.resolve_paths(**so_kwargs)  # throw nice error if data not found
 
         kwargs = self.driver_kwargs.copy()
+        return kwargs
 
-        # read and clip
+    def _read_data(self, **kwargs):
         logger.info(f"DataFrame: Read {self.driver} data.")
-
         if self.driver in ["csv"]:
             df = pd.read_csv(self.path, **kwargs)
         elif self.driver == "parquet":
+            _ = kwargs.pop("index_col", None)
             df = pd.read_parquet(self.path, **kwargs)
         elif self.driver in ["xls", "xlsx", "excel"]:
             df = pd.read_excel(self.path, engine="openpyxl", **kwargs)
@@ -224,16 +267,21 @@ def get_data(
         else:
             raise IOError(f"DataFrame: driver {self.driver} unknown.")
 
-        # rename and select columns
+        return df
+
+    def _rename_vars(self, df, variables):
         if self.rename:
             rename = {k: v for k, v in self.rename.items() if k in df.columns}
             df = df.rename(columns=rename)
+
         if variables is not None:
             if np.any([var not in df.columns for var in variables]):
                 raise ValueError(f"DataFrame: Not all variables found: {variables}")
             df = df.loc[:, variables]
 
-        # nodata and unit conversion for numeric data
+        return df
+
+    def _apply_unit_conversion(self, df):
         if df.index.size == 0:
             logger.warning(f"DataFrame: No data within spatial domain {self.path}.")
         else:
@@ -250,24 +298,18 @@ def get_data(
                         is_nodata = np.isin(df[c], np.atleast_1d(mv))
                         df[c] = np.where(is_nodata, np.nan, df[c])
 
-            # unit conversion
-            unit_names = list(self.unit_mult.keys()) + list(self.unit_add.keys())
-            unit_names = [k for k in unit_names if k in df.columns]
-            if len(unit_names) > 0:
-                logger.debug(f"DataFrame: Convert units for {len(unit_names)} columns.")
-            for name in list(set(unit_names)):  # unique
-                m = self.unit_mult.get(name, 1)
-                a = self.unit_add.get(name, 0)
-                df[name] = df[name] * m + a
-
-        # clip time slice
-        if time_tuple is not None and np.dtype(df.index).type == np.datetime64:
-            logger.debug(f"DataFrame: Slicing time dime {time_tuple}")
-            df = df[df.index.slice_indexer(*time_tuple)]
-            if df.size == 0:
-                raise IndexError("DataFrame: Time slice out of range.")
+        unit_names = list(self.unit_mult.keys()) + list(self.unit_add.keys())
+        unit_names = [k for k in unit_names if k in df.columns]
+        if len(unit_names) > 0:
+            logger.debug(f"DataFrame: Convert units for {len(unit_names)} columns.")
+        for name in list(set(unit_names)):  # unique
+            m = self.unit_mult.get(name, 1)
+            a = self.unit_add.get(name, 0)
+            df[name] = df[name] * m + a
+
+        return df
 
-        # set meta data
+    def _set_meta_data(self, df):
         df.attrs.update(self.meta)
 
         # set column attributes

hydromt/data_adapter/geodataframe.py

@@ -212,83 +212,64 @@ def get_data(
         buffer=0,
         logger=logger,
         variables=None,
-        # **kwargs,  # this is not used, for testing only
     ):
         """Return a clipped and unified GeoDataFrame (vector).
 
         For a detailed description see:
         :py:func:`~hydromt.data_catalog.DataCatalog.get_geodataframe`
         """
-        # If variable is string, convert to list
-        if variables:
-            variables = np.atleast_1d(variables).tolist()
-
-        if "storage_options" in self.driver_kwargs:
-            # not sure if storage options can be passed to fiona.open()
-            # for now throw NotImplemented Error
-            raise NotImplementedError(
-                "Remote file storage_options not implemented for GeoDataFrame"
-            )
-        _ = self.resolve_paths()  # throw nice error if data not found
+        varialbes, clip_str, geom, predicate, kwargs = self._parse_args(
+            variables, geom, bbox, buffer, predicate
+        )
+        gdf = self._load_data(clip_str, geom, predicate, **kwargs)
+        gdf = self.slice_data(gdf, variables, geom, predicate)
+        gdf = self._uniformize_data(gdf)
 
-        kwargs = self.driver_kwargs.copy()
-        # parse geom, bbox and buffer arguments
-        clip_str = ""
-        if geom is None and bbox is not None:
-            # convert bbox to geom with crs EPGS:4326 to apply buffer later
-            geom = gpd.GeoDataFrame(geometry=[box(*bbox)], crs=4326)
-            clip_str = " and clip to bbox (epsg:4326)"
-        elif geom is not None:
-            clip_str = f" and clip to geom (epsg:{geom.crs.to_epsg():d})"
-        if geom is not None:
-            # make sure geom is projected > buffer in meters!
-            if geom.crs.is_geographic and buffer > 0:
-                geom = geom.to_crs(3857)
-            geom = geom.buffer(buffer)  # a buffer with zero fixes some topology errors
-            bbox_str = ", ".join([f"{c:.3f}" for c in geom.total_bounds])
-            clip_str = f"{clip_str} [{bbox_str}]"
-        if kwargs.pop("within", False):  # for backward compatibility
-            predicate = "contains"
+        return gdf
 
-        # read and clip
-        logger.info(f"GeoDataFrame: Read {self.driver} data{clip_str}.")
-        if self.driver in [
-            "csv",
-            "parquet",
-            "xls",
-            "xlsx",
-            "xy",
-            "vector",
-            "vector_table",
-        ]:
-            # "csv", "xls", "xlsx", "xy" deprecated use vector_table instead.
-            # specific driver should be added to open_vector kwargs
-            if "driver" not in kwargs and self.driver in ["csv", "xls", "xlsx", "xy"]:
-                warnings.warn(
-                    "using the driver setting is deprecated. Please use"
-                    "vector_table instead."
-                )
+    @staticmethod
+    def slice_data(gdf, variables, geom, predicate):
+        """Return a clipped GeoDataFrame (vector).
 
-                kwargs.update(driver=self.driver)
-            # Check if file-object is required because of additional options
-            gdf = io.open_vector(
-                self.path, crs=self.crs, geom=geom, predicate=predicate, **kwargs
-            )
-        else:
-            raise ValueError(f"GeoDataFrame: driver {self.driver} unknown.")
+        Arguments
+        ---------
+        geom : geopandas.GeoDataFrame/Series,
+            A geometry defining the area of interest.
+        predicate : {'intersects', 'within', 'contains', 'overlaps',
+            'crosses', 'touches'}, optional If predicate is provided,
+            the GeoDataFrame is filtered by testing the predicate function
+            against each item. Requires bbox or mask. By default 'intersects'
+        variables : str or list of str, optional.
+            Names of GeoDataFrame columns to return.
 
-        # rename and select columns
-        if self.rename:
-            rename = {k: v for k, v in self.rename.items() if k in gdf.columns}
-            gdf = gdf.rename(columns=rename)
+        Returns
+        -------
+        gdf: geopandas.GeoDataFrame
+            GeoDataFrame
+        """
         if variables is not None:
             if np.any([var not in gdf.columns for var in variables]):
                 raise ValueError(f"GeoDataFrame: Not all variables found: {variables}")
             if "geometry" not in variables:  # always keep geometry column
                 variables = variables + ["geometry"]
             gdf = gdf.loc[:, variables]
 
-        # nodata and unit conversion for numeric data
+        if geom is not None:
+            gdf = gdf.sjoin(
+                gpd.GeoDataFrame(geometry=geom).to_crs(gdf.crs), predicate=predicate
+            )
+            if "index_right" in gdf.columns:
+                gdf = gdf.drop("index_right", axis=1)
+
+        return gdf
+
+    def _uniformize_data(self, gdf):
+        # rename and select columns
+        if self.rename:
+            rename = {k: v for k, v in self.rename.items() if k in gdf.columns}
+            gdf = gdf.renae(columns=rename)
+
+        # nodata and unit conversion
         if gdf.index.size == 0:
             logger.warning(f"GeoDataFrame: No data within spatial domain {self.path}.")
         else:
@@ -324,4 +305,96 @@ def get_data(
         for col in self.attrs:
             if col in gdf.columns:
                 gdf[col].attrs.update(**self.attrs[col])
+
         return gdf
+
+    def _load_data(self, clip_str, geom, predicate, **kwargs):
+        # read and clip
+        logger.info(f"GeoDataFrame: Read {self.driver} data{clip_str}.")
+        if self.driver in [
+            "csv",
+            "parquet",
+            "xls",
+            "xlsx",
+            "xy",
+            "vector",
+            "vector_table",
+        ]:
+            # "csv", "xls", "xlsx", "xy" deprecated use vector_table instead.
+            # specific driver should be added to open_vector kwargs
+            if "driver" not in kwargs and self.driver in ["csv", "xls", "xlsx", "xy"]:
+                warnings.warn(
+                    "using the driver setting is deprecated. Please use"
+                    "vector_table instead."
+                )
+
+                kwargs.update(driver=self.driver)
+            # Check if file-object is required because of additional options
+            gdf = io.open_vector(
+                self.path, crs=self.crs, geom=geom, predicate=predicate, **kwargs
+            )
+        else:
+            raise ValueError(f"GeoDataFrame: driver {self.driver} unknown.")
+
+        return gdf
+
+    def _parse_args(self, variables, geom, bbox, buffer, predicate):
+        # If variable is string, convert to list
+        if variables:
+            variables = np.atleast_1d(variables).tolist()
+
+        if "storage_options" in self.driver_kwargs:
+            # not sure if storage options can be passed to fiona.open()
+            # for now throw NotImplemented Error
+            raise NotImplementedError(
+                "Remote file storage_options not implemented for GeoDataFrame"
+            )
+        _ = self.resolve_paths()  # throw nice error if data not found
+
+        kwargs = self.driver_kwargs.copy()
+        geom, clip_str = GeoDataFrameAdapter.parse_geom(geom, bbox, buffer)
+        if kwargs.pop("within", False):  # for backward compatibility
+            predicate = "contains"
+
+        return variables, clip_str, geom, predicate, kwargs
+
+    @staticmethod
+    def parse_geom(geom, bbox, buffer):
+        """Parse geometries.
+
+        The geometry returned by this function can be
+        used by `GeoDataFrameAdapter.slice_data`.
+
+        Arguments
+        ---------
+        geom : geopandas.GeoDataFrame/Series,
+            A geometry defining the area of interest.
+        bbox : array-like of floats
+            (xmin, ymin, xmax, ymax) bounding box of area of interest
+            (in WGS84 coordinates).
+        buffer : float, optional
+            Buffer around the `bbox` or `geom` area of interest in meters. By default 0.
+
+        Returns
+        -------
+        geom: geometry
+            the actual geometry
+        clip_str: str
+            the string representation of the geom to be used in logging.
+        """
+        clip_str = ""
+        if geom is None and bbox is not None:
+            # convert bbox to geom with crs EPGS:4326 to apply buffer later
+            geom = gpd.GeoDataFrame(geometry=[box(*bbox)], crs=4326)
+            clip_str = " and clip to bbox (epsg:4326)"
+        elif geom is not None:
+            clip_str = f" and clip to geom (epsg:{geom.crs.to_epsg():d})"
+        if geom is not None:
+            # make sure geom is projected > buffer in meters!
+            if geom.crs.is_geographic and buffer > 0:
+                geom = geom.to_crs(3857)
+            geom = geom.buffer(buffer)  # a buffer with zero fixes some topology errors
+            bbox_str = ", ".join([f"{c:.3f}" for c in geom.total_bounds])
+            clip_str = f"{clip_str} [{bbox_str}]"
+
+        return geom, clip_str