Refactor kernel

qkernel/kernel.py

@@ -47,8 +47,6 @@ class Kernel:
     def __init__(
         self,
         p_train: np.ndarray,
-        p_test: np.ndarray,
-        excitations: bool = True,
         distance_fn: str = "exp_js",
         distance_kwargs: Optional[Dict[str, float]] = None,
     ) -> None:
@@ -59,28 +57,15 @@ def __init__(
         ----------
         p_train : np.ndarray
             Training probability distributions of shape (n_train, 2^n_features).
-        p_test : np.ndarray
-            Test probability distributions of shape (n_test, 2^n_features).
-        excitations : bool, optional
-            If True, transforms probability vectors into excitation-count
-            representations before computing kernels.
         distance_fn : str, optional
             Key selecting the distance function from DISTANCES_DICT
         distance_kwargs : dict or None, optional
             Additional keyword arguments passed to the distance function.
         """
         self.distance_fn = DISTANCES_DICT[distance_fn]
         self.distance_kwargs = distance_kwargs or {}
-
-        if excitations:
-            self.p_train = compute_excitation_count(p_train)
-            self.p_test = compute_excitation_count(p_test)
-        else:
-            self.p_train = p_train
-            self.p_test = p_test
-
+        self.p_train = p_train
         self.n_train = len(self.p_train)
-        self.n_test = len(self.p_test)
 
     def compute_gram_train(self) -> np.ndarray:
         """
@@ -105,10 +90,14 @@ def compute_gram_train(self) -> np.ndarray:
 
         return gram_train
 
-    def compute_gram_test(self) -> np.ndarray:
+    def compute_gram_test(self, p_test: np.ndarray,) -> np.ndarray:
         """
         Compute the Gram (kernel) matrix between test and training sets.
 
+        Parameters
+        ----------
+        p_test : np.ndarray
+            Test probability distributions of shape (n_test, 2^n_features).
 
         Returns
         -------
@@ -117,12 +106,13 @@ def compute_gram_test(self) -> np.ndarray:
             corresponds to the kernel value between test sample i and
             training sample j.
         """
-        gram_test: np.ndarray = np.ones((self.n_test, self.n_train))
+        n_test = len(p_test)
+        gram_test: np.ndarray = np.ones((n_test, self.n_train))
 
-        for i in range(self.n_test):
+        for i in range(n_test):
             for j in range(self.n_train):
                 gram_test[i, j] = self.distance_fn(
-                    self.p_test[i],
+                    p_test[i],
                     self.p_train[j],
                     **self.distance_kwargs,
                 )

qkernel/model.py

@@ -2,7 +2,7 @@
 
 import numpy as np
 from sklearn.svm import SVC
-
+from sklearn.kernel_ridge import KernelRidge
 
 class Model(ABC):
     """
@@ -113,3 +113,61 @@ def predict(self, gram: np.ndarray) -> np.ndarray:
             Predicted labels.
         """
         return self.model.predict(gram)
+
+
+class KRR(Model):
+    """
+    Quantum-kernel-compatible Kernel Ridge Regression wrapper.
+
+    This class uses scikit-learn's KernelRidge with a precomputed kernel,
+    making it suitable for quantum kernel methods (QKRR-style pipelines).
+
+    Parameters
+    ----------
+    alpha : float, default=1.0
+        Regularization strength. Larger values mean stronger regularization.
+    **kwargs
+        Additional keyword arguments passed to sklearn.kernel_ridge.KernelRidge.
+    """
+
+    def __init__(self, alpha: float = 1.0, **kwargs) -> None:
+        """
+        Initialize the KRR model with a precomputed kernel.
+
+        Parameters
+        ----------
+        alpha : float, default=1.0
+            Regularization strength.
+        **kwargs
+            Keyword arguments forwarded to sklearn's KernelRidge.
+        """
+        self.model = KernelRidge(kernel="precomputed", alpha=alpha, **kwargs)
+
+    def train(self, gram_train: np.ndarray, y_train: np.ndarray) -> None:
+        """
+        Train the KRR model.
+
+        Parameters
+        ----------
+        gram_train : np.ndarray
+            Training Gram matrix of shape (n_train, n_train).
+        y_train : np.ndarray
+            Training targets.
+        """
+        self.model.fit(gram_train, y_train)
+
+    def predict(self, gram: np.ndarray) -> np.ndarray:
+        """
+        Predict target values.
+
+        Parameters
+        ----------
+        gram : np.ndarray
+            Gram matrix of shape (n_samples, n_train) for prediction.
+
+        Returns
+        -------
+        np.ndarray
+            Predicted values.
+        """
+        return self.model.predict(gram)