colapy 0.0.1__tar.gz

colapy-0.0.1/.gitignore

@@ -0,0 +1,18 @@
+CMakeLists.txt.user
+CMakeCache.txt
+CMakeFiles
+CMakeScripts
+Testing
+Makefile
+cmake_install.cmake
+install_manifest.txt
+compile_commands.json
+CTestTestfile.cmake
+_deps
+CMakeUserPresets.json
+
+build/
+.cache/
+*.egg-info
+__pycache__
+.vscode/

colapy-0.0.1/CMakeLists.txt

@@ -0,0 +1,57 @@
+cmake_minimum_required(VERSION 3.15)
+project(
+    ${SKBUILD_PROJECT_NAME}
+    VERSION ${SKBUILD_PROJECT_VERSION}
+    LANGUAGES CXX
+)
+
+find_package(COLA)
+
+if (NOT COLA_FOUND)
+    include(FetchContent)
+    FetchContent_Declare(
+        COLA
+        GIT_REPOSITORY https://github.com/Spectator-matter-group-INR-RAS/COLA.git
+        GIT_TAG master
+    )
+    FetchContent_MakeAvailable(COLA)
+endif()
+
+set(CMAKE_CXX_STANDARD 17)
+set(CMAKE_CXX_STANDARD_REQUIRED ON)
+
+# Python and pybind11 setup
+find_package(pybind11 CONFIG REQUIRED)
+
+# Create Python module
+pybind11_add_module(
+    _cola_impl 
+    src/bindings.cpp
+)
+
+# Include directories
+# target_include_directories(_cola_impl PRIVATE src/cpp)
+target_link_libraries(_cola_impl PUBLIC COLA)
+
+target_compile_definitions(_cola_impl PRIVATE VERSION_INFO=${PROJECT_VERSION})
+
+# Install configuration
+install(TARGETS _cola_impl DESTINATION colapy)
+
+# Installation for Python package
+install(
+    DIRECTORY src/python/colapy/
+    DESTINATION .
+    FILES_MATCHING PATTERN "*.py"
+)
+
+# Testing
+# find_package(Python REQUIRED COMPONENTS Interpreter Development.Module)
+# enable_testing()
+#     if(Python_FOUND)
+#         add_test(
+#             NAME python_tests
+#             COMMAND ${Python_EXECUTABLE} -m pytest tests/
+#             WORKING_DIRECTORY ${CMAKE_CURRENT_BINARY_DIR}
+#         )
+# endif()

colapy-0.0.1/LICENSE

@@ -0,0 +1,21 @@
+MIT License
+
+Copyright (c) 2025 Novikov Artemii
+
+Permission is hereby granted, free of charge, to any person obtaining a copy
+of this software and associated documentation files (the "Software"), to deal
+in the Software without restriction, including without limitation the rights
+to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
+copies of the Software, and to permit persons to whom the Software is
+furnished to do so, subject to the following conditions:
+
+The above copyright notice and this permission notice shall be included in all
+copies or substantial portions of the Software.
+
+THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
+IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
+FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
+AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
+LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
+OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
+SOFTWARE.

colapy-0.0.1/MANIFEST.in

@@ -0,0 +1,4 @@
+include .gitignore
+include LICENSE
+include README.md
+exclude MANIFEST.in

colapy-0.0.1/PKG-INFO

@@ -0,0 +1,20 @@
+Metadata-Version: 2.2
+Name: colapy
+Version: 0.0.1
+Summary: A sample C++ library with Python bindings
+Keywords: cpp,bindings
+Author-Email: Artemii Novikov <artyasen@yandex.ru>
+License: MIT
+Classifier: Intended Audience :: Developers
+Classifier: License :: OSI Approved :: MIT License
+Classifier: Programming Language :: Python :: 3
+Classifier: Programming Language :: Python :: 3.8
+Classifier: Programming Language :: Python :: 3.9
+Classifier: Programming Language :: Python :: 3.10
+Classifier: Programming Language :: Python :: 3.11
+Classifier: Programming Language :: C++
+Project-URL: Repository, https://github.com/ZakayZ/COLA-PY
+Requires-Python: >=3.8
+Description-Content-Type: text/markdown
+
+# COLA-PY

colapy-0.0.1/README.md

@@ -0,0 +1 @@
+# COLA-PY

colapy-0.0.1/pyproject.toml

