add Halko sketcher

Author: estebanag

pyproject.toml

@@ -28,6 +28,7 @@ dependencies = [
     "numpy==2.0.0",
     "scipy==1.13.1",
     "seaborn==0.13.2",
+    "scikit-learn==1.5.0",
 ]
 description = "Quantum-inspired algorithms"
 keywords = []

src/quantum_inspired_algorithms/estimator.py

@@ -11,6 +11,7 @@
 from quantum_inspired_algorithms.quantum_inspired import sample_from_b
 from quantum_inspired_algorithms.quantum_inspired import sample_from_x
 from quantum_inspired_algorithms.sketching import FKV
+from quantum_inspired_algorithms.sketching import Halko
 
 
 class EstimatorError(Exception):
@@ -32,6 +33,7 @@ def __init__(
         n_samples: int,
         random_state: np.random.RandomState,
         sigma_threshold: float = 1e-15,
+        sketcher_name: str = "fkv",
         func: Optional[Callable[[float], float]] = None,
     ) -> None:
         """Init QILinearEstimator.
@@ -46,6 +48,7 @@ def __init__(
             random_state: random state.
             sigma_threshold: the argument `rank` is recomputed in case it is higher
                              the number of singular values below this threhold.
+            sketcher_name: name of sketching method: "fkv" or "halko".
             func: function to transform singular values when estimating lambda coefficients.
                   This can be used for Tikhonov regularization purposes.
         """
