1#!/usr/bin/env python32"""WP44 — Phase 2: assign the 40 ontology categories (multi-label) to every work of the corpus.34Signals per work: OpenAlex topics (ids + subfield ids), arXiv categories (through link_arxiv_oa), DBLP venue stream5(through link_dblp_oa), keyword regexes on title and abstract. Scoring follows ontology.yaml `assignment_rule`:6per signal type the matched weights are summed and capped at 1.0, then summed over signal types; a category is assigned7if score >= threshold; primary = argmax (ties: child before parent, then lower category number).89Step 1 (`--prepare`): DuckDB builds data/processed/classification/input/chunk_*.parquet (id, year, title, abstract,10topic_ids, subfield_ids, arxiv_cats, dblp_streams). Step 2 (`--run`): multiprocessing over chunks →11data/processed/classification/work_categories.parquet (work_id, category_id, score, s_topic, s_subfield, s_arxiv,12s_dblp, s_kw) and work_primary.parquet (work_id, primary_category, n_categories). Deterministic; version-stamped.13"""14import argparse, glob, json, multiprocessing as mp, os, re, time, datetime as dt15import yaml, duckdb16import pyarrow as pa, pyarrow.parquet as pq1718ONT = None192021def load_ontology(path):22 o = yaml.safe_load(open(path))23 cats = o["categories"]24 comp = []25 for c in cats:26 kws = []27 for k in (c.get("keywords") or []):28 kws.append((re.compile(k["pattern"]), float(k["weight"]), k.get("field", "both")))29 exc = [re.compile(p if isinstance(p, str) else p["pattern"]) for p in (c.get("exclusions") or [])]30 comp.append(dict(31 id=c["id"], number=int(c["number"]), parent=tuple(c["child_of"]) if isinstance(c.get("child_of"), list) else ((c["child_of"],) if c.get("child_of") else ()),32 topics={str(t["id"]).lstrip("T"): float(t["weight"]) for t in (c.get("openalex_topics") or [])},33 subfields={str(t["id"]): float(t["weight"]) for t in (c.get("openalex_subfields") or [])},34 arxiv={t["id"]: float(t["weight"]) for t in (c.get("arxiv_categories") or [])},35 dblp={t["key"]: float(t["weight"]) for t in (c.get("dblp_streams") or [])},36 kws=kws, exc=exc))37 rule = o.get("assignment_rule", {})38 return dict(cats=comp, threshold=float(rule.get("threshold", 1.0)), version=o.get("version"))394041def _init(path):42 global ONT43 ONT = load_ontology(path)444546def score_row(title, abstract, topic_ids, subfield_ids, arxiv_cats, dblp_streams):47 title = title or ""; abstract = abstract or ""; both = title + "\n" + abstract48 out = []49 for c in ONT["cats"]:50 st = min(1.0, sum(c["topics"].get(str(t), 0.0) for t in set(topic_ids or [])))51 ss = min(1.0, sum(c["subfields"].get(str(s), 0.0) for s in set(subfield_ids or [])))52 sa = min(1.0, sum(c["arxiv"].get(a, 0.0) for a in set(arxiv_cats or [])))53 sd = min(1.0, sum(c["dblp"].get(d, 0.0) for d in set(dblp_streams or [])))54 sk = 0.055 for rx, w, field in c["kws"]:56 txt = title if field == "title" else abstract if field == "abstract" else both57 if txt and rx.search(txt):58 sk += w59 if sk >= 1.0:60 sk = 1.0; break61 s = st + ss + sa + sd + sk62 if s >= ONT["threshold"] and c["exc"] and any(rx.search(both) for rx in c["exc"]):63 s = 0.064 if s >= ONT["threshold"]:65 out.append((c["id"], c["number"], c["parent"], round(s, 3), st, ss, sa, sd, sk))66 return out676869def process_chunk(args):70 path, out_dir = args71 name = os.path.basename(path).replace(".parquet", "")72 o1 = os.path.join(out_dir, "cats", name + ".parquet"); o2 = os.path.join(out_dir, "primary", name + ".parquet")73 if os.path.exists(o1) and os.path.exists(o2):74 return name, 0, 075 t = pq.read_table(path).to_pydict()76 cw, cc, cs, c1, c2, c3, c4, c5 = [], [], [], [], [], [], [], []77 pw, pc, pn, ps = [], [], [], []78 n = len(t["id"])79 for i in range(n):80 res = score_row(t["title"][i], t["abstract"][i], t["topic_ids"][i], t["subfield_ids"][i], t["arxiv_cats"][i], t["dblp_streams"][i])81 for (cid, num, parent, s, st, ss, sa, sd, sk) in res:82 cw.append(t["id"][i]); cc.append(cid); cs.append(s); c1.append(st); c2.append(ss); c3.append(sa); c4.append(sd); c5.append(sk)83 if res:84 parents = set().union(*[set(r[2]) for r in res])85 # primary: highest score; tie → child before parent (a category that is someone's parent loses), then lower number86 best = sorted(res, key=lambda r: (-r[3], 1 if r[0] in parents else 0, r[1]))[0]87 pw.append(t["id"][i]); pc.append(best[0]); pn.append(len(res)); ps.append(best[3])88 os.makedirs(os.path.dirname(o1), exist_ok=True); os.makedirs(os.path.dirname(o2), exist_ok=True)89 pq.write_table(pa.table({"work_id": pa.array(cw, pa.int64()), "category_id": cc, "score": pa.array(cs, pa.float32()), "s_topic": pa.array(c1, pa.float32()),90 "s_subfield": pa.array(c2, pa.float32()), "s_arxiv": pa.array(c3, pa.float32()), "s_dblp": pa.array(c4, pa.float32()), "s_kw": pa.array(c5, pa.float32())}), o1 + ".part", compression="zstd")91 os.replace(o1 + ".part", o1)92 pq.write_table(pa.table({"work_id": pa.array(pw, pa.int64()), "primary_category": pc, "n_categories": pa.array(pn, pa.int16()), "primary_score": pa.array(ps, pa.float32())}), o2 + ".part", compression="zstd")93 os.replace(o2 + ".part", o2)94 return name, n, len(pw)959697def prepare(root, db, out_dir, chunk_rows):98 con = duckdb.connect(db); con.execute("SET threads TO 24; SET memory_limit='60GB';")99 os.makedirs(os.path.join(out_dir, "input"), exist_ok=True)100 con.execute("""CREATE OR REPLACE TABLE cls_input AS101 SELECT w.id, w.publication_year AS year, w.title, a.abstract,102 (SELECT list(topic_id ORDER BY rank) FROM oa_work_topics t WHERE t.work_id = w.id) AS topic_ids,103 (SELECT list(subfield_id ORDER BY rank) FROM oa_work_topics t WHERE t.work_id = w.id) AS subfield_ids,104 (SELECT flatten(list(x.categories)) FROM link_arxiv_oa l JOIN arxiv_keys x USING(arxiv_id) WHERE l.oa_id = w.id) AS arxiv_cats,105 (SELECT list(DISTINCT d.stream) FROM link_dblp_oa l JOIN dblp_keys d USING(dblp_key) WHERE l.oa_id = w.id) AS dblp_streams106 FROM corpus_works w LEFT JOIN oa_abstracts a ON a.work_id = w.id""")107 n = con.execute("SELECT count(*) FROM cls_input").fetchone()[0]108 con.execute(f"""COPY (SELECT *, (row_number() OVER ()) // {chunk_rows} AS chunk FROM cls_input) TO '{out_dir}/input' (FORMAT PARQUET, PARTITION_BY (chunk), OVERWRITE_OR_IGNORE, COMPRESSION ZSTD)""")109 print("prepared", n, "works")110111112def main():113 ap = argparse.ArgumentParser(); ap.add_argument("--root", required=True); ap.add_argument("--ontology", required=True)114 ap.add_argument("--prepare", action="store_true"); ap.add_argument("--run", action="store_true"); ap.add_argument("--workers", type=int, default=20)115 ap.add_argument("--chunk-rows", type=int, default=100000); ap.add_argument("--limit", type=int, default=0)116 ap.add_argument("--shard", default=None, help="i/N: process only chunks with index %% N == i in this process (no pool; used by run_classify_shards.sh)")117 a = ap.parse_args(); out_dir = os.path.join(a.root, "data/processed/classification"); db = os.path.join(a.root, "data/processed/wp44.duckdb")118 if a.prepare:119 prepare(a.root, db, out_dir, a.chunk_rows)120 if a.run and a.shard:121 i, N = map(int, a.shard.split("/")); _init(a.ontology)122 files = sorted(glob.glob(os.path.join(out_dir, "input", "chunk=*", "*.parquet")), key=lambda f: int(os.path.basename(os.path.dirname(f)).split("=")[1]))123 t0 = time.time(); tot = cl = 0124 for k, f in enumerate(files):125 if k % N != i:126 continue127 name, n, c = process_chunk_named((f, out_dir, os.path.basename(os.path.dirname(f)))); tot += n; cl += c128 print(f"[{time.time()-t0:6.0f}s] shard {i}/{N} {name}: {n} works, {c} classified | {tot} / {cl}", flush=True)129 return130 if a.run:131 files = sorted(glob.glob(os.path.join(out_dir, "input", "chunk=*", "*.parquet")))132 files = [(f, out_dir) for f in files]133 if a.limit:134 files = files[: a.limit]135 # chunk files share the name data_0.parquet inside chunk=N dirs → rename key136 files = [(f, out_dir) for f, _ in files]137 t0 = time.time(); tot = 0; cl = 0138 with mp.Pool(a.workers, initializer=_init, initargs=(a.ontology,)) as pool:139 for name, n, k in pool.imap_unordered(process_chunk_named, [(f, out_dir, os.path.basename(os.path.dirname(f))) for f, _ in files]):140 tot += n; cl += k141 print(f"[{time.time()-t0:6.0f}s] {name}: {n} works, {k} classified | total {tot} / classified {cl}", flush=True)142 json.dump({"ontology": a.ontology, "version": load_ontology(a.ontology)["version"], "works": tot, "classified": cl, "when": dt.datetime.now(dt.UTC).isoformat(),143 "seconds": round(time.time() - t0)}, open(os.path.join(out_dir, "classification_report.json"), "w"), indent=1)144145146def process_chunk_named(args):147 path, out_dir, chunk_name = args148 name = chunk_name.replace("=", "_")149 o1 = os.path.join(out_dir, "cats", name + ".parquet"); o2 = os.path.join(out_dir, "primary", name + ".parquet")150 if os.path.exists(o1) and os.path.exists(o2):151 return name, 0, 0152 # reuse process_chunk with a symlinked name153 tmp = os.path.join(out_dir, "input_named"); os.makedirs(tmp, exist_ok=True)154 link = os.path.join(tmp, name + ".parquet")155 if not os.path.exists(link):156 os.symlink(path, link)157 return process_chunk((link, out_dir))158159160if __name__ == "__main__":161 main()162