xcoll 0.4.0__py3-none-any.whl → 0.5.0__py3-none-any.whl

xcoll/scattering_routines/geant4/collimasim/src/collimasim/BDSXtrackInterface.hh

@@ -0,0 +1,118 @@
+#include "BDSBunchSixTrackLink.hh"
+#include "BDSException.hh"
+#include "BDSIMLink.hh"
+#include "BDSIonDefinition.hh"
+#include "BDSParticleCoordsFull.hh"
+#include "BDSParticleDefinition.hh"
+#include "BDSPhysicsUtilities.hh"
+
+#include "G4Electron.hh"
+#include "G4GenericIon.hh"
+#include "G4IonTable.hh"
+#include "G4ParticleDefinition.hh"
+#include "G4ParticleTable.hh"
+#include "G4Types.hh"
+
+#include "CLHEP/Units/PhysicalConstants.h"
+#include "CLHEP/Units/SystemOfUnits.h"
+
+#include <algorithm>
+#include <cmath>
+#include <iostream>
+#include <set>
+#include <string>
+#include <vector>
+
+#include <pybind11/pybind11.h>
+#include <pybind11/numpy.h>
+
+namespace py = pybind11;
+
+// The struct is only used for inactive particle coodrinates for now
+struct XtrackCoordinates{
+    double x;
+    double y;
+    double px;
+    double py;
+    double zeta;
+    double delta;
+    double chi;
+    double charge_ratio;
+    double s;
+    int64_t pdgid;
+    int64_t trackid;
+    int64_t state;
+    int64_t at_element;
+    int64_t at_turn;
+};
+
+BDSParticleDefinition* PrepareBDSParticleDefition(long long int pdgIDIn, double momentumIn, 
+                                                  double kineticEnergyIn, double ionChargeIn);
+
+/// TODO: make a base class for the interface classes as there is a lot of shared functionality
+class XtrackInterface
+{
+public:
+    XtrackInterface() = delete;  // No default constructor
+
+    XtrackInterface(const std::string&  bdsimConfigFile,
+                    long long int       referencePdgIdIn,
+                    double              referenceEkIn,
+                    double              relativeEnergyCutIn,
+                    int                 seedIn,
+                    int                 referenceIonChargeIn=0,
+                    bool                batchMode=true);
+
+    virtual ~XtrackInterface();
+
+    void addCollimator(const  std::string&   name,
+                       const  std::string&   material,
+                       double lengthIn,
+                       double apertureLeftIn,
+                       double apertureRightIn,
+                       double rotationIn,
+                       double xOffsetIn,
+                       double yOffsetIn,
+                       double jawTiltLeft,
+                       double jawTiltRight,
+                       int    side);
+
+    void addParticles(const py::list& coordinates);
+
+    // This is not exposed to python
+    void addParticle(double xIn,
+                     double yIn,
+                     double pxIn,
+                     double pyIn,
+                     double ctIn,
+                     double deltapIn,
+                     double chiIn,
+                     double chargeRatioIn,
+                     double sIn,
+                     int64_t pdgIDIn,
+                     int64_t trackidIn);
+
+    void collimate();
+    void clearData();
+    void selectCollimator(const std::string& name);
+
+    double getReferenceMass() { return refParticleDefinition->Mass() / CLHEP::GeV; }
+
+    py::dict collimateReturn(const py::list& coordinates);
+
+private:
+    BDSIMLink* bds = nullptr;
+    BDSBunchSixTrackLink* stp = nullptr;
+    std::vector<char *> argv;
+
+    std::vector<bool> particleActiveState;
+    std::vector<XtrackCoordinates*> xtrackParticles;
+
+    BDSParticleDefinition* refParticleDefinition = nullptr;
+
+    std::string currentCollimatorName;
+
+    int64_t maxParticleID=0;
+    double relativeEnergyCut=0.0;
+    int seed=0;
+};

xcoll/scattering_routines/geant4/collimasim/src/collimasim/__init__.py

@@ -0,0 +1,8 @@
+from .g4interface import PyATInterface, XtrackInterface
+
+try:
+    import xtrack
+except ModuleNotFoundError:
+    print("Xtrack not installed, will not import related features")
+else:
+    from .xtrack_collimator import Geant4CollimationManager, Geant4Collimator

xcoll/scattering_routines/geant4/collimasim/src/collimasim/bindings.cpp

