pen-stack 0.1.0__py3-none-any.whl

pen_stack/web/server.py

@@ -0,0 +1,96 @@
+"""The PEN-STACK web gateway, one typed HTTP surface over the engine + the grounded chat.
+
+This is a *gateway, not a rewrite*. It mounts the AI-surface app (``/verify``, ``/safety``, ``/immune``,
+``/generate``, ``/predict``, ``/suggest``, ``/session``, ``/capabilities``, ``/scope`` + the atlas/crosslink
+routes) and adds the human-facing pieces: the grounded ``/chat`` (and a streaming ``/chat/stream``), CORS for
+the React frontend, and static serving of the built frontend when present. Every quantity still comes from the
+engine; the chat's LLM only narrates (see ``pen_stack.web.llm``).
+
+Run: ``uvicorn pen_stack.web.server:app --host 0.0.0.0 --port 8000`` (needs the ``server`` extra).
+"""
+from __future__ import annotations
+
+import json
+from pathlib import Path
+
+try:
+    from fastapi import FastAPI, HTTPException
+    from fastapi.middleware.cors import CORSMiddleware
+    from fastapi.responses import StreamingResponse
+except ImportError as e: # pragma: no cover - server extra optional
+    raise ImportError("FastAPI not installed: pip install 'pen-stack[server]'") from e
+
+from pen_stack import __version__
+from pen_stack.server.api import app as _engine_app # the typed engine surface (reused verbatim)
+
+
+def _require_message(req: dict) -> str:
+    msg = (req or {}).get("message")
+    if not isinstance(msg, str) or not msg.strip():
+        raise HTTPException(422, "field 'message' is required and must be a non-empty string")
+    return msg
+
+app = FastAPI(
+    title="PEN-STACK, Web Platform",
+    version=__version__,
+    description=("The human surface for genome writing: a grounded co-scientist chat + every engine feature, "
+                 "with provenance-first UX (confidence bands, provenance, scope ledger, safety badges). The LLM narrates "
+                 "and routes but never sources a number."),
+)
+
+# CORS, the Vite dev server (5173) and the bundled static frontend both talk to this gateway.
+app.add_middleware(
+    CORSMiddleware,
+    allow_origins=["*"], # self-hosted, single-tenant; tighten via PEN_STACK_CORS if deployed
+    allow_methods=["*"],
+    allow_headers=["*"],
+)
+
+# Mount the entire engine surface under the same app so the frontend has ONE base URL/one OpenAPI.
+app.mount("/api", _engine_app)
+
+
+@app.get("/health")
+def health() -> dict:
+    return {"status": "ok", "version": __version__, "surface": "web"}
+
+
+@app.post("/chat", tags=["co-scientist"])
+def chat_route(req: dict) -> dict:
+    """The grounded co-scientist: the LLM narrates over ENGINE tool outputs; every number comes from the tools
+    and the grounding guard strikes any the model invents. Body: {message, history?}. Returns
+    {reply, tool_results, grounded, backend}."""
+    from pen_stack.web.llm import grounded_reply
+
+    return grounded_reply(_require_message(req), history=req.get("history", []),
+                          allow_llm=bool(req.get("allow_llm", True)))
+
+
+@app.post("/chat/stream", tags=["co-scientist"])
+def chat_stream_route(req: dict) -> StreamingResponse:
+    """Same grounded reply, streamed as Server-Sent Events. The reply is computed grounded first (guard already
+    applied), then emitted word-by-word so the UI can render progressively; a final event carries the dossier."""
+    from pen_stack.web.llm import grounded_reply
+
+    result = grounded_reply(_require_message(req), history=req.get("history", []),
+                            allow_llm=bool(req.get("allow_llm", True)))
+
+    def _events():
+        for word in result["reply"].split(" "):
+            yield f"data: {json.dumps({'token': word + ' '})}\n\n"
+        yield ("event: done\ndata: " + json.dumps(
+            {"tool_results": result.get("tool_results"), "backend": result["backend"],
+             "grounded": result["grounded"], "mode": result.get("mode"),
+             "provenance": result.get("provenance"), "angles": result.get("angles"),
+             "facts": result.get("facts"), "sources": result.get("sources"),
+             "status": result.get("status")}, default=str) + "\n\n")
+
+    return StreamingResponse(_events(), media_type="text/event-stream")
+
+
+# Serve the built React frontend (web/dist) when present, one-command self-host serves UI + API together.
+_DIST = Path(__file__).resolve().parents[2] / "web" / "dist"
+if _DIST.exists(): # pragma: no cover - only when the frontend has been built
+    from fastapi.staticfiles import StaticFiles
+
+    app.mount("/", StaticFiles(directory=str(_DIST), html=True), name="frontend")

