add refill_decay_stack (#161)

When Pythia8 is used for decay of particles, it is required to put the
particle that should be decayed on a `event` stack . Currently it is
only possible to do like this:

```python
for i in range(len(stack)):
    pid = stack.pid[i]
    en = stack.energy[i]
    m = self._pythia.particleData.findParticle(pid).m0
    pz = sqrt((en - m)*(en + m))
    self._pythia.event.append(pid, 91, 0, 0, 0, 0, pz, en, m)
```

So, to put a lot of particles from vector stack, it is needed to call
`append` for every particle. This is time consuming because of Python
`for` loop and the call to `append` itself. The time of filling loop is
the same as the time required for decay in Pythia, which from
calculation point of view are not comparable tasks at all.

The added function `self._pythia.refill_decay_stack(stack.pid,
stack.energy)` accepts 2 numpy arrays with pid and energy, and calls the
above loop in C++. These significanly (~40 times for my test cases)
reduces the time for filling the `event` array.

src/chromo/models/pythia8.py

@@ -71,7 +71,7 @@ class Pythia8(MCRun):
         + "/releases/download/zipped_data_v1.0/Pythia8_v002.zip"
     )
 
-    def __init__(self, evt_kin, *, seed=None, config=None):
+    def __init__(self, evt_kin, *, seed=None, config=None, banner=True):
         """
 
         Parameters
@@ -95,15 +95,30 @@ def __init__(self, evt_kin, *, seed=None, config=None):
         # or init() may fail.
         if "PYTHIA8DATA" in environ:
             del environ["PYTHIA8DATA"]
-        self._pythia = self._lib.Pythia(datdir, True)
+        self._pythia = self._lib.Pythia(datdir, banner)
 
         if config is None:
             self._config = ["SoftQCD:inelastic = on"]
         else:
             self._config = self._parse_config(config)
 
+        # Common settings
+        self._config += [
+            # use our random seed
+            "Random:setSeed = on",
+            # Pythia's RANMAR PRNG accepts only seeds smaller than 900_000_000,
+            # this may change in the future if they switch to a different PRNG
+            f"Random:seed = {self.seed % 900_000_000}",
+            # reduce verbosity
+            "Print:quiet = on",
+        ]
+
         # must come last
-        self.kinematics = evt_kin
+        if evt_kin is None:
+            # Decay mode
+            self._init_pythia(self._config)
+        else:
+            self.kinematics = evt_kin
         self._set_final_state_particles()
 
     def _cross_section(self, kin=None):
@@ -119,22 +134,12 @@ def _cross_section(self, kin=None):
         )
 
     def _set_kinematics(self, kin):
-        pythia = self._pythia
-        pythia.settings.resetAll()
-
         config = self._config[:]
 
         # TODO use numpy PRNG instead of Pythia's
         config += [
-            # use our random seed
-            "Random:setSeed = on",
-            # Pythia's RANMAR PRNG accepts only seeds smaller than 900_000_000,
-            # this may change in the future if they switch to a different PRNG
-            f"Random:seed = {self.seed % 900_000_000}",
             # use center-of-mass frame
             "Beams:frameType = 1",
-            # reduce verbosity
-            "Print:quiet = on",
             # do not print progress
             "Next:numberCount = 0",
         ]
@@ -157,6 +162,12 @@ def _set_kinematics(self, kin):
             f"Beams:eCM = {kin.ecm}",
         ]
 
+        self._init_pythia(config)
+
+    def _init_pythia(self, config):
+        pythia = self._pythia
+        pythia.settings.resetAll()
+
         for line in config:
             if not pythia.readString(line):
                 raise RuntimeError(f"readString({line!r}) failed")

src/cpp/_pythia8.cpp

@@ -103,6 +103,32 @@ py::array_t<int> event_array_daughters(Event &event)
     return py::array_t<int>(shape, strides, ptr1);
 }
 
+// refills "event" stack with particles
+void fill(Event &event,
+          py::array_t<int> &pid,
+          py::array_t<int> &status,
+          py::array_t<double> &px,
+          py::array_t<double> &py,
+          py::array_t<double> &pz,
+          py::array_t<double> &energy,
+          py::array_t<double> &mass) 
+{
+    // Get a raw reference to numpy array
+    auto pid_ = pid.unchecked<1>();
+    auto status_ = status.unchecked<1>();
+    auto px_ = px.unchecked<1>();
+    auto py_ = py.unchecked<1>();
+    auto pz_ = pz.unchecked<1>();
+    auto energy_ = energy.unchecked<1>();
+    auto mass_ = mass.unchecked<1>();
+
+    event.reset();
+    for (int i = 0; i != pid.size(); ++i) {
+        event.append(pid_[i], status_[i], 0, 0,
+                     px_[i], py_[i], pz_[i], energy_[i], mass_[i]);
+    }    
+}
+
 PYBIND11_MODULE(_pythia8, m)
 {
     py::class_<ParticleData>(m, "ParticleData")
@@ -120,7 +146,6 @@ PYBIND11_MODULE(_pythia8, m)
                      pl.append(p.second);
                  return pl;
              })
-
         ;
 
     py::class_<ParticleDataEntry, ParticleDataEntryPtr>(m, "ParticleDataEntry")
@@ -236,7 +261,6 @@ PYBIND11_MODULE(_pythia8, m)
                      *ptr++ = charge_from_pid(self.particleData, pit->id());
                  return result;
              })
-
         ;
 
     py::class_<Event>(m, "Event")
@@ -257,6 +281,6 @@ PYBIND11_MODULE(_pythia8, m)
         .def("daughters", event_array_daughters)
         .def("reset", &Event::reset)
         .def("append", py::overload_cast<int, int, int, int, double, double, double, double, double, double, double>(&Event::append), "pdgid"_a, "status"_a, "col"_a, "acol"_a, "px"_a, "py"_a, "pz"_a, "e"_a, "m"_a = 0, "scale"_a = 0, "pol"_a = 9.)
-
+        .def("fill", &fill)
         ;
 }

tests/test_pythia8_fill.py

@@ -0,0 +1,98 @@
+import chromo
+from chromo.common import EventData
+from .util import run_in_separate_process
+import pytest
+import sys
+import numpy as np
+
+if sys.platform == "win32":
+    pytest.skip("Pythia8 does not run on windows", allow_module_level=True)
+
+
+def init_events(nevents):
+    evt_kin = chromo.kinematics.FixedTarget(1e7, "p", "p")
+    evt_gen = chromo.models.Sibyll23d(evt_kin)
+
+    events = []
+    for event in evt_gen(nevents):
+        event = event.final_state()
+        events.append(event.copy())
+
+    return events
+
+
+def get_pythia_event(pythia, evt_kin):
+    pevent = pythia.event
+    res = EventData(
+        ("pythia", "pythia"),
+        evt_kin,
+        0,
+        0.0,
+        (0, 0),
+        pevent.pid(),
+        pevent.status(),
+        pythia.charge(),
+        pevent.px(),
+        pevent.py(),
+        pevent.pz(),
+        pevent.en(),
+        pevent.m(),
+        pevent.vx(),
+        pevent.vy(),
+        pevent.vz(),
+        pevent.vt(),
+        np.maximum(pevent.mothers() - 1, -1),
+        np.maximum(pevent.daughters() - 1, -1),
+    )
+    return res.copy()
+
+
+def append_result(pythia, event):
+    pythia.event.reset()
+    for i in range(len(event)):
+        pythia.event.append(
+            event.pid[i],
+            91,
+            0,
+            0,
+            event.px[i],
+            event.py[i],
+            event.pz[i],
+            event.en[i],
+            event.m[i],
+        )
+    return get_pythia_event(pythia, event.kin)
+
+
+def fill_result(pythia, event):
+    pythia.event.fill(
+        event.pid, event.status + 90, event.px, event.py, event.pz, event.en, event.m
+    )
+    return get_pythia_event(pythia, event.kin)
+
+
+def run_event_fill():
+    """
+    Tests that `pythia.event.fill` and `pythia.event.append`
+    produce the same event stack
+
+    `pythia.event.fill` is optimized version of `pythia.event.append`
+    `pythia.event.fill` accepts numpy arrays instead of scalar value
+    as for `pythia.event.append`.
+
+    `pythia.event.fill` resets event stack by `pythia.event.reset()`
+    and uses `pythia.event.append` in C++ loop
+    """
+
+    config = ["ProcessLevel:all = off" "ParticleDecays:tau0Max = 1e100"]
+    pythia8 = chromo.models.Pythia8(None, seed=1, config=config, banner=False)
+    pythia = pythia8._pythia
+
+    for event in init_events(10):
+        fill_res = fill_result(pythia, event)
+        append_res = append_result(pythia, event)
+        assert fill_res == append_res
+
+
+def test_event_fill():
+    run_in_separate_process(run_event_fill)