Get example 1 working

python/_util.py

@@ -6,6 +6,12 @@
 import pathlib
 import importlib
 
+import time
+from siren import dataclasses as _dataclasses
+import numpy as np
+import awkward as ak
+import h5py
+
 THIS_DIR = os.path.abspath(os.path.dirname(__file__))
 
 
@@ -661,6 +667,7 @@ def import_resource(resource_type, resource_name):
     fname = os.path.join(abs_dir, f"{resource_type}.py")
     if not os.path.isfile(fname):
         return None
+    print(fname)
     return load_module(f"siren-{resource_type}-{resource_name}", fname, persist=False)
 
 
@@ -835,3 +842,127 @@ def load_detector():
     load_detector.__doc__ = detector_docs(detector_name)
 
     return load_detector
+
+
+
+###### Injector helper functions #######
+
+# Generate events using an injector object
+# Optionally save events to hdf5, parquet, and/or custom SIREN filetypes
+# If the weighter exists, calculate the event weight too
+def GenerateEvents(injector, N=None):
+    if N is None:
+        N = injector.number_of_events
+    count = 0
+    gen_times = []
+    prev_time = time.time()
+    events = []
+    while (injector.injected_events < injector.number_of_events) and (count < N):
+        print("Injecting Event %d/%d  " % (count, N), end="\r")
+        event = injector.generate_event()
+        events.append(event)
+        t = time.time()
+        gen_times.append(t-prev_time)
+        prev_time = t
+        count += 1
+    return events,gen_times
+
+def SaveEvents(events,
+               weighter=None,
+               gen_times=None,
+               save_hdf5=True,
+               save_parquet=True,
+               save_siren_events=True,
+               fid_vol=None,
+               output_filename=None):
+
+
+    # Optionally save things
+    if save_siren_events: _dataclasses.SaveInteractionTrees(events, output_filename)
+    # A dictionary containing each dataset we'd like to save
+    datasets = {
+        "event_weight":[], # weight of entire event
+        "event_gen_time":[], # generation time of each event
+        "event_weight_time":[], # generation time of each event
+        "num_interactions":[], # number of interactions per event
+        "vertex":[], # vertex of each interaction in an event
+        "in_fiducial":[], # whether or not each vertex is in the fiducial volume
+        "primary_type":[], # primary type of each interaction
+        "target_type":[], # target type of each interaction
+        "num_secondaries":[], # number of secondary particles of each interaction
+        "secondary_types":[], # secondary type of each interaction
+        "primary_momentum":[], # primary momentum of each interaction
+        "secondary_momenta":[], # secondary momentum of each interaction
+        "parent_idx":[], # index of the parent interaction
+    }
+    for ie, event in enumerate(events):
+        print("Saving Event %d/%d  " % (ie, len(events)), end="\r")
+        t0 = time.time()
+        datasets["event_weight"].append(weighter(event) if weighter is not None else 0)
+        datasets["event_weight_time"].append(time.time()-t0)
+        datasets["event_gen_time"].append(gen_times[ie])
+        # add empty lists for each per interaction dataset
+        for k in ["vertex",
+                  "in_fiducial",
+                  "primary_type",
+                  "target_type",
+                  "num_secondaries",
+                  "secondary_types",
+                  "primary_momentum",
+                  "secondary_momenta",
+                  "parent_idx"]:
+            datasets[k].append([])
+        # loop over interactions
+        for id, datum in enumerate(event.tree):
+            datasets["vertex"][-1].append(np.array(datum.record.interaction_vertex,dtype=float))
+
+             # primary particle stuff
+            datasets["primary_type"][-1].append(int(datum.record.signature.primary_type))
+            datasets["primary_momentum"][-1].append(np.array(datum.record.primary_momentum, dtype=float))
+
+            # check parent idx; match on secondary momenta
+            if datum.depth()==0:
+                datasets["parent_idx"][-1].append(-1)
+            else:
+                for _id in range(len(datasets["secondary_momenta"][-1])):
+                    for secondary_momentum in datasets["secondary_momenta"][-1][_id]:
+                        if (datasets["primary_momentum"][-1][-1] == secondary_momentum).all():
+                            datasets["parent_idx"][-1].append(_id)
+                            break
+
+            if fid_vol is not None:
+                pos = _math.Vector3D(datasets["vertex"][-1][-1])
+                dir = _math.Vector3D(datasets["primary_momentum"][-1][-1][1:])
+                dir.normalize()
+                datasets["in_fiducial"][-1].append(fid_vol.IsInside(pos,dir))
+            else:
+                datasets["in_fiducial"][-1].append(False)
+
+            # target particle stuff
+            datasets["target_type"][-1].append(int(datum.record.signature.target_type))
+
+            # secondary particle stuff
+            datasets["secondary_types"][-1].append([])
+            datasets["secondary_momenta"][-1].append([])
+            for isec, (sec_type, sec_momenta) in enumerate(zip(datum.record.signature.secondary_types,
+                                                               datum.record.secondary_momenta)):
+                datasets["secondary_types"][-1][-1].append(int(sec_type))
+                datasets["secondary_momenta"][-1][-1].append(np.array(sec_momenta,dtype=float))
+            datasets["num_secondaries"][-1].append(isec+1)
+        datasets["num_interactions"].append(id+1)
+
+    # save events
+    ak_array = ak.Array(datasets)
+    if save_hdf5:
+        fout = h5py.File(output_filename+".hdf5", "w")
+        group = fout.create_group("Events")
+        form, length, container = ak.to_buffers(ak.to_packed(ak_array), container=group)
+        group.attrs["form"] = form.to_json()
+        group.attrs["length"] = length
+        fout.close()
+    if save_parquet:
+        ak.to_parquet(ak_array,output_filename+".parquet")
+
+# Load events from the custom SIREN event format
+def LoadEvents(filename):
+    return _dataclasses.LoadInteractionTrees(filename)

