Kernel function extension

alibi_detect/cd/base.py

@@ -462,7 +462,7 @@ def __init__(
             preprocess_x_ref: bool = True,
             update_x_ref: Optional[Dict[str, int]] = None,
             preprocess_fn: Optional[Callable] = None,
-            sigma: Optional[np.ndarray] = None,
+            # sigma: Optional[np.ndarray] = None,
             configure_kernel_from_x_ref: bool = True,
             n_permutations: int = 100,
             input_shape: Optional[tuple] = None,
@@ -502,12 +502,13 @@ def __init__(
         if p_val is None:
             logger.warning('No p-value set for the drift threshold. Need to set it to detect data drift.')
 
-        self.infer_sigma = configure_kernel_from_x_ref
-        if configure_kernel_from_x_ref and isinstance(sigma, np.ndarray):
-            self.infer_sigma = False
-            logger.warning('`sigma` is specified for the kernel and `configure_kernel_from_x_ref` '
-                           'is set to True. `sigma` argument takes priority over '
-                           '`configure_kernel_from_x_ref` (set to False).')
+        self.infer_parameter = configure_kernel_from_x_ref
+        # self.infer_sigma = configure_kernel_from_x_ref
+        # if configure_kernel_from_x_ref and isinstance(sigma, np.ndarray):
+        #     self.infer_sigma = False
+        #     logger.warning('`sigma` is specified for the kernel and `configure_kernel_from_x_ref` '
+        #                    'is set to True. `sigma` argument takes priority over '
+        #                    '`configure_kernel_from_x_ref` (set to False).')
 
         # optionally already preprocess reference data
         self.p_val = p_val
@@ -612,7 +613,8 @@ def __init__(
             preprocess_x_ref: bool = True,
             update_x_ref: Optional[Dict[str, int]] = None,
             preprocess_fn: Optional[Callable] = None,
-            sigma: Optional[np.ndarray] = None,
+            # kernel: BaseKernel = None,
+            # sigma: Optional[np.ndarray] = None,
             n_permutations: int = 100,
             n_kernel_centers: Optional[int] = None,
             lambda_rd_max: float = 0.2,
@@ -665,7 +667,7 @@ def __init__(
             self.x_ref = preprocess_fn(x_ref)
         else:
             self.x_ref = x_ref
-        self.sigma = sigma
+        # self.sigma = sigma
         self.preprocess_x_ref = preprocess_x_ref
         self.update_x_ref = update_x_ref
         self.preprocess_fn = preprocess_fn

alibi_detect/cd/context_aware.py

@@ -2,6 +2,7 @@
 import numpy as np
 from typing import Callable, Dict, Optional, Union, Tuple
 from alibi_detect.utils.frameworks import has_pytorch, has_tensorflow
+from alibi_detect.utils.pytorch.kernels import BaseKernel
 
 if has_pytorch:
     from alibi_detect.cd.pytorch.context_aware import ContextMMDDriftTorch
@@ -22,8 +23,10 @@ def __init__(
             preprocess_x_ref: bool = True,
             update_ref: Optional[Dict[str, int]] = None,
             preprocess_fn: Optional[Callable] = None,
-            x_kernel: Callable = None,
-            c_kernel: Callable = None,
+            # x_kernel: Callable = None,
+            x_kernel: BaseKernel = None,
+            # c_kernel: Callable = None,
+            c_kernel: BaseKernel = None,
             n_permutations: int = 1000,
             prop_c_held: float = 0.25,
             n_folds: int = 5,

alibi_detect/cd/lsdd.py

@@ -18,7 +18,8 @@ def __init__(
             preprocess_x_ref: bool = True,
             update_x_ref: Optional[Dict[str, int]] = None,
             preprocess_fn: Optional[Callable] = None,
-            sigma: Optional[np.ndarray] = None,
+            # sigma: Optional[np.ndarray] = None,
+            # kernel: BaseKernel = None,
             n_permutations: int = 100,
             n_kernel_centers: Optional[int] = None,
             lambda_rd_max: float = 0.2,

alibi_detect/cd/lsdd_online.py

@@ -17,7 +17,8 @@ def __init__(
             window_size: int,
             backend: str = 'tensorflow',
             preprocess_fn: Optional[Callable] = None,
-            sigma: Optional[np.ndarray] = None,
+            # sigma: Optional[np.ndarray] = None,
+            # kernel: BaseKernel = None,
             n_bootstraps: int = 1000,
             n_kernel_centers: Optional[int] = None,
             lambda_rd_max: float = 0.2,

alibi_detect/cd/mmd.py

@@ -22,7 +22,7 @@ def __init__(
             update_x_ref: Optional[Dict[str, int]] = None,
             preprocess_fn: Optional[Callable] = None,
             kernel: Callable = None,
-            sigma: Optional[np.ndarray] = None,
+            # sigma: Optional[np.ndarray] = None,
             configure_kernel_from_x_ref: bool = True,
             n_permutations: int = 100,
             device: Optional[str] = None,
@@ -84,7 +84,7 @@ def __init__(
                 from alibi_detect.utils.tensorflow.kernels import GaussianRBF
             else:
                 from alibi_detect.utils.pytorch.kernels import GaussianRBF  # type: ignore
-            kwargs.update({'kernel': GaussianRBF})
+            kwargs.update({'kernel': GaussianRBF()})
 
         if backend == 'tensorflow' and has_tensorflow:
             kwargs.pop('device', None)

alibi_detect/cd/mmd_online.py

@@ -4,9 +4,11 @@
 
 if has_pytorch:
     from alibi_detect.cd.pytorch.mmd_online import MMDDriftOnlineTorch
+    from alibi_detect.utils.pytorch.kernels import BaseKernel as BaseKernelTorch
 
 if has_tensorflow:
     from alibi_detect.cd.tensorflow.mmd_online import MMDDriftOnlineTF
+    from alibi_detect.utils.tensorflow.kernels import BaseKernel as BaseKernelTF
 
 
 class MMDDriftOnline:
@@ -17,8 +19,8 @@ def __init__(
             window_size: int,
             backend: str = 'tensorflow',
             preprocess_fn: Optional[Callable] = None,
-            kernel: Callable = None,
-            sigma: Optional[np.ndarray] = None,
+            kernel: Union[BaseKernelTorch, BaseKernelTF] = None,
+            # sigma: Optional[np.ndarray] = None,
             n_bootstraps: int = 1000,
             device: Optional[str] = None,
             verbose: bool = True,
@@ -82,7 +84,7 @@ def __init__(
                 from alibi_detect.utils.tensorflow.kernels import GaussianRBF
             else:
                 from alibi_detect.utils.pytorch.kernels import GaussianRBF  # type: ignore
-            kwargs.update({'kernel': GaussianRBF})
+            kwargs.update({'kernel': GaussianRBF()})
 
         if backend == 'tensorflow' and has_tensorflow:
             kwargs.pop('device', None)

alibi_detect/cd/pytorch/context_aware.py

@@ -4,13 +4,36 @@
 from typing import Callable, Dict, Optional, Tuple, Union
 from alibi_detect.cd.base import BaseContextMMDDrift
 from alibi_detect.utils.pytorch import get_device
-from alibi_detect.utils.pytorch.kernels import GaussianRBF
+from alibi_detect.utils.pytorch.kernels import BaseKernel, GaussianRBF
 from alibi_detect.cd._domain_clf import _SVCDomainClf
 from tqdm import tqdm
 
 logger = logging.getLogger(__name__)
 
 
+def _sigma_median_diag(x: torch.Tensor, y: torch.Tensor, dist: torch.Tensor) -> torch.Tensor:
+    """
+    Private version of the bandwidth estimation function :py:func:`~alibi_detect.utils.pytorch.kernels.sigma_median`,
+    with the +n (and -1) term excluded to account for the diagonal of the kernel matrix.
+
+    Parameters
+    ----------
+    x
+        Tensor of instances with dimension [Nx, features].
+    y
+        Tensor of instances with dimension [Ny, features].
+    dist
+        Tensor with dimensions [Nx, Ny], containing the pairwise distances between `x` and `y`.
+
+    Returns
+    -------
+    The computed bandwidth, `sigma`.
+    """
+    n_median = np.prod(dist.shape) // 2
+    sigma = (.5 * dist.flatten().sort().values[n_median].unsqueeze(dim=-1)) ** .5
+    return sigma
+
+
 class ContextMMDDriftTorch(BaseContextMMDDrift):
     lams: Optional[Tuple[torch.Tensor, torch.Tensor]] = None
 
@@ -22,8 +45,10 @@ def __init__(
             preprocess_x_ref: bool = True,
             update_ref: Optional[Dict[str, int]] = None,
             preprocess_fn: Optional[Callable] = None,
-            x_kernel: Callable = GaussianRBF,
-            c_kernel: Callable = GaussianRBF,
+            # x_kernel: Callable = GaussianRBF,
+            x_kernel: BaseKernel = GaussianRBF(init_fn_sigma=_sigma_median_diag),
+            # c_kernel: Callable = GaussianRBF,
+            c_kernel: BaseKernel = GaussianRBF(init_fn_sigma=_sigma_median_diag),
             n_permutations: int = 1000,
             prop_c_held: float = 0.25,
             n_folds: int = 5,
@@ -98,8 +123,10 @@ def __init__(
         self.device = get_device(device)
 
         # initialize kernel
-        self.x_kernel = x_kernel(init_sigma_fn=_sigma_median_diag) if x_kernel == GaussianRBF else x_kernel
-        self.c_kernel = c_kernel(init_sigma_fn=_sigma_median_diag) if c_kernel == GaussianRBF else c_kernel
+        # self.x_kernel = x_kernel(init_sigma_fn=_sigma_median_diag) if x_kernel == GaussianRBF else x_kernel
+        # self.c_kernel = c_kernel(init_sigma_fn=_sigma_median_diag) if c_kernel == GaussianRBF else c_kernel
+        self.x_kernel = x_kernel
+        self.c_kernel = c_kernel
 
         # Initialize classifier (hardcoded for now)
         self.clf = _SVCDomainClf(self.c_kernel)
@@ -244,26 +271,3 @@ def _pick_lam(self, lams: torch.Tensor, K: torch.Tensor, L: torch.Tensor, n_fold
             kxx = torch.ones_like(lWk).to(lWk.device) * torch.max(K)
             losses += (lWKWl + kxx - 2*lWk).sum(-1)
         return lams[torch.argmin(losses)]
-
-
-def _sigma_median_diag(x: torch.Tensor, y: torch.Tensor, dist: torch.Tensor) -> torch.Tensor:
-    """
-    Private version of the bandwidth estimation function :py:func:`~alibi_detect.utils.pytorch.kernels.sigma_median`,
-    with the +n (and -1) term excluded to account for the diagonal of the kernel matrix.
-
-    Parameters
-    ----------
-    x
-        Tensor of instances with dimension [Nx, features].
-    y
-        Tensor of instances with dimension [Ny, features].
-    dist
-        Tensor with dimensions [Nx, Ny], containing the pairwise distances between `x` and `y`.
-
-    Returns
-    -------
-    The computed bandwidth, `sigma`.
-    """
-    n_median = np.prod(dist.shape) // 2
-    sigma = (.5 * dist.flatten().sort().values[n_median].unsqueeze(dim=-1)) ** .5
-    return sigma

alibi_detect/cd/pytorch/lsdd.py

@@ -3,7 +3,7 @@
 from typing import Callable, Dict, Optional, Tuple, Union
 from alibi_detect.cd.base import BaseLSDDDrift
 from alibi_detect.utils.pytorch import get_device
-from alibi_detect.utils.pytorch.kernels import GaussianRBF
+from alibi_detect.utils.pytorch.kernels import BaseKernel, GaussianRBF
 from alibi_detect.utils.pytorch.distance import permed_lsdds
 
 
@@ -15,7 +15,8 @@ def __init__(
             preprocess_x_ref: bool = True,
             update_x_ref: Optional[Dict[str, int]] = None,
             preprocess_fn: Optional[Callable] = None,
-            sigma: Optional[np.ndarray] = None,
+            # sigma: Optional[np.ndarray] = None,
+            kernel: BaseKernel = GaussianRBF(),
             n_permutations: int = 100,
             n_kernel_centers: Optional[int] = None,
             lambda_rd_max: float = 0.2,
@@ -67,7 +68,8 @@ def __init__(
             preprocess_x_ref=preprocess_x_ref,
             update_x_ref=update_x_ref,
             preprocess_fn=preprocess_fn,
-            sigma=sigma,
+            # sigma=sigma,
+            # kernel=kernel,
             n_permutations=n_permutations,
             n_kernel_centers=n_kernel_centers,
             lambda_rd_max=lambda_rd_max,
@@ -83,24 +85,25 @@ def __init__(
         #  in the method signature, so we can't cast it to torch.Tensor unless we change the signature
         #  to also accept torch.Tensor. We also can't redefine it's type as that would involve enabling
         #  --allow-redefinitions in mypy settings (which we might do eventually).
+        self.kernel = kernel
         if self.preprocess_x_ref or self.preprocess_fn is None:
             x_ref = torch.as_tensor(self.x_ref).to(self.device)  # type: ignore[assignment]
             self._configure_normalization(x_ref)  # type: ignore[arg-type]
             x_ref = self._normalize(x_ref)
-            self._initialize_kernel(x_ref)  # type: ignore[arg-type]
+            # self._initialize_kernel(x_ref)  # type: ignore[arg-type]
             self._configure_kernel_centers(x_ref)  # type: ignore[arg-type]
             self.x_ref = x_ref.cpu().numpy()  # type: ignore[union-attr]
             # For stability in high dimensions we don't divide H by (pi*sigma^2)^(d/2)
             # Results in an alternative test-stat of LSDD*(pi*sigma^2)^(d/2). Same p-vals etc.
             self.H = GaussianRBF(np.sqrt(2.) * self.kernel.sigma)(self.kernel_centers, self.kernel_centers)
 
-    def _initialize_kernel(self, x_ref: torch.Tensor):
-        if self.sigma is None:
-            self.kernel = GaussianRBF()
-            _ = self.kernel(x_ref, x_ref, infer_sigma=True)
-        else:
-            sigma = torch.from_numpy(self.sigma)
-            self.kernel = GaussianRBF(sigma)
+    # def _initialize_kernel(self, x_ref: torch.Tensor):
+    #     if self.sigma is None:
+    #         self.kernel = GaussianRBF()
+    #         _ = self.kernel(x_ref, x_ref, infer_sigma=True)
+    #     else:
+    #         sigma = torch.from_numpy(self.sigma)
+    #         self.kernel = GaussianRBF(sigma)
 
     def _configure_normalization(self, x_ref: torch.Tensor, eps: float = 1e-12):
         x_ref_means = x_ref.mean(0)
@@ -140,7 +143,7 @@ def score(self, x: Union[np.ndarray, list]) -> Tuple[float, float, float]:
         if self.preprocess_fn is not None and self.preprocess_x_ref is False:
             self._configure_normalization(x_ref)  # type: ignore[arg-type]
             x_ref = self._normalize(x_ref)
-            self._initialize_kernel(x_ref)  # type: ignore[arg-type]
+            # self._initialize_kernel(x_ref)  # type: ignore[arg-type]
             self._configure_kernel_centers(x_ref)  # type: ignore[arg-type]
             self.H = GaussianRBF(np.sqrt(2.) * self.kernel.sigma)(self.kernel_centers, self.kernel_centers)
 

alibi_detect/cd/pytorch/lsdd_online.py

@@ -4,7 +4,8 @@
 from typing import Any, Callable, Optional, Union
 from alibi_detect.cd.base_online import BaseMultiDriftOnline
 from alibi_detect.utils.pytorch import get_device
-from alibi_detect.utils.pytorch import GaussianRBF, permed_lsdds, quantile
+from alibi_detect.utils.pytorch import permed_lsdds, quantile
+from alibi_detect.utils.pytorch.kernels import BaseKernel, GaussianRBF
 
 
 class LSDDDriftOnlineTorch(BaseMultiDriftOnline):
@@ -14,7 +15,8 @@ def __init__(
             ert: float,
             window_size: int,
             preprocess_fn: Optional[Callable] = None,
-            sigma: Optional[np.ndarray] = None,
+            # sigma: Optional[np.ndarray] = None,
+            kernel: BaseKernel = GaussianRBF(),
             n_bootstraps: int = 1000,
             n_kernel_centers: Optional[int] = None,
             lambda_rd_max: float = 0.2,
@@ -86,14 +88,15 @@ def __init__(
         self._configure_normalization()
 
         # initialize kernel
-        if sigma is None:
-            x_ref = torch.from_numpy(self.x_ref).to(self.device)  # type: ignore[assignment]
-            self.kernel = GaussianRBF()
-            _ = self.kernel(x_ref, x_ref, infer_sigma=True)
-        else:
-            sigma = torch.from_numpy(sigma).to(self.device) if isinstance(sigma,  # type: ignore[assignment]
-                                                                          np.ndarray) else None
-            self.kernel = GaussianRBF(sigma)  # type: ignore[arg-type]
+        # if sigma is None:
+        #     x_ref = torch.from_numpy(self.x_ref).to(self.device)  # type: ignore[assignment]
+        #     self.kernel = GaussianRBF()
+        #     _ = self.kernel(x_ref, x_ref, infer_sigma=True)
+        # else:
+        #     sigma = torch.from_numpy(sigma).to(self.device) if isinstance(sigma,  # type: ignore[assignment]
+        #                                                                   np.ndarray) else None
+        #     self.kernel = GaussianRBF(sigma)  # type: ignore[arg-type]
+        self.kernel = kernel
 
         if self.n_kernel_centers is None:
             self.n_kernel_centers = 2 * window_size