Comparing MAF for coordinate-wise subsampling vs record-level subsampling

plrvo.py

import math


def logsumexp(xs):
    m = max(xs)
    return m + math.log(sum(math.exp(x - m) for x in xs))


def log_weight(m, q, eta):
    return (
        math.lgamma(m + 1)
        - math.lgamma(eta + 1)
        - math.lgamma(m - eta + 1)
        + (m - eta) * math.log1p(-q)
        + eta * math.log(q)
    )


def log_F_laplace(a, eta):
    """
    F(a, eta) = E_{Y~Lap(0,1)}[(p_1(Y)/p_0(Y))^eta],
    where p_1 is Lap(a,1).

    Closed form:
      F(a,eta) =
      [ eta exp((eta-1)a) + (eta-1) exp(-eta a) ] / (2 eta - 1).
    """
    if eta in (0, 1):
        return 0.0

    term1 = math.log(eta) + (eta - 1) * a
    term2 = math.log(eta - 1) - eta * a
    return logsumexp([term1, term2]) - math.log(2 * eta - 1)


def alpha_shared_record_subsampling(d, a, q, lam):
    """
    Correct MAF for one shared record-level subsampling event:

      log sum_eta w_eta prod_i F_i(eta).

    Here all d coordinates have the same shift a.
    """
    m = lam + 1
    return logsumexp([
        log_weight(m, q, eta) + d * log_F_laplace(a, eta)
        for eta in range(m + 1)
    ])


def alpha_coordinatewise_subsampling(d, a, q, lam):
    """
    MAF for independent coordinate-level subsampling:

      sum_i log sum_eta w_eta F_i(eta).

    This is the coordinate-summed accountant.
    """
    m = lam + 1
    scalar_alpha = logsumexp([
        log_weight(m, q, eta) + log_F_laplace(a, eta)
        for eta in range(m + 1)
    ])
    return d * scalar_alpha


def laplace_pdf(y, mu):
    return 0.5 * math.exp(-abs(y - mu))


def density_shared(z, q, shift):
    """
    Shared record-level subsampling:
      (1-q) prod_i p0_i + q prod_i p1_i.
    """
    prod0 = math.prod(laplace_pdf(zi, 0.0) for zi in z)
    prod1 = math.prod(laplace_pdf(zi, shift) for zi in z)
    return (1 - q) * prod0 + q * prod1


def density_coordinatewise(z, q, shift):
    """
    Independent coordinate-level subsampling:
      prod_i [(1-q) p0_i + q p1_i].
    """
    return math.prod(
        (1 - q) * laplace_pdf(zi, 0.0) + q * laplace_pdf(zi, shift)
        for zi in z
    )


def main():
    d = 2
    q = 0.1
    lam = 3
    a = 1.0

    joint = alpha_shared_record_subsampling(d, a, q, lam)
    coord = alpha_coordinatewise_subsampling(d, a, q, lam)

    print("MAF comparison")
    print(f"shared record-level accountant:       {joint:.12f}")
    print(f"coordinate-summed accountant:         {coord:.12f}")
    print(f"underestimate from coordinate sum:    {joint - coord:.12f}")

    assert joint > coord

    z = (0.0, 1.0)
    p_shared = density_shared(z, q, shift=1.0)
    p_coord = density_coordinatewise(z, q, shift=1.0)

    print("\nDensity comparison at z=(0,1)")
    print(f"shared record-level density:          {p_shared:.12f}")
    print(f"coordinate-wise subsampling density:  {p_coord:.12f}")

if __name__ == "__main__":
    main()