resources/examples/example1/DIS_ATLAS.py

@@ -1,10 +1,6 @@
 import os
 import siren
-try:
-    from tqdm import tqdm as tqdm
-except ImportError:
-    print("Importing tqdm failed, using default range")
-    tqdm = lambda x: x
+from siren._util import GenerateEvents,SaveEvents
 
 seed = 99
 
@@ -86,25 +82,18 @@
 injector.primary_type = primary_type
 injector.primary_interactions = primary_processes[primary_type]
 injector.primary_injection_distributions = primary_injection_distributions
-injector.secondary_interactions = {}
-injector.secondary_injection_distributions = {}
 
-print("Generating events")
-events = [injector.generate_event() for _ in tqdm(range(events_to_inject))]
+events,gen_times = GenerateEvents(injector)
 
 weighter = siren.injection.Weighter()
 weighter.injectors = [injector]
 weighter.detector_model = detector_model
 weighter.primary_type = primary_type
 weighter.primary_interactions = primary_processes[primary_type]
 weighter.primary_physical_distributions = primary_physical_distributions
-weighter.secondary_interactions = {}
-weighter.secondary_physical_distributions = {}
-
-print("Weighting events")
-weights = [weighter(event) for event in tqdm(events)]
 
+# Output events and weights
 os.makedirs("output", exist_ok=True)
-
-#TODO save the events and weights
+print("Saving events")
+SaveEvents(events,weighter,gen_times,output_filename="output/ATLAS")
 

resources/examples/example1/DIS_DUNE.py

@@ -1,79 +1,64 @@
 import os
 import siren
-from siren import utilities
+from siren._util import GenerateEvents,SaveEvents
 
 # Number of events to inject
 events_to_inject = int(1e5)
 
 # Experiment to run
 experiment = "DUNEFD"
-detector_model = utilities.load_detector(experiment)
+detector_model = siren.utilities.load_detector(experiment)
+
+# Particle to inject
 primary_type = siren.dataclasses.Particle.ParticleType.NuMu
-target_type = siren.dataclasses.Particle.ParticleType.Nucleon
 
