pretender.py

import numpy as np
from utils import compute_mmd, rbf_kernel


def linear_minimization_oracle(grad, K_val):
    """
    Solve the linear minimization subproblem for the Frank–Wolfe algorithm.

    Given the gradient vector 'grad', the extreme points of the feasible set
    { s ∈ R^(2m) : s_j ∈ [0, 1/K_val] for all j and Σ_j s_j = 1 } take the form
    of assigning 1/K_val to K_val coordinates and 0 to the others.

    Parameters
    ----------
    grad : np.ndarray, shape (d,)
        Current gradient vector.
    K_val : int
        Number of items to select (K).

    Returns
    -------
    s : np.ndarray, shape (d,)
        The solution (extreme point) of the linear subproblem.
    """
    d = grad.shape[0]
    s = np.zeros(d)
    # Select the indices corresponding to the smallest K_val entries of the gradient
    idx = np.argsort(grad)
    s[idx[:K_val]] = 1.0 / K_val
    return s


def frank_wolfe(K_tt, K_st, K_val, L=100):
    """
    Solve the continuous optimization problem using the Frank–Wolfe algorithm.

    The problem is:
      min_{w ∈ F}  f(w) = w^T K_tt w - (2/n_source) * 1^T K_st w,
    where the feasible set F is defined as
      { w ∈ ℝ^(2m) : 0 ≤ w_j ≤ 1/K_val and Σ_j w_j = 1 }.
    The initial solution is set to the uniform distribution, and the step size is chosen as 1/(t+2).

    Parameters
    ----------
    K_tt : np.ndarray, shape (d, d)
        Kernel matrix computed between target candidates. Here, d is the number of candidates (usually 2m).
    K_st : np.ndarray, shape (d, n_source)
        Kernel matrix computed between target candidates and source items.
    K_val : int
        Number of items to select (K).
    L : int, default=100
        Number of iterations for the Frank–Wolfe algorithm.

    Returns
    -------
    w : np.ndarray, shape (d,)
        The final solution of the continuous optimization problem.
    """
    n_source = K_st.shape[1]
    d = K_tt.shape[0]  # Number of candidates (2m)
    # Initialize w as the uniform distribution over candidates
    w = np.ones(d) / d
    # Iterate for L steps
    for t in range(L):
        # Objective function: f(w) = w^T K_tt w - (2/n_source) * (1^T K_st w)
        # Its gradient is: grad = 2 * K_tt w - (2/n_source) * (K_st @ 1)
        grad = 2 * (K_tt @ w) - (2.0 / n_source) * (K_st @ np.ones(n_source))
        # Solve the linear minimization subproblem to obtain s
        s = linear_minimization_oracle(grad, K_val)
        # Step size γ_t = 1/(t+2)
        gamma = 1.0 / (t + 2)
        # Update w ← (1 - γ_t) * w + γ_t * s
        w = (1 - gamma) * w + gamma * s
        # Clip each coordinate to [0, 1/K_val] and renormalize to ensure that sum(w) == 1
        w = np.clip(w, 0, 1.0 / K_val)
        w = w / np.sum(w)
    return w


def random_selection(w, K_val):
    """
    Perform random selection of candidate items based on the continuous solution w using Bernoulli sampling.

    For each candidate j, a Bernoulli trial is performed with probability p_j = K_val * w_j.
    If the total number of selected items does not equal K_val, a greedy procedure is applied to either add or remove items
    until exactly K_val items are selected.

    Parameters
    ----------
    w : np.ndarray, shape (d,)
        The continuous solution, where each coordinate is in [0, 1/K_val] and the sum is 1.
    K_val : int
        Number of items to select (K).

    Returns
    -------
    selected : np.ndarray, shape (d,), boolean
        Boolean mask indicating the final selected candidates.
    """
    d = w.shape[0]
    # For each candidate, perform a Bernoulli trial with probability p_j = K_val * w_j
    selected = np.random.rand(d) < (K_val * w)
    current_count = np.sum(selected)
    # If the number of selected items is less than K_val
    if current_count < K_val:
        # Among the non-selected candidates, choose those with higher w values to add
        not_selected = np.where(~selected)[0]
        sorted_not_selected = not_selected[np.argsort(-w[not_selected])]
        num_needed = K_val - current_count
        selected[sorted_not_selected[:num_needed]] = True
    # If the number of selected items exceeds K_val
    elif current_count > K_val:
        # Among the selected candidates, remove those with smaller w values
        selected_indices = np.where(selected)[0]
        sorted_selected = selected_indices[np.argsort(w[selected_indices])]
        num_to_remove = current_count - K_val
        selected[sorted_selected[:num_to_remove]] = False
    return selected


