import math
import logging
import warnings
import numpy as np
import pandas as pd
import matplotlib.pyplot as plt
import seaborn as sns
warnings.filterwarnings("ignore", category=RuntimeWarning)
logging.basicConfig(level=logging.INFO, format="%(levelname)s - %(message)s")
logger = logging.getLogger(__name__)
# --------------------------------------------------------------------------- #
# JSD math (1D / 2D, histogram based) #
# --------------------------------------------------------------------------- #
def sqrt_rule(n):
return math.ceil(math.sqrt(n))
def sturges_rule(n):
return math.ceil(math.log(n, 2) + 1)
def resolve_num_bins(bins, sizes, auto_threshold=2500, auto_default=50):
"""
"auto" -> sqrt rule if min(sizes) < auto_threshold, else auto_default.
"sqrt" -> square-root rule.
"sturges" -> Sturges' formula.
int -> used directly.
"""
if isinstance(bins, (int, np.integer)):
return int(bins)
min_n = min(sizes)
if bins == "auto":
return auto_default if min_n >= auto_threshold else sqrt_rule(min_n)
if bins == "sqrt":
return sqrt_rule(min_n)
if bins == "sturges":
return sturges_rule(min_n)
raise ValueError(f"Unknown bins rule: {bins!r}")
def resolve_bins_for_axes(bins, sizes):
"""
Resolve a bins spec for a 2D (joint) feature into (n_bins_x, n_bins_y).
`bins` can be a single value ("auto"/"sqrt"/"sturges"/int, same bin
count on both axes) or a (bins_x, bins_y) 2-tuple (each axis resolved
independently).
"""
if isinstance(bins, (tuple, list)) and len(bins) == 2:
return resolve_num_bins(bins[0], sizes), resolve_num_bins(bins[1], sizes)
n = resolve_num_bins(bins, sizes)
return n, n
def _freqs(x, edges):
h = np.histogram(x, bins=edges)[0].astype(float)
total = h.sum()
return h / total if total > 0 else h
def _kld(p, m):
nz = p > 0
if nz.sum() == 0:
return 0.0
return float(np.sum(p[nz] * np.log(p[nz] / m[nz])))
def jsd_from_freqs(p, q):
m = 0.5 * (p + q)
return 0.5 * (_kld(p, m) + _kld(q, m))
def jsd_1d(p_data, q_data, bins, value_range):
"""1D histogram JSD over a fixed `value_range` (resolved once by the
caller, NOT recomputed per bootstrap draw)."""
n_bins = resolve_num_bins(bins, [p_data.shape[0], q_data.shape[0]])
lo, hi = value_range
edges = np.linspace(lo, hi, n_bins + 1)
return jsd_from_freqs(_freqs(p_data, edges), _freqs(q_data, edges))
def jsd_2d(p_xy, q_xy, bins, x_range, y_range):
"""
2D histogram JSD for a joint pair of columns (e.g. PC1/PC2, phi/psi,
or any 2 precomputed features). Each axis gets its own fixed range
(circular axes -> [-pi, pi], linear axes -> observed pooled min/max),
resolved once by the caller.
"""
n_bins_x, n_bins_y = resolve_bins_for_axes(bins, [p_xy.shape[0], q_xy.shape[0]])
x_edges = np.linspace(x_range[0], x_range[1], n_bins_x + 1)
y_edges = np.linspace(y_range[0], y_range[1], n_bins_y + 1)
h_p = np.histogram2d(p_xy[:, 0], p_xy[:, 1], bins=[x_edges, y_edges])[0].ravel()
h_q = np.histogram2d(q_xy[:, 0], q_xy[:, 1], bins=[x_edges, y_edges])[0].ravel()
h_p = h_p / max(h_p.sum(), 1)
h_q = h_q / max(h_q.sum(), 1)
return jsd_from_freqs(h_p, h_q)
def jsd_score(f1, f2, circular, bins, value_range, is_2d=False):
if is_2d:
rx, ry = value_range
return jsd_2d(f1, f2, bins=bins, x_range=rx, y_range=ry)
return jsd_1d(f1[:, 0], f2[:, 0], bins=bins, value_range=value_range)
def percentile_ci(values, confidence=0.95):
alpha = (1 - confidence) / 2
return (
float(np.percentile(values, alpha * 100)),
float(np.percentile(values, (1 - alpha) * 100)),
)
# --------------------------------------------------------------------------- #
# Angle unit detection + per-group matrix assembly #
# --------------------------------------------------------------------------- #
def _looks_like_degrees(values):
"""Heuristic: angles/dihedrals in radians live in [-pi, pi]
(~[-3.14, 3.14]); degrees commonly reach 180/360. If the observed
amplitude clearly exceeds pi, we assume degrees."""
return np.nanmax(np.abs(values)) > (np.pi * 1.05)
def prepare_feature_matrices(
df, sim_name_col, group_col, feature_cols, angle_flags, angle_units, verbose=True
):
"""
Build, for each group (in order of first appearance), an (n_i, F)
array for the requested feature column(s) (F=1 for 1D, F=2 for a 2D
joint feature), converting angle columns to radians (unit detection
pooled across all groups, or forced unit).
`group_col` is whatever column the caller grouped on -- normally the
per-replica id (group_by="replica", the default and most rigorous
mode) but can also be the system/sim_name column itself
(group_by="system", pooled mode: all replicas of a system are fused
into a single group before histogramming).
Returns
-------
feature_matrices : list of np.ndarray
One (n_i, F) array per group.
labels : list of str
Group labels, in order.
system_indices : dict
{sim_name: [indices into feature_matrices]}.
circular : list of bool
One flag per feature column.
value_range : list of (lo, hi)
One range per feature column (fixed to [-pi, pi] if circular).
"""
grouped = df.groupby(group_col, sort=False)
raw_per_group = []
labels = []
system_of = []
for glabel, g in grouped:
labels.append(str(glabel))
system_of.append(str(g[sim_name_col].iloc[0]))
raw_per_group.append([g[c].to_numpy(dtype=float) for c in feature_cols])
convert = []
circular = []
value_range = []
for ci, col in enumerate(feature_cols):
pooled = np.concatenate([grp[ci] for grp in raw_per_group])
if angle_flags[ci]:
unit = angle_units[ci]
if unit == "auto":
is_deg = _looks_like_degrees(pooled)
elif unit == "deg":
is_deg = True
elif unit == "rad":
is_deg = False
else:
raise ValueError(
f"Invalid angle_unit: {unit!r}, expected 'auto'/'deg'/'rad'."
)
convert.append(is_deg)
circular.append(True)
value_range.append((-np.pi, np.pi))
if verbose:
logger.info(
"Column '%s': detected as %s%s",
col,
"degrees" if is_deg else "radians",
" (auto)" if unit == "auto" else " (forced)",
)
else:
convert.append(False)
circular.append(False)
value_range.append((float(np.nanmin(pooled)), float(np.nanmax(pooled))))
feature_matrices = []
for grp_cols in raw_per_group:
converted = [
np.radians(arr) if convert[ci] else arr for ci, arr in enumerate(grp_cols)
]
feature_matrices.append(np.stack(converted, axis=1))
system_indices = {}
for idx, s in enumerate(system_of):
system_indices.setdefault(s, []).append(idx)
return feature_matrices, labels, system_indices, circular, value_range
# --------------------------------------------------------------------------- #
# Bootstrap engine #
# --------------------------------------------------------------------------- #
def bootstrap_score_matrix(
feature_matrices,
circular,
bins,
iters,
frac,
replace,
seed,
is_2d=False,
value_range=None,
verbose=False,
progress_label="",
):
"""
Bootstrap the JSD score matrix between every pair of groups (replicas,
or systems in pooled mode).
Returns
-------
scores : np.ndarray
(n, n, iters) array: score for each pair of groups, at each
bootstrap draw.
"""
rng = np.random.default_rng(seed)
n = len(feature_matrices)
scores = np.zeros((n, n, iters))
pairs = [(i, j) for i in range(n) for j in range(i, n)]
iterator = pairs
if verbose:
try:
from tqdm.auto import tqdm
iterator = tqdm(
pairs, desc=f"bootstrap JSD {progress_label}".strip(), unit="pair"
)
except ImportError:
logger.info(
"(tip: `pip install tqdm` for a progress bar -- %d pairs x %d draws)",
len(pairs),
iters,
)
for i, j in iterator:
fi, fj = feature_matrices[i], feature_matrices[j]
for k in range(iters):
ni = max(int(fi.shape[0] * frac), 1)
nj = max(int(fj.shape[0] * frac), 1)
idx_i = rng.choice(fi.shape[0], ni, replace=replace)
idx_j = rng.choice(fj.shape[0], nj, replace=replace)
s = jsd_score(
fi[idx_i],
fj[idx_j],
circular=circular,
bins=bins,
value_range=value_range,
is_2d=is_2d,
)
scores[i, j, k] = s
scores[j, i, k] = s
return scores
def _pairs(scores, idx1, idx2, exclude_diag):
vals = []
for i in idx1:
for j in idx2:
if exclude_diag and i == j:
continue
if i < j or not exclude_diag:
vals.append(scores[i, j])
return np.array(vals) if vals else np.empty((0, scores.shape[2]))
def _ci_overlap(ci_a, ci_b):
return not (ci_a[1] < ci_b[0] or ci_b[1] < ci_a[0])
def _bootstrap_ratio(
intra1_pooled, intra2_pooled, inter_pooled, ratio_bootstrap, ratio_seed
):
"""
Bootstrap the ratio inter_mean / max(intra_1_mean, intra_2_mean), by
resampling the already-bootstrapped score arrays. Independent pass
(size `ratio_bootstrap`), just to get a stable CI on the ratio (a
non-linear function of two random quantities).
"""
rng = np.random.default_rng(ratio_seed)
ratios = np.empty(ratio_bootstrap)
has_1, has_2 = intra1_pooled.size > 0, intra2_pooled.size > 0
for k in range(ratio_bootstrap):
it = rng.choice(inter_pooled, size=inter_pooled.size, replace=True).mean()
i1 = (
rng.choice(intra1_pooled, size=intra1_pooled.size, replace=True).mean()
if has_1
else np.nan
)
i2 = (
rng.choice(intra2_pooled, size=intra2_pooled.size, replace=True).mean()
if has_2
else np.nan
)
if has_1 and has_2:
ref = max(i1, i2)
elif has_1:
ref = i1
elif has_2:
ref = i2
else:
ref = np.nan
ratios[k] = (
(it / ref)
if (ref is not None and not np.isnan(ref) and ref > 1e-12)
else np.inf
)
return ratios
def _effect_magnitude_label(ratio_mean, ratio_ci):
"""
Heuristic label from the bootstrap distribution of the ratio
inter / max(intra_1, intra_2). If the CI contains 1, there is no clear
separation from inter-replica sampling noise -- flagged explicitly
regardless of the point estimate.
"""
lo, hi = ratio_ci
if lo <= 1.0 <= hi:
return "no clear separation (ratio CI includes 1)"
if ratio_mean < 2:
return "small"
if ratio_mean < 5:
return "moderate"
if ratio_mean < 10:
return "large"
return "very large"
def _jsd_magnitude_label(frac_of_max):
"""
Heuristic label from `inter_frac_of_max` (inter_mean / ln(2)) -- an
ABSOLUTE measure of JSD magnitude on its own scale [0, ln(2)],
insensitive to the "ratio explodes when intra is tiny" trap above.
"""
if frac_of_max is None:
return "n/a"
if frac_of_max < 0.05:
return "negligible"
if frac_of_max < 0.3:
return "moderate"
return "large"
def summarize_systems(
scores,
system_indices,
confidence_level=0.95,
feature_name="",
metric_max=None,
ratio_bootstrap=5000,
ratio_seed=0,
):
"""
Replica-level mode (group_by="replica"). For each pair of systems
(X, Y): intra-X, intra-Y, inter-XY (bootstrap mean + CI, no
p-value), a `verdict` based on CI non-overlap ("distinct" if the
inter CI overlaps NEITHER intra CI, "similar" if it overlaps both,
"ambiguous" if it overlaps exactly one), PLUS two complementary
magnitude heuristics (`effect_ratio_*`/`effect_magnitude` and
`inter_frac_of_max`/`jsd_magnitude` -- see code for details).
Single-system input -> intra-only row (inter/ratio/magnitude columns
set to None/"n/a").
"""
names = list(system_indices.keys())
intra_cache = {}
intra_pooled_cache = {}
for name in names:
idx = system_indices[name]
pooled = _pairs(scores, idx, idx, exclude_diag=True).reshape(-1)
intra_pooled_cache[name] = pooled
intra_cache[name] = (
pooled.mean() if pooled.size else float("nan"),
(
percentile_ci(pooled, confidence_level)
if pooled.size
else (float("nan"),) * 2
),
)
if len(names) == 1:
name = names[0]
m, ci = intra_cache[name]
return pd.DataFrame(
[
{
"feature": feature_name,
"system_1": name,
"system_2": None,
"intra_1_mean": m,
"intra_1_ci_low": ci[0],
"intra_1_ci_high": ci[1],
"intra_2_mean": None,
"intra_2_ci_low": None,
"intra_2_ci_high": None,
"inter_mean": None,
"inter_ci_low": None,
"inter_ci_high": None,
"verdict": "single system: intra-replica variability only",
"effect_ratio_mean": None,
"effect_ratio_ci_low": None,
"effect_ratio_ci_high": None,
"effect_magnitude": "n/a",
"inter_frac_of_max": None,
"jsd_magnitude": "n/a",
}
]
)
rows = []
for a in range(len(names)):
for b in range(a + 1, len(names)):
g1, g2 = names[a], names[b]
inter_pooled = _pairs(
scores, system_indices[g1], system_indices[g2], exclude_diag=False
).reshape(-1)
inter_mean, inter_ci = inter_pooled.mean(), percentile_ci(
inter_pooled, confidence_level
)
m1, ci1 = intra_cache[g1]
m2, ci2 = intra_cache[g2]
overlaps_1 = _ci_overlap(inter_ci, ci1) if not np.isnan(m1) else None
overlaps_2 = _ci_overlap(inter_ci, ci2) if not np.isnan(m2) else None
if overlaps_1 is None and overlaps_2 is None:
verdict = "n/a (need >=2 replicas per system)"
else:
n_overlap = sum(v for v in [overlaps_1, overlaps_2] if v is not None)
verdict = ["distinct", "ambiguous", "similar"][n_overlap]
ratio_samples = _bootstrap_ratio(
intra_pooled_cache[g1],
intra_pooled_cache[g2],
inter_pooled,
ratio_bootstrap=ratio_bootstrap,
ratio_seed=ratio_seed,
)
finite = ratio_samples[np.isfinite(ratio_samples)]
if finite.size:
ratio_mean = float(finite.mean())
ratio_ci = percentile_ci(finite, confidence_level)
else:
ratio_mean, ratio_ci = float("nan"), (float("nan"), float("nan"))
effect_magnitude = _effect_magnitude_label(ratio_mean, ratio_ci)
frac_of_max = (inter_mean / metric_max) if metric_max else None
jsd_magnitude = _jsd_magnitude_label(frac_of_max)
rows.append(
{
"feature": feature_name,
"system_1": g1,
"system_2": g2,
"intra_1_mean": m1,
"intra_1_ci_low": ci1[0],
"intra_1_ci_high": ci1[1],
"intra_2_mean": m2,
"intra_2_ci_low": ci2[0],
"intra_2_ci_high": ci2[1],
"inter_mean": inter_mean,
"inter_ci_low": inter_ci[0],
"inter_ci_high": inter_ci[1],
"verdict": verdict,
"effect_ratio_mean": ratio_mean,
"effect_ratio_ci_low": ratio_ci[0],
"effect_ratio_ci_high": ratio_ci[1],
"effect_magnitude": effect_magnitude,
"inter_frac_of_max": frac_of_max,
"jsd_magnitude": jsd_magnitude,
}
)
return pd.DataFrame(rows)
def summarize_systems_pooled(
scores,
system_indices,
confidence_level=0.95,
feature_name="",
metric_max=None,
):
"""
System-level mode (group_by="system"). All replicas of each system
were already fused into a single group upstream, so there is no
per-replica intra-system variability left to compare against here --
just a direct inter-system JSD (mean + bootstrap CI, where the
bootstrap only captures binning/sampling noise on the pooled data,
NOT replica-to-replica variability).
This is a quicker, more optimistic view than `summarize_systems`
(group_by="replica"): a "large" pooled JSD here can still turn out
"similar" in the replica-level mode if the replicas themselves are
noisy/divergent. Use both together -- pooled for a fast overview,
replica-level for the statistically defensible conclusion.
"""
names = list(system_indices.keys())
idx_of = {n: system_indices[n][0] for n in names}
if len(names) == 1:
return pd.DataFrame(
[
{
"feature": feature_name,
"system_1": names[0],
"system_2": None,
"inter_mean": None,
"inter_ci_low": None,
"inter_ci_high": None,
"inter_frac_of_max": None,
"jsd_magnitude": "n/a",
}
]
)
rows = []
for a in range(len(names)):
for b in range(a + 1, len(names)):
g1, g2 = names[a], names[b]
vals = scores[idx_of[g1], idx_of[g2]]
mean, ci = vals.mean(), percentile_ci(vals, confidence_level)
frac = (mean / metric_max) if metric_max else None
rows.append(
{
"feature": feature_name,
"system_1": g1,
"system_2": g2,
"inter_mean": mean,
"inter_ci_low": ci[0],
"inter_ci_high": ci[1],
"inter_frac_of_max": frac,
"jsd_magnitude": _jsd_magnitude_label(frac),
}
)
return pd.DataFrame(rows)
def replica_pair_summary(
scores, system_indices, labels, confidence_level=0.95, feature_name=""
):
"""Fine-grained intra-system, replica-vs-replica detail (mean + CI),
NOT pooled -- the finer view behind `intra-X`, useful for spotting a
replica that diverges from its "sisters" (sampling noise, poor
equilibration, a different conformational state...). Only meaningful
in group_by="replica" mode."""
rows = []
for sys_name, idxs in system_indices.items():
for a in range(len(idxs)):
for b in range(a + 1, len(idxs)):
i, j = idxs[a], idxs[b]
pooled = scores[i, j]
mean, ci = pooled.mean(), percentile_ci(pooled, confidence_level)
rows.append(
{
"feature": feature_name,
"system": sys_name,
"replica_1": labels[i],
"replica_2": labels[j],
"mean": mean,
"ci_low": ci[0],
"ci_high": ci[1],
}
)
return pd.DataFrame(rows)
def per_replica_outlier_summary(pair_df, z_threshold=1.2):
"""
For each system, average each replica's score against its "sisters",
then compute a `deviation_score` (z-score-like) and an `assessment`
("worth_checking" if above `z_threshold`, else "typical"). A
descriptive heuristic, not a hypothesis test.
CAUTION -- ceiling effect with few replicas: with N replicas,
`deviation_score` (population z-score, ddof=0) is bounded by
sqrt(N-1). With N=3 (a common case), that ceiling is ~1.414 -- a
threshold of 1.5 would therefore never trigger. `z_threshold=1.2`
remains usable at N=3.
"""
rows = []
for feat in pair_df["feature"].unique():
fsub = pair_df[pair_df["feature"] == feat]
for sys_name in fsub["system"].unique():
ssub = fsub[fsub["system"] == sys_name]
replicas = sorted(set(ssub["replica_1"]) | set(ssub["replica_2"]))
if len(replicas) < 3:
for r in replicas:
rows.append(
{
"feature": feat,
"system": sys_name,
"replica": r,
"avg_pairwise_score": np.nan,
"deviation_score": np.nan,
"assessment": "n/a (<3 replicas)",
}
)
continue
avgs = {}
for r in replicas:
vals = ssub.loc[
(ssub["replica_1"] == r) | (ssub["replica_2"] == r), "mean"
]
avgs[r] = vals.mean()
all_avgs = np.array(list(avgs.values()))
mu, sigma = all_avgs.mean(), all_avgs.std()
for r, a in avgs.items():
z = (a - mu) / sigma if sigma > 0 else 0.0
assessment = "worth_checking" if z > z_threshold else "typical"
rows.append(
{
"feature": feat,
"system": sys_name,
"replica": r,
"avg_pairwise_score": a,
"deviation_score": z,
"assessment": assessment,
}
)
return (
pd.DataFrame(rows)
.sort_values(["feature", "system", "replica"])
.reset_index(drop=True)
)
def aggregate_verdicts(summaries):
"""Combine the outputs of `summarize_systems` (one or several
features) into a cross-feature summary: for each pair of systems,
how many/which features say "distinct", and an overall verdict.
Only meaningful in group_by="replica" mode (needs a `verdict`
column)."""
all_df = pd.concat(summaries, ignore_index=True)
rows = []
for (g1, g2), sub in all_df.groupby(["system_1", "system_2"], dropna=False):
n_feat = len(sub)
distinct = sub[sub["verdict"] == "distinct"]
rows.append(
{
"system_1": g1,
"system_2": g2,
"n_features": n_feat,
"n_distinct": len(distinct),
"distinct_on": ", ".join(distinct["feature"].tolist()) or "-",
"not_distinct_on": ", ".join(
sub.loc[sub["verdict"] != "distinct", "feature"].tolist()
)
or "-",
"overall_verdict": (
"clearly different"
if len(distinct) == n_feat and n_feat > 0
else (
"not different (on any feature tested)"
if len(distinct) == 0
else "partially different (depends on feature)"
)
),
}
)
return pd.DataFrame(rows)
def display_summary(df, float_format="{:.4f}"):
"""Display an HTML table (notebook) or an ASCII table as fallback."""
fdf = df.copy()
for col in fdf.select_dtypes(include=[float, "float64"]).columns:
fdf[col] = fdf[col].map(lambda x: float_format.format(x) if pd.notna(x) else "")
try:
from IPython.display import display
display(fdf)
return
except Exception:
pass
cols = list(fdf.columns)
widths = [
max(len(str(c)), fdf[c].astype(str).map(len).max() if len(fdf) else 0)
for c in cols
]
sep = "+" + "+".join("-" * (w + 2) for w in widths) + "+"
header = "|" + "|".join(f" {c:<{w}} " for c, w in zip(cols, widths)) + "|"
print(sep)
print(header)
print(sep)
for _, row in fdf.iterrows():
print(
"|" + "|".join(f" {str(row[c]):<{w}} " for c, w in zip(cols, widths)) + "|"
)
print(sep)
# --------------------------------------------------------------------------- #
# Human-readable interpretation helper #
# --------------------------------------------------------------------------- #
def _ascii_table(headers, rows, aligns=None):
"""Render a simple '|'/'-' ASCII table (list of header strings, list
of row tuples of strings). Column widths auto-sized."""
n_cols = len(headers)
aligns = aligns or ["<"] * n_cols
widths = [
max(len(headers[c]), max((len(r[c]) for r in rows), default=0))
for c in range(n_cols)
]
sep = "+" + "+".join("-" * (w + 2) for w in widths) + "+"
lines = [sep]
lines.append(
"|" + "|".join(f" {headers[c]:<{widths[c]}} " for c in range(n_cols)) + "|"
)
lines.append(sep)
for r in rows:
lines.append(
"|"
+ "|".join(
f" {r[c]:{aligns[c]}{widths[c]}} " for c in range(n_cols)
)
+ "|"
)
lines.append(sep)
return "\n".join(lines)
def explain_jsd(tables, top_outliers=None):
"""
Print a quick, human-readable interpretation of the dict returned by
`compute_jsd`, as formatted ASCII tables, followed by a legend for
every key in the dict. Purely a convenience layer for a fast read --
does not replace inspecting the DataFrames directly for anything
beyond a first pass.
Works with both group_by modes:
- "replica" (default): full verdict/effect-size table, distinguishing
"DISTINGUISHABLE" (CI-supported difference) from "NOT
distinguishable" (intra-replica noise overlaps the inter-system
signal) from "AMBIGUOUS".
- "system" (pooled): a lighter table with just the inter-system JSD
and its absolute magnitude label (no verdict, since there is no
intra-system spread left to compare against).
Parameters
----------
tables : dict
Return value of `compute_jsd`.
top_outliers : int or None, optional
If set, only print the N most extreme "worth_checking" replicas
(sorted by |deviation_score|) instead of all of them.
"""
summary = tables.get("system_summary")
if summary is None or summary.empty:
print("No system_summary table found -- nothing to interpret.")
return
feat = summary["feature"].iloc[0]
is_replica_mode = "verdict" in summary.columns
title = f" {feat} -- quick interpretation "
print("=" * max(len(title), 60))
print(title.center(max(len(title), 60)))
print("=" * max(len(title), 60))
bins_used = tables.get("bins_used")
if bins_used is not None:
print(f"(bins used: {bins_used})")
print()
if summary["system_2"].isna().all():
r = summary.iloc[0]
if is_replica_mode:
print(
f"Single system ({r['system_1']}): intra-replica variability = "
f"{r['intra_1_mean']:.4f} [{r['intra_1_ci_low']:.4f}, {r['intra_1_ci_high']:.4f}]"
)
else:
print(f"Single system ({r['system_1']}): no pairwise comparison to show.")
print()
elif is_replica_mode:
headers = ["System pair", "Inter JSD [95% CI]", "Magnitude (% of ln2)", "Verdict"]
rows = []
for _, r in summary.iterrows():
if r["verdict"] == "distinct":
verdict_txt = "DISTINGUISHABLE"
elif r["verdict"] == "similar":
verdict_txt = "NOT distinguishable"
elif r["verdict"] == "ambiguous":
verdict_txt = "AMBIGUOUS"
else:
verdict_txt = str(r["verdict"])
pair = f"{r['system_1']} vs {r['system_2']}"
ci_txt = f"{r['inter_mean']:.4f} [{r['inter_ci_low']:.4f}, {r['inter_ci_high']:.4f}]"
frac = r["inter_frac_of_max"]
frac_txt = (
f"{r['jsd_magnitude']} ({frac:.1%})" if frac is not None else "n/a"
)
rows.append((pair, ci_txt, frac_txt, verdict_txt))
print(_ascii_table(headers, rows))
print()
print(
"Reading guide: 'DISTINGUISHABLE' means the inter-system JSD CI does "
"NOT overlap either system's intra-replica (replica-to-replica) CI --\n"
"the systems differ beyond sampling noise. 'NOT distinguishable' means "
"intra-replica noise covers the inter-system signal (systems could be\n"
"identical given how much replicas already vary among themselves). "
"'AMBIGUOUS' overlaps exactly one of the two intra CIs."
)
print()
else:
headers = ["System pair", "JSD [95% CI]", "Magnitude (% of ln2)"]
rows = []
for _, r in summary.iterrows():
pair = f"{r['system_1']} vs {r['system_2']}"
ci_txt = f"{r['inter_mean']:.4f} [{r['inter_ci_low']:.4f}, {r['inter_ci_high']:.4f}]"
frac = r["inter_frac_of_max"]
frac_txt = (
f"{r['jsd_magnitude']} ({frac:.1%})" if frac is not None else "n/a"
)
rows.append((pair, ci_txt, frac_txt))
print(_ascii_table(headers, rows))
print()
print(
"Reading guide (pooled mode -- group_by='system'): this is the raw JSD "
"between systems with all replicas fused together. There is no\n"
"intra-system spread to compare against here, so treat the magnitude "
"label as absolute, not as a statistically tested verdict. Re-run with\n"
"group_by='replica' (the default) for a verdict that accounts for "
"replica-to-replica noise."
)
print()
if "replica_outliers" in tables:
flagged = tables["replica_outliers"]
flagged = flagged[flagged["assessment"] == "worth_checking"].copy()
if len(flagged):
flagged["abs_z"] = flagged["deviation_score"].abs()
flagged = flagged.sort_values("abs_z", ascending=False)
if top_outliers is not None:
flagged = flagged.head(top_outliers)
print("Replicas worth checking (deviate from their sisters):")
headers = ["Replica", "System", "z-score"]
rows = [
(str(r["replica"]), str(r["system"]), f"{r['deviation_score']:.2f}")
for _, r in flagged.iterrows()
]
print(_ascii_table(headers, rows))
print()
print("--- Table legend (see the dict for full detail) ---")
headers = ["Key", "What it contains"]
rows = [
(k, v) for k, v in tables.get("__description__", {}).items()
]
if rows:
print(_ascii_table(headers, rows))
# --------------------------------------------------------------------------- #
# Heatmap #
# --------------------------------------------------------------------------- #
def plot_jsd_heatmap(
matrix_or_scores,
labels,
system_of=None,
confidence_level=0.95,
cmap="seismic",
title=None,
cbar_label="JSD",
context="notebook",
triangular=None,
annot_fontsize=8.5,
ax=None,
figsize=None,
fmt="{:.3f}",
vmin=0.0,
vmax=None,
show_values=True,
):
"""Heatmap of JSD scores between groups (replicas, or systems in
pooled mode).
When `show_values=True`, each cell shows the mean score and its
bootstrap confidence interval (for bootstrapped scores). When
`show_values=False`, the annotations are hidden and the figure is
slightly compacted.
Returns a Matplotlib `Axes` object (`tables["heatmap_ax"]` when
called via `compute_jsd`), so it can be further customized, resized,
or saved after the call (e.g. `ax.figure.savefig(...)`).
"""
if vmax is None:
vmax = math.log(2)
is_boot = np.ndim(matrix_or_scores) == 3
mean_matrix = (
matrix_or_scores.mean(axis=2) if is_boot else np.asarray(matrix_or_scores)
).copy()
n = mean_matrix.shape[0]
np.fill_diagonal(mean_matrix, 0.0)
annot = None
if show_values:
annot = np.empty((n, n), dtype=object)
for i in range(n):
for j in range(n):
if i == j:
annot[i, j] = ""
elif is_boot:
lo, hi = percentile_ci(matrix_or_scores[i, j], confidence_level)
annot[i, j] = (
f"{fmt.format(mean_matrix[i, j])}\n[{fmt.format(lo)}, {fmt.format(hi)}]"
)
else:
annot[i, j] = fmt.format(mean_matrix[i, j])
mask = None
if triangular == "upper":
mask = np.tril(np.ones_like(mean_matrix, dtype=bool), k=-1)
elif triangular == "lower":
mask = np.triu(np.ones_like(mean_matrix, dtype=bool), k=1)
if figsize is None:
if show_values:
figsize = (max(1.1 * n + 3, 7), max(1.0 * n + 2.5, 6))
else:
figsize = (max(0.9 * n + 2.2, 5.2), max(0.8 * n + 2.0, 4.8))
with sns.plotting_context(context):
created_fig = ax is None
if created_fig:
fig, ax = plt.subplots(figsize=figsize)
heatmap_kwargs = dict(
mask=mask,
cmap=cmap,
annot=annot if show_values else False,
fmt="",
xticklabels=labels,
yticklabels=labels,
square=True,
linewidths=0.5,
linecolor="white",
cbar_kws={"label": cbar_label},
ax=ax,
vmin=vmin,
vmax=vmax,
)
if show_values:
heatmap_kwargs["annot_kws"] = {"size": annot_fontsize}
sns.heatmap(
mean_matrix,
**heatmap_kwargs,
)
ax.set_title(title if title is not None else "JSD score matrix")
ax.set_xticklabels(ax.get_xticklabels(), rotation=45, ha="right")
ax.set_yticklabels(ax.get_yticklabels(), rotation=0)
if system_of is not None:
boundaries = [
i for i in range(1, len(system_of)) if system_of[i] != system_of[i - 1]
]
for b in boundaries:
ax.axhline(b, color="black", lw=2)
ax.axvline(b, color="black", lw=2)
if created_fig:
fig.tight_layout()
return ax
# --------------------------------------------------------------------------- #
# Main entry point #
# --------------------------------------------------------------------------- #
[docs]
def compute_jsd(
df,
sim_name_col,
replica_col,
feature_col,
group_by="replica",
is_angle=False,
angle_unit="auto",
bins="auto",
iters=200,
frac=1.0,
replace=True,
seed=42,
confidence_level=0.95,
ratio_bootstrap=5000,
ratio_seed=0,
outlier_z_threshold=1.2,
plot_heatmap=False,
heatmap_kwargs=None,
show_values=True,
verbose=True,
explain=True,
explain_kwargs=None,
):
"""
Compare N systems x M replicas on an already-computed feature (1 or 2
columns), via the Jensen-Shannon Divergence (JSD), with bootstrap
confidence intervals. No hypothesis test / p-value: only bootstrap
means, percentile CIs, and descriptive magnitude/verdict heuristics.
Parameters
----------
df : pandas.DataFrame
"Long" dataframe (one row per frame), with the sim_name_col /
replica_col columns and the feature column(s).
sim_name_col : str
Name of the column identifying the system (e.g. "A", "B").
replica_col : str
Name of the column identifying the replica, unique per system.
Only used when group_by="replica" (the default); ignored (but
still required) when group_by="system".
feature_col : str or (str, str)
Name of the column to compare (1D JSD), or a tuple of 2 column
names for a joint 2D JSD (e.g. ("PC1", "PC2")).
group_by : {"replica", "system"}, optional
"replica" (default): the rigorous mode. Groups by `replica_col`,
so intra-system (replica-to-replica) variability is computed and
used as the noise floor against which the inter-system JSD is
judged (see `summarize_systems` -- produces `verdict`,
`effect_magnitude`, and the full replica-level detail tables).
"system": a quicker, pooled mode. Groups by `sim_name_col`
directly -- all replicas of a system are fused into one series
before histogramming. Produces only a direct inter-system JSD
(see `summarize_systems_pooled`); no intra-system spread, no
verdict, no replica-level detail tables (replica_pairs/
replica_outliers/system_overall are omitted). Useful for a fast
overview, but more optimistic than "replica" mode: it ignores
that some replicas may diverge from each other.
is_angle : bool or (bool, bool), optional
True if the column is an angle/dihedral (circular treatment:
fixed JSD range [-pi, pi], after conversion to radians if
needed). Tuple (bool, bool) to specify each column independently
in 2D. Default False.
angle_unit : str or (str, str), optional
"auto" (automatic degrees/radians detection from value
amplitude), "deg" or "rad" to force. Default "auto".
bins : str or int or (int, int), optional
"auto"/"sqrt"/"sturges", a fixed integer, or (bins_x, bins_y) in
2D. Default "auto".
iters : int, optional
Number of bootstrap draws. Default 200.
frac : float, optional
Fraction of each group resampled at each draw. Default 1.0.
replace : bool, optional
Resample with replacement. Default True.
seed : int, optional
Seed for the main bootstrap. Default 42.
confidence_level : float, optional
Confidence level for percentile CIs. Default 0.95.
ratio_bootstrap : int, optional
Number of draws for the inter/intra ratio CI. Only used in
group_by="replica" mode. Default 5000.
ratio_seed : int, optional
Seed for the ratio bootstrap. Default 0.
outlier_z_threshold : float, optional
`deviation_score` threshold above which a replica is flagged
"worth_checking". Only used in group_by="replica" mode.
Default 1.2.
plot_heatmap : bool, optional
If True, plot the JSD score heatmap (added to the return dict
under the "heatmap_ax" key). Default False.
heatmap_kwargs : dict, optional
Extra arguments passed to `plot_jsd_heatmap` (e.g.
triangular="lower", context="talk", ax=...).
show_values : bool, optional
If True, display the mean JSD value and bootstrap confidence
interval in each heatmap cell. If False, hide the annotations for
a more compact heatmap.
verbose : bool, optional
Show progress (detected angle unit, resolved bins, bootstrap
progress bar if tqdm is installed). Default True.
explain : bool, optional
If True (default), automatically call `explain_jsd(tables)` at
the end and print a formatted, human-readable interpretation
(verdicts / magnitudes / outlier replicas / table legend) right
below the log output and the heatmap. Set to False for a silent
call that only returns the dict (e.g. in a loop over many
features). Set to True and use `explain_kwargs={"top_outliers": N}`
to only show the N most extreme outlier replicas.
explain_kwargs : dict, optional
Extra arguments passed to `explain_jsd` when `explain=True`
(currently just `top_outliers`).
Returns
-------
tables : dict
- "system_summary" : system-vs-system comparison. In
group_by="replica" mode: intra/inter/verdict/effect size (see
`summarize_systems`). In group_by="system" mode: pooled
inter-system JSD only (see `summarize_systems_pooled`).
- "replica_pairs" : replica-vs-replica comparison within each
system (group_by="replica" only).
- "replica_outliers" : per-replica heuristic (deviation_score/
assessment) (group_by="replica" only).
- "system_overall" : aggregated verdict across features, only
if >= 2 systems (group_by="replica" only).
- "heatmap_ax" : heatmap Axes (only if plot_heatmap=True).
- "__description__" : description of each table above.
See also
--------
explain_jsd(tables) : called automatically when explain=True (the
default) -- prints a formatted, human-readable interpretation of
this dict (quick read; the dict itself remains available for
full detail / further analysis / plot tweaking). Can also be
called manually later, e.g. after re-running with different
heatmap_kwargs, or on a tables dict obtained with explain=False.
"""
if group_by not in ("replica", "system"):
raise ValueError(f"group_by must be 'replica' or 'system', got {group_by!r}")
is_2d = isinstance(feature_col, (tuple, list))
feature_cols = list(feature_col) if is_2d else [feature_col]
angle_flags = (
list(is_angle)
if isinstance(is_angle, (tuple, list))
else [bool(is_angle)] * len(feature_cols)
)
angle_units = (
list(angle_unit)
if isinstance(angle_unit, (tuple, list))
else [angle_unit] * len(feature_cols)
)
feature_name = "_vs_".join(feature_cols) if is_2d else feature_cols[0]
if verbose:
logger.info(
"=== %s (%s) | metric=jsd | group_by=%s ===",
feature_name,
", ".join(feature_cols),
group_by,
)
effective_group_col = sim_name_col if group_by == "system" else replica_col
feature_matrices, labels, system_indices, circular, value_range = (
prepare_feature_matrices(
df,
sim_name_col,
effective_group_col,
feature_cols,
angle_flags,
angle_units,
verbose=verbose,
)
)
sizes = [f.shape[0] for f in feature_matrices]
if is_2d:
resolved_bins = resolve_bins_for_axes(bins, sizes)
else:
resolved_bins = resolve_num_bins(bins, sizes)
if verbose:
logger.info("bins: %r -> %s", bins, resolved_bins)
circ = tuple(circular) if is_2d else circular[0]
vrange = tuple(value_range) if is_2d else value_range[0]
scores = bootstrap_score_matrix(
feature_matrices,
circular=circ,
bins=resolved_bins,
iters=iters,
frac=frac,
replace=replace,
seed=seed,
is_2d=is_2d,
value_range=vrange,
verbose=verbose,
progress_label=feature_name,
)
if group_by == "system":
summary = summarize_systems_pooled(
scores,
system_indices,
confidence_level=confidence_level,
feature_name=feature_name,
metric_max=math.log(2),
)
tables = {
"system_summary": summary,
"bins_used": resolved_bins,
"__description__": {
"system_summary": "Pooled mode (group_by='system'): direct JSD "
"between systems (all replicas fused together, one series per "
"system). No intra-system spread, so no distinct/similar verdict "
"-- read inter_mean/jsd_magnitude as an absolute magnitude, not a "
"statistical test result. Switch to group_by='replica' (default) "
"for a verdict that accounts for replica-to-replica noise.",
"bins_used": "Number of histogram bins actually used to compute "
"the JSD (resolved from the `bins` argument -- see resolve_num_bins "
"/ resolve_bins_for_axes). Also shown in the heatmap title and in "
"the log when verbose=True.",
},
}
else:
summary = summarize_systems(
scores,
system_indices,
confidence_level=confidence_level,
feature_name=feature_name,
metric_max=math.log(2),
ratio_bootstrap=ratio_bootstrap,
ratio_seed=ratio_seed,
)
pair_summary = replica_pair_summary(
scores,
system_indices,
labels,
confidence_level=confidence_level,
feature_name=feature_name,
)
outlier_summary = per_replica_outlier_summary(
pair_summary, z_threshold=outlier_z_threshold
)
tables = {
"system_summary": summary,
"replica_pairs": pair_summary,
"replica_outliers": outlier_summary,
"bins_used": resolved_bins,
}
if len(system_indices) >= 2:
tables["system_overall"] = aggregate_verdicts([summary])
tables["__description__"] = {
"system_summary": "SYSTEM-level comparison: intra-1, intra-2, inter "
"(mean+CI), verdict (distinct/ambiguous/similar, CI non-overlap), plus "
"effect_ratio_*/effect_magnitude and inter_frac_of_max/jsd_magnitude "
"(absolute magnitude on the JSD's own [0, ln(2)] scale).",
"replica_pairs": "REPLICA-level comparison, within each system: "
"mean+CI per replica pair -- the detail behind 'intra' in "
"system_summary.",
"replica_outliers": "Per REPLICA (not per pair): avg_pairwise_score, "
"deviation_score (z-score-like vs its sisters), assessment "
"(typical/worth_checking).",
"system_overall": "Cross-feature summary: for each pair of systems, "
"how many features say 'distinct', and an overall verdict.",
"bins_used": "Number of histogram bins actually used to compute the "
"JSD (resolved from the `bins` argument -- see resolve_num_bins / "
"resolve_bins_for_axes). Also shown in the heatmap title and in the "
"log when verbose=True.",
}
if plot_heatmap:
heatmap_kwargs = dict(heatmap_kwargs or {})
heatmap_kwargs.setdefault(
"title", f"{feature_name} (JSD, bins={resolved_bins}, group_by={group_by})"
)
heatmap_kwargs.setdefault("show_values", show_values)
heatmap_kwargs.setdefault("cbar_label", "JSD")
heatmap_kwargs.setdefault("vmin", 0.0)
heatmap_kwargs.setdefault("vmax", math.log(2))
system_of_list = [None] * len(labels)
for s, idxs in system_indices.items():
for i in idxs:
system_of_list[i] = s
ax = plot_jsd_heatmap(
scores,
labels,
system_of=system_of_list if group_by == "replica" else None,
confidence_level=confidence_level,
**heatmap_kwargs,
)
tables["heatmap_ax"] = ax
if explain:
explain_jsd(tables, **(explain_kwargs or {}))
return tables