@@ -0,0 +1,50 @@
+[project]
+name = "colapy"
+version = "0.0.1"
+description = "A sample C++ library with Python bindings"
+readme = "README.md"
+license = {text = "MIT"}
+authors = [
+    {name = "Artemii Novikov", email = "artyasen@yandex.ru"}
+]
+classifiers = [
+    "Intended Audience :: Developers",
+    "License :: OSI Approved :: MIT License",
+    "Programming Language :: Python :: 3",
+    "Programming Language :: Python :: 3.8",
+    "Programming Language :: Python :: 3.9",
+    "Programming Language :: Python :: 3.10",
+    "Programming Language :: Python :: 3.11",
+    "Programming Language :: C++"
+]
+keywords = ["cpp", "bindings"]
+dependencies = [
+]
+requires-python = ">=3.8"
+
+[project.urls]
+# Homepage = "https://github.com/ZakayZ/COLA-PY"
+Repository = "https://github.com/ZakayZ/COLA-PY"
+# Documentation = "https://github.com/yourusername/my-library"
+
+
+[build-system]
+# Указываем, что для сборки нужны scikit-build-core и pybind11 (для CMake find_package)
+requires = [
+    "scikit-build-core>=0.5",
+    "cmake>=3.15",
+    "pybind11",
+]
+build-backend = "scikit_build_core.build"
+
+[tool.scikit-build]
+# Указываем, где лежит корневой CMakeLists.txt
+cmake.source-dir = "."
+# Указываем, куда scikit-build-core должен установить готовый модуль,
+# чтобы он попал в wheel-пакет под правильным именем.
+# В нашем случае имя модуля 'my_python_module' совпадает с именем таргета в CMake.
+# Если бы у вас были еще и чисто Python-файлы, вы бы указали их здесь:
+wheel.packages = [
+    "src/colapy"
+]
+

colapy-0.0.1/src/COLA.hh

