# Author: Simon-Pierre Boucher — simon-pierre.boucher@uqo.ca """Spatial/temporal index over an analytic table, built once and shared by comparable retrieval and market context. * one ``scipy.spatial.cKDTree`` on (x_m, y_m) — EPSG:32198 metres, so Euclidean distance = ground distance; * column arrays (dates as int32 days since epoch, class codes, parcel ids, prices, …) for fast masking; * per-municipality and per-region row indices sorted by date (``searchsorted`` date windows). Nothing here looks at the subject's price; the as-of discipline (date < cutoff) is enforced by the callers through the ``date_days`` array. """ from __future__ import annotations from datetime import date from typing import Any import numpy as np import polars as pl from scipy.spatial import cKDTree from ..data.analytic_schema import PROPERTY_CLASSES CLASS_CODE = {c: i for i, c in enumerate(PROPERTY_CLASSES)} EPOCH = date(1970, 1, 1) def to_days(d: date) -> int: return (d - EPOCH).days def from_days(n: int) -> date: return date.fromordinal(EPOCH.toordinal() + int(n)) class AnalyticIndex: """Immutable view of an analytic table with a KD-tree and grouped date-sorted indices.""" def __init__(self, analytic: pl.DataFrame) -> None: df = analytic self.n = df.height self.tx_id = df["tx_id"].to_numpy().astype(object) self.parcel_id = df["parcel_id"].to_numpy().astype(object) self.date_days = df["date"].cast(pl.Int32).to_numpy().astype(np.int32) # Date → days since epoch self.price = df["price"].to_numpy().astype(np.float64) self.x = df["x_m"].to_numpy().astype(np.float64) self.y = df["y_m"].to_numpy().astype(np.float64) self.lat = df["lat"].to_numpy().astype(np.float64) self.lng = df["lng"].to_numpy().astype(np.float64) self.cls = df["property_class"].replace_strict(CLASS_CODE, default=-1, return_dtype=pl.Int8).to_numpy().astype(np.int8) self.floor_area = df["floor_area_m2"].to_numpy().astype(np.float64) self.land_area = df["land_area_m2"].to_numpy().astype(np.float64) self.year_built = df["year_built"].cast(pl.Float64).to_numpy().astype(np.float64) self.storeys = df["storeys"].to_numpy().astype(np.float64) self.n_units = df["n_units"].cast(pl.Float64).to_numpy().astype(np.float64) self.roll_total = df["roll_total_value"].to_numpy().astype(np.float64) self.subtype = df["subtype"].fill_null("unknown").to_numpy().astype(object) self.mun_code = df["mun_code"].to_numpy().astype(object) self.mun_name = df["mun_name"].to_numpy().astype(object) self.region_name = df["region_name"].to_numpy().astype(object) self.ppm2 = np.where(self.floor_area > 0, self.price / self.floor_area, np.nan) self.ratio = np.where(self.roll_total > 0, self.price / self.roll_total, np.nan) self.tree = cKDTree(np.column_stack([self.x, self.y])) self._mun_groups = self._group_sorted(df, "mun_code") self._region_groups = self._group_sorted(df, "region_name") self._tx_pos = {t: i for i, t in enumerate(self.tx_id)} def _group_sorted(self, df: pl.DataFrame, col: str) -> dict[str, np.ndarray]: idx = df.with_row_index("_i").select(col, "_i", pl.col("date").cast(pl.Int32).alias("_d")).sort([col, "_d"]) groups: dict[str, np.ndarray] = {} for key, g in idx.group_by(col, maintain_order=True): groups[str(key[0])] = g["_i"].to_numpy().astype(np.int64) return groups # ---- lookups -------------------------------------------------------------------------------------------- def position(self, tx_id: str) -> int: return self._tx_pos[tx_id] def row(self, tx_id: str) -> dict[str, Any]: i = self.position(tx_id) return { "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], "lat": self.lat[i], "lng": self.lng[i], "property_class": PROPERTY_CLASSES[self.cls[i]], "subtype": self.subtype[i], "floor_area_m2": None if np.isnan(self.floor_area[i]) else float(self.floor_area[i]), "land_area_m2": None if np.isnan(self.land_area[i]) else float(self.land_area[i]), "year_built": None if np.isnan(self.year_built[i]) else int(self.year_built[i]), "mun_code": self.mun_code[i], "mun_name": self.mun_name[i], "region_name": self.region_name[i], } def within_radius(self, x: float, y: float, radius_m: float) -> np.ndarray: """Row indices within ``radius_m`` of (x, y) (unsorted).""" return np.asarray(self.tree.query_ball_point([x, y], r=radius_m), dtype=np.int64) def group_window(self, groups: dict[str, np.ndarray], key: str, start_days: int, end_days: int) -> np.ndarray: """Rows of a group with start_days <= date <= end_days (group indices are date-sorted).""" idx = groups.get(str(key)) if idx is None or idx.size == 0: return np.empty(0, dtype=np.int64) d = self.date_days[idx] lo = np.searchsorted(d, start_days, side="left") hi = np.searchsorted(d, end_days, side="right") return idx[lo:hi] def municipality_window(self, mun_code: str, start_days: int, end_days: int) -> np.ndarray: return self.group_window(self._mun_groups, mun_code, start_days, end_days) def region_window(self, region_name: str, start_days: int, end_days: int) -> np.ndarray: return self.group_window(self._region_groups, region_name, start_days, end_days) _CACHE: dict[tuple[int, int], AnalyticIndex] = {} def get_index(analytic: pl.DataFrame) -> AnalyticIndex: """Build (or reuse) the index of a given analytic frame. Keyed by object identity + height; callers that mutate a frame in place must pass a new object.""" key = (id(analytic), analytic.height) idx = _CACHE.get(key) if idx is None: idx = AnalyticIndex(analytic) _CACHE.clear() _CACHE[key] = idx return idx def subject_to_dict(subject: Any) -> dict[str, Any]: """Accept a dict, a 1-row polars DataFrame, or a polars row (tuple from iter_rows(named=True)).""" if isinstance(subject, dict): return subject if isinstance(subject, pl.DataFrame): if subject.height != 1: raise ValueError("subject frame must have exactly one row") return subject.row(0, named=True) raise TypeError(f"unsupported subject type {type(subject)!r}")