"""Build analysis datasets from spid.duckdb into research/data/*.parquet. Outputs ------- projects.parquet project cross-section (cleaned, with AI/digital flags, age) mentions.parquet mention-level events with filing_date, form_type panel.parquet balanced firm x year panel (counts, capex, type shares, AI) kg_nodes/edges copied knowledge-graph tables embeddings.parquet project_id + 384-d vector """ from __future__ import annotations import numpy as np import pandas as pd from common import con, DATA, save_results YEAR_MIN, YEAR_MAX = 2010, 2025 # 2026 is partial -> excluded from panel AI_TECH = ("ai", "artificial intelligence", "machine learning", "deep learning", "generative", "llm", "neural") DIGITAL_TYPES = ("ai_initiative", "digital_transformation", "cloud_migration", "erp_implementation", "data_center", "automation", "technology_deployment") def _is_ai(ptype: str, techs) -> bool: if ptype == "ai_initiative": return True for t in (techs or []): tl = str(t).lower() if any(k in tl for k in AI_TECH): return True return False def main() -> None: c = con() # ---- project cross-section ------------------------------------------- proj = c.execute(""" SELECT project_id, cik, ticker, company_name, sector, project_type, project_name, canonical_location, technologies, status, total_amount_usd, n_mentions, n_filings, first_seen, last_seen, first_year, last_year, avg_confidence FROM projects """).df() proj["technologies"] = proj["technologies"].apply(lambda x: list(x) if x is not None else []) proj["is_ai"] = [ _is_ai(t, tech) for t, tech in zip(proj["project_type"], proj["technologies"]) ] proj["is_digital"] = proj["project_type"].isin(DIGITAL_TYPES) | proj["is_ai"] proj["has_amount"] = proj["total_amount_usd"].fillna(0) > 0 proj["log_amount"] = np.log(proj["total_amount_usd"].where(proj["has_amount"])) proj["first_seen"] = pd.to_datetime(proj["first_seen"]) proj["last_seen"] = pd.to_datetime(proj["last_seen"]) proj["duration_days"] = (proj["last_seen"] - proj["first_seen"]).dt.days proj["n_tech"] = proj["technologies"].apply(len) # announcement year proj["ann_year"] = proj["first_year"].fillna(proj["first_seen"].dt.year) proj.to_parquet(DATA / "projects.parquet") print(f"projects: {len(proj):,}") # ---- mention-level ---------------------------------------------------- men = c.execute(""" SELECT mention_id, cik, ticker, sector, project_type, status, amount_usd, form_type, filing_date, fiscal_year, confidence FROM project_mentions """).df() men["filing_date"] = pd.to_datetime(men["filing_date"]) men["year"] = men["filing_date"].dt.year men.to_parquet(DATA / "mentions.parquet") print(f"mentions: {len(men):,}") # ---- firm coverage (years each firm appears in filings) -------------- cov = c.execute(""" SELECT cik, MIN(EXTRACT(year FROM CAST(filing_date AS DATE))) y0, MAX(EXTRACT(year FROM CAST(filing_date AS DATE))) y1 FROM src_filings GROUP BY cik """).df() cov["y0"] = cov["y0"].clip(lower=YEAR_MIN).astype(int) cov["y1"] = cov["y1"].clip(upper=YEAR_MAX).astype(int) # firm meta (modal sector/ticker) meta = (proj.groupby("cik") .agg(ticker=("ticker", lambda s: s.dropna().iloc[0] if s.notna().any() else None), sector=("sector", lambda s: s.dropna().mode().iloc[0] if s.notna().any() else "Unknown"), company_name=("company_name", lambda s: s.dropna().iloc[0] if s.notna().any() else None)) .reset_index()) # ---- build balanced firm x year skeleton ----------------------------- rows = [] for _, r in cov.iterrows(): if r.y1 < r.y0: continue for y in range(int(r.y0), int(r.y1) + 1): rows.append((r.cik, y)) panel = pd.DataFrame(rows, columns=["cik", "year"]) # project announcements aggregated by firm-year pa = proj[(proj["ann_year"] >= YEAR_MIN) & (proj["ann_year"] <= YEAR_MAX)].copy() pa["ann_year"] = pa["ann_year"].astype(int) g = pa.groupby(["cik", "ann_year"]) agg = g.agg( n_projects=("project_id", "count"), n_ai=("is_ai", "sum"), n_digital=("is_digital", "sum"), capex_usd=("total_amount_usd", lambda s: np.nansum(s.values)), n_with_amount=("has_amount", "sum"), mean_conf=("avg_confidence", "mean"), ).reset_index().rename(columns={"ann_year": "year"}) # type shares (selected types) for t in ("ma_integration", "capex_program", "plant_construction", "rd_program", "product_launch", "sustainability"): s = (pa[pa["project_type"] == t].groupby(["cik", "ann_year"])["project_id"] .count().reset_index().rename(columns={"project_id": f"n_{t}", "ann_year": "year"})) agg = agg.merge(s, on=["cik", "year"], how="left") panel = panel.merge(agg, on=["cik", "year"], how="left") cnt_cols = [c_ for c_ in panel.columns if c_.startswith("n_") or c_ == "capex_usd"] panel[cnt_cols] = panel[cnt_cols].fillna(0) panel = panel.merge(meta, on="cik", how="left") panel["log_capex"] = np.log1p(panel["capex_usd"]) panel["any_ai"] = (panel["n_ai"] > 0).astype(int) panel["post_covid"] = (panel["year"] >= 2020).astype(int) panel = panel.sort_values(["cik", "year"]).reset_index(drop=True) # lagged dependent panel["n_projects_lag"] = panel.groupby("cik")["n_projects"].shift(1) panel["cum_ai"] = panel.groupby("cik")["n_ai"].cumsum() panel.to_parquet(DATA / "panel.parquet") print(f"panel: {len(panel):,} firm-years, {panel['cik'].nunique()} firms, " f"{panel['year'].min()}-{panel['year'].max()}") # ---- KG + embeddings -------------------------------------------------- c.execute(f"COPY (SELECT * FROM kg_nodes) TO '{DATA/'kg_nodes.parquet'}' (FORMAT PARQUET)") c.execute(f"COPY (SELECT * FROM kg_edges) TO '{DATA/'kg_edges.parquet'}' (FORMAT PARQUET)") emb = c.execute("SELECT project_id, vector FROM project_embeddings").df() emb.to_parquet(DATA / "embeddings.parquet") print(f"kg_nodes/edges + embeddings: {len(emb):,} vectors") save_results("dataset_summary", { "n_projects": int(len(proj)), "n_mentions": int(len(men)), "n_firms": int(proj["cik"].nunique()), "n_sectors": int(proj["sector"].nunique()), "n_types": int(proj["project_type"].nunique()), "n_locations": int(proj["canonical_location"].nunique()), "panel_firm_years": int(len(panel)), "year_range": [int(men["year"].min()), int(men["year"].max())], "projects_with_amount": int(proj["has_amount"].sum()), "share_ai": float(proj["is_ai"].mean()), "total_capex_usd": float(proj["total_amount_usd"].fillna(0).sum()), }) c.close() if __name__ == "__main__": main()