@@ -0,0 +1,311 @@
+#ifndef COLA_COLA_HH
+#define COLA_COLA_HH
+
+#include <cmath>
+#include <iostream>
+#include <map>
+#include <memory>
+#include <queue>
+#include <vector>
+
+#include "LorentzVector.hh"
+
+namespace cola {
+    using LorentzVector = LorentzVectorImpl<double>;
+
+    /** A typedef representing mass and charge of a nucleon.
+     */
+    using AZ = std::pair<unsigned short, unsigned short>;
+
+    /** PDG code to AZ converter.
+     *  **WARNING:** this function is intended to process heavy ions PDG codes.
+     *  @param pdgCode PDG code of the ion.
+     *  @return AZ of the ion.
+     */
+    AZ pdgToAZ(int pdgCode);
+
+    /** AZ to PDG code convverter.
+     *  @param data AZ of the ion.
+     *  @return PDG code of the ion
+     */
+    int AZToPdg(AZ data);
+
+    /** \defgroup Data Data Classes and supporting methods.
+     * @{
+     */
+
+    /** Particle class by generator output.
+     *  This enum represents various outcomes of a generator event for every particle.
+     */
+    enum class ParticleClass: char {
+        produced,     /**< A particle that was not present in the starting nuclei. */
+        elasticA,     /**< A particle that was present in the projectile nucleus and has experienced only elastic interactions. */
+        elasticB,     /**< A particle that was present in the target nucleus and has experienced only elastic interactions. */
+        nonelasticA,  /**< A particle that was present in the projectile nucleus and has experienced at least one non-elastic interaction. */
+        nonelasticB,  /**< A particle that was present in the target nucleus and has experienced at least one non-elastic interaction. */
+        spectatorA,   /**< A particle that was present in the projectile nucleus and hasn't experienced any interactions. */
+        spectatorB    /**< A particle that was present in the projectile nucleus and hasn't experienced any interactions. */
+    };
+
+    /** Particle data.
+     *  A structure representing data about a single particle
+     */
+    struct Particle {
+        AZ getAZ() const;
+
+        LorentzVector position; /**< Position <t, x, y, z> vector. */
+
+        LorentzVector momentum; /**< Momentum <e, x, y, z> vector. */
+
+        int pdgCode;    /**< PDG code of the particle. */
+        ParticleClass pClass;   /**< Data about particle origin. See ParticleClass for more info.*/
+    };
+
+    /**
+     * Convenient typedef for Particle vector.
+     */
+    using EventParticles = std::vector<Particle>;
+
+    /** Initial state data.
+     *  This structure contains data about initial state of any given event.
+     */
+    struct EventIniState {
+        int pdgCodeA;           /**< PDF code of the projectile. */
+        int pdgCodeB;           /**< PDF code of the target. */
+
+        double pZA;             /** Axial momentum of the projectile */
+        double pZB;             /** Axial momentum of the target */
+        double energy;          /** Incidental energy of the event. Depending on pZB being zero, this is either \f$E/A\f$ of target or \f$\sqrt{s_{NN}}\f$. */
+
+        float sectNN;           /** Nucleon-Nucleon cross section from generator. */
+        float b;                /** Impact parameter of the event. */
+
+        int nColl;              /** Diagnostic. Total number of collisions. */
+        int nCollPP;            /** Diagnostic. Number of proton-proton. */
+        int nCollPN;            /** Diagnostic. Number of proton-neutron collisions. */
+        int nCollNN;            /** Diagnostic. Number of neutron-neutron collisions. */
+        int nPart;              /** Diagnostic. Total number of participants. */
+        int nPartA;             /** Diagnostic. Number of participants from the projectile nucleus. */
+        int nPartB;             /** Diagnostic. Number of participants from the target nucleus. */
+
+        float phiRotA;          /** Diagnostic. Polar angle \f$\phi\f$ of rotation of the projectile nucleon. */
+        float thetaRotA;        /** Diagnostic. Polar angle \f&\Theta\f$ of rotation of the projectile nucleon. */
+        float phiRotB;          /** Diagnostic. Polar angle \f$\phi\f$ of rotation of the target nucleon. */
+        float thetaRotB;        /** Diagnostic. Polar angle \f$\Theta\f$ of rotation of the target nucleon. */
+
+        EventParticles iniStateParticles;   /** The array of all Particles just before the event. */
+    };
+
+    /** A structure combining EventIniState and EventParticles of the event.
+     */
+    struct EventData {
+        EventIniState iniState;
+        EventParticles particles;
+    };
+
+    /** @}
+     * \defgroup Interface Pure abstract classes used for dependency injection.
+     * @{
+     */
+
+    /** A common abstract parent class representing any model in the pipeline.
+     *  A single model in the COLA-driven pipeline is named a Filter.
+     */
+    class VFilter {
+    public:
+        VFilter() = default;
+        VFilter (const VFilter&) = delete;
+        VFilter (VFilter&&) = delete;
+        VFilter& operator=(const VFilter&) = delete;
+        VFilter& operator=(VFilter&&) = delete;
+        virtual ~VFilter() = 0;
+    };
+
+    inline VFilter::~VFilter() = default;
+
+    /** Generator abstract class.
+     *  This is a generator interface. Generators are the first step of the MC simulation: they take data from existing
+     *  files or encapsulate nucleus-nucleus collision generators.
+     */
+    class VGenerator : public VFilter {
+    public:
+        VGenerator() = default;
+        VGenerator (const VGenerator&) = delete;
+        VGenerator (VGenerator&&) = delete;
+        VGenerator& operator=(const VGenerator&) = delete;
+        VGenerator& operator=(VGenerator&&) = delete;
+        ~VGenerator() override = 0;
+
+        /** A method to request one event from the generator.
+         *  Users are supposed to override this method to provide a single event from the generator model.
+         *  @return A pointer to the EventData of the produced event.
+         */
+        virtual std::unique_ptr<EventData> operator()() = 0;
+    };
+
+    inline VGenerator::~VGenerator() = default;
+
+    /** Converter abstract class.
+    *  This is a converter interface. It is inherited by all filters that are in the middle of MC simulation.
+    *  Converters are intended to transform data from previous steps.
+    */
+    class VConverter : public VFilter {
+    public:
+        VConverter() = default;
+        VConverter (const VConverter&) = delete;
+        VConverter (VConverter&&) = delete;
+        VConverter& operator=(const VConverter&) = delete;
+        VConverter& operator=(VConverter&&) = delete;
+        ~VConverter() override = 0;
+
+        /** A method to process one event by the converter model.
+         *  Users are supposed to override this method to process a single event by the converter model.
+         *  @param data A pointer to the EventData to be processed.
+         *  @return A pointer to the EventData of the processed event.
+         */
+        virtual std::unique_ptr<EventData> operator()(std::unique_ptr<EventData>&& data) = 0;
+    };
+
+    inline VConverter::~VConverter() = default;
+
+    /** Writer abstract class.
+     *  This is a writer interface. Writers are what the name suggests - they implement writing results into different
+     *  data formats. Note, that it isn't necessary to *only* write events to memory with this class. If you need to
+     *  process data in a way that the EventData is incompatible with the output (for example, you want to model a
+     *  detector response to the generated event) you can encapsulate these calculations in a writer class.
+     */
+    class VWriter : public VFilter {
+    public:
+        VWriter() = default;
+        VWriter (const VWriter&) = delete;
+        VWriter (VWriter&&) = delete;
+        VWriter& operator=(const VWriter&) = delete;
+        VWriter& operator=(VWriter&&) = delete;
+        ~VWriter() override = 0;
+
+        /** A method to save one event by the writer.
+         *  Users are supposed to override this method to save a single event. Note, that the implementation is left
+         *  to the user: you can save events to buffer with this method and write them to the filesystem in batches if
+         *  you please so.
+         *  @param data A pointer to the EventData to be saved.
+         */
+        virtual void operator()(std::unique_ptr<EventData>&& data) = 0;
+    };
+
+    inline VWriter::~VWriter() = default;
+
+    /** Factory abstract class.
+     * This is a factory interface. It generates a Filter with its VFactory::create method. DI in COLA works via using the
+     * factory classes, which are registered in a MetaProcessor instance.
+     */
+    class VFactory {
+    public:
+        VFactory() = default;
+        VFactory (const VFactory&) = delete;
+        VFactory (VFactory&&) = delete;
+        VFactory& operator=(const VFactory&) = delete;
+        VFactory& operator=(VFactory&&) = delete;
+        virtual ~VFactory() = default;
+
+        /** The method for constructing filters.
+         *  This method is intended to be called in MetaProcessor with key-value pairs obtained from a XML-file.
+         *  @param metaData A dictionary with key-value pairs needed for configuring a model.
+         *  @return A configured class that is a VFilter child.
+         */
+        virtual VFilter* create(const std::map<std::string, std::string>& metaData) = 0;
+    };
+
+
+    std::unique_ptr<EventData> operator|(const std::unique_ptr<VGenerator>&, const std::unique_ptr<VConverter>&);
+    std::unique_ptr<EventData> operator|(std::unique_ptr<EventData>&&, const std::unique_ptr<VConverter>&);
+    void operator|(std::unique_ptr<EventData>&&, const std::unique_ptr<VWriter>&);
+
+    /** An enum for marking Filter types.
+     */
+    enum class FilterType: char {
+        generator,
+        converter,
+        writer
+    };
+    /** @}
+     *  \defgroup Metadata Classes for constructing and running a model.
+     *  @{
+     */
+
+    /** A structure representing the model pipeline.
+     */
+    struct FilterEnsemble {
+        std::unique_ptr<VGenerator> generator;                  /**< Event generator. */
+        std::vector<std::unique_ptr<VConverter>> converters;    /**< Vector of converters, applied step-by-step. */
+        std::unique_ptr<VWriter> writer;                        /**< Writer to save the results. */
+    };
+
+    /** A class for processing meta information.
+     *  This class stores data about all available Filters and corresponding factory classes and uses it to create the
+     *  pipeline using its MetaProcessor:parse method to read all needed information from an XML-file.
+     */
+    class MetaProcessor {
+    public:
+
+        /** Default constructor.
+         */
+        MetaProcessor() = default;
+
+        /** Constructor with immediate factories registration.
+         * Note that unique pointers in the dictionary are invalidated.
+         * @param filterMap A dictionary with all relevant information. Note that unique pointers in the dictionary are invalidated.
+         */
+        explicit MetaProcessor(std::map<std::string, std::pair<std::unique_ptr<VFactory>, FilterType>>& filterMap);
+
+        ~MetaProcessor() = default;
+
+        /** A method for registering new Filters.
+         * @param factory A rvalue-reference to the Filter factory pointer
+         * @param name The name of the Filter.
+         * @param type The type of the Filter. See FilterType.
+         */
+        void reg(std::unique_ptr<VFactory>&& factory, const std::string& name, FilterType type);
+
+        /** A method to parse a XML-file to set up a configured FilterEnsemble.
+         *  This method opens an XML-file @param fName to get the information to set up the model.
+         *  Inside the root element should be one <generator> element followed by any number of <converter> elements
+         *  (none is possible) and, finally, a <writer> element. Each element must have "name" attribute followed by
+         *  any number of additional attributes. These attributes are then passed to the corresponding factory's VFactory::create
+         *  method as a dictionary with keys being attribute names and values - attribute values.
+         *  This method throws an error if a relevant factory isn't found or XML-file is malformed.
+         *  @param fName Name with the configuration XML-file.
+         *  @return A configured FilterEnsemble.
+         */
+        FilterEnsemble parse(const std::string& fName) const;
+
+    private:
+        std::map<std::string, std::unique_ptr<VFactory>> generatorMap;
+        std::map<std::string, std::unique_ptr<VFactory>> converterMap;
+        std::map<std::string, std::unique_ptr<VFactory>> writerMap;
+
+        void regGen(std::unique_ptr<VFactory>&& factory, const std::string& name){ generatorMap.emplace(name, std::move(factory)); }
+        void regConv(std::unique_ptr<VFactory>&& factory, const std::string& name){ converterMap.emplace(name, std::move(factory)); }
+        void regWrite(std::unique_ptr<VFactory>&& factory, const std::string& name){ writerMap.emplace(name, std::move(factory)); }
+    };
+
+    /** Manager class.
+     * Currently more of a boilerplate, but potentially useful to incorporate parallel computing.
+     */
+    class ColaRunManager {
+    public:
+        ColaRunManager() = delete;
+        /** A constructor that moves the configured FilterEnsemble into the manager.
+         * @param ensemble Configured model.
+         */
+        explicit ColaRunManager(FilterEnsemble&& ensemble) : filterEnsemble(std::move(ensemble)) {}
+        ~ColaRunManager() = default;
+        /** A method to run the resulting model @param n times.
+         * @param n Number of runs.
+         */
+        void run(int n = 1) const;
+    private:
+        FilterEnsemble filterEnsemble;
+    };
+} // cola
+
+#endif //COLA_COLA_HH

