Add slicing for TriangularMap class

.github/workflows/build-external-lib-tests.yml

@@ -59,7 +59,7 @@ jobs:
                 -DCMAKE_CXX_STANDARD=17    \
                 -DCMAKE_INSTALL_PREFIX=$GITHUB_WORKSPACE/KOKKOS_INSTALL/ \
                 ../
-          sudo make -j2 install
+          sudo make -j$(nproc) install
 
       - name: Install Conda Dependencies
         shell: bash -l {0}
@@ -74,7 +74,7 @@ jobs:
 
       - name: Build MParT
         shell: bash -l {0}
-        run: cd $GITHUB_WORKSPACE/mpart/build; make -j2
+        run: cd $GITHUB_WORKSPACE/mpart/build; make -j$(nproc)
 
       - name: Run Tests
         shell: bash -l {0}

MParT/AffineMap.h

@@ -59,6 +59,8 @@ class AffineMap : public ConditionalMapBase<MemorySpace>
     void LogDeterminantInputGradImpl(StridedMatrix<const double, MemorySpace> const& pts,
                                      StridedMatrix<double, MemorySpace>              output) override;
 
+    std::shared_ptr<ConditionalMapBase<MemorySpace>> Slice(int a, int b) override;
+
     /** Computes an LU factorization of the matrix A_ */
     void Factorize();
 

MParT/ComposedMap.h

@@ -102,6 +102,8 @@ class ComposedMap : public ConditionalMapBase<MemorySpace>
     void GradientImpl(StridedMatrix<const double, MemorySpace> const& pts,
                       StridedMatrix<const double, MemorySpace> const& sens,
                       StridedMatrix<double, MemorySpace>              output) override;
+
+    std::shared_ptr<ConditionalMapBase<MemorySpace>> Slice(int a, int b) override;
 private:
 
     unsigned int maxChecks_;

MParT/ConditionalMapBase.h

@@ -149,6 +149,11 @@ namespace mpart {
 
         virtual void LogDeterminantInputGradImpl(StridedMatrix<const double, MemorySpace> const& pts,
                                                  StridedMatrix<double, MemorySpace>              output) = 0;
+
+        virtual std::shared_ptr<ConditionalMapBase<MemorySpace>> Slice(int a, int b) {
+            throw std::runtime_error("Slice not implemented for this map type.");
+        }
+
 #if defined(MPART_HAS_CEREAL)
         template<typename Archive>
         void serialize(Archive& ar) {

MParT/IdentityMap.h

@@ -57,6 +57,10 @@ class IdentityMap : public ConditionalMapBase<MemorySpace>
     void LogDeterminantInputGradImpl(StridedMatrix<const double, MemorySpace> const& pts,
                                      StridedMatrix<double, MemorySpace>              output) override;
 
+    // This only slices the tail of the map due to how IdentityMap is defined.
+    // i.e. this is invariant on identical shifting of a,b.
+    std::shared_ptr<ConditionalMapBase<MemorySpace>> Slice(int a, int b) override;
+
 }; // class IdentityMap
 
 }

MParT/MapFactory.h

@@ -4,6 +4,7 @@
 #include "MParT/MapOptions.h"
 
 #include "MParT/ConditionalMapBase.h"
+#include "MParT/TriangularMap.h"
 #include "MParT/SummarizedMap.h"
 //#include "MParT/DebugMap.h"
 #include "MParT/MultiIndices/FixedMultiIndexSet.h"

MParT/MonotoneComponent.h

@@ -1,6 +1,8 @@
 #ifndef MPART_MONOTONECOMPONENT_H
 #define MPART_MONOTONECOMPONENT_H
 
+#include <memory>
+
 #include "MParT/MultiIndices/FixedMultiIndexSet.h"
 #include "MParT/MultiIndices/MultiIndexSet.h"
 
