diff --git a/km3buu/geometry.py b/km3buu/geometry.py
index 27c685ec2d2bdbad50b47a787d644b43e2568d10..e94ad97f14cef7a5e8a88e7e6356cd35726b4b13 100644
--- a/km3buu/geometry.py
+++ b/km3buu/geometry.py
@@ -13,6 +13,14 @@ __status__ = "Development"
import numpy as np
from abc import ABC, abstractmethod
+from collections import defaultdict
+from scipy.spatial.transform import Rotation
+from particle import Particle
+
+import proposal as pp
+from .propagation import *
+
+M_TO_CM = 1e2
class DetectorVolume(ABC):
@@ -26,10 +34,15 @@ class DetectorVolume(ABC):
self._solid_angle = 1
@abstractmethod
- def random_dir(self):
+ def random_dir(self, n=1):
"""
Generate a random direction for the interaction
+ Parameter
+ ---------
+ n: int (default: 1)
+ Number of vertex positions to sample
+
Return
------
tuple [rad, 1] (phi, cos(theta))
@@ -37,17 +50,41 @@ class DetectorVolume(ABC):
pass
@abstractmethod
- def random_pos(self):
+ def random_pos(self, n=1):
"""
Generate a random position in the detector volume based on a uniform
event distribution
+ Parameter
+ ---------
+ n: int (default: 1)
+ Number of vertex directions to sample
+
Returns
-------
tuple [m] (x, y, z)
"""
pass
+ def distribute_event(self, *pargs, **kwargs):
+ """
+ Integrated event distribution method which also handles propagation
+
+ Returns
+ -------
+ vtx_pos: tuple
+ Position of the vertex in the interaction volume
+ vtx_dir: tuple
+ Direction of the vertex in the interaction volume
+ weight_correction: float
+ If the sampling requires a correction to the weight this factor is
+ unequal 1
+ additional_events: awkward.highlevel.Array (
+ E, Px, Py, Pz, x, y, z, pdgid)
+ Propagated / Additional Particles from decays (None if no propagation is done)
+ """
+ return self.random_pos(), self.random_dir(), 1.0, None
+
@abstractmethod
def header_entries(self):
"""
@@ -77,7 +114,7 @@ class DetectorVolume(ABC):
Returns
-------
- float [m^3]
+ float [m^3]
"""
return self._volume
@@ -95,7 +132,7 @@ class DetectorVolume(ABC):
class NoVolume(DetectorVolume):
"""
- Dummy volume to write out data w/o geometry
+ Dummy volume to write out data w/o geometry
"""
def __init__(self):
@@ -110,16 +147,22 @@ class NoVolume(DetectorVolume):
retdct[key] = value
return retdct
- def random_pos(self):
- return (.0, .0, .0)
+ def random_pos(self, n=1):
+ if n == 1:
+ return np.zeros(3)
+ else:
+ return np.zeros((n, 3))
- def random_dir(self):
- return (0, 1)
+ def random_dir(self, n=1):
+ if n == 1:
+ return (0, 1)
+ else:
+ return np.concatenate([np.zeros(n), np.ones(n)]).reshape((2, -1)).T
class CanVolume(DetectorVolume):
"""
- Cylindrical detector geometry
+ Cylindrical detector geometry, only CAN / no propagation
Parameters
----------
@@ -174,9 +217,233 @@ class CanVolume(DetectorVolume):
self._volume, nevents)
retdct[key] = value
key = "fixedcan"
- value = "0 0 {} {} {}".format(self._zmin, self._zmax, self._radius)
+ value = "{} {} {} {} {}".format(detector_center[0], detector_center[1],
+ self._zmin, self._zmax, self._radius)
return retdct
+ def distribute_event(self, evt):
+ vtx_pos = self.random_pos()
+ vtx_dir = self.random_dir()
+ weight = 1
+ evts = None
+ return vtx_pos, vtx_dir, weight, evts
+
+
+class CylindricalVolume(DetectorVolume):
+ """
+ Cylindrical detector geometry
+
+ Parameters
+ ----------
+ radius: float [m] (default: 403.5)
+ Radius of the interaction volume
+ sw_height: float [m] (default: xxx)
+ Height of the seawater in the interaction volume
+ sr_height: float [m] (default: xxx)
+ Height of the seabottom rock in the interaction volume
+ can_radius: float [m] (default: 0.0)
+ Cylinder bottom z position
+ can_zmin: float [m] (default: 0.0)
+ Cylinder bottom z position
+ can_zmax: float [m] (default: 476.5)
+ Cylinder top z position
+ detector_center: tuple [m] (default: (0.0, 0.0) )
+ Detector center position in the xy-plane
+ zenith: float [1] (default: (-1.0, 1.0) )
+ Zenith range given as cos(θ)
+ """
+
+ def __init__(self,
+ radius=500.0,
+ sw_height=650.0,
+ sr_height=100.0,
+ can_radius=200.4,
+ can_zmin=0.0,
+ can_zmax=350,
+ detector_center=(0., 0.),
+ zenith=(-1, 1)):
+ super().__init__()
+ self._detector_center = detector_center
+ # Interaction Volume
+ self._radius = radius
+ self._water_height = sw_height
+ self._rock_height = sr_height
+ # Can Volume
+ self._canradius = can_radius
+ self._canzmin = can_zmin
+ self._canzmax = can_zmax
+ # Direction
+ self._cosZmin = zenith[0]
+ self._cosZmax = zenith[1]
+ # Properties
+ self._solid_angle = 2 * np.pi * (self._cosZmax - self._cosZmin)
+ self._volume = self._calc_volume()
+ # PROPOSAL
+ self._pp_geometry = None
+ self._propagator = None
+
+ def _calc_volume(self):
+ return np.pi * (self._water_height + self._rock_height) * np.power(
+ self._radius, 2)
+
+ def random_pos(self, n=1):
+ r = self._radius * np.sqrt(np.random.uniform(0, 1, n))
+ phi = np.random.uniform(0, 2 * np.pi, n)
+ pos_x = r * np.cos(phi) + self._detector_center[0]
+ pos_y = r * np.sin(phi) + self._detector_center[1]
+ pos_z = np.random.uniform(-self._rock_height, self._water_height, n)
+ pos = np.concatenate([pos_x, pos_y, pos_z]).reshape((3, -1)).T
+ if pos.shape[0] == 1:
+ return pos[0, :]
+ else:
+ return pos
+
+ def in_can(self, pos):
+ """
+ Check if position is inside the CAN
+
+ Parameters
+ ----------
+ pos: np.array
+ The positions which should be checked
+
+ Return
+ ------
+ boolean / np.array
+ """
+ if type(pos) is tuple or pos.ndim == 1:
+ pos = np.reshape(pos, (-1, 3))
+ zmask = (pos[:, 2] >= self._canzmin) & (pos[:, 2] <= self._canzmax)
+ r2 = (pos[:, 0] - self._coord_origin[0])**2 + \
+ (pos[:, 1] - self._coord_origin[1])**2
+ rmask = r2 < (self._canradius**2)
+ mask = zmask & rmask
+ if len(mask) == 1:
+ return mask[0]
+ else:
+ return mask
+
+ def random_dir(self, n=1):
+ phi = np.random.uniform(0, 2 * np.pi, n)
+ cos_theta = np.random.uniform(self._cosZmin, self._cosZmax, n)
+ direction = np.concatenate([phi, cos_theta]).reshape((2, -1)).T
+ if direction.shape[0] == 1:
+ return direction[0, :]
+ else:
+ return direction
+
+ def header_entries(self, nevents=0):
+ retdct = dict()
+ key = "genvol"
+ value = "{} {} {} {} {}".format(-self._rock_height, self._water_height,
+ self._radius, self._volume, nevents)
+ retdct[key] = value
+ key = "fixedcan"
+ value = "{} {} {} {} {}".format(self._detector_center[0],
+ self._detector_center[1],
+ self._canzmin, self._canzmax,
+ self._canradius)
+ return retdct
+
+ def make_proposal_geometries(self):
+ """
+ Setup the geometries for the propagation using PROPOSAL
+ """
+ geometries = dict()
+ # General
+ center_x = self._detector_center[0] * M_TO_CM
+ center_y = self._detector_center[1] * M_TO_CM
+ # StopVolume
+ geometry_stop = pp.geometry.Sphere(pp.Cartesian3D(), 1e20)
+ geometry_stop.hierarchy = PROPOSAL_STOP_HIERARCHY_LEVEL
+ geometries["stop"] = geometry_stop
+ # CAN
+ can_zpos = (self._canzmin + self._canzmax) / 2 * M_TO_CM
+ geometry_can = pp.geometry.Cylinder(
+ pp.Cartesian3D(center_x, center_y, can_zpos),
+ (self._canzmax - self._canzmin) * M_TO_CM,
+ self._canradius * M_TO_CM, 0.)
+ geometry_can.hierarchy = PROPOSAL_CAN_HIERARCHY_LEVEL
+ geometries["can"] = geometry_can
+ # Interaction Volume
+ geometry_mantle = pp.geometry.Cylinder(
+ pp.Cartesian3D(center_x, center_y, can_zpos),
+ (self._canzmax - self._canzmin) * M_TO_CM, self._radius * M_TO_CM,
+ self._canradius * M_TO_CM)
+ geometry_mantle.hierarchy = PROPOSAL_PROPAGATION_HIERARCHY_LEVEL
+ geometries["can_mantle"] = geometry_mantle
+ if self._canzmin > 0:
+ floor_zpos = self._canzmin / 2 * M_TO_CM
+ geometry_floor = pp.geometry.Cylinder(
+ pp.Cartesian3D(center_x, center_y, floor_zpos),
+ (self._canzmax - self._canzmin) * M_TO_CM,
+ self._radius * M_TO_CM, 0.)
+ geometry_floor.hierarchy = PROPOSAL_PROPAGATION_HIERARCHY_LEVEL
+ geometries["can_floor"] = geometry_floor
+ if self._water_height > self._canzmax:
+ ceil_zpos = (self._water_height + self._canzmax) / 2 * M_TO_CM
+ geometry_ceil = pp.geometry.Cylinder(
+ pp.Cartesian3D(center_x, center_y, ceil_zpos),
+ (self._water_height - self._canzmax) * M_TO_CM,
+ self._radius * M_TO_CM, 0.)
+ geometry_ceil.hierarchy = PROPOSAL_PROPAGATION_HIERARCHY_LEVEL
+ geometries["can_ceiling"] = geometry_ceil
+ if self._rock_height > 0.:
+ sr_zpos = -self._rock_height / 2 * M_TO_CM
+ geometry_sr = pp.geometry.Cylinder(
+ pp.Cartesian3D(center_x, center_y, sr_zpos),
+ self._rock_height * M_TO_CM, self._radius * M_TO_CM, 0)
+ geometry_sr.hierarchy = PROPOSAL_PROPAGATION_HIERARCHY_LEVEL
+ geometries["sr"] = geometry_sr
+ return geometries
+
+ @staticmethod
+ def _addparticles(dct, particle_infos):
+ for prtcl in particle_infos:
+ dct['barcode'].append(prtcl.type)
+ dct['E'].append(prtcl.energy)
+ dct['x'].append(prtcl.position.x / 100)
+ dct['y'].append(prtcl.position.y / 100)
+ dct['z'].append(prtcl.position.z / 100)
+ dct['Px'].append(prtcl.direction.x * prtcl.momentum / 1e3)
+ dct['Py'].append(prtcl.direction.y * prtcl.momentum / 1e3)
+ dct['Pz'].append(prtcl.direction.z * prtcl.momentum / 1e3)
+ dct['deltaT'].append(prtcl.time)
+
+ def distribute_event(self, evt):
+ # NC -> doesn't require any propagation
+ if abs(evt.process_ID) == 3 or evt.flavor_ID == 1:
+ vtx_pos = self.random_pos()
+ vtx_dir = self.random_dir()
+ weight = 1
+ evts = None
+ return vtx_pos, vtx_dir, weight, evts
+
+ if not self._pp_geometry:
+ self._pp_geometry = self.make_proposal_geometries()
+
+ if not self._propagator:
+ self._propagator = Propagator([13, -13, 15, -15],
+ self._pp_geometry)
+
+ charged_lepton_type = np.sign(evt.process_ID) * (2 * evt.flavor_ID + 9)
+
+ samples = 0
+ lepout_dir = np.array([evt.lepOut_Px, evt.lepOut_Py, evt.lepOut_Pz])
+
+ while True:
+ samples += 1
+ vtx_pos = self.random_pos()
+ vtx_angles = self.random_dir()
+ if self.in_can(vtx_pos) and evt.flavor_ID == 2:
+ return vtx_pos, vtx_angles, samples, None
+ R = Rotation.from_euler("yz", vtx_angles)
+ particles = self._propagator.propagate_to_can(
+ charged_lepton_type, evt.lepOut_E, vtx_pos,
+ R.apply(lepout_dir))
+ if not particles is None:
+ return vtx_pos, vtx_angles, samples, particles
+
class SphericalVolume(DetectorVolume):
"""
@@ -205,20 +472,29 @@ class SphericalVolume(DetectorVolume):
def _calc_volume(self):
return 4 / 3 * np.pi * np.power(self._radius, 3)
- def random_pos(self):
- r = self._radius * np.power(np.random.random(), 1 / 3)
- phi = np.random.uniform(0, 2 * np.pi)
- cosTheta = np.random.uniform(-1, 1)
- pos_x = r * np.cos(phi) * np.sqrt(1 - np.power(cosTheta, 2))
- pos_y = r * np.sin(phi) * np.sqrt(1 - np.power(cosTheta, 2))
- pos_z = r * cosTheta
- pos = (pos_x, pos_y, pos_z)
- return tuple(np.add(self._coord_origin, pos))
-
- def random_dir(self):
- phi = np.random.uniform(0, 2 * np.pi)
- cos_theta = np.random.uniform(self._cosZmin, self._cosZmax)
- return (phi, cos_theta)
+ def random_pos(self, n=1):
+ r = self._radius * np.power(np.random.random(n), 1 / 3)
+ phi = np.random.uniform(0, 2 * np.pi, n)
+ cosTheta = np.random.uniform(-1, 1, n)
+ pos_x = r * np.cos(phi) * np.sqrt(
+ 1 - np.power(cosTheta, 2)) + self._coord_origin[0]
+ pos_y = r * np.sin(phi) * np.sqrt(
+ 1 - np.power(cosTheta, 2)) + self._coord_origin[1]
+ pos_z = r * cosTheta + self._coord_origin[2]
+ pos = np.concatenate([pos_x, pos_y, pos_z]).reshape((3, -1)).T
+ if pos.shape[0] == 1:
+ return pos[0, :]
+ else:
+ return pos
+
+ def random_dir(self, n=1):
+ phi = np.random.uniform(0, 2 * np.pi, n)
+ cos_theta = np.random.uniform(self._cosZmin, self._cosZmax, n)
+ direction = np.concatenate([phi, cos_theta]).reshape((2, -1)).T
+ if direction.shape[0] == 1:
+ return direction[0, :]
+ else:
+ return direction
def header_entries(self, nevents=0):
retdct = dict()
diff --git a/km3buu/output.py b/km3buu/output.py
index 357cc10014e5f0c8f31bbaae3adef5baea8aa30f..109098616a69e7c5e0a5cc0abe0e945de1d2f7b3 100644
--- a/km3buu/output.py
+++ b/km3buu/output.py
@@ -30,7 +30,7 @@ import km3io
from .physics import visible_energy_fraction
from .jobcard import Jobcard, read_jobcard, PDGID_LOOKUP
-from .geometry import DetectorVolume, CanVolume
+from .geometry import *
from .config import Config, read_default_media_compositions
from .__version__ import version
@@ -540,7 +540,6 @@ class GiBUUOutput:
self.flux_data["energy"][mask],
self.flux_data["flux"][mask],
p0=[1, -1])
-
self._flux_index = popt[1]
@property
@@ -552,7 +551,7 @@ def write_detector_file(gibuu_output,
ofile="gibuu.offline.root",
no_files=1,
run_number=1,
- geometry=CanVolume(),
+ geometry=CylindricalVolume(),
livetime=3.156e7,
propagate_tau=True): # pragma: no cover
"""
@@ -576,8 +575,8 @@ def write_detector_file(gibuu_output,
if not isinstance(geometry, DetectorVolume):
raise TypeError("Geometry needs to be a DetectorVolume")
- def add_particles(particle_data, pos_offset, rotation, start_mc_trk_id,
- timestamp, status):
+ def add_particles(particle_data, start_mc_trk_id, timestamp, status,
+ mother_id):
nonlocal evt
for i in range(len(particle_data.E)):
trk = ROOT.Trk()
@@ -585,19 +584,13 @@ def write_detector_file(gibuu_output,
mom = np.array([
particle_data.Px[i], particle_data.Py[i], particle_data.Pz[i]
])
- p_dir = rotation.apply(mom / np.linalg.norm(mom))
- try:
- prtcl_pos = np.array([
- particle_data.x[i], particle_data.y[i], particle_data.z[i]
- ])
- prtcl_pos = rotation.apply(prtcl_pos)
- trk.pos.set(*np.add(pos_offset, prtcl_pos))
- trk.t = timestamp + particle_data.deltaT[i] * 1e9
- except AttributeError:
- trk.pos.set(*pos_offset)
- trk.t = timestamp
+ p_dir = mom / np.linalg.norm(mom)
trk.dir.set(*p_dir)
- trk.mother_id = 0
+ prtcl_pos = np.array(
+ [particle_data.x[i], particle_data.y[i], particle_data.z[i]])
+ trk.pos.set(*prtcl_pos)
+ trk.t = timestamp + particle_data.deltaT[i]
+ trk.mother_id = mother_id
trk.type = int(particle_data.barcode[i])
trk.E = particle_data.E[i]
trk.status = status
@@ -627,11 +620,6 @@ def write_detector_file(gibuu_output,
bjorkenx = event_data.Bx
bjorkeny = event_data.By
- tau_secondaries = None
- if propagate_tau and abs(nu_type) == 16 and ichan == 2:
- from .propagation import propagate_lepton
- tau_secondaries = propagate_lepton(event_data, np.sign(nu_type) * 15)
-
media = read_default_media_compositions()
density = media["SeaWater"]["density"] # [g/cm^3]
element = mendeleev.element(gibuu_output.Z)
@@ -663,6 +651,8 @@ def write_detector_file(gibuu_output,
header_dct["primary"] = "{:d}".format(nu_type)
header_dct["start_run"] = str(run_number)
+ event_times = np.random.uniform(0, livetime, len(event_data))
+
for i in range(no_files):
start_id = 0
stop_id = len(event_data)
@@ -677,15 +667,20 @@ def write_detector_file(gibuu_output,
outfile = ROOT.TFile.Open(tmp_filename, "RECREATE")
tree = ROOT.TTree("E", "KM3NeT Evt Tree")
tree.Branch("Evt", evt, 32000, 4)
- mc_trk_id = 0
-
- for mc_event_id, event in enumerate(event_data[start_id:stop_id]):
+ from tqdm import tqdm
+ for mc_event_id, event in tqdm(enumerate(
+ event_data[start_id:stop_id])):
+ mc_trk_id = 0
+ total_id = start_id + mc_event_id
evt.clear()
evt.id = mc_event_id
evt.mc_run_id = mc_event_id
+ # Vertex Positioning & Propagation
+ vtx_pos, vtx_angles, samples, prop_particles = geometry.distribute_event(
+ event)
# Weights
evt.w.push_back(geometry.volume) # w1 (can volume)
- evt.w.push_back(w2[start_id + mc_event_id]) # w2
+ evt.w.push_back(w2[total_id] / samples) # w2
evt.w.push_back(-1.0) # w3 (= w2*flux)
# Event Information (w2list)
evt.w2list.resize(W2LIST_LENGTH)
@@ -711,10 +706,9 @@ def write_detector_file(gibuu_output,
evt.w2list[km3io.definitions.w2list_km3buu["W2LIST_KM3BUU_P_EARTH"]] = np.nan
evt.w2list[km3io.definitions.w2list_km3buu["W2LIST_KM3BUU_WATER_INT_LEN"]] = np.nan
- # Vertex Position
- vtx_pos = np.array(geometry.random_pos())
+ timestamp = event_times[total_id]
# Direction
- phi, cos_theta = geometry.random_dir()
+ phi, cos_theta = vtx_angles
sin_theta = np.sqrt(1 - cos_theta**2)
dir_x = np.cos(phi) * sin_theta
@@ -725,8 +719,6 @@ def write_detector_file(gibuu_output,
theta = np.arccos(cos_theta)
R = Rotation.from_euler("yz", [theta, phi])
- timestamp = np.random.uniform(0, livetime)
-
nu_in_trk = ROOT.Trk()
nu_in_trk.id = mc_trk_id
mc_trk_id += 1
@@ -736,7 +728,8 @@ def write_detector_file(gibuu_output,
nu_in_trk.dir.set(*direction)
nu_in_trk.E = event.lepIn_E
nu_in_trk.t = timestamp
- nu_in_trk.status = km3io.definitions.trkmembers["TRK_ST_PRIMARYNEUTRINO"]
+ nu_in_trk.status = km3io.definitions.trkmembers[
+ "TRK_ST_PRIMARYNEUTRINO"]
evt.mc_trks.push_back(nu_in_trk)
lep_out_trk = ROOT.Trk()
@@ -751,21 +744,40 @@ def write_detector_file(gibuu_output,
lep_out_trk.E = event.lepOut_E
lep_out_trk.t = timestamp
- if tau_secondaries is not None:
- lep_out_trk.status = km3io.definitions.trkmembers["TRK_ST_UNDEFINED"]
+ if prop_particles is not None:
+ lep_out_trk.status = km3io.definitions.trkmembers[
+ "TRK_ST_PROPLEPTON"]
+ generator_particle_state = km3io.definitions.trkmembers[
+ "TRK_ST_UNDEFINED"]
else:
- lep_out_trk.status = km3io.definitions.trkmembers["TRK_ST_FINALSTATE"]
+ lep_out_trk.status = km3io.definitions.trkmembers[
+ "TRK_ST_FINALSTATE"]
+ generator_particle_state = km3io.definitions.trkmembers[
+ "TRK_ST_FINALSTATE"]
evt.mc_trks.push_back(lep_out_trk)
- if tau_secondaries is not None:
- event_tau_sec = tau_secondaries[mc_event_id]
- add_particles(event_tau_sec, vtx_pos, R, mc_trk_id, timestamp,
- km3io.definitions.trkmembers["TRK_ST_FINALSTATE"])
- mc_trk_id += len(event_tau_sec.E)
+ event.x = np.ones(len(event.E)) * vtx_pos[0]
+ event.y = np.ones(len(event.E)) * vtx_pos[1]
+ event.z = np.ones(len(event.E)) * vtx_pos[2]
+
+ directions = R.apply(np.array([event.Px, event.Py, event.Pz]).T)
+ event.Px = directions[:, 0]
+ event.Py = directions[:, 1]
+ event.Pz = directions[:, 2]
+
+ event.deltaT = np.zeros(len(event.E))
+
+ add_particles(event, mc_trk_id, timestamp,
+ generator_particle_state, 0)
+ mc_trk_id += len(event.E)
+
+ if prop_particles is not None:
+ add_particles(
+ prop_particles, mc_trk_id, timestamp,
+ km3io.definitions.trkmembers["TRK_ST_FINALSTATE"], 1)
+ mc_trk_id += len(prop_particles.E)
- add_particles(event, vtx_pos, R, mc_trk_id, timestamp,
- km3io.definitions.trkmembers["TRK_ST_FINALSTATE"])
tree.Fill()
for k, v in geometry.header_entries(mc_event_id + 1).items():
diff --git a/km3buu/propagation.py b/km3buu/propagation.py
index 61b59a9cab0137a0ec79af8a2dfc3f9a0aa669a3..b1874ba1a198abb72399ad9506c74a992959c64a 100644
--- a/km3buu/propagation.py
+++ b/km3buu/propagation.py
@@ -22,6 +22,12 @@ import pathlib
from .config import Config
+M_TO_CM = 1e2
+
+itp_table_path = Config().proposal_itp_tables
+pathlib.Path(itp_table_path).mkdir(parents=True, exist_ok=True)
+pp.InterpolationSettings.tables_path = itp_table_path
+
PROPOSAL_LEPTON_DEFINITIONS = {
11: pp.particle.EMinusDef,
-11: pp.particle.EPlusDef,
@@ -37,16 +43,21 @@ PROPOSAL_LEPTON_DEFINITIONS = {
-16: pp.particle.NuTauBarDef
}
-PROPOSAL_TARGET = pp.medium.AntaresWater()
+PROPOSAL_TARGET_WATER = pp.medium.AntaresWater()
+PROPOSAL_TARGET_ROCK = pp.medium.StandardRock()
+
+PROPOSAL_STOP_HIERARCHY_LEVEL = 0
+PROPOSAL_PROPAGATION_HIERARCHY_LEVEL = 4
+PROPOSAL_CAN_HIERARCHY_LEVEL = 2
+PROPOSAL_MUON_STOP_CONDITION = 3
+PROPOSAL_TAU_STOP_CONDITION = 1
-def _setup_propagator(max_distance, particle_definition):
- itp_table_path = Config().proposal_itp_tables
- pathlib.Path(itp_table_path).mkdir(parents=True, exist_ok=True)
- pp.InterpolationSettings.tables_path = itp_table_path
+
+def _setup_utility(particle_definition, target):
crosssection = pp.crosssection.make_std_crosssection(
particle_def=particle_definition,
- target=PROPOSAL_TARGET,
+ target=target,
interpolate=True,
cuts=pp.EnergyCutSettings(np.inf, 0.05, False))
collection = pp.PropagationUtilityCollection()
@@ -54,73 +65,120 @@ def _setup_propagator(max_distance, particle_definition):
collection.interaction = pp.make_interaction(crosssection, True)
collection.decay = pp.make_decay(crosssection, particle_definition, True)
collection.time = pp.make_time(crosssection, particle_definition, True)
+ return pp.PropagationUtility(collection=collection)
- utility = pp.PropagationUtility(collection=collection)
- geometry = pp.geometry.Sphere(pp.Cartesian3D(), max_distance)
- density_distr = pp.density_distribution.density_homogeneous(
- PROPOSAL_TARGET.mass_density)
+def setup_propagator(particle_definition, proposal_geometries):
+ utility_sw = _setup_utility(particle_definition, PROPOSAL_TARGET_WATER)
+ utility_sr = _setup_utility(
+ particle_definition,
+ PROPOSAL_TARGET_ROCK) if "sr" in proposal_geometries.keys() else None
- propagator = pp.Propagator(particle_definition,
- [(geometry, utility, density_distr)])
+ density_sw = pp.density_distribution.density_homogeneous(
+ PROPOSAL_TARGET_WATER.mass_density)
+ density_sr = pp.density_distribution.density_homogeneous(
+ PROPOSAL_TARGET_ROCK.mass_density)
- return propagator
+ sectors = [(geometry, utility_sw, density_sw)
+ for k, geometry in proposal_geometries.items() if k != "sr"]
+ if "sr" in proposal_geometries.keys():
+ sectors.append((proposal_geometries["sr"], utility_sr, density_sr))
+ propagator = pp.Propagator(particle_definition, sectors)
-def propagate_lepton(lepout_data, pdgid):
- """
- Lepton propagation based on PROPOSAL
-
- Parameters
- ----------
- lepout_data: awkward.highlevel.Array
- Lepton data in the GiBUU output shape containing the fields
- 'lepOut_E, lepOut_Px, lepOut_Py, lepOut_Pz'
- pdgid:
- The pdgid of the propagated lepton
-
- Returns
- -------
- awkward.highlevel.Array (E, Px, Py, Pz, x, y, z)
- """
- lepton_info = Particle.from_pdgid(pdgid)
- prop_range = const.c * lepton_info.lifetime * 1e11 * np.max(
- lepout_data.lepOut_E) / lepton_info.mass #[cm]
-
- lepton_def = PROPOSAL_LEPTON_DEFINITIONS[pdgid]()
+ return propagator
- propagator = _setup_propagator(10 * prop_range, lepton_def)
- propagated_data = defaultdict(list)
+class Propagator(object):
+
+ def __init__(self, pdgids, geometry):
+ self._geometry = geometry
+ self._pdgids = pdgids
+ self._pp_propagators = dict()
+ self._setup()
+
+ def _setup(self):
+ for p in self._pdgids:
+ self._pp_propagators[p] = setup_propagator(
+ PROPOSAL_LEPTON_DEFINITIONS[p](), self._geometry)
+
+ # @geometry.setter
+ # def geometry(self, geodct):
+ # self._geometry = geodct
+ # # Update propagators
+ # self._setup()
+
+ @staticmethod
+ def _addparticles(dct, particle_infos):
+ for prtcl in particle_infos:
+ dct['barcode'].append(prtcl.type)
+ dct['E'].append(prtcl.energy)
+ dct['x'].append(prtcl.position.x / M_TO_CM)
+ dct['y'].append(prtcl.position.y / M_TO_CM)
+ dct['z'].append(prtcl.position.z / M_TO_CM)
+ dct['Px'].append(prtcl.direction.x * prtcl.momentum / 1e3)
+ dct['Py'].append(prtcl.direction.y * prtcl.momentum / 1e3)
+ dct['Pz'].append(prtcl.direction.z * prtcl.momentum / 1e3)
+ dct['deltaT'].append(prtcl.time * 1e-9)
+
+ def propagate_particle(self, prtcl_state):
+ stop_condition = PROPOSAL_MUON_STOP_CONDITION if abs(
+ prtcl_state.type) == 13 else PROPOSAL_TAU_STOP_CONDITION
+ return self._pp_propagators[prtcl_state.type].propagate(
+ prtcl_state, hierarchy_condition=stop_condition)
+
+ def propagate_to_can(self, lep_pdgid, lep_E, lep_pos, lep_dir):
+ """
+ Lepton propagation to can based on PROPOSAL
+
+ Parameters
+ ----------
+ vtx_pos: np.array
+ Vertex positions of the given events
+ lep_E: np.array
+ Outgoing lepton energy
+ lep_dir: np.array
+ Lepton direction positions of the given events
+ pdgid:
+ The pdgid of the propagated lepton
+
+ Returns
+ -------
+ awkward.highlevel.Array (barcode, deltaT, E, Px, Py, Pz, x, y, z)
+ """
+ if not lep_pdgid in self._pdgids:
+ return None
- for entry in lepout_data:
init_state = pp.particle.ParticleState()
- init_state.energy = entry.lepOut_E * 1e3
- init_state.position = pp.Cartesian3D(0, 0, 0)
- init_state.direction = pp.Cartesian3D(entry.lepOut_Px, entry.lepOut_Py,
- entry.lepOut_Pz)
+ init_state.type = lep_pdgid
+ init_state.energy = lep_E * 1e3
+ init_state.direction = pp.Cartesian3D(lep_dir)
init_state.direction.normalize()
- track = propagator.propagate(init_state, 5 * prop_range)
- secondaries = track.decay_products()
-
- E = [s.energy / 1e3 for s in secondaries]
- pdgid = [p.type for p in secondaries]
- Px = [p.direction.x * p.momentum / 1e3 for p in secondaries]
- Py = [p.direction.y * p.momentum / 1e3 for p in secondaries]
- Pz = [p.direction.z * p.momentum / 1e3 for p in secondaries]
- x = [p.position.x / 100 for p in secondaries]
- y = [p.position.y / 100 for p in secondaries]
- z = [p.position.z / 100 for p in secondaries]
- dt = [p.time for p in secondaries]
-
- propagated_data["E"].append(E)
- propagated_data["Px"].append(Px)
- propagated_data["Py"].append(Py)
- propagated_data["Pz"].append(Pz)
- propagated_data["barcode"].append(pdgid)
- propagated_data["x"].append(x)
- propagated_data["y"].append(y)
- propagated_data["z"].append(z)
- propagated_data["deltaT"].append(dt)
-
- return ak.Array(propagated_data)
+ init_state.position = pp.Cartesian3D(*(lep_pos * M_TO_CM))
+
+ particle_dct = defaultdict(list)
+
+ track = self.propagate_particle(init_state)
+ fsi = track.final_state()
+ decay = track.decay_products()
+
+ if self._geometry['can'].is_inside(fsi.position, fsi.direction):
+ if len(decay) > 0:
+ self._addparticles(particle_dct, decay)
+ else:
+ self._addparticles(particle_dct, [fsi])
+ else:
+ for d in decay:
+ if d.type in self._pp_propagators.keys():
+ track = self.propagate_particle(d)
+ fsi = track.final_state()
+ decay = track.decay_products()
+ if self._geometry['can'].is_inside(fsi.position,
+ fsi.direction):
+ if len(decay) > 0:
+ self._addparticles(particle_dct, decay)
+ else:
+ self._addparticles(particle_dct, [fsi])
+ if len(particle_dct["E"]) == 0:
+ return None
+ return ak.Array(particle_dct)