diff --git a/giant/cli.py b/giant/cli.py index 3e7bf5f..2f98f9d 100644 --- a/giant/cli.py +++ b/giant/cli.py @@ -350,6 +350,118 @@ def predict( total_rows = sum(pq.ParquetFile(path).metadata.num_rows for path in files) chunk_iter = iter_file_chunks if coord == Coord.local else iter_cond_chunks + def _concat( + a: dict[str, np.ndarray], b: dict[str, np.ndarray] + ) -> dict[str, np.ndarray]: + return {k: np.concatenate([a[k], b[k]], axis=0) for k in a} + + def _process(piece: dict[str, np.ndarray]) -> None: + nonlocal writer, total + + if coord == Coord.local: + cond_cont, cond_cat, target_raw, _, _ = build_features( + piece, pdg_map, mat_map + ) + cond_cont = cond_norm.transform(cond_cont) + else: + cond_cont, cond_cat = build_cond_features( + piece, pdg_map, mat_map, cond_norm + ) + + cc = torch.from_numpy(cond_cont).float().to(_device) + ck = torch.from_numpy(cond_cat).long().to(_device) + pred = sample_flow(model, cc, ck, steps=steps).cpu().numpy() # normalised + + # Inverse-normalise → local frame, log-scaled scalars + raw = tgt_norm.inverse_transform(pred) + + if coord == Coord.local: + table = pa.table( + { + "event_id": piece["event_id"], + "pdg": piece["pdg"], + "pre_x": piece["pre_pos"][:, 0], + "pre_y": piece["pre_pos"][:, 1], + "pre_z": piece["pre_pos"][:, 2], + "pre_E": piece["pre_E"], + "pre_dx": piece["pre_dir"][:, 0], + "pre_dy": piece["pre_dir"][:, 1], + "pre_dz": piece["pre_dir"][:, 2], + "material": piece["material"], + "layer_id": piece["layer_id"], + "n_sec": piece["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]) + edep = inv_log_transform(raw[:, 2]) + + # Normalise predicted direction then rotate back to world frame + 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(piece["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. + travel_dir_local = raw[:, 6:9].copy() + norms = np.linalg.norm(travel_dir_local, axis=1, keepdims=True) + travel_dir_local /= np.where(norms < 1e-8, 1.0, norms) + post_pos_world = reconstruct_post_pos( + piece["pre_pos"], piece["pre_dir"], step_length, travel_dir_local + ) + + table = pa.table( + { + "event_id": piece["event_id"], + "pdg": piece["pdg"], + "pre_x": piece["pre_pos"][:, 0], + "pre_y": piece["pre_pos"][:, 1], + "pre_z": piece["pre_pos"][:, 2], + "pre_E": piece["pre_E"], + "pre_dx": piece["pre_dir"][:, 0], + "pre_dy": piece["pre_dir"][:, 1], + "pre_dz": piece["pre_dir"][:, 2], + "material": piece["material"], + "layer_id": piece["layer_id"], + "n_sec": piece["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, + } + ) + + if writer is None: + writer = pq.ParquetWriter(out, table.schema) + writer.write_table(table) + total += len(piece["event_id"]) + + # Buffer rows across row-group boundaries so the inference batch size + # isn't capped by however the source file happens to be chunked. + buffer: dict[str, np.ndarray] | None = None + bar = tqdm(total=total_rows, desc="predict", unit="row", dynamic_ncols=True) for path in files: for chunk in chunk_iter(path): @@ -360,119 +472,19 @@ def predict( unknown_pdg_counts.update(int(p) for p in chunk["pdg"][~pdg_mask]) chunk = {k: v[pdg_mask] for k, v in chunk.items()} - N = len(chunk["event_id"]) - skipped += N_in - N - if N == 0: - bar.update(N_in) + skipped += N_in - len(chunk["event_id"]) + bar.update(N_in) + if len(chunk["event_id"]) == 0: continue - if coord == Coord.local: - 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 - ) + buffer = chunk if buffer is None else _concat(buffer, chunk) + while len(buffer["event_id"]) >= bs: + piece = {k: v[:bs] for k, v in buffer.items()} + buffer = {k: v[bs:] for k, v in buffer.items()} + _process(piece) - # Inference in batch_size slices - pred_parts = [] - for start in range(0, N, bs): - end = min(start + bs, N) - cc = torch.from_numpy(cond_cont[start:end]).float().to(_device) - ck = torch.from_numpy(cond_cat[start:end]).long().to(_device) - pred_parts.append(sample_flow(model, cc, ck, steps=steps).cpu().numpy()) - pred = np.concatenate(pred_parts, axis=0) # (N, 9) normalised - - # Inverse-normalise → local frame, log-scaled scalars - 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) - }, - } - ) - else: - step_length = inv_log_transform(raw[:, 0]) - delta_e = inv_log_transform(raw[:, 1]) - edep = inv_log_transform(raw[:, 2]) - - # Normalise predicted direction then rotate back to world frame - 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 - ) - - # Same for the travel direction, then reconstruct post_pos from - # the single shared step_length so the two stay consistent. - travel_dir_local = raw[:, 6:9].copy() - norms = np.linalg.norm(travel_dir_local, axis=1, keepdims=True) - travel_dir_local /= np.where(norms < 1e-8, 1.0, norms) - post_pos_world = reconstruct_post_pos( - 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 = 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) - writer.write_table(table) - total += N - bar.update(N_in) + if buffer is not None and len(buffer["event_id"]) > 0: + _process(buffer) bar.close() if writer is not None: