Change lejasampler to set init_pivot the way choleskysampler does

Author: jdjakem

pyapprox/surrogates/affine/activelearning.py

@@ -84,24 +84,20 @@ def _generate_new_samples(self, nsamples: int) -> Array:
     def _update_ntrain_samples(self):
         self._ntrain_samples = self._train_samples.shape[1]
 
-    def set_seq_initial_pivots(self, init_seq_pivots: Array):
-        self._init_seq_pivots = init_seq_pivots
-        print(self._init_seq_pivots, self._init_seq_pivots.shape)
-
-    def _get_seq_init_pivots(self) -> Array:
-        if hasattr(self, "_init_seq_pivots"):
-            if (
-                self._init_seq_pivots.shape[0]
-                > self._surrogate.basis().nterms()
-            ):
+    def set_initial_pivots(self, init_pivots: Array):
+        self._init_pivots = init_pivots
+
+    def _get_init_pivots(self) -> Array:
+        if hasattr(self, "_init_pivots"):
+            if self._init_pivots.shape[0] > self._surrogate.basis().nterms():
                 # this conditioned can be removed if code is changed
                 # to allow init pivots to be passed to update pivots
                 # but I do not think this case needs to be supported
                 raise ValueError(
-                    "too many init_seq_pivots given. "
+                    "too many init_pivots given. "
                     "Increase nterms of initial basis"
                 )
-            return self._init_seq_pivots
+            return self._init_pivots
         return None
 
     def _init_factorization(self, nprev_train_samples: int, nsamples: int):
@@ -115,7 +111,7 @@ def _init_factorization(self, nprev_train_samples: int, nsamples: int):
         self._LUFactorizer = PivotedLUFactorizer(mat, bkd=self._bkd)
 
         self._L, self._U = self._LUFactorizer.factorize(
-            nsamples, init_seq_pivots=self._get_seq_init_pivots()
+            nsamples, init_pivots=self._get_init_pivots()
         )
         return self._LUFactorizer.pivots()[nprev_train_samples:nsamples]
 

pyapprox/surrogates/affine/tests/test_activelearning.py

@@ -125,7 +125,6 @@ def fun(samples):
         gen1 = ExpandingMarginGenerator(basis1.nvars(), init_degree, 1.0, 4)
         basis1.set_indices(gen1.get_indices())
 
-        print("####")
         np.random.seed(1)
         bexp1 = AdaptivePolynomialChaosExpansion(
             basis1,
@@ -135,31 +134,24 @@ def fun(samples):
             # max_nsamples=6,
             nqoi=nqoi,
         )
-        import torch
-
-        torch.set_printoptions(linewidth=1000)
         bexp1.build(fun)
 
-        print("$$$$")
         init_degree = 3
         basis2 = OrthonormalPolynomialBasis(bases_1d)
         gen2 = ExpandingMarginGenerator(basis2.nvars(), init_degree, 1.0, 4)
         basis2.set_indices(gen1.get_indices())
         np.random.seed(
             1
-        )  # sue same seed to make sure candidate samples are the same
+        )  # use the same seed to make sure candidate samples are the same
         bexp2 = AdaptivePolynomialChaosExpansion(
             basis2,
             variable,
             gen2,
             max_nsamples=10,
-            # max_nsamples=6,
             nqoi=nqoi,
         )
-        bexp2.sampler().set_seq_initial_pivots(
-            bexp1.sampler()._LUFactorizer.seq_pivots()[
-                : bexp1.sampler().pivots().shape[0]
-            ]
+        bexp2.sampler().set_initial_pivots(
+            bexp1.sampler()._LUFactorizer.pivots()
         )
         bexp2.build(fun)
 

pyapprox/util/linalg.py

@@ -246,11 +246,13 @@ def solve_linear_system(self, rhs: Array) -> Array:
 
 class PivotedLUFactorizer(PivotedFactorizer):
     def _best_pivot(self, it: int) -> int:
-        if (
-            self._init_seq_pivots is not None
-            and it < self._init_seq_pivots.shape[0]
-        ):
-            return self._init_seq_pivots[it]
+        if self._init_pivots is not None and it < len(self._init_pivots):
+            return self._bkd.where(self._pivots == self._init_pivots[it])[0][0]
+        # if (
+        #     self._init_seq_pivots is not None
+        #     and it < self._init_seq_pivots.shape[0]
+        # ):
+        #     return self._init_seq_pivots[it]
         else:
             return (
                 self._bkd.argmax(self._bkd.abs(self._LU_factor[it:, it])) + it
@@ -293,23 +295,24 @@ def _split_lu(self, LU_factor: Array, npivots: int) -> Tuple[Array, Array]:
     def factorize(
         self,
         npivots: int,
-        init_seq_pivots: Array = None,
+        init_pivots: Array = None,
         pivot_weights: Array = None,
     ) -> Tuple[Array, Array]:
-        self._init_seq_pivots = init_seq_pivots
+        self._init_pivots = init_pivots
         self._ncompleted_pivots = 0
         self._LU_factor = self._bkd.copy(self._Amat)
         self._seq_pivots = self._bkd.arange(self._Amat.shape[0])
         return self.update(npivots)
 
     def update(self, npivots: int) -> Tuple[Array, Array]:
         npivots = min(npivots, min(*self._Amat.shape))
-        for it in range(self._ncompleted_pivots, npivots):
+        for it in self._bkd.arange(self._ncompleted_pivots, npivots):
             pivot = self._best_pivot(it)
             # update pivots vector
             self._seq_pivots[it] = pivot
             # apply pivots(swap rows) in L factorization
             swap_rows(self._LU_factor, it, pivot)
+            swap_rows(self._pivots, it, pivot, self._bkd)
             if self._terminate(it):
                 self._termination_flag = 1
                 return self._split_lu(self._LU_factor, it)
@@ -340,6 +343,15 @@ def add_rows(self, new_rows: Array):
                 ),
             )
         )
