Kernel function extension

alibi_detect/cd/base.py

@@ -462,7 +462,6 @@ 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,
             configure_kernel_from_x_ref: bool = True,
             n_permutations: int = 100,
             input_shape: Optional[tuple] = None,
@@ -502,12 +501,7 @@ 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
 
         # optionally already preprocess reference data
         self.p_val = p_val
@@ -612,7 +606,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 +660,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,8 @@ 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: BaseKernel = GaussianRBF(init_fn_sigma=_sigma_median_diag),
+            c_kernel: BaseKernel = GaussianRBF(init_fn_sigma=_sigma_median_diag),
             n_permutations: int = 1000,
             prop_c_held: float = 0.25,
             n_folds: int = 5,
@@ -97,9 +120,8 @@ def __init__(
         # set device
         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
+        self.c_kernel = c_kernel
 
         # Initialize classifier (hardcoded for now)
         self.clf = _SVCDomainClf(self.c_kernel)
@@ -244,26 +266,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

@@ -15,7 +15,6 @@ 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,
             n_permutations: int = 100,
             n_kernel_centers: Optional[int] = None,
             lambda_rd_max: float = 0.2,
@@ -67,7 +66,6 @@ def __init__(
             preprocess_x_ref=preprocess_x_ref,
             update_x_ref=update_x_ref,
             preprocess_fn=preprocess_fn,
-            sigma=sigma,
             n_permutations=n_permutations,
             n_kernel_centers=n_kernel_centers,
             lambda_rd_max=lambda_rd_max,
@@ -83,25 +81,17 @@ 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 = GaussianRBF()
         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._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 _configure_normalization(self, x_ref: torch.Tensor, eps: float = 1e-12):
         x_ref_means = x_ref.mean(0)
         x_ref_stds = x_ref.std(0)
@@ -140,7 +130,6 @@ 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._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 GaussianRBF
 
 
 class LSDDDriftOnlineTorch(BaseMultiDriftOnline):
@@ -14,7 +15,6 @@ def __init__(
             ert: float,
             window_size: int,
             preprocess_fn: Optional[Callable] = None,
-            sigma: Optional[np.ndarray] = None,
             n_bootstraps: int = 1000,
             n_kernel_centers: Optional[int] = None,
             lambda_rd_max: float = 0.2,
@@ -85,15 +85,7 @@ 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]
+        self.kernel = GaussianRBF()
 
         if self.n_kernel_centers is None:
             self.n_kernel_centers = 2 * window_size

alibi_detect/cd/pytorch/mmd.py

@@ -5,7 +5,7 @@
 from alibi_detect.cd.base import BaseMMDDrift
 from alibi_detect.utils.pytorch import get_device
 from alibi_detect.utils.pytorch.distance import mmd2_from_kernel_matrix
-from alibi_detect.utils.pytorch.kernels import GaussianRBF
+from alibi_detect.utils.pytorch.kernels import BaseKernel, GaussianRBF
 
 logger = logging.getLogger(__name__)
 
@@ -18,8 +18,7 @@ def __init__(
             preprocess_x_ref: bool = True,
             update_x_ref: Optional[Dict[str, int]] = None,
             preprocess_fn: Optional[Callable] = None,
-            kernel: Callable = GaussianRBF,
-            sigma: Optional[np.ndarray] = None,
+            kernel: BaseKernel = GaussianRBF(),
             configure_kernel_from_x_ref: bool = True,
             n_permutations: int = 100,
             device: Optional[str] = None,
@@ -66,7 +65,6 @@ def __init__(
             preprocess_x_ref=preprocess_x_ref,
             update_x_ref=update_x_ref,
             preprocess_fn=preprocess_fn,
-            sigma=sigma,
             configure_kernel_from_x_ref=configure_kernel_from_x_ref,
             n_permutations=n_permutations,
             input_shape=input_shape,
@@ -77,22 +75,19 @@ def __init__(
         # set device
         self.device = get_device(device)
 
-        # initialize kernel
-        sigma = torch.from_numpy(sigma).to(self.device) if isinstance(sigma,  # type: ignore[assignment]
-                                                                      np.ndarray) else None
-        self.kernel = kernel(sigma) if kernel == GaussianRBF else kernel
+        self.kernel = kernel
 
         # compute kernel matrix for the reference data
-        if self.infer_sigma or isinstance(sigma, torch.Tensor):
+        if self.infer_parameter:
             x = torch.from_numpy(self.x_ref).to(self.device)
-            self.k_xx = self.kernel(x, x, infer_sigma=self.infer_sigma)
-            self.infer_sigma = False
+            self.k_xx = self.kernel(x, x, infer_parameter=self.infer_parameter)
+            self.infer_parameter = False
         else:
-            self.k_xx, self.infer_sigma = None, True
+            self.k_xx, self.infer_parameter = None, True
 
     def kernel_matrix(self, x: torch.Tensor, y: torch.Tensor) -> torch.Tensor:
         """ Compute and return full kernel matrix between arrays x and y. """
-        k_xy = self.kernel(x, y, self.infer_sigma)
+        k_xy = self.kernel(x, y, self.infer_parameter)
         k_xx = self.k_xx if self.k_xx is not None and self.update_x_ref is None else self.kernel(x, x)
         k_yy = self.kernel(y, y)
         kernel_mat = torch.cat([torch.cat([k_xx, k_xy], 1), torch.cat([k_xy.T, k_yy], 1)], 0)

alibi_detect/cd/pytorch/mmd_online.py

@@ -4,7 +4,7 @@
 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.kernels import GaussianRBF
+from alibi_detect.utils.pytorch.kernels import BaseKernel, GaussianRBF
 from alibi_detect.utils.pytorch import zero_diag, quantile
 
 
@@ -15,8 +15,7 @@ def __init__(
             ert: float,
             window_size: int,
             preprocess_fn: Optional[Callable] = None,
-            kernel: Callable = GaussianRBF,
-            sigma: Optional[np.ndarray] = None,
+            kernel: BaseKernel = GaussianRBF(),
             n_bootstraps: int = 1000,
             device: Optional[str] = None,
             verbose: bool = True,
@@ -74,14 +73,11 @@ def __init__(
         # set device
         self.device = get_device(device)
 
-        # initialize kernel
-        sigma = torch.from_numpy(sigma).to(self.device) if isinstance(sigma,  # type: ignore[assignment]
-                                                                      np.ndarray) else None
-        self.kernel = kernel(sigma) if kernel == GaussianRBF else kernel
+        self.kernel = kernel
 
         # compute kernel matrix for the reference data
         self.x_ref = torch.from_numpy(self.x_ref).to(self.device)
-        self.k_xx = self.kernel(self.x_ref, self.x_ref, infer_sigma=(sigma is None))
+        self.k_xx = self.kernel(self.x_ref, self.x_ref, infer_parameter=self.kernel.init_required)
 
         self._configure_thresholds()
         self._initialise()