utils.py

import cvxpy as cp
import numpy as np
from scipy.stats import ortho_group

from distributions import Distribution


def generate_y_true_from_triplets_under_kernel(
    triplets_train, triplets_tests, kernel, K, rng
) -> np.ndarray:

    # project triplets into reprsentation space
    varphi_triplets_tests = []
    for triplets_test in triplets_tests:
        varphi_triplets_train, varphi_triplets_test = (
            kernel.generate_projection_triplets(triplets_train, triplets_test)
        )
        varphi_triplets_tests.append(varphi_triplets_test)

    # generate y from projected triplets
    varphi_y_true_train, dists = generate_y_true_from_triplets(
        K, varphi_triplets_train, rng, deterministic=False
    )
    varphi_y_true_tests = []
    for varphi_triplets_test in varphi_triplets_tests:
        varphi_y_true_test, _ = generate_y_true_from_triplets(
            K, varphi_triplets_test, rng, deterministic=True
        )
        varphi_y_true_tests.append(varphi_y_true_test)

    return varphi_y_true_train, varphi_y_true_tests, dists


def generate_all_triplets(n: int) -> np.ndarray:
    T = []
    for i in range(n):
        for j in list(set(range(n)) - {i}):
            for k in list(set(range(j + 1, n)) - {i}):
                T.append((i, j, k))
    return np.array(T)


def generate_triplets(
    rng: np.random.Generator,
    n: int,
    size: int = 10000,
    mode: str = "jain",
    distribution: Distribution = None,
) -> np.ndarray:
    """
    Generate triplets according to https://arxiv.org/pdf/1709.06171
    Each triplet is a tuple (i, j, k) where i, j, k are indices of the triplets
    is drawn uniformly at random from the full set of n * [(n - 1) choose 2] triplets.
    """
    if mode == "jain":
        T = []
        for _ in range(size):
            i = rng.integers(0, n, size=1)[0]
            j, k = rng.choice(list(set(range(n)) - {i}), size=2, replace=False)
            j, k = min(j, k), max(j, k)
            T.append((i, j, k))
            assert 0 <= 1 < n and i != j and j != k and i != k and j < k
        T = np.array(T)
        return T
    elif mode == "tatli":
        assert distribution is not None
        assert rng is None
        assert n is None

        triplets_size = (size, 3)
        triplets = distribution.sample(triplets_size)
        return triplets
    else:
        raise ValueError(f"Invalid mode: {mode}")


def generate_K_star(rng: np.random.Generator, p: int, d: int) -> dict:
    r"""
    Generate K^\star according to https://arxiv.org/pdf/1709.06171
    K^\star = p / \sqrt{d} UU^\top
    where U is a random p x d matrix with orthonormal columns.
    """
    assert p >= d

    U = ortho_group(dim=p, seed=rng).rvs()[:, :d]
    K_star_dense = p / np.sqrt(d) * U @ U.T
    K_star_sparse = p / np.sqrt(d) * U @ U.T
    mask = rng.choice(a=np.arange(p), replace=False, size=p - d)
    K_star_sparse[mask] = 0
    K_star_sparse = K_star_sparse.T
    K_star_sparse[mask] = 0
    K_star_sparse = K_star_sparse.T

    return {
        "dense": K_star_dense,
        "sparse": K_star_sparse,
        "U": U,
    }


def M_t(
    K,
    x_i,
    x_j,
    x_k,
):
    return np.dot(x_i - x_j, np.dot(K, x_i - x_j)) - np.dot(
        x_i - x_k, np.dot(K, x_i - x_k)
    )


def generate_y_true_from_triplets(
    K: np.ndarray,
    triplets: np.ndarray,
    rng: np.random.Generator,
    deterministic: bool = False,
) -> np.ndarray:
    y_true_s = []
    dist_diff_s = []
    for x_i, x_j, x_k in triplets:
        dist_diff = M_t(K, x_i, x_j, x_k)

        if deterministic:
            y_true = -1 if dist_diff < 0 else 1
        else:
            q_t = 1 / (1 + np.exp(dist_diff))
            y_true = -1 if rng.uniform() < q_t else 1

        y_true_s.append(y_true)
        dist_diff_s.append(dist_diff)
    y_true_s = np.array(y_true_s)
    dist_diff_s = np.array(dist_diff_s)
    return y_true_s, dist_diff_s


def generate_y_true(
    K: np.ndarray,
    X: np.ndarray,
    triplets: np.ndarray,
    rng: np.random.Generator,
    deterministic: bool = False,
) -> np.ndarray:
    y_true_s = []
    dist_diff_s = []

    for x_i, x_j, x_k in X[triplets]:

        dist_diff = M_t(K, x_i, x_j, x_k)

        if deterministic:
            y_true = -1 if dist_diff < 0 else 1
        else:
            q_t = 1 / (1 + np.exp(dist_diff))
            y_true = -1 if rng.uniform() < q_t else 1

        dist_diff_s.append(dist_diff)
        y_true_s.append(y_true)

    y_true_s = np.array(y_true_s)
    dist_diff_s = np.array(dist_diff_s)
    return y_true_s, dist_diff_s