+        self._pivots = self._bkd.hstack(
+            (
+                self._pivots,
+                self._bkd.arange(
+                    self._Amat.shape[0],
+                    self._Amat.shape[0] + new_rows.shape[0],
+                ),
+            )
+        )
         self._Amat = self._bkd.vstack((self._Amat, new_rows))
         LU_factor_extra = self._bkd.copy(new_rows)
         for it in range(self._ncompleted_pivots):
@@ -394,29 +406,6 @@ def add_columns(self, new_cols: Array):
             )
         self._LU_factor = self._bkd.hstack((self._LU_factor, new_cols))
 
-    def pivots(self) -> Array:
-        """
-        Return the vector that changes the original array to the final
-        permuted array in 1 shot.
-
-        Return
-        ------
-        pivots: Array (ncompleted_pivots,)
-        """
-        return get_final_pivots_from_sequential_pivots(
-            self._seq_pivots, bkd=self._bkd
-        )[: self._ncompleted_pivots]
-
-    def seq_pivots(self) -> Array:
-        """
-        Return the pivot vector obtained by inserting pivot at each iteratation
-
-        Return
-        ------
-        seq_pivots: Array (ncandidates,)
-        """
-        return self._seq_pivots
-
     def _undo_preconditioning(
         self,
         precond_weights: Array,

pyapprox/util/tests/test_linalg.py

@@ -312,7 +312,7 @@ def test_add_columns_to_pivoted_lu_factorization(self):
 
         factorizer.update(3)
         full_factorizer = PivotedLUFactorizer(A, bkd=bkd)
-        full_factorizer.factorize(3, init_seq_pivots=factorizer.seq_pivots())
+        full_factorizer.factorize(3, init_pivots=factorizer.pivots())
         assert bkd.allclose(factorizer._LU_factor, full_factorizer._LU_factor)
 
     def test_add_rows_to_pivoted_lu_factorization(self):
@@ -369,7 +369,7 @@ def test_add_rows_to_pivoted_lu_factorization(self):
         factorizer3.update(npivots + 1)
         full_factorizer = PivotedLUFactorizer(A_reordered, bkd=bkd)
         full_factorizer.factorize(
-            npivots + 1, init_seq_pivots=factorizer3.seq_pivots()
+            npivots + 1, init_pivots=factorizer3.pivots()
         )
         assert bkd.allclose(factorizer3._LU_factor, full_factorizer._LU_factor)
 
@@ -389,9 +389,7 @@ def test_undo_lu_preconditioning(self):
         )
 
         factorizer2 = PivotedLUFactorizer(A, bkd=bkd)
-        L, U = factorizer2.factorize(
-            npivots, init_seq_pivots=factorizer1.seq_pivots()
-        )
+        L, U = factorizer2.factorize(npivots, init_pivots=factorizer1.pivots())
         assert bkd.allclose(factorizer2.pivots(), factorizer1.pivots())
         assert bkd.allclose(L @ U, A[factorizer2.pivots()])
 
@@ -426,9 +424,7 @@ def test_undo_lu_preconditioning(self):
         L_precond, U_precond = factorizer1.factorize(npivots)
 
         factorizer2 = PivotedLUFactorizer(A, bkd=bkd)
-        L, U = factorizer2.factorize(
-            npivots, init_seq_pivots=factorizer1.seq_pivots()
-        )
+        L, U = factorizer2.factorize(npivots, init_pivots=factorizer1.pivots())
 
         L_adjusted, U_adjusted = factorizer1.undo_preconditioning(
             precond_weights, npivots, update_internal_state=True
@@ -451,9 +447,7 @@ def test_update_lu_preconditioning(self):
         )
         new_precond_weights = precond_weights**2
         factorizer2 = PivotedLUFactorizer(new_precond_weights * A, bkd=bkd)
-        L, U = factorizer2.factorize(
-            npivots, init_seq_pivots=factorizer1.seq_pivots()
-        )
+        L, U = factorizer2.factorize(npivots, init_pivots=factorizer1.pivots())
         assert bkd.allclose(factorizer2.pivots(), factorizer1.pivots())
         assert bkd.allclose(
             L @ U, (precond_weights**2 * A)[factorizer2.pivots()]
@@ -475,9 +469,7 @@ def test_update_lu_preconditioning(self):
         L_precond, U_precond = factorizer1.factorize(npivots)
 
         factorizer2 = PivotedLUFactorizer(new_precond_weights * A, bkd=bkd)
-        L, U = factorizer2.factorize(
-            npivots, init_seq_pivots=factorizer1.seq_pivots()
-        )
+        L, U = factorizer2.factorize(npivots, init_pivots=factorizer1.pivots())
 
         L_adjusted, U_adjusted = factorizer1.update_preconditioning(
             precond_weights,