Move nonzero diagonal functionality to a limiter

MeasureTransport · Feb 1, 2024 · e815a31 · e815a31
1 parent aac4825
commit e815a31
Show file tree

Hide file tree

Showing 4 changed files with 30 additions and 38 deletions.
diff --git a/MParT/MultiIndices/FixedMultiIndexSet.h b/MParT/MultiIndices/FixedMultiIndexSet.h
@@ -59,9 +59,6 @@ class FixedMultiIndexSet
     // Returns the indices of multiindices with nonzero elements in dimension length
     std::vector<unsigned int> NonzeroDiagonalEntries() const;
 
-    // Creates a new FixedMultiIndexSet with all entries from this except those with zero in last element
-    FixedMultiIndexSet<MemorySpace> RemoveZeroDiagonalEntries() const;
-
     // Returns the linear index of a given multiindex.  Returns -1 if not found.
     int MultiToIndex(std::vector<unsigned int> const& multi) const;
 

diff --git a/MParT/MultiIndices/MultiIndexLimiter.h b/MParT/MultiIndices/MultiIndexLimiter.h
@@ -50,6 +50,29 @@ namespace MultiIndexLimiter{
 
     };
 
+    /** @class NonzeroDiagTotalOrderLimiter
+        @brief Same as TotalOrder, except without any term that has nonzero diagonal entries
+        @details This limter only allows terms that satisfy
+        \f$\|\mathbf{j}\|_1\leq p_U\f$, where \f$\mathbf{j}\f$
+        is the multiindex, and \f$p_U\f$ is a nonnegative integer passed to the
+        constructor of this class.
+    */
+    class NonzeroDiagTotalOrderLimiter{
+
+    public:
+
+        NonzeroDiagTotalOrderLimiter(unsigned int totalOrderIn) : totalOrder(totalOrderIn){};
+
+        bool operator()(MultiIndex const& multi){
+          unsigned int sum = multi.Sum();
+          return (sum <= totalOrder) && (multi.HasNonzeroEnd());
+        };
+
+    private:
+        const unsigned int totalOrder;
+
+    };
+
 
     /** @class Dimension
     @ingroup MultiIndices

diff --git a/src/MultiIndices/FixedMultiIndexSet.cpp b/src/MultiIndices/FixedMultiIndexSet.cpp
@@ -340,35 +340,6 @@ MultiIndexSet FixedMultiIndexSet<MemorySpace>::Unfix() const
     return output;
 }
 
-template<typename MemorySpace>
-FixedMultiIndexSet<MemorySpace> FixedMultiIndexSet<MemorySpace>::RemoveZeroDiagonalEntries() const
-{
-    std::vector<unsigned int> diagEntries = NonzeroDiagonalEntries();
-    Kokkos::View<unsigned int*, MemorySpace> diagEntryView = VecToKokkos<unsigned int, MemorySpace>(diagEntries);
-    Kokkos::View<unsigned int*, MemorySpace> newStarts("newStarts", diagEntries.size()+1);
-    unsigned int numNz = 0;
-    Kokkos::parallel_scan(diagEntries.size(), KOKKOS_LAMBDA(const int i, unsigned int& lnumNz, const bool final){
-        unsigned int localNz = nzStarts(diagEntryView(i)+1) - nzStarts(diagEntryView(i));
-        if(final){
-            newStarts(i) = lnumNz;
-            if(i == diagEntryView.extent(0)-1) newStarts(i+1) = lnumNz + localNz;
-        }
-        lnumNz += localNz;
-    }, numNz);
-    Kokkos::View<unsigned int*, MemorySpace> newDims("newDims", numNz);
-    Kokkos::View<unsigned int*, MemorySpace> newOrders("newOrders", numNz);
-    Kokkos::parallel_for(diagEntries.size(), KOKKOS_CLASS_LAMBDA(const int i){
-        unsigned int numNz = nzStarts(diagEntryView(i)+1) - nzStarts(diagEntryView(i));
-        for(unsigned int j = 0; j < numNz; j++){
-            unsigned int oldIdx = nzStarts(diagEntryView(i)) + j;
-            unsigned int newIdx = newStarts(i) + j;
-            newDims(newIdx) = nzDims(oldIdx);
-            newOrders(newIdx) = nzOrders(oldIdx);
-        }
-    });
-    return FixedMultiIndexSet<MemorySpace>(dim, newStarts, newDims, newOrders);
-}
-
 template<typename MemorySpace>
 void FixedMultiIndexSet<MemorySpace>::Print() const
 {

diff --git a/tests/MultiIndices/Test_MultiIndexSet.cpp b/tests/MultiIndices/Test_MultiIndexSet.cpp
@@ -28,12 +28,6 @@ TEST_CASE( "Testing the FixedMultiIndexSet class", "[FixedMultiIndexSet]" ) {
     std::vector<unsigned int> diagonal_idxs_ref = multiSet_reconstructed.NonzeroDiagonalEntries();
     std::vector<unsigned int> diagonal_idxs = multiSet_fixed.NonzeroDiagonalEntries();
     REQUIRE(diagonal_idxs_ref == diagonal_idxs);
-    FixedMultiIndexSet<Kokkos::HostSpace> multiSet_allDiag = multiSet_fixed.RemoveZeroDiagonalEntries();
-    std::vector<unsigned int> diagonal_idxs_allDiag = multiSet_allDiag.NonzeroDiagonalEntries();
-    REQUIRE(diagonal_idxs_allDiag.size() == multiSet_allDiag.Size());
-    for(int i = 0; i < diagonal_idxs_allDiag.size(); i++) {
-        REQUIRE(diagonal_idxs_allDiag[i] == i);
-    }
 }
 
 TEST_CASE( "Testing dimension sorting in the FixedMultiIndexSet class", "[FixedMultiIndexSetSorting]" ) {
@@ -505,6 +499,13 @@ TEST_CASE("Testing the MultiIndexSet class", "[MultiIndexSet]" ) {
             expected_size += multi.HasNonzeroEnd();
         }
         REQUIRE( expected_size == inds.size() );
+        MultiIndexSet full_set = MultiIndexSet::CreateTotalOrder(dim, maxOrder, MultiIndexLimiter::NonzeroDiagTotalOrderLimiter(maxOrder));
+        inds = full_set.NonzeroDiagonalEntries();
+        REQUIRE( inds.size() == full_set.Size() );
+        for(int i = 0; i < full_set.Size(); ++i) {
+            REQUIRE(inds[i] == i);
+            REQUIRE( full_set.at(i).HasNonzeroEnd() );
+        }
     }
 }