Add type hints to LayeredActiveSpace

nps_active_space/active_space/layered_active_space.py

@@ -1,3 +1,5 @@
+from __future__ import annotations
+
 import numpy as np
 import os
 import geopandas as gpd
@@ -10,54 +12,56 @@
 
 
 class LayeredActiveSpace():
-    def __init__(self, designator, layer_dirs, study_area, gain=None, crs="epsg:4326"):
-        self.designator = designator
-        self.layer_dirs = dict(sorted(layer_dirs.items()))
-        self.study_area = study_area
-        self.activespaces = {}
-        self.all_activespaces = {}  # set by self.preload_all_activespaces()
-        self.crs = crs
-        self.gain = None
+    def __init__(self, designator: str, layer_dirs: dict[int, str], study_area: gpd.GeoDataFrame,
+                 gain: float | None = None, crs: str = "epsg:4326") -> None:
+        self.designator: str = designator
+        self.layer_dirs: dict[int, str] = dict(sorted(layer_dirs.items()))
+        self.study_area: gpd.GeoDataFrame = study_area
+        self.activespaces: dict[int, gpd.GeoDataFrame] = {}
+        self.all_activespaces: dict[float, dict[int, gpd.GeoDataFrame] | None] = {}  # set by self.preload_all_activespaces()
+        self.crs: str = crs
+        self.gain: float | None = None
         if gain is not None:
             self.set_gain(gain)
-        self.fit_pbar = None
+        self.fit_pbar: tqdm | None = None
 
         # determine min and max gain - import here to avoid circular import
         from nps_active_space.utils.helpers import omni_to_gain
         first_layer_dir = list(self.layer_dirs.values())[0]
         active_names = glob.glob(os.path.join(first_layer_dir, "*_O_*.geojson"))
         gains = list(map(lambda f: omni_to_gain(f), active_names))
-        self.min_gain = min(gains)
-        self.max_gain = max(gains)
+        self.min_gain: float = min(gains)
+        self.max_gain: float = max(gains)
 
-    def load_activespaces(self, gain):
+    def load_activespaces(self, gain: float) -> dict[int, gpd.GeoDataFrame] | None:
         if gain in self.all_activespaces:
             return self.all_activespaces[gain]
 
         sign = "-" if gain < 0 else "+"
         gain_string = str(np.abs(int(10*gain))).zfill(3)
 
-        activespaces = {}
+        activespaces: dict[int, gpd.GeoDataFrame] = {}
         for altitude, dir in self.layer_dirs.items():
             glob_result = glob.glob(os.path.join(dir, f"*_O_{sign}{gain_string}.geojson"))
             if len(glob_result) == 0:
                 print(f"Couldn't find active space for gain {gain} in {dir}")
-                return
+                return None
             activespace_file = glob_result[0]
             activespaces[altitude] = gpd.read_file(activespace_file).to_crs(self.crs)
 
         return activespaces
 
-    def preload_all_activespaces(self):
+    def preload_all_activespaces(self) -> None:
         self.all_activespaces = {}
         for gain in tqdm(np.arange(self.min_gain, self.max_gain + 0.5, 0.5), desc="Loading all active spaces"):
             self.all_activespaces[gain] = self.load_activespaces(gain)
 
-    def set_gain(self, gain):
+    def set_gain(self, gain: float) -> None:
         self.activespaces = self.load_activespaces(gain)
         self.gain = gain
 
-    def fit(self, annotations, beta=1., plot=True, plot_savepath=None):
+    def fit(self, annotations: gpd.GeoDataFrame, beta: float = 1., plot: bool = True,
+            plot_savepath: str | None = None) -> pd.Series:
 
         # Extract all valid points from their LineStrings. These will be needed for calculating fbeta scores later.
         valid_points_lst = []
@@ -69,7 +73,8 @@ def fit(self, annotations, beta=1., plot=True, plot_savepath=None):
         result["Number of valid annotated segments"] = len(annotations)
         return result
 
-    def fit_points(self, points, min_gain=-10., max_gain=40., beta=1., plot=True, plot_savepath=None):
+    def fit_points(self, points: gpd.GeoDataFrame, min_gain: float = -10., max_gain: float = 40.,
+                   beta: float = 1., plot: bool = True, plot_savepath: str | None = None) -> pd.Series:
         assert min_gain <= max_gain
 
         # Reduce point density to median density, so very dense areas (e.g. airports) don't skew the fit
@@ -142,20 +147,20 @@ def fit_points(self, points, min_gain=-10., max_gain=40., beta=1., plot=True, pl
             f"F{beta}": best[f"F{beta}"]
         })
 
-    def assign_layers(self, points):
+    def assign_layers(self, points: gpd.GeoDataFrame) -> gpd.GeoDataFrame:
         """
         Returns a copy of points with a new column added representing the active space layer
         each point belongs to (which layer is closest to the point's z value)
         """
         points = points.copy()
         altitudes = list(self.layer_dirs.keys())
-        def closest_layer(z):
+        def closest_layer(z: float) -> int:
             return min(altitudes, key=lambda alt: abs(alt - z))
         points["layer"] = points.geometry.z.apply(closest_layer)
         return points
 
 
-    def predict(self, points):
+    def predict(self, points: gpd.GeoDataFrame) -> pd.Series:
         """Given a GeoDataFrame of 3D points, predict whether they are audible.
 
         Returns
@@ -179,4 +184,4 @@ def predict(self, points):
             in_AS_gdf = gpd.clip(points[layer_mask], activespace)
             points.loc[in_AS_gdf.index, "in_AS"] = True
 
-        return points["in_AS"]
+        return points["in_AS"]