spb/uqo-research-db

1"""Shared paths, plotting style, and helpers for the SPID research project."""2from __future__ import annotations34import json5from pathlib import Path67import matplotlib8matplotlib.use("Agg")9import matplotlib.pyplot as plt10import matplotlib as mpl1112ROOT = Path(__file__).resolve().parent.parent          # research/13DB = ROOT.parent / "data" / "db" / "spid.duckdb"        # project-db/data/db/spid.duckdb14DATA = ROOT / "data"15FIG = ROOT / "figures"16TAB = ROOT / "tables"17RES = ROOT / "results"18for d in (DATA, FIG, TAB, RES):19    d.mkdir(parents=True, exist_ok=True)2021# ---- publication plotting style -------------------------------------------22PALETTE = ["#1b3a5b", "#c44e52", "#2e8b57", "#dd8452", "#8172b3",23           "#937860", "#da8bc3", "#8c8c8c", "#ccb974", "#64b5cd"]2425def set_style() -> None:26    mpl.rcParams.update({27        "figure.dpi": 130,28        "savefig.dpi": 200,29        "font.size": 11,30        "font.family": "serif",31        "axes.titlesize": 12,32        "axes.titleweight": "bold",33        "axes.labelsize": 11,34        "axes.spines.top": False,35        "axes.spines.right": False,36        "axes.grid": True,37        "grid.alpha": 0.25,38        "grid.linewidth": 0.6,39        "legend.frameon": False,40        "axes.prop_cycle": mpl.cycler(color=PALETTE),41    })4243def savefig(fig, name: str) -> Path:44    p = FIG / name45    fig.tight_layout()46    fig.savefig(p, bbox_inches="tight")47    fig.savefig(p.with_suffix(".png"), bbox_inches="tight")48    plt.close(fig)49    print(f"  [fig] {p.name}")50    return p5152def con(read_only: bool = True):53    import duckdb54    return duckdb.connect(str(DB), read_only=read_only)5556def save_results(name: str, obj: dict) -> None:57    p = RES / f"{name}.json"58    p.write_text(json.dumps(obj, indent=2, default=str))59    print(f"  [res] {p.name}")6061def write_table(name: str, latex: str) -> None:62    p = TAB / f"{name}.tex"63    p.write_text(latex)64    print(f"  [tab] {p.name}")6566def esc(s) -> str:67    """Escape a string for LaTeX."""68    s = str(s)69    for a, b in [("\\", r"\textbackslash{}"), ("&", r"\&"), ("%", r"\%"),70                 ("$", r"\$"), ("#", r"\#"), ("_", r"\_"), ("{", r"\{"),71                 ("}", r"\}"), ("~", r"\textasciitilde{}"), ("^", r"\textasciicircum{}")]:72        s = s.replace(a, b)73    return s74