diff --git a/python/LIDarkNews.py b/python/LIDarkNews.py index da91ad95..31701af3 100644 --- a/python/LIDarkNews.py +++ b/python/LIDarkNews.py @@ -831,7 +831,7 @@ def SampleRecordFromDarkNews(self, record, random): np.expand_dims(np.array(record.primary_momentum), 0), ) - secondaries = record.GetSecondaryParticles() + secondaries = record.GetSecondaryParticleRecords() if type(self.dec_case) == FermionSinglePhotonDecay: gamma_idx = 0 @@ -843,13 +843,13 @@ def SampleRecordFromDarkNews(self, record, random): print("No gamma found in the list of secondaries!") exit(0) nu_idx = 1 - gamma_idx - secondaries[gamma_idx].momentum = list( - np.squeeze(four_momenta["P_decay_photon"]) - ) + secondaries[gamma_idx].four_momentum = np.squeeze(four_momenta["P_decay_photon"]) secondaries[gamma_idx].mass = 0 - secondaries[nu_idx].momentum = list( - np.squeeze(four_momenta["P_decay_N_daughter"]) - ) + secondaries[nu_idx].four_momentum = np.squeeze(four_momenta["P_decay_N_daughter"]) + secondaries[nu_idx].mass = 0 + + print("P_gamma", secondaries[gamma_idx].four_momentum) + print("P_nu", secondaries[nu_idx].four_momentum) elif type(self.dec_case) == FermionDileptonDecay: lepminus_idx = -1 @@ -876,14 +876,13 @@ def SampleRecordFromDarkNews(self, record, random): print("Couldn't find two leptons and a neutrino in the final state!") exit(0) secondary_momenta = [] - seconaries[lepminus_idx].momentum = list( + seconaries[lepminus_idx].four_momentum = ( np.squeeze(four_momenta["P_decay_ell_minus"]) ) - secondaries[lepplus_idx].momentum = list( + secondaries[lepplus_idx].four_momentum = ( np.squeeze(four_momenta["P_decay_ell_plus"]) ) - secondaries[nu_idx].momentum = list( + secondaries[nu_idx].four_momentum = ( np.squeeze(four_momenta["P_decay_N_daughter"]) ) - record.SetSecondaryParticles(secondaries) return record diff --git a/resources/Examples/Example1/DIS_IceCube.py b/resources/Examples/Example1/DIS_IceCube.py index 99ef3a75..302175e6 100644 --- a/resources/Examples/Example1/DIS_IceCube.py +++ b/resources/Examples/Example1/DIS_IceCube.py @@ -1,15 +1,33 @@ import os import sys import numpy as np +import functools import leptoninjector as LI from leptoninjector import _util from leptoninjector.LIController import LIController +@functools.wraps(LIController.GenerateEvents) +def GenerateEvents(self, N=None): + if N is None: + N = self.events_to_inject + count = 0 + while (self.injector.InjectedEvents() < self.events_to_inject) and (count < N): + print("Injecting Event", count, end="\r") + tree = self.injector.GenerateEvent() + self.weighter.EventWeight(tree) + self.events.append(tree) + count += 1 + #if hasattr(self, "DN_processes"): + # self.DN_processes.SaveCrossSectionTables() + return self.events + +LIController.GenerateEvents = GenerateEvents + resources_dir = _util.resource_package_dir() # Number of events to inject -events_to_inject = 1000 +events_to_inject = int(1e6) # Expeirment to run experiment = "IceCube"