pen_stack/web/tools.py

@@ -0,0 +1,197 @@
+"""Engine tool-runner for the grounded co-scientist (PEN-STACK).
+
+The chat NEVER sources a number. This module is where the ENGINE computes everything: it parses a plain-language
+goal, runs the validated tools (verify -> legality + safety + calibrated confidence + immune profile; the scope
+matcher for known-unknowns), and returns a structured "dossier" of grounded facts. `extract_grounded_numbers`
+returns the allow-list of values the LLM is permitted to cite, anything else it emits is stripped by the
+grounding guard. No LLM is involved here; this is deterministic.
+"""
+from __future__ import annotations
+
+import re
+from typing import Any
+
+# tiny, transparent keyword maps (the engine grounds everything; this only routes a plain-language goal).
+_VEHICLES = {"aav": "AAV_single", "aav single": "AAV_single", "aav dual": "AAV_dual", "dual aav": "AAV_dual",
+             "lentivir": "lentivirus", "lnp": "lnp_mrna", "mrna": "lnp_mrna", "adenovir": "helper_dependent_adenovirus",
+             "hsv": "hsv_amplicon", "electroporat": "electroporation"}
+_INTENTS = {"safe harbour": "safe_harbour_insertion", "safe harbor": "safe_harbour_insertion",
+            "knock-in": "knock_in_with_disruption", "knock in": "knock_in_with_disruption",
+            "knockin": "knock_in_with_disruption", "durab": "high_durability_insertion",
+            "excis": "regulatory_element_excision", "regulator": "regulatory_element_excision",
+            "landing pad": "landing_pad_insertion", "repeat": "repeat_excision"}
+_CELLS = {"liver": "hepg2", "hepato": "hepg2", "hepg2": "hepg2", "hspc": "hspc", "stem cell": "h1_hesc",
+          "ipsc": "ipsc", "k562": "k562", "t cell": "cd8_t", "t-cell": "cd8_t", "car-t": "cd8_t", "pbmc": "pbmc"}
+_GENE_RE = re.compile(r"\b([A-Z][A-Z0-9]{1,7})\b") # crude gene-symbol token (AAVS1, TRAC, FIX, ...)
+_KB_RE = re.compile(r"(\d+(?:\.\d+)?)\s*kb", re.I)
+_BP_RE = re.compile(r"(\d{3,6})\s*bp", re.I)
+# uppercase tokens that look like a gene but are not one (jargon / vehicle / form abbreviations)
+_GENE_STOP = {"DNA", "RNA", "MRNA", "AAV", "LNP", "HSV", "CAR", "RNP", "PCR", "WT", "KO", "KI", "ITR",
+              "ORF", "UTR", "CDS", "GFP", "ID", "QC", "VG", "MOI", "HLA", "MHC", "CRISPR", "CAS", "TF"}
+# safe-harbour locus nicknames whose text collides with a vehicle keyword (AAVS1 ⊃ "aav"), stripped before
+# vehicle matching so the user's stated vehicle (e.g. lentivirus) is not overridden by the locus name.
+_SAFE_HARBOUR_RE = re.compile(r"\b(aavs1|h11|hipp11|rosa26)\b", re.I)
+
+
+def _first(text: str, table: dict, default):
+    low = text.lower()
+    for key, val in table.items():
+        if key in low:
+            return val
+    return default
+
+
+def _resolve_chrom(gene: str) -> str | None:
+    """The real chromosome for a gene symbol (or a safe-harbour nickname like AAVS1); None if not resolvable
+    offline. So a chat goal about ITGB2 carries chr21, not a hardcoded default. Atlas-gated, never fabricates."""
+    try:
+        from pen_stack.planner.optimize import gene_region, resolve_gene
+        reg = gene_region(resolve_gene(gene))
+        return reg[0] if reg else None
+    except Exception: # noqa: BLE001 - data/atlas absent (CI/offline) -> caller falls back to the default
+        return None
+
+
+def parse_goal(message: str) -> dict:
+    """Best-effort parse of a plain-language goal into a Design/Goal dict. The engine grounds everything; this
+    just picks a starting point (with sensible defaults) so the tools can run."""
+    cargo = 3000
+    if (m := _KB_RE.search(message)):
+        cargo = int(float(m.group(1)) * 1000)
+    elif (m := _BP_RE.search(message)):
+        cargo = int(m.group(1))
+    genes = [g for g in _GENE_RE.findall(message) if g not in _GENE_STOP]
+    gene = genes[0] if genes else "AAVS1"
+    # Vehicle matching is substring-based (so "lentivir" catches "lentiviral"); but the safe-harbour nickname
+    # "AAVS1" contains "aav", which would wrongly match the AAV vehicle even when the user said lentivirus/LNP.
+    # Strip those nicknames from the vehicle-search text first so the stated vehicle wins.
+    veh_text = _SAFE_HARBOUR_RE.sub(" ", message.lower())
+    return {"write_type": "insertion", "gene": gene, "chrom": _resolve_chrom(gene) or "chr19",
+            "edit_intent": _first(message, _INTENTS, "safe_harbour_insertion"),
+            "delivery_vehicle": _first(veh_text, _VEHICLES, "AAV_single"), "cargo_bp": cargo,
+            "cell_type": _first(message, _CELLS, "k562"),
+            # the user's plain-language goal IS the cargo-function description the Guardian must screen, so a
+            # message like "express a ricin toxin" is biosecurity-screened, not silently passed as benign.
+            "cargo_function": message.strip()}
+
+
+# the chat's free-text intents map onto the planner's EditIntent enum (landing-pad/regulatory have nearest valid).
+_INTENT_TO_ENUM = {"safe_harbour_insertion": "safe_harbour_insertion",
+                   "high_durability_insertion": "high_durability_insertion",
+                   "knock_in_with_disruption": "knock_in_with_disruption",
+                   "regulatory_element_excision": "regulatory_excision",
+                   "repeat_excision": "repeat_excision",
+                   "landing_pad_insertion": "safe_harbour_insertion"}
+
+# plain-language reading of each immune axis (0-1, higher = safer), so a value is self-explanatory in the reply.
+_AXIS_DIRECTION = {
+    "genotoxicity": "higher = safer (less integration-site oncogene risk; 1.0 = episomal / non-integrating)",
+    "cd8_epitope": "higher = fewer strong CD8/MHC-I capsid epitopes (1.0 = non-viral / none)",
+    "innate": "higher = less innate (CpG/TLR9, RIG-I) sensing of the delivered cargo",
+    "preexisting_nab": "higher = fewer patients excluded by pre-existing neutralizing antibodies to the vector",
+    "anti_peg": "higher = lower pre-existing anti-PEG barrier to re-dosing (PEGylated vehicles only)",
+}
+
+
+def _band_word(v: float) -> str:
+    return "favourable" if v >= 0.7 else ("moderate" if v >= 0.4 else "a concern")
+
+
+def axis_meaning(name: str, value, validation: str | None) -> str:
+    """A self-explanatory, plain-language reading of one immune axis: what the number means + the proxy caveat."""
+    if value is None:
+        return "out of scope for this design, not predicted (no applicable mechanism)."
+    direction = _AXIS_DIRECTION.get(name, "0-1, higher = safer")
+    s = f"{float(value):.2f} on a 0-1 scale, {_band_word(float(value))}; {direction}."
+    if validation and ("proxy" in validation.lower() or "not outcome-validated" in validation.lower()):
+        s += (" This is a mechanistically/population-computed PROXY, it is not validated against a measured "
+              "clinical outcome, so read it as a directional estimate, not a guaranteed result.")
+    return s
+
+
+def _run_planner(design: dict) -> dict[str, Any]:
+    """The ACTUAL writer/site recommendation for the goal (atlas-gated). This is what makes a 'which writer can
+    integrate N kb in GENE' question real: a named writer family, the top site, cargo-capacity fit, delivery,
+    all engine-computed. Computed when the gene isn't in the atlas or the atlas isn't mounted; NEVER fabricates."""
+    try:
+        from pen_stack.planner.optimize import EditIntent
+        from pen_stack.planner.pipeline import plan_write
+    except Exception: # noqa: BLE001
+        return {"available": False, "why": "planner unavailable in this environment"}
+    intent = _INTENT_TO_ENUM.get(design.get("edit_intent"), "safe_harbour_insertion")
+    try:
+        plans = plan_write(design["gene"], EditIntent(intent), int(design["cargo_bp"]),
+                           design.get("cell_type", "k562"), k=3)
+    except FileNotFoundError:
+        return {"available": False, "why": "writability atlas not mounted (planner runs on the live app only)"}
+    except Exception as e: # noqa: BLE001 - never let a planner error break the chat dossier
+        return {"available": False, "why": f"planner error: {type(e).__name__}"}
+    if not plans:
+        return {"available": True, "found": False, "gene": design["gene"],
+                "why": (f"no writable plan for '{design['gene']}', it is not in the writability atlas. Check it "
+                        "is an HGNC gene symbol or a known safe-harbour nickname (AAVS1, H11/HIPP11).")}
+    top = plans[0]
+    cargo = top.get("cargo") or {}
+    return {"available": True, "found": True, "n_plans": len(plans),
+            "recommended_writer": top.get("writer"), "site": top.get("site"),
+            "safety": top.get("safety"), "durability": top.get("durability"), "score": top.get("score"),
+            "writer_activity": top.get("writer_activity"), "reachability_tier": top.get("reachability_tier"),
+            "cargo_capacity_bp": cargo.get("cargo_capacity_bp"), "assembled_bp": cargo.get("assembled_bp"),
+            "cargo_fits_single_vector": cargo.get("size_ok"),
+            "delivery": (top.get("delivery") or {}).get("delivery"),
+            # distinct OTHER writer families (the top-k can repeat one family across sites, don't list it twice)
+            "alternative_writers": sorted({p.get("writer") for p in plans[1:]
+                                           if p.get("writer") and p.get("writer") != top.get("writer")})}
+
+
+def run_tools(message: str, history: list | None = None) -> dict[str, Any]:
+    """Run the validated engine over a plain-language message and return a grounded dossier. EVERY number here
+    is computed by the engine (verify / planner / scope), no fabrication, no LLM."""
+    from pen_stack.agent.scope import match_scope
+    from pen_stack.verify import verify
+
+    design = parse_goal(message)
+    v = verify(dict(design), question=message)
+    imm = v.immune_profile or {}
+    axes = {k: {"value": a.get("value"), "uncertainty": a.get("uncertainty"),
+                "validation": a.get("validation"), "in_scope": a.get("in_scope"),
+                "meaning": axis_meaning(k, a.get("value"), a.get("validation"))}
+            for k, a in (imm.get("axes") or {}).items()}
+    oos = match_scope(message) # is the QUESTION out of scope (a known-unknown)?
+    return {
+        "parsed_design": design,
+        "plan": _run_planner(design), # the actual writer/site recommendation (atlas-gated)
+        "verdict": {"legal": v.legal, "confidence": v.confidence, "interval": v.interval,
+                    "epistemic_status": v.epistemic_status,
+                    "violations": [x.get("rule_id") for x in v.violations]},
+        "safety": {"decision": (v.safety.decision if v.safety else None),
+                   "reason": (v.safety.reason if v.safety else None)},
+        "immune_profile": {"axes": axes, "collapsed_score": imm.get("collapsed_score"),
+                           "known_unknowns": imm.get("known_unknowns")},
+        "scope": ({"out_of_scope": True, "id": oos["id"], "title": oos["title"], "why": oos.get("deferral")}
+                  if oos else {"out_of_scope": False}),
+        "disclaimer": "Decision-support only; not a clinical directive. Every number is tool-sourced.",
+    }
+
+
+_NUM_RE = re.compile(r"-?\d+(?:\.\d+)?")
+
+
+def extract_grounded_numbers(tool_results: dict) -> set[str]:
+    """The allow-list: every numeric string that appears in the engine's tool results. The grounding guard
+    permits the LLM to cite ONLY these; any other number it emits is stripped."""
+    import json
+    text = json.dumps(tool_results, default=str)
+    grounded = set(_NUM_RE.findall(text))
+    # also allow the common normalised forms (e.g. 0.5 / .5 / 50%) of each grounded value
+    extra = set()
+    for n in list(grounded):
+        try:
+            f = float(n)
+        except ValueError:
+            continue
+        extra.add(str(int(f)) if f.is_integer() else str(f))
+        extra.add(f"{f:.2f}")
+        if 0 <= f <= 1:
+            extra.add(str(round(f * 100))) # percent form of a [0,1] score
+    return grounded | extra

