-
Notifications
You must be signed in to change notification settings - Fork 12
Commit
This commit does not belong to any branch on this repository, and may belong to a fork outside of the repository.
- Loading branch information
1 parent
ce009b3
commit 4d9ff79
Showing
5 changed files
with
332 additions
and
56 deletions.
There are no files selected for viewing
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
| Original file line number | Diff line number | Diff line change |
|---|---|---|
|
|
@@ -28,6 +28,7 @@ python_requires = >=3.6 | |
|
|
||
| [options.extras_require] | ||
| test = | ||
| numpy | ||
| pytest | ||
| pytest-cov | ||
| graphviz | ||
|
|
||
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
| Original file line number | Diff line number | Diff line change |
|---|---|---|
| @@ -0,0 +1,182 @@ | ||
| #include "HepMC3/GenEvent.h" | ||
| #include "HepMC3/GenParticle.h" | ||
| #include "HepMC3/GenVertex.h" | ||
| #include "pybind.h" | ||
|
|
||
| #include <array> | ||
| #include <cassert> | ||
| #include <map> | ||
| #include <utility> | ||
| #include <vector> | ||
|
|
||
| // template <> | ||
| // struct std::hash<std::pair<int, int>> { | ||
| // std::size_t operator()(const std::pair<int, int>& p) const noexcept { | ||
| // auto h1 = std::hash<int>{}(p.first); | ||
| // auto h2 = std::hash<int>{}(p.second); | ||
| // return h1 ^ (h2 << 1); // or use boost::hash_combine | ||
| // } | ||
| // }; | ||
|
|
||
| template <> | ||
| struct std::less<std::pair<int, int>> { | ||
| bool operator()(const std::pair<int, int>& a, | ||
| const std::pair<int, int>& b) const noexcept { | ||
| if (a.first == b.first) return a.second < b.second; | ||
| return a.first < b.first; | ||
| } | ||
| }; | ||
|
|
||
| void normalize(int& m1, int& m2) { | ||
| // 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 && m2 < m1) std::swap(m1, m2); | ||
|
|
||
| if (m1 > 0 && m2 == 0) m2 = m1; | ||
|
|
||
| --m1; // fortran to c index | ||
| } | ||
|
|
||
| namespace HepMC3 { | ||
|
|
||
| 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) { | ||
|
|
||
| if (parents_or_children.request().ndim != 2) | ||
| throw std::runtime_error("parents or children must be 2D"); | ||
|
|
||
| auto rco = parents_or_children.unchecked<2>(); | ||
|
|
||
| const std::vector<GenParticlePtr>& particles = event.particles(); | ||
| const int n = particles.size(); | ||
| if (rco.shape(0) != n || rco.shape(1) != 2) | ||
| throw std::runtime_error("parents or children must have shape (N, 2)"); | ||
|
|
||
| // find unique vertices: | ||
| // particles with same parents or children share one vertex | ||
| // if parents: map from parents to children | ||
| // if children: map from children to parents | ||
| std::map<std::pair<int, int>, std::vector<int>> vmap; | ||
| for (int i = 0; i < n; ++i) { | ||
| if (rco(i, 0) == 0 && rco(i, 1) == 0) continue; | ||
| vmap[std::make_pair(rco(i, 0), rco(i, 1))].push_back(i); | ||
| } | ||
|
|
||
| const int has_vertex = !vx.is_none() + !vy.is_none() + !vz.is_none() + !vt.is_none(); | ||
|
|
||
| if (has_vertex && has_vertex != 4) | ||
| throw std::runtime_error("if one of vx, vy, vz, vt is set, all must be set"); | ||
|
|
||
| py::array_t<double> avx; | ||
| py::array_t<double> avy; | ||
| py::array_t<double> avz; | ||
| py::array_t<double> avt; | ||
|
|
||
| if (has_vertex) { | ||
| avx = py::cast<py::array_t<double>>(vx); | ||
| avy = py::cast<py::array_t<double>>(vy); | ||
| avz = py::cast<py::array_t<double>>(vz); | ||
| avt = py::cast<py::array_t<double>>(vt); | ||
|
|
||
| if (avx.request().ndim != 1 || avy.request().ndim != 1 || avz.request().ndim != 1 || | ||
| avt.request().ndim != 1) | ||
| throw std::runtime_error("vx, vy, vz, vt must be 1D"); | ||
| } | ||
|
|
||
| auto x = avx.unchecked<1>(); | ||
| auto y = avy.unchecked<1>(); | ||
| auto z = avz.unchecked<1>(); | ||
| auto t = avt.unchecked<1>(); | ||
|
|
||
| if (has_vertex && | ||
| (x.shape(0) != n || y.shape(0) != n || z.shape(0) != n || t.shape(0) != n)) | ||
| throw std::runtime_error("vx, vy, vz, vt must have same length"); | ||
|
|
||
| for (const auto& vi : vmap) { | ||
|
|
||
| int m1 = vi.first.first; | ||
| int m2 = vi.first.second; | ||
| // there must be at least one parent or child when we arrive here... | ||
| normalize(m1, m2); | ||
| assert(m1 < m2); | ||
|
|
||
| // ...with at least one child or parent | ||
| const auto& co = vi.second; | ||
| assert(!co.empty()); | ||
|
|
||
| FourVector pos; | ||
| if (has_vertex) { | ||
| // we assume this is a production vertex | ||
| // if parent, co.front() is location of production vertex of first child | ||
| // if child, we use location of first child m1 | ||
| const int i = parents ? co.front() : m1; | ||
| pos.set(x(i), y(i), z(i), t(i)); | ||
| } | ||
|
|
||
| GenVertexPtr v{new GenVertex(pos)}; | ||
|
|
||
| if (parents) { | ||
| for (int k = m1; k < m2; ++k) v->add_particle_in(particles.at(k)); | ||
| for (const auto k : co) v->add_particle_out(particles.at(k)); | ||
| } else { | ||
| for (int k = m1; k < m2; ++k) v->add_particle_out(particles.at(k)); | ||
| for (const auto k : co) v->add_particle_in(particles.at(k)); | ||
| } | ||
|
|
||
| event.add_vertex(v); | ||
| } | ||
| } | ||
|
|
||
| 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) { | ||
| 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"); | ||
| if (en.request().ndim != 1) throw std::runtime_error("en must be 1D"); | ||
| if (m.request().ndim != 1) throw std::runtime_error("m must be 1D"); | ||
| if (pid.request().ndim != 1) throw std::runtime_error("pid must be 1D"); | ||
| if (status.request().ndim != 1) throw std::runtime_error("status must be 1D"); | ||
|
|
||
| auto rpx = px.unchecked<1>(); | ||
| auto rpy = py.unchecked<1>(); | ||
| auto rpz = pz.unchecked<1>(); | ||
| auto ren = en.unchecked<1>(); | ||
| auto rm = m.unchecked<1>(); | ||
| auto rpid = pid.unchecked<1>(); | ||
| auto rsta = status.unchecked<1>(); | ||
|
|
||
| const int n = pid.shape(0); | ||
| if (rpx.shape(0) != n || rpy.shape(0) != n || rpz.shape(0) != n || | ||
| ren.shape(0) != n || rm.shape(0) != n) | ||
| throw std::runtime_error("px, py, pz, en, m, pid, status must have same length"); | ||
|
|
||
| event.clear(); | ||
| event.reserve(n); | ||
| event.set_event_number(event_number); | ||
|
|
||
| for (int i = 0; i < n; ++i) { | ||
| GenParticlePtr p{ | ||
| new GenParticle(FourVector(rpx(i), rpy(i), rpz(i), ren(i)), rpid(i), rsta(i))}; | ||
| p->set_generated_mass(rm(i)); | ||
| event.add_particle(p); | ||
| } | ||
|
|
||
| const bool have_parents = !parents.is_none(); | ||
| const bool have_children = !children.is_none(); | ||
| 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); | ||
| } | ||
|
|
||
| } // namespace HepMC3 |
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Oops, something went wrong.