Optimal-Schemes-for-Discrete-Distribution-Estimation-under-Locally-Differential-Privacy-/data.py at main · LXRbecky/Optimal-Schemes-for-Discrete-Distribution-Estimation-under-Locally-Differential-Privacy- · GitHub

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
import numpy as np


# =========================================================
# normalization helper
# =========================================================
def normalize(P):
    P = np.array(P, dtype=float)
    P = np.maximum(P, 0)
    return P / P.sum()


# =========================================================
# Bernoulli-style random simplex (Dirichlet)
# =========================================================
def bernoulli_dist(k):
    return normalize(np.random.dirichlet(np.ones(k)))


# =========================================================
# Binomial distribution (discretized)
# =========================================================
def binomial_dist(k, n_bin=10, p=0.5):
    from math import comb

    P = np.zeros(k)

    for i in range(k):
        if i <= n_bin:
            P[i] = comb(n_bin, i) * (p ** i) * ((1 - p) ** (n_bin - i))

    return normalize(P)


# =========================================================
# Truncated Poisson
# =========================================================
def truncated_poisson(k, lam=3.0):

    P = np.zeros(k)

    for i in range(k):
        P[i] = np.exp(-lam) * lam ** i / np.math.factorial(i)

    return normalize(P)


# =========================================================
# Geometric distribution
# =========================================================
def geometric_dist(k, p=0.4):

    P = np.array([(1 - p) ** i * p for i in range(k)])

    return normalize(P)


# =========================================================
# NEW: structured adversarial distribution
# =========================================================
def structured_dist(k, delta=0.2):

    """
    Your definition:

    Even k:
        V = {-1,1}^{k/2}
        p_v = p_U + (delta/k)[v, -v]

    Odd k:
        V = {-1,1}^{(k-1)/2}
        p_v = 1/(k-1)[1,...,1,0] + (delta/(k-1))[v,-v,0]
    """

    if k < 2:
        raise ValueError("k must be >= 2")

    # uniform base
    p = np.ones(k) / k

    if k % 2 == 0:
        m = k // 2

        v = np.random.choice([-1, 1], size=m)

        perturb = np.concatenate([v, -v])  # length k

        P = p + (delta / k) * perturb

    else:
        m = (k - 1) // 2

        v = np.random.choice([-1, 1], size=m)

        perturb = np.concatenate([v, -v, [0]])  # length k

        base = np.concatenate([np.ones(k - 1), [0]]) / (k - 1)

        P = base + (delta / (k - 1)) * perturb

    return normalize(P)


# =========================================================
# main interface
# =========================================================
def generate_distribution(k, dist_type="bernoulli"):

    if dist_type == "bernoulli":
        return bernoulli_dist(k)

    elif dist_type == "binomial":
        return binomial_dist(k)

    elif dist_type == "poisson":
        return truncated_poisson(k)

    elif dist_type == "geometric":
        return geometric_dist(k)

    elif dist_type == "structured":
        return structured_dist(k)

    else:
        raise ValueError("Unknown distribution type")