Fix cross_section function (#130)

These are small fixes to solve the issues #128 and #129

src/chromo/common.py

@@ -95,6 +95,14 @@ def __eq__(self, other):
     def __ne__(self, other):
         return not self == other
 
+    def _mul_radd(self, factor, other):
+        for field in dataclasses.fields(self):
+            setattr(
+                self,
+                field.name,
+                getattr(self, field.name) + factor * getattr(other, field.name),
+            )
+
 
 # Do we need EventData.n_spectators in addition to EventData.n_wounded?
 # n_spectators can be computed from n_wounded as number_of_nucleons - n_wounded
@@ -732,9 +740,7 @@ 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
-                    cs = self.cross_section(kin3)
-                    for i, val in enumerate(dataclasses.astuple(cs)):
-                        cross_section[i] += fraction * val
+                    cross_section._mul_radd(fraction, self.cross_section(kin3))
                 return cross_section
             else:
                 return self._cross_section(kin)

src/chromo/models/dpmjetIII.py

@@ -113,15 +113,11 @@ def __init__(self, evt_kin, *, seed=None):
         self._lib.pydat1.mstj[22 - 1] = 2
         self._set_final_state_particles()
 
-    def _cross_section(self, kin=None, precision=None):
+    def _cross_section(self, kin=None):
         kin = self.kinematics if kin is None else kin
         # we override to set precision
         if (kin.p1.A and kin.p1.A > 1) or kin.p2.A > 1:
             assert kin.p2.A >= 1, "DPMJET requires nucleons or nuclei on side 2."
-            if precision is not None:
-                saved = self._lib.dtglgp.jstatb
-                # Set number of trials for Glauber model integration
-                self._lib.dtglgp.jstatb = precision
             self._lib.dt_xsglau(
                 kin.p1.A or 1,
                 kin.p2.A or 1,
@@ -135,8 +131,6 @@ def _cross_section(self, kin=None, precision=None):
                 1,
                 1,
             )
-            if precision is not None:
-                self._lib.dtglgp.jstatb = saved
             return CrossSectionData(inelastic=self._lib.dtglxs.xspro[0, 0, 0])
         else:
             stot, sela = self._lib.dt_xshn(
@@ -146,6 +140,21 @@ def _cross_section(self, kin=None, precision=None):
 
         # TODO set more cross-sections
 
+    @property
+    def glauber_trials(self):
+        """Number of trials for Glauber model integration
+
+        Default is 1000 (set at model initialisation).
+        Larger number of `ntrials` reduces the fluctuations in the cross section,
+        thus, making it more smooth. Smaller number of `ntrials` makes calculations of
+        cross section faster.
+        """
+        return self._lib.dtglgp.jstatb
+
+    @glauber_trials.setter
+    def glauber_trials(self, ntrials):
+        self._lib.dtglgp.jstatb = ntrials
+
     def _set_kinematics(self, kin):
         # Save maximal mass that has been initialized
         # (DPMJET sometimes crashes if higher mass requested than initialized)

tests/test_dpmjetIII.py

@@ -0,0 +1,41 @@
+import numpy as np
+import chromo
+from .util import run_in_separate_process
+
+
+def run_cross_section_ntrials():
+    event_kin = chromo.kinematics.FixedTarget(1e4, "proton", "O16")
+    event_generator = chromo.models.DpmjetIII193(event_kin)
+
+    air = chromo.util.CompositeTarget([("N", 0.78), ("O", 0.22)])
+    event_kin = chromo.kinematics.FixedTarget(1e3, "proton", air)
+
+    default_precision = 1000
+    # Check the default precision
+    assert event_generator.glauber_trials == default_precision
+
+    # Set a new one
+    other_precision = 58
+    event_generator.glauber_trials = other_precision
+    assert event_generator.glauber_trials == other_precision
+
+    trials = 10
+
+    # With small precision
+    event_generator.glauber_trials = 1
+    cross_section_run1 = np.empty(trials, dtype=np.float64)
+    for i in range(trials):
+        cross_section_run1[i] = event_generator.cross_section(event_kin).inelastic
+
+    # With default precision
+    event_generator.glauber_trials = 1000
+    cross_section_run2 = np.empty(trials, dtype=np.float64)
+    for i in range(trials):
+        cross_section_run2[i] = event_generator.cross_section(event_kin).inelastic
+
+    # Standard deviation should be large for small precision
+    assert np.std(cross_section_run1) > np.std(cross_section_run2)
+
+
+def test_cross_section_ntrials():
+    run_in_separate_process(run_cross_section_ntrials)