Switch to generalized python reader

Author: dobrienphd

include/particlezoo/PDGParticleCodes.h

@@ -2,6 +2,7 @@
 
 #include <cstdint>
 #include <string_view>
+#include <vector>
 
 namespace ParticleZoo {
 
@@ -911,6 +912,27 @@ namespace ParticleZoo {
         return static_cast<std::int32_t>(type);
     }
 
+    /**
+     * @brief Get a complete list of particle names and types.
+     *
+     * Exposes all ParticleType values generated by PARTICLE_LIST for use in bindings.
+     * Returns a stable reference to a static vector of pairs (name, ParticleType).
+     */
+    inline const std::vector<std::pair<std::string_view, ParticleType>>& getAllParticleTypes() {
+        static const std::vector<std::pair<std::string_view, ParticleType>> list = []{
+            std::vector<std::pair<std::string_view, ParticleType>> v;
+            v.reserve(1024);
+            #define X(name, code) v.emplace_back(std::string_view(#name), ParticleType::name);
+            PARTICLE_LIST
+            #undef X
+            // Include sentinel values as well
+            v.emplace_back(std::string_view("Unsupported"), ParticleType::Unsupported);
+            v.emplace_back(std::string_view("PseudoParticle"), ParticleType::PseudoParticle);
+            return v;
+        }();
+        return list;
+    }
+
     #undef PARTICLE_LIST
 
 } // namespace ParticleZoo

python/README.md

@@ -7,6 +7,15 @@ This directory contains a Python package that exposes the C++ ParticleZoo API us
 - A C++20 compiler (g++ or clang++)
 - pybind11 (will be installed automatically by pip)
 
+## Setup
+Create a new Python virtual environment and activate it by running the following commands:
+
+```bash
+python -m venv .venv
+. .venv/bin/activate
+python -m pip install -U pip setuptools wheel pybind11
+```
+
 ## Build and install (editable)
 From the repository root or from this folder:
 

python/examples/read_phase_space.py

@@ -0,0 +1,101 @@
+#!/usr/bin/env python3
+"""
+Example: List supported formats and read a phase space file.
+
+Usage:
+  python read_phase_space.py /path/to/file.IAEAphsp [--format IAEA] [--limit N]
+  python read_phase_space.py --formats-only
+
+Arguments:
+  path                Path to a phase space file to read. If omitted with --formats-only,
+                      the script just lists formats and exits.
+
+Options:
+  --format NAME       Explicit format name (e.g., IAEA, EGS, TOPAS, penEasy).
+  --limit N           Print at most N particles (default: all).
+  --formats-only      Only list formats and exit.
+"""
+
+from __future__ import annotations
+import argparse
+import sys
+from typing import Optional
+
+import particlezoo as pz
+
+
+def print_supported_formats() -> None:
+    fmts = pz.FormatRegistry.supported_formats()
+    print("Supported formats ({}):".format(len(fmts)))
+    for f in fmts:
+        print("- {name:<8} ext={ext:<12} {desc}".format(
+            name=f.name, ext=f.file_extension, desc=f.description
+        ))
+
+
+def create_reader(path: str, fmt: Optional[str]) -> pz.PhaseSpaceFileReader:
+    if fmt:
+        return pz.create_reader_for_format(fmt, path)
+    return pz.create_reader(path)
+
+
+def print_particle_summary(p: pz.Particle, idx: int) -> None:
+    tname = pz.get_particle_type_name(p.type)
+    print(
+        f"[{idx}] type={tname:<20} KE={p.kinetic_energy:.6g}"
+        f" pos=({p.x:.6g},{p.y:.6g},{p.z:.6g})"
+        f" dir=({p.px:.6g},{p.py:.6g},{p.pz:.6g})"
+        f" w={p.weight:.6g} newHist={p.is_new_history}"
+    )
+
+
+def main(argv: list[str]) -> int:
+    ap = argparse.ArgumentParser(description="List formats and read a phase space file")
+    ap.add_argument("path", nargs="?", help="Path to phase space file")
+    ap.add_argument("--format", dest="format", help="Explicit format name to use")
+    ap.add_argument("--limit", type=int, default=None, help="Print at most N particles")
+    ap.add_argument("--formats-only", action="store_true", help="Only list formats and exit")
+    args = ap.parse_args(argv)
+
+    print_supported_formats()
+    if args.formats_only and not args.path:
+        return 0
+
+    if not args.path:
+        ap.error("path is required unless --formats-only is provided")
+
+    # Create reader
+    reader = create_reader(args.path, args.format)
+
+    # File summary
+    try:
+        n = reader.get_number_of_particles()
+    except Exception:
+        # Some formats may compute this lazily: fall back to counting
+        n = None
+    try:
+        h = reader.get_number_of_original_histories()
+    except Exception:
+        h = None
+    print(
+        f"\nReading: {reader.get_file_name()} (format={reader.get_phsp_format()})\n"
+        f"File size: {reader.get_file_size()} bytes\n"
+        f"Particles: {n if n is not None else 'unknown'}\n"
+        f"Original histories: {h if h is not None else 'unknown'}\n"
+    )
+
+    # Iterate particles
+    count = 0
+    for idx, p in enumerate(reader):
+        print_particle_summary(p, idx)
+        count += 1
+        if args.limit is not None and count >= args.limit:
+            break
+
+    print(f"\nRead {count} particles. Histories read: {reader.get_histories_read()}.")
+    reader.close()
+    return 0
+
+
+if __name__ == "__main__":
+    raise SystemExit(main(sys.argv[1:]))

