merging with main

.docker/environment.yml

@@ -12,3 +12,4 @@ dependencies:
   - gcc
   - gxx
   - make
+  - dill

.github/environment.yml

@@ -14,4 +14,4 @@ dependencies:
   - nlopt >= 2.7
   - pytorch
   - cxx-compiler
-
+  - dill

.github/workflows/build-bindings.yml

@@ -3,6 +3,7 @@ name: binding-tests
 on:
   push:
     branches:
+      - release
       - main
   pull_request: {}
 
@@ -47,7 +48,7 @@ jobs:
 
       - name: Setup Julia
         run: |
-          julia -e "using Pkg; Pkg.add([\"CxxWrap\",\"TestReports\"])"
+          julia -e "using Pkg; Pkg.add([Pkg.PackageSpec(;name=\"CxxWrap\",version=v\"0.14.2\"),Pkg.PackageSpec(;name=\"TestReports\")])"
           export GITHUB_JULIA_PATH=$(julia -e "println(DEPOT_PATH[1])")
           echo -n $'[bee5971c-294f-5168-9fcd-9fb3c811d495]\nMParT = \"' >> $GITHUB_JULIA_PATH/artifacts/Overrides.toml
           echo -n $GITHUB_WORKSPACE                                     >> $GITHUB_JULIA_PATH/artifacts/Overrides.toml

.github/workflows/build-doc.yml

@@ -4,6 +4,7 @@ on:
   push:
     branches:
       - main
+      - release
 
 jobs:
   build-docs:

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

@@ -4,6 +4,7 @@ on:
   push:
     branches:
       - main
+      - release
   pull_request: {}
 
 env:

.github/workflows/build-push-docker.yml

@@ -4,6 +4,7 @@ on:
   push:
     branches:
       - main
+      - release
 
 jobs:
   docker:

.github/workflows/build-tests.yml

@@ -4,6 +4,7 @@ on:
   push:
     branches:
       - main
+      - release
   pull_request: {}
 
 jobs:

MParT/ComposedMap.h

@@ -50,12 +50,12 @@ class ComposedMap : public ConditionalMapBase<MemorySpace>
                   start at index \f$\sum_{j=1}^{k-1} C_j\f$.
     */
     using ConditionalMapBase<MemorySpace>::SetCoeffs;
-    void SetCoeffs(Kokkos::View<const double*, Kokkos::HostSpace> coeffs) override;
-    void WrapCoeffs(Kokkos::View<double*, Kokkos::HostSpace> coeffs) override;
+    using ConditionalMapBase<MemorySpace>::WrapCoeffs;
+    void SetCoeffs(Kokkos::View<const double*, MemorySpace> coeffs) override;
+    void WrapCoeffs(Kokkos::View<double*, MemorySpace> coeffs) override;
 
     #if defined(MPART_ENABLE_GPU)
-    void SetCoeffs(Kokkos::View<const double*, Kokkos::DefaultExecutionSpace::memory_space> coeffs) override;
-    void WrapCoeffs(Kokkos::View<double*, mpart::DeviceSpace> coeffs) override;
+    void SetCoeffs(Kokkos::View<const double*, std::conditional_t<std::is_same_v<MemorySpace, Kokkos::HostSpace>, mpart::DeviceSpace, Kokkos::HostSpace>> coeffs) override;
     #endif
 
     virtual std::shared_ptr<ConditionalMapBase<MemorySpace>> GetComponent(unsigned int i){ return maps_.at(i);}

MParT/Distributions/DensityBase.h

