"""Scoring a structure under a fitted model: energy, ``log Z``, likelihood, contact probabilities.
Every quantity here is per structure and needs the model's parameters plus that structure's own
available-pair set. Nothing is fitted.
"""
from __future__ import annotations
import hashlib
import os
from concurrent.futures import ProcessPoolExecutor
import numpy as np
import polars as pl
from .kernel import bucket_edges, colour, coupling_matrix, edges
from .model import PottsModel, kernel_names
from .sample import (ais_log_z, count_free_energy, delta_f_empty, delta_f_threshold, energy,
factorised_log_z, gibbs, mu_star)
from .sites import available_pairs, eta, site_codes
def _rng_for(seed: int, pdb_id) -> np.random.Generator:
"""A generator determined by ``(seed, pdb_id)`` alone — never by table position.
Every sampler below used to share one generator across its loop over structures, so a
structure's ``log Z`` depended on **how many structures preceded it in the table**: the same
PDB scored on its own, in a reordered frame, or in a subset came back with a different value.
Deriving the stream from the structure's own identifier makes every sampled quantity a function
of the structure, which is what makes a per-structure score reproducible. ``blake2b`` rather
than :func:`hash`, whose seed is randomised per interpreter.
"""
h = int.from_bytes(hashlib.blake2b(str(pdb_id).encode(), digest_size=8).digest(), "big")
return np.random.default_rng(np.random.SeedSequence([int(seed), h]))
def _prepare(sites: pl.DataFrame, model: PottsModel):
"""Per-structure arrays: ``(pdb ids, starts, eta, sigma, edge table, offsets, coefficients)``."""
# Sorted on the site's own identity, not on arrival order. The colouring, the edge indices and
# therefore every sampled quantity are functions of this order, so leaving it to the caller made
# `log Z` depend on how the pair table happened to be concatenated. With `_rng_for` below, this
# is the second half of "a per-structure score is a function of the structure".
sites = (sites.filter(pl.col("d_ca") <= model.radius)
.sort([c for c in ("pdb.id", "chain.id.rec", "chain.rec", "region.rec",
"pos.rec", "cls", "chain.id.par", "role.par", "pos.par")
if c in sites.columns]))
codes, _, q = site_codes(sites, model)
e = eta(codes, model)
sigma = q["sigma"].to_numpy()
upid, gid = np.unique(q["pdb.id"].to_numpy(), return_inverse=True)
starts = np.searchsorted(gid, np.arange(len(upid) + 1))
kv = np.asarray(model.kernel, dtype=float)
edge_lists = edges(q, joint=model.joint, far=model.far) if kv.size else []
E, eoff = bucket_edges(edge_lists, starts[:-1], kv) if edge_lists else (
np.zeros((0, 3), np.int64), np.zeros(len(upid) + 1, np.int64))
return upid, starts, e, sigma, E, eoff, kv, q
def _one(lo: int, hi: int, e, sigma, E, eoff, kv, s: int):
"""``(eta, sigma, A, colours)`` for one structure, with the colouring asserted valid."""
n = hi - lo
Es = E[eoff[s]:eoff[s + 1]]
ea, eb, ec = Es[:, 0] - lo, Es[:, 1] - lo, Es[:, 2]
A = coupling_matrix(n, ea, eb, ec, kv) if len(ea) else np.zeros((n, n))
cols = colour(n, ea, eb) if len(ea) else ([np.arange(n)] if n else [])
for cidx in cols:
assert not A[np.ix_(cidx, cidx)].any(), "same-colour sites are coupled"
return e[lo:hi], sigma[lo:hi], A, cols
#: Filled once per worker process by :func:`_init_workers`, so the prepared arrays cross the
#: process boundary once rather than once per structure.
_SHARED: dict = {}
def _init_workers(payload: dict) -> None:
_SHARED.clear()
_SHARED.update(payload)
def _run_chunk(args):
body, idx, kw = args
if body not in _BODIES:
# A spawned worker imports this module, not whatever registered the body. The ensemble
# bodies register themselves into `_BODIES` on import, so import them on demand rather
# than at module scope, which would be a cycle (`ensemble` imports from here).
from . import ensemble # noqa: F401
return [_BODIES[body](s, **kw) for s in idx]
def _map_structures(body: str, n: int, payload: dict, kw: dict, workers: int | None):
"""Run one per-structure body over every structure, in contiguous chunks.
Safe to parallelise only because each result is a function of ``(seed, pdb.id)`` and of that
structure's own sites — never of its position in the frame — so the output is identical however
the work is split. Chunks are contiguous and as few as there are workers: a pool of thousands of
one-structure tasks spends its time in dispatch and pickling.
Measured on an M3 (16 cores), 64 structures of 20--80 sites: serial 48.9 ms per structure,
3.6x at 8 processes. Threads make it *slower* (0.33x) — the arrays are too small for numpy to
release the GIL usefully.
"""
if workers is None:
workers = os.cpu_count() or 1
if workers <= 1 or n < 2 * workers:
_init_workers(payload)
return _run_chunk((body, list(range(n)), kw))
bounds = np.linspace(0, n, workers + 1).astype(int)
chunks = [(body, list(range(bounds[i], bounds[i + 1])), kw)
for i in range(workers) if bounds[i + 1] > bounds[i]]
with ProcessPoolExecutor(max_workers=workers, initializer=_init_workers,
initargs=(payload,)) as ex:
return [row for part in ex.map(_run_chunk, chunks) for row in part]
def _payload(upid, starts, e, sigma, E, eoff, kv) -> dict:
return {"upid": upid, "starts": starts, "e": e, "sigma": sigma, "E": E, "eoff": eoff, "kv": kv}
def _unpack(s: int):
"""``(pdb id, lo, hi, rng, eta, sigma, A, colours)`` for structure ``s`` from the shared arrays."""
d = _SHARED
lo, hi = d["starts"][s], d["starts"][s + 1]
et, sg, A, cols = _one(lo, hi, d["e"], d["sigma"], d["E"], d["eoff"], d["kv"], s)
return d["upid"][s], lo, hi, et, sg, A, cols
def _body_score_sites(s: int, *, particles: int, steps: int, seed: int) -> dict:
pid, lo, hi, et, sg, A, cols = _unpack(s)
log_z, ess = ais_log_z(et, A, cols, _rng_for(seed, pid), particles=particles, steps=steps)
en = energy(sg, et, A)
f = et + A @ sg
pll = float(np.sum(sg * f - np.logaddexp(0.0, f)))
return {"pdb.id": pid, "n_sites": int(hi - lo), "n_contacts": int(sg.sum()),
"energy": en, "neg_energy": -en,
"log_z": log_z, "log_z0": factorised_log_z(et),
"log_lik": -en - log_z, "psi": (-en - log_z) / (hi - lo),
"pseudo_log_lik": pll, "psi_pseudo": pll / (hi - lo),
"ais_ess": ess, "n_colours": len(cols)}
def _body_bound_unbound(s: int, *, threshold, chains: int, burn: int, draws: int, thin: int,
particles: int, steps: int, seed: int) -> dict:
pid, lo, hi, et, sg, A, cols = _unpack(s)
rng = _rng_for(seed, pid)
_, tot = gibbs(et, A, cols, rng, chains=chains, burn=burn, draws=draws, thin=thin)
lz, _ = ais_log_z(et, A, cols, rng, particles=particles, steps=steps)
n_obs = float(sg.sum())
return {"pdb.id": pid, "n_sites": hi - lo, "n_contacts": n_obs,
"neg_energy": -energy(sg, et, A),
"log_z": lz, "df_empty": delta_f_empty(lz),
"df_threshold": (float("nan") if threshold is None
else delta_f_threshold(tot, threshold)),
"mu_star": mu_star(tot, n_obs),
"n_lo": float(tot.min()) if len(tot) else float("nan"),
"n_hi": float(tot.max()) if len(tot) else float("nan"),
"n_mean": float(tot.mean()) if len(tot) else float("nan"),
"n_var": float(tot.var()) if len(tot) else float("nan")}
def _body_contact_probabilities(s: int, *, chains: int, burn: int, draws: int, thin: int,
seed: int) -> dict:
pid, lo, hi, et, sg, A, cols = _unpack(s)
occ, _ = gibbs(et, A, cols, _rng_for(seed, pid), chains=chains, burn=burn, draws=draws,
thin=thin)
# `lo`/`hi` travel back with the arrays so the caller scatters them into the right slice of
# the full-table buffer; a worker cannot see the frame it came from.
return {"lo": int(lo), "hi": int(hi), "p_model": occ,
"p_conditional": 1.0 / (1.0 + np.exp(-(et + A @ sg)))}
_BODIES = {"score_sites": _body_score_sites, "bound_unbound": _body_bound_unbound,
"contact_probabilities": _body_contact_probabilities}
[docs]
def score_sites(sites: pl.DataFrame, model: PottsModel, *, particles: int = 64,
steps: int = 256, seed: int = 0, workers: int | None = None) -> pl.DataFrame:
"""Energy, ``log Z`` and likelihoods for every structure in a table of available pairs.
Columns, one row per structure:
``n_sites`` / ``n_contacts``
available pairs and the contacts among them.
``energy`` / ``neg_energy``
``E(sigma_obs)``, the Hamiltonian of the observed contact map, and its negation. Lower
``energy`` is more favourable, so ``neg_energy`` is the one that runs the same way as
every other score here: higher is more native-like. It is the :math:`\\Pi` block of
:func:`tcren.reliability.s_score`.
``log_z`` / ``log_z0``
the coupled partition function by AIS, and the factorised one in closed form. ``log_z0``
is the free energy of the available set *before* any contact is observed.
``log_lik`` / ``psi``
``log P(sigma_obs)`` and the same divided by ``n_sites``, so it compares across
interfaces of different size.
``pseudo_log_lik`` / ``psi_pseudo``
the exact, MCMC-free cross-check.
``ais_ess``
effective sample size of the AIS weights, out of ``particles``. Close to ``particles``
means the annealing schedule was long enough.
Args:
sites: Rows from :func:`tcren.potts.available_pairs`, one or many structures.
model: A fitted :class:`PottsModel`.
particles: AIS particles per structure.
steps: AIS annealing steps.
seed: Seed for the sampler.
workers: Processes to split the structures across. ``None`` takes every core, ``1``
runs serially. The result is identical either way — each structure's numbers
depend on ``(seed, pdb.id)`` alone, never on how the work was split.
"""
upid, starts, e, sigma, E, eoff, kv, _ = _prepare(sites, model)
rows = _map_structures("score_sites", len(upid), _payload(upid, starts, e, sigma, E, eoff, kv),
{"particles": particles, "steps": steps, "seed": seed}, workers)
return pl.DataFrame(rows)
#: What each ``by=`` groups on, beyond ``pdb.id``. ``"pair"`` is the ungrouped passthrough.
_MAP_KEYS = {"pair": (),
"loop": ("region.rec", "pos.par", "aa.par"),
"position": ("pos.par", "aa.par")}
[docs]
def peptide_free_energy(sites: pl.DataFrame, model: PottsModel, *, coupled: bool = False,
marginals: pl.DataFrame | None = None, chains: int = 64, burn: int = 100,
draws: int = 100, thin: int = 3, seed: int = 0,
workers: int | None = None) -> pl.DataFrame:
r"""Historical log-partition substitution readout, not a calibrated binding free energy.
:func:`contact_map` reads how engaged a position is expected to be *before any residue identity
is scored*. This reads what happens when the identity changes. The partner residue enters the
one-body field twice — through the partner propensity :math:`h^{\mathrm{par}}` and through the
pair term :math:`J` — so substituting position :math:`i` shifts :math:`\eta` at every available
pair carrying that position, and the uncoupled log-partition changes with it:
.. math::
\Phi^{\mathrm{Potts}}(x) \;=\; \log Z_0\big(\eta(x)\big)
\;=\; \sum_a \log\!\big(1 + e^{\eta_a(x)}\big),
\qquad
\Delta\Phi_i(a) \;=\; \Phi^{\mathrm{Potts}}(x_{i \to a})
\;-\; \tfrac{1}{20}\sum_b \Phi^{\mathrm{Potts}}(x_{i \to b})
Higher is more favourable: :math:`\log Z_0` is the interface's capacity to make contacts at all,
so a residue that raises it engages more. The reference is the **equimolar** one — the mean over
the twenty residues at that position, not the residue the structure happens to carry — which is
the null a positional-scanning library actually holds the other positions at.
``log Z0`` is dimensionless, in nats. The uncoupled statistical free energy is
``-log Z0``, with the opposite sign; neither is a calibrated binding free energy.
These fitted-model readouts are historical, excluded from final fit-free scoring.
Two readings, from the same fields:
``coupled=False`` (default)
:math:`\log Z_0` for the coupling-free model, which is exact and closed form — no sampling.
``coupled=True``
linear response about the observed sequence. Since
:math:`\partial \log Z / \partial \eta_a = \langle\sigma_a\rangle`, the coupled log-partition
changes to first order as :math:`\Delta \log Z \approx \sum_a p_a \Delta\eta_a` with :math:`p_a` the marginal
of :func:`contact_probabilities` — one Gibbs pass, then a dot product per cell.
This is not an exact finite-substitution coupled log-partition difference.
The retained ``log_z0`` column then stores the frozen-marginal linear surrogate,
not an estimate of the coupled normalizer.
Only ``aa.par`` changes: the backbone, the Cα distances, the receptor residues and the partner
roles are the structure's own and are held fixed, which is the same fixed-backbone approximation
every threading score in the package makes.
Args:
sites: Rows from :func:`tcren.potts.available_pairs`.
model: A fitted :class:`PottsModel`.
coupled: Take the linear-response path against the coupled marginals.
marginals: A frame from :func:`contact_probabilities` to reuse when ``coupled=True``.
``None`` computes it.
chains, burn, draws, thin: Gibbs settings, used only when ``coupled=True``.
seed: Seed for the sampler.
workers: Processes to split the structures across when ``coupled=True``.
Returns:
One row per (``pdb.id``, ``pos.par``, ``aa.par``) with ``log_z0`` (the whole-interface
:math:`\log Z_0` under that substitution), ``dF`` (its equimolar-referenced effect),
``n_pairs`` (available pairs carrying that position) and ``is_observed`` (1 for the residue
the structure carries). ``dF`` sums to zero over the twenty residues at every position.
Example:
>>> from tcren.potts import available_pairs, peptide_free_energy, PottsModel # doctest: +SKIP
>>> peptide_free_energy(available_pairs(structure), PottsModel.bundled()) # doctest: +SKIP
"""
if sites.is_empty():
return pl.DataFrame(schema={"pdb.id": pl.String, "pos.par": pl.Int64, "aa.par": pl.String,
"log_z0": pl.Float64, "dF": pl.Float64,
"n_pairs": pl.Int64, "is_observed": pl.Int64})
aa = tuple(model.alphabet)
sites = sites.filter(pl.col("d_ca") <= model.radius)
codes, _sizes, q = site_codes(sites, model)
# A partner position is one residue of one chain, so every site carrying it must agree on the
# identity. Averaging over a disagreement would silently score a sequence that does not exist.
clash = (q.group_by("pdb.id", "pos.par").agg(pl.col("aa.par").n_unique().alias("n"))
.filter(pl.col("n") > 1))
if not clash.is_empty():
raise ValueError(f"{clash.height} partner position(s) carry more than one residue, so "
f"there is no sequence to substitute into: {clash.head(3).to_dicts()}")
# eta splits into a part the partner residue does not touch and a part it does. Recovering the
# first as eta(observed) minus the second keeps ONE definition of eta in the package: any change
# to `sites.eta` propagates here rather than being silently re-implemented.
J, h_par = model.coupling_array(), np.asarray(model.h_par)
rest = eta(codes, model) - h_par[codes[1]] - J[codes[0], codes[1]]
# (n_sites, 20): the field at every site under every candidate partner residue
eta_all = rest[:, None] + h_par[None, :] + J[codes[0], :]
if coupled:
marg = (marginals if marginals is not None else
contact_probabilities(sites, model, chains=chains, burn=burn, draws=draws,
thin=thin, seed=seed, workers=workers))
keys = [c for c in ("chain.id.rec", "chain.id.par") if c in sites.columns] + ["pdb.id", "chain.rec", "region.rec", "pos.rec", "cls", "role.par", "pos.par"]
if all(k in marg.columns for k in keys):
# contact_probabilities sorts sites; align by identity, never arrival order.
aligned = (q.select(keys).with_row_index("_order")
.join(marg.select(*keys, "p_model"), on=keys, how="left",
validate="1:1").sort("_order"))
if aligned["p_model"].null_count():
raise ValueError("marginals must match every site by identity")
p = aligned["p_model"].to_numpy()
elif marg.columns == ["p_model"]:
# Retain the historical explicit positional-vector input.
p = marg["p_model"].to_numpy()
else:
raise ValueError("marginals must carry site identity columns or only p_model")
if p.shape != (q.height,):
raise ValueError(f"marginals must have one row per site ({q.height}), got {p.shape}")
# log Z is linear in eta to first order with slope <sigma>; the constant cancels in dF
per_site = p[:, None] * eta_all
else:
per_site = np.logaddexp(0.0, eta_all)
# what each site contributes at the residue the structure actually carries there
at_observed = per_site[np.arange(q.height), codes[1]]
out = []
for (pid,), g in q.with_row_index("_r").group_by(["pdb.id"], maintain_order=True):
whole = float(at_observed[g["_r"].to_numpy()].sum())
for pos, gp in g.group_by(["pos.par"], maintain_order=True):
rp = gp["_r"].to_numpy()
# sites away from this position keep their own residues and contribute the same term
# whatever this position carries, so they enter as one constant and drop out of dF
here = per_site[rp].sum(axis=0)
z = whole - float(at_observed[rp].sum()) + here
obs = gp["aa.par"][0]
out += [{"pdb.id": pid, "pos.par": int(pos[0]), "aa.par": a,
"log_z0": float(z[k]), "dF": float(here[k] - here.mean()),
"n_pairs": len(rp), "is_observed": int(a == obs)}
for k, a in enumerate(aa)]
return pl.DataFrame(out).sort("pdb.id", "pos.par", "aa.par")
[docs]
def bound_unbound(sites: pl.DataFrame, model: PottsModel, *, threshold: int | None = None,
chains: int = 64, burn: int = 100, draws: int = 100, thin: int = 3,
particles: int = 64, steps: int = 256, seed: int = 0,
workers: int | None = None) -> pl.DataFrame:
"""The whole-interface two-state free energy, in its three readings.
A single site has two states, ``sigma_a = 0`` and ``1``, and ``eta_a`` is the free-energy
difference between them. The same contrast for the whole interface needs a macrostate, and the
contact count ``N(sigma)`` is the coordinate that defines one. Because every tilt in ``N`` is
an exponential family, one Gibbs pass serves all three readings.
Columns, one row per structure:
``neg_energy``
:math:`-E(\\sigma_{\\mathrm{obs}})`, higher is more native-like. The :math:`\\Pi` block of
:func:`tcren.reliability.s_score`, emitted here so one call supplies both the free-energy
readings and the energy term the reliability score needs.
``df_empty``
``log[P(N >= 1) / P(N = 0)] = log(Z - 1)``, exact, from the AIS ``log Z``. ``E(empty) = 0``,
so the empty configuration needs no separate estimate. This is the *capacity* of the
interface: what it would gain by making any contact at all.
``df_threshold``
``log[P(N >= x) / P(N < x)]`` at ``threshold``, from the sampled histogram. ``Z`` cancels,
so this needs no AIS — but it is only finite while ``x`` sits inside the sampled range.
``mu_star``
the tilt at which ``<N>_mu`` equals the observed count: how much extra chemical potential
the model needs to explain the map it was shown. Positive means the structure made more
contacts than its fields and couplings warrant. ``nan`` outside the sampled support.
``n_lo``, ``n_hi``, ``n_mean``, ``n_var``
the sampled contact-count range, mean and variance, so every ``nan`` above is auditable.
Note:
``df_empty`` and ``df_threshold`` are not two estimates of one number. The unbound basin of
a *docked* pose is astronomically improbable — the model is conditioned on an available set
that already holds the receptor against the peptide — so no sampler reaches ``N = 0`` and
only the ``log Z`` route gives it. The threshold reading is a local contrast inside the
bound basin.
Args:
sites: Rows from :func:`tcren.potts.available_pairs`.
model: A fitted :class:`PottsModel`.
threshold: ``x`` for ``df_threshold``. ``None`` disables that column.
chains, burn, draws, thin: Gibbs settings; see :func:`tcren.potts.gibbs`.
particles, steps: AIS settings; see :func:`tcren.potts.ais_log_z`.
seed: Seed for both samplers.
workers: Processes to split the structures across. ``None`` takes every core, ``1``
runs serially. The result is identical either way — each structure's numbers
depend on ``(seed, pdb.id)`` alone, never on how the work was split.
Returns:
One row per structure.
"""
upid, starts, e, sigma, E, eoff, kv, _ = _prepare(sites, model)
out = _map_structures("bound_unbound", len(upid), _payload(upid, starts, e, sigma, E, eoff, kv),
{"threshold": threshold, "chains": chains, "burn": burn, "draws": draws,
"thin": thin, "particles": particles, "steps": steps, "seed": seed},
workers)
return pl.DataFrame(out)
[docs]
def count_profile(sites: pl.DataFrame, model: PottsModel, *, chains: int = 64, burn: int = 100,
draws: int = 100, thin: int = 3, seed: int = 0) -> pl.DataFrame:
"""Pooled ``F(N) = -log p(N)`` over every structure in ``sites``, plus the observed counts.
The free-energy profile along the contact-count coordinate. Use it to see whether the model's
contact-count landscape has a barrier -- if it does not, a threshold reading of the two-state
contrast has nothing to key on and ``mu_star`` is the meaningful statistic.
"""
upid, starts, e, sigma, E, eoff, kv, _ = _prepare(sites, model)
keep, obs = [], []
for s in range(len(upid)):
lo, hi = starts[s], starts[s + 1]
rng = _rng_for(seed, upid[s])
et, sg, A, cols = _one(lo, hi, e, sigma, E, eoff, kv, s)
_, tot = gibbs(et, A, cols, rng, chains=chains, burn=burn, draws=draws, thin=thin)
keep.append(tot)
obs.append(sg.sum())
n, f = count_free_energy(np.concatenate(keep))
o = np.asarray(obs, float)
c = np.array([(o == k).sum() for k in n], float)
with np.errstate(divide="ignore"):
f_obs = -np.log(c / max(c.sum(), 1.0))
return pl.DataFrame({"n_contacts": n, "f_model": f, "f_observed": f_obs,
"count_observed": c})
[docs]
def sample_maps(sites: pl.DataFrame, model: PottsModel, *, chains: int = 64, burn: int = 100,
draws: int = 100, thin: int = 3, seed: int = 0) -> pl.DataFrame:
"""Contact totals of maps drawn from the model, one row per structure per draw.
The generative check: a model that reproduces a real interface must reproduce the *spread* of
its contact count, not only the mean.
"""
upid, starts, e, sigma, E, eoff, kv, _ = _prepare(sites, model)
out = []
for s in range(len(upid)):
lo, hi = starts[s], starts[s + 1]
rng = _rng_for(seed, upid[s])
et, sg, A, cols = _one(lo, hi, e, sigma, E, eoff, kv, s)
_, totals = gibbs(et, A, cols, rng, chains=chains, burn=burn, draws=draws, thin=thin)
out.append(pl.DataFrame({"pdb.id": [upid[s]] * len(totals),
"n_contacts_observed": [float(sg.sum())] * len(totals),
"n_contacts_sampled": totals}))
return pl.concat(out) if out else pl.DataFrame()
[docs]
def connected_correlations(sites: pl.DataFrame, model: PottsModel, *, chains: int = 64,
burn: int = 100, draws: int = 100, thin: int = 3,
seed: int = 0) -> pl.DataFrame:
r"""The generative test: two-point correlations in the data against in maps sampled from it.
For each coupling class, pooled over every edge in it and every structure,
.. math::
c_\Delta = \langle \sigma_a \sigma_b \rangle
- \langle \sigma_a \rangle \langle \sigma_b \rangle
computed once from the observed contact maps and once from maps drawn from the model. This is
the criterion Boltzmann-machine DCA trains to, and a pseudolikelihood fit never sees it — the
conditionals it maximises are one-site quantities — so agreement here is a real test rather
than a restatement of the objective.
Returns:
One row per coupling class: ``class``, ``n_edges``, ``c_data``, ``c_model``.
"""
upid, starts, e, sigma, E, eoff, kv, q = _prepare(sites, model)
n_k = len(kv)
if not n_k:
return pl.DataFrame(schema={"class": pl.Utf8, "n_edges": pl.Int64,
"c_data": pl.Float64, "c_model": pl.Float64})
acc = {k: np.zeros(n_k) for k in ("m2", "m1a", "m1b", "d2", "d1a", "d1b", "dn")}
m_n = np.zeros(n_k)
for s in range(len(upid)):
lo, hi = starts[s], starts[s + 1]
rng = _rng_for(seed, upid[s])
et, sg, A, cols = _one(lo, hi, e, sigma, E, eoff, kv, s)
Es = E[eoff[s]:eoff[s + 1]]
ea, eb, ec = Es[:, 0] - lo, Es[:, 1] - lo, Es[:, 2]
if not len(ea):
continue
acc["d2"] += np.bincount(ec, weights=sg[ea] * sg[eb], minlength=n_k)
acc["d1a"] += np.bincount(ec, weights=sg[ea], minlength=n_k)
acc["d1b"] += np.bincount(ec, weights=sg[eb], minlength=n_k)
acc["dn"] += np.bincount(ec, minlength=n_k)
two, one_a, one_b, tot = _two_point(et, A, cols, rng, ea, eb, ec, n_k,
chains, burn, draws, thin)
acc["m2"] += two
acc["m1a"] += one_a
acc["m1b"] += one_b
m_n += np.bincount(ec, minlength=n_k) * tot
dn = np.where(acc["dn"] > 0, acc["dn"], np.nan)
mn = np.where(m_n > 0, m_n, np.nan)
return pl.DataFrame({
"class": kernel_names(model.joint, model.far)[:n_k],
"n_edges": acc["dn"].astype(np.int64),
"c_data": acc["d2"] / dn - (acc["d1a"] / dn) * (acc["d1b"] / dn),
"c_model": acc["m2"] / mn - (acc["m1a"] / mn) * (acc["m1b"] / mn),
})
def _two_point(eta_s, A, cols, rng, ea, eb, ec, n_k, chains, burn, draws, thin):
"""Accumulate ``sum sigma_a sigma_b`` per class over sampled maps, without storing them."""
n = len(eta_s)
sig = (rng.random((chains, n)) < 1.0 / (1.0 + np.exp(-eta_s))).astype(np.float64)
two, one_a, one_b, tot = np.zeros(n_k), np.zeros(n_k), np.zeros(n_k), 0
coupled = A.any()
for it in range(burn + draws * thin):
for cidx in cols:
f = eta_s[cidx] + (sig @ A[:, cidx] if coupled else 0.0)
sig[:, cidx] = (rng.random((chains, len(cidx)))
< 1.0 / (1.0 + np.exp(-f))).astype(np.float64)
if it >= burn and (it - burn) % thin == 0:
np.add.at(two, ec, (sig[:, ea] * sig[:, eb]).sum(0))
np.add.at(one_a, ec, sig[:, ea].sum(0))
np.add.at(one_b, ec, sig[:, eb].sum(0))
tot += chains
return two, one_a, one_b, tot
[docs]
def score_structure(structure, model: PottsModel | None = None, *, partner: str = "peptide",
**kwargs) -> dict:
"""Convenience: enumerate one structure's available pairs and score them.
Args:
structure: A chain-typed structure (MHC-annotated as well, for ``partner="mhc"``).
model: A fitted model; defaults to the bundled TCR:peptide one.
partner: ``"peptide"`` or ``"mhc"``.
**kwargs: Passed to :func:`score_sites`.
"""
model = model or PottsModel.bundled()
pairs = available_pairs(structure, partner, radius=model.radius, cutoff=model.cutoff)
if pairs.is_empty():
return {"pdb.id": structure.pdb_id, "n_sites": 0, "n_contacts": 0}
return score_sites(pairs, model, **kwargs).to_dicts()[0]
[docs]
def map_agreement(sites: pl.DataFrame, *, column: str = "p_model",
observed: str = "sigma") -> pl.DataFrame:
r"""How well a predicted contact map agrees with the one a structure actually made.
Two numbers, because they answer different questions and no single one answers both:
``js``
**Placement.** Jensen--Shannon similarity between the two maps read as distributions over
sites, :math:`q_a = \sigma_a / \sum \sigma` and :math:`r_a = p_a / \sum p`:
.. math:: \mathrm{js} = 1 - \tfrac12 D_{\mathrm{KL}}(q \| m)
- \tfrac12 D_{\mathrm{KL}}(r \| m),
\qquad m = \tfrac12 (q + r)
in bits, so it lies in :math:`[0, 1]`: 1 is identical placement, 0 disjoint support. It is
**scale-free in the level** by construction --- normalising divides the total out --- which
is why ``ratio`` is reported beside it and not folded in.
:math:`\mathrm{js} = 1` requires the model to put *all* its mass on exactly the contacting
cells in equal proportion. This also holds for :math:`p_a=c\sigma_a`,
with any :math:`0<c<1`; it does not require a deterministic map. The level
ratio distinguishes these proportional maps.
``ratio``
**Level.** :math:`\sum_a p_a / \sum_a \sigma_a`. One is the count reproduced.
``auc`` is reported too, as the ordering read alone.
**Not R².** On a binary outcome ``1 - SS_res/SS_tot`` *is* the Brier skill score --- they agree
to three decimals on every map we have checked --- so reporting both is reporting one number
twice. It is also dominated by the many easy zeros and scales with prevalence, so two maps with
different site counts are not comparable under it. That is exactly the comparison this function
exists for, which is why the placement read-out here is normalised and symmetric instead.
Args:
sites: A scored site table --- from :func:`contact_probabilities`,
or ``tcren potts contacts``.
Grouped by ``pdb.id``; rows with a null ``observed`` are read as non-contacts.
column: The probability column to score.
observed: The observed indicator column.
Returns:
One row per ``pdb.id``: ``sites``, ``contacts``, ``expected``, ``ratio``, ``js``, ``auc``.
"""
if column not in sites.columns:
raise ValueError(f"no {column!r} column; scored tables carry p_model, p_independent and "
f"p_conditional")
out = []
for pid, d in sites.group_by("pdb.id", maintain_order=True):
y = d[observed].fill_null(0).to_numpy().astype(float)
p = d[column].fill_null(0.0).to_numpy().astype(float)
n_obs, n_exp = float(y.sum()), float(p.sum())
out.append({
"pdb.id": pid[0] if isinstance(pid, tuple) else pid,
"sites": int(d.height), "contacts": int(n_obs), "expected": n_exp,
"ratio": n_exp / n_obs if n_obs else float("nan"),
"js": _js_similarity(y, p),
"auc": _rank_auc(p, y > 0.5),
})
return pl.DataFrame(out)
def _js_similarity(y: np.ndarray, p: np.ndarray, eps: float = 1e-12) -> float:
"""1 - JSD in bits between the two maps normalised to distributions over sites."""
if y.sum() <= 0 or p.sum() <= 0:
return float("nan")
q, r = y / y.sum(), p / p.sum()
m = 0.5 * (q + r)
def kl(a, b):
return float(np.sum(np.where(a > 0, a * np.log2((a + eps) / (b + eps)), 0.0)))
return 1.0 - 0.5 * (kl(q, m) + kl(r, m))
def _rank_auc(score: np.ndarray, truth: np.ndarray) -> float:
"""ROC-AUC by rank sum, ties averaged. A flat predictor must read 0.5, not better."""
pos, neg = int(truth.sum()), int((~truth).sum())
if not pos or not neg:
return float("nan")
order = np.argsort(score, kind="mergesort")
ranks = np.empty(len(score), float)
ranks[order] = np.arange(1, len(score) + 1)
srt = score[order]
i = 0
while i < len(srt):
j = i
while j + 1 < len(srt) and srt[j + 1] == srt[i]:
j += 1
if j > i:
ranks[order[i:j + 1]] = ranks[order[i:j + 1]].mean()
i = j + 1
return float((ranks[truth].sum() - pos * (pos + 1) / 2) / (pos * neg))