Apply ruff format and document lint/type tooling in CLAUDE.md
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:
Claude Sonnet 4.6
+160
-80
@@ -54,9 +54,16 @@ class Coord(str, Enum):
|
||||
|
||||
@app.command()
|
||||
def train(
|
||||
data: Annotated[Path, typer.Argument(help="Parquet file or directory of parquet files")],
|
||||
config: Annotated[Optional[Path], typer.Option(help="TOML config file (overridden by explicit flags)")] = None,
|
||||
mode: Annotated[Optional[Mode], typer.Option(help="Generative model: flow matching or DDPM")] = None,
|
||||
data: Annotated[
|
||||
Path, typer.Argument(help="Parquet file or directory of parquet files")
|
||||
],
|
||||
config: Annotated[
|
||||
Optional[Path],
|
||||
typer.Option(help="TOML config file (overridden by explicit flags)"),
|
||||
] = None,
|
||||
mode: Annotated[
|
||||
Optional[Mode], typer.Option(help="Generative model: flow matching or DDPM")
|
||||
] = None,
|
||||
epochs: Annotated[Optional[int], typer.Option()] = None,
|
||||
batch_size: Annotated[Optional[int], typer.Option()] = None,
|
||||
lr: Annotated[Optional[float], typer.Option()] = None,
|
||||
@@ -64,25 +71,55 @@ def train(
|
||||
n_blocks: Annotated[Optional[int], typer.Option()] = None,
|
||||
emb_dim: Annotated[Optional[int], typer.Option()] = None,
|
||||
val_fraction: Annotated[Optional[float], typer.Option()] = None,
|
||||
seed: Annotated[Optional[int], typer.Option(help="Random seed for reproducibility")] = None,
|
||||
validate_every: Annotated[Optional[int], typer.Option(help="Run marginal+KL validation every N epochs (0 disables)")] = None,
|
||||
shuffle_buffer: Annotated[int, typer.Option(help="Rows held in RAM per worker for shuffling")] = 65536,
|
||||
out: Annotated[Optional[Path], typer.Option(help="Checkpoint dir (default: auto from hyperparams)")] = None,
|
||||
device: Annotated[Optional[str], typer.Option(help="cpu | cuda | mps (default: auto)")] = None,
|
||||
seed: Annotated[
|
||||
Optional[int], typer.Option(help="Random seed for reproducibility")
|
||||
] = None,
|
||||
validate_every: Annotated[
|
||||
Optional[int],
|
||||
typer.Option(help="Run marginal+KL validation every N epochs (0 disables)"),
|
||||
] = None,
|
||||
shuffle_buffer: Annotated[
|
||||
int, typer.Option(help="Rows held in RAM per worker for shuffling")
|
||||
] = 65536,
|
||||
out: Annotated[
|
||||
Optional[Path],
|
||||
typer.Option(help="Checkpoint dir (default: auto from hyperparams)"),
|
||||
] = None,
|
||||
device: Annotated[
|
||||
Optional[str], typer.Option(help="cpu | cuda | mps (default: auto)")
|
||||
] = None,
|
||||
num_workers: Annotated[Optional[int], typer.Option()] = None,
|
||||
resume: Annotated[Optional[Path], typer.Option(help="Checkpoint .pt to resume training from")] = None,
|
||||
resume: Annotated[
|
||||
Optional[Path], typer.Option(help="Checkpoint .pt to resume training from")
|
||||
] = None,
|
||||
) -> None:
|
||||
"""Train the GIANT surrogate model."""
|
||||
cli_train = {k: v for k, v in {
|
||||
"mode": mode.value if mode is not None else None,
|
||||
"epochs": epochs, "batch_size": batch_size, "lr": lr,
|
||||
"val_fraction": val_fraction, "num_workers": num_workers, "seed": seed,
|
||||
"validate_every": validate_every,
|
||||
}.items() if v is not None}
|
||||
cli_model = {k: v for k, v in {
|
||||
"hidden_dim": hidden_dim, "n_blocks": n_blocks, "emb_dim": emb_dim,
|
||||
}.items() if v is not None}
|
||||
cfg = gconfig.merge_cli_overrides(gconfig.DEFAULT_CONFIG, config, cli_train, cli_model)
|
||||
cli_train = {
|
||||
k: v
|
||||
for k, v in {
|
||||
"mode": mode.value if mode is not None else None,
|
||||
"epochs": epochs,
|
||||
"batch_size": batch_size,
|
||||
"lr": lr,
|
||||
"val_fraction": val_fraction,
|
||||
"num_workers": num_workers,
|
||||
"seed": seed,
|
||||
"validate_every": validate_every,
|
||||
}.items()
|
||||
if v is not None
|
||||
}
|
||||
cli_model = {
|
||||
k: v
|
||||
for k, v in {
|
||||
"hidden_dim": hidden_dim,
|
||||
"n_blocks": n_blocks,
|
||||
"emb_dim": emb_dim,
|
||||
}.items()
|
||||
if v is not None
|
||||
}
|
||||
cfg = gconfig.merge_cli_overrides(
|
||||
gconfig.DEFAULT_CONFIG, config, cli_train, cli_model
|
||||
)
|
||||
t, m = cfg["train"], cfg["model"]
|
||||
|
||||
_device = torch.device(device) if device else gconfig.auto_device()
|
||||
@@ -99,26 +136,45 @@ def train(
|
||||
typer.echo(f"out_dir: {out_dir}")
|
||||
|
||||
run_train_job(
|
||||
data=data, cfg=cfg, out_dir=out_dir, device=_device,
|
||||
shuffle_buffer=shuffle_buffer, num_workers=t["num_workers"],
|
||||
resume=resume, echo=typer.echo,
|
||||
data=data,
|
||||
cfg=cfg,
|
||||
out_dir=out_dir,
|
||||
device=_device,
|
||||
shuffle_buffer=shuffle_buffer,
|
||||
num_workers=t["num_workers"],
|
||||
resume=resume,
|
||||
echo=typer.echo,
|
||||
)
|
||||
|
||||
|
||||
@app.command()
|
||||
def predict(
|
||||
data: Annotated[Path, typer.Argument(help="Parquet file or directory of parquet files")],
|
||||
checkpoint: Annotated[Path, typer.Option(help="Path to checkpoint .pt file (best.pt or last.pt)")],
|
||||
coord: Annotated[Coord, typer.Option(
|
||||
help="global: full physical units, world frame (default). "
|
||||
"local: raw 9D model output (denormalised only, local frame, "
|
||||
"log-scaled scalars) alongside the matching ground-truth target "
|
||||
"for the same input file — requires post-step columns."
|
||||
)] = Coord.global_,
|
||||
out: Annotated[Optional[Path], typer.Option(help="Output parquet path (default: <data>_predicted[_local].parquet)")] = None,
|
||||
data: Annotated[
|
||||
Path, typer.Argument(help="Parquet file or directory of parquet files")
|
||||
],
|
||||
checkpoint: Annotated[
|
||||
Path, typer.Option(help="Path to checkpoint .pt file (best.pt or last.pt)")
|
||||
],
|
||||
coord: Annotated[
|
||||
Coord,
|
||||
typer.Option(
|
||||
help="global: full physical units, world frame (default). "
|
||||
"local: raw 9D model output (denormalised only, local frame, "
|
||||
"log-scaled scalars) alongside the matching ground-truth target "
|
||||
"for the same input file — requires post-step columns."
|
||||
),
|
||||
] = Coord.global_,
|
||||
out: Annotated[
|
||||
Optional[Path],
|
||||
typer.Option(
|
||||
help="Output parquet path (default: <data>_predicted[_local].parquet)"
|
||||
),
|
||||
] = None,
|
||||
batch_size: Annotated[int, typer.Option(help="Inference batch size")] = 4096,
|
||||
steps: Annotated[int, typer.Option(help="Flow matching ODE steps")] = 10,
|
||||
device: Annotated[Optional[str], typer.Option(help="cpu | cuda | mps (default: auto)")] = None,
|
||||
device: Annotated[
|
||||
Optional[str], typer.Option(help="cpu | cuda | mps (default: auto)")
|
||||
] = None,
|
||||
) -> None:
|
||||
"""Run trained model on a parquet file and save predictions."""
|
||||
_device = torch.device(device) if device else gconfig.auto_device()
|
||||
@@ -127,7 +183,10 @@ def predict(
|
||||
# --- Load checkpoint ---
|
||||
ckpt = torch.load(checkpoint, map_location="cpu", weights_only=False)
|
||||
if "model_config" not in ckpt:
|
||||
typer.echo("error: checkpoint has no model_config — retrain with the current code", err=True)
|
||||
typer.echo(
|
||||
"error: checkpoint has no model_config — retrain with the current code",
|
||||
err=True,
|
||||
)
|
||||
raise typer.Exit(1)
|
||||
|
||||
model_cfg = ckpt["model_config"]
|
||||
@@ -145,7 +204,9 @@ def predict(
|
||||
# --- Output path ---
|
||||
if out is None:
|
||||
stem = data.stem if data.is_file() else data.name
|
||||
suffix = "_predicted_local.parquet" if coord == Coord.local else "_predicted.parquet"
|
||||
suffix = (
|
||||
"_predicted_local.parquet" if coord == Coord.local else "_predicted.parquet"
|
||||
)
|
||||
out = data.parent / f"{stem}{suffix}"
|
||||
typer.echo(f"output: {out}")
|
||||
|
||||
@@ -162,10 +223,14 @@ def predict(
|
||||
N = len(chunk["event_id"])
|
||||
|
||||
if coord == Coord.local:
|
||||
cond_cont, cond_cat, target_raw, _, _ = build_features(chunk, pdg_map, mat_map)
|
||||
cond_cont, cond_cat, target_raw, _, _ = build_features(
|
||||
chunk, pdg_map, mat_map
|
||||
)
|
||||
cond_cont = cond_norm.transform(cond_cont)
|
||||
else:
|
||||
cond_cont, cond_cat = build_cond_features(chunk, pdg_map, mat_map, cond_norm)
|
||||
cond_cont, cond_cat = build_cond_features(
|
||||
chunk, pdg_map, mat_map, cond_norm
|
||||
)
|
||||
|
||||
# Inference in batch_size slices
|
||||
pred_parts = []
|
||||
@@ -180,22 +245,30 @@ def predict(
|
||||
raw = tgt_norm.inverse_transform(pred)
|
||||
|
||||
if coord == Coord.local:
|
||||
table = pa.table({
|
||||
"event_id": chunk["event_id"],
|
||||
"pdg": chunk["pdg"],
|
||||
"pre_x": chunk["pre_pos"][:, 0],
|
||||
"pre_y": chunk["pre_pos"][:, 1],
|
||||
"pre_z": chunk["pre_pos"][:, 2],
|
||||
"pre_E": chunk["pre_E"],
|
||||
"pre_dx": chunk["pre_dir"][:, 0],
|
||||
"pre_dy": chunk["pre_dir"][:, 1],
|
||||
"pre_dz": chunk["pre_dir"][:, 2],
|
||||
"material": chunk["material"],
|
||||
"layer_id": chunk["layer_id"],
|
||||
"n_sec": chunk["n_sec"],
|
||||
**{f"pred_{name}": raw[:, j] for j, name in enumerate(LOCAL_TARGET_NAMES)},
|
||||
**{f"true_{name}": target_raw[:, j] for j, name in enumerate(LOCAL_TARGET_NAMES)},
|
||||
})
|
||||
table = pa.table(
|
||||
{
|
||||
"event_id": chunk["event_id"],
|
||||
"pdg": chunk["pdg"],
|
||||
"pre_x": chunk["pre_pos"][:, 0],
|
||||
"pre_y": chunk["pre_pos"][:, 1],
|
||||
"pre_z": chunk["pre_pos"][:, 2],
|
||||
"pre_E": chunk["pre_E"],
|
||||
"pre_dx": chunk["pre_dir"][:, 0],
|
||||
"pre_dy": chunk["pre_dir"][:, 1],
|
||||
"pre_dz": chunk["pre_dir"][:, 2],
|
||||
"material": chunk["material"],
|
||||
"layer_id": chunk["layer_id"],
|
||||
"n_sec": chunk["n_sec"],
|
||||
**{
|
||||
f"pred_{name}": raw[:, j]
|
||||
for j, name in enumerate(LOCAL_TARGET_NAMES)
|
||||
},
|
||||
**{
|
||||
f"true_{name}": target_raw[:, j]
|
||||
for j, name in enumerate(LOCAL_TARGET_NAMES)
|
||||
},
|
||||
}
|
||||
)
|
||||
else:
|
||||
step_length = inv_log_transform(raw[:, 0])
|
||||
delta_e = inv_log_transform(raw[:, 1])
|
||||
@@ -205,7 +278,9 @@ def predict(
|
||||
post_dir_local = raw[:, 3:6].copy()
|
||||
norms = np.linalg.norm(post_dir_local, axis=1, keepdims=True)
|
||||
post_dir_local /= np.where(norms < 1e-8, 1.0, norms)
|
||||
post_dir_world = inv_local_frame_rotation(chunk["pre_dir"], post_dir_local)
|
||||
post_dir_world = inv_local_frame_rotation(
|
||||
chunk["pre_dir"], post_dir_local
|
||||
)
|
||||
|
||||
# Same for the travel direction, then reconstruct post_pos from
|
||||
# the single shared step_length so the two stay consistent.
|
||||
@@ -216,34 +291,38 @@ def predict(
|
||||
chunk["pre_pos"], chunk["pre_dir"], step_length, travel_dir_local
|
||||
)
|
||||
|
||||
table = pa.table({
|
||||
"event_id": chunk["event_id"],
|
||||
"pdg": chunk["pdg"],
|
||||
"pre_x": chunk["pre_pos"][:, 0],
|
||||
"pre_y": chunk["pre_pos"][:, 1],
|
||||
"pre_z": chunk["pre_pos"][:, 2],
|
||||
"pre_E": chunk["pre_E"],
|
||||
"pre_dx": chunk["pre_dir"][:, 0],
|
||||
"pre_dy": chunk["pre_dir"][:, 1],
|
||||
"pre_dz": chunk["pre_dir"][:, 2],
|
||||
"material": chunk["material"],
|
||||
"layer_id": chunk["layer_id"],
|
||||
"n_sec": chunk["n_sec"],
|
||||
"step_length": step_length,
|
||||
"delta_e": delta_e,
|
||||
"edep": edep,
|
||||
"post_dx": post_dir_world[:, 0],
|
||||
"post_dy": post_dir_world[:, 1],
|
||||
"post_dz": post_dir_world[:, 2],
|
||||
"post_x": post_pos_world[:, 0],
|
||||
"post_y": post_pos_world[:, 1],
|
||||
"post_z": post_pos_world[:, 2],
|
||||
})
|
||||
table = pa.table(
|
||||
{
|
||||
"event_id": chunk["event_id"],
|
||||
"pdg": chunk["pdg"],
|
||||
"pre_x": chunk["pre_pos"][:, 0],
|
||||
"pre_y": chunk["pre_pos"][:, 1],
|
||||
"pre_z": chunk["pre_pos"][:, 2],
|
||||
"pre_E": chunk["pre_E"],
|
||||
"pre_dx": chunk["pre_dir"][:, 0],
|
||||
"pre_dy": chunk["pre_dir"][:, 1],
|
||||
"pre_dz": chunk["pre_dir"][:, 2],
|
||||
"material": chunk["material"],
|
||||
"layer_id": chunk["layer_id"],
|
||||
"n_sec": chunk["n_sec"],
|
||||
"step_length": step_length,
|
||||
"delta_e": delta_e,
|
||||
"edep": edep,
|
||||
"post_dx": post_dir_world[:, 0],
|
||||
"post_dy": post_dir_world[:, 1],
|
||||
"post_dz": post_dir_world[:, 2],
|
||||
"post_x": post_pos_world[:, 0],
|
||||
"post_y": post_pos_world[:, 1],
|
||||
"post_z": post_pos_world[:, 2],
|
||||
}
|
||||
)
|
||||
|
||||
table = table.replace_schema_metadata({
|
||||
PREDICT_COORD_METADATA_KEY: coord.value,
|
||||
PREDICT_SCHEMA_VERSION_KEY: PREDICT_SCHEMA_VERSION,
|
||||
})
|
||||
table = table.replace_schema_metadata(
|
||||
{
|
||||
PREDICT_COORD_METADATA_KEY: coord.value,
|
||||
PREDICT_SCHEMA_VERSION_KEY: PREDICT_SCHEMA_VERSION,
|
||||
}
|
||||
)
|
||||
|
||||
if writer is None:
|
||||
writer = pq.ParquetWriter(out, table.schema)
|
||||
@@ -255,5 +334,6 @@ def predict(
|
||||
|
||||
typer.echo(f"wrote {total:,} rows → {out}")
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
app()
|
||||
|
||||
Reference in New Issue
Block a user