1# Author: Simon-Pierre Boucher — simon-pierre.boucher@uqo.ca2"""Spatial/temporal index over an analytic table, built once and shared by comparable retrieval and market context.34* one ``scipy.spatial.cKDTree`` on (x_m, y_m) — EPSG:32198 metres, so Euclidean distance = ground distance;5* column arrays (dates as int32 days since epoch, class codes, parcel ids, prices, …) for fast masking;6* per-municipality and per-region row indices sorted by date (``searchsorted`` date windows).78Nothing here looks at the subject's price; the as-of discipline (date < cutoff) is enforced by the callers9through the ``date_days`` array.10"""11from __future__ import annotations1213from datetime import date14from typing import Any1516import numpy as np17import polars as pl18from scipy.spatial import cKDTree1920from ..data.analytic_schema import PROPERTY_CLASSES2122CLASS_CODE = {c: i for i, c in enumerate(PROPERTY_CLASSES)}23EPOCH = date(1970, 1, 1)242526def to_days(d: date) -> int:27 return (d - EPOCH).days282930def from_days(n: int) -> date:31 return date.fromordinal(EPOCH.toordinal() + int(n))323334class AnalyticIndex:35 """Immutable view of an analytic table with a KD-tree and grouped date-sorted indices."""3637 def __init__(self, analytic: pl.DataFrame) -> None:38 df = analytic39 self.n = df.height40 self.tx_id = df["tx_id"].to_numpy().astype(object)41 self.parcel_id = df["parcel_id"].to_numpy().astype(object)42 self.date_days = df["date"].cast(pl.Int32).to_numpy().astype(np.int32) # Date → days since epoch43 self.price = df["price"].to_numpy().astype(np.float64)44 self.x = df["x_m"].to_numpy().astype(np.float64)45 self.y = df["y_m"].to_numpy().astype(np.float64)46 self.lat = df["lat"].to_numpy().astype(np.float64)47 self.lng = df["lng"].to_numpy().astype(np.float64)48 self.cls = df["property_class"].replace_strict(CLASS_CODE, default=-1, return_dtype=pl.Int8).to_numpy().astype(np.int8)49 self.floor_area = df["floor_area_m2"].to_numpy().astype(np.float64)50 self.land_area = df["land_area_m2"].to_numpy().astype(np.float64)51 self.year_built = df["year_built"].cast(pl.Float64).to_numpy().astype(np.float64)52 self.storeys = df["storeys"].to_numpy().astype(np.float64)53 self.n_units = df["n_units"].cast(pl.Float64).to_numpy().astype(np.float64)54 self.roll_total = df["roll_total_value"].to_numpy().astype(np.float64)55 self.subtype = df["subtype"].fill_null("unknown").to_numpy().astype(object)56 self.mun_code = df["mun_code"].to_numpy().astype(object)57 self.mun_name = df["mun_name"].to_numpy().astype(object)58 self.region_name = df["region_name"].to_numpy().astype(object)59 self.ppm2 = np.where(self.floor_area > 0, self.price / self.floor_area, np.nan)60 self.ratio = np.where(self.roll_total > 0, self.price / self.roll_total, np.nan)61 self.tree = cKDTree(np.column_stack([self.x, self.y]))62 self._mun_groups = self._group_sorted(df, "mun_code")63 self._region_groups = self._group_sorted(df, "region_name")64 self._tx_pos = {t: i for i, t in enumerate(self.tx_id)}6566 def _group_sorted(self, df: pl.DataFrame, col: str) -> dict[str, np.ndarray]:67 idx = df.with_row_index("_i").select(col, "_i", pl.col("date").cast(pl.Int32).alias("_d")).sort([col, "_d"])68 groups: dict[str, np.ndarray] = {}69 for key, g in idx.group_by(col, maintain_order=True):70 groups[str(key[0])] = g["_i"].to_numpy().astype(np.int64)71 return groups7273 # ---- lookups --------------------------------------------------------------------------------------------74 def position(self, tx_id: str) -> int:75 return self._tx_pos[tx_id]7677 def row(self, tx_id: str) -> dict[str, Any]:78 i = self.position(tx_id)79 return {80 "tx_id": tx_id, "parcel_id": self.parcel_id[i], "date": from_days(self.date_days[i]), "x_m": self.x[i], "y_m": self.y[i],81 "lat": self.lat[i], "lng": self.lng[i], "property_class": PROPERTY_CLASSES[self.cls[i]], "subtype": self.subtype[i],82 "floor_area_m2": None if np.isnan(self.floor_area[i]) else float(self.floor_area[i]),83 "land_area_m2": None if np.isnan(self.land_area[i]) else float(self.land_area[i]),84 "year_built": None if np.isnan(self.year_built[i]) else int(self.year_built[i]),85 "mun_code": self.mun_code[i], "mun_name": self.mun_name[i], "region_name": self.region_name[i],86 }8788 def within_radius(self, x: float, y: float, radius_m: float) -> np.ndarray:89 """Row indices within ``radius_m`` of (x, y) (unsorted)."""90 return np.asarray(self.tree.query_ball_point([x, y], r=radius_m), dtype=np.int64)9192 def group_window(self, groups: dict[str, np.ndarray], key: str, start_days: int, end_days: int) -> np.ndarray:93 """Rows of a group with start_days <= date <= end_days (group indices are date-sorted)."""94 idx = groups.get(str(key))95 if idx is None or idx.size == 0:96 return np.empty(0, dtype=np.int64)97 d = self.date_days[idx]98 lo = np.searchsorted(d, start_days, side="left")99 hi = np.searchsorted(d, end_days, side="right")100 return idx[lo:hi]101102 def municipality_window(self, mun_code: str, start_days: int, end_days: int) -> np.ndarray:103 return self.group_window(self._mun_groups, mun_code, start_days, end_days)104105 def region_window(self, region_name: str, start_days: int, end_days: int) -> np.ndarray:106 return self.group_window(self._region_groups, region_name, start_days, end_days)107108109_CACHE: dict[tuple[int, int], AnalyticIndex] = {}110111112def get_index(analytic: pl.DataFrame) -> AnalyticIndex:113 """Build (or reuse) the index of a given analytic frame. Keyed by object identity + height; callers that114 mutate a frame in place must pass a new object."""115 key = (id(analytic), analytic.height)116 idx = _CACHE.get(key)117 if idx is None:118 idx = AnalyticIndex(analytic)119 _CACHE.clear()120 _CACHE[key] = idx121 return idx122123124def subject_to_dict(subject: Any) -> dict[str, Any]:125 """Accept a dict, a 1-row polars DataFrame, or a polars row (tuple from iter_rows(named=True))."""126 if isinstance(subject, dict):127 return subject128 if isinstance(subject, pl.DataFrame):129 if subject.height != 1:130 raise ValueError("subject frame must have exactly one row")131 return subject.row(0, named=True)132 raise TypeError(f"unsupported subject type {type(subject)!r}")133