@@ -47,8 +47,10 @@ class DensityBase {
      * @param X The points where we want to evaluate the log density.
      * @return Matrix \f$A\in\mathbb{R}^{m\times n}\f$ containing the log density at each point.
      */
-    template<typename AnyMemorySpace>
-    StridedVector<double, AnyMemorySpace> LogDensity(StridedMatrix<const double, AnyMemorySpace> const &X);
+    StridedVector<double, MemorySpace> LogDensity(StridedMatrix<const double, MemorySpace> const &X);
+    #if defined(MPART_ENABLE_GPU)
+    StridedVector<double, std::conditional_t< std::is_same_v<MemorySpace, Kokkos::HostSpace> , DeviceSpace, Kokkos::HostSpace> > LogDensity(StridedMatrix<const double, std::conditional_t< std::is_same_v<MemorySpace, Kokkos::HostSpace> , DeviceSpace, Kokkos::HostSpace > > const &X);
+    #endif
 
     /**
      * @brief Computes the log density at the given points.
@@ -66,8 +68,10 @@ class DensityBase {
      * @param X The points where we want to evaluate the gradient log density.
      * @return A matrix \f$A\in\mathbb{R}^{m\times n}\f$ containing the gradient of the log density at each point.
      */
-    template<typename AnyMemorySpace>
-    StridedMatrix<double, AnyMemorySpace> LogDensityInputGrad(StridedMatrix<const double, AnyMemorySpace> const &X);
+    StridedMatrix<double, MemorySpace> LogDensityInputGrad(StridedMatrix<const double, MemorySpace> const &X);
+    #if defined(MPART_ENABLE_GPU)
+    StridedMatrix<double, std::conditional_t< std::is_same_v<MemorySpace, Kokkos::HostSpace> , DeviceSpace, Kokkos::HostSpace> > LogDensityInputGrad(StridedMatrix<const double, std::conditional_t< std::is_same_v<MemorySpace, Kokkos::HostSpace> , DeviceSpace, Kokkos::HostSpace > > const &X);
+    #endif
 
     /**
      * @brief Returns the input dimension of the density
@@ -81,4 +85,4 @@ class DensityBase {
 };
 }
 
-#endif // MPART_DensityBase_H
+#endif // MPART_DensityBase_H

MParT/Distributions/PushforwardDensity.h

@@ -46,7 +46,8 @@ class PushforwardDensity: public DensityBase<MemorySpace> {
         StridedMatrix<double, MemorySpace> mappedPts = map_->Inverse(prefix_null, pts);
         density_->LogDensityImpl(mappedPts, output);
         StridedVector<double, MemorySpace> logJacobian = map_->LogDeterminant(mappedPts);
-        Kokkos::parallel_for("Subtract logJac", output.extent(0), KOKKOS_LAMBDA(const unsigned int i){
+	Kokkos::RangePolicy<typename MemoryToExecution<MemorySpace>::Space> policy{0lu, output.extent(0)};
+        Kokkos::parallel_for("Subtract logJac", policy, KOKKOS_LAMBDA(const unsigned int i){
             output(i) -= logJacobian(i);
         });
     };
@@ -71,4 +72,4 @@ class PushforwardDensity: public DensityBase<MemorySpace> {
 
 } // namespace mpart
 
-#endif //MPART_PUSHFORWARDDENSITY_H
+#endif //MPART_PUSHFORWARDDENSITY_H

MParT/MapFactory.h

@@ -140,25 +140,27 @@ namespace mpart{
             unsigned int inputDim, unsigned int totalOrder, StridedVector<const double, MemorySpace> centers,
             MapOptions opts);
 
-        template<typename MemorySpace, std::enable_if_t<std::is_same_v<MemorySpace, Kokkos::HostSpace>, bool> = true>
-        std::shared_ptr<ConditionalMapBase<Kokkos::HostSpace>> CreateSigmoidComponent(
+        template<typename MemorySpace>
+        std::shared_ptr<ConditionalMapBase<MemorySpace>> CreateSigmoidComponent(
             unsigned int inputDim, unsigned int totalOrder, Eigen::Ref<const Eigen::RowVectorXd> centers,
             MapOptions opts) {
             StridedVector<const double, Kokkos::HostSpace> centersVec = ConstVecToKokkos<double, Kokkos::HostSpace>(centers);
-            return CreateSigmoidComponent<Kokkos::HostSpace>(inputDim, totalOrder, centersVec, opts);
+            Kokkos::View<const double*, MemorySpace> centers_d = Kokkos::create_mirror_view_and_copy(MemorySpace(), centersVec);
+            return CreateSigmoidComponent<MemorySpace>(inputDim, totalOrder, centers_d, opts);
         }
 
         template<typename MemorySpace>
         std::shared_ptr<ConditionalMapBase<MemorySpace>> CreateSigmoidComponent(
             FixedMultiIndexSet<MemorySpace> mset_offdiag, FixedMultiIndexSet<MemorySpace> mset_diag,
             StridedVector<const double, MemorySpace> centers, MapOptions opts);
 
-        template<typename MemorySpace, std::enable_if_t<std::is_same_v<MemorySpace, Kokkos::HostSpace>, bool> = true>
+        template<typename MemorySpace>
         std::shared_ptr<ConditionalMapBase<MemorySpace>> CreateSigmoidComponent(
             FixedMultiIndexSet<MemorySpace> mset_offdiag, FixedMultiIndexSet<MemorySpace> mset_diag,
             Eigen::Ref<const Eigen::RowVectorXd> centers, MapOptions opts) {
             StridedVector<const double, Kokkos::HostSpace> centersVec = ConstVecToKokkos<double, Kokkos::HostSpace>(centers);
-            return CreateSigmoidComponent<Kokkos::HostSpace>(mset_offdiag, mset_diag, centersVec, opts);
+            Kokkos::View<const double*, MemorySpace> centers_d = Kokkos::create_mirror_view_and_copy(MemorySpace(), centersVec);
+            return CreateSigmoidComponent<MemorySpace>(mset_offdiag, mset_diag, centers_d, opts);
         }
 
         template<typename MemorySpace>
@@ -180,13 +182,14 @@ namespace mpart{
             return CreateSigmoidTriangular<MemorySpace>(inputDim, outputDim, totalOrder, centersVecs, opts);
         }
 
-        template<typename MemorySpace, std::enable_if_t<std::is_same_v<MemorySpace, Kokkos::HostSpace>, bool> = true>
-        std::shared_ptr<ConditionalMapBase<Kokkos::HostSpace>> CreateSigmoidTriangular(
+        template<typename MemorySpace>
+        std::shared_ptr<ConditionalMapBase<MemorySpace>> CreateSigmoidTriangular(
             unsigned int inputDim, unsigned int outputDim, unsigned int totalOrder,
             Eigen::Ref<const Eigen::RowMatrixXd> const& centers, MapOptions opts
         ) {
             StridedMatrix<const double, Kokkos::HostSpace> centersMat = ConstRowMatToKokkos<double,Kokkos::HostSpace>(centers);
-            return CreateSigmoidTriangular<Kokkos::HostSpace>(inputDim, outputDim, totalOrder, centersMat, opts);
+            Kokkos::View<const double**, MemorySpace> centers_d = Kokkos::create_mirror_view_and_copy(MemorySpace(), centersMat);
+            return CreateSigmoidTriangular<MemorySpace>(inputDim, outputDim, totalOrder, centers_d, opts);
         }
 
         /** This struct is used to map the options to functions that can create a map component with types corresponding

MParT/MonotoneComponent.h

@@ -141,7 +141,7 @@ class MonotoneComponent : public ConditionalMapBase<MemorySpace>
             "sens: (" << sensRows << "," << sensCols << "), expected: " << this->outputDim << ", " << ptsCols << "), "
             << "pts: (" << ptsRows << "," << ptsCols << "), expected: (" << this->inputDim << "," << ptsCols << "), "
             << "output: (" << outputRows << "," << outputCols << "), expected: (" << expectedOutputRows << "," << ptsCols << ")";
-            ProcAgnosticError<MemorySpace,std::invalid_argument>::error(ss.str().c_str());
+            ProcAgnosticError<std::invalid_argument>(ss.str().c_str());
         }
     }
 
@@ -262,7 +262,7 @@ class MonotoneComponent : public ConditionalMapBase<MemorySpace>
             std::stringstream ss;
             ss << "EvaluateImpl: output has incorrect number of columns. "
             << "Expected: " << pts.extent(1) << ", got " << output.extent(0);
-            ProcAgnosticError<MemorySpace,std::invalid_argument>::error(ss.str().c_str());
+            ProcAgnosticError<std::invalid_argument>(ss.str().c_str());
         }
 
         // Ask the expansion how much memory it would like for its one-point cache
@@ -663,7 +663,7 @@ class MonotoneComponent : public ConditionalMapBase<MemorySpace>
                << "jacobian: (" << jacRows << "," << jacCols << "), expected: (" << expectJacRows << "," << expectJacCols << "), ";
             if(expectEvalRows > 0)
                 ss << "evaluations: (" << evalRows << "), expected: (" << expectEvalRows << ")";
-            ProcAgnosticError<MemorySpace,std::invalid_argument>::error(ss.str().c_str());
+            ProcAgnosticError<std::invalid_argument>(ss.str().c_str());
         }
     }
 
@@ -1082,9 +1082,9 @@ class MonotoneComponent : public ConditionalMapBase<MemorySpace>
         ExpansionType expansion;
         double nugget;
 
-        SingleEvaluator(double* workspace_, double* cache_, PointType pt_, CoeffType coeffs_, QuadratureType quad_, ExpansionType expansion_, double nugget_):
+        KOKKOS_FUNCTION SingleEvaluator(double* workspace_, double* cache_, PointType pt_, CoeffType coeffs_, QuadratureType quad_, ExpansionType expansion_, double nugget_):
             workspace(workspace_), cache(cache_), pt(pt_), coeffs(coeffs_), quad(quad_), expansion(expansion_), nugget(nugget_) {};
-        double operator()(double x) {
+        KOKKOS_INLINE_FUNCTION double operator()(double x) {
             return EvaluateSingle(cache, workspace, pt, x, coeffs, quad, expansion, nugget);
         }
     };

MParT/MonotoneIntegrand.h

@@ -183,7 +183,7 @@ class MonotoneIntegrand{
         // Check for infs or nans
         if(std::isinf(gf)){
             if(failOnNaN) {
-                ProcAgnosticError<MemorySpace,std::runtime_error>::error("MonotoneIntegrand: nan was encountered in value of g(df(...)). Use MonotoneIntegrand::setFailOnNaN for enabling NaN propagation.");
+                ProcAgnosticError<std::runtime_error>("MonotoneIntegrand: nan was encountered in value of g(df(...)). Use MonotoneIntegrand::setFailOnNaN for enabling NaN propagation.");
             } else {
                 printf("\nERROR: In MonotoneIntegrand, value of g(df(...)) is inf.  The value of df(...) is %0.4f, and the value of f(df(...)) is %0.4f.\n\n", df, gf);
             }

MParT/MultivariateExpansionWorker.h

@@ -339,7 +339,6 @@ class MultivariateExpansionWorker
     {
         const unsigned int numTerms = multiSet_.Size();
         double df = 0.0;
-        unsigned int posInd;
 
         for(int wrt=-1; wrt<int(dim_); ++wrt){
             if(wrt>=0){
@@ -538,7 +537,6 @@ class MultivariateExpansionWorker
     KOKKOS_FUNCTION double GetTermValMixedCoeffDeriv(unsigned int termInd, int wrt, const double* cache, unsigned int posIndex) const {
         // Compute the value of this term in the expansion
         double termVal = 1.0;
-        bool hasDeriv = false;
         if(multiSet_.nzStarts(termInd)==multiSet_.nzStarts(termInd+1)) return 0.; // Value is zero if constant in last dimension
         unsigned int end_idx = multiSet_.nzStarts(termInd+1)-1; // last index in loop
         if(multiSet_.nzDims(end_idx)!=dim_-1) return 0.; // Value is zero if constant in last dimension

MParT/ParameterizedFunctionBase.h

@@ -51,23 +51,21 @@ namespace mpart {
             @param coeffs A view containing the coefficients to copy.
         */
        virtual void SetCoeffs(Kokkos::View<const double*, MemorySpace> coeffs);
+               void SetCoeffs(Kokkos::View<      double*, MemorySpace> coeffs);
 
-       #if defined(MPART_ENABLE_GPU)
-       void SetCoeffs(Kokkos::View<const double*, std::conditional_t<std::is_same_v<MemorySpace,Kokkos::HostSpace>,mpart::DeviceSpace,Kokkos::HostSpace>> coeffs);
-       #endif
 
         /** @brief Wrap the internal coefficient view around another view.
             @details Performs a shallow copy of the input coefficients to the internally stored coefficients.
             If values in the view passed to this function are changed, the values will also change in the
             internally stored view.
             @param coeffs A view containing the coefficients we want to wrap.
         */
-        virtual void WrapCoeffs(Kokkos::View<double*, Kokkos::HostSpace> coeffs);
+        virtual void WrapCoeffs(Kokkos::View<double*, MemorySpace> coeffs);
 
         #if defined(MPART_ENABLE_GPU)
-        virtual void SetCoeffs(Kokkos::View<const double*, mpart::DeviceSpace> coeffs);
-        virtual void WrapCoeffs(Kokkos::View<double*, mpart::DeviceSpace> coeffs);
-        #endif
+        virtual void SetCoeffs(Kokkos::View<const double*, std::conditional_t<std::is_same_v<Kokkos::HostSpace,MemorySpace>, mpart::DeviceSpace, Kokkos::HostSpace>> coeffs);
+	            void SetCoeffs(Kokkos::View<      double*, std::conditional_t<std::is_same_v<Kokkos::HostSpace,MemorySpace>, mpart::DeviceSpace, Kokkos::HostSpace>> coeffs);
+	    #endif
 
         /** SetCoeffs function with conversion from Eigen to Kokkos vector types.*/
         virtual void SetCoeffs(Eigen::Ref<Eigen::VectorXd> coeffs);
@@ -249,4 +247,4 @@ namespace mpart {
     }; // class ParameterizedFunctionBase
 }
 
