simplify decay handler interface

src/chromo/common.py

@@ -752,9 +752,7 @@ def _update_final_state_particles(self, pdgid, stable):
             self._final_state_particles = self._final_state_particles[is_stable]
 
         if self._apply_decay_handler:
-            self._decay_handler.set_stable_decaying(
-                stable_pids=self._final_state_particles
-            )
+            self._decay_handler.set_stable(self._final_state_particles)
 
     def set_stable(self, pdgid, stable=True):
         """Prevent decay of unstable particles
@@ -838,7 +836,7 @@ def _set_decay_handler(self, seed=None):
             return
 
         try:
-            self._decay_handler = Pythia8DecayHandler(seed=seed)
+            self._decay_handler = Pythia8DecayHandler([], seed=seed)
         except ModuleNotFoundError as ex:
             import warnings
 

src/chromo/decay_handler.py

@@ -5,79 +5,33 @@
 class Pythia8DecayHandler:
     def __init__(
         self,
+        stable_pids,
         seed=None,
-        extra_stable_pids=None,
-        extra_decaying_pids=None,
-        stable_pids=None,
-        decaying_pids=None,
     ):
         """
         Initialize the Pythia8DecayHandler with Pythia8 and set particle types for decay.
 
         Parameters:
             seed (int): Random seed for Pythia8 initialization.
-            extra_stable_pids (list[int]): List of additional PDG IDs
-            for stable particles.
-            extra_decaying_pids (list[int]): List of additional PDG IDs
-            for decaying particles.
             stable_pids (list[int]): List of PDG IDs
-            for exclusively stable particles (all other are decaying if they can decay).
-            decaying_pids (list[int]): List of PDG IDs
-            for exclusively decaying particles (all other are stable).
-
-        Notes:
-            - If both `extra_stable_pids` and `extra_decaying_pids`
-            are provided, check for duplicates.
-            - If `stable_pids` or `decaying_pids`
-            are provided, ensure exclusivity.
+            for stable particles (all other are decaying).
+
         """
         # Initialize Pythia8 with specific configurations
-        config = ["ProcessLevel:all = off", "ParticleDecays:tau0Max = 1e100"]
+        config = ["ProcessLevel:all = off"]
         pythia8 = chromo.models.Pythia8(None, seed=seed, config=config, banner=False)
         self.pythia = pythia8._pythia
 
-        self.set_stable_decaying(
-            extra_stable_pids,
-            extra_decaying_pids,
-            stable_pids,
-            decaying_pids,
-        )
+        self.set_stable(stable_pids)
 
-    def set_stable_decaying(
-        self,
-        extra_stable_pids=None,
-        extra_decaying_pids=None,
-        stable_pids=None,
-        decaying_pids=None,
-    ):
-        self.extra_stable_pids = extra_stable_pids
-        self.extra_decaying_pids = extra_decaying_pids
-        self.stable_pids = stable_pids
-        self.decaying_pids = decaying_pids
-        # Check for duplicates in extra_stable_pids and extra_decaying_pids
-        if (extra_stable_pids is not None) and (extra_decaying_pids is not None):
-            dup_pids = set(extra_stable_pids) & set(extra_decaying_pids)
-            if dup_pids:
-                raise ValueError(
-                    f"Duplicates found in {dup_pids} between "
-                    "extra_stable_pids and extra_decaying_pids"
-                )
-
-        # Set particles as stable or decaying based on the provided lists
-        if extra_stable_pids is not None:
-            for pid in extra_stable_pids:
-                self.pythia.particleData.mayDecay(pid, False)
-
-        if extra_decaying_pids is not None:
-            for pid in extra_decaying_pids:
-                self.pythia.particleData.mayDecay(pid, True)
-
-        # Handle exclusive stability options
-        if stable_pids is not None:
-            self._set_exclusive_stability(stable_pids, False)
-
-        if decaying_pids is not None:
-            self._set_exclusive_stability(decaying_pids, True)
+    def set_stable(self, stable_pids):
+        # Set all particles as decaying
+        for pentry in self.pythia.particleData.all():
+            self.pythia.particleData.mayDecay(pentry.id, True)
+
+        # Set stable_pids as decaying
+        for pid in stable_pids:
+            self.pythia.particleData.mayDecay(pid, False)
 
         # Store stable and decaying PDG IDs for future reference
         all_stable_pids = []
@@ -95,41 +49,6 @@ def set_stable_decaying(
         self.all_decaying_pids = np.array(all_decaying_pids, dtype=np.int64)
         self.all_stable_pids = np.array(all_stable_pids, dtype=np.int64)
 
-    def _set_exclusive_stability(self, pids, is_decaying):
-        """
-        Set particles exclusively as stable or decaying based on the provided list.
-
-        Parameters:
-            pids (list[int]): List of PDG IDs for particles to be exclusively set.
-            is_decaying (bool): True if the particles should be set as decaying,
-            False for stable.
-
-        Raises:
-            ValueError: If exclusivity conditions are not met.
-        """
-        if (
-            sum(
-                x is not None
-                for x in [
-                    self.extra_stable_pids,
-                    self.extra_decaying_pids,
-                    self.stable_pids,
-                    self.decaying_pids,
-                ]
-            )
-            != 1
-        ):
-            raise ValueError(
-                "stable_pids or decaying_pids can be set "
-                "if all other options are None"
-            )
-
-        # Set particles exclusively as stable or decaying
-        for pentry in self.pythia.particleData.all():
-            self.pythia.particleData.mayDecay(pentry.id, not is_decaying)
-        for pid in pids:
-            self.pythia.particleData.mayDecay(pid, is_decaying)
-
     def __call__(self, event):
         """
         Decay particles in the provided `event` that are marked as `decaying`.

