Added additional bindings an a sample phase space writing script

Author: dobrienphd

.gitignore

@@ -2,5 +2,24 @@ config.status
 build
 tests
 .vscode
+.venv
 docs/latex
-docs/html
+docs/html
+python/build/
+python/dist/
+python/**/*.egg-info/
+python/**/__pycache__/
+python/**/*.pyc
+python/**/*.pyo
+python/**/*.so
+python/**/*.dylib
+python/**/*.dll
+python/.eggs/
+python/*.egg
+python/.pytest_cache/
+python/.coverage
+python/htmlcov/
+python/.tox/
+python/.venv/
+python/venv/
+python/ENV/

python/examples/read_phase_space.py

@@ -3,23 +3,29 @@
 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 <file> [options]
   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.
+Examples:
+  python read_phase_space.py file.IAEAphsp
+  python read_phase_space.py file.root --ROOT-format OpenGATE
+  python read_phase_space.py file.root --ROOT-format TOPAS --limit 10
 
-Options:
-  --format NAME       Explicit format name (e.g., IAEA, EGS, TOPAS, penEasy).
+This script passes all command line arguments through to ParticleZoo's C++ argument
+parser, which handles format-specific options automatically.
+
+Note: The filename must come first, before any optional arguments.
+
+Common options:
   --limit N           Print at most N particles (default: all).
   --formats-only      Only list formats and exit.
+  
+Format-specific options are automatically supported (e.g., --ROOT-format, --iaea-ignore-zlast, etc.)
+Run with --help to see all available options.
 """
 
 from __future__ import annotations
-import argparse
 import sys
-from typing import Optional
 
 import particlezoo as pz
 
@@ -33,39 +39,68 @@ def print_supported_formats() -> None:
         ))
 
 
-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)
+    # Convert to human-readable units for display
+    ke_mev = p.kinetic_energy / pz.MeV
+    x_cm = p.x / pz.cm
+    y_cm = p.y / pz.cm
+    z_cm = p.z / pz.cm
+    
     print(
-        f"[{idx}] type={tname:<20} KE={p.kinetic_energy:.6g}"
-        f" pos=({p.x:.6g},{p.y:.6g},{p.z:.6g})"
+        f"[{idx}] type={tname:<20} KE={ke_mev:.6g} MeV"
+        f" pos=({x_cm:.6g},{y_cm:.6g},{z_cm:.6g}) cm"
         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:
+def main(argv: list[str]) -> int:
+    # Parse arguments using ParticleZoo's C++ argument parser
+    # This will handle --help, --formats, and all format-specific options automatically
+    usage_message = (
+        "Usage: read_phase_space.py <file> [options]\n\n"
+        "Example script to read phase space files.\n\n"
+        "Options:\n"
+        "  --limit N           Print at most N particles\n"
+        "  --formats-only      Only list formats and exit\n"
+    )
+    
+    # Check for --formats-only flag manually before parsing
+    if "--formats-only" in argv:
+        print_supported_formats()
         return 0
+    
+    # Extract the limit option manually (not a format-specific option)
+    limit = None
+    clean_argv = []
+    skip_next = False
+    for i, arg in enumerate(argv):
+        if skip_next:
+            skip_next = False
+            continue
+        if arg == "--limit":
+            if i + 1 < len(argv):
+                try:
+                    limit = int(argv[i + 1])
+                    skip_next = True
+                    continue
+                except ValueError:
+                    pass
+        clean_argv.append(arg)
+    
+    # Parse remaining arguments through C++ parser (requires at least 1 positional arg: the file path)
+    options = pz.ArgParser.parse_args(clean_argv, usage_message, min_positional_args=1)
+    
+    # Extract the file path (first positional argument)
+    filepath = options.extract_positional(0)
+    
+    print_supported_formats()
+    print()
 
-    if not args.path:
-        ap.error("path is required unless --formats-only is provided")
-
-    # Create reader
-    reader = create_reader(args.path, args.format)
+    # Create reader with parsed options - format detection and options handled automatically
+    reader = pz.create_reader(filepath, options)
 
     # File summary
     try:
@@ -89,7 +124,7 @@ def main(argv: list[str]) -> int:
     for idx, p in enumerate(reader):
         print_particle_summary(p, idx)
         count += 1
-        if args.limit is not None and count >= args.limit:
+        if limit is not None and count >= limit:
             break
 
     print(f"\nRead {count} particles. Histories read: {reader.get_histories_read()}.")

python/examples/write_phase_space.py

@@ -0,0 +1,121 @@
+#!/usr/bin/env python3
+"""
+Example: Write a phase space file.
+
+Usage:
+  python write_phase_space.py <output_file> [options]
+
+Examples:
+  python write_phase_space.py output.IAEAphsp
+  python write_phase_space.py output.root --ROOT-format OpenGATE
+  python write_phase_space.py output.phsp --TOPAS-format binary
+
+This script demonstrates how to create a phase space file writer and write particles.
+"""
+
+from __future__ import annotations
+import sys
+import math
+
+import particlezoo as pz
+
+
+def create_sample_particles(count: int = 100) -> list[pz.Particle]:
+    """Create a list of sample particles for demonstration."""
+    particles = []
+    
+    for i in range(count):
+        # Create electrons with varying energies and positions
+        # Simulate a simple beam diverging from the origin
+        angle = (i / count) * 2 * pz.PI
+        energy = (1.0 + (i % 10) * 0.5) * pz.MeV  # 1.0 to 5.5 MeV
+        
+        # Position on a plane at z=0
+        x = math.cos(angle) * 0.5 * pz.cm  # 0.5 cm radius
+        y = math.sin(angle) * 0.5 * pz.cm
+        z = 0.0 * pz.cm
+        
+        # Direction slightly diverging
+        theta = (i / count) * 0.1 * pz.radian  # small divergence angle
+        phi = angle
+        dir_x = math.sin(theta) * math.cos(phi)
+        dir_y = math.sin(theta) * math.sin(phi)
+        dir_z = math.cos(theta)
+        
+        # Create electron (PDG code 11)
+        particle = pz.particle_from_pdg(
+            pdg=11,
+            kineticEnergy=energy,
+            x=x, y=y, z=z,
+            px=dir_x, py=dir_y, pz=dir_z,
+            isNewHistory=(i % 5 == 0),  # New history every 5 particles
+            weight=1.0
+        )
+        particles.append(particle)
+    
+    return particles
+
+
+def main(argv: list[str]) -> int:
+    usage_message = (
+        "Usage: write_phase_space.py <output_file> [options]\n\n"
+        "Example script to write phase space files.\n"
+    )
+    
+    # Extract particle count option manually (not a format-specific option)
+    particle_count = 100
+    clean_argv = []
+    skip_next = False
+    for i, arg in enumerate(argv):
+        if skip_next:
+            skip_next = False
+            continue
+        if arg == "--count":
+            if i + 1 < len(argv):
+                try:
+                    particle_count = int(argv[i + 1])
+                    skip_next = True
+                    continue
+                except ValueError:
+                    pass
+        clean_argv.append(arg)
+    
+    # Parse remaining arguments through C++ parser (requires at least 1 positional arg: the file path)
+    options = pz.ArgParser.parse_args(clean_argv, usage_message, min_positional_args=1)
+    
+    # Extract the file path (first positional argument)
+    output_path = options.extract_positional(0)
+    
+    print(f"Creating writer for: {output_path}")
+    
+    # Optional: Set up fixed values (constant properties for all particles)
+    # This is useful for optimization when certain values don't change
+    fixed_values = pz.FixedValues()
+    # Example: If all particles have the same Z position
+    fixed_values.z_is_constant = True
+    fixed_values.constant_z = 0.0 * pz.cm
+    
+    # Create writer with parsed options - format detection and options handled automatically
+    writer = pz.create_writer(output_path, options, fixed_values)
+    
+    print(f"Format: {writer.get_phsp_format()}")
+    print(f"Maximum supported particles: {writer.get_maximum_supported_particles()}")
+    print(f"\nGenerating and writing {particle_count} particles...")
+    
+    # Create and write sample particles
+    particles = create_sample_particles(particle_count)
+    for particle in particles:
+        writer.write_particle(particle)
+    
+    # Close the writer (flushes buffers and finalizes the file)
+    writer.close()
+    
+    print(f"\nSuccessfully wrote {writer.get_particles_written()} particles")
+    print(f"Total histories: {writer.get_histories_written()}")
+    print(f"Output file: {writer.get_file_name()}")
+    
+    return 0
+
+
+if __name__ == "__main__":
+    raise SystemExit(main(sys.argv[1:]))

python/pyproject.toml

@@ -8,7 +8,7 @@ build-backend = "setuptools.build_meta"
 
 [project]
 name = "particlezoo"
-version = "0.1.0"
+version = "1.1.0"
 description = "Python bindings for the ParticleZoo C++ library"
 readme = "README.md"
 requires-python = ">=3.8"

python/setup.py

@@ -1,4 +1,6 @@
 from pathlib import Path
+import subprocess
+import re
 from setuptools import setup
 from setuptools.command.build_ext import build_ext
 from pybind11.setup_helpers import Pybind11Extension, build_ext as build_ext_pybind
@@ -23,19 +25,82 @@
     str(Path("..") / "src" / "IAEA" / "IAEAphspFile.cc"),
 ]
 
+define_macros = [("PYBIND11_DETAILED_ERROR_MESSAGES", "1")]
+extra_compile_args = ["-O3", "-fvisibility=hidden", "-std=c++20"]
+extra_link_args = []
+
+# Try to read ROOT configuration from config.status
+config_status = proj / "config.status"
+use_root = False
+
+if config_status.exists():
+    print(f"Reading configuration from {config_status}")
+    config_vars = {}
+    with open(config_status, "r") as f:
+        for line in f:
+            # Parse Makefile-style variable assignments: VAR = value
+            match = re.match(r'^(\w+)\s*=\s*(.*)$', line.strip())
+            if match:
+                var_name, var_value = match.groups()
+                config_vars[var_name] = var_value.strip()
+    
+    if config_vars.get("USE_ROOT") == "1":
+        use_root = True
+        root_cflags = config_vars.get("ROOT_CFLAGS", "").split()
+        root_libs = config_vars.get("ROOT_LIBS", "").split()
+        
+        if root_cflags or root_libs:
+            print("ROOT support enabled (from config.status)")
+            define_macros.append(("USE_ROOT", "1"))
+            sources.append(str(Path("..") / "src" / "ROOT" / "ROOTphsp.cc"))
+            extra_compile_args.extend(root_cflags)
+            extra_link_args.extend(root_libs)
+        else:
+            print("WARNING: USE_ROOT=1 but no ROOT flags found")
+            use_root = False
+
+if not use_root and not config_status.exists():
+    # Fallback: try to detect ROOT directly if config.status doesn't exist
+    print("config.status not found, attempting to detect ROOT...")
+    try:
+        root_cflags = subprocess.check_output(
+            ["root-config", "--cflags"], 
+            stderr=subprocess.DEVNULL, 
+            text=True
+        ).strip().split()
+        root_libs = subprocess.check_output(
+            ["root-config", "--libs"], 
+            stderr=subprocess.DEVNULL, 
+            text=True
+        ).strip().split()
+        
+        if root_cflags and root_libs:
+            print("ROOT detected - enabling ROOT support")
+            define_macros.append(("USE_ROOT", "1"))
+            sources.append(str(Path("..") / "src" / "ROOT" / "ROOTphsp.cc"))
+            extra_compile_args.extend(root_cflags)
+            extra_link_args.extend(root_libs)
+        else:
+            print("ROOT found but flags empty - building without ROOT support")
+    except (subprocess.CalledProcessError, FileNotFoundError):
+        print("ROOT not found - building without ROOT support")
+
+if not use_root and config_status.exists():
+    print("ROOT support disabled (per config.status)")
+
 ext_modules = [
     Pybind11Extension(
         "particlezoo._pz",
         sources=[str(Path("src/pybind/module.cpp"))] + sources,
         include_dirs=include_dirs,
-        # Use C++20
         cxx_std=20,
-        define_macros=[("PYBIND11_DETAILED_ERROR_MESSAGES", "1")],
-        extra_compile_args=["-O3", "-fvisibility=hidden", "-std=c++20"],
+        define_macros=define_macros,
+        extra_compile_args=extra_compile_args,
+        extra_link_args=extra_link_args,
     )
 ]
 
 setup(
     cmdclass={"build_ext": build_ext_pybind},
     ext_modules=ext_modules,
-)
+)