Merge branch 'main' into accelerate_RMVE

MeasureTransport · May 10, 2024 · 0dc9511 · 0dc9511
2 parents 8bfd229 + c4e70db
commit 0dc9511
Show file tree

Hide file tree

Showing 14 changed files with 543 additions and 91 deletions.
diff --git a/MParT/BasisEvaluator.h b/MParT/BasisEvaluator.h
@@ -284,51 +284,113 @@ class BasisEvaluator<BasisHomogeneity::OffdiagHomogeneous,
   DiagEvaluatorType diag_;
 };
 
+
+
 /**
  * @brief Basis Evaluator to eval different basis fcns for arbitrary inputs
  *
  * @tparam CommonBasisEvaluatorType Supertype to univariate basis eval types
  */
-template <typename CommonBasisEvaluatorType, typename Rectifier>
+template <typename Rectifier, typename ...Bases>
 class BasisEvaluator<BasisHomogeneity::Heterogeneous,
-                     std::vector<std::shared_ptr<CommonBasisEvaluatorType>>, Rectifier> {
+                     std::tuple<Bases...>, Rectifier> {
   public:
+
+  template<typename I, std::enable_if_t<std::is_integral_v<I>, bool> = true>
   BasisEvaluator(
-      unsigned int dim,
-      std::vector<std::shared_ptr<CommonBasisEvaluatorType>> const &basis1d)
+      std::array<I, sizeof...(Bases)> const &dims,
+      std::tuple<Bases...> basis1d)
       : basis1d_(basis1d) {
-    assert(dim == basis1d.size());  // TODO: Fix
-  }
+        for(unsigned int i = 0; i < sizeof...(Bases); i++){
+            dims_[i] = dims[i];
+        }
+      }
+
   // EvaluateAll(dim, output, max_order, input)
   // dim is zero-based indexing
   KOKKOS_INLINE_FUNCTION void EvaluateAll(unsigned int dim, double *output,
                                           int max_order, double input) const {
-    basis1d_[dim]->EvaluateAll(output, max_order, input);
+    basis1d_Evaluator<FunctionEvalType::EvaluateAll, 0>(dim, output, max_order, input);
   }
   // EvaluateDerivatives(dim, output_eval, output_deriv, max_order, input)
   KOKKOS_INLINE_FUNCTION void EvaluateDerivatives(unsigned int dim, double *output,
                                                   double *output_diff,
                                                   int max_order,
                                                   double input) const {
-    basis1d_[dim]->EvaluateDerivatives(output, output_diff, max_order, input);
+    basis1d_Evaluator<FunctionEvalType::EvaluateDerivatives, 0>(dim, output, output_diff, max_order, input);
   }
   // EvaluateSecondDerivatives(dim, output_eval, output_deriv, max_order, input)
   KOKKOS_INLINE_FUNCTION void EvaluateSecondDerivatives(unsigned int dim, double *output,
                                                         double *output_diff,
                                                         double *output_diff2,
                                                         int max_order,
                                                         double input) const {
-    basis1d_[dim]->EvaluateSecondDerivatives(output, output_diff, output_diff2,
-                                             max_order, input);
+    basis1d_Evaluator<FunctionEvalType::EvaluateSecondDerivatives, 0>(dim, output, output_diff, output_diff2, max_order, input);
+  }
+
+  private:
+  constexpr static unsigned int MaxIndex = std::integral_constant<unsigned int, sizeof...(Bases)>::value;
+
+  enum class FunctionEvalType { EvaluateAll, EvaluateDerivatives, EvaluateSecondDerivatives };
+
+  template<FunctionEvalType evalType, unsigned int idx, typename ...Args>
+  KOKKOS_INLINE_FUNCTION void basis1d_Evaluator(unsigned int dim_eval, Args... args) const {
+    // Ensure that idx is within bounds (prevent overflow)
+    if constexpr(idx >= MaxIndex) {
+      assert(false); // TODO: use ProcAgnosticError
+      return;
+    } else {
+      // Check if idx gives the right basis function to evaluate
+      if(dim_eval >= std::get<idx>(dims_)){
+        basis1d_Evaluator<evalType, idx + 1, Args...>(dim_eval - std::get<idx>(dims_), args...);
+        return;
+      }
+      // Evaluate the basis function
+      if constexpr (evalType == FunctionEvalType::EvaluateAll) {
+        std::get<idx>(basis1d_).EvaluateAll(args...);
+      } else if constexpr (evalType == FunctionEvalType::EvaluateDerivatives) {
+        std::get<idx>(basis1d_).EvaluateDerivatives(args...);
+      } else if constexpr (evalType == FunctionEvalType::EvaluateSecondDerivatives) {
+        std::get<idx>(basis1d_).EvaluateSecondDerivatives(args...);
+      }
+    }
   }
