Adapt to zero based indices (#69)

The method `GenEvent.from_hepevt` now has an additional argument
`fortran` which allows one to define whether the input uses 1-based
indices or 0-based indices.

---------

Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com>

Author: HDembinski

.pre-commit-config.yaml

@@ -35,7 +35,7 @@ repos:
 
 # Python formatting
 - repo: https://github.com/psf/black-pre-commit-mirror
-  rev: 23.1.0
+  rev: 23.7.0
   hooks:
   - id: black
 
@@ -48,14 +48,14 @@ repos:
 
 # Python linter (Flake8)
 - repo: https://github.com/pycqa/flake8
-  rev: 6.0.0
+  rev: 6.1.0
   hooks:
   - id: flake8
     exclude: ^(docs/.*|tests/.*)$
 
 # C++ formatting
 - repo: https://github.com/pre-commit/mirrors-clang-format
-  rev: v15.0.7
+  rev: v16.0.6
   hooks:
   - id: clang-format
 
@@ -70,7 +70,7 @@ repos:
 
 # Python type checking
 - repo: https://github.com/pre-commit/mirrors-mypy
-  rev: 'v0.991'
+  rev: 'v1.5.1'
   hooks:
   - id: mypy
     # additional_dependencies: [numpy]

docs/examples/basics.ipynb

@@ -1,5 +1,5 @@
 {
- "cells": [
+  "cells": [
   {
    "cell_type": "markdown",
    "metadata": {},

docs/examples/processing.ipynb

@@ -1,5 +1,5 @@
 {
- "cells": [
+  "cells": [
   {
    "attachments": {},
    "cell_type": "markdown",

src/core.cpp

@@ -63,7 +63,7 @@ void from_hepevt(GenEvent& event, int event_number, py::array_t<double> px,
                  py::array_t<double> py, py::array_t<double> pz, py::array_t<double> en,
                  py::array_t<double> m, py::array_t<int> pid, py::array_t<int> status,
                  py::object parents, py::object children, py::object vx, py::object vy,
-                 py::object vz, py::object vt);
+                 py::object vz, py::object vt, bool fortran);
 
 } // namespace HepMC3
 
@@ -522,7 +522,8 @@ PYBIND11_MODULE(_core, m) {
       .def("from_hepevt", from_hepevt, "event_number"_a, "px"_a, "py"_a, "pz"_a, "en"_a,
            "m"_a, "pid"_a, "status"_a, "parents"_a = py::none(),
            "children"_a = py::none(), "vx"_a = py::none(), "vy"_a = py::none(),
-           "vz"_a = py::none(), "vt"_a = py::none(), DOC(GenEvent.from_hepevt))
+           "vz"_a = py::none(), "vt"_a = py::none(), "fortran"_a = true,
+           DOC(GenEvent.from_hepevt))
       .def("write_data", &GenEvent::write_data, "data"_a, DOC(GenEvent.write_data))
       .def("read_data", &GenEvent::read_data, "data"_a, DOC(GenEvent.read_data))
       .def_property_readonly("numpy", [](py::object self) { return NumpyAPI(self); })

src/from_hepevt.cpp

@@ -18,21 +18,33 @@ struct std::less<std::pair<int, int>> {
   }
 };
 
-void normalize(int& m1, int& m2) {
+void normalize(int& m1, int& m2, bool fortran) {
   // normalize mother range, see
   // https://pythia.org/latest-manual/ParticleProperties.html
-  // m1 == m2 == 0 : no mothers
-  // m1 == m2 > 0 : elastic scattering
-  // m1 > 0, m2 == 0 : one mother
-  // m1 < m2, both > 0: interaction
-  // m2 < m1, both > 0: same, needs swapping
-
-  if (m1 > 0 && m2 == 0)
+  //   m1 == m2 == 0 : no mothers
+  //   m1 == m2 > 0 : elastic scattering
+  //   m1 > 0, m2 == 0 : one mother
+  //   m1 < m2, both > 0: interaction
+  //   m2 < m1, both > 0: same, needs swapping
+  // If the input is coming from C, all indices one less.
+  // After normalization, we want m1, m2 to be a valid
+  // c-style integer range, where m2 points one past the
+  // last included index or m1 == m2 (empty range).
+
+  const int invalid = fortran ? 0 : -1;
+
+  if (m1 > invalid && m2 == invalid)
     m2 = m1;
   else if (m2 < m1)
     std::swap(m1, m2);
 
-  --m1; // fortran to c index
+  if (fortran)
+    --m1;
+  else
+    ++m2;
+
+  // postcondition after normalize
+  assert((m1 == 0 && m2 == 0) || m1 < m2);
 }
 
 namespace HepMC3 {
@@ -52,7 +64,8 @@ void modded_add_particle_in(int vid, GenVertexPtr v, GenParticlePtr p) {
 
 void connect_parents_and_children(GenEvent& event, bool parents,
                                   py::array_t<int> parents_or_children, py::object vx,
-                                  py::object vy, py::object vz, py::object vt) {
+                                  py::object vy, py::object vz, py::object vt,
+                                  bool fortran) {
 
   if (parents_or_children.request().ndim != 2)
     throw std::runtime_error("parents or children must be 2D");
@@ -69,8 +82,9 @@ void connect_parents_and_children(GenEvent& event, bool parents,
   // if parents: map from parents to children
   // if children: map from children to parents
   std::map<std::pair<int, int>, std::vector<int>> vmap;
+  const int invalid = fortran ? 0 : -1;
   for (int i = 0; i < n; ++i) {
-    if (rco(i, 0) == 0 && rco(i, 1) == 0) continue;
+    if (rco(i, 0) <= invalid && rco(i, 1) <= invalid) continue;
     vmap[std::make_pair(rco(i, 0), rco(i, 1))].push_back(i);
   }
 
@@ -110,8 +124,7 @@ void connect_parents_and_children(GenEvent& event, bool parents,
     int m2 = vi.first.second;
 
     // there must be at least one parent or child when we arrive here...
-    normalize(m1, m2);
-    assert(m1 < m2); // postcondition after normalize
+    normalize(m1, m2, fortran);
 
     if (m1 < 0 || m2 > n) {
       std::ostringstream os;
@@ -158,7 +171,7 @@ void from_hepevt(GenEvent& event, int event_number, py::array_t<double> px,
                  py::array_t<double> py, py::array_t<double> pz, py::array_t<double> en,
                  py::array_t<double> m, py::array_t<int> pid, py::array_t<int> status,
                  py::object parents, py::object children, py::object vx, py::object vy,
-                 py::object vz, py::object vt) {
+                 py::object vz, py::object vt, bool fortran) {
   if (px.request().ndim != 1) throw std::runtime_error("px must be 1D");
   if (py.request().ndim != 1) throw std::runtime_error("py must be 1D");
   if (pz.request().ndim != 1) throw std::runtime_error("pz must be 1D");
@@ -196,7 +209,8 @@ void from_hepevt(GenEvent& event, int event_number, py::array_t<double> px,
   if (have_parents || have_children)
     connect_parents_and_children(
         event, have_parents,
-        py::cast<py::array_t<int>>(have_parents ? parents : children), vx, vy, vz, vt);
+        py::cast<py::array_t<int>>(have_parents ? parents : children), vx, vy, vz, vt,
+        fortran);
 }
 
 } // namespace HepMC3

src/numpy_api.cpp

@@ -26,28 +26,28 @@ void register_numpy_api(py::module& m) {
   py::class_<ParticlesAPI>(m, "ParticlesAPI")
       .def(py::init<py::object>())
       // clang-format off
-        NP_ARRAY(ParticlesAPI, particles, int, id)
-        NP_ARRAY(ParticlesAPI, particles, int, pid)
-        NP_ARRAY(ParticlesAPI, particles, int, status)
-        NP_ARRAY(ParticlesAPI, particles, bool, is_generated_mass_set)
-        NP_ARRAY(ParticlesAPI, particles, double, generated_mass)
-        NP_ARRAY2(ParticlesAPI, particles, double, momentum, px)
-        NP_ARRAY2(ParticlesAPI, particles, double, momentum, py)
-        NP_ARRAY2(ParticlesAPI, particles, double, momentum, pz)
-        NP_ARRAY2(ParticlesAPI, particles, double, momentum, e)
+      NP_ARRAY(ParticlesAPI, particles, int, id)
+      NP_ARRAY(ParticlesAPI, particles, int, pid)
+      NP_ARRAY(ParticlesAPI, particles, int, status)
+      NP_ARRAY(ParticlesAPI, particles, bool, is_generated_mass_set)
+      NP_ARRAY(ParticlesAPI, particles, double, generated_mass)
+      NP_ARRAY2(ParticlesAPI, particles, double, momentum, px)
+      NP_ARRAY2(ParticlesAPI, particles, double, momentum, py)
+      NP_ARRAY2(ParticlesAPI, particles, double, momentum, pz)
+      NP_ARRAY2(ParticlesAPI, particles, double, momentum, e)
       // clang-format on
       ;
 
   py::class_<VerticesAPI>(m, "VerticesAPI")
       .def(py::init<py::object>())
       // clang-format off
-        NP_ARRAY(VerticesAPI, vertices, int, id)
-        NP_ARRAY(VerticesAPI, vertices, int, status)
-        NP_ARRAY(VerticesAPI, vertices, bool, has_set_position)
-        NP_ARRAY2(VerticesAPI, vertices, double, position, x)
-        NP_ARRAY2(VerticesAPI, vertices, double, position, y)
-        NP_ARRAY2(VerticesAPI, vertices, double, position, z)
-        NP_ARRAY2(VerticesAPI, vertices, double, position, t)
+      NP_ARRAY(VerticesAPI, vertices, int, id)
+      NP_ARRAY(VerticesAPI, vertices, int, status)
+      NP_ARRAY(VerticesAPI, vertices, bool, has_set_position)
+      NP_ARRAY2(VerticesAPI, vertices, double, position, x)
+      NP_ARRAY2(VerticesAPI, vertices, double, position, y)
+      NP_ARRAY2(VerticesAPI, vertices, double, position, z)
+      NP_ARRAY2(VerticesAPI, vertices, double, position, t)
       // clang-format on
       ;
 

src/pyhepmc/_doc.py

@@ -54,7 +54,11 @@
     status: array-like
         Status codes of particles.
     parents: array-like or None, optional
-        Array of int with shape (N, 2). Start and stop index (inclusive) in the record for the parents of each particle in Fortran numbering (starting with 1). No parents for a particle are indicated by the pair (0, 0) or (-1, -1). Single parents are indicated either by (N, 0) or (N, N) with N > 0.
+        Array of int with shape (N, 2). Start and stop index (inclusive) in the record
+        for the parents of each particle. Indices must be zero-based (C) or one-based
+        (Fortran). In the latter case, keyword `subtract_one` must be set to true. No
+        parents for a particle are indicated by the pair (-1, -1). Single parents are
+        indicated either by (N, -1) or (N, N) with N > 0.
     children: array-like or None, optional
         Like parents but for the children of this particle.
     vx : array-like or None, optional
@@ -65,6 +69,9 @@
         Z-component of location of production vertex of particles in mm.
     vt : array-like or None, optional
         Time (ct) of production vertex of particles in mm.
+    fortran : bool, optional
+        If True (default), the source indices are 1-based (Fortran, Pythia8). Set this
+        to False, if the indices are 0-based (C-style).
     """,
     "GenEvent.weight": """Get event weight accessed by index (or the canonical/first one if there is no argument) or name.
 

src/pyhepmc/io.py

@@ -212,6 +212,9 @@ class HepMCFile:
     IOError if reading or writing fails.
     """
 
+    _reader: Optional[ReaderMixin]
+    _writer: Optional[Union[WriterAscii, WriterAsciiHepMC2, WriterHEPEVT]]
+
     def __init__(
         self,
         fileobj: Filename,
@@ -246,8 +249,7 @@ def __init__(
 
                 mode += "b"
 
-            def open_file() -> Any:
-                return open(fn, mode)
+            open_file = lambda: open(fn, mode)  # noqa: E731
 
             self._close_file = True
 
@@ -310,6 +312,7 @@ def open_file() -> Any:
     __exit__ = _exit_close
 
     def __iter__(self) -> Any:
+        assert self._reader is not None
         return self._reader.__iter__()
 
     def flush(self) -> None:

tests/test_basic.py

@@ -1,6 +1,7 @@
 import pytest
 import pyhepmc as hep
 from numpy.testing import assert_equal
+import numpy as np
 
 
 def make_evt():
@@ -325,7 +326,8 @@ def test_GenEvent_generated_mass():
 
 
 @pytest.mark.parametrize("use_parent", (True, False))
-def test_GenEvent_from_hepevt(use_parent, evt):
+@pytest.mark.parametrize("fortran", (True, False))
+def test_GenEvent_from_hepevt(use_parent, fortran, evt):
     status = [p.status for p in evt.particles]
     pid = [p.pid for p in evt.particles]
     px = [p.momentum[0] for p in evt.particles]
@@ -346,9 +348,14 @@ def test_GenEvent_from_hepevt(use_parent, evt):
     #       /            p6
     #     p2
 
+    # fortran style
     parents = [(0, 0), (0, 0), (1, 1), (2, 2), (3, 4), (3, 4), (5, 5), (5, 5)]
     children = [(3, 0), (4, 0), (5, 6), (5, 6), (7, 8), (0, 0), (0, 0), (0, 0)]
 
+    if not fortran:
+        parents = np.subtract(parents, 1)
+        children = np.subtract(children, 1)
+
     ev = hep.GenEvent()
     if use_parent:
         ev.from_hepevt(
@@ -366,6 +373,7 @@ def test_GenEvent_from_hepevt(use_parent, evt):
             vy,
             vz,
             vt,
+            fortran=fortran,
         )
     else:
         ev.from_hepevt(
@@ -383,6 +391,7 @@ def test_GenEvent_from_hepevt(use_parent, evt):
             vy,
             vz,
             vt,
+            fortran=fortran,
         )
 
     # cannot be taken from HepEvt record, but is set for evt