def generate_y_from_K_and_triplets(K: np.ndarray, triplets: np.ndarray) -> np.ndarray:
    y = []
    for x_i, x_j, x_k in triplets:

        dist_diff = M_t(K, x_i, x_j, x_k)
        y_hat = -1 if dist_diff < 0 else 1
        y.append(y_hat)

    y = np.array(y)
    return y


def generate_y_from_K(K: np.ndarray, X: np.ndarray, triplets: np.ndarray) -> np.ndarray:
    y = []

    for x_i, x_j, x_k in X[triplets]:

        dist_diff = M_t(K, x_i, x_j, x_k)
        y_hat = -1 if dist_diff < 0 else 1
        y.append(y_hat)

    y = np.array(y)
    return y


def risk_fn(K, K_star, triplets, X) -> float:
    i_indices, j_indices, k_indices = zip(*triplets)
    i_indices, j_indices, k_indices = (
        np.array(i_indices),
        np.array(j_indices),
        np.array(k_indices),
    )
    diff_i_j = X[i_indices] - X[j_indices]
    diff_i_k = X[i_indices] - X[k_indices]

    M_t_K_s = np.sum(K @ diff_i_j.T * diff_i_j.T, axis=0) - np.sum(
        K @ diff_i_k.T * diff_i_k.T, axis=0
    )
    M_t_K_star_s = np.sum(K_star @ diff_i_j.T * diff_i_j.T, axis=0) - np.sum(
        K_star @ diff_i_k.T * diff_i_k.T, axis=0
    )

    f_M_t_K_s = 1 / (1 + np.exp(M_t_K_s))
    f_M_t_K_star_s = 1 / (1 + np.exp(M_t_K_star_s))

    risk = np.sum(
        f_M_t_K_star_s * np.log(1 / (f_M_t_K_s + 1e-10))
        + (1 - f_M_t_K_star_s) * np.log(1 / (1 - f_M_t_K_s + 1e-10))
    )

    risk /= len(triplets)

    # for i, j, k in tqdm(triplets):

    #     M_t_K = np.dot(X[i] - X[j], np.dot(K, X[i] - X[j])) - np.dot(
    #         X[i] - X[k], np.dot(K, X[i] - X[k])
    #     )
    #     M_t_K_star = np.dot(X[i] - X[j], np.dot(K_star, X[i] - X[j])) - np.dot(
    #         X[i] - X[k], np.dot(K_star, X[i] - X[k])
    #     )

    #     f_M_t_K = 1 / (1 + np.exp(M_t_K))
    #     f_M_t_K_star = 1 / (1 + np.exp(M_t_K_star))

    #     term_1 = f_M_t_K_star * np.log(1 / (f_M_t_K + 1e-10))
    #     term_2 = (1 - f_M_t_K_star) * np.log(1 / (1 - f_M_t_K + 1e-10))

    #     risk += term_1 + term_2

    return risk


def relative_excess_risk_fn(K, K_star, triplets, X) -> float:
    risk_K = risk_fn(K, K_star, triplets, X)
    risk_K_star = risk_fn(K_star, K_star, triplets, X)
    return (risk_K - risk_K_star) / (risk_K_star + 1e-10)


def relative_squared_recovery_error_fn(K, K_star) -> float:
    return np.linalg.norm(K - K_star) ** 2 / np.linalg.norm(K_star) ** 2


def solve_for_K(
    p: int,
    triplets: np.ndarray,
    y_true: np.ndarray,
    gamma: float,
    lambda_: float,
    loss_type: str,
    constraint_type: str,
    solver: str,
    verbose: bool,
) -> np.ndarray:
    diff_i_j = triplets[:, 0, :] - triplets[:, 1, :]
    diff_i_k = triplets[:, 0, :] - triplets[:, 2, :]

    K_hat = cp.Variable((p, p), PSD=True)
    M_t_K_s = cp.sum(cp.multiply(K_hat @ diff_i_j.T, diff_i_j.T), axis=0) - cp.sum(
        cp.multiply(K_hat @ diff_i_k.T, diff_i_k.T), axis=0
    )

    if loss_type == "logistic":
        risk = cp.sum(cp.logistic(-1 * cp.multiply(y_true, M_t_K_s))) / len(triplets)
    elif loss_type == "hinge":
        risk = cp.sum(cp.pos(1 - cp.multiply(y_true, M_t_K_s))) / len(triplets)
    else:
        raise ValueError(f"Invalid loss_type: {loss_type}")

    if constraint_type == "nuc":
        constraints = [cp.norm(K_hat, p="nuc") <= lambda_]
    elif constraint_type == "fro":
        constraints = [cp.norm(K_hat, p="fro") <= lambda_]
    elif constraint_type == "l_12":
        constraints = [cp.mixed_norm(K_hat, p=2, q=1) <= lambda_]
    else:
        raise ValueError(f"Invalid constraint_type: {constraint_type}")
    constraints.append(cp.max(cp.abs(M_t_K_s)) <= gamma)

    problem = cp.Problem(cp.Minimize(risk), constraints=constraints)

    problem.solve(solver=solver, verbose=verbose)

    return problem, K_hat.value