@@ -37,7 +39,7 @@ T(x_1, x_2, ..., x_D) = f(x_1,x_2,..., x_{D-1}, 0) + \int_0^{x_D}  g\left( \part
 @tparam QuadratureType A class defining the integration scheme used to approximate \f$\int_0^{x_N}  g\left( \frac{\partial f}{\partial x_d}(x_1,x_2,..., x_{N-1}, t) \right) dt\f$.  The type must have a function `Integrate(f,lb,ub)` that accepts a functor `f`, a double lower bound `lb`, a double upper bound `ub`, and returns a double with an estimate of the integral.   The MParT::AdaptiveSimpson and MParT::RecursiveQuadrature classes provide this interface.
 */
 template<class ExpansionType, class PosFuncType, class QuadratureType, typename MemorySpace>
-class MonotoneComponent : public ConditionalMapBase<MemorySpace>
+class MonotoneComponent : public ConditionalMapBase<MemorySpace>, public std::enable_shared_from_this<MonotoneComponent<ExpansionType, PosFuncType, QuadratureType, MemorySpace>>
 {
 
 public:
@@ -1078,6 +1080,11 @@ class MonotoneComponent : public ConditionalMapBase<MemorySpace>
         return 0.5*(xub+xlb);
     }
 
+    std::shared_ptr<ConditionalMapBase<MemorySpace>> Slice(int a, int b) override {
+        assert(a == 0 && b == 1);
+        return MonotoneComponent<ExpansionType, PosFuncType, QuadratureType, MemorySpace>::shared_from_this();
+    }
+
     /** Give access to the underlying FixedMultiIndexSet
      * @return The FixedMultiIndexSet
      */

MParT/TriangularMap.h

@@ -60,6 +60,18 @@ class TriangularMap : public ConditionalMapBase<MemorySpace>{
     void WrapCoeffs(Kokkos::View<double*, mpart::DeviceSpace> coeffs) override;
     #endif
 
+    std::shared_ptr<ConditionalMapBase<MemorySpace>> Slice(int a, int b) override;
+
+    /** @brief Returns a subsection of the map according to index of components
+     * @details This function returns a subsection (or "slice") of a block triangular map by simply just taking a subset of the blocks.
+     *          The returned map will have the same number of inputs as the original map, and this is equivalent to {@link #Slice Slice} when
+     *          the number of outputs of each component is 1.
+     * @param a The first component block to take
+     * @param b The index AFTER the last component block to take
+     * @return A shared pointer to a new ConditionalMapBase object containing the subsection of the map.
+     */
+    virtual std::shared_ptr<ConditionalMapBase<MemorySpace>> BlockSlice(int a, int b);
+
     virtual std::shared_ptr<ConditionalMapBase<MemorySpace>> GetComponent(unsigned int i){ return comps_.at(i);}
 
     /** @brief Computes the log determinant of the Jacobian matrix of this map.

bindings/julia/src/ConditionalMapBase.cpp

@@ -39,6 +39,7 @@ void mpart::binding::ConditionalMapBaseWrapper(jlcxx::Module &mod) {
             map.InverseImpl(JuliaToKokkos(x1), JuliaToKokkos(r), JuliaToKokkos(output));
             return output;
         })
+        .method("Slice", [](ConditionalMapBase<Kokkos::HostSpace>& map, int begin, int end){ return map.Slice(begin-1, end); })
         ;
     jlcxx::stl::apply_stl<ConditionalMapBase<Kokkos::HostSpace>*>(mod);
 }

bindings/julia/src/TriangularMap.cpp

@@ -14,6 +14,6 @@ void mpart::binding::TriangularMapWrapper(jlcxx::Module &mod) {
        .method("InverseInplace", [](TriangularMap<Kokkos::HostSpace> &map, jlcxx::ArrayRef<double,2> x, jlcxx::ArrayRef<double,2> r){
             map.InverseInplace(JuliaToKokkos(x), JuliaToKokkos(r));
        })
-       .method("GetComponent", &TriangularMap<Kokkos::HostSpace>::GetComponent)
+       .method("GetComponent", [](TriangularMap<Kokkos::HostSpace>& map, int i){ return map.GetComponent(i-1); })
     ;
 }

bindings/python/src/ConditionalMapBase.cpp

@@ -23,6 +23,12 @@ void mpart::binding::ConditionalMapBaseWrapper(py::module &m)
         .def("LogDeterminantCoeffGrad", static_cast<Eigen::RowMatrixXd (ConditionalMapBase<MemorySpace>::*)(Eigen::Ref<const Eigen::RowMatrixXd> const&)>(&ConditionalMapBase<MemorySpace>::LogDeterminantCoeffGrad))
         .def("LogDeterminantInputGrad", static_cast<Eigen::RowMatrixXd (ConditionalMapBase<MemorySpace>::*)(Eigen::Ref<const Eigen::RowMatrixXd> const&)>(&ConditionalMapBase<MemorySpace>::LogDeterminantInputGrad))
         .def("GetBaseFunction", &ConditionalMapBase<MemorySpace>::GetBaseFunction)
+        .def("__getitem__", [](ConditionalMapBase<MemorySpace> &m, const py::slice &s) {
+            size_t start = 0, stop = 0, step = 0, slicelength = 0;
+            s.compute(m.outputDim, &start, &stop, &step, &slicelength);
+            if(step != 1) throw std::runtime_error("Map slice step size must be 1");
+            return m.Slice(start, stop);
+        })
         ;
 
 }

bindings/python/src/TriangularMap.cpp

@@ -20,6 +20,8 @@ void mpart::binding::TriangularMapWrapper(py::module &m)
     py::class_<TriangularMap<MemorySpace>, ConditionalMapBase<MemorySpace>, std::shared_ptr<TriangularMap<MemorySpace>>>(m, tName.c_str())
         .def(py::init<std::vector<std::shared_ptr<ConditionalMapBase<MemorySpace>>>, bool>(), py::arg("comps"), py::arg("moveCoeffs") = false)
         .def("GetComponent", &TriangularMap<MemorySpace>::GetComponent)
+        .def("Slice", &TriangularMap<MemorySpace>::Slice)
+        .def("__getitem__", &TriangularMap<MemorySpace>::GetComponent)
         ;
 
 }

src/AffineMap.cpp

@@ -3,13 +3,13 @@
 
 #include "MParT/Utilities/ArrayConversions.h"
 #include "MParT/Initialization.h"
- 
+
 
 using namespace mpart;
 
 template<typename MemorySpace>
-AffineMap<MemorySpace>::AffineMap(StridedVector<double,MemorySpace> b) : ConditionalMapBase<MemorySpace>(b.size(),b.size(),0), 
-                                                                  b_("b",b.layout()), 
+AffineMap<MemorySpace>::AffineMap(StridedVector<double,MemorySpace> b) : ConditionalMapBase<MemorySpace>(b.size(),b.size(),0),
+                                                                  b_("b",b.layout()),
                                                                   logDet_(0.0)
 {
     Kokkos::deep_copy(b_, b);
@@ -30,7 +30,7 @@ AffineMap<MemorySpace>::AffineMap(StridedMatrix<double,MemorySpace> A,
                                   StridedVector<double,MemorySpace> b) : ConditionalMapBase<MemorySpace>(A.extent(1),A.extent(0),0),
                                                                   A_("A", A.layout()),
                                                                   b_("b",b.layout())
-{   
+{
     Kokkos::deep_copy(A_, A);
     Kokkos::deep_copy(b_, b);
     assert(A_.extent(0)<=A_.extent(1));
@@ -91,15 +91,15 @@ void AffineMap<MemorySpace>::InverseImpl(StridedMatrix<const double, MemorySpace
     }else{
         outLeft = output;
     }
-    
+
     // Bias part
     if(b_.size()>0){
 
         Kokkos::MDRangePolicy<Kokkos::Rank<2>, typename MemoryToExecution<MemorySpace>::Space> policy({{0, 0}}, {{numPts, this->outputDim}});
 
         // After this for loop, we will have out = r - b
         Kokkos::parallel_for(policy, KOKKOS_CLASS_LAMBDA(const int& j, const int& i) {
-            outLeft(i,j) = r(i,j) - b_(i); 
+            outLeft(i,j) = r(i,j) - b_(i);
         });
 
     }else{
@@ -115,7 +115,7 @@ void AffineMap<MemorySpace>::InverseImpl(StridedMatrix<const double, MemorySpace
         // If the matrix is rectangular, treat it as the lower part of a block triangular matrix
         // The value of x1 contains the compute inverse for the first block
         if(nrows != ncols){
-            
+
             StridedMatrix<const double, MemorySpace> xsub = Kokkos::subview(x1, std::make_pair(0,ncols-nrows), Kokkos::ALL());
             StridedMatrix<const double, MemorySpace> Asub = Kokkos::subview(A_, Kokkos::ALL(), std::make_pair(0,ncols-nrows));
 
@@ -129,13 +129,13 @@ void AffineMap<MemorySpace>::InverseImpl(StridedMatrix<const double, MemorySpace
         if(copyOut)
             Kokkos::deep_copy(output, outLeft);
     }
-    
+
 }
 
 
 
 template<typename MemorySpace>
-void AffineMap<MemorySpace>::LogDeterminantCoeffGradImpl(StridedMatrix<const double, MemorySpace> const& pts, 
+void AffineMap<MemorySpace>::LogDeterminantCoeffGradImpl(StridedMatrix<const double, MemorySpace> const& pts,
                                                          StridedMatrix<double, MemorySpace>              output)
 {
     return;
@@ -151,7 +151,7 @@ void AffineMap<MemorySpace>::EvaluateImpl(StridedMatrix<const double, MemorySpac
     }else{
         Kokkos::deep_copy(output, pts);
     }
-    
+
     // Bias part
     if(b_.size()>0){
 
@@ -165,22 +165,22 @@ void AffineMap<MemorySpace>::EvaluateImpl(StridedMatrix<const double, MemorySpac
 }
 
 template<typename MemorySpace>
-void AffineMap<MemorySpace>::CoeffGradImpl(StridedMatrix<const double, MemorySpace> const& pts,  
+void AffineMap<MemorySpace>::CoeffGradImpl(StridedMatrix<const double, MemorySpace> const& pts,
                                            StridedMatrix<const double, MemorySpace> const& sens,
                                            StridedMatrix<double, MemorySpace>              output)
 {
     return;
 }
 
-template<typename MemorySpace> 
-void AffineMap<MemorySpace>::LogDeterminantInputGradImpl(StridedMatrix<const double, MemorySpace> const& pts, 
+template<typename MemorySpace>
+void AffineMap<MemorySpace>::LogDeterminantInputGradImpl(StridedMatrix<const double, MemorySpace> const& pts,
                                                          StridedMatrix<double, MemorySpace>              output)
 {
     return;
 }
 
 template<typename MemorySpace>
-void AffineMap<MemorySpace>::GradientImpl(StridedMatrix<const double, MemorySpace> const& pts,  
+void AffineMap<MemorySpace>::GradientImpl(StridedMatrix<const double, MemorySpace> const& pts,
                                           StridedMatrix<const double, MemorySpace> const& sens,
                                           StridedMatrix<double, MemorySpace>              output)
 {
@@ -192,8 +192,16 @@ void AffineMap<MemorySpace>::GradientImpl(StridedMatrix<const double, MemorySpac
     }
 }
 
+template<typename MemorySpace>
+std::shared_ptr<ConditionalMapBase<MemorySpace>> AffineMap<MemorySpace>::Slice(int a, int b) {
+    StridedMatrix<double, MemorySpace> A_new;
+    StridedVector<double, MemorySpace> b_new;
+    if(A_.extent(0) > 0) A_new = Kokkos::subview(A_, std::make_pair(a,b), Kokkos::ALL());
+    if(b_.size() > 0) b_new = Kokkos::subview(b_, std::make_pair(a,b));
+    return std::make_shared<AffineMap<MemorySpace>>(A_new, b_new);
+}
 
 template class mpart::AffineMap<Kokkos::HostSpace>;
 #if defined(MPART_ENABLE_GPU)
     template class mpart::AffineMap<mpart::DeviceSpace>;
-#endif 
+#endif