add function to sample scintillation

pyproject.toml

@@ -31,6 +31,7 @@ classifiers = [
 ]
 requires-python = ">=3.9"
 dependencies = [
+    "numba",
     "numpy",
     "pint >= 0.24.1",
     "pyg4ometry",

src/legendoptics/lar.py

@@ -34,12 +34,8 @@
 import pint
 from pint import Quantity
 
-from legendoptics.utils import (
-    InterpolatingGraph,
-    ScintConfig,
-    ScintParticle,
-    readdatafile,
-)
+from legendoptics.scintillate import ScintConfig, ScintParticle
+from legendoptics.utils import InterpolatingGraph, readdatafile
 
 log = logging.getLogger(__name__)
 u = pint.get_application_registry()
@@ -52,6 +48,9 @@ class ArScintLiftime(NamedTuple):
     singlet: Quantity
     triplet: Quantity
 
+    def as_tuple(self) -> tuple[Quantity, Quantity]:
+        return (self.singlet, self.triplet)
+
 
 def lar_dielectric_constant_bideau_mehu(
     λ: Quantity,
@@ -262,7 +261,9 @@ def lar_lifetimes(
     return ArScintLiftime(singlet=5.95 * u.ns, triplet=triplet)
 
 
-def lar_scintillation_params(flat_top_yield: Quantity = 31250 / u.MeV) -> ScintConfig:
+def lar_scintillation_params(
+    flat_top_yield: Quantity = 31250 / u.MeV,
+) -> ScintConfig:
     r"""Scintillation yield (approx. inverse of the mean energy to produce a UV photon).
 
     This depends on the nature of the impinging particles, the field configuration
@@ -287,9 +288,14 @@ def lar_scintillation_params(flat_top_yield: Quantity = 31250 / u.MeV) -> ScintC
     Excitation ratio:
     For example, for nuclear recoils it should be 0.75
     nominal value for electrons and gammas: 0.23 (WArP data)
+
+    See Also
+    --------
+    .lar_fano_factor
     """
     return ScintConfig(
         flat_top=flat_top_yield,
+        fano_factor=lar_fano_factor(),
         particles=[
             ScintParticle("electron", yield_factor=0.8, exc_ratio=0.23),
             ScintParticle("alpha", yield_factor=0.7, exc_ratio=1),
@@ -425,7 +431,6 @@ def pyg4_lar_attach_scintillation(
     See Also
     --------
     .lar_scintillation_params
-    .lar_fano_factor
     .lar_emission_spectrum
     .lar_lifetimes
     """
@@ -449,9 +454,11 @@ def pyg4_lar_attach_scintillation(
     lar_mat.addConstPropertyPint("SCINTILLATIONTIMECONSTANT1", lifetimes.singlet)
     lar_mat.addConstPropertyPint("SCINTILLATIONTIMECONSTANT2", lifetimes.triplet)
 
+    scint_params = lar_scintillation_params(flat_top_yield)
+
     # the fano factor is the ratio between variance and mean. Geant4 calculates
     # σ = RESOLUTIONSCALE × √mean, so we have to take the root of the fano factor
     # here to keep it consistent.
-    lar_mat.addConstPropertyPint("RESOLUTIONSCALE", np.sqrt(lar_fano_factor()))
+    lar_mat.addConstPropertyPint("RESOLUTIONSCALE", np.sqrt(scint_params.fano_factor))
 
-    pyg4_def_scint_by_particle_type(lar_mat, lar_scintillation_params(flat_top_yield))
+    pyg4_def_scint_by_particle_type(lar_mat, scint_params)

src/legendoptics/pen.py

@@ -19,12 +19,8 @@
 import pint
 from pint import Quantity
 
-from legendoptics.utils import (
-    InterpolatingGraph,
-    ScintConfig,
-    ScintParticle,
-    readdatafile,
-)
+from legendoptics.scintillate import ScintConfig, ScintParticle
+from legendoptics.utils import InterpolatingGraph, readdatafile
 
 log = logging.getLogger(__name__)
 u = pint.get_application_registry()
@@ -99,6 +95,23 @@ def pen_wls_absorption() -> tuple[Quantity, Quantity]:
     return wvl, absorp
 
 
+def pen_scintillation_params() -> ScintConfig:
+    """Get a :class:`ScintConfig` object for PEN.
+
+    See Also
+    --------
+    .pen_scint_light_yield
+    """
+    # setup scintillation response just for electrons.
+    return ScintConfig(
+        flat_top=pen_scint_light_yield(),
+        fano_factor=None,
+        particles=[
+            ScintParticle("electron", yield_factor=1, exc_ratio=None),
+        ],
+    )
+
+
 def pyg4_pen_attach_rindex(mat, reg) -> None:
     """Attach the refractive index to the given PEN material instance.
 
@@ -193,11 +206,4 @@ def pyg4_pen_attach_scintillation(mat, reg) -> None:
     # so we set it to 1; otherwise Geant4 will crash.
     mat.addConstPropertyPint("RESOLUTIONSCALE", 1)
 
-    # setup scintillation response just for electrons.
-    scint_config = ScintConfig(
-        flat_top=pen_scint_light_yield(),
-        particles=[
-            ScintParticle("electron", yield_factor=1, exc_ratio=None),
-        ],
-    )
-    pyg4_def_scint_by_particle_type(mat, scint_config)
+    pyg4_def_scint_by_particle_type(mat, pen_scintillation_params())

src/legendoptics/pyg4utils.py

@@ -7,7 +7,7 @@
 import pyg4ometry.geant4 as g4
 from pint import Quantity
 
-from .utils import ScintConfig
+from .scintillate import ScintConfig
 
 log = logging.getLogger(__name__)
 ureg = pint.get_application_registry().get()
@@ -55,6 +55,12 @@ def _def_scint_particle(
 def pyg4_def_scint_by_particle_type(mat, scint_cfg: ScintConfig) -> None:
     """Define a full set of particles for scintillation."""
     for particle in scint_cfg.particles:
+        if not particle.valid_geant_particle():
+            msg = (
+                f"Unknown particle type {particle.name} for ScintillationByParticleType"
+            )
+            raise ValueError(msg)
+
         _def_scint_particle(
             mat,
             particle.name.upper(),

src/legendoptics/scintillate.py

@@ -0,0 +1,206 @@
+from __future__ import annotations
+
+from typing import Literal, NamedTuple, NewType, Optional, get_args, get_type_hints
+
+import numpy as np
+import pint
+from numba import njit
+from pint import Quantity
+
+u = pint.get_application_registry()
+
+
+class ScintParticle(NamedTuple):
+    """Configuration for the scintillation yield relative to the flat-top yield."""
+
+    name: Literal["deuteron", "triton", "alpha", "ion", "electron"]
+    yield_factor: float
+    exc_ratio: Optional[float]  # noqa: UP007
+
+    def valid_geant_particle(self) -> bool:
+        return self.name.lower() in get_args(get_type_hints(ScintParticle)["name"])
+
+
+class ScintConfig(NamedTuple):
+    """Scintillation yield parameters, depending on the particle types."""
+
+    flat_top: Quantity
+    fano_factor: Optional[float]  # noqa: UP007
+    particles: list[ScintParticle]
+
+    def get_particle(self, name: str) -> Optional[ScintParticle]:  # noqa: UP007
+        for p in self.particles:
+            if p.name.upper() == name.upper():
+                return p
+        return None
+
+
+ParticleIndex = NewType("ParticleIndex", int)
+ComputedScintParams = NewType(
+    "ComputedScintParams", tuple[float, float, np.ndarray, np.ndarray]
+)
+PARTICLE_INDICES = {"electron": 0, "alpha": 1, "ion": 2, "deuteron": 3, "triton": 4}
+
+
+def precompute_scintillation_params(
+    scint_config: ScintConfig,
+    time_components: tuple[Quantity, ...],
+) -> ComputedScintParams:
+    # validate that we have a valid configuration of scintillation parameters.
+    part_with_exc_ratio = [1 for p in scint_config.particles if p.exc_ratio is not None]
+    if len(part_with_exc_ratio) not in (0, len(scint_config.particles)):
+        msg = "Not all particles have exc_ratio defined"
+        raise ValueError(msg)
+    has_2_timeconstants = len(part_with_exc_ratio) > 0
+    if len(time_components) != (1 + int(has_2_timeconstants)):
+        msg = "Number of time_components is not as required"
+        raise ValueError(msg)
+
+    particles = np.zeros((len(PARTICLE_INDICES), (2 + int(has_2_timeconstants))))
+
+    electron = scint_config.get_particle("electron")
+    if electron is None:
+        raise ValueError()
+    for k, v in PARTICLE_INDICES.items():
+        p = scint_config.get_particle(k)
+        p = electron if p is None else p
+        if has_2_timeconstants:
+            particles[v] = np.array([p.yield_factor, p.exc_ratio, 1 - p.exc_ratio])
+        else:
+            particles[v] = np.array([p.yield_factor, 1])
+
+    time_components = np.array([t.to(u.nanosecond).m for t in time_components])
+    fano = scint_config.fano_factor if scint_config.fano_factor is not None else 1
+
+    return scint_config.flat_top.to("1/keV").m, fano, time_components, particles
+
+
+def particle_to_index(particle: str) -> ParticleIndex:
+    """Converts the given G4 scintillation particle name to a module-internal index."""
+    return PARTICLE_INDICES[particle.lower()]
+
+
+@njit
+def scintillate_local(
+    params: ComputedScintParams,
+    particle: ParticleIndex,
+    edep_keV: float,
+    rng: np.random.Generator,
+) -> np.ndarray:
+    """Generates a Poisson/Gauss-distributed number of photons according to the
+    scintillation yield formula, as implemented in Geant4.
+
+    This function only calculates the local part of scintillation.
+
+    Parameters
+    ----------
+    params
+        scintillation parameter tuple, as created by :meth:`precompute_scintillation_params`.
+    particle
+        module-internal particle index, see :meth:`particle_to_index`.
+    edep_keV
+        energy deposition along this step, in units pf keV.
+
+    Returns
+    -------
+    array of scintillation time stamps in nanoseconds, relative to the point in
+    time of energy deposition.
+    """
+    flat_top, fano, time_components, particles = params
+
+    # get the particle scintillation info.
+    if particle < 0 or particle > particles.shape[0]:
+        msg = "unknown particle index"
+        raise IndexError(msg)
+    part = particles[particle]
+    yield_factor = part[0]
+    yields = part[1:]
+
+    mean_num_phot = flat_top * yield_factor * edep_keV
+
+    # derive the actual number of photons generated in this step.
+    if mean_num_phot > 10:
+        sigma = np.sqrt(fano * mean_num_phot)
+        num_photons = int(rng.normal(mean_num_phot, sigma) + 0.5)
+    else:
+        num_photons = rng.poisson(mean_num_phot)
+
+    if num_photons <= 0:
+        return np.empty((0,))
+
+    # derive number of photons for all time components.
+    yields = (num_photons * yields).astype(np.int64)
+    yields[-1] = num_photons - np.sum(yields[0:-1])  # to keep the sum constant.
+
+    # now, calculate the timestamps of each generated photon.
+    times = np.log(rng.uniform(size=num_photons))
+    start = 0
+    for num_phot, scint_t in zip(yields, time_components):
+        times[start : start + num_phot] *= -scint_t
+        start += num_phot
+
+    return times
+
+
+@njit
+def scintillate(
+    params: ComputedScintParams,
+    x0_m: np.ndarray,
+    x1_m: np.ndarray,
+    v0_mpns: float,
+    v1_mpns: float,
+    t0_ns: float,
+    particle: ParticleIndex,
+    particle_charge: int,
+    edep_keV: float,
+    rng: np.random.Generator,
+):
+    """Generates a Poisson/Gauss-distributed number of photons according to the
+    scintillation yield formula, as implemented in Geant4, along the line segment
+    between x0 and x1.
+
+    Parameters
+    ----------
+    params
+        scintillation parameter tuple, as created by :meth:`precompute_scintillation_params`.
+    x0_m
+        three-vector of pre step point position (in units of meter).
+    x1_m
+        three-vector of post step point position (in units of meter).
+    v0_mpns
+        velocity of particle before step (in units of meter per nanosecond).
+    v1_mpns
+        velocity of particle after step (in units of meter per nanosecond).
+    t0_ns
+        global time offset of step start in scintillator, in nanoseconds.
+    particle
+        module-internal particle index, see :meth:`particle_to_index`.
+    particle_charge
+        charge of the particle, in units of the elementary charge.
+    edep_keV
+        energy deposition along this step, in units pf keV.
+
+    Returns
+    -------
+    array of four-vectors containing time stamps (in nanoseconds) and global scintillation
+    positions (in meter).
+    """
+    delta_t_scint = scintillate_local(params, particle, edep_keV, rng)
+    # emission position for each single photon.
+    if particle_charge != 0:
+        λ = rng.uniform(size=delta_t_scint.shape[0])
+    else:
+        λ = np.ones(shape=delta_t_scint.shape[0])
+
+    x = np.empty((delta_t_scint.shape[0], 4))
+    # spatial components along the line segment between x0 and x1.
+    x[:, 1] = x0_m[0] + λ * (x1_m[0] - x0_m[0])
+    x[:, 2] = x0_m[1] + λ * (x1_m[1] - x0_m[1])
+    x[:, 3] = x0_m[2] + λ * (x1_m[2] - x0_m[2])
+
+    # time component along the velocity decrease between x0 and x1.
+    x[:, 0] = np.linalg.norm(x1_m - x0_m) / (v0_mpns + λ * (v1_mpns - v0_mpns) / 2)
+    # add the global time offset and the emission time offset.
+    x[:, 0] += t0_ns + delta_t_scint
+
+    return x

src/legendoptics/utils.py

@@ -1,7 +1,5 @@
 from __future__ import annotations
 
-from typing import NamedTuple
-
 import numpy as np
 import pint
 import scipy.interpolate
@@ -104,18 +102,3 @@ def __call__(self, pts: Quantity) -> Quantity:
         if pts > self.d_max:
             return self.vals.iloc[-1]
         return self.fn(pts.to(self.idx.u).m) * self.vals.u
-
-
-class ScintParticle(NamedTuple):
-    """Configuration for the scintillation yield relative to the flat-top yield."""
-
-    name: str
-    yield_factor: float
-    exc_ratio: float | None
-
-
-class ScintConfig(NamedTuple):
-    """Scintillation yield parameters, depending on the particle types."""
-
-    flat_top: Quantity
-    particles: list[ScintParticle]