Feature python interface

CMakeLists.txt

@@ -3,6 +3,16 @@ set(CMAKE_CXX_STANDARD 17)
 
 project(nuTens)
 
+# Changes default install path to be a subdirectory of the build dir.
+# Can set build dir at configure time with -DCMAKE_INSTALL_PREFIX=/install/path
+if(CMAKE_INSTALL_PREFIX STREQUAL "" OR CMAKE_INSTALL_PREFIX STREQUAL
+  "/usr/local")
+  set(CMAKE_INSTALL_PREFIX "${CMAKE_BINARY_DIR}")
+elseif(NOT DEFINED CMAKE_INSTALL_PREFIX)
+  set(CMAKE_INSTALL_PREFIX "${CMAKE_BINARY_DIR}")
+endif()
+
+
 # Need to add some special compile flags to check the code test coverage 
 OPTION(NT_TEST_COVERAGE "produce code coverage reports when running tests" OFF)
 IF(NT_TEST_COVERAGE)
@@ -46,6 +56,19 @@ ELSE()
     message("Won't benchmark")
 ENDIF()
 
+# If user wants to enable python interface we need to include pybind
+OPTION(NT_ENABLE_PYTHON "enable python interface" OFF)
+IF(NT_ENABLE_PYTHON)
+    message("Enabling python")
+    CPMAddPackage(
+        GITHUB_REPOSITORY "pybind/pybind11"
+        VERSION 2.13.5 
+    )
+
+ELSE()
+    message("Won't enable python interface")
+ENDIF()
+
 
 ## check build times
 ## have this optional as it's not supported on all CMake platforms
@@ -66,6 +89,13 @@ IF(NT_ENABLE_BENCHMARKING)
     add_subdirectory(benchmarks)
 ENDIF()
 
+IF(NT_ENABLE_PYTHON)
+    set_property( TARGET tensor PROPERTY POSITION_INDEPENDENT_CODE ON )
+    set_property( TARGET propagator PROPERTY POSITION_INDEPENDENT_CODE ON )
+    set_property( TARGET spdlog PROPERTY POSITION_INDEPENDENT_CODE ON )
+    set_property( TARGET logging PROPERTY POSITION_INDEPENDENT_CODE ON )
+    add_subdirectory(python)
+ENDIF()
 
 # Print out a handy message to more easily see the config options
 message( STATUS "The following variables have been used to configure the build: " )

README.md

