Encode edep/secondary/post energy as a conservation-constrained simplex
Replaces the independent log_delta_e/log_edep targets with 2 additive-log-ratio coordinates over the deposit/secondary/post-energy simplex (fractions of pre_E summing to 1), so edep + e_sec + post_E == pre_E holds by construction after decoding (softmax) rather than being learned approximately. Requires e_sec (secondary energy) as a new conditioning input and a steps_to_parquet.py pass to derive it from child track first-step energies.
This commit is contained in:
+109
-43
@@ -110,6 +110,7 @@ from giant.data.loader import find_parquet_files, load_steps
|
||||
from giant.data.transforms import (
|
||||
Normalizer,
|
||||
build_features,
|
||||
energy_simplex_decode,
|
||||
inv_log_transform,
|
||||
reconstruct_post_pos,
|
||||
)
|
||||
@@ -118,13 +119,63 @@ from giant.model.schedule import CosineSchedule
|
||||
from giant.sample import sample_ddim, sample_ddpm, sample_flow
|
||||
from giant.validate import _histogram_kl
|
||||
|
||||
_N_LOG_DIMS = 3 # log_step_length, log_delta_e, log_edep are the first 3 target dims
|
||||
RAW_TARGET_NAMES = [n.removeprefix("log_") for n in LOCAL_TARGET_NAMES]
|
||||
# The first 3 target dims are the scalar (non-direction) outputs. In raw/physical
|
||||
# space they are step_length, delta_e, edep; in the model's native target space the
|
||||
# first is log_step_length and the next two are the deposit/secondary ALR energy
|
||||
# logits (see giant.constants.LOCAL_TARGET_NAMES and energy_simplex_encode).
|
||||
_N_SCALAR_DIMS = 3
|
||||
RAW_TARGET_NAMES = [
|
||||
"step_length",
|
||||
"delta_e",
|
||||
"edep",
|
||||
"post_dx",
|
||||
"post_dy",
|
||||
"post_dz",
|
||||
"travel_dx",
|
||||
"travel_dy",
|
||||
"travel_dz",
|
||||
]
|
||||
_LOG_EPS = (
|
||||
1e-8 # mirrors giant.data.transforms._EPS, duplicated for use in polars exprs
|
||||
)
|
||||
|
||||
|
||||
def _decode_raw_targets(target_local: np.ndarray, pre_E: np.ndarray) -> np.ndarray:
|
||||
"""Map a model-native target array (N, 9) to physical raw units.
|
||||
|
||||
Column 0 (log_step_length) is de-logged; columns 1–2 (the deposit/secondary
|
||||
ALR energy logits) are decoded against `pre_E` into physical `delta_e` and
|
||||
`edep` (energy_simplex_decode); the six direction components pass through
|
||||
unchanged. Output columns therefore line up with `RAW_TARGET_NAMES`.
|
||||
"""
|
||||
raw = target_local.astype(np.float32).copy()
|
||||
raw[:, 0] = inv_log_transform(target_local[:, 0])
|
||||
edep, _e_sec, _post_E, delta_e = energy_simplex_decode(target_local[:, 1:3], pre_E)
|
||||
raw[:, 1] = delta_e
|
||||
raw[:, 2] = edep
|
||||
return raw
|
||||
|
||||
|
||||
def _edep_pl(prefix: str) -> pl.Expr:
|
||||
"""Physical edep from a `{prefix}_edep_logit`/`{prefix}_sec_logit` pair + `pre_E`.
|
||||
|
||||
Polars equivalent of `energy_simplex_decode(...)[0]` (the deposit component):
|
||||
a softmax over `[z_edep, z_sec, 0]` times pre_E.
|
||||
"""
|
||||
z1, z2 = pl.col(f"{prefix}_edep_logit"), pl.col(f"{prefix}_sec_logit")
|
||||
m = pl.max_horizontal(z1, z2, pl.lit(0.0))
|
||||
e1, e2, e3 = (z1 - m).exp(), (z2 - m).exp(), (pl.lit(0.0) - m).exp()
|
||||
return (e1 / (e1 + e2 + e3)) * pl.col("pre_E")
|
||||
|
||||
|
||||
def _delta_e_pl(prefix: str) -> pl.Expr:
|
||||
"""Physical delta_e (= edep + e_sec = pre_E - post_E) from the ALR logits + pre_E."""
|
||||
z1, z2 = pl.col(f"{prefix}_edep_logit"), pl.col(f"{prefix}_sec_logit")
|
||||
m = pl.max_horizontal(z1, z2, pl.lit(0.0))
|
||||
e1, e2, e3 = (z1 - m).exp(), (z2 - m).exp(), (pl.lit(0.0) - m).exp()
|
||||
return ((e1 + e2) / (e1 + e2 + e3)) * pl.col("pre_E")
|
||||
|
||||
|
||||
def _hist_edges(*arrays: np.ndarray, bins: int) -> np.ndarray:
|
||||
"""Bin edges that don't blow up on near-constant data (e.g. a tight unit-norm cluster).
|
||||
|
||||
@@ -240,10 +291,11 @@ def make_val_loader(
|
||||
return DataLoader(val_ds, batch_size=batch_size, shuffle=False)
|
||||
|
||||
|
||||
def _to_raw_targets(target_norm: np.ndarray, normalizer: Normalizer) -> np.ndarray:
|
||||
def _to_raw_targets(
|
||||
target_norm: np.ndarray, normalizer: Normalizer, pre_E: np.ndarray
|
||||
) -> np.ndarray:
|
||||
raw = normalizer.inverse_transform(target_norm)
|
||||
raw[:, :_N_LOG_DIMS] = inv_log_transform(raw[:, :_N_LOG_DIMS])
|
||||
return raw
|
||||
return _decode_raw_targets(raw, pre_E)
|
||||
|
||||
|
||||
@torch.no_grad()
|
||||
@@ -298,8 +350,12 @@ def collect_samples(
|
||||
cond_cont_raw=cond_cont_raw,
|
||||
pdg=pdg,
|
||||
material=material,
|
||||
real_raw=_to_raw_targets(real_norm, bundle.target_normalizer),
|
||||
gen_raw=_to_raw_targets(gen_norm, bundle.target_normalizer),
|
||||
real_raw=_to_raw_targets(
|
||||
real_norm, bundle.target_normalizer, cond_cont_raw[:, 3]
|
||||
),
|
||||
gen_raw=_to_raw_targets(
|
||||
gen_norm, bundle.target_normalizer, cond_cont_raw[:, 3]
|
||||
),
|
||||
real_norm=real_norm,
|
||||
gen_norm=gen_norm,
|
||||
)
|
||||
@@ -386,19 +442,15 @@ def load_predicted_local(
|
||||
gen_log_local = df.select(pred_cols).to_numpy().astype(np.float32)
|
||||
real_log_local = df.select(true_cols).to_numpy().astype(np.float32)
|
||||
|
||||
def to_raw(log_local: np.ndarray) -> np.ndarray:
|
||||
raw = log_local.copy()
|
||||
raw[:, :_N_LOG_DIMS] = inv_log_transform(raw[:, :_N_LOG_DIMS])
|
||||
return raw
|
||||
|
||||
cond_cont_raw = df.select(_COND_CONT_COLS).to_numpy().astype(np.float32)
|
||||
pre_E = cond_cont_raw[:, 3] # pre_E is already physical in predict output
|
||||
|
||||
return SampleCollection(
|
||||
cond_cont_raw=cond_cont_raw,
|
||||
pdg=df["pdg"].to_numpy(),
|
||||
material=df["material"].to_numpy(),
|
||||
real_raw=to_raw(real_log_local),
|
||||
gen_raw=to_raw(gen_log_local),
|
||||
real_raw=_decode_raw_targets(real_log_local, pre_E),
|
||||
gen_raw=_decode_raw_targets(gen_log_local, pre_E),
|
||||
)
|
||||
|
||||
|
||||
@@ -499,10 +551,20 @@ def _histogram_kl_pl(
|
||||
return float(np.sum(p_hist * np.log(p_hist / q_hist)))
|
||||
|
||||
|
||||
def _raw_expr(col: str, j: int) -> pl.Expr:
|
||||
"""`pred_*`/`true_*` columns are log-scaled for the first `_N_LOG_DIMS` dims."""
|
||||
e = pl.col(col)
|
||||
return (e.exp() - _LOG_EPS) if j < _N_LOG_DIMS else e
|
||||
def _raw_dim_expr(prefix: str, j: int) -> pl.Expr:
|
||||
"""Physical raw value of target dim j (`RAW_TARGET_NAMES[j]`) from a predict parquet.
|
||||
|
||||
Dim 0 is de-logged step_length; dims 1–2 are the physical `delta_e`/`edep`
|
||||
decoded from the deposit/secondary ALR logits against `pre_E`; dims ≥3 are
|
||||
direction components, used as-is. `prefix` is "true" or "pred".
|
||||
"""
|
||||
if j == 0:
|
||||
return pl.col(f"{prefix}_log_step_length").exp() - _LOG_EPS
|
||||
if j == 1:
|
||||
return _delta_e_pl(prefix)
|
||||
if j == 2:
|
||||
return _edep_pl(prefix)
|
||||
return pl.col(f"{prefix}_{LOCAL_TARGET_NAMES[j]}")
|
||||
|
||||
|
||||
def _scan_predicted_local(source: str | Path | pl.LazyFrame) -> pl.LazyFrame:
|
||||
@@ -554,8 +616,6 @@ def marginal_table_pl(
|
||||
see the module-level note above on why this stays lazy.
|
||||
"""
|
||||
lf = _scan_predicted_local(source)
|
||||
pred_cols = [f"pred_{name}" for name in LOCAL_TARGET_NAMES]
|
||||
true_cols = [f"true_{name}" for name in LOCAL_TARGET_NAMES]
|
||||
|
||||
rows = []
|
||||
for label, cond in _group_filters_pl(lf, group_by, n_energy_bins):
|
||||
@@ -566,8 +626,8 @@ def marginal_table_pl(
|
||||
for j, name in enumerate(RAW_TARGET_NAMES):
|
||||
pair = glf.select(
|
||||
[
|
||||
_raw_expr(true_cols[j], j).alias("real"),
|
||||
_raw_expr(pred_cols[j], j).alias("gen"),
|
||||
_raw_dim_expr("true", j).alias("real"),
|
||||
_raw_dim_expr("pred", j).alias("gen"),
|
||||
]
|
||||
).collect()
|
||||
real_s, gen_s = pair["real"], pair["gen"]
|
||||
@@ -875,7 +935,7 @@ def constraint_report(
|
||||
"mean_abs_error": float(np.mean(np.abs(travel_norm - 1))),
|
||||
},
|
||||
]
|
||||
for j, name in enumerate(RAW_TARGET_NAMES[:_N_LOG_DIMS]):
|
||||
for j, name in enumerate(RAW_TARGET_NAMES[:_N_SCALAR_DIMS]):
|
||||
rows.append(
|
||||
{
|
||||
"check": f"{name} >= 0",
|
||||
@@ -900,7 +960,7 @@ def constraint_report_pl(
|
||||
|
||||
post_norm = sum(pl.col(pred_cols[k]) ** 2 for k in range(3, 6)).sqrt()
|
||||
travel_norm = sum(pl.col(pred_cols[k]) ** 2 for k in range(6, 9)).sqrt()
|
||||
raw_log_dims = [_raw_expr(pred_cols[j], j) for j in range(_N_LOG_DIMS)]
|
||||
raw_log_dims = [_raw_dim_expr("pred", j) for j in range(_N_SCALAR_DIMS)]
|
||||
|
||||
agg = (
|
||||
lf.select(
|
||||
@@ -915,11 +975,15 @@ def constraint_report_pl(
|
||||
(travel_norm - 1).abs().mean().alias("travel_dir_mean_abs_error"),
|
||||
*[
|
||||
(raw < 0).mean().alias(f"{name}_violation_rate")
|
||||
for raw, name in zip(raw_log_dims, RAW_TARGET_NAMES[:_N_LOG_DIMS])
|
||||
for raw, name in zip(
|
||||
raw_log_dims, RAW_TARGET_NAMES[:_N_SCALAR_DIMS]
|
||||
)
|
||||
],
|
||||
*[
|
||||
raw.clip(upper_bound=0).abs().mean().alias(f"{name}_mean_abs_error")
|
||||
for raw, name in zip(raw_log_dims, RAW_TARGET_NAMES[:_N_LOG_DIMS])
|
||||
for raw, name in zip(
|
||||
raw_log_dims, RAW_TARGET_NAMES[:_N_SCALAR_DIMS]
|
||||
)
|
||||
],
|
||||
]
|
||||
)
|
||||
@@ -939,7 +1003,7 @@ def constraint_report_pl(
|
||||
"mean_abs_error": agg["travel_dir_mean_abs_error"],
|
||||
},
|
||||
]
|
||||
for name in RAW_TARGET_NAMES[:_N_LOG_DIMS]:
|
||||
for name in RAW_TARGET_NAMES[:_N_SCALAR_DIMS]:
|
||||
rows.append(
|
||||
{
|
||||
"check": f"{name} >= 0",
|
||||
@@ -956,7 +1020,7 @@ def plot_constraint_violations(collection: SampleCollection):
|
||||
post_norm = np.linalg.norm(gen[:, 3:6], axis=1)
|
||||
travel_norm = np.linalg.norm(gen[:, 6:9], axis=1)
|
||||
|
||||
n_panels = 2 + _N_LOG_DIMS
|
||||
n_panels = 2 + _N_SCALAR_DIMS
|
||||
fig, axes = plt.subplots(1, n_panels, figsize=(4 * n_panels, 3.5))
|
||||
for ax, norm, title in [
|
||||
(axes[0], post_norm, "||post_dir||"),
|
||||
@@ -966,7 +1030,7 @@ def plot_constraint_violations(collection: SampleCollection):
|
||||
ax.set_yscale("log")
|
||||
ax.axvline(1.0, color="k", linestyle="--", linewidth=1)
|
||||
ax.set_title(title)
|
||||
for k, name in enumerate(RAW_TARGET_NAMES[:_N_LOG_DIMS]):
|
||||
for k, name in enumerate(RAW_TARGET_NAMES[:_N_SCALAR_DIMS]):
|
||||
ax = axes[2 + k]
|
||||
ax.hist(gen[:, k], bins=_hist_edges(gen[:, k], bins=50), histtype="step")
|
||||
ax.set_yscale("log")
|
||||
@@ -1134,6 +1198,7 @@ def compute_event_observables_pl(
|
||||
"pre_x",
|
||||
"pre_y",
|
||||
"pre_z",
|
||||
"pre_E", # needed to decode the energy simplex into physical edep
|
||||
"pre_dx",
|
||||
"pre_dy",
|
||||
"pre_dz",
|
||||
@@ -1154,10 +1219,12 @@ def compute_event_observables_pl(
|
||||
.astype(np.float32)
|
||||
)
|
||||
|
||||
pre_E = batch_df["pre_E"].to_numpy().astype(np.float32)
|
||||
|
||||
def _reconstruct(cols: list[str]) -> tuple[np.ndarray, np.ndarray, np.ndarray]:
|
||||
raw = batch_df.select(cols).to_numpy().astype(np.float32)
|
||||
step_length = inv_log_transform(raw[:, 0])
|
||||
edep = inv_log_transform(raw[:, 2])
|
||||
edep, _e_sec, _post_E, _delta_e = energy_simplex_decode(raw[:, 1:3], pre_E)
|
||||
travel_dir_local = raw[:, 6:9]
|
||||
post_pos = reconstruct_post_pos(
|
||||
pre_pos, pre_dir, step_length, travel_dir_local
|
||||
@@ -1226,19 +1293,18 @@ def compute_event_observables_pl(
|
||||
medians = (
|
||||
lf.select(
|
||||
"event_id",
|
||||
"true_log_edep",
|
||||
"pred_log_edep",
|
||||
"true_edep_logit",
|
||||
"true_sec_logit",
|
||||
"pred_edep_logit",
|
||||
"pred_sec_logit",
|
||||
"pre_E",
|
||||
"true_log_step_length",
|
||||
"pred_log_step_length",
|
||||
)
|
||||
.group_by("event_id")
|
||||
.agg(
|
||||
(pl.col("true_log_edep").exp() - _LOG_EPS)
|
||||
.median()
|
||||
.alias("real_median_edep"),
|
||||
(pl.col("pred_log_edep").exp() - _LOG_EPS)
|
||||
.median()
|
||||
.alias("gen_median_edep"),
|
||||
_edep_pl("true").median().alias("real_median_edep"),
|
||||
_edep_pl("pred").median().alias("gen_median_edep"),
|
||||
(pl.col("true_log_step_length").exp() - _LOG_EPS)
|
||||
.median()
|
||||
.alias("real_median_length"),
|
||||
@@ -1545,9 +1611,9 @@ def pdg_contribution_table_pl(source: str | Path | pl.LazyFrame) -> pl.DataFrame
|
||||
"""Total edep / step_length contributed by each pdg species, real vs generated.
|
||||
|
||||
One row per pdg code, sorted by pdg. Pure lazy polars `group_by` over the
|
||||
whole file — `edep`/`step_length` are scalars unaffected by the
|
||||
local-frame rotation, so this only needs the same `exp(...) - eps` de-log
|
||||
transform `inv_log_transform` does, expressed directly as a polars expr.
|
||||
whole file — `edep`/`step_length` are scalars unaffected by the local-frame
|
||||
rotation. `step_length` is a simple `exp(...) - eps` de-log; `edep` is decoded
|
||||
from the deposit/secondary energy logits against pre_E (`_edep_pl`).
|
||||
"""
|
||||
lf = _scan_predicted_local(source)
|
||||
|
||||
@@ -1557,8 +1623,8 @@ def pdg_contribution_table_pl(source: str | Path | pl.LazyFrame) -> pl.DataFrame
|
||||
return (
|
||||
lf.group_by("pdg")
|
||||
.agg(
|
||||
_delog("true_log_edep").sum().alias("real_total_edep"),
|
||||
_delog("pred_log_edep").sum().alias("gen_total_edep"),
|
||||
_edep_pl("true").sum().alias("real_total_edep"),
|
||||
_edep_pl("pred").sum().alias("gen_total_edep"),
|
||||
_delog("true_log_step_length").sum().alias("real_total_length"),
|
||||
_delog("pred_log_step_length").sum().alias("gen_total_length"),
|
||||
)
|
||||
|
||||
+8
-4
@@ -27,6 +27,7 @@ from giant.data.loader import (
|
||||
from giant.data.transforms import (
|
||||
build_features,
|
||||
build_cond_features,
|
||||
energy_simplex_decode,
|
||||
inv_local_frame_rotation,
|
||||
inv_log_transform,
|
||||
reconstruct_post_pos,
|
||||
@@ -283,8 +284,7 @@ def predict(
|
||||
batch_size_value = int(batch_size)
|
||||
except ValueError:
|
||||
typer.echo(
|
||||
f"error: --batch-size must be an integer or 'auto', "
|
||||
f"got {batch_size!r}",
|
||||
f"error: --batch-size must be an integer or 'auto', got {batch_size!r}",
|
||||
err=True,
|
||||
)
|
||||
raise typer.Exit(1)
|
||||
@@ -402,8 +402,12 @@ def predict(
|
||||
)
|
||||
else:
|
||||
step_length = inv_log_transform(raw[:, 0])
|
||||
delta_e = inv_log_transform(raw[:, 1])
|
||||
edep = inv_log_transform(raw[:, 2])
|
||||
# Columns 1:3 are ALR coords of the deposit/secondary/post energy
|
||||
# simplex; decode them against pre_E so edep + e_sec + post_E == pre_E
|
||||
# (hence delta_e == edep + e_sec) holds by construction.
|
||||
edep, _e_sec, _post_E, delta_e = energy_simplex_decode(
|
||||
raw[:, 1:3], piece["pre_E"]
|
||||
)
|
||||
|
||||
# Normalise predicted direction then rotate back to world frame
|
||||
post_dir_local = raw[:, 3:6].copy()
|
||||
|
||||
+14
-3
@@ -1,9 +1,20 @@
|
||||
X_DIM = 9
|
||||
|
||||
# Conditioning continuous-feature width: pre_pos(3), log(pre_E)(1), pre_dir(3),
|
||||
# layer_id(1), n_sec(1), log(e_sec)(1). One wider than X_DIM because e_sec
|
||||
# (secondary energy) is a conditioning input in the energy-conservation PoC.
|
||||
COND_DIM = 10
|
||||
|
||||
# The two energy columns are additive-log-ratio (ALR) coordinates of the
|
||||
# deposit/secondary/post energy simplex (fractions of pre_E that sum to 1),
|
||||
# referenced to the post-energy fraction — see giant.data.transforms
|
||||
# .energy_simplex_encode/.energy_simplex_decode. They replace the former
|
||||
# independent log_delta_e / log_edep targets so energy conservation holds by
|
||||
# construction after decoding.
|
||||
LOCAL_TARGET_NAMES = [
|
||||
"log_step_length",
|
||||
"log_delta_e",
|
||||
"log_edep",
|
||||
"edep_logit",
|
||||
"sec_logit",
|
||||
"post_dx",
|
||||
"post_dy",
|
||||
"post_dz",
|
||||
@@ -17,4 +28,4 @@ LOCAL_TARGET_NAMES = [
|
||||
# guessing from its column names.
|
||||
PREDICT_COORD_METADATA_KEY = "giant.predict.coord"
|
||||
PREDICT_SCHEMA_VERSION_KEY = "giant.predict.schema_version"
|
||||
PREDICT_SCHEMA_VERSION = "1"
|
||||
PREDICT_SCHEMA_VERSION = "2"
|
||||
|
||||
@@ -26,7 +26,9 @@ def _df_to_dict(df: pd.DataFrame) -> dict[str, np.ndarray]:
|
||||
"material": df["material"].to_numpy(dtype=object),
|
||||
"layer_id": df["layer_id"].to_numpy(dtype=np.int32),
|
||||
"n_sec": df["child_track_ids"].apply(len).to_numpy(dtype=np.int32),
|
||||
"e_sec": df["e_sec"].to_numpy(dtype=np.float32),
|
||||
"step_length": df["step_length"].to_numpy(dtype=np.float32),
|
||||
"post_E": df["post_E"].to_numpy(dtype=np.float32),
|
||||
"delta_e": (df["pre_E"] - df["post_E"]).to_numpy(dtype=np.float32),
|
||||
"edep": df["edep"].to_numpy(dtype=np.float32),
|
||||
"post_dir": df[["post_dx", "post_dy", "post_dz"]].to_numpy(dtype=np.float32),
|
||||
@@ -63,6 +65,7 @@ _COND_COLS = [
|
||||
"material",
|
||||
"layer_id",
|
||||
"child_track_ids",
|
||||
"e_sec",
|
||||
]
|
||||
|
||||
|
||||
@@ -76,6 +79,7 @@ def _cond_df_to_dict(df: pd.DataFrame) -> dict[str, np.ndarray]:
|
||||
"material": df["material"].to_numpy(dtype=object),
|
||||
"layer_id": df["layer_id"].to_numpy(dtype=np.int32),
|
||||
"n_sec": df["child_track_ids"].apply(len).to_numpy(dtype=np.int32),
|
||||
"e_sec": df["e_sec"].to_numpy(dtype=np.float32),
|
||||
}
|
||||
|
||||
|
||||
|
||||
@@ -2,6 +2,12 @@ import numpy as np
|
||||
|
||||
_EPS = 1e-8
|
||||
|
||||
# Floor added to each energy fraction before taking log-ratios so the simplex
|
||||
# coordinates stay finite on the heavily populated boundary: e_sec is 0 in ~half
|
||||
# of all steps and post_E is 0 at every track end. Kept tiny (1e-5 of pre_E) so
|
||||
# the conservation it slightly softens is physically negligible (~0.001%).
|
||||
_SIMPLEX_FLOOR = 1e-5
|
||||
|
||||
|
||||
def log_transform(x: np.ndarray, eps: float = _EPS) -> np.ndarray:
|
||||
return np.log(np.asarray(x, dtype=np.float32) + eps)
|
||||
@@ -11,6 +17,69 @@ def inv_log_transform(y: np.ndarray, eps: float = _EPS) -> np.ndarray:
|
||||
return np.exp(np.asarray(y, dtype=np.float32)) - eps
|
||||
|
||||
|
||||
def energy_simplex_encode(
|
||||
edep: np.ndarray,
|
||||
e_sec: np.ndarray,
|
||||
post_E: np.ndarray,
|
||||
pre_E: np.ndarray,
|
||||
) -> np.ndarray:
|
||||
"""Encode (edep, e_sec, post_E) as 2 additive-log-ratio (ALR) coordinates.
|
||||
|
||||
The three energies are first expressed as fractions of pre_E that sum to 1:
|
||||
post_E is kept exact (so the particle's retained energy — the next step's
|
||||
input during a rollout — is preserved), and the energy actually lost,
|
||||
`delta_e = pre_E - post_E`, is split between local deposit and secondaries
|
||||
in the recorded edep:e_sec ratio (re-attributing any sub-threshold / rest-mass
|
||||
leakage proportionally so the three fractions sum to exactly 1). A small floor
|
||||
(`_SIMPLEX_FLOOR`) keeps the log-ratios finite where a fraction is 0.
|
||||
|
||||
Returns an (N, 2) array of ALR coordinates referenced to the post fraction;
|
||||
`energy_simplex_decode` is its inverse (up to the floor softening).
|
||||
"""
|
||||
pre_E = np.maximum(np.asarray(pre_E, dtype=np.float32), _EPS)
|
||||
post_E = np.clip(np.asarray(post_E, dtype=np.float32), 0.0, pre_E)
|
||||
delta_e = pre_E - post_E
|
||||
lost = np.asarray(edep, dtype=np.float32) + np.asarray(e_sec, dtype=np.float32)
|
||||
has_loss = lost > _EPS
|
||||
scale = np.where(has_loss, delta_e / np.maximum(lost, _EPS), 0.0)
|
||||
# Where nothing was recorded as deposited/secondary but energy was lost,
|
||||
# attribute all of delta_e to local deposit.
|
||||
edep_r = np.where(has_loss, edep * scale, delta_e)
|
||||
e_sec_r = np.where(has_loss, e_sec * scale, 0.0)
|
||||
|
||||
f = np.stack([edep_r, e_sec_r, post_E], axis=-1) / pre_E[:, None] # (N,3), sums≈1
|
||||
f = (f + _SIMPLEX_FLOOR) / (1.0 + 3.0 * _SIMPLEX_FLOOR)
|
||||
log_f = np.log(f)
|
||||
z = log_f[:, :2] - log_f[:, 2:3] # ALR vs post reference
|
||||
return z.astype(np.float32)
|
||||
|
||||
|
||||
def energy_simplex_decode(
|
||||
z: np.ndarray, pre_E: np.ndarray
|
||||
) -> tuple[np.ndarray, np.ndarray, np.ndarray, np.ndarray]:
|
||||
"""Inverse of `energy_simplex_encode`: ALR coords + pre_E → physical energies.
|
||||
|
||||
A softmax over `[z_edep, z_sec, 0]` recovers the three simplex fractions, so
|
||||
`edep + e_sec + post_E == pre_E` holds by construction (the energy-conservation
|
||||
inductive bias). Returns `(edep, e_sec, post_E, delta_e)` in physical units.
|
||||
"""
|
||||
z = np.asarray(z, dtype=np.float32)
|
||||
pre_E = np.asarray(pre_E, dtype=np.float32)
|
||||
logits = np.concatenate([z, np.zeros((len(z), 1), dtype=np.float32)], axis=1)
|
||||
logits = logits - logits.max(axis=1, keepdims=True)
|
||||
f = np.exp(logits)
|
||||
f /= f.sum(axis=1, keepdims=True) # sums to 1 exactly → exact conservation
|
||||
E = f * pre_E[:, None]
|
||||
edep, e_sec, post_E = E[:, 0], E[:, 1], E[:, 2]
|
||||
delta_e = pre_E - post_E
|
||||
return (
|
||||
edep.astype(np.float32),
|
||||
e_sec.astype(np.float32),
|
||||
post_E.astype(np.float32),
|
||||
delta_e.astype(np.float32),
|
||||
)
|
||||
|
||||
|
||||
def local_frame_rotation(pre_dir: np.ndarray, post_dir: np.ndarray) -> np.ndarray:
|
||||
"""Rotate post_dir into the local frame where pre_dir maps to ẑ (Rodrigues).
|
||||
|
||||
@@ -172,6 +241,7 @@ def build_cond_features(
|
||||
data["pre_dir"],
|
||||
data["layer_id"].astype(np.float32),
|
||||
data["n_sec"].astype(np.float32),
|
||||
log_transform(data["e_sec"]),
|
||||
]
|
||||
).astype(np.float32)
|
||||
|
||||
@@ -202,11 +272,14 @@ def build_features(
|
||||
data["pre_dir"], travel_direction(data["pre_pos"], data["post_pos"])
|
||||
)
|
||||
|
||||
energy_z = energy_simplex_encode(
|
||||
data["edep"], data["e_sec"], data["post_E"], data["pre_E"]
|
||||
) # (N, 2): ALR coords of the deposit/secondary/post energy simplex
|
||||
|
||||
target = np.column_stack(
|
||||
[
|
||||
log_transform(data["step_length"]),
|
||||
log_transform(data["delta_e"]),
|
||||
log_transform(data["edep"]),
|
||||
energy_z,
|
||||
post_dir_local,
|
||||
travel_dir_local,
|
||||
]
|
||||
@@ -219,8 +292,9 @@ def build_features(
|
||||
data["pre_dir"],
|
||||
data["layer_id"].astype(np.float32),
|
||||
data["n_sec"].astype(np.float32),
|
||||
log_transform(data["e_sec"]),
|
||||
]
|
||||
).astype(np.float32) # (N, 9)
|
||||
).astype(np.float32) # (N, COND_DIM)
|
||||
|
||||
pdg_idx = np.array([pdg_map[int(p)] for p in data["pdg"]], dtype=np.int64)
|
||||
mat_idx = np.array([mat_map[str(m)] for m in data["material"]], dtype=np.int64)
|
||||
|
||||
@@ -3,7 +3,7 @@ import math
|
||||
import torch
|
||||
import torch.nn as nn
|
||||
|
||||
from giant.constants import X_DIM
|
||||
from giant.constants import COND_DIM, X_DIM
|
||||
|
||||
|
||||
class SinusoidalEmbedding(nn.Module):
|
||||
@@ -29,7 +29,7 @@ class ConditionEncoder(nn.Module):
|
||||
self,
|
||||
pdg_vocab: int,
|
||||
mat_vocab: int,
|
||||
cont_dim: int = 9,
|
||||
cont_dim: int = COND_DIM,
|
||||
emb_dim: int = 16,
|
||||
out_dim: int = 128,
|
||||
) -> None:
|
||||
|
||||
+2
-2
@@ -5,7 +5,7 @@ import torch
|
||||
from torch.utils.data import DataLoader
|
||||
|
||||
from giant import config
|
||||
from giant.constants import X_DIM
|
||||
from giant.constants import COND_DIM, X_DIM
|
||||
from giant.data.loader import (
|
||||
find_parquet_files,
|
||||
load_event_ids,
|
||||
@@ -55,7 +55,7 @@ def run_train_job(
|
||||
echo(f" {len(pdg_map)} PDG codes | {len(mat_map)} materials")
|
||||
|
||||
echo("fitting normalizer (streaming) …")
|
||||
cond_acc = _WelfordAccumulator(X_DIM)
|
||||
cond_acc = _WelfordAccumulator(COND_DIM)
|
||||
tgt_acc = _WelfordAccumulator(X_DIM)
|
||||
for path in files:
|
||||
for chunk in iter_file_chunks(path):
|
||||
|
||||
@@ -19,6 +19,43 @@ import uproot
|
||||
ParquetCompression = Literal["lz4", "uncompressed", "snappy", "gzip", "brotli", "zstd"]
|
||||
|
||||
|
||||
def _add_secondary_energy(df: pl.DataFrame) -> pl.DataFrame:
|
||||
"""Add per-step `e_sec`: total initial kinetic energy of the secondaries born in it.
|
||||
|
||||
Each secondary's creation energy is the `pre_E` of that child track's first step
|
||||
(min `step_no`) in the same event, so for a parent step
|
||||
`e_sec = Σ over child_track_ids of the child track's first-step pre_E`. Steps that
|
||||
spawn nothing get 0.0. The full event must be present in `df` (it is — the writer
|
||||
concatenates every batch before this runs), since a child track's first step can
|
||||
live in a different read batch than its parent step.
|
||||
"""
|
||||
first_E = (
|
||||
df.sort("step_no")
|
||||
.group_by(["event_id", "track_id"])
|
||||
.agg(pl.col("pre_E").first().alias("child_E"))
|
||||
.rename({"track_id": "child_track_id"})
|
||||
)
|
||||
|
||||
exploded = (
|
||||
df.select(["event_id", "child_track_ids"])
|
||||
.with_row_index("_step_row")
|
||||
.explode("child_track_ids")
|
||||
.rename({"child_track_ids": "child_track_id"})
|
||||
.drop_nulls("child_track_id") # steps with no children explode to a null row
|
||||
)
|
||||
summed = (
|
||||
exploded.join(first_E, on=["event_id", "child_track_id"], how="left")
|
||||
.group_by("_step_row")
|
||||
.agg(pl.col("child_E").sum().alias("e_sec"))
|
||||
)
|
||||
return (
|
||||
df.with_row_index("_step_row")
|
||||
.join(summed, on="_step_row", how="left")
|
||||
.with_columns(pl.col("e_sec").fill_null(0.0).cast(pl.Float64))
|
||||
.drop("_step_row")
|
||||
)
|
||||
|
||||
|
||||
def _batch_to_polars(batch: ak.Array) -> pl.DataFrame:
|
||||
"""Convert one awkward-array batch to a Polars DataFrame.
|
||||
|
||||
@@ -77,8 +114,15 @@ def convert_steps_to_parquet(
|
||||
rows_done += len(batch)
|
||||
print(f" {rows_done:,} / {n_entries:,} rows read", end="\r", flush=True)
|
||||
|
||||
df = pl.concat(batches)
|
||||
# Steps tree carries the parent→child links needed to derive secondary energy;
|
||||
# other trees (e.g. Hits) don't, so only augment when the column is present.
|
||||
if "child_track_ids" in df.columns:
|
||||
print("\nComputing per-step secondary energy (e_sec) …", end=" ", flush=True)
|
||||
df = _add_secondary_energy(df)
|
||||
|
||||
print(f"\nWriting {output_path} …", end=" ", flush=True)
|
||||
pl.concat(batches).write_parquet(output_path, compression=compression)
|
||||
df.write_parquet(output_path, compression=compression)
|
||||
print(f"done ({output_path.stat().st_size / 1e6:.1f} MB)")
|
||||
return output_path
|
||||
|
||||
|
||||
+50
-20
@@ -41,7 +41,12 @@ from giant.constants import (
|
||||
PREDICT_SCHEMA_VERSION,
|
||||
PREDICT_SCHEMA_VERSION_KEY,
|
||||
)
|
||||
from giant.data.transforms import inv_log_transform, log_transform, reconstruct_post_pos
|
||||
from giant.data.transforms import (
|
||||
energy_simplex_decode,
|
||||
inv_log_transform,
|
||||
log_transform,
|
||||
reconstruct_post_pos,
|
||||
)
|
||||
|
||||
|
||||
def _unit_vectors(rng, n):
|
||||
@@ -186,13 +191,17 @@ def test_plot_kl_bars_caps_groups_by_pdg():
|
||||
|
||||
|
||||
def _write_predicted_local_parquet(path, n=50, metadata=None, rng=None):
|
||||
"""Mimic `giant predict --coord local`'s output schema for the loader tests."""
|
||||
"""Mimic `giant predict --coord local`'s output schema for the loader tests.
|
||||
|
||||
Column 0 is a log-scaled step_length; columns 1–2 are the deposit/secondary
|
||||
ALR energy logits (unconstrained reals, decoded against pre_E); columns 3–8
|
||||
are direction components.
|
||||
"""
|
||||
rng = rng or np.random.default_rng(0)
|
||||
true_log_local = rng.standard_normal((n, 9)).astype(np.float32)
|
||||
true_log_local[:, :3] = log_transform(
|
||||
rng.uniform(0.1, 5.0, (n, 3)).astype(np.float32)
|
||||
)
|
||||
true_log_local[:, 0] = log_transform(rng.uniform(0.1, 5.0, n).astype(np.float32))
|
||||
pred_log_local = true_log_local + rng.normal(0, 0.01, (n, 9)).astype(np.float32)
|
||||
pre_E = rng.uniform(1.0, 100.0, n).astype(np.float32)
|
||||
|
||||
table = pa.table(
|
||||
{
|
||||
@@ -201,7 +210,7 @@ def _write_predicted_local_parquet(path, n=50, metadata=None, rng=None):
|
||||
"pre_x": rng.standard_normal(n).astype(np.float32),
|
||||
"pre_y": rng.standard_normal(n).astype(np.float32),
|
||||
"pre_z": rng.standard_normal(n).astype(np.float32),
|
||||
"pre_E": rng.uniform(1.0, 100.0, n).astype(np.float32),
|
||||
"pre_E": pre_E,
|
||||
"pre_dx": rng.standard_normal(n).astype(np.float32),
|
||||
"pre_dy": rng.standard_normal(n).astype(np.float32),
|
||||
"pre_dz": rng.standard_normal(n).astype(np.float32),
|
||||
@@ -221,12 +230,12 @@ def _write_predicted_local_parquet(path, n=50, metadata=None, rng=None):
|
||||
if metadata is not None:
|
||||
table = table.replace_schema_metadata(metadata)
|
||||
pq.write_table(table, path)
|
||||
return true_log_local, pred_log_local
|
||||
return true_log_local, pred_log_local, pre_E
|
||||
|
||||
|
||||
def test_load_predicted_local_round_trips_values(tmp_path):
|
||||
path = tmp_path / "predicted_local.parquet"
|
||||
true_log_local, pred_log_local = _write_predicted_local_parquet(
|
||||
true_log_local, pred_log_local, pre_E = _write_predicted_local_parquet(
|
||||
path,
|
||||
metadata={
|
||||
PREDICT_COORD_METADATA_KEY: "local",
|
||||
@@ -236,12 +245,20 @@ def test_load_predicted_local_round_trips_values(tmp_path):
|
||||
|
||||
collection = load_predicted_local(path)
|
||||
|
||||
expected_real = true_log_local.copy()
|
||||
expected_real[:, :3] = np.exp(expected_real[:, :3]) - 1e-8
|
||||
expected_gen = pred_log_local.copy()
|
||||
expected_gen[:, :3] = np.exp(expected_gen[:, :3]) - 1e-8
|
||||
np.testing.assert_allclose(collection.real_raw, expected_real, atol=1e-4)
|
||||
np.testing.assert_allclose(collection.gen_raw, expected_gen, atol=1e-4)
|
||||
def expected_raw(log_local):
|
||||
raw = log_local.copy()
|
||||
raw[:, 0] = np.exp(log_local[:, 0]) - 1e-8
|
||||
edep, _e_sec, _post_E, delta_e = energy_simplex_decode(log_local[:, 1:3], pre_E)
|
||||
raw[:, 1] = delta_e
|
||||
raw[:, 2] = edep
|
||||
return raw
|
||||
|
||||
np.testing.assert_allclose(
|
||||
collection.real_raw, expected_raw(true_log_local), atol=1e-4
|
||||
)
|
||||
np.testing.assert_allclose(
|
||||
collection.gen_raw, expected_raw(pred_log_local), atol=1e-4
|
||||
)
|
||||
assert collection.real_norm is None
|
||||
assert collection.gen_norm is None
|
||||
|
||||
@@ -321,13 +338,21 @@ def test_marginal_table_pl_matches_numpy_version(tmp_path, group_by):
|
||||
|
||||
assert list(expected["group"]) == list(actual["group"])
|
||||
assert list(expected["n"]) == list(actual["n"])
|
||||
for col in ["real_mean", "gen_mean", "real_std", "gen_std", "kl_real_gen"]:
|
||||
for col in ["real_mean", "gen_mean", "real_std", "gen_std"]:
|
||||
np.testing.assert_allclose(
|
||||
expected[col].to_numpy(),
|
||||
actual[col].to_numpy(),
|
||||
atol=1e-4,
|
||||
rtol=1e-4,
|
||||
)
|
||||
# KL uses np.histogram (numpy path) vs polars Series.hist (lazy path); the two
|
||||
# backends bin the boundary (min/max) sample differently, so allow a small
|
||||
# absolute discrepancy rather than requiring bit-identical estimates.
|
||||
np.testing.assert_allclose(
|
||||
expected["kl_real_gen"].to_numpy(),
|
||||
actual["kl_real_gen"].to_numpy(),
|
||||
atol=2e-2,
|
||||
)
|
||||
|
||||
|
||||
@pytest.mark.parametrize("group_by", [None, "pdg", "material", "energy"])
|
||||
@@ -387,7 +412,8 @@ def _make_event_level_arrays(rng):
|
||||
|
||||
def _local_block():
|
||||
block = rng.standard_normal((n, 9)).astype(np.float32)
|
||||
block[:, :3] = log_transform(rng.uniform(0.1, 5.0, (n, 3)).astype(np.float32))
|
||||
block[:, 0] = log_transform(rng.uniform(0.1, 5.0, n).astype(np.float32))
|
||||
# cols 1–2 stay as random ALR energy logits (decoded against pre_E)
|
||||
block[:, 3:6] = _unit_vectors(rng, n)
|
||||
block[:, 6:9] = _unit_vectors(rng, n)
|
||||
return block
|
||||
@@ -451,7 +477,9 @@ def _expected_event_table(
|
||||
|
||||
def agg(log_local):
|
||||
step_length = inv_log_transform(log_local[mask, 0])
|
||||
edep = inv_log_transform(log_local[mask, 2])
|
||||
edep, _e_sec, _post_E, _delta_e = energy_simplex_decode(
|
||||
log_local[mask, 1:3], pre_E[mask]
|
||||
)
|
||||
travel_dir_local = log_local[mask, 6:9]
|
||||
post_pos = reconstruct_post_pos(
|
||||
pre_pos[mask], pre_dir[mask], step_length, travel_dir_local
|
||||
@@ -574,10 +602,12 @@ def test_event_level_plots_run_without_error(tmp_path):
|
||||
|
||||
def test_pdg_contribution_table_pl_matches_manual_sums(tmp_path):
|
||||
path = tmp_path / "event_level.parquet"
|
||||
_, _, _, _, true_log_local, pred_log_local, pdg = _write_event_level_parquet(path)
|
||||
_, _, _, pre_E, true_log_local, pred_log_local, pdg = _write_event_level_parquet(
|
||||
path
|
||||
)
|
||||
|
||||
real_edep = inv_log_transform(true_log_local[:, 2])
|
||||
gen_edep = inv_log_transform(pred_log_local[:, 2])
|
||||
real_edep = energy_simplex_decode(true_log_local[:, 1:3], pre_E)[0]
|
||||
gen_edep = energy_simplex_decode(pred_log_local[:, 1:3], pre_E)[0]
|
||||
real_length = inv_log_transform(true_log_local[:, 0])
|
||||
gen_length = inv_log_transform(pred_log_local[:, 0])
|
||||
|
||||
|
||||
+4
-3
@@ -1,4 +1,5 @@
|
||||
import torch
|
||||
from giant.constants import COND_DIM
|
||||
from giant.model.network import DenoisingMLP
|
||||
from giant.model.schedule import CosineSchedule, flow_matching_loss
|
||||
from giant.sample import sample_flow, sample_ddim
|
||||
@@ -10,7 +11,7 @@ def _small_model():
|
||||
|
||||
def _batch(B=8):
|
||||
x1 = torch.randn(B, 9)
|
||||
cond_cont = torch.randn(B, 9)
|
||||
cond_cont = torch.randn(B, COND_DIM)
|
||||
cond_cat = torch.zeros(B, 2, dtype=torch.long)
|
||||
return x1, cond_cont, cond_cat
|
||||
|
||||
@@ -36,7 +37,7 @@ def test_flow_matching_loss_has_grad():
|
||||
|
||||
def test_sample_flow_shape():
|
||||
B = 6
|
||||
cond_cont = torch.randn(B, 9)
|
||||
cond_cont = torch.randn(B, COND_DIM)
|
||||
cond_cat = torch.zeros(B, 2, dtype=torch.long)
|
||||
out = sample_flow(_small_model(), cond_cont, cond_cat, steps=5)
|
||||
assert out.shape == (B, 9)
|
||||
@@ -52,7 +53,7 @@ def test_ddpm_loss_nonneg():
|
||||
def test_sample_ddim_shape():
|
||||
B = 4
|
||||
schedule = CosineSchedule(T=50)
|
||||
cond_cont = torch.randn(B, 9)
|
||||
cond_cont = torch.randn(B, COND_DIM)
|
||||
cond_cat = torch.zeros(B, 2, dtype=torch.long)
|
||||
out = sample_ddim(_small_model(), cond_cont, cond_cat, schedule, steps=5)
|
||||
assert out.shape == (B, 9)
|
||||
|
||||
@@ -1,4 +1,5 @@
|
||||
import torch
|
||||
from giant.constants import COND_DIM
|
||||
from giant.model.network import DenoisingMLP, SinusoidalEmbedding
|
||||
|
||||
|
||||
@@ -19,7 +20,7 @@ def test_denoising_mlp_output_shape():
|
||||
model = DenoisingMLP(pdg_vocab=5, mat_vocab=3)
|
||||
x_t = torch.randn(B, 9)
|
||||
t = torch.rand(B)
|
||||
cond_cont = torch.randn(B, 9)
|
||||
cond_cont = torch.randn(B, COND_DIM)
|
||||
cond_cat = torch.stack(
|
||||
[
|
||||
torch.randint(0, 5, (B,)),
|
||||
@@ -36,7 +37,7 @@ def test_denoising_mlp_gradients_flow():
|
||||
model = DenoisingMLP(pdg_vocab=3, mat_vocab=2, hidden_dim=32, n_blocks=2)
|
||||
x_t = torch.randn(B, 9)
|
||||
t = torch.rand(B)
|
||||
cond_cont = torch.randn(B, 9)
|
||||
cond_cont = torch.randn(B, COND_DIM)
|
||||
cond_cat = torch.zeros(B, 2, dtype=torch.long)
|
||||
loss = model(x_t, t, cond_cont, cond_cat).sum()
|
||||
loss.backward()
|
||||
|
||||
@@ -0,0 +1,50 @@
|
||||
import importlib.util
|
||||
from pathlib import Path
|
||||
|
||||
import polars as pl
|
||||
|
||||
# scripts/ is not an installed package — load the module straight from its path.
|
||||
_SPEC = importlib.util.spec_from_file_location(
|
||||
"steps_to_parquet",
|
||||
Path(__file__).resolve().parents[1] / "scripts" / "steps_to_parquet.py",
|
||||
)
|
||||
steps_to_parquet = importlib.util.module_from_spec(_SPEC)
|
||||
_SPEC.loader.exec_module(steps_to_parquet)
|
||||
|
||||
|
||||
def _frame() -> pl.DataFrame:
|
||||
# event 0: step (1,0) spawns track 2 (first-step pre_E=15) → e_sec=15.
|
||||
# event 1: step (1,0) spawns tracks 2 & 3 (20 + 30) → e_sec=50.
|
||||
return pl.DataFrame(
|
||||
{
|
||||
"event_id": [0, 0, 0, 1, 1, 1],
|
||||
"track_id": [1, 1, 2, 1, 2, 3],
|
||||
"step_no": [0, 1, 0, 0, 0, 0],
|
||||
"pre_E": [100.0, 80.0, 15.0, 200.0, 20.0, 30.0],
|
||||
"child_track_ids": [[2], [], [], [2, 3], [], []],
|
||||
}
|
||||
)
|
||||
|
||||
|
||||
def test_e_sec_sums_child_first_step_energy():
|
||||
out = steps_to_parquet._add_secondary_energy(_frame())
|
||||
e_sec = dict(
|
||||
zip(zip(out["track_id"], out["step_no"], out["event_id"]), out["e_sec"])
|
||||
)
|
||||
assert e_sec[(1, 0, 0)] == 15.0 # one child, first-step pre_E 15
|
||||
assert e_sec[(1, 0, 1)] == 50.0 # two children, 20 + 30
|
||||
|
||||
|
||||
def test_e_sec_zero_when_no_children():
|
||||
out = steps_to_parquet._add_secondary_energy(_frame())
|
||||
childless = out.filter(
|
||||
(pl.col("event_id") == 0) & (pl.col("track_id") == 1) & (pl.col("step_no") == 1)
|
||||
)
|
||||
assert childless["e_sec"].item() == 0.0
|
||||
|
||||
|
||||
def test_e_sec_preserves_row_count_and_order():
|
||||
df = _frame()
|
||||
out = steps_to_parquet._add_secondary_energy(df)
|
||||
assert out.height == df.height
|
||||
assert out["pre_E"].to_list() == df["pre_E"].to_list()
|
||||
@@ -1,5 +1,7 @@
|
||||
import numpy as np
|
||||
from giant.data.transforms import (
|
||||
energy_simplex_decode,
|
||||
energy_simplex_encode,
|
||||
inv_log_transform,
|
||||
local_frame_rotation,
|
||||
log_transform,
|
||||
@@ -111,6 +113,51 @@ def test_reconstruct_post_pos_general_roundtrip():
|
||||
np.testing.assert_allclose(reconstructed, post_pos, atol=1e-4)
|
||||
|
||||
|
||||
def test_energy_simplex_conservation():
|
||||
"""Decoding any ALR coords yields energies that sum to pre_E exactly."""
|
||||
rng = np.random.default_rng(11)
|
||||
N = 500
|
||||
z = rng.standard_normal((N, 2)).astype(np.float32) * 3.0
|
||||
pre_E = rng.uniform(1.0, 100.0, N).astype(np.float32)
|
||||
edep, e_sec, post_E, delta_e = energy_simplex_decode(z, pre_E)
|
||||
np.testing.assert_allclose(edep + e_sec + post_E, pre_E, rtol=1e-5, atol=1e-4)
|
||||
np.testing.assert_allclose(delta_e, edep + e_sec, rtol=1e-5, atol=1e-4)
|
||||
assert np.all(edep >= 0) and np.all(e_sec >= 0) and np.all(post_E >= 0)
|
||||
|
||||
|
||||
def test_energy_simplex_roundtrip():
|
||||
"""Encode → decode recovers energies whose lost part already sums to delta_e."""
|
||||
rng = np.random.default_rng(12)
|
||||
N = 500
|
||||
pre_E = rng.uniform(1.0, 100.0, N).astype(np.float32)
|
||||
post_E = (pre_E * rng.uniform(0.0, 1.0, N)).astype(np.float32)
|
||||
delta_e = pre_E - post_E
|
||||
g = rng.uniform(0.0, 1.0, N).astype(np.float32)
|
||||
edep = (g * delta_e).astype(np.float32)
|
||||
e_sec = ((1.0 - g) * delta_e).astype(np.float32)
|
||||
|
||||
z = energy_simplex_encode(edep, e_sec, post_E, pre_E)
|
||||
edep_r, e_sec_r, post_E_r, _ = energy_simplex_decode(z, pre_E)
|
||||
# Tolerance reflects the tiny simplex floor (~1e-5 of pre_E).
|
||||
np.testing.assert_allclose(edep_r, edep, atol=5e-3)
|
||||
np.testing.assert_allclose(e_sec_r, e_sec, atol=5e-3)
|
||||
np.testing.assert_allclose(post_E_r, post_E, atol=5e-3)
|
||||
|
||||
|
||||
def test_energy_simplex_handles_boundary_zeros():
|
||||
"""e_sec=0 (no secondaries) and post_E=0 (track end) stay finite and decode near 0."""
|
||||
pre_E = np.array([10.0, 50.0, 100.0], dtype=np.float32)
|
||||
edep = np.array([4.0, 50.0, 0.0], dtype=np.float32)
|
||||
e_sec = np.array([0.0, 0.0, 0.0], dtype=np.float32) # no secondaries
|
||||
post_E = np.array([6.0, 0.0, 100.0], dtype=np.float32) # row 1: track ends
|
||||
|
||||
z = energy_simplex_encode(edep, e_sec, post_E, pre_E)
|
||||
assert np.all(np.isfinite(z))
|
||||
_, e_sec_r, post_E_r, _ = energy_simplex_decode(z, pre_E)
|
||||
np.testing.assert_allclose(e_sec_r, 0.0, atol=1e-2)
|
||||
assert post_E_r[1] < 1e-2 # the absorbed track decodes to ~0 post energy
|
||||
|
||||
|
||||
def test_normalizer_roundtrip():
|
||||
rng = np.random.default_rng(3)
|
||||
X = rng.standard_normal((200, 9)).astype(np.float32)
|
||||
|
||||
Reference in New Issue
Block a user