Refactor MonotoneComponent and revise tests

MeasureTransport · Sep 10, 2023 · e52fb11 · e52fb11
1 parent 85fead6
commit e52fb11
Show file tree

Hide file tree

Showing 2 changed files with 18 additions and 12 deletions.
diff --git a/MParT/MonotoneComponent.h b/MParT/MonotoneComponent.h
@@ -400,13 +400,8 @@ class MonotoneComponent : public ConditionalMapBase<MemorySpace>
 
                 // Compute the inverse
                 Kokkos::View<double*,MemorySpace> workspace(team_member.thread_scratch(1), workspaceSize);
-                auto evaluate_lambda = KOKKOS_LAMBDA(double x){
-                    return EvaluateSingle(workspace.data(), cache.data(), pt, x, coeffs, quad_, expansion_);
-                };
-                // OLD
-                // output(ptInd) = InverseSingleBracket(workspace.data(), cache.data(), pt, ys(ptInd), coeffs, xtol, ytol, quad_, expansion_);
-                // NEW
-                output(ptInd) = RootFinding::InverseSingleBracket<MemorySpace>(ys(ptInd), evaluate_lambda, pt(pt.extent(0)-1), xtol, ytol);
+                auto eval = SingleEvaluator(workspace.data(), cache.data(), pt, coeffs, quad_, expansion_);
+                output(ptInd) = RootFinding::InverseSingleBracket<MemorySpace>(ys(ptInd), eval, pt(pt.extent(0)-1), xtol, ytol);
             }
         };
 
@@ -1130,6 +1125,21 @@ class MonotoneComponent : public ConditionalMapBase<MemorySpace>
     QuadratureType quad_;
     const unsigned int dim_;
     bool useContDeriv_;
+    template<typename PointType, typename CoeffType>
+    struct SingleEvaluator {
+        double* workspace;
+        double* cache;
+        PointType pt;
+        CoeffType coeffs;
+        QuadratureType quad;
+        ExpansionType expansion;
+
+        SingleEvaluator(double* workspace_, double* cache_, PointType pt_, CoeffType coeffs_, QuadratureType quad_, ExpansionType expansion_):
+            workspace(workspace_), cache(cache_), pt(pt_), coeffs(coeffs_), quad(quad_), expansion(expansion_) {};
+        double operator()(double x) {
+            return EvaluateSingle(cache, workspace, pt, x, coeffs, quad, expansion);
+        }
+    };
 };
 
 } // namespace mpart

diff --git a/tests/Test_MonotoneComponent.cpp b/tests/Test_MonotoneComponent.cpp
@@ -353,7 +353,7 @@ TEST_CASE( "Testing monotone component evaluation in 1d", "[MonotoneComponent1d]
 
 TEST_CASE( "Testing bracket-based inversion of monotone component", "[MonotoneBracketInverse]" ) {
 
-    const double testTol = 1e-7;
+    const double testTol = 1e-6;
     unsigned int dim = 1;
 
     // Create points evently space on [lb,ub]
@@ -390,12 +390,8 @@ TEST_CASE( "Testing bracket-based inversion of monotone component", "[MonotoneBr
         comp.EvaluateImpl(evalPts, coeffs, ys);
 
         Kokkos::View<double*, HostSpace> testInverse("Test output", numPts);
-        auto start = std::chrono::high_resolution_clock::now();
         for(int i = 0; i < 100; i++)
             comp.InverseImpl(evalPts, ys, coeffs, testInverse);
-        auto end = std::chrono::high_resolution_clock::now();
-        auto duration = std::chrono::duration_cast<std::chrono::microseconds>(end - start);
-        std::cout << "Inversion time taken: " << static_cast<double>(duration.count())*1e-9 << " microseconds" << std::endl;
 
         for(unsigned int i=0; i<numPts; ++i){
             CHECK(testInverse(i) == Approx(evalPts(0,i)).epsilon(testTol));