pen_stack/wgenome/__init__.py

@@ -0,0 +1 @@
+"""pen_stack.wgenome - see the PEN-STACK program doc."""

pen_stack/wgenome/chromatin_seq.py

@@ -0,0 +1,83 @@
+"""Sequence-derived chromatin tracks: map AlphaGenome predictions onto the measured-atlas schema
+and recompute writability/safety/durability from predicted tracks.
+
+Two details:
+  * Unit handling. AlphaGenome track outputs are in the model's own units, not the measured ENCODE scale the
+    safety/durability models were trained on. Per-track agreement is therefore reported with Spearman (rank,
+    unit-free) alongside Pearson. For the *score-level* recompute we quantile-map each predicted track onto
+    the measured track's marginal (a standard rank-preserving calibration), so the recomputed scores test
+    whether AlphaGenome's RANKING of the epigenome recovers the measured-track scores - not a unit accident.
+  * Coverage. AlphaGenome predicts H3K9me3 for HepG2 but NOT K562; missing marks come back NaN and are
+    excluded from per-track correlation and passed as NaN to the (NaN-native) durability model.
+"""
+from __future__ import annotations
+
+from pathlib import Path
+
+import numpy as np
+import pandas as pd
+
+from pen_stack.wgenome.providers import AlphaGenomeProvider, _HISTONES, TRACK_NAMES
+
+_ROOT = Path(__file__).resolve().parents[2]
+_P1_FEAT = _ROOT.parent / "phase_1" / "features"
+_P1_OUT = _ROOT.parent / "phase_1" / "out"
+
+
+def predicted_tracks_frame(ct: str, bins: pd.DataFrame, provider: AlphaGenomeProvider | None = None,
+                           offline: bool = False) -> pd.DataFrame:
+    """Predicted 7-track values for the given (chrom, bin) rows. Cached per bin in the provider."""
+    provider = provider or AlphaGenomeProvider(assembly="hg38")
+    rows = []
+    for r in bins.itertuples():
+        rec = provider.tracks(r.chrom, int(r.bin), ct, offline=offline)
+        if rec.get("available"):
+            rows.append({"chrom": r.chrom, "bin": int(r.bin),
+                         **{t: rec.get(t, np.nan) for t in TRACK_NAMES}})
+    return pd.DataFrame(rows)
+
+
+def quantile_map(pred: pd.Series, measured: pd.Series) -> pd.Series:
+    """Map `pred` onto `measured`'s marginal by matching ranks (rank-preserving calibration)."""
+    pred = pred.astype(float)
+    if pred.notna().sum() < 2 or measured.notna().sum() < 2:
+        return pred
+    ranks = pred.rank(pct=True, na_option="keep")
+    q = np.nanpercentile(measured.to_numpy(dtype=float), np.clip(ranks.to_numpy() * 100, 0, 100))
+    return pd.Series(q, index=pred.index)
+
+
+def _load_models(ct: str):
+    from pen_stack.wgenome.writability import load_pickle
+    safety = load_pickle(str(_P1_OUT / f"safety_{ct}.pkl"))
+    dur = load_pickle(str(_P1_OUT / "durability.pkl"))
+    return safety, dur
+
+
+def recompute_scores(matrix: pd.DataFrame, ct: str) -> pd.DataFrame:
+    """Apply the trained safety + durability models to a feature matrix; return writability components."""
+    from pen_stack.wgenome.writability import build_writability
+    safety, dur = _load_models(ct)
+    return build_writability(matrix, safety, dur)
+
+
+def build_predicted_matrix(measured_matrix: pd.DataFrame, predicted: pd.DataFrame, ct: str) -> pd.DataFrame:
+    """Substitute quantile-mapped predicted tracks into a copy of the measured feature matrix.
+
+    Distance/integration features are genomic (not predicted) and are kept as-is; only the chromatin tracks
+    (atac/dnase/5 histones -> accessibility + marks) are replaced, then `accessibility` is rederived.
+    """
+    from pen_stack.wgenome.features import add_accessibility
+    m = measured_matrix.merge(predicted, on=["chrom", "bin"], how="inner", suffixes=("", "_pred"))
+    for t in TRACK_NAMES:
+        pc = f"{t}_pred"
+        if pc in m.columns and t in m.columns:
+            m[t] = quantile_map(m[pc], m[t]) # map predicted onto this sample's measured marginal
+    m = m.drop(columns=[c for c in m.columns if c.endswith("_pred")])
+    m = m.drop(columns=["accessibility"], errors="ignore")
+    return add_accessibility(m)
+
+
+def histone_marks_for(ct: str) -> list[str]:
+    """Marks AlphaGenome actually predicts for this cell type (K562 lacks H3K9me3)."""
+    return [m for m in _HISTONES if not (ct.lower() == "k562" and m == "H3K9me3")]

