Compare commits
3 Commits
| Author | SHA1 | Date | |
|---|---|---|---|
| 313373cc10 | |||
| 98b09d2b4b | |||
| d0381728ee |
@@ -29,8 +29,17 @@ from giant.constants import TERM_ESCAPED
|
||||
|
||||
|
||||
def _bin_expr(value: pl.Expr, lo: float, hi: float, nbins: int) -> pl.Expr:
|
||||
"""Uniform bin index of ``value`` over ``[lo, hi]`` into ``nbins`` bins."""
|
||||
return ((value - lo) / (hi - lo) * nbins).floor().cast(pl.Int32).clip(0, nbins - 1)
|
||||
"""Uniform bin index of ``value`` over ``[lo, hi]`` into ``nbins`` bins.
|
||||
|
||||
Out-of-range values clamp into the edge bins, and the clamp deliberately
|
||||
happens in f64 *before* the integer cast: a rollout is free to emit a wildly
|
||||
out-of-range outlier (a step_length of 1e10 mm, say) or an inf, whose
|
||||
unclamped bin index overflows i32 and makes the cast fail outright. NaN has
|
||||
no edge to clamp to, so it becomes null and is dropped by the callers below
|
||||
— the same thing ``np.histogram`` does with it.
|
||||
"""
|
||||
idx = ((value - lo) / (hi - lo) * nbins).floor().clip(0, nbins - 1)
|
||||
return pl.when(idx.is_nan()).then(None).otherwise(idx).cast(pl.Int32)
|
||||
|
||||
|
||||
def hist1d(
|
||||
@@ -50,6 +59,7 @@ def hist1d(
|
||||
group = pl.lit(0, dtype=pl.Int64) if group is None else group
|
||||
res = (
|
||||
lf.select(group.alias("_g"), _bin_expr(value, lo, hi, nbins).alias("_b"))
|
||||
.drop_nulls("_b")
|
||||
.group_by("_g", "_b")
|
||||
.agg(pl.len().alias("_n"))
|
||||
.collect(engine="streaming")
|
||||
@@ -190,6 +200,7 @@ def profile_partial(
|
||||
_bin_expr(coord, lo, hi, nbins).alias("_b"),
|
||||
weight.alias("_w"),
|
||||
)
|
||||
.drop_nulls("_b")
|
||||
.group_by("event_id", "_b")
|
||||
.agg(pl.col("_w").sum().alias("_ws"))
|
||||
.collect(engine="streaming")
|
||||
|
||||
@@ -10,6 +10,14 @@ _EPS = 1e-8
|
||||
# the conservation it slightly softens is physically negligible (~0.001%).
|
||||
_SIMPLEX_FLOOR = 1e-5
|
||||
|
||||
# Upper clip for a raw predicted log_mass before inv_log_transform: exp(y)
|
||||
# must stay well inside float32 range (~3.4e38, i.e. y < ~88.7) or it
|
||||
# overflows to inf, which — like the negative-mass case below — blows up the
|
||||
# next log_transform call once that mass is fed back in as conditioning.
|
||||
# 80.0 leaves comfortable headroom while still being far beyond any physical
|
||||
# particle mass a converged model would ever predict.
|
||||
_LOG_MASS_MAX = 80.0
|
||||
|
||||
|
||||
def log_transform(x: np.ndarray, eps: float = _EPS) -> np.ndarray:
|
||||
x = np.asarray(x, dtype=np.float32)
|
||||
@@ -629,9 +637,18 @@ def decode_secondaries(
|
||||
pre_dir[valid], dir_local[valid, i]
|
||||
)
|
||||
|
||||
# mass is non-negative by construction (inv_log_transform of a real
|
||||
# number is always > 0); clip to 0 for padded/invalid slots rather than
|
||||
# leaving a spurious small positive floor from the log inverse.
|
||||
# log_mass is a raw model prediction, not itself the output of
|
||||
# log_transform, so it can land far outside the range that round-trips
|
||||
# cleanly through inv_log_transform: too negative and exp(log_mass)
|
||||
# undershoots _EPS, making inv_log_transform go slightly negative; too
|
||||
# positive and exp(log_mass) overflows float32 to inf. Either one then
|
||||
# blows up the next log_transform call on this track's mass once it's
|
||||
# fed back in as conditioning for a further rollout step
|
||||
# (giant/rollout.py -> build_cond_features -> _physical_cond_columns).
|
||||
# Clip log_mass to a range whose inverse is guaranteed finite and >= 0
|
||||
# before that can happen; clip to 0 separately for padded/invalid slots
|
||||
# rather than leaving a spurious small positive floor.
|
||||
log_mass = np.clip(log_mass, np.log(_EPS), _LOG_MASS_MAX)
|
||||
sec_mass = np.where(sec_valid, inv_log_transform(log_mass), 0.0).astype(np.float32)
|
||||
sec_charge = np.where(sec_valid, charge, 0.0).astype(np.float32)
|
||||
|
||||
|
||||
@@ -102,6 +102,33 @@ def test_hist1d_overall_and_grouped():
|
||||
assert hg[11].sum() == 4
|
||||
|
||||
|
||||
def test_hist1d_clamps_extreme_values_and_drops_nan():
|
||||
# A rollout can emit a wildly out-of-range step_length (or an inf/NaN); the
|
||||
# fixed-edge binning must clamp rather than overflow the i32 bin cast.
|
||||
lf = pl.DataFrame(
|
||||
{"x": [5.0, 1.0725e10, float("inf"), -float("inf"), float("nan"), None]}
|
||||
).lazy()
|
||||
edges = np.linspace(0.0, 50.0, 6) # width 10
|
||||
h = R.hist1d(lf, pl.col("x"), edges)
|
||||
# 5 -> bin 0; 1e10 and +inf -> top bin; -inf -> bin 0; NaN/null dropped
|
||||
assert h[0].tolist() == [2, 0, 0, 0, 2]
|
||||
|
||||
|
||||
def test_profile_partial_clamps_extreme_values_and_drops_nan():
|
||||
lf = pl.DataFrame(
|
||||
{
|
||||
"event_id": [1, 1, 1, 1],
|
||||
"z": [5.0, 1.0725e10, float("nan"), 45.0],
|
||||
"w": [1.0, 2.0, 4.0, 8.0],
|
||||
}
|
||||
).lazy()
|
||||
edges = np.linspace(0.0, 50.0, 6)
|
||||
ev, mat = R.profile_partial(lf, pl.col("z"), edges, pl.col("w"))
|
||||
assert ev.tolist() == [1]
|
||||
# 1e10 clamps into the top bin alongside 45; the NaN row's weight is dropped
|
||||
assert mat[0].tolist() == [1.0, 0.0, 0.0, 0.0, 10.0]
|
||||
|
||||
|
||||
def test_physical_steps_drops_synthetic_rollout_rows_only():
|
||||
lf = _rollout_frame()
|
||||
phys = physical_steps(lf, Side.rollout).collect()
|
||||
|
||||
@@ -540,3 +540,53 @@ def test_decode_secondaries_mass_charge_round_trip_with_normalizer():
|
||||
)
|
||||
assert sec_mass[0, 0] == pytest.approx(938.27208943, abs=1e-2)
|
||||
assert sec_charge[0, 0] == pytest.approx(1.0, abs=1e-4)
|
||||
|
||||
|
||||
def test_decode_secondaries_extreme_negative_log_mass_stays_nonnegative():
|
||||
from giant.data.transforms import decode_secondaries, log_transform
|
||||
|
||||
N = 1
|
||||
e_sec = np.array([5.0], dtype=np.float32)
|
||||
n_sec = np.array([1])
|
||||
pre_dir = np.array([[0.0, 0.0, 1.0]], dtype=np.float32)
|
||||
|
||||
sec_cont = np.zeros((N, K_MAX, 6), dtype=np.float32)
|
||||
sec_cont[0, 0, 0] = 10.0 # stick logit -> ~all of e_sec
|
||||
sec_cont[0, 0, 1:4] = [0, 0, 1]
|
||||
sec_cont[0, 0, 4] = -50.0 # raw model prediction: extremely negative log_mass
|
||||
sec_cont[0, 0, 5] = 1.0
|
||||
|
||||
_, _, sec_mass, _, _ = decode_secondaries(sec_cont, n_sec, e_sec, pre_dir)
|
||||
|
||||
# A raw model prediction isn't itself the output of log_transform, so
|
||||
# naively applying inv_log_transform can undershoot zero (see
|
||||
# decode_secondaries) — which then crashes the next log_transform call
|
||||
# once this mass is fed back in as conditioning during rollout. The
|
||||
# float32 residual from clipping can land a hair below zero, but must
|
||||
# stay well above -eps so log_transform(mass) stays finite.
|
||||
assert sec_mass[0, 0] > -1e-8
|
||||
log_transform(sec_mass[0, 0])
|
||||
|
||||
|
||||
def test_decode_secondaries_extreme_positive_log_mass_stays_finite():
|
||||
from giant.data.transforms import decode_secondaries, log_transform
|
||||
|
||||
N = 1
|
||||
e_sec = np.array([5.0], dtype=np.float32)
|
||||
n_sec = np.array([1])
|
||||
pre_dir = np.array([[0.0, 0.0, 1.0]], dtype=np.float32)
|
||||
|
||||
sec_cont = np.zeros((N, K_MAX, 6), dtype=np.float32)
|
||||
sec_cont[0, 0, 0] = 10.0 # stick logit -> ~all of e_sec
|
||||
sec_cont[0, 0, 1:4] = [0, 0, 1]
|
||||
sec_cont[0, 0, 4] = 200.0 # raw model prediction: extremely positive log_mass
|
||||
sec_cont[0, 0, 5] = 1.0
|
||||
|
||||
_, _, sec_mass, _, _ = decode_secondaries(sec_cont, n_sec, e_sec, pre_dir)
|
||||
|
||||
# Mirror image of the extreme-negative case above: exp(log_mass)
|
||||
# overflows float32 to inf for an unclipped raw prediction this large,
|
||||
# which then crashes the next log_transform call the same way a
|
||||
# negative mass would.
|
||||
assert np.isfinite(sec_mass[0, 0])
|
||||
log_transform(sec_mass[0, 0])
|
||||
|
||||
Reference in New Issue
Block a user