lars/giant
1
0
Fork 0
Code Issues 2 Pull Requests 2 Actions Packages Projects Releases Wiki Activity

Apply ruff format and document lint/type tooling in CLAUDE.md

Browse Source 
First repo-wide ruff format pass, plus a note in CLAUDE.md to run
ruff and ty periodically.

Co-Authored-By: Claude Sonnet 4.6 <noreply@anthropic.com>
This commit is contained in:
Lars Bogner
2026-06-18 17:40:27 +02:00
parent 53fd2e4405
commit 8cebc4809d
20 changed files with 668 additions and 307 deletions
Download Patch File Download Diff File Expand all files Collapse all files
tests/test_analysis.py
+65 -35
View File
@@ -39,27 +39,36 @@ def _unit_vectors(rng, n):
def _make_collection(n=200, seed=0, gen_offset=0.0) -> SampleCollection:
rng = np.random.default_rng(seed)
real = np.column_stack([
rng.uniform(0.1, 5.0, n), # step_length
rng.uniform(0.1, 5.0, n), # delta_e
rng.uniform(0.1, 5.0, n), # edep
_unit_vectors(rng, n), # post_dir
_unit_vectors(rng, n), # travel_dir
]).astype(np.float32)
real = np.column_stack(
[
rng.uniform(0.1, 5.0, n), # step_length
rng.uniform(0.1, 5.0, n), # delta_e
rng.uniform(0.1, 5.0, n), # edep
_unit_vectors(rng, n), # post_dir
_unit_vectors(rng, n), # travel_dir
]
).astype(np.float32)
gen = real + gen_offset
pre_E = rng.uniform(1.0, 100.0, n).astype(np.float32)
cond_cont_raw = np.column_stack([
rng.standard_normal((n, 3)), pre_E, rng.standard_normal((n, 3)),
rng.integers(0, 5, n), rng.integers(0, 3, n),
]).astype(np.float32)
cond_cont_raw = np.column_stack(
[
rng.standard_normal((n, 3)),
pre_E,
rng.standard_normal((n, 3)),
rng.integers(0, 5, n),
rng.integers(0, 3, n),
]
).astype(np.float32)
return SampleCollection(
cond_cont_raw=cond_cont_raw,
pdg=rng.choice([11, -11, 22], size=n),
material=rng.choice(["W", "Pb"], size=n),
real_raw=real, gen_raw=gen,
real_norm=real, gen_norm=gen,
real_raw=real,
gen_raw=gen,
real_norm=real,
gen_norm=gen,
)
@@ -150,25 +159,35 @@ def _write_predicted_local_parquet(path, n=50, metadata=None, rng=None):
"""Mimic `giant predict --coord local`'s output schema for the loader tests."""
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[:, :3] = log_transform(
rng.uniform(0.1, 5.0, (n, 3)).astype(np.float32)
)
pred_log_local = true_log_local + rng.normal(0, 0.01, (n, 9)).astype(np.float32)
table = pa.table({
"event_id": rng.integers(0, 10, n),
"pdg": rng.choice([11, -11, 22], n),
"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_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),
"material": rng.choice(["W", "Pb"], n),
"layer_id": rng.integers(0, 10, n).astype(np.int32),
"n_sec": rng.integers(0, 3, n).astype(np.int32),
**{f"pred_{name}": pred_log_local[:, j] for j, name in enumerate(LOCAL_TARGET_NAMES)},
**{f"true_{name}": true_log_local[:, j] for j, name in enumerate(LOCAL_TARGET_NAMES)},
})
table = pa.table(
{
"event_id": rng.integers(0, 10, n),
"pdg": rng.choice([11, -11, 22], n),
"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_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),
"material": rng.choice(["W", "Pb"], n),
"layer_id": rng.integers(0, 10, n).astype(np.int32),
"n_sec": rng.integers(0, 3, n).astype(np.int32),
**{
f"pred_{name}": pred_log_local[:, j]
for j, name in enumerate(LOCAL_TARGET_NAMES)
},
**{
f"true_{name}": true_log_local[:, j]
for j, name in enumerate(LOCAL_TARGET_NAMES)
},
}
)
if metadata is not None:
table = table.replace_schema_metadata(metadata)
pq.write_table(table, path)
@@ -263,14 +282,21 @@ def test_marginal_table_pl_matches_numpy_version(tmp_path, group_by):
path = _predicted_local_path(tmp_path)
collection = load_predicted_local(path)
expected = marginal_table(collection, group_by=group_by).sort_values(["group", "dim"])
actual = marginal_table_pl(path, group_by=group_by).sort(["group", "dim"]).to_pandas()
expected = marginal_table(collection, group_by=group_by).sort_values(
["group", "dim"]
)
actual = (
marginal_table_pl(path, group_by=group_by).sort(["group", "dim"]).to_pandas()
)
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"]:
np.testing.assert_allclose(
expected[col].to_numpy(), actual[col].to_numpy(), atol=1e-4, rtol=1e-4,
expected[col].to_numpy(),
actual[col].to_numpy(),
atol=1e-4,
rtol=1e-4,
)
@@ -290,8 +316,12 @@ def test_constraint_report_pl_matches_numpy_version(tmp_path):
assert list(expected["check"]) == list(actual["check"])
np.testing.assert_allclose(
expected["violation_rate"].to_numpy(), actual["violation_rate"].to_numpy(), atol=1e-6,
expected["violation_rate"].to_numpy(),
actual["violation_rate"].to_numpy(),
atol=1e-6,
)
np.testing.assert_allclose(
expected["mean_abs_error"].to_numpy(), actual["mean_abs_error"].to_numpy(), atol=1e-4,
expected["mean_abs_error"].to_numpy(),
actual["mean_abs_error"].to_numpy(),
atol=1e-4,
)