spb/uqo-research-db
Public
1# ============================================================2# Project: EDGAR Research Database3# Description: SEC EDGAR Research Infrastructure4#5# Authors:6# - Manel Kammoun — manel.kammoun@uqo.ca7# - Djerrycharli Tandjambianda — djerrycharli.tandjambianda@uqo.ca8# - Simon-Pierre Boucher — simon-pierre.boucher@uqo.ca9#10# Institution:11# Universite du Quebec en Outaouais (UQO)12#13# License:14# Research / Academic Use15# ============================================================1617from __future__ import annotations1819from typing import TYPE_CHECKING2021import numpy as np22import pandas as pd23import pyarrow.dataset as ds24from tqdm import tqdm2526from src.utils.config import get_config27from src.utils.constants import DATA_DIRS28from src.utils.logging import get_logger2930if TYPE_CHECKING:31 from src.storage.database import DatabaseManager3233logger = get_logger("factors.longitudinal")343536def _firm_year_centroids(chunk_to_firmyear: dict) -> dict[tuple[str, int], np.ndarray]:37 """Mean (then unit-normalised) embedding per (ticker, year), streamed from the parquet store."""38 from src.factors.scoring import _vectors_from_column3940 keys = sorted(set(chunk_to_firmyear.values()), key=lambda k: (str(k[0]), int(k[1]) if k[1] is not None else -1))41 key_index = {k: i for i, k in enumerate(keys)}42 chunk_index = {cid: key_index[k] for cid, k in chunk_to_firmyear.items()}4344 sums: np.ndarray | None = None45 counts = np.zeros(len(keys), dtype=np.int64)4647 emb_dir = DATA_DIRS["parquet"] / "embeddings"48 dataset = ds.dataset(str(emb_dir), format="parquet", partitioning="hive")49 scanner = dataset.scanner(columns=["source_id", "vector"], batch_size=20_000)50 for batch in tqdm(scanner.to_batches(), desc="Accumulating firm-year centroids", unit="batch"):51 if batch.num_rows == 0:52 continue53 ids = batch.column("source_id").to_pylist()54 idx = np.fromiter((chunk_index.get(i, -1) for i in ids), dtype=np.int64, count=len(ids))55 mask = idx >= 056 if not mask.any():57 continue58 vectors = _vectors_from_column(batch.column("vector"))[mask].astype(np.float64)59 if sums is None:60 sums = np.zeros((len(keys), vectors.shape[1]), dtype=np.float64)61 np.add.at(sums, idx[mask], vectors)62 np.add.at(counts, idx[mask], 1)6364 centroids: dict[tuple[str, int], np.ndarray] = {}65 if sums is None:66 return centroids67 for k, i in key_index.items():68 if counts[i] == 0:69 continue70 c = (sums[i] / counts[i]).astype(np.float32)71 n = np.linalg.norm(c)72 centroids[k] = c / n if n > 0 else c73 return centroids747576class LongitudinalBuilder:7778 def __init__(self, db: DatabaseManager) -> None:79 self.db = db80 cfg = get_config("factors").get("factors", {}).get("longitudinal", {})81 self._rolling_window = cfg.get("rolling_window", 3)82 self._min_years_drift = cfg.get("min_years_for_drift", 2)8384 def add_yoy_changes(85 self,86 panel: pd.DataFrame,87 anchor_ids: list[str],88 ) -> pd.DataFrame:89 panel = panel.sort_values(["ticker", "fiscal_year"]).copy()90 for aid in anchor_ids:91 col = f"{aid}_mean"92 if col not in panel.columns:93 continue94 panel[f"{aid}_delta"] = panel.groupby("ticker")[col].diff(1)95 panel[f"{aid}_pct_change"] = panel.groupby("ticker")[col].pct_change(1)96 logger.info("Added YoY changes for %d anchors", len(anchor_ids))97 return panel9899 def add_acceleration(100 self,101 panel: pd.DataFrame,102 anchor_ids: list[str],103 ) -> pd.DataFrame:104 panel = panel.sort_values(["ticker", "fiscal_year"]).copy()105 for aid in anchor_ids:106 delta_col = f"{aid}_delta"107 if delta_col not in panel.columns:108 continue109 panel[f"{aid}_accel"] = panel.groupby("ticker")[delta_col].diff(1)110 logger.info("Added acceleration for %d anchors", len(anchor_ids))111 return panel112113 def add_rolling_stats(114 self,115 panel: pd.DataFrame,116 anchor_ids: list[str],117 window: int | None = None,118 ) -> pd.DataFrame:119 w = window or self._rolling_window120 panel = panel.sort_values(["ticker", "fiscal_year"]).copy()121 for aid in anchor_ids:122 col = f"{aid}_mean"123 if col not in panel.columns:124 continue125 grp = panel.groupby("ticker")[col]126 panel[f"{aid}_roll{w}_mean"] = grp.transform(127 lambda x: x.rolling(w, min_periods=2).mean()128 )129 panel[f"{aid}_roll{w}_std"] = grp.transform(130 lambda x: x.rolling(w, min_periods=2).std()131 )132 logger.info("Added rolling stats (window=%d) for %d anchors", w, len(anchor_ids))133 return panel134135 def compute_embedding_drift(self) -> pd.DataFrame:136 logger.info("Computing firm-year embedding centroids for drift...")137138 chunk_years = self.db.fetch_df("""139 SELECT140 c.chunk_id,141 f.ticker,142 EXTRACT(YEAR FROM f.filing_date)::INTEGER AS fiscal_year143 FROM filing_chunks c144 JOIN filings f ON c.accession_number = f.accession_number145 """)146147 chunk_to_firmyear = dict(148 zip(chunk_years["chunk_id"], zip(chunk_years["ticker"], chunk_years["fiscal_year"]))149 )150151 centroids = _firm_year_centroids(chunk_to_firmyear)152 del chunk_to_firmyear153154 tickers = sorted(set(k[0] for k in centroids))155 drift_records = []156 for ticker in tickers:157 years = sorted(y for t, y in centroids if t == ticker)158 for i in range(1, len(years)):159 y_prev, y_curr = years[i - 1], years[i]160 if y_curr - y_prev > 2:161 continue162 c_prev = centroids[(ticker, y_prev)]163 c_curr = centroids[(ticker, y_curr)]164 sim = float(np.dot(c_prev, c_curr))165 drift_records.append({166 "ticker": ticker,167 "fiscal_year": y_curr,168 "embedding_drift": 1.0 - sim,169 })170171 drift_df = pd.DataFrame(drift_records)172 logger.info("Computed embedding drift for %d firm-year observations", len(drift_df))173 return drift_df174175 def compute_novelty(self) -> pd.DataFrame:176 logger.info("Computing novelty scores (distance from sector centroid)...")177178 companies = self.db.fetch_df("SELECT cik, ticker, sector FROM companies")179 ticker_to_sector = dict(zip(companies["ticker"], companies["sector"]))180181 annual = self.db.fetch_df("SELECT ticker, fiscal_year FROM annual_panel")182183 chunk_years = self.db.fetch_df("""184 SELECT c.chunk_id, f.ticker,185 EXTRACT(YEAR FROM f.filing_date)::INTEGER AS fiscal_year186 FROM filing_chunks c187 JOIN filings f ON c.accession_number = f.accession_number188 """)189 chunk_to_firmyear = dict(190 zip(chunk_years["chunk_id"], zip(chunk_years["ticker"], chunk_years["fiscal_year"]))191 )192193 firm_centroids = _firm_year_centroids(chunk_to_firmyear)194 del chunk_to_firmyear195196 sector_year_sums: dict[tuple[str, int], np.ndarray] = {}197 sector_year_counts: dict[tuple[str, int], int] = {}198 for (ticker, year), centroid in firm_centroids.items():199 sector = ticker_to_sector.get(ticker, "Unknown")200 key = (sector, year)201 if key not in sector_year_sums:202 sector_year_sums[key] = np.zeros_like(centroid, dtype=np.float64)203 sector_year_counts[key] = 0204 sector_year_sums[key] += centroid.astype(np.float64)205 sector_year_counts[key] += 1206207 sector_centroids: dict[tuple[str, int], np.ndarray] = {}208 for key, s in sector_year_sums.items():209 centroid = (s / sector_year_counts[key]).astype(np.float32)210 norm = np.linalg.norm(centroid)211 if norm > 0:212 centroid = centroid / norm213 sector_centroids[key] = centroid214215 novelty_records = []216 for (ticker, year), firm_c in firm_centroids.items():217 sector = ticker_to_sector.get(ticker, "Unknown")218 sec_c = sector_centroids.get((sector, year))219 if sec_c is None:220 continue221 sim = float(np.dot(firm_c, sec_c))222 novelty_records.append({223 "ticker": ticker,224 "fiscal_year": year,225 "novelty_score": 1.0 - sim,226 })227228 novelty_df = pd.DataFrame(novelty_records)229 logger.info("Computed novelty for %d firm-year observations", len(novelty_df))230 return novelty_df231