Merge branch 'main' into sigflex

MeasureTransport · May 10, 2024 · 9359ac4 · 9359ac4
2 parents 0520e57 + c4e70db
commit 9359ac4
Show file tree

Hide file tree

Showing 12 changed files with 452 additions and 49 deletions.
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);
+        }
+
     }
 }
 

diff --git a/src/MapObjective.cpp b/src/MapObjective.cpp
@@ -6,6 +6,7 @@ template<typename MemorySpace>
 double MapObjective<MemorySpace>::operator()(unsigned int n, const double* coeffs, double* grad, std::shared_ptr<ConditionalMapBase<MemorySpace>> map) {
 
     Kokkos::View<const double*, MemorySpace> coeffView = ToConstKokkos<double,MemorySpace>(coeffs, n);
+    if (grad == nullptr) return ObjectiveImpl(train_, map);
     StridedVector<double, MemorySpace> gradView = ToKokkos<double,MemorySpace>(grad, n);
     map->SetCoeffs(coeffView);
     return ObjectivePlusCoeffGradImpl(train_, gradView, map);
@@ -89,6 +90,149 @@ double KLObjective<MemorySpace>::ObjectivePlusCoeffGradImpl(StridedMatrix<const
     return sumDensity/N_samps;
 }
 
+template<typename MemorySpace>
+FastGaussianReverseKLObjective<MemorySpace>::FastGaussianReverseKLObjective(StridedMatrix<const double, MemorySpace> train, StridedMatrix<const double, MemorySpace> test, unsigned int numCoeffs) : MapObjective<MemorySpace>(train, test) {
+    unsigned int N = std::max(train.extent(1), test.extent(1));
+    eval_space_ = Kokkos::View<double**, MemorySpace>("eval_space", 1, N);
+    logdet_space_ = Kokkos::View<double*, MemorySpace>("logdet_space", N);
+    grad_space_ = Kokkos::View<double**, MemorySpace>("grad_space", numCoeffs, N);
+    logdet_grad_space_ = Kokkos::View<double**, MemorySpace>("logdet_grad_space", numCoeffs, N);
+}
+
+template<typename MemorySpace>
+FastGaussianReverseKLObjective<MemorySpace>::FastGaussianReverseKLObjective(StridedMatrix<const double, MemorySpace> train, unsigned int numCoeffs) : FastGaussianReverseKLObjective<MemorySpace>(train, Kokkos::View<const double**, MemorySpace>(), numCoeffs) {}
+
+template<typename MemorySpace>
+template<unsigned int Type_idx>
+void FastGaussianReverseKLObjective<MemorySpace>::FillSpaces(std::shared_ptr<ConditionalMapBase<MemorySpace>> map, StridedMatrix<const double, MemorySpace> data) const {
+    constexpr ObjectiveType Type = static_cast<ObjectiveType>(Type_idx);
+    unsigned int N_points = data.extent(1);
+    if(N_points > eval_space_.extent(1)){
+        std::stringstream ss;
+        ss << "FastGaussianReverseKLObjective: Not enough space allocated for evaluation!" <<
+            "Need " << N_points << " points, storing " << eval_space_.extent(1) << " points.";
+        throw std::invalid_argument(ss.str().c_str());
+    }
+    StridedMatrix<double, MemorySpace> eval_space_view = Kokkos::subview(eval_space_, Kokkos::ALL(), Kokkos::make_pair(0u, N_points));
+    map->EvaluateImpl(data, eval_space_view);
+    if constexpr((Type == ObjectiveType::Eval) || (Type == ObjectiveType::EvalGrad)) {
+        StridedVector<double, MemorySpace> logdet_space_view = Kokkos::subview(logdet_space_, Kokkos::make_pair(0u, N_points));
+        map->LogDeterminantImpl(data, logdet_space_view);
+    }
+    if constexpr((Type == ObjectiveType::Grad) || (Type == ObjectiveType::EvalGrad)) {
+        StridedMatrix<double, MemorySpace> grad_space_view = Kokkos::subview(grad_space_, Kokkos::ALL(), Kokkos::make_pair(0u, N_points));
+        StridedMatrix<double, MemorySpace> logdet_grad_space_view = Kokkos::subview(logdet_grad_space_, Kokkos::ALL(), Kokkos::make_pair(0u, N_points));
+        Kokkos::fence();
+        map->CoeffGradImpl(data, eval_space_view, grad_space_view);
+        map->LogDeterminantCoeffGradImpl(data, logdet_grad_space_view);
+    }
+    Kokkos::fence();
+}
+
+template<typename MemorySpace>
+template<unsigned int Type_idx>
+void FastGaussianReverseKLObjective<MemorySpace>::ClearSpaces() const {
+    constexpr ObjectiveType Type = static_cast<ObjectiveType>(Type_idx);
+    if constexpr((Type == ObjectiveType::Eval) || (Type == ObjectiveType::EvalGrad)) {
+        Kokkos::deep_copy(eval_space_, 0.);
+        Kokkos::deep_copy(logdet_space_, 0.);
+    }
+    if constexpr((Type == ObjectiveType::Grad) || (Type == ObjectiveType::EvalGrad)) {
+        Kokkos::deep_copy(grad_space_, 0.);
+        Kokkos::deep_copy(logdet_grad_space_, 0.);
+    }
+}
+
+template<typename EvalSpaceType, typename LogDetSpaceType>
+class EvaluationFunctor {
+    public:
+    EvaluationFunctor(const EvalSpaceType& eval_space_view, const LogDetSpaceType& logdet_space_view): eval_space_view_(eval_space_view), logdet_space_view_(logdet_space_view){}
+    KOKKOS_FUNCTION void operator()(const unsigned int j, double& loss) const {
+        loss += 0.5*eval_space_view_(0, j)*eval_space_view_(0, j) - logdet_space_view_(j);
+    }
+    private:
+    const EvalSpaceType eval_space_view_;
+    const LogDetSpaceType logdet_space_view_;
+};
+
+template<typename KLGradSpaceType, typename GradSpaceType, typename GradLogDetSpaceType>
+class GradientFunctor {
+    using team_handle=typename Kokkos::TeamPolicy<typename GradSpaceType::execution_space>::member_type;
+    public:
+    GradientFunctor(KLGradSpaceType& kl_grad, unsigned int N_points, const GradSpaceType& grad_space_view, const GradLogDetSpaceType& logdet_grad_space_view):
+        kl_grad_(kl_grad), N_points_(N_points), grad_space_view_(grad_space_view), logdet_grad_space_view_(logdet_grad_space_view) {
+    }
+    KOKKOS_FUNCTION void operator()(const team_handle& team) const {
+        int d = team.league_rank();
+        double grad_d_ = 0.;
+        Kokkos::parallel_reduce(Kokkos::TeamThreadRange(team, N_points_), [&](const unsigned int j, double& grad_d_tmp) {
+            grad_d_tmp += grad_space_view_(d, j) - logdet_grad_space_view_(d, j);
+        }, grad_d_);
+        kl_grad_(d) = grad_d_/N_points_;
+    }
+
+    private:
+    mutable KLGradSpaceType kl_grad_;
+    const unsigned int N_points_;
+    const GradSpaceType grad_space_view_;
+    const GradLogDetSpaceType logdet_grad_space_view_;
+};
+
+template<typename MemorySpace>
+template<unsigned int Type_idx>
+double FastGaussianReverseKLObjective<MemorySpace>::CommonEval(StridedMatrix<const double, MemorySpace> points, StridedVector<double, MemorySpace> kl_grad, std::shared_ptr<ConditionalMapBase<MemorySpace>> map) const {
+    if(map->outputDim != 1){
+        throw std::invalid_argument("FastGaussianReverseKLObjective: Map output dimension must be 1! Found dimension " + std::to_string(map->outputDim) + ".");
+    }
+    constexpr ObjectiveType Type = static_cast<ObjectiveType>(Type_idx);
+    unsigned int N_points = points.extent(1);
+    unsigned int numCoeffs = map->numCoeffs;
+    if(Type != ObjectiveType::Eval && numCoeffs != kl_grad.extent(0)){
+        std::stringstream ss;
+        ss << "FastGaussianReverseKLObjective: Gradient vector has incorrect size!"
+           << "Given size " << kl_grad.extent(0) << ", expected size " << numCoeffs
+           << ".";
+        throw std::invalid_argument(ss.str().c_str());
+    }
+    FillSpaces<Type_idx>(map, points);
+
+    double kl_loss = 0.;
+    if constexpr((Type == ObjectiveType::Eval) || (Type == ObjectiveType::EvalGrad)) {
+        Kokkos::RangePolicy<ExecSpace> eval_policy(0, N_points);
+        EvaluationFunctor eval_kernel(eval_space_, logdet_space_);
+        Kokkos::parallel_reduce("FastGaussianReverseKL Evaluate", eval_policy, eval_kernel, kl_loss);
+        kl_loss /= N_points;
+    }
+    if constexpr((Type == ObjectiveType::Grad) || (Type == ObjectiveType::EvalGrad)) {
+        Kokkos::TeamPolicy<ExecSpace> grad_policy(numCoeffs, Kokkos::AUTO());
+        GradientFunctor gradient_functor(kl_grad, N_points, grad_space_, logdet_grad_space_);
+
+        Kokkos::parallel_for("FastGaussianReverseKL Gradient", grad_policy, gradient_functor);
+        Kokkos::fence();
+    }
+    ClearSpaces<Type_idx>();
+    return kl_loss;
+}
+
+template<typename MemorySpace>
+double FastGaussianReverseKLObjective<MemorySpace>::ObjectivePlusCoeffGradImpl(StridedMatrix<const double, MemorySpace> data, StridedVector<double, MemorySpace> grad, std::shared_ptr<ConditionalMapBase<MemorySpace>> map) const {
+    constexpr unsigned int EvalGrad_idx = static_cast<unsigned int>(ObjectiveType::EvalGrad);
+    return CommonEval<EvalGrad_idx>(data, grad, map);
+}
+
+template<typename MemorySpace>
+double FastGaussianReverseKLObjective<MemorySpace>::ObjectiveImpl(StridedMatrix<const double, MemorySpace> data, std::shared_ptr<ConditionalMapBase<MemorySpace>> map) const {
+    Kokkos::View<double*, MemorySpace> grad_holder;
+    constexpr unsigned int Eval_idx = static_cast<unsigned int>(ObjectiveType::Eval);
+    return CommonEval<Eval_idx>(data, grad_holder, map);
+}
+
+template<typename MemorySpace>
+void FastGaussianReverseKLObjective<MemorySpace>::CoeffGradImpl(StridedMatrix<const double, MemorySpace> data, StridedVector<double, MemorySpace> grad, std::shared_ptr<ConditionalMapBase<MemorySpace>> map) const {
+    constexpr unsigned int Grad_idx = static_cast<unsigned int>(ObjectiveType::Grad);
+    CommonEval<Grad_idx>(data, grad, map);
+}
+
 template<typename MemorySpace>
 double KLObjective<MemorySpace>::ObjectiveImpl(StridedMatrix<const double, MemorySpace> data, std::shared_ptr<ConditionalMapBase<MemorySpace>> map) const {
     unsigned int N_samps = data.extent(1);
@@ -107,7 +251,7 @@ void KLObjective<MemorySpace>::CoeffGradImpl(StridedMatrix<const double, MemoryS
     using ExecSpace = typename MemoryToExecution<MemorySpace>::Space;
     unsigned int N_samps = data.extent(1);
     unsigned int grad_dim = grad.extent(0);
-    PullbackDensity<MemorySpace> pullback {map, density_};
+    PullbackDensity<MemorySpace> pullback {map, density_}; 
     StridedMatrix<double, MemorySpace> densityGradX = pullback.LogDensityCoeffGrad(data);
 
     double scale = -1.0/((double) N_samps);
@@ -131,11 +275,13 @@ void KLObjective<MemorySpace>::CoeffGradImpl(StridedMatrix<const double, MemoryS
 // Explicit template instantiation
 template class mpart::MapObjective<Kokkos::HostSpace>;
 template class mpart::KLObjective<Kokkos::HostSpace>;
+template class mpart::FastGaussianReverseKLObjective<Kokkos::HostSpace>;
 template std::shared_ptr<MapObjective<Kokkos::HostSpace>> mpart::ObjectiveFactory::CreateGaussianKLObjective<Kokkos::HostSpace>(StridedMatrix<const double, Kokkos::HostSpace>, unsigned int);
 template std::shared_ptr<MapObjective<Kokkos::HostSpace>> mpart::ObjectiveFactory::CreateGaussianKLObjective<Kokkos::HostSpace>(StridedMatrix<const double, Kokkos::HostSpace>, StridedMatrix<const double, Kokkos::HostSpace>, unsigned int);
 #if defined(MPART_ENABLE_GPU)
     template class mpart::MapObjective<DeviceSpace>;
     template class mpart::KLObjective<DeviceSpace>;
+    template class mpart::FastGaussianReverseKLObjective<DeviceSpace>;
     template std::shared_ptr<MapObjective<DeviceSpace>> mpart::ObjectiveFactory::CreateGaussianKLObjective<DeviceSpace>(StridedMatrix<const double, DeviceSpace>, unsigned int);
     template std::shared_ptr<MapObjective<DeviceSpace>> mpart::ObjectiveFactory::CreateGaussianKLObjective<DeviceSpace>(StridedMatrix<const double, DeviceSpace>, StridedMatrix<const double, DeviceSpace>, unsigned int);
 #endif