scripts/augment_gamma_channel.py

"""Augment the empirical corpus with a non-redundant physics channel: the UNIFAC
activity-coefficient (gamma) interaction term.

Motivation (E2 finding): the existing 2-channel physics signal (x_liquid,
log10 OAV) is largely REDUNDANT with the model's existing molecular-weight
input (volatility is ~76% MW-recoverable). The genuinely non-redundant physics
is the *mixture interaction*: the activity coefficient gamma_i(t), which depends
on the whole liquid composition and is emphatically not contained in any
single-molecule embedding.

gamma is recovered without re-running UNIFAC. The engine's headspace obeys
non-ideal Raoult:  P_i = gamma_i * x_i * P_sat_i,  y_i = P_i / sum_j(P_j).
Hence  gamma_i ∝ y_i / (x_i * P_sat_i).  The shared normalizer cancels in any
ratio, and the FiLM block standardizes channels, so we store log10 of the
proportional gamma (a sufficient, monotone encoding of the interaction term).

Output: a new physics tensor per formula of shape (T, S, 3) with channels
[x_liquid, log10 OAV, log10 gamma_prop], written alongside the original record
so train.py can opt into state_dim=3.
"""
from __future__ import annotations

import json
import sqlite3
from pathlib import Path

import numpy as np

ROOT = Path(__file__).resolve().parent.parent
DB = ROOT / "src" / "pino" / "registry.db"


def load_psat() -> dict[str, float]:
    con = sqlite3.connect(DB)
    con.row_factory = sqlite3.Row
    out = {}
    for r in con.execute("SELECT cas, vapor_pressure_pa FROM aroma_chemicals WHERE vapor_pressure_pa > 0"):
        out[r["cas"]] = float(r["vapor_pressure_pa"])
    con.close()
    return out


def gamma_prop_channel(step: dict, cass: list[str], psat: dict[str, float]) -> np.ndarray:
    """Return (S,) log10 proportional-gamma for one timestep.

    gamma_i ∝ y_i / (x_i * P_sat_i). Components with missing P_sat or zero
    x_liquid get gamma=1 (log10=0, the ideal limit).
    """
    xl = step.get("x_liquid", {})
    yg = step.get("y_gas", {})
    logg = np.zeros(len(cass), dtype=np.float32)
    for j, cas in enumerate(cass):
        x = float(xl.get(cas, 0.0))
        y = float(yg.get(cas, 0.0))
        p = psat.get(cas.replace("NATURAL:", ""))
        if p is None or x <= 1e-9 or y <= 0.0:
            logg[j] = 0.0  # ideal / no signal
        else:
            logg[j] = np.log10(max(y / (x * p), 1e-12))
    # center so the channel carries only relative (mixture-interaction) info
    if len(cass) > 0:
        logg = logg - logg.mean()
    return logg


def augment_record(rec: dict, psat: dict[str, float]) -> dict:
    formula = rec.get("formula", [])
    traj = rec.get("trajectory", [])
    cass = [str(c.get("cas", f"ing_{j}")) for j, c in enumerate(formula)]
    # build (T, S, 1) gamma channel aligned with the existing physics tensor
    gam = np.stack([gamma_prop_channel(st, cass, psat) for st in traj], axis=0)  # (T, S)
    rec["physics_gamma"] = gam.tolist()
    return rec


def main(in_path: str, out_path: str, max_records: int | None = None):
    psat = load_psat()
    print(f"registry P_sat for {len(psat)} CAS")
    n = 0
    with open(in_path) as fin, open(out_path, "w") as fout:
        for line in fin:
            if not line.strip():
                continue
            rec = json.loads(line)
            if "formula" in rec and "trajectory" in rec:
                rec = augment_record(rec, psat)
            fout.write(json.dumps(rec) + "\n")
            n += 1
            if max_records and n >= max_records:
                break
            if n % 1000 == 0:
                print(f"  {n} records...")
    print(f"wrote {n} augmented records -> {out_path}")


if __name__ == "__main__":
    import argparse
    ap = argparse.ArgumentParser()
    ap.add_argument("--in", dest="inp", default=str(ROOT / "data/empirical_dataset_v9_plus_wisemoor.jsonl"))
    ap.add_argument("--out", dest="out", default=str(ROOT / "data/empirical_dataset_v9_gamma.jsonl"))
    ap.add_argument("--max", type=int, default=None)
    a = ap.parse_args()
    main(a.inp, a.out, a.max)