def pretender_mmd(X_source, X_candidate, K_val, L=100, rounding_trial=1, sigma=1.0, return_continuous=False):
    """
    Main function to implement the MMD-based version of Pretender.

    Given a set of target candidates (each represented by a feature vector) and a set of source items,
    this function first solves the continuous optimization problem via the Frank–Wolfe algorithm,
    and then performs Bernoulli sampling followed by postprocessing to select exactly K_val candidates.

    Parameters
    ----------
    X_source : np.ndarray, shape (n_source, d)
        Source item set.
    X_candidate : np.ndarray, shape (num_candidates, d)
        Target candidate set. Each candidate is represented by its feature vector.
        Typically, the number of candidates is 2m.
    K_val : int
        Number of user interactions (K).
    L : int, default=100
        Number of iterations for the Frank–Wolfe algorithm.
    sigma : float, default=1.0
        Parameter for the RBF kernel.

    Returns
    -------
    selected_mask : np.ndarray, shape (num_candidates,), boolean
        Boolean mask indicating the final selected candidates.
    """
    # Compute the kernel matrix among target candidates, K_tt
    K_tt = rbf_kernel(X_candidate, X_candidate, sigma=sigma)
    # Compute the kernel matrix between target candidates and source items, K_st
    K_st = rbf_kernel(X_candidate, X_source, sigma=sigma)
    n_source = X_source.shape[0]

    # Solve the continuous optimization problem using the Frank–Wolfe algorithm
    w = frank_wolfe(K_tt, K_st, K_val, L=L)

    # For the source measure, assume uniform weights
    w_source = np.ones(n_source) / n_source

    # Compute MMD for the continuous solution:
    # The target measure is given by the continuous weights w on X_candidate.
    mmd_cont = compute_mmd(X_source, X_candidate, w1=w_source, w2=w, sigma=sigma)

    best_mmd = None
    best_selected_mask = None
    for trial in range(rounding_trial):
        # Based on the continuous solution w, perform Bernoulli sampling to discretize the selection
        selected_mask = random_selection(w, K_val)

        # Construct discrete weight vector: for each candidate j, weight = 1/K_val if selected, else 0.
        discrete_w = np.where(selected_mask, 1.0 / K_val, 0.0)

        mmd_disc = compute_mmd(X_source, X_candidate, w1=w_source, w2=discrete_w, sigma=sigma)

        if best_mmd is None or mmd_disc < best_mmd:
            best_mmd = mmd_disc
            best_selected_mask = selected_mask

    print("MMD for continuous solution: ", mmd_cont)
    print("MMD for discrete solution: ", best_mmd)

    if return_continuous:
        return best_selected_mask, w

    return best_selected_mask


# Example usage of the implementation.
if __name__ == "__main__":
    np.random.seed(0)
    # Set the number of target items, m (example: m = 50)
    m = 50
    # Set the dimension of feature vectors, d (example: d = 10)
    d = 10
    # Generate random feature vectors for target items (in actual applications, appropriate features should be used)
    X_target = np.random.randn(m, d)

    # Prepare the candidate set for items.
    # For each target item, two candidates are generated corresponding to label 0 and label 1.
    # In this example, the label information is appended to the feature vector.
    X_candidate_list = []
    for i in range(m):
        # Candidate for label 0: append 0 to the feature vector
        x0 = np.concatenate([X_target[i], np.array([0.0])])
        # Candidate for label 1: append 1 to the feature vector
        x1 = np.concatenate([X_target[i], np.array([1.0])])
        X_candidate_list.append(x0)
        X_candidate_list.append(x1)
    X_candidate = np.vstack(X_candidate_list)  # Shape is (2m, d+1)

    # Generate the source item set (example: n_source = 100)
    n_source = 100
    # For simplicity, source item feature vectors have the same dimension as candidates (d+1)
    X_source = np.random.randn(n_source, d + 1)

    # Set the number of items the user will interact with, K (example: K = 10)
    K_val = 10

    # Execute the Pretender (MMD version) process
    selected_mask = pretender_mmd(X_candidate, X_source, K_val, L=100, sigma=1.0)

    # Retrieve the indices of the selected candidates
    selected_indices = np.where(selected_mask)[0]
    print("Selected candidate indices:", selected_indices)