refactor: optim loss of MultivariateFailure

skmodel.MultivariateFailure
Implement vectorized calculations to replace nested loops for efficiency.
Utilize logsumexp for better numerical stability in exponential calculations.
test code
import numpy as np
import jax.numpy as jnp
from jax.scipy.special import logsumexp
import time

def multivariate_failure_objective(params, X, y, delta, n, K):
    Xbeta = jnp.matmul(X, params)
    tmp = jnp.ones((n, K))
    for i in range(n):
        for k in range(K):
            tmp = tmp.at[i, k].set(Xbeta[i] - jnp.log(jnp.matmul(y[:, k] >= y[i, k], jnp.exp(Xbeta))))
    loss = -jnp.mean(tmp * delta)
    return loss

def multivariate_failure_objective_vectorized_logsumexp(params, X, y, delta, n, K):
    Xbeta_expanded = jnp.matmul(X, params)[:, None]
    sum_exp_Xbeta = logsumexp(Xbeta_expanded + jnp.log(y >= y[:, None, :]), axis=1)
    loss = -jnp.mean((Xbeta_expanded - sum_exp_Xbeta) * delta)
    return loss

def make_Clayton2_data(n, theta=15, lambda1=1, lambda2=1, c1=1, c2=1):
    u1 = np.random.uniform(0, 1, n)
    u2 = np.random.uniform(0, 1, n)
    time2 = -np.log(1 - u2) / lambda2
    time1 = (np.log(1 - np.power((1 - u2), -theta) + np.power((1 - u1), -theta / (1 + theta)) * np.power((1 - u2), -theta)) / theta / lambda1)
    ctime1 = np.random.uniform(0, c1, n)
    ctime2 = np.random.uniform(0, c2, n)
    delta1 = (time1 < ctime1) * 1
    delta2 = (time2 < ctime2) * 1
    time1 = np.minimum(time1, ctime1)
    time2 = np.minimum(time2, ctime2)
    y = np.hstack((time1.reshape((-1, 1)), time2.reshape((-1, 1))))
    delta = np.hstack((delta1.reshape((-1, 1)), delta2.reshape((-1, 1))))
    return y, delta

def test(seed):
    np.random.seed(seed)
    n, p, s, rho = 100, 100, 10, 0.5
    K = 2

    beta = np.zeros(p)
    beta[:s] = 5
    Sigma = np.power(rho, np.abs(np.linspace(1, p, p) - np.linspace(1, p, p).reshape(p, 1)))
    X = np.random.multivariate_normal(mean=np.zeros(p), cov=Sigma, size=(n,))
    lambda1 = 1 * np.exp(np.matmul(X, beta))
    lambda2 = 10 * np.exp(np.matmul(X, beta))

    y, delta = make_Clayton2_data(n, theta=50, lambda1=lambda1, lambda2=lambda2, c1=5, c2=5)

    # Convert numpy arrays to jax numpy arrays
    X_jax = jnp.array(X)
    y_jax = jnp.array(y)
    delta_jax = jnp.array(delta)

    # Generate random parameters for testing
    params = jnp.array(np.random.randn(p))

    # Calculate loss using both methods
    t1 = time.time()
    loss_original = multivariate_failure_objective(params, X_jax, y_jax, delta_jax, n, K)
    t2 = time.time()
    loss_vectorized = multivariate_failure_objective_vectorized_logsumexp(params, X_jax, y_jax, delta_jax, n, K)
    t3 = time.time()

    return loss_original, t2 - t1, loss_vectorized, t3 - t2

if __name__ == "__main__":
    for i in range(10):
        loss_original, time_original, loss_vectorized, time_vectorized = test(i)
        print("loss_original:   ", loss_original, "time_original: ", time_original)
        print("loss_vectorized: ", loss_vectorized, "time_vectorized: ", time_vectorized)
abess-team / skscope

refactor: optim loss of MultivariateFailure #75

Codecov Report
