lars
a14a4f973a
Whitespace-only reflow (line wrapping, blank lines between defs); no logic changes. Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com>
70 lines
2.2 KiB
Python
70 lines
2.2 KiB
Python
"""Per-step marginal KL(real||gen) per target dim, 10 vs 20 ODE steps.
|
|
|
|
Streaming histogram over shared bin edges (computed from the true distribution),
|
|
so the two runs are directly comparable dim-by-dim.
|
|
"""
|
|
|
|
import numpy as np
|
|
import polars as pl
|
|
|
|
from giant.analysis import RAW_TARGET_NAMES, _kl_from_counts, _raw_dim_expr
|
|
|
|
FILES = {
|
|
"10-step": "/home/lars/Programming/giant/9879e806-5e88-4b06-b1fa-0e61de9cda6f.parquet",
|
|
"20-step": "/home/lars/Programming/giant/0e236919-65ae-4b9b-9957-d31a8211aca4.parquet",
|
|
}
|
|
BINS = 100
|
|
|
|
# Fixed shared edges from the true distribution (identical across both files), using
|
|
# robust quantiles to avoid a few outliers dominating the range.
|
|
base = FILES["10-step"]
|
|
edges = {}
|
|
for j, name in enumerate(RAW_TARGET_NAMES):
|
|
lo, hi = (
|
|
pl.scan_parquet(base)
|
|
.select(
|
|
_raw_dim_expr("true", j).quantile(0.001).alias("lo"),
|
|
_raw_dim_expr("true", j).quantile(0.999).alias("hi"),
|
|
)
|
|
.collect(engine="streaming")
|
|
.row(0)
|
|
)
|
|
if not (hi - lo > 1e-9):
|
|
lo, hi = lo - 0.5, hi + 0.5
|
|
edges[name] = np.linspace(lo, hi, BINS + 1)
|
|
|
|
|
|
def counts(file, prefix, j, e):
|
|
vals = (
|
|
pl.scan_parquet(file)
|
|
.select(_raw_dim_expr(prefix, j).alias("v"))
|
|
.collect(engine="streaming")["v"]
|
|
.to_numpy()
|
|
)
|
|
c, _ = np.histogram(vals, bins=e)
|
|
return c
|
|
|
|
|
|
kls = {name: {} for name in FILES}
|
|
for j, name in enumerate(RAW_TARGET_NAMES):
|
|
e = edges[name]
|
|
real_c = counts(base, "true", j, e) # identical true dist across files
|
|
for run, f in FILES.items():
|
|
gen_c = counts(f, "pred", j, e)
|
|
kls[run][name] = _kl_from_counts(real_c, gen_c)
|
|
|
|
print(f"{'dim':<14}{'KL 10-step':>14}{'KL 20-step':>14}{'ratio 20/10':>14}")
|
|
print("-" * 56)
|
|
tot = {"10-step": 0.0, "20-step": 0.0}
|
|
for name in RAW_TARGET_NAMES:
|
|
a, b = kls["10-step"][name], kls["20-step"][name]
|
|
tot["10-step"] += a
|
|
tot["20-step"] += b
|
|
print(f"{name:<14}{a:>14.5f}{b:>14.5f}{b / a if a else float('nan'):>14.2f}")
|
|
print("-" * 56)
|
|
print(
|
|
f"{'SUM':<14}{tot['10-step']:>14.5f}{tot['20-step']:>14.5f}"
|
|
f"{tot['20-step'] / tot['10-step']:>14.2f}"
|
|
)
|
|
print(f"{'MEAN':<14}{tot['10-step'] / 9:>14.5f}{tot['20-step'] / 9:>14.5f}")
|