def simple_tmle()

in causalml/inference/meta/tmle.py [0:0]

• 33 lines of code
• 1 McCabe index (conditional complexity)

def simple_tmle(y, w, q0w, q1w, p, alpha=0.0001):
    """Calculate the ATE and variances with the simplified TMLE method.

    Args:
        y (numpy.array): an outcome vector
        w (numpy.array): a treatment vector
        q0w (numpy.array): an outcome prediction vector given no treatment
        q1w (numpy.array): an outcome prediction vector given treatment
        p (numpy.array): a propensity score vector
        alpha (float, optional): a clipping threshold for predictions

    Returns:
        (tuple)

            - ate (float): ATE
            - se (float): The standard error of ATE
    """
    scaler = MinMaxScaler()
    ystar = scaler.fit_transform(y.reshape(-1, 1)).flatten()

    q0 = np.clip(scaler.transform(q0w.reshape(-1, 1)).flatten(), alpha, 1 - alpha)
    q1 = np.clip(scaler.transform(q1w.reshape(-1, 1)).flatten(), alpha, 1 - alpha)
    qaw = q0 * (1 - w) + q1 * w
    intercept = logit(qaw)

    h1 = w / p
    h0 = (1 - w) / (1 - p)
    sol = minimize(
        logit_tmle,
        np.zeros(2),
        args=(ystar, intercept, h0, h1),
        method="Newton-CG",
        jac=logit_tmle_grad,
        hess=logit_tmle_hess,
    )

    qawstar = scaler.inverse_transform(
        expit(intercept + sol.x[0] * h0 + sol.x[1] * h1).reshape(-1, 1)
    ).flatten()
    q0star = scaler.inverse_transform(
        expit(logit(q0) + sol.x[0] / (1 - p)).reshape(-1, 1)
    ).flatten()
    q1star = scaler.inverse_transform(
        expit(logit(q1) + sol.x[1] / p).reshape(-1, 1)
    ).flatten()

    ic = (
        (w / p - (1 - w) / (1 - p)) * (y - qawstar)
        + q1star
        - q0star
        - np.mean(q1star - q0star)
    )

    return np.mean(q1star - q0star), np.sqrt(np.var(ic) / np.size(y))