logistic_regression.py

import numpy as np
from scipy.optimize import minimize
from scipy.special import expit


class LogRegExperiment:
    def __init__(self,beta_true, intercept, coef_index, fit_intercept=True, reg_lambda=1e-5, fisher_n=None):
        self.beta_true = beta_true
        self.intercept = intercept
        self.coef_index = coef_index
        self.fit_intercept = fit_intercept
        self.reg_lambda = reg_lambda
        self.fisher_n=fisher_n
        self.true_sd=None

def _fit_logistic_ridge(X, y, reg_lambda=1.0, fit_intercept=True):
    X = np.asarray(X)
    if X.ndim == 1:
        X = X[:, None]

    n= X.shape[0]
    X1 = np.column_stack([np.ones(n), X]) if fit_intercept else X
    # X1 /= np.sqrt(2)

    def obj(beta):
        eta = X1 @ beta
        loss = np.mean(np.logaddexp(0.0, eta) - y * eta)
        penalty = reg_lambda * np.sum(beta ** 2)
        return loss + penalty

    def grad(beta):
        eta = X1 @ beta
        p = expit(eta)
        g = (X1.T @ (p - y)) / n
        if reg_lambda > 0:
            g += 2.0 * reg_lambda * beta
        return g

    p = X1.shape[1]
    beta0 = np.zeros(p)
    res = minimize(obj, beta0, jac=grad, method="BFGS")
    if not res.success:
        raise RuntimeError(res.message)
    return res.x

def private_logreg_subset_estimator(X, y, coef_index=1, fit_intercept=True, reg_lambda=1.0, noise_func=None):
    beta_hat = _fit_logistic_ridge(X, y, reg_lambda=reg_lambda, fit_intercept=fit_intercept)
    if noise_func is not None:
        beta_hat = beta_hat + noise_func(c=1, size=beta_hat.shape)
    return beta_hat[coef_index]

# def estimate_logreg_on_subsets(X_batches, y_batches, subset_indices, coef_index=1,
#                                fit_intercept=True, reg_lambda=1.0, noise_func=None):
#     B, T, m = subset_indices.shape
#     est = np.empty(B * T, dtype=float)
#
#     flat_idx = subset_indices.reshape(B * T, m)
#
#     for k, idx in enumerate(flat_idx):
#         b = k // T
#         est[k] = private_logreg_subset_estimator( X_batches[b, idx],y_batches[b, idx],coef_index=coef_index,fit_intercept=fit_intercept,
#                                                   reg_lambda=reg_lambda, noise_func=noise_func,)
#
#     return est.reshape(B, T)
#
# def estimate_logreg_center(X_batches, y_batches, coef_index=1, fit_intercept=True, reg_lambda=1.0, noise_func=None):
#     B = X_batches.shape[0]
#     center = np.empty(B, dtype=float)
#
#     for b in range(B):
#         center[b] = private_logreg_subset_estimator(X_batches[b],y_batches[b], coef_index=coef_index,
#             fit_intercept=fit_intercept,reg_lambda=reg_lambda, noise_func=noise_func,)
#     return center


def make_logreg_scalar_estimator(log_exp_setting, noise_func=None):
    reg_lambda, fit_intercept = log_exp_setting.reg_lambda, log_exp_setting.fit_intercept
    def estimator(X, y):
        beta_hat = _fit_logistic_ridge(X, y, reg_lambda=reg_lambda, fit_intercept=fit_intercept)
        if noise_func is not None:
            noise = noise_func(c=1, size=beta_hat.shape)
            beta_hat = beta_hat + noise
        return beta_hat[log_exp_setting.coef_index]

    return estimator


def compute_single_n_logreg(n, num_experiments, algs, log_exp_setting, settings, dist,):
    X_batches, y_batches = generate_logreg_data(num_experiments=num_experiments, n=n, beta=log_exp_setting.beta_true,
                                                intercept=log_exp_setting.intercept, dist=dist)
    params = settings.base_ci_params(n)

    B = X_batches.shape[0]
    results = []

    for alg in algs:
        res = alg.run_3D(X=X_batches, y=y_batches, settings=settings, log_exp_setting=log_exp_setting, B=B, n=n, base_params=params)
        results.append(res)
    return results

def sigmoid(z):
    return 1.0 / (1.0 + np.exp(-z))


def generate_logreg_data(num_experiments, n, beta, intercept, dist):
    X = dist.rvs(size=(num_experiments, n))

    logits = intercept + beta[0]*X
    probs = expit(logits)
    y = np.random.binomial(1, probs)

    return X, y


def full_sample_indices(B, n):
    idx = np.arange(n, dtype=int)[None, None, :]
    return np.repeat(idx, B, axis=0)

def evaluate_estimator_on_index_sets(X_batches, y_batches, index_sets, estimator_fn):
    B, T, m = index_sets.shape
    out = np.empty(B*T, dtype=float)

    flat_idx = index_sets.reshape(B * T, m)

    # map flat job k -> (b, t)
    for k, idx in enumerate(flat_idx):
        b = k // T
        out[k] = estimator_fn(X_batches[b, idx], y_batches[b, idx])

    return out.reshape(B, T)

def evaluate_full_sample_estimator(X_batches, y_batches, estimator_fn):
    B = X_batches.shape[0]
    out = np.empty(B, dtype=float)
    for b in range(B):
        out[b] = estimator_fn(X_batches[b], y_batches[b])
    return out