-#endif
+#endif

MParT/Quadrature.h

@@ -9,6 +9,8 @@
 
 #include <Eigen/Core>
 
+#include "MParT/Utilities/KokkosSpaceMappings.h"
+
 #if defined(MPART_HAS_CEREAL)
 #include <cereal/access.hpp>
 #include <cereal/types/base_class.hpp>
@@ -353,8 +355,8 @@ struct GetRuleFunctor{
 template<typename MemorySpace>
 inline ClenshawCurtisQuadrature<MemorySpace>::ClenshawCurtisQuadrature(unsigned int numPts, unsigned int maxDim) : QuadratureBase<MemorySpace>(maxDim,maxDim),  pts_("Points", numPts), wts_("Weights", numPts), numPts_(numPts)
 {
-    // TODO: Add parallel for loop here with one thread to make sure rule is filled in the correct space
-    Kokkos::parallel_for(1, GetRuleFunctor<MemorySpace>(numPts, pts_.data(), wts_.data()));
+    Kokkos::RangePolicy<typename MemoryToExecution<MemorySpace>::Space> policy {0,1};
+    Kokkos::parallel_for(policy, GetRuleFunctor<MemorySpace>(numPts, pts_.data(), wts_.data()));
 };
 
 template<>
@@ -367,8 +369,8 @@ inline ClenshawCurtisQuadrature<Kokkos::HostSpace>::ClenshawCurtisQuadrature(uns
 template<typename MemorySpace>
 inline ClenshawCurtisQuadrature<MemorySpace>::ClenshawCurtisQuadrature(unsigned int numPts, unsigned int maxDim, double* workspace) : QuadratureBase<MemorySpace>(maxDim,maxDim,workspace),  pts_("Points", numPts), wts_("Weights", numPts), numPts_(numPts)
 {
-    // TODO: Add parallel for loop here with one thread to make sure rule is filled in the correct space
-    Kokkos::parallel_for(1, GetRuleFunctor<MemorySpace>(numPts, pts_.data(), wts_.data()));
+    Kokkos::RangePolicy<typename MemoryToExecution<MemorySpace>::Space> policy {0,1};
+    Kokkos::parallel_for(policy, GetRuleFunctor<MemorySpace>(numPts, pts_.data(), wts_.data()));
 };
 
 template<>
@@ -1020,9 +1022,9 @@ inline AdaptiveClenshawCurtis<MemorySpace>::AdaptiveClenshawCurtis(unsigned int
                                                             fineWts_("Coarse Pts", std::pow(2,level+1)+1)
 {
     assert(std::pow(2,level)+1 >=3);
-
-    Kokkos::parallel_for(1, GetRuleFunctor<MemorySpace>(std::pow(2,level)+1, coarsePts_.data(), coarseWts_.data()));
-    Kokkos::parallel_for(1, GetRuleFunctor<MemorySpace>(std::pow(2,level+1)+1, finePts_.data(), fineWts_.data()));
+    Kokkos::RangePolicy<typename MemoryToExecution<MemorySpace>::Space> policy {0,1};
+    Kokkos::parallel_for(policy, GetRuleFunctor<MemorySpace>(std::pow(2,level)+1, coarsePts_.data(), coarseWts_.data()));
+    Kokkos::parallel_for(policy, GetRuleFunctor<MemorySpace>(std::pow(2,level+1)+1, finePts_.data(), fineWts_.data()));
 };
 
 template<>
@@ -1060,9 +1062,9 @@ inline AdaptiveClenshawCurtis<MemorySpace>::AdaptiveClenshawCurtis(unsigned int
                                                             fineWts_("Coarse Pts", std::pow(2,level+1)+1)
 {
     assert(std::pow(2,level)+1 >=3);
-
-    Kokkos::parallel_for(1, GetRuleFunctor<MemorySpace>(std::pow(2,level)+1, coarsePts_.data(), coarseWts_.data()));
-    Kokkos::parallel_for(1, GetRuleFunctor<MemorySpace>(std::pow(2,level+1)+1, finePts_.data(), fineWts_.data()));
+    Kokkos::RangePolicy<typename MemoryToExecution<MemorySpace>::Space> policy {0, 1};
+    Kokkos::parallel_for(policy, GetRuleFunctor<MemorySpace>(std::pow(2,level)+1, coarsePts_.data(), coarseWts_.data()));
+    Kokkos::parallel_for(policy, GetRuleFunctor<MemorySpace>(std::pow(2,level+1)+1, finePts_.data(), fineWts_.data()));
 };
 
 template<>

MParT/RectifiedMultivariateExpansion.h

@@ -215,7 +215,7 @@ namespace mpart{
             unsigned int cacheSize = worker.CacheSize();
 
             // Take logdet of diagonal expansion
-            auto functor = KOKKOS_LAMBDA (typename Kokkos::TeamPolicy<ExecutionSpace>::member_type team_member) {
+            auto functor = KOKKOS_CLASS_LAMBDA (typename Kokkos::TeamPolicy<ExecutionSpace>::member_type team_member) {
 
                 unsigned int ptInd = team_member.league_rank () * team_member.team_size () + team_member.team_rank ();
 
@@ -313,7 +313,7 @@ namespace mpart{
             unsigned int cacheSize = worker.CacheSize();
 
             // Take logdet of diagonal expansion
-            auto functor = KOKKOS_LAMBDA (typename Kokkos::TeamPolicy<ExecutionSpace>::member_type team_member) {
+            auto functor = KOKKOS_CLASS_LAMBDA (typename Kokkos::TeamPolicy<ExecutionSpace>::member_type team_member) {
 
                 unsigned int ptInd = team_member.league_rank () * team_member.team_size () + team_member.team_rank ();
 
@@ -362,7 +362,7 @@ namespace mpart{
             unsigned int cacheSize = worker.CacheSize();
 
             // Take logdet of diagonal expansion
-            auto functor = KOKKOS_LAMBDA (typename Kokkos::TeamPolicy<ExecutionSpace>::member_type team_member) {
+            auto functor = KOKKOS_CLASS_LAMBDA (typename Kokkos::TeamPolicy<ExecutionSpace>::member_type team_member) {
 
                 unsigned int ptInd = team_member.league_rank () * team_member.team_size () + team_member.team_rank ();
 
@@ -413,7 +413,7 @@ namespace mpart{
 
             SingleWorkerEvaluator(double* cache_, PointType pt_, CoeffType coeffs_, Worker_T worker_):
                 cache(cache_), pt(pt_), coeffs(coeffs_), worker(worker_) {}
-            double operator()(double x) {
+            KOKKOS_INLINE_FUNCTION double operator()(double x) {
                 worker.FillCache2(cache, pt, x, DerivativeFlags::None);
                 return worker.Evaluate(cache, coeffs);
             }