import logging
import warnings
import numpy as np
import pandas as pd
import matplotlib.pyplot as plt
import seaborn as sns
from scipy.stats import wasserstein_distance
warnings.filterwarnings("ignore", category=RuntimeWarning)
logging.basicConfig(level=logging.INFO, format="%(levelname)s - %(message)s")
logger = logging.getLogger(__name__)
# --------------------------------------------------------------------------- #
# Wasserstein math (1D only; circular approximation via grid search) #
# --------------------------------------------------------------------------- #
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)),
)
def _circular_wasserstein_1d(a, b, n_shifts=180):
"""
Approximates the 1D Wasserstein distance between two circular
(angular, radians) samples: resamples both to the same size, grid
searches over rotations of `b` for the minimal linear Wasserstein
distance -- a standard trick to approximate optimal circular
transport in 1D (the classic Wasserstein distance is not rotation
invariant around the +-pi wrap-around point).
"""
n = min(a.shape[0], b.shape[0])
rng = np.random.default_rng(0)
a_s = rng.choice(a, n, replace=(a.shape[0] < n))
b_s = rng.choice(b, n, replace=(b.shape[0] < n))
shifts = np.linspace(-np.pi, np.pi, n_shifts, endpoint=False)
best = np.inf
for shift in shifts:
b_shifted = np.mod(b_s + shift + np.pi, 2 * np.pi) - np.pi
d = wasserstein_distance(a_s, b_shifted)
if d < best:
best = d
return float(best)
def wasserstein_score(f1, f2, circular, bins=None, value_range=None, is_2d=False):
"""`bins`/`value_range` are ignored (Wasserstein does not rely on a
histogram -- it is computed exactly from the raw samples via optimal
transport) -- kept in the signature so this stays interchangeable
with `jsd_score`."""
if is_2d:
raise NotImplementedError(
"2D Wasserstein is not implemented; use compute_jsd for a 2-column "
"feature, or plug in scipy.stats.wasserstein_distance_nd/POT."
)
if circular:
return _circular_wasserstein_1d(f1[:, 0], f2[:, 0])
return float(wasserstein_distance(f1[:, 0], f2[:, 0]))
# --------------------------------------------------------------------------- #
# 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 (F=1 -- 2D Wasserstein is not
implemented), 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 computing the distance).
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.
"""
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 = []
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)
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)
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
# --------------------------------------------------------------------------- #
# Bootstrap engine #
# --------------------------------------------------------------------------- #
def bootstrap_score_matrix(
feature_matrices,
circular,
iters,
frac,
replace,
seed,
verbose=False,
progress_label="",
):
"""
Bootstrap the Wasserstein 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 Wasserstein {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 = wasserstein_score(fi[idx_i], fj[idx_j], circular=circular, is_2d=False)
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.
"""
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.
"""
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 summarize_systems(
scores,
system_indices,
confidence_level=0.95,
feature_name="",
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"/"ambiguous"/"similar"), plus `effect_ratio_*`/
`effect_magnitude` (inter / max(intra_1, intra_2)).
Unlike the JSD, the Wasserstein distance has no fixed upper bound: no
`inter_frac_of_max`/absolute-magnitude column here (kept as None/
"n/a" so the column schema stays comparable with `compute_jsd`'s
output).
Single-system input -> intra-only row.
"""
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,
"wasserstein_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)
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": None,
"wasserstein_magnitude": "n/a",
}
)
return pd.DataFrame(rows)
def summarize_systems_pooled(
scores,
system_indices,
confidence_level=0.95,
feature_name="",
):
"""
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 Wasserstein distance (mean + bootstrap
CI, where the bootstrap only captures 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 distance 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,
}
]
)
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)
rows.append(
{
"feature": feature_name,
"system_1": g1,
"system_2": g2,
"inter_mean": mean,
"inter_ci_low": ci[0],
"inter_ci_high": ci[1],
}
)
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. 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"). See
`compute_jsd.per_replica_outlier_summary` for the note on the
ceiling effect with few replicas.
"""
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` into a cross-feature
summary. 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_wasserstein(tables, top_outliers=None):
"""
Print a quick, human-readable interpretation of the dict returned by
`compute_wasserstein`, 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
distance (no verdict, since there is no intra-system spread left
to compare against).
Note on units: unlike the JSD (bounded in [0, ln(2)]), the
Wasserstein distance is expressed in the feature's own physical
units and is NOT bounded -- there is no universal "large/moderate/
negligible" scale for it. Judge the absolute inter_mean value against
what a meaningful difference looks like for that specific feature
(e.g. Angstroms for a distance, kcal/mol for an energy).
Parameters
----------
tables : dict
Return value of `compute_wasserstein`.
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 (Wasserstein) "
print("=" * max(len(title), 60))
print(title.center(max(len(title), 60)))
print("=" * max(len(title), 60))
print("(exact optimal-transport distance -- no histogram binning involved)")
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 distance [95% CI]", "Effect ratio (95% CI)", "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}]"
ratio_txt = (
f"{r['effect_ratio_mean']:.2f} [{r['effect_ratio_ci_low']:.2f}, "
f"{r['effect_ratio_ci_high']:.2f}]"
if r["effect_ratio_mean"] is not None and not pd.isna(r["effect_ratio_mean"])
else "n/a"
)
rows.append((pair, ci_txt, ratio_txt, verdict_txt))
print(_ascii_table(headers, rows))
print()
print(
"Reading guide: 'DISTINGUISHABLE' means the inter-system distance 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. 'AMBIGUOUS' overlaps\n"
"exactly one of the two intra CIs. The effect ratio is "
"inter_mean / max(intra_1_mean, intra_2_mean) -- how many times larger the\n"
"inter-system distance is than the noisier of the two systems' own "
"replica-to-replica spread. Values are in the feature's native units, not\n"
"a bounded [0,1]-style scale -- judge magnitude against what matters "
"physically for this feature."
)
print()
else:
headers = ["System pair", "Distance [95% CI]"]
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}]"
rows.append((pair, ci_txt))
print(_ascii_table(headers, rows))
print()
print(
"Reading guide (pooled mode -- group_by='system'): this is the raw "
"Wasserstein distance between systems with all replicas fused together.\n"
"There is no intra-system spread to compare against here, so there is no "
"verdict -- judge the absolute inter_mean value against what matters\n"
"physically for this feature. Re-run with 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_wasserstein_heatmap(
matrix_or_scores,
labels,
system_of=None,
confidence_level=0.95,
cmap="seismic",
title=None,
cbar_label="Wasserstein distance",
context="notebook",
triangular=None,
annot_fontsize=8.5,
ax=None,
figsize=None,
fmt="{:.3f}",
vmin=None,
vmax=None,
):
"""Heatmap of Wasserstein scores between groups (replicas, or
systems in pooled mode). No fixed vmin/vmax by default (unbounded
metric, in the feature's own units) -- pass them explicitly if you
want a comparable scale across several heatmaps.
Returns a Matplotlib `Axes` object (`tables["heatmap_ax"]` when
called via `compute_wasserstein`), so it can be further customized,
resized, or saved after the call (e.g. `ax.figure.savefig(...)`).
"""
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 = 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:
figsize = (max(1.1 * n + 3, 7), max(1.0 * n + 2.5, 6))
heatmap_kwargs = {}
if vmin is not None:
heatmap_kwargs["vmin"] = vmin
if vmax is not None:
heatmap_kwargs["vmax"] = vmax
with sns.plotting_context(context):
created_fig = ax is None
if created_fig:
fig, ax = plt.subplots(figsize=figsize)
sns.heatmap(
mean_matrix,
mask=mask,
cmap=cmap,
annot=annot,
fmt="",
annot_kws={"size": annot_fontsize},
xticklabels=labels,
yticklabels=labels,
square=True,
linewidths=0.5,
linecolor="white",
cbar_kws={"label": cbar_label},
ax=ax,
**heatmap_kwargs,
)
ax.set_title(title if title is not None else "Wasserstein 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_wasserstein(
df,
sim_name_col,
replica_col,
feature_col,
group_by="replica",
is_angle=False,
angle_unit="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,
verbose=True,
explain=True,
explain_kwargs=None,
):
"""
Compare N systems x M replicas on an already-computed feature (1
column -- 2D Wasserstein is not implemented, use `compute_jsd` for a
2-column feature), via the (1st order) Wasserstein distance, with
bootstrap confidence intervals. No hypothesis test / p-value.
Unlike the JSD, the Wasserstein distance is NOT bounded: it is
expressed in the feature's own physical units, computed exactly from
the raw samples (no histogram/binning step at all -- there is no
"number of bins" to report or tune for this metric), and has no
universal absolute-magnitude scale (only the inter/intra ratio is
reported, no "jsd_magnitude"-style column).
Parameters
----------
df : pandas.DataFrame
"Long" dataframe (one row per frame), with the sim_name_col /
replica_col columns and the feature column.
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
Name of the column to compare.
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 distance
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 computing the distance. Produces only a direct
inter-system distance (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, optional
True if the column is an angle/dihedral (circular approximation
of the Wasserstein distance via a rotation grid search). Default
False.
angle_unit : str, optional
"auto"/"deg"/"rad". 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 Wasserstein score heatmap (added to the return
dict under the "heatmap_ax" key). Default False.
heatmap_kwargs : dict, optional
Extra arguments passed to `plot_wasserstein_heatmap`.
verbose : bool, optional
Show progress. Default True.
explain : bool, optional
If True (default), automatically call `explain_wasserstein(tables)`
at the end and print a formatted, human-readable interpretation
(verdicts / 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). Use
`explain_kwargs={"top_outliers": N}` to only show the N most
extreme outlier replicas.
explain_kwargs : dict, optional
Extra arguments passed to `explain_wasserstein` 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 ratio (see
`summarize_systems`). In group_by="system" mode: pooled
inter-system distance 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_wasserstein(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}")
if isinstance(feature_col, (tuple, list)):
raise NotImplementedError(
"2D Wasserstein is not implemented; pass a single column name "
"(str), or use compute_jsd for a 2-column feature."
)
feature_cols = [feature_col]
angle_flags = [bool(is_angle)]
angle_units = [angle_unit]
feature_name = feature_col
if verbose:
logger.info(
"=== %s | metric=wasserstein (exact, no binning) | group_by=%s ===",
feature_name,
group_by,
)
effective_group_col = sim_name_col if group_by == "system" else replica_col
feature_matrices, labels, system_indices, circular = prepare_feature_matrices(
df,
sim_name_col,
effective_group_col,
feature_cols,
angle_flags,
angle_units,
verbose=verbose,
)
scores = bootstrap_score_matrix(
feature_matrices,
circular=circular[0],
iters=iters,
frac=frac,
replace=replace,
seed=seed,
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,
)
tables = {
"system_summary": summary,
"__description__": {
"system_summary": "Pooled mode (group_by='system'): direct "
"Wasserstein distance between systems (all replicas fused "
"together, one series per system). No intra-system spread, so no "
"distinct/similar verdict -- read inter_mean as an absolute "
"distance in the feature's own units, not a statistical test "
"result. Switch to group_by='replica' (default) for a verdict "
"that accounts for replica-to-replica noise.",
},
}
else:
summary = summarize_systems(
scores,
system_indices,
confidence_level=confidence_level,
feature_name=feature_name,
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,
}
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 (inter/max(intra_1,intra_2)). No "
"absolute magnitude column (Wasserstein is unbounded).",
"replica_pairs": "REPLICA-level comparison, within each system: "
"mean+CI per replica pair.",
"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.",
}
if plot_heatmap:
heatmap_kwargs = dict(heatmap_kwargs or {})
heatmap_kwargs.setdefault(
"title",
f"{feature_name} (Wasserstein distance, exact -- no binning, "
f"group_by={group_by})",
)
system_of_list = [None] * len(labels)
for s, idxs in system_indices.items():
for i in idxs:
system_of_list[i] = s
ax = plot_wasserstein_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_wasserstein(tables, **(explain_kwargs or {}))
return tables