pen_stack/wgenome/durability.py

@@ -0,0 +1,108 @@
+"""Durability layer - the conditional chromatin-context model.
+
+Learns ONE function: `local chromatin features -> (expression level, silenced/stable)` on TRIP
+integrations. The model never sees a coordinate, so it is cell-type-agnostic in function: to score a
+new cell type you supply its chromatin tracks. This is the layer no safe-harbour resource provides,
+and TRIP supervises exactly the writing-relevant quantity (position effect on an integrated cassette).
+"""
+from __future__ import annotations
+
+import os
+from pathlib import Path
+
+import lightgbm as lgb
+import numpy as np
+import pandas as pd
+import pyBigWig
+from scipy.stats import spearmanr
+from sklearn.metrics import roc_auc_score
+from sklearn.model_selection import GroupKFold
+
+# canonical chromatin feature names (must match across mouse training + human application)
+CHROMATIN = ["atac", "dnase", "H3K27ac", "H3K4me1", "H3K4me3", "H3K9me3", "H3K27me3"]
+
+
+def liftover_positions(df: pd.DataFrame, chain_file: str) -> pd.DataFrame:
+    """Lift (chrom,pos) with a UCSC chain (e.g. mm9->mm10). Drops positions that fail to lift."""
+    from pyliftover import LiftOver
+    lo = LiftOver(chain_file)
+    out = []
+    for _, r in df.iterrows():
+        c = lo.convert_coordinate(r["chrom"], int(r["pos"]))
+        if c:
+            row = r.to_dict()
+            row["chrom"], row["pos"] = c[0][0], c[0][1]
+            out.append(row)
+    return pd.DataFrame(out)
+
+
+def extract_chromatin_at(df: pd.DataFrame, panel: dict, raw_dir: str, download_fn,
+                         window: int = 2500) -> pd.DataFrame:
+    """Point-query each bigWig's mean signal in +/-window around each integration position.
+    Only the integration sites are queried (no genome-wide binning needed)."""
+    out = df.copy()
+    for name, rec in panel.items():
+        path = download_fn(rec["href"], os.path.join(raw_dir, f"mES_{name}_{rec['accession']}.bigWig"))
+        bw = pyBigWig.open(path)
+        chroms = set(bw.chroms().keys())
+        vals = []
+        for chrom, pos in zip(out["chrom"], out["pos"]):
+            key = chrom if chrom in chroms else chrom.replace("chr", "")
+            if key not in chroms:
+                vals.append(0.0)
+                continue
+            try:
+                v = bw.stats(key, max(0, pos - window), pos + window, type="mean")[0]
+            except (RuntimeError, IndexError):
+                v = None
+            vals.append(0.0 if v is None else float(v))
+        out[name] = vals
+        bw.close()
+        os.remove(path)
+        print(f" extracted {name} at {len(out)} sites", flush=True)
+    return out
+
+
+def train_durability(trip_df: pd.DataFrame, seed: int = 42) -> dict:
+    feats = [c for c in CHROMATIN if c in trip_df.columns]
+    df = trip_df.dropna(subset=feats + ["expression"]).copy()
+    X = df[feats].astype("float32").fillna(0.0)
+    y_expr = df["expression"].to_numpy()
+    y_sil = df["silenced"].astype(int).to_numpy()
+    groups = df["chrom"].astype("category").cat.codes.to_numpy()
+
+    gkf = GroupKFold(n_splits=min(5, len(np.unique(groups))))
+    oof_expr = np.zeros(len(df))
+    oof_sil = np.zeros(len(df))
+    for tr, te in gkf.split(X, y_expr, groups):
+        reg = lgb.LGBMRegressor(n_estimators=500, learning_rate=0.03, num_leaves=31,
+                                subsample=0.8, random_state=seed, n_jobs=-1, verbosity=-1).fit(X.iloc[tr], y_expr[tr])
+        oof_expr[te] = reg.predict(X.iloc[te])
+        clf = lgb.LGBMClassifier(n_estimators=500, learning_rate=0.03, num_leaves=31,
+                                 subsample=0.8, random_state=seed, n_jobs=-1, verbosity=-1).fit(X.iloc[tr], y_sil[tr])
+        oof_sil[te] = clf.predict_proba(X.iloc[te])[:, 1]
+
+    rho = float(spearmanr(oof_expr, y_expr).statistic)
+    auroc = float(roc_auc_score(y_sil, oof_sil))
+    # baseline: H3K9me3 (heterochromatin) alone as a silencing predictor, and LAD-like (low ATAC) for expression
+    base_sil = roc_auc_score(y_sil, df["H3K9me3"].fillna(0)) if "H3K9me3" in df else float("nan")
+    base_expr = spearmanr(df.get("atac", pd.Series(0, index=df.index)).fillna(0), y_expr).statistic
+
+    final_reg = lgb.LGBMRegressor(n_estimators=500, learning_rate=0.03, num_leaves=31,
+                                  random_state=seed, n_jobs=-1, verbosity=-1).fit(X, y_expr)
+    imp = dict(sorted(zip(feats, final_reg.feature_importances_.tolist()), key=lambda kv: kv[1], reverse=True))
+    return {
+        "n": int(len(df)), "features": feats,
+        "expr_spearman": rho, "expr_baseline_atac_spearman": float(base_expr),
+        "silenced_auroc": auroc, "silenced_baseline_h3k9me3_auroc": float(base_sil),
+        "feature_importance": imp, "reg": final_reg,
+        "clf": lgb.LGBMClassifier(n_estimators=500, learning_rate=0.03, num_leaves=31,
+                                  random_state=seed, n_jobs=-1, verbosity=-1).fit(X, y_sil),
+    }
+
+
+def save_models(res: dict, out_dir: str, tag: str = "durability") -> None:
+    import pickle
+    Path(out_dir).mkdir(parents=True, exist_ok=True)
+    with open(f"{out_dir}/{tag}.pkl", "wb") as fh:
+        pickle.dump({"reg": res["reg"], "clf": res["clf"], "features": res["features"]}, fh)