@@ -0,0 +1,63 @@
+#include <pybind11/pybind11.h>
+#include <pybind11/iostream.h>
+
+#include "BDSPyATInterface.hh"
+#include "BDSXtrackInterface.hh"
+
+#define STRINGIFY(x) #x
+#define MACRO_STRINGIFY(x) STRINGIFY(x)
+
+namespace py = pybind11;
+using namespace pybind11::literals;
+
+PYBIND11_MODULE(g4interface, m) {
+  m.doc() = R"pbdoc(Python interface to BDSIM (Geant4). The main purpose of the interface is to enable collimation studies, including particle-matter interaction in collimators, for pure particle tracking codes.)pbdoc";
+
+  py::class_<PyATInterface>(m, "PyATInterface")
+            .def(py::init<const std::string&, int, double, double, int, bool>(),
+                    "bdsimConfigFile"_a, "referencePdgId"_a, "referenceEk"_a,
+                    "relativeEnergyCut"_a, "seed"_a, "batchMode"_a=true)
+            .def("addCollimator", &PyATInterface::addCollimator,
+                 "name"_a, "material"_a, "length"_a, "aperture"_a, "rotation"_a,
+                 "xOffset"_a, "yOffset"_a,  "jawTiltLeft"_a, "jawTiltRight"_a, "side"_a,
+                 py::call_guard<py::scoped_ostream_redirect, py::scoped_estream_redirect>())
+            .def("addParticle",
+                 static_cast<void (PyATInterface::*)(const py::array_t<double>&)>(&PyATInterface::addParticle),
+                 "coordinate_array"_a)
+            .def("collimate", &PyATInterface::collimate,
+                 py::call_guard<py::scoped_ostream_redirect, py::scoped_estream_redirect>())
+            .def("selectCollimator", &PyATInterface::selectCollimator, "collimatorName"_a)
+            .def("collimateReturn", &PyATInterface::collimateReturn,
+                 py::call_guard<py::scoped_ostream_redirect, py::scoped_estream_redirect>())
+            .def("GetEnergyDifferential", &PyATInterface::GetEnergyDifferential)
+            .def("clearData", &PyATInterface::clearData);
+
+  py::class_<XtrackInterface>(m, "XtrackInterface")
+            .def(py::init<const std::string&, int, double, double, int, int, bool>(),
+                 "bdsimConfigFile"_a, "referencePdgId"_a, "referenceEk"_a,
+                 "relativeEnergyCut"_a, "seed"_a, "referenceIonCharge"_a=0,"batchMode"_a=true,
+                 py::call_guard<py::scoped_ostream_redirect, py::scoped_estream_redirect>())
+            .def("addCollimator", &XtrackInterface::addCollimator,
+                 "name"_a, "material"_a, "length"_a, "apertureLeft"_a, "apertureRight"_a, "rotation"_a,
+                 "xOffset"_a, "yOffset"_a, "jawTiltLeft"_a, "jawTiltRight"_a, "side"_a,
+                 py::call_guard<py::scoped_ostream_redirect, py::scoped_estream_redirect>())
+            /// The cast here to disambiguate overloaded methods. The other versions of
+            /// addParticles were removed, but keep the cast for reference for now.
+            .def("addParticles", static_cast<void (XtrackInterface::*)
+            (const py::list& coordinates)>(&XtrackInterface::addParticles), "coordinates"_a)
+            /// The C++ ostream and estream are redirected to the Python streams
+            /// so all io can be handled on the Python side
+            .def("collimate", &XtrackInterface::collimate,
+                 py::call_guard<py::scoped_ostream_redirect, py::scoped_estream_redirect>())
+            .def("selectCollimator", &XtrackInterface::selectCollimator, "collimatorName"_a)
+            .def("collimateReturn", &XtrackInterface::collimateReturn,
+                 py::call_guard<py::scoped_ostream_redirect, py::scoped_estream_redirect>())
+            .def("clearData", &XtrackInterface::clearData)
+            .def("getReferenceMass", &XtrackInterface::getReferenceMass);
+
+#ifdef VERSION_INFO
+  m.attr("__version__") = MACRO_STRINGIFY(VERSION_INFO);
+#else
+  m.attr("__version__") = "dev";
+#endif
+}

xcoll/scattering_routines/geant4/collimasim/src/collimasim/pyCollimatorPass.py