@@ -55,6 +58,7 @@ def __init__(
         self.n_samples = n_samples
         self.random_state = random_state
         self.sigma_threshold = sigma_threshold
+        self.sketcher_name = sketcher_name
         self.func = func
 
     def fit(
@@ -84,24 +88,37 @@ def fit(
         # 1. Generate length-square probability distributions to sample from matrix `A`
         logging.info("1. Generate length-square probability distributions to sample from matrix `A`")
         (
-            self.A_ls_prob_rows_,
-            self.A_ls_prob_columns_2d_,
+            A_ls_prob_rows,
+            A_ls_prob_columns_2d,
+            A_ls_prob_columns,
             _,
-            _,
-            self.A_frobenius_,
+            A_frobenius,
         ) = compute_ls_probs(A)
 
         # 2. Build matrix `C`
         logging.info("2. Build matrix `C`")
-        self.sketcher_ = FKV(
-            A,
-            self.r,
-            self.c,
-            self.A_ls_prob_rows_,
-            self.A_ls_prob_columns_2d_,
-            self.A_frobenius_,
-            self.random_state,
-        )
+        if self.sketcher_name == "fkv":
+            self.sketcher_ = FKV(
+                A,
+                self.r,
+                self.c,
+                A_ls_prob_rows,
+                A_ls_prob_columns_2d,
+                A_frobenius,
+                self.random_state,
+            )
+        elif self.sketcher_name == "halko":
+            self.sketcher_ = Halko(
+                A,
+                self.r,
+                self.c,
+                A_ls_prob_rows,
+                A_ls_prob_columns,
+                self.random_state,
+            )
+        else:
+            raise ValueError('`sketcher_name` should be either "fkv" or "halko"')
+
         C = self.sketcher_.right_project(self.sketcher_.left_project(A))
 
         # 3. Compute the SVD of `C`
@@ -135,9 +152,9 @@ def func_(arg: float) -> float:
             self.w_left_,
             self.sigma_,
             self.sketcher_,
-            self.A_ls_prob_rows_,
-            self.A_ls_prob_columns_2d_,
-            self.A_frobenius_,
+            A_ls_prob_rows,
+            A_ls_prob_columns_2d,
+            A_frobenius,
             self.random_state,
             func,
         )
@@ -147,9 +164,6 @@ def func_(arg: float) -> float:
     def _check_is_fitted(self):
         """Check if the `fit` method has been called."""
         for attribute_name in [
-            "A_ls_prob_rows_",
-            "A_ls_prob_columns_2d_",
-            "A_frobenius_",
             "sketcher_",
             "w_left_",
             "sigma_",

src/quantum_inspired_algorithms/sketching.py

@@ -3,6 +3,7 @@
 import numpy as np
 from numpy import linalg as la
 from numpy.typing import NDArray
+from sklearn.utils.extmath import randomized_range_finder
 
 
 class Sketcher(ABC):
@@ -46,7 +47,7 @@ def __init__(
         frobenius: NDArray[np.float64],
         rng: np.random.RandomState,
     ) -> None:
-        """Init QILinearEstimator.
+        """Init FKV.
 
         Note: LS stands for length-square.
 
@@ -130,3 +131,89 @@ def sample_row_idx(self, rng: np.random.RandomState) -> int:
         sample_i = rng.choice(self._n_rows, 1, p=self._C_ls_prob_rows[:, sample_j])[0]
 
         return sample_i
+
+
+class Halko(Sketcher):
+    """Halko sketching."""
+
+    def __init__(
+        self,
+        A: NDArray[np.float64],
+        r: int,
+        c: int,
+        ls_prob_rows: NDArray[np.float64],
+        ls_prob_columns: NDArray[np.float64],
+        rng: np.random.RandomState,
+    ) -> None:
+        """Init Halko.
+
+        Note: LS stands for length-square.
+
+        Args:
+            A: coefficient matrix.
+            r: number of rows for left projection matrix.
+            c: number of columns for right projection matrix.
+            ls_prob_rows: row LS probability distribution of `A`.
+            ls_prob_columns: column LS probability distribution of `A`.
+            rng: random state.
+        """
+        self._Q_left = Halko._get_low_dimensional_projector(A, axis=0, n_components=r, random_state=rng)
+        self._Q_right = Halko._get_low_dimensional_projector(A, axis=1, n_components=c, random_state=rng)
+        self._ls_prob_rows = ls_prob_rows
+        self._ls_prob_columns = ls_prob_columns
+
+    @classmethod
+    def _get_low_dimensional_projector(
+        cls,
+        M: NDArray[np.float64],
+        axis: int,
+        n_components: int,
+        random_state: np.random.RandomState,
+    ) -> NDArray[np.float128]:
+        """Find random matrix to reduce dimensionality of axis of `M`."""
+        n_oversamples = 10
+        n_random = n_components + n_oversamples
+        n_iter = 7 if n_components < 0.1 * min(M.shape) else 4
+        if axis == 1:
+            M = M.T
+        Q = np.asarray(
+            randomized_range_finder(
+                M,
+                size=n_random,
+                n_iter=n_iter,
+                power_iteration_normalizer="auto",
+                random_state=random_state,
+            )
+        )
+        if axis == 0:
+            Q = Q.T
+
+        return np.asarray(Q)
+
+    def left_project(self, M: NDArray[np.float64]) -> NDArray[np.float64]:
+        """Define left projector."""
+        return self._Q_left @ M
+
+    def right_project(self, M: NDArray[np.float64]) -> NDArray[np.float64]:
+        """Define right projector."""
+        return M @ self._Q_right
+
+    def set_up_column_sampler(self, A: NDArray[np.float64]) -> None:
+        """No setup required."""
+
+    def set_up_row_sampler(self, A: NDArray[np.float64]) -> None:
+        """No setup required."""
+
+    def sample_column_idx(self, rng: np.random.RandomState) -> int:
+        """Sample a column index."""
+        n_cols = self._ls_prob_columns.size
+        sample_j = rng.choice(n_cols, 1, p=self._ls_prob_columns)[0]
+
+        return sample_j
+
+    def sample_row_idx(self, rng: np.random.RandomState) -> int:
+        """Sample a row index."""
+        n_rows = self._ls_prob_rows.size
+        sample_i = rng.choice(n_rows, 1, p=self._ls_prob_rows)[0]
+
+        return sample_i

tests/test_estimator.py

@@ -141,7 +141,8 @@ def func(arg: float) -> float:
     )
 
 
-def test_finding_largest_entries_b_underdetermined():
+@pytest.mark.parametrize("sketcher_name,n_matches_expected", [("fkv", 48), ("halko", 49)])
+def test_finding_largest_entries_b_underdetermined(sketcher_name: str, n_matches_expected: int):
     """Test quantum-inspired least squares."""
     # Load data
     A, b, _ = _load_data(underdetermined=True)
@@ -154,7 +155,7 @@ def test_finding_largest_entries_b_underdetermined():
     n_samples = 100
     n_entries_b = 1000
     rng = np.random.RandomState(111)
-    qi = QILinearEstimator(r, c, rank, n_samples, rng)
+    qi = QILinearEstimator(r, c, rank, n_samples, rng, sketcher_name=sketcher_name)
     qi = qi.fit(A, b)
     sampled_indices, sampled_b = qi.predict_b(A, n_entries_b)
 
@@ -169,14 +170,14 @@ def test_finding_largest_entries_b_underdetermined():
     n_matches = plot_solution(
         b,
         b_idx,
-        "test_finding_largest_entries_b_underdetermined",
+        f"test_finding_largest_entries_b_underdetermined_{sketcher_name}",
         expected_solution=b[unique_sampled_indices],
         solution=unique_sampled_b,
         expected_counts=n_entries_b * np.abs(b / norm(b))[unique_sampled_indices] ** 2,
         counts=np.squeeze(np.round(df_counts.values)),
     )
 
-    assert n_matches == 48
+    assert n_matches == n_matches_expected
 
 
 def test_finding_largest_entries_b():

tests/test_sketching.py

@@ -3,10 +3,11 @@
 from numpy.typing import NDArray
 from quantum_inspired_algorithms.quantum_inspired import compute_ls_probs
 from quantum_inspired_algorithms.sketching import FKV
+from quantum_inspired_algorithms.sketching import Halko
 
 
 def _get_FKV_sketcher(A: NDArray[np.float64], r: int, c: int) -> FKV:
-    """Load dummy FKV sketcher."""
+    """Load dummy sketcher."""
     A_ls_prob_rows, A_ls_prob_columns_2d, _, _, A_frobenius = compute_ls_probs(A)
     r = 30
     c = 40
@@ -22,8 +23,24 @@ def _get_FKV_sketcher(A: NDArray[np.float64], r: int, c: int) -> FKV:
     )
 
 
+def _get_Halko_sketcher(A: NDArray[np.float64], r: int, c: int) -> Halko:
+    """Load dummy sketcher."""
+    A_ls_prob_rows, _, A_ls_prob_columns, _, _ = compute_ls_probs(A)
+    r = 30
+    c = 40
+    random_state = np.random.RandomState(7)
+    return Halko(
+        A,
+        r,
+        c,
+        A_ls_prob_rows,
+        A_ls_prob_columns,
+        random_state,
+    )
+
+
 def test_FKV_dimensions():
-    """Test dimensions of FKV-based sketches."""
+    """Test dimensions of sketches."""
     A = np.arange(100 * 100, dtype=np.float64).reshape((100, 100))
     r = 30
     c = 40
@@ -38,6 +55,22 @@ def test_FKV_dimensions():
     assert left_right_sketch_matrix.shape == (30, 40)
 
 
+def test_Halko_dimensions():
+    """Test dimensions of sketches."""
+    A = np.arange(100 * 100, dtype=np.float64).reshape((100, 100))
+    r = 30
+    c = 40
+    sketcher = _get_Halko_sketcher(A, r, c)
+
+    left_sketch_matrix = sketcher.left_project(A)
+    right_sketch_matrix = sketcher.right_project(A)
+    left_right_sketch_matrix = sketcher.right_project(sketcher.left_project(A))
+
+    assert left_sketch_matrix.shape == (30 + 10, 100)
+    assert right_sketch_matrix.shape == (100, 40 + 10)
+    assert left_right_sketch_matrix.shape == (30 + 10, 40 + 10)
+
+
 def test_FKV_samplers():
     """Test FKV-based samplers."""
     A = np.arange(100 * 100, dtype=np.float64).reshape((100, 100))