pen_stack/wgenome/export_tracks.py

@@ -0,0 +1,52 @@
+"""Export the Writable Genome atlas as genome-browser tracks.
+
+Writes per-cell-type BigWig tracks (writability, safety, p_durable) loadable in IGV/UCSC, plus a BED
+of the top-writable loci. The atlas parquet stays the queryable source; these are the shareable tracks.
+"""
+from __future__ import annotations
+
+from pathlib import Path
+
+import pandas as pd
+import pyBigWig
+
+from pen_stack.data.genome import MAIN_CHROMS, load_chrom_sizes
+
+BIN_BP = 1000
+TRACKS = ["writability", "safety", "p_durable"]
+
+
+def write_bigwig(df: pd.DataFrame, col: str, chrom_sizes: dict[str, int], out_bw: str) -> None:
+    bw = pyBigWig.open(out_bw, "w")
+    # header must be sorted; keep canonical chrom order with sizes
+    chroms = [(c, chrom_sizes[c]) for c in MAIN_CHROMS if c in chrom_sizes]
+    bw.addHeader(chroms)
+    for chrom, _ in chroms:
+        g = df[df["chrom"] == chrom].sort_values("bin")
+        if g.empty:
+            continue
+        starts = (g["bin"].to_numpy() * BIN_BP).astype("int64")
+        vals = g[col].astype("float64").fillna(0.0).to_numpy()
+        bw.addEntries(chrom, list(starts), values=list(vals), span=BIN_BP, step=BIN_BP)
+    bw.close()
+
+
+def export_atlas(atlas_parquet: str, chrom_sizes_tsv: str, out_dir: str, ct: str,
+                 top_n: int = 5000) -> dict:
+    df = pd.read_parquet(atlas_parquet)
+    sizes = load_chrom_sizes(chrom_sizes_tsv)
+    Path(out_dir).mkdir(parents=True, exist_ok=True)
+    written = {}
+    for col in TRACKS:
+        if col in df.columns:
+            out_bw = f"{out_dir}/atlas_{ct}_{col}.bw"
+            write_bigwig(df, col, sizes, out_bw)
+            written[col] = out_bw
+    # top-writable loci BED
+    top = df.nlargest(top_n, "writability")[["chrom", "bin", "writability"]].copy()
+    top["start"] = top["bin"] * BIN_BP
+    top["end"] = top["start"] + BIN_BP
+    bed = f"{out_dir}/atlas_{ct}_top{top_n}.bed"
+    top[["chrom", "start", "end", "writability"]].to_csv(bed, sep="\t", header=False, index=False)
+    written["top_bed"] = bed
+    return written