@@ -0,0 +1,142 @@
+import numpy
+
+import collimasim
+
+'''
+This is a specialised pass method that performs a full simulation of physical
+interactions in collimators using Geant4. The interface to BDSIM (Geant4) is
+provided by the collimasim package.
+
+The interface must be initialised first, with a list of parameters and a file
+containing the collimator specification. All collimators that will be used
+must be provided in the collimator file and are accessed by name. Currently,
+collimators are limited to block jaw collimators, and the geometry model is
+built automatically.
+
+The passmethods contains a function trackFunction(), used for tracking by pyat,
+and has to be in the python path or in at.integrators
+
+The element creation is as follows:
+
+collimator = elements.Element('MyColl', 1.0, PassMethod='pyCollimatorPass')
+
+To be noted:
+The name of a collimator using the pyCollimatorPass method must match a
+collimator defined during initialisation. For collimators where a physical
+interaction simulation is not required, please specify a different pass method,
+e.g. drift.
+'''
+
+bdsimlink = None # Keep this as a global variable for now
+
+def load_collgaps(collgaps_file):
+    coll_dict = {}
+
+    coll_data = numpy.genfromtxt(collgaps_file, dtype=None, encoding=None, comments='#')
+
+    for coll in coll_data:
+        name = coll[1].lower()
+        one_coll = {}
+
+        one_coll["material"] = coll[6]
+        one_coll["length"] = float(coll[7])
+        one_coll["angle"] = float(coll[2])
+        one_coll["halfgap"] = float(coll[5])
+        # The collgaps do not specify the offsets
+        # or if a collimator is one-sided
+        # TODO: add the closed orbit at the collimator
+        one_coll["xoffs"] = 0
+        one_coll["yoffs"] = 0
+        one_coll["side"] = 0 # 0 is two sided
+
+        coll_dict[name] = one_coll
+
+    return coll_dict
+
+
+def initialise(bdsim_config_file, collimator_info, reference_pdg_id,
+               reference_Ek, relative_energy_cut, seed):
+
+    global bdsimlink
+
+    # Initialise a code interface
+    if bdsimlink is not None:
+        return
+
+    bdsimlink = collimasim.PyATInterface(bdsimConfigFile=bdsim_config_file,
+                                         referencePdgId=reference_pdg_id,
+                                         referenceEk=reference_Ek,
+                                         relativeEnergyCut=relative_energy_cut,
+                                         seed=seed, batchMode=True)
+
+    if isinstance(collimator_info, dict):
+        collimators = collimator_info
+    else:
+        collimators = load_collgaps(collimator_info)
+
+    for name in collimators:
+        cdict = collimators[name]
+        material = cdict["material"]
+        length = cdict["length"]
+        angle = cdict["angle"]
+        halfgap = cdict["halfgap"]
+        xoffs = cdict["xoffs"]
+        yoffs = cdict["yoffs"]
+        side = cdict["side"]
+
+        bdsimlink.addCollimator(name, material, length, aperture=halfgap*2,
+                                rotation=angle, xOffset=xoffs, yOffset=yoffs,
+                                jawTiltLeft=0.0, jawTiltRight=0.0,
+                                side=side)
+
+
+def process_particles(collimator_name):
+    # MAD-X convention uses . in the names, while pyAT convention uses _
+    # Provide handling for MAD-X name to avoid trivial issues
+    try:
+        bdsimlink.selectCollimator(collimator_name)
+    except RuntimeError:
+        try:
+            collimator_name_madx = collimator_name.replace("_", ".").lower()
+            bdsimlink.selectCollimator(collimator_name_madx)
+        except RuntimeError:
+            raise ValueError(f"No collimator {collimator_name} (or "
+                             f"{collimator_name_madx} defined)")
+
+    print(f"Processing collimator: {collimator_name}")
+    # Perform the Geant4 simulation
+    bdsimlink.collimate()
+
+    # Collect the surviving particles - a 1D array with n * 6 elements
+    arr_out = bdsimlink.collimateReturn()
+
+    return arr_out
+
+
+def trackFunction(rin,elem=None):
+    if bdsimlink is None:
+        raise ValueError("The Geant4 interface is not initialised. Call"
+                         " pyCollimatorPass.initialise(...) before tracking")
+
+    bdsimlink.clearData() # Clear the old data
+
+    name = elem.FamName
+    # rin is a (6,n) array
+    for rtmp in numpy.atleast_2d(rin.T):
+        if hasattr(elem,'T1'): rtmp += elem.T1
+        if hasattr(elem,'R1'): rtmp[:] = numpy.dot(elem.R1,rtmp)
+        bdsimlink.addParticle(rtmp)
+
+    rout = process_particles(name)
+
+    # The returned array may be longer than the original
+    # due to secondary particles. Access by index and only
+    # change only the primary coordinates
+    for i in range(len(rin.T)):
+        part = rout.T[i, :]
+        if hasattr(elem,'R2'): part[:] = numpy.dot(elem.R2, part)
+        if hasattr(elem,'T2'): part += elem.T2
+
+        rin.T[i, :] = part
+
+    #return rin # This return value doesn't seem to be used by pyAT