@@ -77,9 +77,9 @@ nuTens uses [Googles benchmark library](https://github.com/google/benchmark) to
 - [x] Add suite of benchmarking tests
 - [x] Integrate benchmarks into CI ( maybe use [hyperfine](https://github.com/sharkdp/hyperfine) and [bencher](https://bencher.dev/) for this? )
 - [ ] Add proper unit tests
-- [ ] Expand CI to include more platforms
+- [x] Expand CI to include more platforms
 - [ ] Add support for modules (see [PyTorch doc](https://pytorch.org/cppdocs/api/classtorch_1_1nn_1_1_module.html))
 - [ ] Propagation in variable matter density
 - [ ] Add support for Tensorflow backend
-- [ ] Add python interface 
+- [x] Add python interface 
 

nuTens/CMakeLists.txt

@@ -5,6 +5,7 @@
 add_library(logging logging.hpp)
 target_link_libraries(logging spdlog::spdlog)
 set_target_properties(logging PROPERTIES LINKER_LANGUAGE CXX)
+target_include_directories(logging PUBLIC "${SPDLOG_INCLUDE_DIRS}")
 
 set( NT_LOG_LEVEL "INFO" CACHE STRING "the level of detail to log to the console" )
 

nuTens/propagator/const-density-solver.cpp

@@ -10,8 +10,8 @@ void ConstDensityMatterSolver::calculateEigenvalues(const Tensor &energies, Tens
         for (int j = 0; j < nGenerations; j++)
         {
             hamiltonian.setValue({"...", i, j},
-                                 Tensor::div(diagMassMatrix.getValue({i, j}), energies.getValue({"...", 0})) -
-                                     electronOuter.getValue({i, j}));
+                                 Tensor::div(diagMassMatrix.getValues({i, j}), energies.getValues({"...", 0})) -
+                                     electronOuter.getValues({i, j}));
         }
     }
 

nuTens/propagator/const-density-solver.hpp

@@ -51,8 +51,9 @@
 
         // construct the outer product of the electron neutrino row of the PMNS
         // matrix used to construct the hamiltonian
-        electronOuter = Tensor::scale(Tensor::outer(PMNS.getValue({0, 0, "..."}), PMNS.getValue({0, 0, "..."}).conj()),
-                                      Constants::Groot2 * density);
+        electronOuter =
+            Tensor::scale(Tensor::outer(PMNS.getValues({0, 0, "..."}), PMNS.getValues({0, 0, "..."}).conj()),
+                          Constants::Groot2 * density);
     };
 
     /// @brief Set new mass eigenvalues for this solver
@@ -64,7 +65,7 @@
 
         masses = newMasses;
 
-        Tensor m = masses.getValue({0, "..."});
+        Tensor m = masses.getValues({0, "..."});
         Tensor diag = Tensor::scale(Tensor::mul(m, m), 0.5);
 
         // construct the diagonal mass^2 matrix used in the hamiltonian
@@ -79,7 +80,7 @@
     /// shape should look like {Nbatches, 1, 1}.
     /// @param[out] eigenvectors The returned eigenvectors
     /// @param[out] eigenvalues The corresponding eigenvalues
     void calculateEigenvalues(const Tensor &energies, Tensor &eigenvectors, Tensor &eigenvalues);
 
   private:
     Tensor PMNS;

nuTens/propagator/propagator.cpp

@@ -44,7 +44,7 @@ Tensor Propagator::_calculateProbs(const Tensor &energies, const Tensor &massesS
     {
         for (int j = 0; j < _nGenerations; j++)
         {
-            weightMatrix.setValue({"...", i, j}, weightVector.getValue({"...", j}));
+            weightMatrix.setValue({"...", i, j}, weightVector.getValues({"...", j}));
         }
     }
     weightMatrix.requiresGrad(true);

nuTens/propagator/propagator.hpp

@@ -40,7 +40,7 @@ class Propagator
 
     /// @brief Set a matter solver to use to deal with matter effects
     /// @param newSolver A derivative of BaseMatterSolver
-    inline void setMatterSolver(std::unique_ptr<BaseMatterSolver> &newSolver)
+    inline void setMatterSolver(std::shared_ptr<BaseMatterSolver> &newSolver)
     {
         NT_PROFILE();
         _matterSolver = std::move(newSolver);
@@ -113,5 +113,5 @@ class Propagator
     int _nGenerations;
     float _baseline;
 
-    std::unique_ptr<BaseMatterSolver> _matterSolver;
+    std::shared_ptr<BaseMatterSolver> _matterSolver;
 };

nuTens/tensors/CMakeLists.txt

@@ -9,6 +9,8 @@ IF(NT_USE_PCH)
     target_link_libraries(tensor PUBLIC nuTens-pch)
 
 ELSE()
+    target_link_libraries(tensor PUBLIC logging)
+
     IF(TORCH_FOUND)
 
         target_link_libraries(tensor PUBLIC "${TORCH_LIBRARIES}")

nuTens/tensors/dtypes.hpp

@@ -1,5 +1,9 @@
 #pragma once
 
+#if USE_PYTORCH
+#include <torch/torch.h>
+#endif
+
 /*!
  * @file dtypes.hpp
  * @brief Defines various datatypes used in the project
@@ -25,4 +29,23 @@ enum deviceType
     kGPU
 };
 
+#if USE_PYTORCH
+/// map between the data types used in nuTens and those used by pytorch
+const static std::map<scalarType, c10::ScalarType> scalarTypeMap = {{kFloat, torch::kFloat},
+                                                                    {kDouble, torch::kDouble},
+                                                                    {kComplexFloat, torch::kComplexFloat},
+                                                                    {kComplexDouble, torch::kComplexDouble}};
+
+/// inverse map between the data types used in nuTens and those used by pytorch
+const static std::map<c10::ScalarType, scalarType> invScalarTypeMap = {{torch::kFloat, kFloat},
+                                                                       {torch::kDouble, kDouble},
+                                                                       {torch::kComplexFloat, kComplexFloat},
+                                                                       {torch::kComplexDouble, kComplexDouble}};
+
+// map between the device types used in nuTens and those used by pytorch
+const static std::map<deviceType, c10::DeviceType> deviceTypeMap = {{kCPU, torch::kCPU}, {kGPU, torch::kCUDA}};
+
+// inverse map between the device types used in nuTens and those used by pytorch
+const static std::map<c10::DeviceType, deviceType> invDeviceTypeMap = {{torch::kCPU, kCPU}, {torch::kCUDA, kGPU}};
+#endif
 } // namespace NTdtypes

nuTens/tensors/tensor.hpp

@@ -6,6 +6,7 @@
 #include <iostream>
 #include <map>
 #include <nuTens/instrumentation.hpp>
+#include <nuTens/logging.hpp>
 #include <nuTens/tensors/dtypes.hpp>
 #include <variant>
 #include <vector>
@@ -42,14 +43,19 @@
      */
 
   public:
+    /// Holds the possible "index" types, this allows us to pass integers OR strings as index values which allows us to
+    /// do some basic slicing of tensors similar to python
     using indexType = std::variant<int, std::string>;
 
+    /// Container that holds all allowed types that can be returned by a tensor
+    using variantType = std::variant<int, float, double, std::complex<float>, std::complex<double>>;
+
     /// @name Constructors
     /// Use these methods to construct tensors
     /// @{
 
     /// @brief Default constructor with no initialisation
     Tensor(){};
 
     /// @brief Construct a 1-d array with specified values
     /// @arg values The values to include in the tensor
@@ -233,13 +239,37 @@
     /// @brief Get elementwise phases of a complex tensor
     [[nodiscard]] Tensor angle() const;
 
-    /// @brief Get the result of summing this tensor over some dimension
+    /// @brief Get the cumulative sum over some dimension
     /// @param dim The dimension to sum over
     [[nodiscard]] Tensor cumsum(int dim) const;
 
     /// @brief Get the result of summing this tensor over all dimensions
     [[nodiscard]] Tensor sum() const;
 
+    /// @brief Get the result of summing this tensor over all dimensions
+    /// @param dims The dimensions to sum over
+    [[nodiscard]] Tensor sum(const std::vector<long int> &dims) const;
+
+    /// @brief Get the cumulative sum over some dimension
+    /// @param dim The dimension to sum over
+    static inline Tensor cumsum(const Tensor &t, int dim)
+    {
+        return t.cumsum(dim);
+    }
+
+    /// @brief Get the result of summing this tensor over all dimensions
+    static inline Tensor sum(const Tensor &t)
+    {
+        return t.sum();
+    }
+
+    /// @brief Get the result of summing this tensor over all dimensions
+    /// @param dims The dimensions to sum over
+    static inline Tensor sum(const Tensor &t, const std::vector<long int> &dims)
+    {
+        return t.sum(dims);
+    }
+
     /// @name Gradients
     /// @{
 
@@ -280,18 +310,24 @@
     /// @arg indices The indices of the value to set
     /// @arg value The value to set it to
     void setValue(const Tensor &indices, const Tensor &value);
-    void setValue(const std::vector<std::variant<int, std::string>> &indices, const Tensor &value);
+    void setValue(const std::vector<indexType> &indices, const Tensor &value);
     void setValue(const std::vector<int> &indices, float value);
     void setValue(const std::vector<int> &indices, std::complex<float> value);
 
     /// @brief Get the value at a certain entry in the tensor
     /// @param indices The index of the entry to get
-    [[nodiscard]] Tensor getValue(const std::vector<std::variant<int, std::string>> &indices) const;
+    [[nodiscard]] Tensor getValues(const std::vector<indexType> &indices) const;
+
+    /// @brief Get the value at a certain entry in the tensor as an std::variant
+    /// @details This mainly exists so we can get the values of a tensor in python as pybind11 DOES NOT like templated
+    /// functions If using the c++ interface it is probably easier, faster and safer to use the templated getValue()
+    /// function.
+    [[nodiscard]] variantType getVariantValue(const std::vector<int> &indices) const;
 
     /// @brief Get the number of dimensions in the tensor
     [[nodiscard]] size_t getNdim() const;
 
-    /// @brief Get the batch dimension size of the tensor
+    /// @brief Get the size of the batch dimension of the tensor
     [[nodiscard]] int getBatchDim() const;
 
     /// @brief Get the shape of the tensor
@@ -302,6 +338,8 @@
 
   private:
     bool _hasBatchDim = false;
+    NTdtypes::scalarType _dType;
+    NTdtypes::deviceType _device;
 
     // ###################################################
     // ########## Tensor library specific stuff ##########
@@ -312,37 +350,84 @@
   public:
     /// @brief Get the value at a particular index of the tensor
     /// @arg indices The indices of the value to set
-    template <typename T> inline T getValue(const std::vector<int> &indices)
+    template <typename T> inline T getValue(const std::vector<int> &indices) const
     {
         NT_PROFILE();
 
-        std::vector<at::indexing::TensorIndex> indicesVec;
-        indicesVec.reserve(indices.size());
-        for (const int &i : indices)
-        {
-            indicesVec.push_back(at::indexing::TensorIndex(i));
-        }
-
-        return _tensor.index(indicesVec).item<T>();
+        return _tensor.index(convertIndices(indices)).item<T>();
     }
 
     /// Get the value of a size 0 tensor (scalar)
-    template <typename T> inline T getValue()
+    template <typename T> inline T getValue() const
     {
-
         NT_PROFILE();
 
         return _tensor.item<T>();
     }
 
+    // return the underlying tensor with type determined by the backend library
     [[nodiscard]] inline const torch::Tensor &getTensor() const
     {
-
         NT_PROFILE();
 
         return _tensor;
     }
 
+  private:
+    /// Set the underlying tensor, setting the relevant information like _dtype and _device
+    inline void setTensor(torch::Tensor tensor)
+    {
+        NT_PROFILE();
+
+        _tensor = std::move(tensor);
+        _dType = NTdtypes::invScalarTypeMap.at(tensor.scalar_type());
+        _device = NTdtypes::invDeviceTypeMap.at(tensor.device().type());
+    }
+
+    /// Utility function to convert from a vector of ints to a vector of a10 tensor indices, which is needed for
+    /// accessing values of a tensor.
+    [[nodiscard]] inline std::vector<at::indexing::TensorIndex> convertIndices(const std::vector<int> &indices) const
+    {
+        NT_PROFILE();
+
+        std::vector<at::indexing::TensorIndex> indicesVec;
+        indicesVec.reserve(indices.size());
+        for (const int &i : indices)
+        {
+            indicesVec.push_back(at::indexing::TensorIndex(i));
+        }
+
+        return indicesVec;
+    }
+
+    /// Utility function to convert from a vector of ints to a vector of a10 tensor indices, which is needed for
+    /// accessing values of a tensor.
+    [[nodiscard]] inline std::vector<at::indexing::TensorIndex> convertIndices(
+        const std::vector<Tensor::indexType> &indices) const
+    {
+        NT_PROFILE();
+
+        std::vector<at::indexing::TensorIndex> indicesVec;
+        for (const Tensor::indexType &i : indices)
+        {
+            if (const int *index = std::get_if<int>(&i))
+            {
+                indicesVec.push_back(at::indexing::TensorIndex(*index));
+            }
+            else if (const std::string *index = std::get_if<std::string>(&i))
+            {
+                indicesVec.push_back(at::indexing::TensorIndex((*index).c_str()));
+            }
+            else
+            {
+                assert(false && "ERROR: Unsupported index type");
+                throw;
+            }
+        }
+
+        return indicesVec;
+    }
+
   private:
     torch::Tensor _tensor;
 #endif