-# Primary interactions and distributions
-primary_processes, _ = utilities.load_processes(
-    "CSMSDISSplines",          # model_name
+# Cross-section model to use
+cross_section_model = "CSMSDISSplines"
+
+# Load the cross-section model
+primary_processes, _ = siren.utilities.load_processes(
+    cross_section_model,
     primary_types=[primary_type],
-    target_types=[target_type],
-    isoscalar=True,             # for isoscalar splines
-    process_types=["CC"]         # specify the process type, e.g., "CC" for charged current
+    target_types=[siren.dataclasses.Particle.ParticleType.Nucleon],
+    isoscalar=True,
+    process_types=["CC"]
 )
 
-# Energy distribution
-edist = siren.distributions.PowerLaw(1, 1e3, 1e6)
-
-# Direction distribution
-direction_distribution = siren.distributions.IsotropicDirection()
-
-# Position distribution
-muon_range_func = siren.distributions.LeptonDepthFunction()
-position_distribution = siren.distributions.ColumnDepthPositionDistribution(
-    60, 60.0, muon_range_func)
-
-
-# Define injection distributions
-primary_injection_distributions = {
-    "energy": edist,
-    "direction": direction_distribution,
-    "position": position_distribution
-}
+# Extract the primary cross-sections for the primary type
+primary_cross_sections = primary_processes[primary_type]
 
 # Set up the Injector
 injector = siren.injection.Injector()
 injector.number_of_events = events_to_inject
 injector.detector_model = detector_model
 injector.primary_type = primary_type
-injector.primary_interactions = primary_processes[primary_type]
+injector.primary_interactions = primary_cross_sections
+
+# Directly set the distributions
 injector.primary_injection_distributions = [
-    siren.distributions.PrimaryMass(0),
-    edist,
-    direction_distribution,
-    position_distribution
+    siren.distributions.PrimaryMass(0),  # Mass distribution
+    siren.distributions.PowerLaw(1, 1e3, 1e6),  # Energy distribution
+    siren.distributions.IsotropicDirection(),  # Direction distribution
+    siren.distributions.ColumnDepthPositionDistribution(60, 60.0, siren.distributions.LeptonDepthFunction())  # Position distribution
 ]
 
 # Generate events
-event = injector.generate_event()
+events,gen_times = GenerateEvents(injector)
 
-# Set up the Weighter for event weighting
+# Set up the Weighter for event weighting (without position distribution)
 weighter = siren.injection.Weighter()
 weighter.injectors = [injector]
 weighter.detector_model = detector_model
 weighter.primary_type = primary_type
-weighter.primary_interactions = primary_processes[primary_type]
+weighter.primary_interactions = primary_cross_sections
 weighter.primary_physical_distributions = [
-    edist,
-    direction_distribution,
+    siren.distributions.PowerLaw(1, 1e3, 1e6),  # Energy distribution
+    siren.distributions.IsotropicDirection()  # Direction distribution
 ]
 
-# Compute weight
-weight = weighter(event)
 
 # Output events and weights
 os.makedirs("output", exist_ok=True)
-print(str(event))
-print(f"Event weight: {weight}")
+SaveEvents(events,weighter,gen_times,output_filename="output/DUNE")
 
-# Save events
-# TODO
 

resources/examples/example1/DIS_IceCube.py

@@ -1,13 +1,13 @@
 import os
 import siren
-from siren import utilities
+from siren._util import GenerateEvents,SaveEvents
 
 # Number of events to inject
 events_to_inject = int(1e5)
 
 # Experiment to run
 experiment = "IceCube"
-detector_model = utilities.load_detector(experiment)
+detector_model = siren.utilities.load_detector(experiment)
 
 # Particle to inject
 primary_type = siren.dataclasses.Particle.ParticleType.NuMu
@@ -16,7 +16,7 @@
 cross_section_model = "CSMSDISSplines"
 
 # Load the cross-section model
-primary_processes, _ = utilities.load_processes(
+primary_processes, _ = siren.utilities.load_processes(
     cross_section_model,
     primary_types=[primary_type],
     target_types=[siren.dataclasses.Particle.ParticleType.Nucleon],
@@ -43,7 +43,7 @@
 ]
 
 # Generate events
-event = injector.generate_event()
+events,gen_times = GenerateEvents(injector)
 
 # Set up the Weighter for event weighting (without position distribution)
 weighter = siren.injection.Weighter()
@@ -56,14 +56,7 @@
     siren.distributions.IsotropicDirection()  # Direction distribution
 ]
 
-# Compute weight
-weight = weighter(event)
 
 # Output events and weights
 os.makedirs("output", exist_ok=True)
-print(str(event))
-print(f"Event weight: {weight}")
-
-# Save events
-# injector.SaveEvents("output/IceCube_DIS")
-
+SaveEvents(events,weighter,gen_times,output_filename="output/IceCube")