Generate cross section information in (py)hepmc files (#198)

Add inelastic cross section to hepmc output.

Since chromo generates only successful events, the generated cross
section is equal to the estimated one.

setup.cfg

@@ -1,6 +1,6 @@
 [metadata]
 name = chromo
-version = 0.5.0
+version = 0.5.1
 author = Anatoli Fedynitch
 maintainer_email = [email protected]
 description = Hadronic Interaction Model interface in Python

src/chromo/common.py

@@ -10,8 +10,8 @@
 
 import copy
 import dataclasses
-import warnings
 import importlib
+import warnings
 from abc import ABC, abstractmethod
 from contextlib import contextmanager
 from typing import Optional, Tuple
@@ -158,6 +158,9 @@ class EventData:
         Impact parameter for nuclear collisions in mm.
     n_wounded: (int, int)
         Number of wounded nucleons on sides A and B.
+    production_cross_section: float
+        Production cross section in mb; inelastic for photon-/hadron-hadron
+        and production cross section for hadron-/nucleus-nucleus collisions.
     pid: 1D array of int
         PDG IDs of the particles.
     status: 1D array of int
@@ -196,6 +199,7 @@ class EventData:
     nevent: int
     impact_parameter: float
     n_wounded: Tuple[int, int]
+    production_cross_section: float
     pid: np.ndarray
     status: np.ndarray
     charge: np.ndarray
@@ -250,6 +254,7 @@ def __getstate__(self):
             self.nevent,
             self.impact_parameter,
             self.n_wounded,
+            self.production_cross_section,
             self.pid.copy(),
             self.status.copy(),
             self.charge.copy(),
@@ -329,14 +334,14 @@ def _select(self, arg, update_mothers):
         # This selection is faster than __getitem__, because we skip
         # parent selection, which is just wasting time if we select only
         # final state particles.
-        pid = self.pid[arg]
         return EventData(
             self.generator,
             self.kin,
             self.nevent,
             self.impact_parameter,
             self.n_wounded,
-            pid,
+            self.production_cross_section,
+            self.pid[arg],
             self.status[arg],
             self.charge[arg],
             self.px[arg],
@@ -472,6 +477,11 @@ def to_hepmc3(self, genevent=None):
             vt=ev.vt,
             fortran=False,
         )
+        # Convert cross section from mb to pb
+        genevent.cross_section = pyhepmc.GenCrossSection()
+        genevent.cross_section.set_cross_section(
+            self.production_cross_section * 1e9, 0, -1, -1
+        )
 
         return genevent
 
@@ -537,6 +547,7 @@ def __init__(self, generator):
             int(getattr(evt, self._nevhep)),
             self._get_impact_parameter(),
             self._get_n_wounded(),
+            generator._inel_or_prod_cross_section,
             getattr(evt, self._idhep)[sel],
             getattr(evt, self._isthep)[sel],
             self._charge_init(npart),
@@ -603,6 +614,8 @@ class MCRun(ABC):
     _projectiles = standard_projectiles
     _targets = Nuclei()
     _ecm_min = 10 * GeV  # default for many models
+    # Corresponds to current cross section in mb, updated when kinematics is set
+    _inel_or_prod_cross_section = None
     _restore_beam_and_history = True
     nevents = 0  # number of generated events so far
     _unstable_pids = set(all_unstable_pids)
@@ -737,12 +750,9 @@ def random_state(self):
     def random_state(self, rng_state):
         self._rng.__setstate__(rng_state)
 
-    @property
-    def kinematics(self):
-        return self._kinematics
+    def _check_kinematics(self, kin):
+        """Check if kinematics are allowed for this generator."""
 
-    @kinematics.setter
-    def kinematics(self, kin):
         if abs(kin.p1) not in self._projectiles:
             raise ValueError(
                 f"projectile {pdg2name(kin.p1)}[{int(kin.p1)}] is not allowed, "
@@ -758,17 +768,34 @@ def kinematics(self, kin):
                 f"center-of-mass energy {kin.ecm/GeV} GeV < "
                 f"minimum energy {self._ecm_min/GeV} GeV"
             )
+
+    @property
+    def kinematics(self):
+        return self._kinematics
+
+    @kinematics.setter
+    def kinematics(self, kin):
+        self._check_kinematics(kin)
         self._kinematics = kin
         self._set_kinematics(kin)
 
-    def cross_section(self, kin=None):
+        if (kin.p1.is_nucleus and kin.p1.A > 1) or (kin.p2.is_nucleus and kin.p2.A > 1):
+            self._inel_or_prod_cross_section = self.cross_section().prod
+        else:
+            self._inel_or_prod_cross_section = self.cross_section().inelastic
+
+    def cross_section(self, kin=None, max_info=False):
         """Cross sections according to current setup.
 
         Parameters
         ----------
         kin : EventKinematics, optional
             If provided, calculate cross-section for EventKinematics.
             Otherwise return values for current setup.
+        max_info : bool, optional
+            Return full maximal information about interaction cross sections for
+            nucleus-nucleus case. Slow - uses Monte Carlo integration. Number of
+            trials is controled separately with `generator.n_trials` attribute.
         """
         with self._temporary_kinematics(kin):
             kin2 = self.kinematics
@@ -778,10 +805,12 @@ def cross_section(self, kin=None):
                 for component, fraction in zip(kin2.p2.components, kin2.p2.fractions):
                     kin3.p2 = component
                     # this calls cross_section recursively, which is fine
-                    cross_section._mul_radd(fraction, self.cross_section(kin3))
+                    cross_section._mul_radd(
+                        fraction, self._cross_section(kin3, max_info=max_info)
+                    )
                 return cross_section
             else:
-                return self._cross_section(kin)
+                return self._cross_section(kin, max_info=max_info)
 
     @abstractmethod
     def _cross_section(self, kin):

src/chromo/models/dpmjetIII.py

@@ -142,12 +142,13 @@ def __init__(self, evt_kin, *, seed=None):
 
         self._set_final_state_particles()
 
-    def _cross_section(self, kin=None, photon_x=0):
+    def _cross_section(self, kin=None, photon_x=0, max_info=False):
         kin = self.kinematics if kin is None else kin
         # we override to set precision
-        if (kin.p1.is_nucleus and kin.p1.A > 1) or kin.p2.A > 1:
+        if (
+            (kin.p1.is_nucleus and kin.p1.A > 1) or (kin.p2.is_nucleus and kin.p2.A > 1)
+        ) and max_info:
             assert kin.p2.A >= 1, "DPMJET requires nucleons or nuclei on side 2."
-            print("case 1")
             # Enable total and elastic cross section calculation
             self._lib.dtglgp.lprod = False
             self._lib.dt_xsglau(
@@ -184,6 +185,14 @@ def _generate():
                 + glxs.xsqe2[0, 0, 0]
                 + glxs.xsela[0, 0, 0],
             )
+        elif (kin.p1.is_nucleus and kin.p1.A > 1) or (
+            kin.p2.is_nucleus and kin.p2.A > 1
+        ):
+            glxs = self._lib.dtglxs
+
+            return CrossSectionData(
+                prod=glxs.xspro[0, 0, 0],
+            )
         elif kin.p1 == 22 and kin.p2.A == 1:
             stot, sine, _ = self._lib.dt_siggp(photon_x, kin.virt_p1, kin.ecm, 0)
             return CrossSectionData(total=stot, inelastic=sine, elastic=stot - sine)

src/chromo/models/epos.py

@@ -203,7 +203,7 @@ def __init__(self, evt_kin, *, seed=None):
         self._lib.charge_vect = np.vectorize(self._lib.getcharge, otypes=[np.float32])
         self.kinematics = evt_kin
 
-    def _cross_section(self, kin=None):
+    def _cross_section(self, kin=None, max_info=False):
         total, inel, el, dd, sd, _ = self._lib.xsection()
         return CrossSectionData(
             total=total,

src/chromo/models/phojet.py

@@ -1,6 +1,5 @@
 from chromo.common import MCRun, MCEvent, CrossSectionData
 from chromo.util import fortran_chars, _cached_data_dir
-from chromo.kinematics import EventFrame
 from chromo.constants import standard_projectiles
 from particle import literals as lp
 import warnings
@@ -97,6 +96,7 @@ class PHOJETRun(MCRun):
     _projectiles = standard_projectiles
     _targets = {lp.proton.pdgid, lp.neutron.pdgid}
     _param_file_name = "dpmjpar.dat"
+    _pho_event_init_ran = False
     _data_url = (
         "https://github.com/impy-project/chromo"
         + "/releases/download/zipped_data_v1.0/dpm3191_v001.zip"
@@ -170,15 +170,10 @@ def __init__(self, evt_kin, *, seed=None):
 
         self._set_final_state_particles()
 
-        self.kinematics = evt_kin
-
         # Initialize kinematics and tables (only once needed)
-        if self._lib.pho_event(-1, self.p1, self.p2)[1]:
-            raise RuntimeError(
-                "initialization failed with the current event kinematics"
-            )
+        self.kinematics = evt_kin
 
-    def _cross_section(self, kin=None):
+    def _cross_section(self, kin=None, max_info=False):
         kin = self.kinematics if kin is None else kin
         self._lib.pho_setpar(1, kin.p1, 0, 0.0)
         self._lib.pho_setpar(2, kin.p2, 0, 0.0)
@@ -208,15 +203,16 @@ def _set_stable(self, pdgid, stable):
         self._lib.pydat3.mdcy[kc - 1, 0] = not stable
 
     def _set_kinematics(self, k):
-        if k.frame == EventFrame.FIXED_TARGET:
-            self._lib.dtflg1.iframe = 1
-            self._frame = EventFrame.FIXED_TARGET
-        else:
-            self._lib.dtflg1.iframe = 2
-            self._frame = EventFrame.CENTER_OF_MASS
+        self._frame = k.frame
         self._lib.pho_setpar(1, k.p1, 0, k.virt_p1)
         self._lib.pho_setpar(2, k.p2, 0, k.virt_p2)
         self.p1, self.p2 = k.beams
+        if not self._pho_event_init_ran:
+            if self._lib.pho_event(-1, self.p1, self.p2)[1]:
+                raise RuntimeError(
+                    "initialization failed with the current event kinematics"
+                )
+            self._pho_event_init_ran = True
 
     def _generate(self):
         return not self._lib.pho_event(1, self.p1, self.p2)[1]

src/chromo/models/pythia6.py

@@ -62,7 +62,7 @@ def __init__(self, evt_kin, *, seed=None, new_mpi=False):
 
         self._set_final_state_particles()
 
-    def _cross_section(self, kin=None):
+    def _cross_section(self, kin=None, max_info=False):
         s = self._lib.pyint7.sigt[0, 0]
         c = CrossSectionData(
             total=s[0],

src/chromo/models/pythia8.py

@@ -21,6 +21,7 @@ def __init__(self, generator):
             generator.nevents,
             self._get_impact_parameter(pythia),
             self._get_n_wounded(pythia),
+            generator._inel_or_prod_cross_section,
             event.pid(),
             event.status(),
             pythia.charge(),
@@ -149,7 +150,7 @@ def __init__(self, evt_kin, *, seed=None, config=None, banner=True):
             self.kinematics = evt_kin
         self._set_final_state_particles()
 
-    def _cross_section(self, kin=None):
+    def _cross_section(self, kin=None, max_info=False):
         st = self._pythia.info.sigmaTot
         return CrossSectionData(
             total=st.sigmaTot,

src/chromo/models/qgsjet.py

@@ -104,17 +104,17 @@ def _tabulated_cross_section(self, kin=None):
             np.log(qgsgrid), np.log(cross_section), ext="extrapolate", s=0, k=1
         )
 
-        inel = np.exp(spl(np.log(kin.elab)))
-        return CrossSectionData(inelastic=inel)
+        prod = np.exp(spl(np.log(kin.elab)))
+        return CrossSectionData(prod=prod)
 
-    def _cross_section(self, kin=None, full_AA=False):
+    def _cross_section(self, kin=None, max_info=False):
         """
         Calculate cross section for QGSJET-01c event generator.
 
         Parameters
         ----------
         kin : EventKinematics, optional
-        full_AA : bool, optional
+        max_info : bool, optional
             Flag indicating whether to calculate the cross section for
             AA collisions using MC. Returns the cross sections for
             specific processes. Default is False and only inelastic or
@@ -129,7 +129,7 @@ def _cross_section(self, kin=None, full_AA=False):
         -----
         This method calculates the cross section for the QGSJET-01c
         event generator. The cross section is calculated based on the
-        provided kinematics and collision type. If `full_AA is True`
+        provided kinematics and collision type. If `max_info is True`
         and the projectile particle is a nucleus (A > 1), the cross
         section is calculated using the Glauber model. Otherwise, the
         cross section is calculated using the tabulated cross section.
@@ -139,7 +139,7 @@ def _cross_section(self, kin=None, full_AA=False):
         """
         kin = self.kinematics if kin is None else kin
 
-        if full_AA and (kin.p1.is_nucleus and kin.p1.A > 1):
+        if max_info and (kin.p1.is_nucleus and kin.p1.A > 1):
             gtot, gprod, _, gdd, gqel, gcoh = self._lib.crossc(self.glauber_trials)
             return CrossSectionData(
                 total=gtot,
@@ -216,28 +216,22 @@ def _tabulated_cross_section(self, kin=None):
         }.get(
             abs(kin.p1), 2
         )  # 2 is correct for nuclei
-        inel = self._lib.qgsect(
+        prod = self._lib.qgsect(
             kin.elab,
             projectile_id,
             kin.p1.A or 1,
             kin.p2.A,
         )
-        inel = self._lib.qgsect(
-            kin.elab,
-            projectile_id,
-            kin.p1.A or 1,
-            kin.p2.A,
-        )
-        return CrossSectionData(inelastic=inel)
+        return CrossSectionData(prod=prod)
 
-    def _cross_section(self, kin=None, full_AA=False):
+    def _cross_section(self, kin=None, max_info=False):
         """
         Calculate cross section for QGSJET-II-03/4 event generators.
 
         Parameters
         ----------
         kin : EventKinematics, optional
-        full_AA : bool, optional
+        max_info : bool, optional
             Flag indicating whether to calculate the cross section for
             AA collisions using MC. Returns the cross sections for
             specific processes. Default is False and only inelastic or
@@ -252,7 +246,7 @@ def _cross_section(self, kin=None, full_AA=False):
         -----
         This method calculates the cross section for the QGSJET-II-0X
         event generator. The cross section is calculated based on the
-        provided kinematics and collision type. If `full_AA is True`
+        provided kinematics and collision type. If `max_info is True`
         and the projectile particle is a nucleus (A > 1), the cross
         section is calculated using the Glauber model. Otherwise, the
         cross section is calculated using the tabulated cross section.
@@ -262,7 +256,7 @@ def _cross_section(self, kin=None, full_AA=False):
         """
         kin = self.kinematics if kin is None else kin
 
-        if full_AA and (kin.p1.is_nucleus and kin.p1.A > 1):
+        if max_info and (kin.p1.is_nucleus and kin.p1.A > 1):
             gtot, gprod, _, gdd, gqel, gcoh = self._lib.qgcrossc(self.glauber_trials)
             return CrossSectionData(
                 total=gtot,

src/chromo/models/sibyll.py

@@ -201,7 +201,7 @@ def __init__(self, evt_kin, *, seed=None):
 
         super()._set_final_state_particles()
 
-    def _cross_section(self, kin=None):
+    def _cross_section(self, kin=None, max_info=False):
         kin = self.kinematics if kin is None else kin
         if kin.p1.A is not None and kin.p1.A > 1:
             warnings.warn(
@@ -234,7 +234,11 @@ def _cross_section(self, kin=None):
             nsig = self._lib.nucsig
             return CrossSectionData(
                 total=float(nsig.sigt),
-                prod=float(nsig.sigt - nsig.sigqe),
+                prod=(
+                    float(nsig.sigt - nsig.sigqe)
+                    if not np.isnan(nsig.sigqe)
+                    else float(nsig.siginel)
+                ),
                 quasielastic=float(nsig.sigqe),
                 inelastic=float(nsig.siginel),
                 diffractive_sum=float(nsig.sigqsd),