colapy-0.0.1/src/LorentzVector.hh

@@ -0,0 +1,199 @@
+#ifndef COLA_LORENTZVECTOR_HH
+#define COLA_LORENTZVECTOR_HH
+
+#include <array>
+#include <cmath>
+#include <iostream>
+
+namespace cola {
+
+    template <typename Type = double>
+    struct Vector3 {
+        Type x, y, z;
+    };
+
+    template <typename Type = double>
+    class LorentzVectorImpl {
+    public:
+        union {
+            Type e;
+            Type t;
+        };
+        Type x;
+        Type y;
+        Type z;
+
+    private:
+        using FieldPtr = Type LorentzVectorImpl::*;
+
+        static inline constexpr std::array<FieldPtr, 4> Fields_ = {
+            &LorentzVectorImpl::e,
+            &LorentzVectorImpl::x,
+            &LorentzVectorImpl::y,
+            &LorentzVectorImpl::z
+        };
+
+    public:
+        const Type& operator[](int i) const { return this->*Fields_[i]; }
+        Type& operator[](int i) { return this->*Fields_[i]; }
+
+        LorentzVectorImpl& operator+=(const LorentzVectorImpl& other) {
+            for (size_t i = 0; i < Fields_.size(); ++i) {
+                this->*Fields_[i] += other.*Fields_[i];
+            }
+            return *this;
+        }
+
+        LorentzVectorImpl& operator-=(const LorentzVectorImpl& other) {
+            for (size_t i = 0; i < Fields_.size(); ++i) {
+                this->*Fields_[i] -= other.*Fields_[i];
+            }
+            return *this;
+        }
+
+        LorentzVectorImpl& operator*=(Type scalar) {
+            for (size_t i = 0; i < Fields_.size(); ++i) {
+                this->*Fields_[i] *= scalar;
+            }
+            return *this;
+        }
+
+        LorentzVectorImpl& operator/=(Type scalar) {
+            for (size_t i = 0; i < Fields_.size(); ++i) {
+                this->*Fields_[i] /= scalar;
+            }
+            return *this;
+        }
+
+        // bx by and bz are projections of beta
+        LorentzVectorImpl& boost(Type bx, Type by, Type bz) {
+
+        Type b2 = bx * bx + by * by + bz * bz;
+
+        if (b2 >= 1)
+            throw std::runtime_error("Boost faster than speed of light.");
+        if (b2 <= .95 or not isSpaceLike()) {
+            Type ggamma = 1.0 / std::sqrt(1.0 - b2);
+            Type bp = bx * x + by * y + bz * z;
+            Type gamma2 = b2 > 0 ? (ggamma - 1.0) / b2 : 0.0;
+
+            x = x + gamma2 * bp * bx + ggamma * bx * t;
+            y = y + gamma2 * bp * by + ggamma * by * t;
+            z = z + gamma2 * bp * bz + ggamma * bz * t;
+            t = ggamma * (t + bp);
+
+            // for big betas use rapidities instead (B = R^-1B'R, where B' is axis boost (Oz here) and R is Rotation matrix for spatial part)
+        } else {
+            // calculate direction vector coordinates
+            Type b1 = std::sqrt(b2);
+            Vector3<Type> rVec{bx / b1, by / b1, bz / b1};
+            Vector3<Type> rBack = rotateUz({0, 0, 1}, rVec);
+
+            // rotate space vector so that boost direction is {0, 0, 1} in new coordinates
+            auto newCoord = rotateUz({x, y, z}, rVec);
+            x = newCoord.x, y = newCoord.y, z = newCoord.z;
+
+            boostAxisRapidity(std::atanh(b1)); // boost along Oz
+            // rotate back
+            newCoord = rotateUz({x, y, z}, rBack);
+            x = newCoord.x, y = newCoord.y, z = newCoord.z;
+        }
+
+        return *this;
+        }
+
+        // axis from 1 to 3 correspond to x-y-z. Note: this gives correct results only if other axes components are zero. TODO: Throw an error otherwise
+        LorentzVectorImpl& boostAxisRapidity(Type rapidity, uint axis=3u) {
+        if (axis > 3 or axis < 1)
+            throw std::runtime_error("Wrong axis in boostAxis. 1 for x, 2 for y, 3 for z.");
+        if (not isSpaceLike()) {
+            throw std::runtime_error("Rapidity calculation only viable for space-like 4-vectors. Use boost() instead");
+        }
+        Type inv = std::sqrt(e*e - this->*Fields_[axis]*this->*Fields_[axis]);
+        Type rRapidity = rapidity + .5 * (std::log(e + this->*Fields_[axis]) - std::log(e - this->*Fields_[axis]));
+        e = inv * std::cosh(rRapidity);
+        this->*Fields_[axis] = inv * std::sinh(rRapidity);
+
+        return *this;
+        }
+
+        Type mag2() const { return t*t - (x*x + y*y + z*z); }
+        Type mag() const { return std::sqrt(mag2()); }
+
+        bool isSpaceLike() const { return mag2() > 0; }
+        bool isLightLike() const { return mag2() == 0; }
+        bool isTimeLike() const { return mag2() < 0; }
+    };
+
+    template <typename Type>
+    LorentzVectorImpl<Type> operator+(const LorentzVectorImpl<Type>& a, const LorentzVectorImpl<Type>& b) {
+        auto res = a;
+        res += b;
+        return res;
+    }
+
+    template <typename Type>
+    LorentzVectorImpl<Type> operator-(const LorentzVectorImpl<Type>& a, const LorentzVectorImpl<Type>& b) {
+        auto res = a;
+        res -= b;
+        return res;
+    }
+
+    template <typename Type, typename Scalar>
+    LorentzVectorImpl<Type> operator*(const LorentzVectorImpl<Type>& vec, Scalar scalar) {
+        auto res = vec;
+        res *= scalar;
+        return res;
+    }
+
+    template <typename Type, typename Scalar>
+    LorentzVectorImpl<Type> operator*(Scalar scalar, const LorentzVectorImpl<Type>& vec) {
+        return vec * scalar;
+    }
+
+    template <typename Type, typename Scalar>
+    LorentzVectorImpl<Type> operator/(const LorentzVectorImpl<Type>& vec, Scalar scalar) {
+        auto res = vec;
+        res /= scalar;
+        return res;
+    }
+
+    template <typename Type>
+    bool operator==(const LorentzVectorImpl<Type>& a, const LorentzVectorImpl<Type>& b) {
+        return a.e == b.e && a.x == b.x && a.y == b.y && a.z == b.z;
+    }
+
+    template <typename Type>
+    bool operator!=(const LorentzVectorImpl<Type>& a, const LorentzVectorImpl<Type>& b) {
+        return !(a == b);
+    }
+
+    template <typename Type>
+    std::ostream& operator<<(std::ostream& out, const LorentzVectorImpl<Type>& vec) {
+        out << "(" << vec.e << ", " << vec.x << ", " << vec.y << ", " << vec.z << ")";
+        return out;
+    }
+
+    template <typename Type>
+    Vector3<Type> rotateUz(const Vector3<Type> stVec, const Vector3<Type> uzVec) {
+        // NewUzVector must be normalized !
+
+        Vector3<Type> resVec;
+        double up = uzVec.x*uzVec.x + uzVec.y*uzVec.y;
+
+        if (up > 0) {
+            up = std::sqrt(up);
+            resVec.x = (uzVec.x * uzVec.z * stVec.x - uzVec.y * stVec.y) / up + uzVec.x * stVec.z;
+            resVec.y = (uzVec.y * uzVec.z * stVec.x + uzVec.x * stVec.y) / up + uzVec.y * stVec.z;
+            resVec.z = -up * stVec.x + uzVec.z * stVec.z;
+        } else if (uzVec.z < 0.) {
+            resVec.x = -stVec.x;
+            resVec.y = stVec.y;
+            resVec.z = -stVec.z;
+        } // phi=0  teta=pi
+
+        return resVec;
+    }
+} // namespace cola
+
+#endif // COLA_LORENTZVECTOR_HH

