Streaming Polyphase Factor Decomposition / SKAT

Download original file

This is the original source text. Reading it does not run or verify the program.

Read the typeset paper

1 from __future__ import annotations
2 
3 from dataclasses import dataclass
4 import math
5 from typing import Iterable
6 
7 import numpy as np
8 from scipy import optimize
9 
10 
11 @dataclass(frozen=True)
12 class SpectralPassport:
13     exact: np.ndarray
14     centers: np.ndarray
15     moments: np.ndarray
16     max_abs_delta: np.ndarray
17     min_eigenvalue: np.ndarray
18     max_eigenvalue: np.ndarray
19     order: int
20     original_rank: int
21 
22     @property
23     def packet_count(self) -> int:
24         return int(self.centers.size)
25 
26     @property
27     def exact_count(self) -> int:
28         return int(self.exact.size)
29 
30     @property
31     def lambda_max(self) -> float:
32         vals = []
33         if self.exact.size:
34             vals.append(float(self.exact.max()))
35         if self.max_eigenvalue.size:
36             vals.append(float(self.max_eigenvalue.max()))
37         return max(vals, default=0.0)
38 
39 
40 def compile_centered_passport(
41     eigenvalues: Iterable[float],
42     *,
43     order: int = 8,
44     packet_count: int = 32,
45     exact_top: int = 8,
46     relative_floor: float = 1e-14,
47 ) -> SpectralPassport:
48     lam = np.asarray(tuple(eigenvalues), dtype=np.float64)
49     if lam.ndim != 1 or lam.size == 0:
50         raise ValueError("eigenvalues must be nonempty and one-dimensional")
51     if np.any(lam < 0.0):
52         raise ValueError("eigenvalues must be nonnegative")
53     lam = np.sort(lam)[::-1]
54     if lam[0] <= 0.0:
55         raise ValueError("at least one eigenvalue must be positive")
56     lam = lam[lam > lam[0] * relative_floor]
57 
58     order = int(order)
59     packet_count = int(packet_count)
60     exact_top = min(max(int(exact_top), 0), lam.size)
61     if order < 1 or packet_count < 1:
62         raise ValueError("order and packet_count must be positive")
63 
64     exact = lam[:exact_top].copy()
65     rest = lam[exact_top:]
66     rows = []
67 
68     if rest.size:
69         for block in np.array_split(rest, min(packet_count, rest.size)):
70             if block.size == 0:
71                 continue
72             mu = float(block.mean())
73             delta = block - mu
74             moments = np.empty(order + 1, dtype=np.float64)
75             moments[0] = block.size
76             for n in range(1, order + 1):
77                 moments[n] = np.sum(delta ** n)
78             rows.append(
79                 (
80                     mu,
81                     moments,
82                     float(np.max(np.abs(delta))),
83                     float(block.min()),
84                     float(block.max()),
85                 )
86             )
87 
88     if rows:
89         centers = np.array([x[0] for x in rows], dtype=np.float64)
90         moments = np.stack([x[1] for x in rows])
91         max_abs_delta = np.array([x[2] for x in rows], dtype=np.float64)
92         min_eigenvalue = np.array([x[3] for x in rows], dtype=np.float64)
93         max_eigenvalue = np.array([x[4] for x in rows], dtype=np.float64)
94     else:
95         centers = np.empty(0, dtype=np.float64)
96         moments = np.empty((0, order + 1), dtype=np.float64)
97         max_abs_delta = np.empty(0, dtype=np.float64)
98         min_eigenvalue = np.empty(0, dtype=np.float64)
99         max_eigenvalue = np.empty(0, dtype=np.float64)
100 
101     return SpectralPassport(
102         exact=exact,
103         centers=centers,
104         moments=moments,
105         max_abs_delta=max_abs_delta,
106         min_eigenvalue=min_eigenvalue,
107         max_eigenvalue=max_eigenvalue,
108         order=order,
109         original_rank=int(lam.size),
110     )
111 
112 
113 def exact_cgf(s, eigenvalues):
114     lam = np.asarray(eigenvalues, dtype=np.float64)
115     values = np.atleast_1d(np.asarray(s, dtype=np.complex128))
116     out = -0.5 * np.sum(
117         np.log(1.0 - 2.0 * values[:, None] * lam[None, :]),
118         axis=1,
119     )
120     return out if np.ndim(s) else out[0]
121 
122 
123 def passport_cgf(s, passport: SpectralPassport):
124     values = np.atleast_1d(np.asarray(s, dtype=np.complex128))
125     out = np.zeros(values.shape, dtype=np.complex128)
126 
127     if passport.exact.size:
128         out += -0.5 * np.sum(
129             np.log(1.0 - 2.0 * values[:, None] * passport.exact[None, :]),
130             axis=1,
131         )
132 
133     if passport.centers.size:
134         d = 1.0 - 2.0 * values[:, None] * passport.centers[None, :]
135         alpha = 2.0 * values[:, None] / d
136         out += np.sum(
137             -0.5 * passport.moments[:, 0][None, :] * np.log(d),
138             axis=1,
139         )
140         for n in range(1, passport.order + 1):
141             out += np.sum(
142                 0.5 * alpha**n * passport.moments[:, n][None, :] / n,
143                 axis=1,
144             )
145 
146     return out if np.ndim(s) else out[0]
147 
148 
149 def packet_log_cgf_remainder_bound(s, passport: SpectralPassport):
150     values = np.atleast_1d(np.asarray(s, dtype=np.complex128))
151     if not passport.centers.size:
152         return np.zeros(values.size, dtype=np.float64)
153 
154     d = 1.0 - 2.0 * values[:, None] * passport.centers[None, :]
155     alpha = 2.0 * values[:, None] / d
156     rho = np.abs(alpha) * passport.max_abs_delta[None, :]
157     p = passport.order
158     denom = np.maximum(1.0 - rho, np.finfo(float).tiny)
159 
160     terms = (
161         0.5
162         * passport.moments[:, 0][None, :]
163         * rho ** (p + 1)
164         / ((p + 1) * denom)
165     )
166     bound = np.sum(terms, axis=1)
167     bound[np.any(rho >= 1.0, axis=1)] = np.inf
168     return bound
169 
170 
171 def exact_cgf_derivative(theta, eigenvalues, order):
172     lam = np.asarray(eigenvalues, dtype=np.float64)
173     r = int(order)
174     prefactor = (2 ** (r - 1)) * math.factorial(r - 1)
175     return float(
176         prefactor
177         * np.sum(lam**r / (1.0 - 2.0 * theta * lam) ** r)
178     )
179 
180 
181 def passport_cgf_derivative(theta, passport: SpectralPassport, order):
182     r = int(order)
183     if r < 1:
184         raise ValueError("order must be at least one")
185     prefactor = (2 ** (r - 1)) * math.factorial(r - 1)
186     total = 0.0
187 
188     if passport.exact.size:
189         total += prefactor * np.sum(
190             passport.exact**r
191             / (1.0 - 2.0 * theta * passport.exact) ** r
192         )
193 
194     if passport.centers.size:
195         mu = passport.centers
196         d = 1.0 - 2.0 * theta * mu
197         alpha = 2.0 * theta / d
198         expansion = np.zeros_like(mu)
199 
200         for n in range(passport.order + 1):
201             coefficient = np.zeros_like(mu)
202             for j in range(0, min(r, n) + 1):
203                 k = n - j
204                 coefficient += (
205                     math.comb(r, j)
206                     * mu ** (r - j)
207                     * math.comb(r + k - 1, k)
208                     * alpha**k
209                 )
210             expansion += coefficient * passport.moments[:, n]
211 
212         total += prefactor * np.sum(expansion / d**r)
213 
214     return float(total)
215 
216 
217 def exact_saddlepoint(q, eigenvalues):
218     lam = np.asarray(eigenvalues, dtype=np.float64)
219     mean = exact_cgf_derivative(0.0, lam, 1)
220     if q <= mean:
221         return 0.0
222     upper = (1.0 / (2.0 * float(lam.max()))) * (1.0 - 1e-12)
223     return float(
224         optimize.brentq(
225             lambda t: exact_cgf_derivative(t, lam, 1) - q,
226             0.0,
227             upper,
228             xtol=1e-14,
229             rtol=1e-13,
230         )
231     )
232 
233 
234 def passport_saddlepoint(q, passport: SpectralPassport):
235     mean = passport_cgf_derivative(0.0, passport, 1)
236     if q <= mean:
237         return 0.0
238     upper = (1.0 / (2.0 * passport.lambda_max)) * (1.0 - 1e-10)
239     return float(
240         optimize.brentq(
241             lambda t: passport_cgf_derivative(t, passport, 1) - q,
242             0.0,
243             upper,
244             xtol=1e-13,
245             rtol=1e-12,
246         )
247     )
248 
249 
250 def standard_cos_upper_tail(q, eigenvalues, *, K=256, L=12.0):
251     """Legacy-style global COS CDF followed by 1-F(q)."""
252     lam = np.asarray(eigenvalues, dtype=np.float64)
253     c1 = exact_cgf_derivative(0.0, lam, 1)
254     c2 = exact_cgf_derivative(0.0, lam, 2)
255     c4 = exact_cgf_derivative(0.0, lam, 4)
256     width = L * math.sqrt(c2 + math.sqrt(c4))
257     a = max(0.0, c1 - width)
258     b = c1 + width
259 
260     u = np.arange(K, dtype=np.float64) * math.pi / (b - a)
261     phi = np.exp(exact_cgf(1j * u, lam))
262     coeff = 2.0 / (b - a) * np.real(phi * np.exp(-1j * u * a))
263     coeff[0] *= 0.5
264 
265     integ = np.empty_like(u)
266     integ[0] = q - a
267     integ[1:] = np.sin(u[1:] * (q - a)) / u[1:]
268     return float(1.0 - np.sum(coeff * integ))
269 
270 
271 def tilted_cos_upper_tail_exact(q, eigenvalues, *, K=128, L=12.0):
272     lam = np.asarray(eigenvalues, dtype=np.float64)
273     theta = exact_saddlepoint(q, lam)
274     if theta <= 0.0:
275         raise ValueError("q must exceed the mean for this tilted path")
276 
277     c2 = exact_cgf_derivative(theta, lam, 2)
278     c4 = exact_cgf_derivative(theta, lam, 4)
279     width = L * math.sqrt(c2 + math.sqrt(c4))
280     a = max(0.0, q - width)
281     b = q + width
282 
283     u = np.arange(K, dtype=np.float64) * math.pi / (b - a)
284     log_m = exact_cgf(theta, lam)
285     phi = np.exp(exact_cgf(theta + 1j * u, lam) - log_m)
286 
287     coeff = 2.0 / (b - a) * np.real(phi * np.exp(-1j * u * a))
288     coeff[0] *= 0.5
289 
290     ell = b - q
291     payoff_integral = np.real(
292         np.exp(1j * u * (q - a))
293         * (1.0 - np.exp((-theta + 1j * u) * ell))
294         / (theta - 1j * u)
295     )
296     tail = (
297         math.exp(float(np.real(log_m)) - theta * q)
298         * float(np.sum(coeff * payoff_integral))
299     )
300     return float(tail), {"theta": float(theta), "a": float(a), "b": float(b)}
301 
302 
303 def tilted_cos_upper_tail_passport(q, passport: SpectralPassport, *, K=128, L=12.0):
304     theta = passport_saddlepoint(q, passport)
305     if theta <= 0.0:
306         raise ValueError("q must exceed the mean for this tilted path")
307 
308     c2 = passport_cgf_derivative(theta, passport, 2)
309     c4 = passport_cgf_derivative(theta, passport, 4)
310     width = L * math.sqrt(c2 + math.sqrt(max(c4, 0.0)))
311     a = max(0.0, q - width)
312     b = q + width
313 
314     u = np.arange(K, dtype=np.float64) * math.pi / (b - a)
315     s_grid = theta + 1j * u
316     log_m = passport_cgf(theta, passport)
317     phi = np.exp(passport_cgf(s_grid, passport) - log_m)
318 
319     coeff = 2.0 / (b - a) * np.real(phi * np.exp(-1j * u * a))
320     coeff[0] *= 0.5
321 
322     ell = b - q
323     payoff_integral = np.real(
324         np.exp(1j * u * (q - a))
325         * (1.0 - np.exp((-theta + 1j * u) * ell))
326         / (theta - 1j * u)
327     )
328     tail = (
329         math.exp(float(np.real(log_m)) - theta * q)
330         * float(np.sum(coeff * payoff_integral))
331     )
332 
333     remainder = packet_log_cgf_remainder_bound(s_grid, passport)
334     if passport.centers.size:
335         d = 1.0 - 2.0 * s_grid[:, None] * passport.centers[None, :]
336         rho = (
337             np.abs(2.0 * s_grid[:, None] / d)
338             * passport.max_abs_delta[None, :]
339         )
340         max_rho = float(np.max(rho))
341     else:
342         max_rho = 0.0
343 
344     return float(tail), {
345         "theta": float(theta),
346         "a": float(a),
347         "b": float(b),
348         "max_packet_rho": max_rho,
349         "max_log_cgf_remainder_bound": float(np.max(remainder)),
350         "packet_count": passport.packet_count,
351         "exact_count": passport.exact_count,
352         "original_rank": passport.original_rank,
353     }

Back to source controls

Back to Streaming Polyphase Factor Decomposition / SKAT