kernel.py

from time import time
from typing import Optional

import numpy as np
import scipy
from scipy.sparse import bsr_array, coo_array
from sklearn.kernel_approximation import Nystroem
from tqdm import tqdm


class Kernel:

    def __init__(
        self,
        type: str,
        eigen_cutoff: int,
        eigen_threshold: Optional[float] = None,
        sigma: Optional[float] = None,
        coef: Optional[float] = None,
        degree: Optional[int] = None,
        alpha: Optional[float] = None,
    ):
        assert type in [
            "linear",
            "gaussian",
            "poly",
            "sigmoid",
            "laplacian",
            "cosine",
            "chi",
        ]
        if type == "gaussian":
            assert sigma is not None
        if type == "poly":
            assert coef is not None
            assert degree is not None
        if type == "sigmoid":
            assert coef is not None
            assert alpha is not None
        if type == "laplacian":
            assert alpha is not None

        self.eigen_cutoff = eigen_cutoff
        self.eigen_threshold = eigen_threshold
        self.sigma = sigma
        self.type = type
        self.coef = coef
        self.degree = degree
        self.alpha = alpha

    def k(self, x: np.ndarray, y: np.ndarray) -> float:
        if self.type == "linear":
            return x @ y
        elif self.type == "gaussian":
            return np.exp(-np.linalg.norm(x - y) ** 2 / (2 * self.sigma**2))
        elif self.type == "poly":
            return (self.coef + x @ y) ** self.degree
        elif self.type == "sigmoid":
            return np.tanh(self.coef + self.alpha * x @ y)
        elif self.type == "laplacian":
            return np.exp(self.alpha * -np.linalg.norm(x - y, ord=1))
        else:
            raise ValueError("Invalid kernel type")

    def gram_matrix(self, X: np.ndarray, Y: np.ndarray = None):
        if Y is None:
            Y = X
        if self.type == "linear":
            return np.dot(X, Y.T)
        elif self.type == "gaussian":
            X_norm = np.sum(X**2, axis=1)[:, np.newaxis]  # shape: (n, 1)
            X_test_norm = np.sum(Y**2, axis=1)[np.newaxis, :]  # shape: (1, m)
            dist_matrix = X_norm + X_test_norm - 2 * np.dot(X, Y.T)  # shape: (n, m)
            return np.exp(-dist_matrix / (2 * self.sigma**2))
        elif self.type == "poly":
            return (self.coef + X @ Y.T) ** self.degree
        elif self.type == "sigmoid":
            K = np.dot(X, Y.T)
            K = np.tanh(self.coef + self.alpha * K)
            return K
        elif self.type == "laplacian":
            X1 = X[:, np.newaxis, :]  # Shape: (n_samples, 1, n_features)
            X2 = Y[np.newaxis, :, :]  # Shape: (1, n_samples, n_features)
            # Compute the L1 distance
            L1_distance = np.sum(
                np.abs(X1 - X2), axis=2
            )  # Shape: (n_samples, n_samples)
            # Compute the Laplacian kernel matrix
            K = np.exp(-L1_distance * self.alpha)
            return K
        elif self.type == "cosine":
            return np.dot(X, Y.T) / (
                np.linalg.norm(X, axis=1)[:, np.newaxis]
                * np.linalg.norm(Y, axis=1)[np.newaxis, :]
            )
        elif self.type == "chi":
            from sklearn.metrics.pairwise import chi2_kernel

            if min(np.min(X), np.min(Y)) < 0:
                m = np.min([np.min(X), np.min(Y)])
                X += abs(m)
                Y += abs(m)

            return chi2_kernel(X, Y, gamma=self.alpha)
        else:
            raise NotImplementedError

    def generate_grammian(self, X: np.ndarray) -> np.ndarray:
        if self.type == "linear":
            return X @ X.T
        elif self.type == "gaussian":
            X_norm = np.sum(X**2, axis=-1)
            distances = X_norm[:, None] + X_norm[None, :] - 2 * np.dot(X, X.T)
            kernel_matrix = np.exp(-distances / (2 * self.sigma**2))
            return kernel_matrix
        elif self.type == "poly":
            return (self.coef + X @ X.T) ** self.degree
        elif self.type == "sigmoid":
            K = np.dot(X, X.T)
            K = np.tanh(self.coef + self.alpha * K)
            return K
        elif self.type == "laplacian":
            X1 = X[:, np.newaxis, :]  # Shape: (n_samples, 1, n_features)
            X2 = X[np.newaxis, :, :]  # Shape: (1, n_samples, n_features)
            # Compute the L1 distance
            L1_distance = np.sum(
                np.abs(X1 - X2), axis=2
            )  # Shape: (n_samples, n_samples)
            # Compute the Laplacian kernel matrix
            K = np.exp(-L1_distance * self.alpha)
            return K
        elif self.type == "cosine":
            return self.gram_matrix(X)
        elif self.type == "chi":
            return self.gram_matrix(X)

        n = X.shape[0]
        K = np.zeros((n, n))
        for i in tqdm(range(n), total=n):
            for j in range(i, n):
                K[i, j] = self.k(X[i], X[j])
                K[j, i] = K[i, j]
        return K

    def generate_projection_triplets(
        self, triplets_train: np.ndarray, triplets_tests: np.ndarray
    ) -> np.ndarray:
        n_train, three, d = triplets_train.shape

        X_train = triplets_train.reshape(n_train * three, d)

        X_tests = []
        for triplets_test in triplets_tests:
            n_test, three, d = triplets_test.shape
            X_test = triplets_test.reshape(n_test * three, d)
            X_tests.append(X_test)

        varphi_train, varphi_tests = self.generate_projection(X_train, X_tests)

        varphi_triplets_train = varphi_train.reshape(n_train, three, -1)
        varphi_triplets_tests = []
        for varphi_test in varphi_tests:
            varphi_triplets_tests.append(varphi_test.reshape(n_test, three, -1))

        return varphi_triplets_train, varphi_triplets_tests

    def generate_projection(self, X: np.ndarray, X_tests: np.ndarray) -> np.ndarray:
        """
        Generate the projections
        Each row of the output matrix is the projection of the corresponding row of the input matrix
        """
        X = X.astype(np.float32)
        n = X.shape[0]

        # form the Grammian matrix
        print("Generating Grammian matrix")
        if self.type == "gaussian":
            feature_map_nystroem = Nystroem(
                kernel="rbf",
                gamma=1 / (2 * self.sigma),
                random_state=1,
                n_components=500,
            )
        elif self.type == "poly":
            feature_map_nystroem = Nystroem(
                kernel="poly",
                coef0=self.coef,
                degree=self.degree,
                gamma=1,
                random_state=1,
                n_components=500,
            )
        elif self.type == "sigmoid":
            feature_map_nystroem = Nystroem(
                kernel="sigmoid",
                gamma=self.alpha,
                coef0=self.coef,
                random_state=1,
                n_components=500,
            )
        elif self.type == "laplacian":
            feature_map_nystroem = Nystroem(
                kernel="laplacian",
                gamma=self.alpha,
                random_state=1,
                n_components=500,
            )
        elif self.type == "linear":
            feature_map_nystroem = Nystroem(
                kernel="linear",
                random_state=1,
                n_components=500,
            )
        else:
            raise NotImplementedError
        K = feature_map_nystroem.fit_transform(X)
        print("Centering K")
        K_centered = K - 1 / n * np.ones((n, 1)) @ (np.ones((1, n)) @ K)
        print("Singular Value Decomposition")
        U, _, _ = np.linalg.svd(K_centered, full_matrices=False)
        # K = K @ K.T
        # K = self.generate_grammian(X)

        # center the Grammian matrix
        # print("Centering Grammian matrix")
        # row_mean = np.mean(K, axis=0)
        # col_mean = np.mean(K, axis=1).reshape(-1, 1)
        # grand_mean = np.mean(K)
        # K_bar = K - row_mean - row_mean.reshape(-1, 1) + grand_mean

        # compute the eigenvectors, form A
        # print("Computing eigenvectors and eigenvalues")
        # if self.eigen_threshold is not None:
        #     print(f"Thresholding eigenvalues")
        #     eigvals, eigvecs = scipy.linalg.eigh(
        #         K_bar,
        #         subset_by_value=[self.eigen_threshold, np.inf],
        #         check_finite=False,
        #     )
        # else:
        #     print(f"Using top {self.eigen_cutoff} eigenvalues")
        #     print(f"{K_bar.shape = }")
        #     eigvals, eigvecs = scipy.sparse.linalg.eigsh(K_bar, k=self.eigen_cutoff)
        # idx = np.argsort(eigvals)[::-1]

        # eigvals = eigvals[idx]
        # A = eigvecs[:, idx]

        A = U[:, : self.eigen_cutoff]

        # compute projections
        print("Computing projections")
        phi = K @ (K.T @ A)

        print("Computing test projections")
        phi_tests = []
        for i in range(len(X_tests)):
            X_test = X_tests[i].astype(np.float32)
            phi_test = self.gram_matrix(X_test, X) @ A
            phi_tests.append(phi_test)

        return phi, phi_tests


def main():
    lk = Kernel(type="gaussian", sigma=1)
    A = lk.generate_projection(np.array([[1, 1, 1], [3, 3, 3], [5, 5, 5]]))


if __name__ == "__main__":
    main()