diff --git a/nps_active_space/active_space/layered_active_space.py b/nps_active_space/active_space/layered_active_space.py index dadd285..0a15d8a 100644 --- a/nps_active_space/active_space/layered_active_space.py +++ b/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"] \ No newline at end of file + return points["in_AS"]