colapy-0.0.1/src/bindings.cpp

@@ -0,0 +1,134 @@
+#include <pybind11/pybind11.h>
+#include <pybind11/stl.h>
+#include <pybind11/functional.h>
+#include <pybind11/operators.h>
+#include <pybind11/attr.h>
+#include <pybind11/native_enum.h>
+
+#include <COLA.hh>
+
+#include <sstream>
+#include <memory>
+
+#define STRINGIFY(x) #x
+#define MACRO_STRINGIFY(x) STRINGIFY(x)
+
+namespace py = pybind11;
+using namespace pybind11::literals;
+
+PYBIND11_MODULE(_cola_impl, m) {
+    m.doc() = "COLA wrapper library";
+
+#ifdef VERSION_INFO
+    m.attr("__version__") = MACRO_STRINGIFY(VERSION_INFO);
+#else
+    m.attr("__version__") = "dev";
+#endif
+
+    py::native_enum<cola::ParticleClass>(m, "ParticleClass", "enum.Enum")
+        .value("PRODUCED", cola::ParticleClass::produced)
+        .value("ELASTIC_A", cola::ParticleClass::elasticA)
+        .value("ELASTIC_B", cola::ParticleClass::elasticB)
+        .value("NON_ELASTIC_A", cola::ParticleClass::nonelasticA)
+        .value("NON_ELASTIC_B", cola::ParticleClass::nonelasticB)
+        .value("SPECTATOR_A", cola::ParticleClass::spectatorA)
+        .value("SPECTATOR_B", cola::ParticleClass::spectatorB)
+        .export_values()
+        .finalize()
+        ;
+
+    py::class_<cola::LorentzVector>(m, "LorentzVector")
+        .def(py::init<double, double, double, double>(), "e"_a = 0., "x"_a = 0., "y"_a = 0., "z"_a = 0.)
+        .def(py::init<cola::LorentzVector>())
+        .def_readwrite("x", &cola::LorentzVector::x, "x vector component")
+        .def_readwrite("y", &cola::LorentzVector::y, "y vector component")
+        .def_readwrite("z", &cola::LorentzVector::z, "z vector component")
+        .def_readwrite("e", &cola::LorentzVector::e, "e vector component")
+        .def_readwrite("t", &cola::LorentzVector::t, "t vector component")
+        .def(py::self + py::self)
+        .def(py::self - py::self)
+        .def(py::self += py::self)
+        .def(py::self -= py::self)
+        .def(py::self *= double())
+        .def(py::self /= double())
+        .def(double() * py::self)
+        .def(py::self * double())
+        .def(py::self / double())
+        .def(py::self == py::self)
+        // .def(-py::self)
+        .def("__copy__", [](const cola::LorentzVector& vec) {
+            return vec;
+        })
+        .def("__deepcopy__", [](const cola::LorentzVector& vec, py::dict) {
+            return vec;
+        }, "memo"_a)
+        .def("__repr__", [](const cola::LorentzVector& p) -> std::string {
+            auto ss = std::stringstream();
+            ss << p;
+            return std::move(ss).str();
+        })
+        ;
+
+    py::class_<cola::Particle>(m, "Particle")
+        .def(py::init<cola::LorentzVector, cola::LorentzVector, int, cola::ParticleClass>(),
+            "position"_a, "momentum"_a, "pdg_code"_a, "p_class"_a)
+        .def(py::init<cola::Particle>())
+        .def_readwrite("position", &cola::Particle::position)
+        .def_readwrite("momentum", &cola::Particle::momentum)
+        .def_readwrite("pdg_code", &cola::Particle::pdgCode)
+        .def_readwrite("p_class", &cola::Particle::pClass)
+        .def("get_az", [](const cola::Particle& p) -> std::pair<int, int> { return p.getAZ(); })
+        .def("__repr__", [](const cola::Particle& p) -> std::string {
+            auto ss = std::stringstream();
+            ss << "Particle(pdg=" << p.pdgCode << ", class=" << static_cast<int>(p.pClass) << ")";
+            return std::move(ss).str();
+        })
+        ;
+
+    py::class_<cola::EventIniState>(m, "EventInitialState")
+        .def(py::init<>())
+        .def_readwrite("pdg_code_a", &cola::EventIniState::pdgCodeA,
+                      "PDG code of the projectile")
+        .def_readwrite("pdg_code_b", &cola::EventIniState::pdgCodeB,
+                      "PDG code of the target")
+        .def_readwrite("pz_a", &cola::EventIniState::pZA,
+                      "Axial momentum of the projectile")
+        .def_readwrite("pz_b", &cola::EventIniState::pZB,
+                      "Axial momentum of the target")
+        .def_readwrite("energy", &cola::EventIniState::energy,
+                      "Incident energy of the event")
+        .def_readwrite("sect_nn", &cola::EventIniState::sectNN,
+                      "Nucleon-Nucleon cross section from generator")
+        .def_readwrite("b", &cola::EventIniState::b,
+                      "Impact parameter of the event")
+        .def_readwrite("n_coll", &cola::EventIniState::nColl,
+                      "Total number of collisions")
+        .def_readwrite("n_coll_pp", &cola::EventIniState::nCollPP,
+                      "Number of proton-proton collisions")
+        .def_readwrite("n_coll_pn", &cola::EventIniState::nCollPN,
+                      "Number of proton-neutron collisions")
+        .def_readwrite("n_coll_nn", &cola::EventIniState::nCollNN,
+                      "Number of neutron-neutron collisions")
+        .def_readwrite("n_part", &cola::EventIniState::nPart,
+                      "Total number of participants")
+        .def_readwrite("n_part_a", &cola::EventIniState::nPartA,
+                      "Number of participants from projectile nucleus")
+        .def_readwrite("n_part_b", &cola::EventIniState::nPartB,
+                      "Number of participants from target nucleus")
+        .def_readwrite("phi_rot_a", &cola::EventIniState::phiRotA,
+                      "Polar angle φ of rotation of projectile nucleon")
+        .def_readwrite("theta_rot_a", &cola::EventIniState::thetaRotA,
+                      "Polar angle Θ of rotation of projectile nucleon")
+        .def_readwrite("phi_rot_b", &cola::EventIniState::phiRotB,
+                      "Polar angle φ of rotation of target nucleon")
+        .def_readwrite("theta_rot_b", &cola::EventIniState::thetaRotB,
+                      "Polar angle Θ of rotation of target nucleon")
+        .def_readwrite("initial_particles", &cola::EventIniState::iniStateParticles,
+                      "Array of all particles just before the event")
+        ;
+
+    py::class_<cola::EventData>(m, "EventData")
+        .def(py::init<>())
+        .def_readwrite("initial_state", &cola::EventData::iniState)
+        .def_readwrite("particles", &cola::EventData::particles);
+}

