spb/uqo-research-db

1"""Build analysis datasets from spid.duckdb into research/data/*.parquet.23Outputs4-------5projects.parquet   project cross-section (cleaned, with AI/digital flags, age)6mentions.parquet   mention-level events with filing_date, form_type7panel.parquet      balanced firm x year panel (counts, capex, type shares, AI)8kg_nodes/edges     copied knowledge-graph tables9embeddings.parquet project_id + 384-d vector10"""11from __future__ import annotations1213import numpy as np14import pandas as pd1516from common import con, DATA, save_results1718YEAR_MIN, YEAR_MAX = 2010, 2025  # 2026 is partial -> excluded from panel1920AI_TECH = ("ai", "artificial intelligence", "machine learning", "deep learning",21           "generative", "llm", "neural")22DIGITAL_TYPES = ("ai_initiative", "digital_transformation", "cloud_migration",23                 "erp_implementation", "data_center", "automation",24                 "technology_deployment")252627def _is_ai(ptype: str, techs) -> bool:28    if ptype == "ai_initiative":29        return True30    for t in (techs or []):31        tl = str(t).lower()32        if any(k in tl for k in AI_TECH):33            return True34    return False353637def main() -> None:38    c = con()3940    # ---- project cross-section -------------------------------------------41    proj = c.execute("""42        SELECT project_id, cik, ticker, company_name, sector, project_type,43               project_name, canonical_location, technologies, status,44               total_amount_usd, n_mentions, n_filings,45               first_seen, last_seen, first_year, last_year, avg_confidence46        FROM projects47    """).df()48    proj["technologies"] = proj["technologies"].apply(lambda x: list(x) if x is not None else [])49    proj["is_ai"] = [50        _is_ai(t, tech) for t, tech in zip(proj["project_type"], proj["technologies"])51    ]52    proj["is_digital"] = proj["project_type"].isin(DIGITAL_TYPES) | proj["is_ai"]53    proj["has_amount"] = proj["total_amount_usd"].fillna(0) > 054    proj["log_amount"] = np.log(proj["total_amount_usd"].where(proj["has_amount"]))55    proj["first_seen"] = pd.to_datetime(proj["first_seen"])56    proj["last_seen"] = pd.to_datetime(proj["last_seen"])57    proj["duration_days"] = (proj["last_seen"] - proj["first_seen"]).dt.days58    proj["n_tech"] = proj["technologies"].apply(len)59    # announcement year60    proj["ann_year"] = proj["first_year"].fillna(proj["first_seen"].dt.year)61    proj.to_parquet(DATA / "projects.parquet")62    print(f"projects: {len(proj):,}")6364    # ---- mention-level ----------------------------------------------------65    men = c.execute("""66        SELECT mention_id, cik, ticker, sector, project_type, status,67               amount_usd, form_type, filing_date, fiscal_year, confidence68        FROM project_mentions69    """).df()70    men["filing_date"] = pd.to_datetime(men["filing_date"])71    men["year"] = men["filing_date"].dt.year72    men.to_parquet(DATA / "mentions.parquet")73    print(f"mentions: {len(men):,}")7475    # ---- firm coverage (years each firm appears in filings) --------------76    cov = c.execute("""77        SELECT cik, MIN(EXTRACT(year FROM CAST(filing_date AS DATE))) y0,78                    MAX(EXTRACT(year FROM CAST(filing_date AS DATE))) y179        FROM src_filings GROUP BY cik80    """).df()81    cov["y0"] = cov["y0"].clip(lower=YEAR_MIN).astype(int)82    cov["y1"] = cov["y1"].clip(upper=YEAR_MAX).astype(int)8384    # firm meta (modal sector/ticker)85    meta = (proj.groupby("cik")86            .agg(ticker=("ticker", lambda s: s.dropna().iloc[0] if s.notna().any() else None),87                 sector=("sector", lambda s: s.dropna().mode().iloc[0] if s.notna().any() else "Unknown"),88                 company_name=("company_name", lambda s: s.dropna().iloc[0] if s.notna().any() else None))89            .reset_index())9091    # ---- build balanced firm x year skeleton -----------------------------92    rows = []93    for _, r in cov.iterrows():94        if r.y1 < r.y0:95            continue96        for y in range(int(r.y0), int(r.y1) + 1):97            rows.append((r.cik, y))98    panel = pd.DataFrame(rows, columns=["cik", "year"])99100    # project announcements aggregated by firm-year101    pa = proj[(proj["ann_year"] >= YEAR_MIN) & (proj["ann_year"] <= YEAR_MAX)].copy()102    pa["ann_year"] = pa["ann_year"].astype(int)103    g = pa.groupby(["cik", "ann_year"])104    agg = g.agg(105        n_projects=("project_id", "count"),106        n_ai=("is_ai", "sum"),107        n_digital=("is_digital", "sum"),108        capex_usd=("total_amount_usd", lambda s: np.nansum(s.values)),109        n_with_amount=("has_amount", "sum"),110        mean_conf=("avg_confidence", "mean"),111    ).reset_index().rename(columns={"ann_year": "year"})112    # type shares (selected types)113    for t in ("ma_integration", "capex_program", "plant_construction",114              "rd_program", "product_launch", "sustainability"):115        s = (pa[pa["project_type"] == t].groupby(["cik", "ann_year"])["project_id"]116             .count().reset_index().rename(columns={"project_id": f"n_{t}", "ann_year": "year"}))117        agg = agg.merge(s, on=["cik", "year"], how="left")118119    panel = panel.merge(agg, on=["cik", "year"], how="left")120    cnt_cols = [c_ for c_ in panel.columns if c_.startswith("n_") or c_ == "capex_usd"]121    panel[cnt_cols] = panel[cnt_cols].fillna(0)122    panel = panel.merge(meta, on="cik", how="left")123    panel["log_capex"] = np.log1p(panel["capex_usd"])124    panel["any_ai"] = (panel["n_ai"] > 0).astype(int)125    panel["post_covid"] = (panel["year"] >= 2020).astype(int)126    panel = panel.sort_values(["cik", "year"]).reset_index(drop=True)127    # lagged dependent128    panel["n_projects_lag"] = panel.groupby("cik")["n_projects"].shift(1)129    panel["cum_ai"] = panel.groupby("cik")["n_ai"].cumsum()130    panel.to_parquet(DATA / "panel.parquet")131    print(f"panel: {len(panel):,} firm-years, {panel['cik'].nunique()} firms, "132          f"{panel['year'].min()}-{panel['year'].max()}")133134    # ---- KG + embeddings --------------------------------------------------135    c.execute(f"COPY (SELECT * FROM kg_nodes) TO '{DATA/'kg_nodes.parquet'}' (FORMAT PARQUET)")136    c.execute(f"COPY (SELECT * FROM kg_edges) TO '{DATA/'kg_edges.parquet'}' (FORMAT PARQUET)")137    emb = c.execute("SELECT project_id, vector FROM project_embeddings").df()138    emb.to_parquet(DATA / "embeddings.parquet")139    print(f"kg_nodes/edges + embeddings: {len(emb):,} vectors")140141    save_results("dataset_summary", {142        "n_projects": int(len(proj)),143        "n_mentions": int(len(men)),144        "n_firms": int(proj["cik"].nunique()),145        "n_sectors": int(proj["sector"].nunique()),146        "n_types": int(proj["project_type"].nunique()),147        "n_locations": int(proj["canonical_location"].nunique()),148        "panel_firm_years": int(len(panel)),149        "year_range": [int(men["year"].min()), int(men["year"].max())],150        "projects_with_amount": int(proj["has_amount"].sum()),151        "share_ai": float(proj["is_ai"].mean()),152        "total_capex_usd": float(proj["total_amount_usd"].fillna(0).sum()),153    })154    c.close()155156157if __name__ == "__main__":158    main()159