+
 #if defined(MPART_HAS_CEREAL)
   template <typename Archive>
   void serialize(Archive &ar) {
+    ar(dims_);
     ar(basis1d_);
   }
 #endif
-  // NOTE: shared_ptr is necessary to avoid type slicing
-  std::vector<std::shared_ptr<CommonBasisEvaluatorType>> basis1d_;
+
+  /// @brief Number of input dimensions for multivariate basis
+  std::tuple<Bases...> basis1d_;
+  std::array<unsigned int, sizeof...(Bases)> dims_;
 };
+
+
+template<typename It, typename... Args>
+constexpr std::tuple<Args...> CreateHeterogeneousBasisEvaluatorHelper(It ret, Args... args) {
+    return std::make_tuple(args...);
+}
+
+template<typename It, typename I, std::enable_if_t<std::is_integral_v<I>, bool> = true, typename... Args>
+constexpr auto CreateHeterogeneousBasisEvaluatorHelper(It ret, I next, Args... args) {
+    *ret = next;
+    return CreateHeterogeneousBasisEvaluatorHelper(ret + 1, args...);
+}
+
+template<typename Rectifier, typename I, std::enable_if_t<std::is_integral_v<I>,bool> = true, typename... Args, typename std::enable_if_t<sizeof...(Args) % 2 == 1, bool> = true>
+constexpr auto CreateHeterogeneousBasisEvaluator(I idx1, Args... args) {
+    std::array<I, (sizeof...(Args) + 1)/2> arr;
+    arr[0] = idx1;
+    auto basis1d = CreateHeterogeneousBasisEvaluatorHelper(arr.begin() + 1, args...);
+    return BasisEvaluator<BasisHomogeneity::Heterogeneous, decltype(basis1d), Rectifier>(arr, basis1d);
+}
+
+template<typename Rectifier, typename ...Bases>
+using HeterogeneousBasisEvaluator = BasisEvaluator<BasisHomogeneity::Heterogeneous, std::tuple<Bases...>, Rectifier>;
+
 }  // namespace mpart
 #endif // MPART_BASISEVALUATOR_H
diff --git a/MParT/Distributions/GaussianSamplerDensity.h b/MParT/Distributions/GaussianSamplerDensity.h
@@ -54,6 +54,8 @@ class GaussianSamplerDensity: public SampleGenerator<MemorySpace>, public Densit
 
     unsigned int Dim() const override { return dim_; };
 
+    bool IsStandardNormal() const { return idCov_ && mean_.size() == 0; }
+
     protected:
     using GeneratorType = typename SampleGenerator<MemorySpace>::PoolType::generator_type;
     using SampleGenerator<MemorySpace>::rand_pool;

