Research-Stack/4-Infrastructure/shim/arxiv_equation_pattern_extractor.py
2026-05-11 22:18:31 -05:00

488 lines
18 KiB
Python

#!/usr/bin/env python3
"""Extract equation-pattern features from a bounded arXiv source sample.
The extractor downloads source packages from arXiv e-print URLs into an external
cache, reads TeX files, and emits equation hashes plus structural features. It
does not store equation text in the repo.
"""
from __future__ import annotations
import argparse
import gzip
import hashlib
import io
import json
import re
import tarfile
import time
import zipfile
from collections import Counter, defaultdict
from pathlib import Path
from typing import Any
from urllib.error import HTTPError, URLError
from urllib.request import Request, urlopen
REPO = Path(__file__).resolve().parents[2]
DEFAULT_ARCHIVE = Path("/home/allaun/Documents/ingest/kaggledataset.zip")
DEFAULT_CACHE = Path("/home/allaun/Documents/ingest/arxiv_sources")
OUT_DIR = REPO / "shared-data" / "data" / "math_research_databases"
EQUATION_JSONL = OUT_DIR / "arxiv_equation_pattern_records.jsonl"
PAPER_JSONL = OUT_DIR / "arxiv_equation_pattern_papers.jsonl"
COMMAND_CSV = OUT_DIR / "arxiv_equation_pattern_commands.csv"
SUMMARY = OUT_DIR / "arxiv_equation_pattern_summary.json"
RECEIPT = OUT_DIR / "arxiv_equation_pattern_receipt.json"
MATH_PREFIXES = (
"math",
"stat",
"cs.",
"nlin",
"q-bio",
"q-fin",
"quant-ph",
"hep-th",
"math-ph",
)
DISPLAY_ENV_NAMES = (
"equation",
"equation*",
"align",
"align*",
"gather",
"gather*",
"multline",
"multline*",
"eqnarray",
"eqnarray*",
"displaymath",
"flalign",
"flalign*",
"alignat",
"alignat*",
)
DISPLAY_ENV_RE = re.compile(
r"\\begin\{(" + "|".join(re.escape(name) for name in DISPLAY_ENV_NAMES) + r")\}(.*?)\\end\{\1\}",
re.DOTALL,
)
BRACKET_DISPLAY_RE = re.compile(r"\\\[(.*?)\\\]", re.DOTALL)
DOUBLE_DOLLAR_RE = re.compile(r"\$\$(.*?)\$\$", re.DOTALL)
INLINE_DOLLAR_RE = re.compile(r"(?<!\\)\$(?!\$)(.{1,500}?)(?<!\\)\$", re.DOTALL)
COMMAND_RE = re.compile(r"\\[A-Za-z]+|\\.")
OPERATOR_RE = re.compile(r"(?:<=|>=|!=|:=|->|=>|[=<>+\-*/^_&|])")
GREEK_COMMANDS = {
"\\alpha",
"\\beta",
"\\gamma",
"\\delta",
"\\epsilon",
"\\varepsilon",
"\\zeta",
"\\eta",
"\\theta",
"\\vartheta",
"\\iota",
"\\kappa",
"\\lambda",
"\\mu",
"\\nu",
"\\xi",
"\\pi",
"\\rho",
"\\sigma",
"\\tau",
"\\upsilon",
"\\phi",
"\\varphi",
"\\chi",
"\\psi",
"\\omega",
"\\Gamma",
"\\Delta",
"\\Theta",
"\\Lambda",
"\\Xi",
"\\Pi",
"\\Sigma",
"\\Phi",
"\\Psi",
"\\Omega",
}
FAMILY_COMMANDS = {
"fraction": {"\\frac", "\\dfrac", "\\tfrac"},
"integral": {"\\int", "\\iint", "\\iiint", "\\oint"},
"summation": {"\\sum", "\\prod", "\\coprod"},
"limit": {"\\lim", "\\sup", "\\inf", "\\max", "\\min"},
"set": {"\\in", "\\subset", "\\subseteq", "\\cup", "\\cap", "\\setminus", "\\emptyset"},
"arrow": {"\\to", "\\rightarrow", "\\leftarrow", "\\mapsto", "\\Rightarrow", "\\Longrightarrow"},
"inequality": {"\\le", "\\leq", "\\ge", "\\geq", "\\neq", "\\sim", "\\simeq", "\\approx"},
"font": {"\\mathbb", "\\mathcal", "\\mathbf", "\\mathrm", "\\mathfrak", "\\operatorname"},
"root": {"\\sqrt"},
"matrix": {"\\matrix", "\\pmatrix", "\\bmatrix", "\\begin"},
}
def stable_json(obj: Any) -> str:
return json.dumps(obj, sort_keys=True, separators=(",", ":"), ensure_ascii=True)
def sha256_bytes(data: bytes) -> str:
return hashlib.sha256(data).hexdigest()
def sha256_text(text: str) -> str:
return hashlib.sha256(text.encode("utf-8", errors="replace")).hexdigest()
def categories_of(record: dict[str, Any]) -> list[str]:
raw = record.get("categories")
if not isinstance(raw, str):
return []
return [part for part in raw.split() if part]
def is_math_bearing(categories: list[str]) -> bool:
return any(cat.startswith(MATH_PREFIXES) for cat in categories)
def text_field(record: dict[str, Any], key: str) -> str:
value = record.get(key)
return value if isinstance(value, str) else ""
def abstract_marker_score(record: dict[str, Any]) -> int:
text = "\n".join([text_field(record, "title"), text_field(record, "abstract")])
return len(COMMAND_RE.findall(text)) + len(INLINE_DOLLAR_RE.findall(text))
def iter_candidate_records(archive: Path, max_scan: int, max_papers: int, min_abstract_markers: int) -> list[dict[str, Any]]:
candidates: list[dict[str, Any]] = []
with zipfile.ZipFile(archive) as zf:
infos = zf.infolist()
json_members = [info for info in infos if info.filename.endswith(".json")]
selected_member = json_members[0] if json_members else infos[0]
with zf.open(selected_member) as handle:
for seen, raw_line in enumerate(handle, start=1):
if max_scan and seen > max_scan:
break
record = json.loads(raw_line)
cats = categories_of(record)
if not is_math_bearing(cats):
continue
score = abstract_marker_score(record)
if score < min_abstract_markers:
continue
candidates.append(
{
"id": record.get("id"),
"title": record.get("title"),
"categories": cats,
"primary_category": cats[0] if cats else "uncategorized",
"abstract_marker_score": score,
}
)
if len(candidates) >= max_papers:
break
return candidates
def records_by_id(archive: Path, wanted_ids: set[str], max_scan: int) -> list[dict[str, Any]]:
found: list[dict[str, Any]] = []
with zipfile.ZipFile(archive) as zf:
infos = zf.infolist()
json_members = [info for info in infos if info.filename.endswith(".json")]
selected_member = json_members[0] if json_members else infos[0]
with zf.open(selected_member) as handle:
for seen, raw_line in enumerate(handle, start=1):
if max_scan and seen > max_scan:
break
record = json.loads(raw_line)
arxiv_id = str(record.get("id", ""))
if arxiv_id not in wanted_ids:
continue
cats = categories_of(record)
found.append(
{
"id": arxiv_id,
"title": record.get("title"),
"categories": cats,
"primary_category": cats[0] if cats else "uncategorized",
"abstract_marker_score": abstract_marker_score(record),
"doi": record.get("doi"),
"versions": record.get("versions"),
}
)
if len(found) >= len(wanted_ids):
break
missing = wanted_ids - {str(row["id"]) for row in found}
for arxiv_id in sorted(missing):
found.append(
{
"id": arxiv_id,
"title": None,
"categories": [],
"primary_category": "uncategorized",
"abstract_marker_score": 0,
"metadata_missing": True,
}
)
return found
def fetch_source(arxiv_id: str, cache_dir: Path, delay_seconds: float) -> tuple[Path | None, str | None]:
cache_dir.mkdir(parents=True, exist_ok=True)
safe_id = arxiv_id.replace("/", "_")
cached = cache_dir / f"{safe_id}.src"
if cached.exists() and cached.stat().st_size > 0:
return cached, None
url = f"https://arxiv.org/e-print/{arxiv_id}"
req = Request(url, headers={"User-Agent": "ResearchStack equation-pattern extractor (local bounded study)"})
try:
with urlopen(req, timeout=60) as response:
data = response.read()
except (HTTPError, URLError, TimeoutError) as exc:
return None, repr(exc)
cached.write_bytes(data)
if delay_seconds > 0:
time.sleep(delay_seconds)
return cached, None
def strip_comments(tex: str) -> str:
lines = []
for line in tex.splitlines():
out = []
escaped = False
for ch in line:
if ch == "%" and not escaped:
break
out.append(ch)
escaped = ch == "\\" and not escaped
lines.append("".join(out))
return "\n".join(lines)
def read_tex_sources(path: Path, max_tex_bytes: int) -> tuple[list[tuple[str, str]], str]:
data = path.read_bytes()
tex_files: list[tuple[str, str]] = []
mode = "plain"
try:
with tarfile.open(fileobj=io.BytesIO(data), mode="r:*") as tf:
mode = "tar"
for member in tf.getmembers():
if not member.isfile() or not member.name.lower().endswith((".tex", ".ltx")):
continue
if member.size > max_tex_bytes:
continue
extracted = tf.extractfile(member)
if extracted is None:
continue
text = extracted.read(max_tex_bytes).decode("utf-8", errors="replace")
tex_files.append((member.name, text))
return tex_files, mode
except tarfile.TarError:
pass
try:
text = gzip.decompress(data).decode("utf-8", errors="replace")
mode = "gzip_plain"
except OSError:
text = data.decode("utf-8", errors="replace")
if "\\documentclass" in text or "\\begin" in text:
tex_files.append((path.name, text[:max_tex_bytes]))
return tex_files, mode
def extract_equations(tex: str) -> list[tuple[str, str]]:
clean = strip_comments(tex)
equations: list[tuple[str, str]] = []
for match in DISPLAY_ENV_RE.finditer(clean):
equations.append((match.group(1), match.group(2)))
for match in BRACKET_DISPLAY_RE.finditer(clean):
equations.append(("bracket_display", match.group(1)))
for match in DOUBLE_DOLLAR_RE.finditer(clean):
equations.append(("double_dollar", match.group(1)))
for match in INLINE_DOLLAR_RE.finditer(clean):
equations.append(("inline_dollar", match.group(1)))
return equations
def equation_features(equation: str) -> dict[str, Any]:
normalized = " ".join(equation.split())
commands = COMMAND_RE.findall(normalized)
operators = OPERATOR_RE.findall(normalized)
command_counts = Counter(commands)
families = {
family: sum(command_counts.get(cmd, 0) for cmd in cmds)
for family, cmds in FAMILY_COMMANDS.items()
}
greek_count = sum(command_counts.get(cmd, 0) for cmd in GREEK_COMMANDS)
signature_parts = []
for family, count in sorted(families.items()):
if count:
signature_parts.append(f"{family}:{count}")
if greek_count:
signature_parts.append(f"greek:{greek_count}")
signature_parts.append(f"cmds:{len(commands)}")
signature_parts.append(f"ops:{len(operators)}")
signature = "|".join(signature_parts)
return {
"equation_sha256": sha256_text(normalized),
"char_length": len(normalized),
"command_count": len(commands),
"operator_count": len(operators),
"greek_command_count": greek_count,
"family_counts": families,
"top_commands": command_counts.most_common(20),
"signature": signature,
"signature_hash": sha256_text(signature),
}
def csv_escape(value: Any) -> str:
text = str(value).replace('"', '""')
return f'"{text}"'
def main() -> None:
parser = argparse.ArgumentParser()
parser.add_argument("--archive", default=str(DEFAULT_ARCHIVE))
parser.add_argument("--cache-dir", default=str(DEFAULT_CACHE))
parser.add_argument("--max-scan", type=int, default=5000)
parser.add_argument("--max-papers", type=int, default=12)
parser.add_argument("--min-abstract-markers", type=int, default=2)
parser.add_argument("--delay-seconds", type=float, default=3.0)
parser.add_argument("--max-tex-bytes", type=int, default=5_000_000)
parser.add_argument("--max-equations-per-paper", type=int, default=1000)
parser.add_argument("--ids", default="")
args = parser.parse_args()
archive = Path(args.archive).expanduser().resolve()
cache_dir = Path(args.cache_dir).expanduser().resolve()
OUT_DIR.mkdir(parents=True, exist_ok=True)
wanted_ids = {part.strip() for part in args.ids.split(",") if part.strip()}
if wanted_ids:
candidates = records_by_id(archive=archive, wanted_ids=wanted_ids, max_scan=args.max_scan)
else:
candidates = iter_candidate_records(
archive=archive,
max_scan=args.max_scan,
max_papers=args.max_papers,
min_abstract_markers=args.min_abstract_markers,
)
equation_records: list[dict[str, Any]] = []
paper_records: list[dict[str, Any]] = []
command_totals: Counter[str] = Counter()
signature_totals: Counter[str] = Counter()
family_totals: Counter[str] = Counter()
env_totals: Counter[str] = Counter()
fetch_status: Counter[str] = Counter()
for candidate in candidates:
arxiv_id = str(candidate["id"])
source_path, error = fetch_source(arxiv_id, cache_dir, args.delay_seconds)
if source_path is None:
fetch_status["fetch_error"] += 1
paper_records.append({**candidate, "decision": "HOLD", "error": error})
continue
source_sha = sha256_bytes(source_path.read_bytes())
tex_files, source_mode = read_tex_sources(source_path, args.max_tex_bytes)
paper_eq_count = 0
paper_command_count = 0
for tex_name, tex in tex_files:
for env, equation in extract_equations(tex)[: args.max_equations_per_paper]:
features = equation_features(equation)
record = {
"schema": "arxiv_equation_pattern_record_v1",
"arxiv_id": arxiv_id,
"primary_category": candidate["primary_category"],
"categories": candidate["categories"],
"source_sha256": source_sha,
"tex_file": tex_name,
"environment": env,
**features,
}
record["packet_hash"] = sha256_text(stable_json(record))
equation_records.append(record)
paper_eq_count += 1
paper_command_count += features["command_count"]
env_totals[env] += 1
signature_totals[features["signature_hash"]] += 1
for family, count in features["family_counts"].items():
family_totals[family] += count
for command, count in features["top_commands"]:
command_totals[command] += count
fetch_status["fetched"] += 1
paper = {
"schema": "arxiv_equation_pattern_paper_v1",
**candidate,
"source_cache_path": str(source_path),
"source_sha256": source_sha,
"source_mode": source_mode,
"tex_file_count": len(tex_files),
"equation_count": paper_eq_count,
"equation_command_count": paper_command_count,
"decision": "HOLD",
}
paper["packet_hash"] = sha256_text(stable_json(paper))
paper_records.append(paper)
EQUATION_JSONL.write_text("\n".join(stable_json(r) for r in equation_records) + "\n", encoding="utf-8")
PAPER_JSONL.write_text("\n".join(stable_json(r) for r in paper_records) + "\n", encoding="utf-8")
command_lines = ["command,count"]
for command, count in command_totals.most_common(300):
command_lines.append(f"{csv_escape(command)},{count}")
COMMAND_CSV.write_text("\n".join(command_lines) + "\n", encoding="utf-8")
summary = {
"schema": "arxiv_equation_pattern_summary_v1",
"archive_path": str(archive),
"source_cache_dir": str(cache_dir),
"candidate_count": len(candidates),
"paper_count": len(paper_records),
"fetch_status": dict(fetch_status),
"equation_count": len(equation_records),
"unique_equation_hashes": len({r["equation_sha256"] for r in equation_records}),
"unique_signature_hashes": len(signature_totals),
"top_environments": env_totals.most_common(30),
"top_commands": command_totals.most_common(80),
"family_totals": dict(family_totals),
"top_signature_hashes": signature_totals.most_common(40),
"equation_records": str(EQUATION_JSONL.relative_to(REPO)),
"paper_records": str(PAPER_JSONL.relative_to(REPO)),
"command_csv": str(COMMAND_CSV.relative_to(REPO)),
"decision": "HOLD",
"claim_boundary": (
"Bounded arXiv source sample. Repo stores equation hashes and structural "
"features only, not equation text. This is compression-surface evidence, "
"not theorem verification or full-corpus coverage."
),
}
summary["packet_hash"] = sha256_text(stable_json(summary))
SUMMARY.write_text(json.dumps(summary, indent=2, sort_keys=True) + "\n", encoding="utf-8")
receipt = {
"schema": "arxiv_equation_pattern_receipt_v1",
"summary": str(SUMMARY.relative_to(REPO)),
"candidate_count": summary["candidate_count"],
"paper_count": summary["paper_count"],
"fetch_status": summary["fetch_status"],
"equation_count": summary["equation_count"],
"unique_equation_hashes": summary["unique_equation_hashes"],
"unique_signature_hashes": summary["unique_signature_hashes"],
"top_environments": summary["top_environments"][:10],
"top_commands": summary["top_commands"][:20],
"family_totals": summary["family_totals"],
"packet_hash": summary["packet_hash"],
"decision": summary["decision"],
"claim_boundary": summary["claim_boundary"],
}
receipt["receipt_hash"] = sha256_text(stable_json(receipt))
RECEIPT.write_text(json.dumps(receipt, indent=2, sort_keys=True) + "\n", encoding="utf-8")
print(json.dumps(receipt, indent=2, sort_keys=True))
if __name__ == "__main__":
main()