File size: 2,485 Bytes
c93e291 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 | #!/usr/bin/env python3
"""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())
|