python/src/particlezoo/__init__.py

@@ -4,7 +4,12 @@
     IntPropertyType,
     FloatPropertyType,
     BoolPropertyType,
-    IAEAReader,
+    PhaseSpaceFileReader,
+    SupportedFormat,
+    FormatRegistry,
+    create_reader,
+    create_reader_for_format,
+    all_particle_types,
     get_particle_type_from_pdgid,
     get_pdgid,
     get_particle_type_name,
@@ -17,7 +22,12 @@
     "IntPropertyType",
     "FloatPropertyType",
     "BoolPropertyType",
-    "IAEAReader",
+    "PhaseSpaceFileReader",
+    "SupportedFormat",
+    "FormatRegistry",
+    "create_reader",
+    "create_reader_for_format",
+    "all_particle_types",
     "get_particle_type_from_pdgid",
     "get_pdgid",
     "get_particle_type_name",

python/src/pybind/module.cpp

@@ -3,18 +3,34 @@
 
 #include "particlezoo/Particle.h"
 #include "particlezoo/PDGParticleCodes.h"
-#include "particlezoo/IAEA/IAEAphspFile.h"
+#include "particlezoo/PhaseSpaceFileReader.h"
+#include "particlezoo/utilities/formats.h"
 
 namespace py = pybind11;
 using namespace ParticleZoo;
 
 PYBIND11_MODULE(_pz, m) {
     m.doc() = "Python bindings for ParticleZoo core and IAEA reader";
 
-    // ParticleType enum (expose only sentinel values; use helpers for full mapping)
-    py::enum_<ParticleType>(m, "ParticleType")
-        .value("Unsupported", ParticleType::Unsupported)
-        .value("PseudoParticle", ParticleType::PseudoParticle);
+    // Ensure built-in formats are registered before any factory calls
+    try { FormatRegistry::RegisterStandardFormats(); } catch (...) { /* idempotent; ignore */ }
+
+    // ParticleType enum (full mapping from X-macro via helper in header)
+    auto pyParticleType = py::enum_<ParticleType>(m, "ParticleType");
+    for (const auto &entry : getAllParticleTypes()) {
+        pyParticleType.value(std::string(entry.first).c_str(), entry.second);
+    }
+    // ensure enum is closed/usable
+    pyParticleType.export_values();
+
+    // Convenience: return all particle types as a dict name -> enum value
+    m.def("all_particle_types", [](){
+        py::dict d;
+        for (const auto &entry : getAllParticleTypes()) {
+            d[py::str(entry.first)] = entry.second;
+        }
+        return d;
+    }, "Return a mapping of ParticleType names to enum values");
 
     // Property enums
     py::enum_<IntPropertyType>(m, "IntPropertyType")
@@ -75,26 +91,66 @@ PYBIND11_MODULE(_pz, m) {
         .def("set_incremental_histories", &Particle::setIncrementalHistories)
         ;
 
-    // IAEA Reader
-    using IAEAReader = IAEAphspFile::Reader;
-    py::class_<IAEAReader>(m, "IAEAReader")
-        .def(py::init<const std::string&, const UserOptions&>(),
-             py::arg("filename"), py::arg("options") = UserOptions{})
-        .def("get_number_of_particles", static_cast<std::uint64_t (IAEAReader::*)() const>(&IAEAReader::getNumberOfParticles))
-        .def("get_number_of_particles_by_type", static_cast<std::uint64_t (IAEAReader::*)(ParticleType) const>(&IAEAReader::getNumberOfParticles), py::arg("particle_type"))
-    .def("get_number_of_original_histories", &IAEAReader::getNumberOfOriginalHistories)
-    .def("has_more_particles", &IAEAReader::hasMoreParticles)
-    .def("get_next_particle", static_cast<Particle (IAEAReader::*)()>(&IAEAReader::getNextParticle))
-    .def("get_file_size", &IAEAReader::getFileSize)
-    .def("get_file_name", &IAEAReader::getFileName)
-    .def("close", &IAEAReader::close)
-        .def("__iter__", [](IAEAReader &self) -> IAEAReader& { return self; }, py::return_value_policy::reference_internal)
-        .def("__next__", [](IAEAReader &self) {
+    // PhaseSpaceFileReader (abstract base; created via FormatRegistry factories)
+    py::class_<PhaseSpaceFileReader, std::unique_ptr<PhaseSpaceFileReader>>(m, "PhaseSpaceFileReader")
+        .def("get_number_of_particles", &PhaseSpaceFileReader::getNumberOfParticles)
+        .def("get_number_of_original_histories", &PhaseSpaceFileReader::getNumberOfOriginalHistories)
+        .def("get_histories_read", &PhaseSpaceFileReader::getHistoriesRead)
+        .def("get_particles_read", static_cast<std::uint64_t (PhaseSpaceFileReader::*)()>(&PhaseSpaceFileReader::getParticlesRead))
+        .def("has_more_particles", &PhaseSpaceFileReader::hasMoreParticles)
+        .def("get_next_particle", static_cast<Particle (PhaseSpaceFileReader::*)()>(&PhaseSpaceFileReader::getNextParticle))
+        .def("get_file_size", &PhaseSpaceFileReader::getFileSize)
+        .def("get_file_name", &PhaseSpaceFileReader::getFileName)
+        .def("get_phsp_format", &PhaseSpaceFileReader::getPHSPFormat)
+        .def("move_to_particle", &PhaseSpaceFileReader::moveToParticle, py::arg("index"))
+        .def("is_x_constant", &PhaseSpaceFileReader::isXConstant)
+        .def("is_y_constant", &PhaseSpaceFileReader::isYConstant)
+        .def("is_z_constant", &PhaseSpaceFileReader::isZConstant)
+        .def("is_px_constant", &PhaseSpaceFileReader::isPxConstant)
+        .def("is_py_constant", &PhaseSpaceFileReader::isPyConstant)
+        .def("is_pz_constant", &PhaseSpaceFileReader::isPzConstant)
+        .def("is_weight_constant", &PhaseSpaceFileReader::isWeightConstant)
+        .def("get_constant_x", &PhaseSpaceFileReader::getConstantX)
+        .def("get_constant_y", &PhaseSpaceFileReader::getConstantY)
+        .def("get_constant_z", &PhaseSpaceFileReader::getConstantZ)
+        .def("get_constant_px", &PhaseSpaceFileReader::getConstantPx)
+        .def("get_constant_py", &PhaseSpaceFileReader::getConstantPy)
+        .def("get_constant_pz", &PhaseSpaceFileReader::getConstantPz)
+        .def("get_constant_weight", &PhaseSpaceFileReader::getConstantWeight)
+        .def("close", &PhaseSpaceFileReader::close)
+        .def("__iter__", [](PhaseSpaceFileReader &self) -> PhaseSpaceFileReader& { return self; }, py::return_value_policy::reference_internal)
+        .def("__next__", [](PhaseSpaceFileReader &self) {
             if (!self.hasMoreParticles()) throw py::stop_iteration();
             return self.getNextParticle();
         })
         ;
 
+    // SupportedFormat struct
+    py::class_<SupportedFormat>(m, "SupportedFormat")
+        .def_property_readonly("name", [](const SupportedFormat& f){ return f.name; })
+        .def_property_readonly("description", [](const SupportedFormat& f){ return f.description; })
+        .def_property_readonly("file_extension", [](const SupportedFormat& f){ return f.fileExtension; })
+        .def_property_readonly("file_extension_can_have_suffix", [](const SupportedFormat& f){ return f.fileExtensionCanHaveSuffix; })
+        ;
+
+    // FormatRegistry static interface
+    py::class_<FormatRegistry>(m, "FormatRegistry")
+        .def_static("register_standard_formats", &FormatRegistry::RegisterStandardFormats)
+        .def_static("supported_formats", &FormatRegistry::SupportedFormats)
+        .def_static("formats_for_extension", &FormatRegistry::FormatsForExtension, py::arg("extension"))
+        .def_static("extension_for_format", &FormatRegistry::ExtensionForFormat, py::arg("format_name"))
+        ;
+
+    // Convenience factory functions for readers
+    m.def("create_reader", [](const std::string& filename){
+        FormatRegistry::RegisterStandardFormats();
+        return FormatRegistry::CreateReader(filename, UserOptions{});
+    }, py::arg("filename"));
+    m.def("create_reader_for_format", [](const std::string& format, const std::string& filename){
+        FormatRegistry::RegisterStandardFormats();
+        return FormatRegistry::CreateReader(format, filename, UserOptions{});
+    }, py::arg("format_name"), py::arg("filename"));
+
     // Expose a simple alias to create a Particle by PDG code
     m.def("particle_from_pdg", [](int pdg, float kineticEnergy, float x, float y, float z, float px, float py, float pz, bool isNewHistory, float weight){
         ParticleType t = getParticleTypeFromPDGID(static_cast<std::int32_t>(pdg));