| |
| """Combine CP-even and CP-odd artifacts into labeled binary artifacts.""" |
|
|
| from __future__ import annotations |
|
|
| import argparse |
| from pathlib import Path |
|
|
| import numpy as np |
|
|
| from validation.artifacts import ValidationArtifact, load_artifact, save_artifact |
|
|
|
|
| def _combine(even: ValidationArtifact, odd: ValidationArtifact) -> ValidationArtifact: |
| node_shift = len(even.node_features_flat) |
| edge_shift = len(even.edge_src_flat) |
| sample_id = np.concatenate( |
| [ |
| np.char.add("ttH_CP_even:", even.sample_id), |
| np.char.add("ttH_CP_odd:", odd.sample_id), |
| ] |
| ) |
| labels = np.concatenate( |
| [ |
| np.zeros(even.event_count, dtype=np.int64), |
| np.ones(odd.event_count, dtype=np.int64), |
| ] |
| ) |
| node_offsets = np.concatenate( |
| [even.node_offsets, odd.node_offsets[1:] + node_shift] |
| ) |
| edge_offsets = np.concatenate( |
| [even.edge_offsets, odd.edge_offsets[1:] + edge_shift] |
| ) |
| return ValidationArtifact( |
| sample_id=sample_id, |
| labels=labels, |
| folds=np.concatenate([even.folds, odd.folds]), |
| weights=np.concatenate([even.weights, odd.weights]), |
| globals=np.concatenate([even.globals, odd.globals]), |
| node_features_flat=np.concatenate( |
| [even.node_features_flat, odd.node_features_flat] |
| ), |
| node_offsets=node_offsets, |
| edge_src_flat=np.concatenate([even.edge_src_flat, odd.edge_src_flat]), |
| edge_dst_flat=np.concatenate([even.edge_dst_flat, odd.edge_dst_flat]), |
| edge_features_flat=np.concatenate( |
| [even.edge_features_flat, odd.edge_features_flat] |
| ), |
| edge_offsets=edge_offsets, |
| manifest={ |
| "task": "ttH_cp_even_vs_cp_odd", |
| "label_map": {"CP_even": 0, "CP_odd": 1}, |
| "selection": "none", |
| "sources": {"CP_even": even.manifest, "CP_odd": odd.manifest}, |
| }, |
| ) |
|
|
|
|
| def main() -> int: |
| parser = argparse.ArgumentParser() |
| parser.add_argument("--even", type=Path, required=True) |
| parser.add_argument("--odd", type=Path, required=True) |
| parser.add_argument("--output", type=Path, required=True) |
| args = parser.parse_args() |
| artifact = _combine(load_artifact(args.even), load_artifact(args.odd)) |
| save_artifact(artifact, args.output) |
| print(f"wrote {args.output} ({artifact.event_count} events)") |
| return 0 |
|
|
|
|
| if __name__ == "__main__": |
| raise SystemExit(main()) |
|
|