pen_stack/wgenome/features.py

@@ -0,0 +1,82 @@
+"""Assemble the per-cell-type training/scoring matrix.
+
+Joins the cell-type chromatin feature store + the (cell-type-agnostic) safety-annotation store
+(+ integration-outcome store when available) on (chrom, bin) into one matrix the safety and
+durability layers consume. Keeps feature provenance explicit.
+"""
+from __future__ import annotations
+
+from pathlib import Path
+
+import numpy as np
+import pandas as pd
+
+# Unified chromatin feature set: ONE accessibility feature (ATAC where present, else DNase) +
+# the 5 core histone marks. This makes every cell type share an IDENTICAL schema, so a cell type
+# that lacks a specific accessibility assay (e.g. CD34+ HSPC has DNase but no ATAC) is fully
+# specified rather than "partial" - ATAC and DNase are interchangeable open-chromatin assays.
+CHROMATIN_TRACKS = ["accessibility", "H3K27ac", "H3K4me1", "H3K4me3", "H3K9me3", "H3K27me3"]
+ACCESS_SOURCES = ["atac", "dnase"]
+SAFETY_DIST = ["dist_oncogene", "dist_tsg", "dist_essential", "dist_tss"]
+
+
+def add_accessibility(m: pd.DataFrame) -> pd.DataFrame:
+    """Derive the unified `accessibility` column: prefer ATAC, fall back to DNase."""
+    if "accessibility" not in m.columns:
+        if "atac" in m.columns:
+            m["accessibility"] = m["atac"]
+            if "dnase" in m.columns: # fill any ATAC gaps with DNase
+                m["accessibility"] = m["accessibility"].fillna(m["dnase"])
+        elif "dnase" in m.columns:
+            m["accessibility"] = m["dnase"]
+    return m
+
+
+def _log_dist(s: pd.Series) -> pd.Series:
+    # large/Inf for "no feature on chromosome" -> log1p of a capped distance; NaN -> max
+    v = s.fillna(s.max() if s.notna().any() else 1e8).clip(lower=0)
+    return np.log1p(v)
+
+
+def assemble_matrix(chromatin_parquet: str, safety_parquet: str,
+                    integration_parquet: str | None = None,
+                    out_parquet: str | None = None) -> pd.DataFrame:
+    chrom = pd.read_parquet(chromatin_parquet)
+    safe = pd.read_parquet(safety_parquet)
+    m = chrom.merge(safe, on=["chrom", "bin"], how="inner")
+    m = add_accessibility(m) # unify ATAC/DNase -> accessibility
+
+    # log-scaled distance features (raw kept too, for transparency)
+    for d in SAFETY_DIST:
+        if d in m.columns:
+            m[f"log_{d}"] = _log_dist(m[d])
+
+    if integration_parquet and Path(integration_parquet).exists():
+        integ = pd.read_parquet(integration_parquet)
+        m = m.merge(integ, on=["chrom", "bin"], how="left")
+        for c in [c for c in integ.columns if c not in ("chrom", "bin")]:
+            m[c] = m[c].fillna(0)
+
+    if out_parquet:
+        Path(out_parquet).parent.mkdir(parents=True, exist_ok=True)
+        m.to_parquet(out_parquet, index=False)
+    return m
+
+
+def resolve_integration(feat_dir: str, ct: str) -> str | None:
+    """Integration-feature parquet for a cell type: prefer the cell-type-specific MLV set
+    (LaFave K562/HepG2); fall back to the cell-type-agnostic VISDB retroviral-propensity track so a
+    cell type without its own integration assay (e.g. CD34+ HSPC) still gets an integration feature."""
+    ct_specific = Path(feat_dir) / f"integration_{ct}.parquet"
+    if ct_specific.exists():
+        return str(ct_specific)
+    fallback = Path(feat_dir) / "integration_density.parquet"
+    return str(fallback) if fallback.exists() else None
+
+
+def feature_columns(df: pd.DataFrame) -> list[str]:
+    """The model feature set: chromatin marks + log-distances + any integration features."""
+    feats = [c for c in CHROMATIN_TRACKS if c in df.columns]
+    feats += [c for c in df.columns if c.startswith("log_dist_")]
+    feats += [c for c in df.columns if c.startswith("integ_")]
+    return feats