tests/test_decay_handler.py

@@ -3,6 +3,7 @@
 import pytest
 import sys
 
+from chromo.constants import long_lived_for
 from chromo.decay_handler import Pythia8DecayHandler
 from chromo.util import get_all_models
 from .util import run_in_separate_process
@@ -11,13 +12,9 @@
     pytest.skip("Pythia8 does not run on windows", allow_module_level=True)
 
 
-def run_decay_handler(model, evt_kin, stable_particles, decaying_particles):
+def run_decay_handler(model, evt_kin, stable_particles):
     seed = 1
-    decay_handler = Pythia8DecayHandler(
-        seed=seed,
-        extra_stable_pids=stable_particles,
-        extra_decaying_pids=decaying_particles,
-    )
+    decay_handler = Pythia8DecayHandler(stable_particles, seed)
     gen = model(evt_kin, seed=seed)
 
     for event in gen(10):
@@ -65,14 +62,13 @@ def run_decay_handler(model, evt_kin, stable_particles, decaying_particles):
 
 @pytest.mark.parametrize("Model", get_all_models())
 def test_decay_handler(Model):
-    stable_particles = [111]
-    decaying_particles = [-211, 211, 13, -13, 2112]
+    stable_particles = long_lived_for(30e-12)
 
     if Model.name in ["UrQMD", "QGSJet"]:
         # UrQMD produce neutrons = 2112 with wrong masses
         # QGSJet produce neutrons with proton mass
         # Pythia8 ignores them for decay
-        decaying_particles = [-211, 211, 13, -13]
+        stable_particles = stable_particles
 
     if Model.name in ["PhoJet", "Pythia"]:
         evt_kin = chromo.kinematics.FixedTarget(100, "p", "p")
@@ -82,6 +78,4 @@ def test_decay_handler(Model):
     else:
         evt_kin = chromo.kinematics.FixedTarget(100, "p", "O")
 
-    run_in_separate_process(
-        run_decay_handler, Model, evt_kin, stable_particles, decaying_particles
-    )
+    run_in_separate_process(run_decay_handler, Model, evt_kin, stable_particles)