colapy-0.0.1/src/colapy/__init__.py

@@ -0,0 +1,29 @@
+import typing as tp
+
+
+from ._cola_impl import (
+    __doc__,
+    __version__,
+    LorentzVector,
+    Particle,
+    ParticleClass,
+    EventInitialState,
+    EventData,
+)
+
+
+class AZ(tp.NamedTuple):
+    A: int
+    Z: int
+
+
+__all__ = [
+    '__doc__',
+    '__version__',
+    'AZ',
+    'LorentzVector',
+    'Particle',
+    'ParticleClass',
+    'EventInitialState',
+    'EventData',
+]

colapy-0.0.1/src/colapy/__init__.pyi

@@ -0,0 +1,111 @@
+import enum
+import typing as tp
+
+
+from . import AZ
+
+
+EventParticles: tp.TypeAlias = tp.List[Particle]
+
+
+class LorentzVector:
+    x: float
+    y: float
+    z: float
+    e: float
+    t: float
+
+    @tp.overload
+    def __init__(self, other: 'LorentzVector') -> None: ...
+    
+    @tp.overload
+    def __init__(
+        self,
+        e: float = 0.0,
+        x: float = 0.0,
+        y: float = 0.0,
+        z: float = 0.0,
+    ) -> None: ...
+
+    def __add__(self, other: 'LorentzVector') -> 'LorentzVector': ...
+    def __sub__(self, other: 'LorentzVector') -> 'LorentzVector': ...
+    def __iadd__(self, other: 'LorentzVector') -> 'LorentzVector': ...
+    def __isub__(self, other: 'LorentzVector') -> 'LorentzVector': ...
+
+    def __mul__(self, scalar: float) -> 'LorentzVector': ...
+    def __rmul__(self, scalar: float) -> 'LorentzVector': ...
+    def __truediv__(self, scalar: float) -> 'LorentzVector': ...
+    def __imul__(self, scalar: float) -> 'LorentzVector': ...
+    def __itruediv__(self, scalar: float) -> 'LorentzVector': ...
+
+    def __eq__(self, other: object) -> bool: ...
+
+    def __copy__(self) -> 'LorentzVector': ...
+    def __deepcopy__(self, memo: object) -> 'LorentzVector': ...
+
+    def __repr__(self) -> str: ...
+
+
+class ParticleClass(enum.Enum):
+    PRODUCED: int
+    ELASTIC_A: int
+    ELASTIC_B: int
+    NON_ELASTIC_A: int
+    NON_ELASTIC_B: int
+    SPECTATOR_A: int
+    SPECTATOR_B: int
+
+
+class Particle:
+    position: LorentzVector
+    momentum: LorentzVector
+    pdg_code: int
+    p_class: ParticleClass
+
+    @tp.overload
+    def __init__(self, other: 'Particle') -> None: ...
+
+    @tp.overload
+    def __init__(
+        self,
+        position: LorentzVector = ...,
+        momentum: LorentzVector = ...,
+        pdg_code: int = ...,
+        p_class: ParticleClass = ..., 
+    ) -> None: ...
+    def get_az(self) -> AZ: ...
+
+
+class EventInitialState:
+    pdg_code_a: int
+    pdg_code_b: int
+    pz_a: float
+    pz_b: float
+    energy: float
+    sect_nn: float
+    b: float
+    n_coll: int
+    n_coll_pp: int
+    n_coll_pn: int
+    n_coll_nn: int
+    n_part: int
+    n_part_a: int
+    n_part_b: int
+    phi_rot_a: float
+    theta_rot_a: float
+    phi_rot_b: float
+    theta_rot_b: float
+    initial_particles: EventParticles
+
+    def __init__(self) -> None: ...
+
+
+class EventData:
+    initial_state: EventInitialState
+    particles: EventParticles
+
+    def __init__(self) -> None: ...
+
+
+def pdg_to_az(pdg_code: int) -> AZ: ...
+def az_to_pdg(az: AZ) -> int: ...