GeospatialScenario instance kdtree

Author: tobin-ford

pvdeg/scenario.py

@@ -1110,6 +1110,7 @@ def __init__(
         self.hpc = hpc
         self.func = func
         self.template = template
+        self.kdtree = None # sklearn kdtree
 
     def __eq__(self, other):
         raise NotImplementedError("""
@@ -1285,14 +1286,22 @@ def location_bounding_box(
 
         self.meta_data = self.meta_data.loc[bbox_gids]
 
+    def set_kdtree(self, kdtree = None) -> None:
+        """Initialize a kidtree and save it to the GeospatialScenario"""
+        if kdtree is None:
+            self.kdtree = pvdeg.geospatial.meta_KDtree(meta_df=self.meta_data)
+        else:
+            self.kdtree = kdtree
+
+
     def classify_mountains_radii(
         self,
-        kdtree,
         rad_1: Union[float, int] = 12,
         rad_2: Union[float, int] = 1,
         threshold_factor: Union[float, int] = 1.25,
         elevation_floor: Union[float, int] = 0,
         bbox_kwarg: Optional[dict] = {},
+        kdtree = None,
     ):
         """
         Find mountains from elevation metadata using sklearn kdtree for fast lookup.
@@ -1306,9 +1315,6 @@ def classify_mountains_radii(
         -----------
         meta_df : pd.DataFrame
             Dataframe of metadata as generated by pvdeg.weather.get for geospatial
-        kdtree : sklearn.neighbors.KDTree
-            kdtree containing latitude-longitude pairs for quick lookups
-            Generate using ``pvdeg.geospatial.meta_KDTree``
         rad_1 : float
             radius of the larger search area whose elevations are compared against
             the smaller search area. controls the kdtree query region.
@@ -1325,15 +1331,21 @@ def classify_mountains_radii(
         elevation_floor : int
             minimum inclusive elevation in meters. If a point has smaller location
             it will be clipped from result.
+        kdtree : sklearn.neighbors.KDTree
+            Generated automatically but can be provided externally.
+            kdtree containing latitude-longitude pairs for quick lookups
+            Generate using ``pvdeg.geospatial.meta_KDTree``
 
         Returns:
         --------
         None, strictly updates meta_data attribute of geospatial scenari instance.
         """
 
+        self.set_kdtree(kdtree=kdtree)
+
         gids = pvdeg.geospatial.identify_mountains_radii(
             meta_df=self.meta_data,
-            kdtree=kdtree,
+            kdtree=self.kdtree,
             rad_1=rad_1,
             rad_2=rad_2,
             threshold_factor=threshold_factor,
@@ -1346,12 +1358,12 @@ def classify_mountains_radii(
 
     def classify_mountains_weights(
         self,
-        kdtree,
         threshold: int = 0,
         percentile: int = 75,
         k_neighbors: int = 3,
         method: str = "mean",
         normalization: str = "linear",
+        kdtree = None,
     ):
         """
         Add a column to the scenario meta_data dataframe containing a boolean
@@ -1360,10 +1372,6 @@ def classify_mountains_weights(
 
         Parameters:
         -----------
-        kdtree : sklearn.neighbors.KDTree or str
-            kdtree containing latitude-longitude pairs for quick lookups
-            Generate using ``pvdeg.geospatial.meta_KDTree``. Can take a pickled
-            kdtree as a path to the .pkl file.
         threshold : float
             minimum weight that a mountain can be identifed.
             value between `[0,1]` (inclusive)
@@ -1379,6 +1387,11 @@ def classify_mountains_weights(
         normalization : str, (default = 'linear')
             function to apply when normalizing weights. Logarithmic uses log_e/ln
             options : `'linear'`, `'logarithmic'`, '`exponential'`
+        kdtree : sklearn.neighbors.KDTree or str
+            Generated automatically but can be provided externally.
+            kdtree containing latitude-longitude pairs for quick lookups
+            Generate using ``pvdeg.geospatial.meta_KDTree``. Can take a pickled
+            kdtree as a path to the .pkl file.
 
         Returns:
         --------
@@ -1388,9 +1401,12 @@ def classify_mountains_weights(
         ---------
         `pvdeg.geospatial.identify_mountains_weights`
         """
+
+        self.set_kdtree(kdtree=kdtree)
+
         gids = pvdeg.geospatial.identify_mountains_weights(
             meta_df=self.meta_data,
-            kdtree=kdtree,
+            kdtree=self.kdtree,
             threshold=threshold,
             percentile=percentile,
             k_neighbors=k_neighbors,
@@ -1403,17 +1419,13 @@ def classify_mountains_weights(
 
     def classify_feature(
         self,
-        kdtree=None,
         feature_name=None,
         resolution="10m",
         radius=None,
+        kdtree=None,
         bbox_kwarg={},
     ):
         """
-        kdtree : sklearn.neighbors.KDTree or str
-            kdtree containing latitude-longitude pairs for quick lookups
-            Generate using ``pvdeg.geospatial.meta_KDTree``. Can take a pickled
-            kdtree as a path to the .pkl file.
         feature_name : str
             cartopy.feature.NaturalEarthFeature feature key.
             Options: ``'lakes'``, ``'rivers_lake_centerlines'``, ``'coastline'``
@@ -1424,6 +1436,11 @@ def classify_feature(
             Area around feature coordinates to include in the downsampled result.
             Bigger area means larger radius and more samples included.
             pass
+        kdtree : sklearn.neighbors.KDTree or str
+            Generated automatically but can be provided externally.
+            kdtree containing latitude-longitude pairs for quick lookups
+            Generate using ``pvdeg.geospatial.meta_KDTree``. Can take a pickled
+            kdtree as a path to the .pkl file.
 
         Returns:
         --------
@@ -1434,9 +1451,11 @@ def classify_feature(
         `pvdeg.geospatial.feature_downselect`
         """
 
+        self.set_kdtree(kdtree=kdtree)
+
         feature_gids = pvdeg.geospatial.feature_downselect(
             meta_df=self.meta_data,
-            kdtree=kdtree,
+            kdtree=self.kdtree,
             feature_name=feature_name,
             resolution=resolution,
             radius=radius,
@@ -1448,22 +1467,18 @@ def classify_feature(
 
     def downselect_elevation_stochastic(
         self,
-        kdtree,
         downselect_prop,
         k_neighbors=3,
         method="mean",
         normalization="linear",
+        kdtree = None,
     ):
         """
         Prefenetially downselect data points based on elevation and update
         scenario metadata.
 
         Parameters:
         -----------
-        kdtree : sklearn.neighbors.KDTree or str
-            kdtree containing latitude-longitude pairs for quick lookups
-            Generate using ``pvdeg.geospatial.meta_KDTree``. Can take a pickled
-            kdtree as a path to the .pkl file.
         downselect_prop : float
             proportion of original datapoints to keep in output gids list
         k_neighbors : int, (default = 3)
@@ -1474,6 +1489,11 @@ def downselect_elevation_stochastic(
         normalization : str, (default = 'linear')
             function to apply when normalizing weights. Logarithmic uses $log_e$, $ln$
             options : `'linear'`, `'log'`, '`exp'`, `'invert-linear'`
+        kdtree : sklearn.neighbors.KDTree or str
+            Generated automatically but can be provided externally.
+            kdtree containing latitude-longitude pairs for quick lookups
+            Generate using ``pvdeg.geospatial.meta_KDTree``. Can take a pickled
+            kdtree as a path to the .pkl file.
 
         Returns:
         --------
@@ -1483,9 +1503,12 @@ def downselect_elevation_stochastic(
         ---------
         `pvdeg.geospatial.elevation_stochastic_downselect` for more info/docs
         """
+
+        self.set_kdtree(kdtree=kdtree)
+
         gids = pvdeg.geospatial.elevation_stochastic_downselect(
             meta_df=self.meta_data,
-            kdtree=kdtree,
+            kdtree=self.kdtree,
             downselect_prop=downselect_prop,
             k_neighbors=k_neighbors,
             method=method,