diff --git a/MParT/MapObjective.h b/MParT/MapObjective.h
@@ -178,12 +178,44 @@ class KLObjective: public MapObjective<MemorySpace> {
     double ObjectiveImpl(StridedMatrix<const double, MemorySpace> data, std::shared_ptr<ConditionalMapBase<MemorySpace>> map) const override;
     void CoeffGradImpl(StridedMatrix<const double, MemorySpace> data, StridedVector<double, MemorySpace> grad, std::shared_ptr<ConditionalMapBase<MemorySpace>> map) const override;
     unsigned int MapOutputDim() const override {return density_->Dim();}
-    private:
+
     /**
      * @brief Density \f$\mu\f$ to calculate the KL with respect to (i.e. \f$D(\cdot||\mu)\f$ )
      *
      */
-    std::shared_ptr<DensityBase<MemorySpace>> density_;
+    const std::shared_ptr<DensityBase<MemorySpace>> density_;
+};
+
+template<typename MemorySpace>
+class FastGaussianReverseKLObjective: public MapObjective<MemorySpace> {
+    public:
+    FastGaussianReverseKLObjective(StridedMatrix<const double, MemorySpace> train, unsigned int numCoeffs);
+    FastGaussianReverseKLObjective(StridedMatrix<const double, MemorySpace> train, StridedMatrix<const double, MemorySpace> test, unsigned int numCoeffs);
+
+    double ObjectivePlusCoeffGradImpl(StridedMatrix<const double, MemorySpace> data, StridedVector<double, MemorySpace> grad, std::shared_ptr<ConditionalMapBase<MemorySpace>> map) const override;
+
+    double ObjectiveImpl(StridedMatrix<const double, MemorySpace> data, std::shared_ptr<ConditionalMapBase<MemorySpace>> map) const override;
+
+    void CoeffGradImpl(StridedMatrix<const double, MemorySpace> data, StridedVector<double, MemorySpace> grad, std::shared_ptr<ConditionalMapBase<MemorySpace>> map) const override;
+
+    private:
+    using ExecSpace = typename MemoryToExecution<MemorySpace>::Space;
+    enum class ObjectiveType {Eval = 0, Grad = 1, EvalGrad = 2};
+
+    template<unsigned int Type_idx>
+    void FillSpaces(std::shared_ptr<ConditionalMapBase<MemorySpace>> map, StridedMatrix<const double, MemorySpace> data) const;
+
+    template<unsigned int Type_idx>
+    double CommonEval(StridedMatrix<const double, MemorySpace> data, StridedVector<double, MemorySpace> grad, std::shared_ptr<ConditionalMapBase<MemorySpace>> map) const;
+
+    template<unsigned int Type_idx>
+    void ClearSpaces() const;
+
+    mutable Kokkos::View<double**, MemorySpace> eval_space_;
+    mutable Kokkos::View<double*, MemorySpace> logdet_space_;
+
+    mutable Kokkos::View<double**, MemorySpace> grad_space_;
+    mutable Kokkos::View<double**, MemorySpace> logdet_grad_space_;
 };
 
 namespace ObjectiveFactory {

diff --git a/MParT/TrainMap.h b/MParT/TrainMap.h
@@ -5,6 +5,7 @@
 #include <nlopt.hpp>
 #include <iostream>
 #include "MParT/ConditionalMapBase.h"
+#include "MParT/TriangularMap.h"
 #include "MParT/MapObjective.h"
 
 namespace mpart {

diff --git a/MParT/TriangularMap.h b/MParT/TriangularMap.h
@@ -64,6 +64,7 @@ class TriangularMap : public ConditionalMapBase<MemorySpace>{
     void WrapCoeffs(Kokkos::View<double*, MemorySpace> coeffs) override;
 
     virtual std::shared_ptr<ConditionalMapBase<MemorySpace>> GetComponent(unsigned int i){ return comps_.at(i);}
+    unsigned int NumComponents() const { return comps_.size();}
 
     /** @brief Computes the log determinant of the Jacobian matrix of this map.
 

diff --git a/MParT/Utilities/ArrayConversions.h b/MParT/Utilities/ArrayConversions.h
@@ -148,9 +148,10 @@ namespace mpart{
     template<typename ScalarType, class... ViewTraits>
     std::vector<typename std::remove_const<ScalarType>::type> KokkosToStd(Kokkos::View<ScalarType*,ViewTraits...> const& view)
     {
+        auto view_copy = Kokkos::create_mirror_view_and_copy(Kokkos::HostSpace(), view);
         std::vector<typename std::remove_const<ScalarType>::type> output(view.extent(0));
         for(unsigned int i=0; i<view.extent(0); ++i)
-            output[i] = view(i);
+            output[i] = view_copy(i);
         return output;
     }
 

diff --git a/src/Distributions/GaussianSamplerDensity.cpp b/src/Distributions/GaussianSamplerDensity.cpp
@@ -30,9 +30,7 @@ void GaussianSamplerDensity<MemorySpace>::LogDensityImpl(StridedMatrix<const dou
     Kokkos::View<double**, Kokkos::LayoutLeft, MemorySpace> diff ("diff", M, N);
 
     if(mean_.extent(0) == 0){
-        Kokkos::parallel_for(policy, KOKKOS_CLASS_LAMBDA(const int& j, const int& i) {
-            diff(i,j) = pts(i,j);
-        });
+        Kokkos::deep_copy(diff, pts);
     }
     else {
         Kokkos::parallel_for(policy, KOKKOS_CLASS_LAMBDA(const int& j, const int& i) {
@@ -75,7 +73,17 @@ void GaussianSamplerDensity<MemorySpace>::LogDensityInputGradImpl(StridedMatrix<
     }
 
     if(!idCov_) {
-        covChol_.solveInPlace(output);
+
+        if(output.stride_0()==1){
+            covChol_.solveInPlace(output);
+        }else{
+            StridedMatrix<double,MemorySpace> outLeft;
+            outLeft = Kokkos::View<double**, Kokkos::LayoutLeft, MemorySpace>("OutLeft", output.extent(0), output.extent(1));
+            Kokkos::deep_copy(outLeft, output);
+            covChol_.solveInPlace(outLeft);
+            Kokkos::deep_copy(output, outLeft);
+        }
+
     }
 }