freealg 0.7.11__py3-none-any.whl → 0.7.14__py3-none-any.whl
This diff represents the content of publicly available package versions that have been released to one of the supported registries. The information contained in this diff is provided for informational purposes only and reflects changes between package versions as they appear in their respective public registries.
- freealg/__init__.py +2 -2
- freealg/__version__.py +1 -1
- freealg/_algebraic_form/__init__.py +2 -1
- freealg/_algebraic_form/_constraints.py +53 -12
- freealg/_algebraic_form/_cusp.py +357 -0
- freealg/_algebraic_form/_cusp_wrap.py +268 -0
- freealg/_algebraic_form/_decompress.py +330 -381
- freealg/_algebraic_form/_decompress2.py +120 -0
- freealg/_algebraic_form/_decompress4.py +739 -0
- freealg/_algebraic_form/_decompress5.py +738 -0
- freealg/_algebraic_form/_decompress6.py +492 -0
- freealg/_algebraic_form/_decompress7.py +355 -0
- freealg/_algebraic_form/_decompress8.py +369 -0
- freealg/_algebraic_form/_decompress9.py +363 -0
- freealg/_algebraic_form/_decompress_new.py +431 -0
- freealg/_algebraic_form/_decompress_new_2.py +1631 -0
- freealg/_algebraic_form/_decompress_util.py +172 -0
- freealg/_algebraic_form/_edge.py +46 -68
- freealg/_algebraic_form/_homotopy.py +62 -30
- freealg/_algebraic_form/_homotopy2.py +289 -0
- freealg/_algebraic_form/_homotopy3.py +215 -0
- freealg/_algebraic_form/_homotopy4.py +320 -0
- freealg/_algebraic_form/_homotopy5.py +185 -0
- freealg/_algebraic_form/_moments.py +43 -57
- freealg/_algebraic_form/_support.py +132 -177
- freealg/_algebraic_form/algebraic_form.py +163 -30
- freealg/distributions/__init__.py +3 -1
- freealg/distributions/_compound_poisson.py +464 -0
- freealg/distributions/_deformed_marchenko_pastur.py +51 -0
- freealg/distributions/_deformed_wigner.py +44 -0
- {freealg-0.7.11.dist-info → freealg-0.7.14.dist-info}/METADATA +2 -1
- {freealg-0.7.11.dist-info → freealg-0.7.14.dist-info}/RECORD +36 -20
- {freealg-0.7.11.dist-info → freealg-0.7.14.dist-info}/WHEEL +1 -1
- {freealg-0.7.11.dist-info → freealg-0.7.14.dist-info}/licenses/AUTHORS.txt +0 -0
- {freealg-0.7.11.dist-info → freealg-0.7.14.dist-info}/licenses/LICENSE.txt +0 -0
- {freealg-0.7.11.dist-info → freealg-0.7.14.dist-info}/top_level.txt +0 -0
|
@@ -0,0 +1,739 @@
|
|
|
1
|
+
# SPDX-FileCopyrightText: Copyright 2026, Siavash Ameli <sameli@berkeley.edu>
|
|
2
|
+
# SPDX-License-Identifier: BSD-3-Clause
|
|
3
|
+
# SPDX-FileType: SOURCE
|
|
4
|
+
#
|
|
5
|
+
# This program is free software: you can redistribute it and/or modify it under
|
|
6
|
+
# the terms of the license found in the LICENSE.txt file in the root directory
|
|
7
|
+
# of this source tree.
|
|
8
|
+
|
|
9
|
+
|
|
10
|
+
# =======
|
|
11
|
+
# Imports
|
|
12
|
+
# =======
|
|
13
|
+
|
|
14
|
+
import numpy
|
|
15
|
+
from ._continuation_algebraic import powers
|
|
16
|
+
|
|
17
|
+
__all__ = ['decompress_newton']
|
|
18
|
+
|
|
19
|
+
|
|
20
|
+
|
|
21
|
+
# ===============
|
|
22
|
+
# eval P partials
|
|
23
|
+
# ===============
|
|
24
|
+
|
|
25
|
+
def eval_P_partials(z, m, a_coeffs):
|
|
26
|
+
"""
|
|
27
|
+
Evaluate P(z,m) and its partial derivatives dP/dz and dP/dm.
|
|
28
|
+
|
|
29
|
+
This assumes P is represented by `a_coeffs` in the monomial basis
|
|
30
|
+
|
|
31
|
+
P(z, m) = sum_{j=0..s} a_j(z) * m^j,
|
|
32
|
+
a_j(z) = sum_{i=0..deg_z} a_coeffs[i, j] * z^i.
|
|
33
|
+
|
|
34
|
+
The function returns P, dP/dz, dP/dm with broadcasting over z and m.
|
|
35
|
+
|
|
36
|
+
Parameters
|
|
37
|
+
----------
|
|
38
|
+
z : complex or array_like of complex
|
|
39
|
+
First argument to P.
|
|
40
|
+
m : complex or array_like of complex
|
|
41
|
+
Second argument to P. Must be broadcast-compatible with `z`.
|
|
42
|
+
a_coeffs : ndarray, shape (deg_z+1, s+1)
|
|
43
|
+
Coefficient matrix for P in the monomial basis.
|
|
44
|
+
|
|
45
|
+
Returns
|
|
46
|
+
-------
|
|
47
|
+
P : complex or ndarray of complex
|
|
48
|
+
Value P(z,m).
|
|
49
|
+
Pz : complex or ndarray of complex
|
|
50
|
+
Partial derivative dP/dz evaluated at (z,m).
|
|
51
|
+
Pm : complex or ndarray of complex
|
|
52
|
+
Partial derivative dP/dm evaluated at (z,m).
|
|
53
|
+
|
|
54
|
+
Notes
|
|
55
|
+
-----
|
|
56
|
+
For scalar (z,m), this uses Horner evaluation for a_j(z) and then Horner
|
|
57
|
+
in m. For array inputs, it uses precomputed power tables via `_powers` for
|
|
58
|
+
simplicity.
|
|
59
|
+
|
|
60
|
+
Examples
|
|
61
|
+
--------
|
|
62
|
+
.. code-block:: python
|
|
63
|
+
|
|
64
|
+
P, Pz, Pm = eval_P_partials(1.0 + 1j, 0.2 + 0.3j, a_coeffs)
|
|
65
|
+
"""
|
|
66
|
+
|
|
67
|
+
z = numpy.asarray(z, dtype=complex)
|
|
68
|
+
m = numpy.asarray(m, dtype=complex)
|
|
69
|
+
|
|
70
|
+
deg_z = int(a_coeffs.shape[0] - 1)
|
|
71
|
+
s = int(a_coeffs.shape[1] - 1)
|
|
72
|
+
|
|
73
|
+
if (z.ndim == 0) and (m.ndim == 0):
|
|
74
|
+
zz = complex(z)
|
|
75
|
+
mm = complex(m)
|
|
76
|
+
|
|
77
|
+
a = numpy.empty(s + 1, dtype=complex)
|
|
78
|
+
ap = numpy.empty(s + 1, dtype=complex)
|
|
79
|
+
|
|
80
|
+
for j in range(s + 1):
|
|
81
|
+
c = a_coeffs[:, j]
|
|
82
|
+
|
|
83
|
+
val = 0.0 + 0.0j
|
|
84
|
+
for i in range(deg_z, -1, -1):
|
|
85
|
+
val = val * zz + c[i]
|
|
86
|
+
a[j] = val
|
|
87
|
+
|
|
88
|
+
dval = 0.0 + 0.0j
|
|
89
|
+
for i in range(deg_z, 0, -1):
|
|
90
|
+
dval = dval * zz + (i * c[i])
|
|
91
|
+
ap[j] = dval
|
|
92
|
+
|
|
93
|
+
p = a[s]
|
|
94
|
+
pm = 0.0 + 0.0j
|
|
95
|
+
for j in range(s - 1, -1, -1):
|
|
96
|
+
pm = pm * mm + p
|
|
97
|
+
p = p * mm + a[j]
|
|
98
|
+
|
|
99
|
+
pz = ap[s]
|
|
100
|
+
for j in range(s - 1, -1, -1):
|
|
101
|
+
pz = pz * mm + ap[j]
|
|
102
|
+
|
|
103
|
+
return p, pz, pm
|
|
104
|
+
|
|
105
|
+
shp = numpy.broadcast(z, m).shape
|
|
106
|
+
zz = numpy.broadcast_to(z, shp).ravel()
|
|
107
|
+
mm = numpy.broadcast_to(m, shp).ravel()
|
|
108
|
+
|
|
109
|
+
zp = powers(zz, deg_z)
|
|
110
|
+
mp = powers(mm, s)
|
|
111
|
+
|
|
112
|
+
dzp = numpy.zeros_like(zp)
|
|
113
|
+
for i in range(1, deg_z + 1):
|
|
114
|
+
dzp[:, i] = i * zp[:, i - 1]
|
|
115
|
+
|
|
116
|
+
P = numpy.zeros(zz.size, dtype=complex)
|
|
117
|
+
Pz = numpy.zeros(zz.size, dtype=complex)
|
|
118
|
+
Pm = numpy.zeros(zz.size, dtype=complex)
|
|
119
|
+
|
|
120
|
+
for j in range(s + 1):
|
|
121
|
+
aj = zp @ a_coeffs[:, j]
|
|
122
|
+
P += aj * mp[:, j]
|
|
123
|
+
|
|
124
|
+
ajp = dzp @ a_coeffs[:, j]
|
|
125
|
+
Pz += ajp * mp[:, j]
|
|
126
|
+
|
|
127
|
+
if j >= 1:
|
|
128
|
+
Pm += (j * aj) * mp[:, j - 1]
|
|
129
|
+
|
|
130
|
+
return P.reshape(shp), Pz.reshape(shp), Pm.reshape(shp)
|
|
131
|
+
|
|
132
|
+
|
|
133
|
+
# ==========
|
|
134
|
+
# fd solve w
|
|
135
|
+
# ==========
|
|
136
|
+
|
|
137
|
+
# def fd_solve_w(z, t, a_coeffs, w_init, max_iter=50, tol=1e-12,
|
|
138
|
+
# armijo=1e-4, min_lam=1e-6, w_min=1e-14):
|
|
139
|
+
# """
|
|
140
|
+
# Solve for w = m(t,z) from the implicit FD equation using damped Newton.
|
|
141
|
+
#
|
|
142
|
+
# We solve in w the equation
|
|
143
|
+
#
|
|
144
|
+
# F(w) = P(z + alpha/w, tau*w) = 0,
|
|
145
|
+
#
|
|
146
|
+
# where tau = exp(t) and alpha = 1 - 1/tau.
|
|
147
|
+
#
|
|
148
|
+
# A backtracking (Armijo) line search is used to stabilize Newton updates.
|
|
149
|
+
# When Im(z) > 0, the iterate is constrained to remain in the upper
|
|
150
|
+
# half-plane (Im(w) > 0), enforcing the Herglotz branch.
|
|
151
|
+
#
|
|
152
|
+
# Parameters
|
|
153
|
+
# ----------
|
|
154
|
+
# z : complex
|
|
155
|
+
# Query point in the complex plane.
|
|
156
|
+
# t : float
|
|
157
|
+
# Time parameter (tau = exp(t)).
|
|
158
|
+
# a_coeffs : ndarray
|
|
159
|
+
# Coefficients defining P(zeta,y) in the monomial basis.
|
|
160
|
+
# w_init : complex
|
|
161
|
+
# Initial guess for w.
|
|
162
|
+
# max_iter : int, optional
|
|
163
|
+
# Maximum number of Newton iterations.
|
|
164
|
+
# tol : float, optional
|
|
165
|
+
# Residual tolerance on |F(w)|.
|
|
166
|
+
# armijo : float, optional
|
|
167
|
+
# Armijo parameter for backtracking sufficient decrease.
|
|
168
|
+
# min_lam : float, optional
|
|
169
|
+
# Minimum damping factor allowed in backtracking.
|
|
170
|
+
# w_min : float, optional
|
|
171
|
+
# Minimum |w| allowed to avoid singularity in z + alpha/w.
|
|
172
|
+
#
|
|
173
|
+
# Returns
|
|
174
|
+
# -------
|
|
175
|
+
# w : complex
|
|
176
|
+
# The computed solution (last iterate if not successful).
|
|
177
|
+
# success : bool
|
|
178
|
+
# True if convergence criteria were met, False otherwise.
|
|
179
|
+
#
|
|
180
|
+
# Notes
|
|
181
|
+
# -----
|
|
182
|
+
# This function does not choose the correct branch globally by itself; it
|
|
183
|
+
# relies on a good initialization strategy (e.g. time continuation and/or
|
|
184
|
+
# x-sweeps) to avoid converging to a different valid root of the implicit
|
|
185
|
+
# equation.
|
|
186
|
+
#
|
|
187
|
+
# Examples
|
|
188
|
+
# --------
|
|
189
|
+
# .. code-block:: python
|
|
190
|
+
#
|
|
191
|
+
# w, ok = fd_solve_w(
|
|
192
|
+
# z=0.5 + 1e-6j, t=2.0, a_coeffs=a_coeffs, w_init=m1_fn(0.5 + 1e-6j),
|
|
193
|
+
# max_iter=50, tol=1e-12
|
|
194
|
+
# )
|
|
195
|
+
# """
|
|
196
|
+
#
|
|
197
|
+
# z = complex(z)
|
|
198
|
+
# w = complex(w_init)
|
|
199
|
+
#
|
|
200
|
+
# tau = float(numpy.exp(t))
|
|
201
|
+
# alpha = 1.0 - 1.0 / tau
|
|
202
|
+
#
|
|
203
|
+
# want_pos_imag = (z.imag > 0.0)
|
|
204
|
+
#
|
|
205
|
+
# for _ in range(max_iter):
|
|
206
|
+
# if not numpy.isfinite(w.real) or not numpy.isfinite(w.imag):
|
|
207
|
+
# return w, False
|
|
208
|
+
# if abs(w) < w_min:
|
|
209
|
+
# return w, False
|
|
210
|
+
# if want_pos_imag and (w.imag <= 0.0):
|
|
211
|
+
# return w, False
|
|
212
|
+
#
|
|
213
|
+
# zeta = z + alpha / w
|
|
214
|
+
# y = tau * w
|
|
215
|
+
#
|
|
216
|
+
# F, Pz, Py = eval_P_partials(zeta, y, a_coeffs)
|
|
217
|
+
# F = complex(F)
|
|
218
|
+
# Pz = complex(Pz)
|
|
219
|
+
# Py = complex(Py)
|
|
220
|
+
#
|
|
221
|
+
# if abs(F) <= tol:
|
|
222
|
+
# return w, True
|
|
223
|
+
#
|
|
224
|
+
# dF = (-alpha / (w * w)) * Pz + tau * Py
|
|
225
|
+
# if dF == 0.0:
|
|
226
|
+
# return w, False
|
|
227
|
+
#
|
|
228
|
+
# step = -F / dF
|
|
229
|
+
#
|
|
230
|
+
# lam = 1.0
|
|
231
|
+
# F_abs = abs(F)
|
|
232
|
+
# ok = False
|
|
233
|
+
#
|
|
234
|
+
# while lam >= min_lam:
|
|
235
|
+
# w_new = w + lam * step
|
|
236
|
+
# if abs(w_new) < w_min:
|
|
237
|
+
# lam *= 0.5
|
|
238
|
+
# continue
|
|
239
|
+
# if want_pos_imag and (w_new.imag <= 0.0):
|
|
240
|
+
# lam *= 0.5
|
|
241
|
+
# continue
|
|
242
|
+
#
|
|
243
|
+
# zeta_new = z + alpha / w_new
|
|
244
|
+
# y_new = tau * w_new
|
|
245
|
+
#
|
|
246
|
+
# F_new = eval_P_partials(zeta_new, y_new, a_coeffs)[0]
|
|
247
|
+
# F_new = complex(F_new)
|
|
248
|
+
#
|
|
249
|
+
# if abs(F_new) <= (1.0 - armijo * lam) * F_abs:
|
|
250
|
+
# w = w_new
|
|
251
|
+
# ok = True
|
|
252
|
+
# break
|
|
253
|
+
#
|
|
254
|
+
# lam *= 0.5
|
|
255
|
+
#
|
|
256
|
+
# if not ok:
|
|
257
|
+
# return w, False
|
|
258
|
+
#
|
|
259
|
+
# F_end = eval_P_partials(z + alpha / w, tau * w, a_coeffs)[0]
|
|
260
|
+
# return w, (abs(F_end) <= 10.0 * tol)
|
|
261
|
+
|
|
262
|
+
def fd_solve_w(z, t, a_coeffs, w_init, max_iter=50, tol=1e-12,
|
|
263
|
+
armijo=1e-4, min_lam=1e-6, w_min=1e-14):
|
|
264
|
+
"""
|
|
265
|
+
Solve for w = m(t,z) from the implicit FD equation using damped Newton.
|
|
266
|
+
|
|
267
|
+
We solve in w the equation
|
|
268
|
+
|
|
269
|
+
F(w) = P(z + alpha/w, tau*w) = 0,
|
|
270
|
+
|
|
271
|
+
where tau = exp(t) and alpha = 1 - 1/tau.
|
|
272
|
+
|
|
273
|
+
A backtracking (Armijo) line search is used to stabilize Newton updates.
|
|
274
|
+
When Im(z) > 0, the iterate is constrained to remain in the upper
|
|
275
|
+
half-plane (Im(w) > 0), enforcing the Herglotz branch.
|
|
276
|
+
|
|
277
|
+
Parameters
|
|
278
|
+
----------
|
|
279
|
+
z : complex
|
|
280
|
+
Query point in the complex plane.
|
|
281
|
+
t : float
|
|
282
|
+
Time parameter (tau = exp(t)).
|
|
283
|
+
a_coeffs : ndarray
|
|
284
|
+
Coefficients defining P(zeta,y) in the monomial basis.
|
|
285
|
+
w_init : complex
|
|
286
|
+
Initial guess for w.
|
|
287
|
+
max_iter : int, optional
|
|
288
|
+
Maximum number of Newton iterations.
|
|
289
|
+
tol : float, optional
|
|
290
|
+
Residual tolerance on |F(w)|.
|
|
291
|
+
armijo : float, optional
|
|
292
|
+
Armijo parameter for backtracking sufficient decrease.
|
|
293
|
+
min_lam : float, optional
|
|
294
|
+
Minimum damping factor allowed in backtracking.
|
|
295
|
+
w_min : float, optional
|
|
296
|
+
Minimum |w| allowed to avoid singularity in z + alpha/w.
|
|
297
|
+
|
|
298
|
+
Returns
|
|
299
|
+
-------
|
|
300
|
+
w : complex
|
|
301
|
+
The computed solution (last iterate if not successful).
|
|
302
|
+
success : bool
|
|
303
|
+
True if convergence criteria were met, False otherwise.
|
|
304
|
+
|
|
305
|
+
Notes
|
|
306
|
+
-----
|
|
307
|
+
This function does not choose the correct branch globally by itself; it
|
|
308
|
+
relies on a good initialization strategy (e.g. time continuation and/or
|
|
309
|
+
x-sweeps) to avoid converging to a different valid root of the implicit
|
|
310
|
+
equation.
|
|
311
|
+
|
|
312
|
+
Examples
|
|
313
|
+
--------
|
|
314
|
+
.. code-block:: python
|
|
315
|
+
|
|
316
|
+
w, ok = fd_solve_w(
|
|
317
|
+
z=0.5 + 1e-6j, t=2.0, a_coeffs=a_coeffs, w_init=m1_fn(0.5 + 1e-6j),
|
|
318
|
+
max_iter=50, tol=1e-12
|
|
319
|
+
)
|
|
320
|
+
"""
|
|
321
|
+
|
|
322
|
+
z = complex(z)
|
|
323
|
+
w = complex(w_init)
|
|
324
|
+
|
|
325
|
+
tau = float(numpy.exp(t))
|
|
326
|
+
alpha = 1.0 - 1.0 / tau
|
|
327
|
+
|
|
328
|
+
want_pos_imag = (z.imag > 0.0)
|
|
329
|
+
|
|
330
|
+
for _ in range(max_iter):
|
|
331
|
+
|
|
332
|
+
# ----------------
|
|
333
|
+
|
|
334
|
+
# if not numpy.isfinite(w.real) or not numpy.isfinite(w.imag):
|
|
335
|
+
# return w, False
|
|
336
|
+
# if abs(w) < w_min:
|
|
337
|
+
# return w, False
|
|
338
|
+
# if want_pos_imag and (w.imag <= 0.0):
|
|
339
|
+
# return w, False
|
|
340
|
+
#
|
|
341
|
+
# zeta = z + alpha / w
|
|
342
|
+
# y = tau * w
|
|
343
|
+
#
|
|
344
|
+
# F, Pz, Py = eval_P_partials(zeta, y, a_coeffs)
|
|
345
|
+
# F = complex(F)
|
|
346
|
+
# Pz = complex(Pz)
|
|
347
|
+
# Py = complex(Py)
|
|
348
|
+
#
|
|
349
|
+
# if abs(F) <= tol:
|
|
350
|
+
# return w, True
|
|
351
|
+
#
|
|
352
|
+
# dF = (-alpha / (w * w)) * Pz + tau * Py
|
|
353
|
+
# if dF == 0.0:
|
|
354
|
+
# return w, False
|
|
355
|
+
#
|
|
356
|
+
# step = -F / dF
|
|
357
|
+
#
|
|
358
|
+
# lam = 1.0
|
|
359
|
+
# F_abs = abs(F)
|
|
360
|
+
# ok = False
|
|
361
|
+
#
|
|
362
|
+
# while lam >= min_lam:
|
|
363
|
+
# w_new = w + lam * step
|
|
364
|
+
# if abs(w_new) < w_min:
|
|
365
|
+
# lam *= 0.5
|
|
366
|
+
# continue
|
|
367
|
+
# if want_pos_imag and (w_new.imag <= 0.0):
|
|
368
|
+
# lam *= 0.5
|
|
369
|
+
# continue
|
|
370
|
+
#
|
|
371
|
+
# zeta_new = z + alpha / w_new
|
|
372
|
+
# y_new = tau * w_new
|
|
373
|
+
#
|
|
374
|
+
# F_new = eval_P_partials(zeta_new, y_new, a_coeffs)[0]
|
|
375
|
+
# F_new = complex(F_new)
|
|
376
|
+
#
|
|
377
|
+
# if abs(F_new) <= (1.0 - armijo * lam) * F_abs:
|
|
378
|
+
# w = w_new
|
|
379
|
+
# ok = True
|
|
380
|
+
# break
|
|
381
|
+
#
|
|
382
|
+
# lam *= 0.5
|
|
383
|
+
#
|
|
384
|
+
# if not ok:
|
|
385
|
+
# return w, False
|
|
386
|
+
|
|
387
|
+
# ---------------
|
|
388
|
+
|
|
389
|
+
# TEST
|
|
390
|
+
|
|
391
|
+
# -------------------------
|
|
392
|
+
# Polynomial root selection
|
|
393
|
+
# -------------------------
|
|
394
|
+
# We solve: P(z + alpha/w, tau*w) = 0.
|
|
395
|
+
# Let y = tau*w. Then alpha/w = alpha*tau/y = (tau - 1)/y.
|
|
396
|
+
# So we solve in y:
|
|
397
|
+
# P(z + beta/y, y) = 0, beta = tau - 1.
|
|
398
|
+
# Multiply by y^deg_z to clear denominators and get a polynomial in y.
|
|
399
|
+
|
|
400
|
+
a = numpy.asarray(a_coeffs, dtype=numpy.complex128)
|
|
401
|
+
deg_z = a.shape[0] - 1
|
|
402
|
+
deg_m = a.shape[1] - 1
|
|
403
|
+
|
|
404
|
+
beta = tau - 1.0
|
|
405
|
+
|
|
406
|
+
# poly_y[p] stores coeff of y^p after clearing denominators
|
|
407
|
+
poly_y = numpy.zeros(deg_z + deg_m + 1, dtype=numpy.complex128)
|
|
408
|
+
|
|
409
|
+
# Build polynomial: sum_{i,j} a[i,j] (z + beta/y)^i y^j * y^{deg_z}
|
|
410
|
+
# Expand (z + beta/y)^i = sum_{k=0}^i C(i,k) z^{i-k} (beta/y)^k
|
|
411
|
+
# Term contributes to power p = deg_z + j - k.
|
|
412
|
+
from math import comb
|
|
413
|
+
for i in range(deg_z + 1):
|
|
414
|
+
for j in range(deg_m + 1):
|
|
415
|
+
aij = a[i, j]
|
|
416
|
+
if aij == 0:
|
|
417
|
+
continue
|
|
418
|
+
for k in range(i + 1):
|
|
419
|
+
p = deg_z + j - k
|
|
420
|
+
poly_y[p] += aij * comb(i, k) * (z ** (i - k)) * (beta ** k)
|
|
421
|
+
|
|
422
|
+
# numpy.roots expects highest degree first
|
|
423
|
+
coeffs = poly_y[::-1]
|
|
424
|
+
|
|
425
|
+
# If leading coefficients are ~0, trim (rare but safe)
|
|
426
|
+
nz_lead = numpy.flatnonzero(numpy.abs(coeffs) > 0)
|
|
427
|
+
if nz_lead.size == 0:
|
|
428
|
+
return w, False
|
|
429
|
+
coeffs = coeffs[nz_lead[0]:]
|
|
430
|
+
|
|
431
|
+
roots_y = numpy.roots(coeffs)
|
|
432
|
+
|
|
433
|
+
# Pick root with Im(w)>0 (if z in upper half-plane), closest to time seed
|
|
434
|
+
y_seed = tau * w_init
|
|
435
|
+
best = None
|
|
436
|
+
best_score = None
|
|
437
|
+
|
|
438
|
+
for y in roots_y:
|
|
439
|
+
if not numpy.isfinite(y.real) or not numpy.isfinite(y.imag):
|
|
440
|
+
continue
|
|
441
|
+
|
|
442
|
+
w_cand = y / tau
|
|
443
|
+
|
|
444
|
+
if want_pos_imag and (w_cand.imag <= 0.0):
|
|
445
|
+
continue
|
|
446
|
+
|
|
447
|
+
if abs(w_cand) < w_min:
|
|
448
|
+
continue
|
|
449
|
+
|
|
450
|
+
# score: stick to time continuation
|
|
451
|
+
score = abs(y - y_seed)
|
|
452
|
+
|
|
453
|
+
if (best_score is None) or (score < best_score):
|
|
454
|
+
best = w_cand
|
|
455
|
+
best_score = score
|
|
456
|
+
|
|
457
|
+
if best is None:
|
|
458
|
+
return w, False
|
|
459
|
+
|
|
460
|
+
w = complex(best)
|
|
461
|
+
|
|
462
|
+
# final residual check
|
|
463
|
+
F_end = eval_P_partials(z + alpha / w, tau * w, a_coeffs)[0]
|
|
464
|
+
return w, (abs(F_end) <= 1e3 * tol)
|
|
465
|
+
|
|
466
|
+
# -------------------
|
|
467
|
+
|
|
468
|
+
|
|
469
|
+
F_end = eval_P_partials(z + alpha / w, tau * w, a_coeffs)[0]
|
|
470
|
+
return w, (abs(F_end) <= 10.0 * tol)
|
|
471
|
+
|
|
472
|
+
|
|
473
|
+
# ============
|
|
474
|
+
# NEW FUNCTION
|
|
475
|
+
# ============
|
|
476
|
+
|
|
477
|
+
def fd_candidates_w(z, t, a_coeffs, w_min=1e-14):
|
|
478
|
+
"""
|
|
479
|
+
Return candidate roots w solving P(z + alpha/w, tau*w)=0 with Im(w)>0 (if Im(z)>0).
|
|
480
|
+
"""
|
|
481
|
+
z = complex(z)
|
|
482
|
+
tau = float(numpy.exp(t))
|
|
483
|
+
alpha = 1.0 - 1.0 / tau
|
|
484
|
+
want_pos_imag = (z.imag > 0.0)
|
|
485
|
+
|
|
486
|
+
a = numpy.asarray(a_coeffs, dtype=numpy.complex128)
|
|
487
|
+
deg_z = a.shape[0] - 1
|
|
488
|
+
deg_m = a.shape[1] - 1
|
|
489
|
+
|
|
490
|
+
beta = tau - 1.0 # since alpha/w = (tau-1)/(tau*w) = beta / y with y=tau*w
|
|
491
|
+
|
|
492
|
+
poly_y = numpy.zeros(deg_z + deg_m + 1, dtype=numpy.complex128)
|
|
493
|
+
|
|
494
|
+
from math import comb
|
|
495
|
+
for i in range(deg_z + 1):
|
|
496
|
+
for j in range(deg_m + 1):
|
|
497
|
+
aij = a[i, j]
|
|
498
|
+
if aij == 0:
|
|
499
|
+
continue
|
|
500
|
+
for k in range(i + 1):
|
|
501
|
+
p = deg_z + j - k
|
|
502
|
+
poly_y[p] += aij * comb(i, k) * (z ** (i - k)) * (beta ** k)
|
|
503
|
+
|
|
504
|
+
coeffs = poly_y[::-1]
|
|
505
|
+
nz_lead = numpy.flatnonzero(numpy.abs(coeffs) > 0)
|
|
506
|
+
if nz_lead.size == 0:
|
|
507
|
+
return []
|
|
508
|
+
|
|
509
|
+
coeffs = coeffs[nz_lead[0]:]
|
|
510
|
+
roots_y = numpy.roots(coeffs)
|
|
511
|
+
|
|
512
|
+
cands = []
|
|
513
|
+
for y in roots_y:
|
|
514
|
+
if not numpy.isfinite(y.real) or not numpy.isfinite(y.imag):
|
|
515
|
+
continue
|
|
516
|
+
w = y / tau
|
|
517
|
+
if abs(w) < w_min:
|
|
518
|
+
continue
|
|
519
|
+
if want_pos_imag and (w.imag <= 0.0):
|
|
520
|
+
continue
|
|
521
|
+
# residual filter (optional but helps)
|
|
522
|
+
# -------------
|
|
523
|
+
# TEST
|
|
524
|
+
# F = eval_P_partials(z + alpha / w, tau * w, a_coeffs)[0]
|
|
525
|
+
# if abs(F) < 1e-6:
|
|
526
|
+
# cands.append(complex(w))
|
|
527
|
+
# ---------------
|
|
528
|
+
# TEST
|
|
529
|
+
cands.append(complex(w))
|
|
530
|
+
# ------------------
|
|
531
|
+
|
|
532
|
+
return cands
|
|
533
|
+
|
|
534
|
+
|
|
535
|
+
# =================
|
|
536
|
+
# decompress newton
|
|
537
|
+
# =================
|
|
538
|
+
|
|
539
|
+
def decompress_newton(z_list, t_grid, a_coeffs, w0_list=None,
|
|
540
|
+
dt_max=0.1, sweep=True, time_rel_tol=5.0,
|
|
541
|
+
active_imag_eps=None, sweep_pad=20,
|
|
542
|
+
max_iter=50, tol=1e-12, armijo=1e-4,
|
|
543
|
+
min_lam=1e-6, w_min=1e-14,
|
|
544
|
+
viterbi_opt=None):
|
|
545
|
+
"""
|
|
546
|
+
Evolve w = m(t,z) on a fixed z grid and time grid using FD.
|
|
547
|
+
|
|
548
|
+
This implementation uses a global 1D Viterbi/DP branch-tracker along the
|
|
549
|
+
spatial grid at every time step to avoid local root mis-selection (multi-bulk
|
|
550
|
+
stability). The inputs sweep/time_rel_tol/active_imag_eps/sweep_pad are kept
|
|
551
|
+
for backward compatibility but are ignored by the Viterbi tracker.
|
|
552
|
+
|
|
553
|
+
Parameters
|
|
554
|
+
----------
|
|
555
|
+
z_list : array_like of complex
|
|
556
|
+
Query points z (typically x + 1j*eta with eta > 0), ordered along x.
|
|
557
|
+
t_grid : array_like of float
|
|
558
|
+
Strictly increasing time grid.
|
|
559
|
+
a_coeffs : ndarray
|
|
560
|
+
Coefficients defining P(z,m) in the monomial basis.
|
|
561
|
+
w0_list : array_like of complex
|
|
562
|
+
Initial values w(t0,z) at t_grid[0].
|
|
563
|
+
|
|
564
|
+
viterbi_opt : dict or None
|
|
565
|
+
Options for the Viterbi tracker. Keys (all optional):
|
|
566
|
+
lam_space : float (default 1.0)
|
|
567
|
+
lam_time : float (default 0.25)
|
|
568
|
+
lam_im : float (default 1e3) penalty = lam_im / max(|Im(w)|, eps)
|
|
569
|
+
tol_im : float (default 1e-12) Herglotz sign tolerance
|
|
570
|
+
edge_k : int (default 3) # of points at each end with asym penalty
|
|
571
|
+
lam_asym : float (default 0.2) penalty = lam_asym * |z*w + 1|
|
|
572
|
+
refine_newton : bool (default True) refine chosen path with fd_solve_w
|
|
573
|
+
|
|
574
|
+
Returns
|
|
575
|
+
-------
|
|
576
|
+
W : ndarray, shape (len(t_grid), len(z_list))
|
|
577
|
+
Evolved values w(t,z).
|
|
578
|
+
ok : ndarray of bool, same shape as W
|
|
579
|
+
Convergence flags from the accepted solve at each point.
|
|
580
|
+
"""
|
|
581
|
+
|
|
582
|
+
z_list = numpy.asarray(z_list, dtype=complex).ravel()
|
|
583
|
+
t_grid = numpy.asarray(t_grid, dtype=float).ravel()
|
|
584
|
+
nt = t_grid.size
|
|
585
|
+
nz = z_list.size
|
|
586
|
+
|
|
587
|
+
if w0_list is None:
|
|
588
|
+
raise ValueError("w0_list must be provided (initial m(z) at t_grid[0]).")
|
|
589
|
+
|
|
590
|
+
w0_list = numpy.asarray(w0_list, dtype=complex).ravel()
|
|
591
|
+
if w0_list.size != nz:
|
|
592
|
+
raise ValueError("w0_list must have the same size as z_list.")
|
|
593
|
+
|
|
594
|
+
if nt == 0:
|
|
595
|
+
return numpy.empty((0, nz), dtype=complex), numpy.empty((0, nz), dtype=bool)
|
|
596
|
+
|
|
597
|
+
# Viterbi options
|
|
598
|
+
opt = {} if viterbi_opt is None else dict(viterbi_opt)
|
|
599
|
+
lam_space = float(opt.get('lam_space', 1.0))
|
|
600
|
+
lam_time = float(opt.get('lam_time', 0.25))
|
|
601
|
+
lam_im = float(opt.get('lam_im', 1.0e3))
|
|
602
|
+
tol_im = float(opt.get('tol_im', 1.0e-12))
|
|
603
|
+
edge_k = int(opt.get('edge_k', 3))
|
|
604
|
+
lam_asym = float(opt.get('lam_asym', 0.2))
|
|
605
|
+
refine_newton = bool(opt.get('refine_newton', True))
|
|
606
|
+
|
|
607
|
+
W = numpy.empty((nt, nz), dtype=complex)
|
|
608
|
+
ok = numpy.zeros((nt, nz), dtype=bool)
|
|
609
|
+
|
|
610
|
+
W[0, :] = w0_list
|
|
611
|
+
ok[0, :] = True
|
|
612
|
+
w_prev = W[0, :].copy()
|
|
613
|
+
|
|
614
|
+
# -----------------
|
|
615
|
+
# helper: candidates
|
|
616
|
+
# -----------------
|
|
617
|
+
|
|
618
|
+
def _candidates(iz, t):
|
|
619
|
+
cands = fd_candidates_w(z_list[iz], t, a_coeffs, w_min=w_min)
|
|
620
|
+
if len(cands) == 0:
|
|
621
|
+
# fallback: carry previous value as a candidate
|
|
622
|
+
return [complex(w_prev[iz])]
|
|
623
|
+
return cands
|
|
624
|
+
|
|
625
|
+
# -------------------------
|
|
626
|
+
# helper: unary / transition
|
|
627
|
+
# -------------------------
|
|
628
|
+
|
|
629
|
+
def _want_pos_imag(z):
|
|
630
|
+
return (complex(z).imag > 0.0)
|
|
631
|
+
|
|
632
|
+
def _herglotz_ok(w, z):
|
|
633
|
+
z = complex(z)
|
|
634
|
+
w = complex(w)
|
|
635
|
+
if not _want_pos_imag(z):
|
|
636
|
+
return True
|
|
637
|
+
return (w.imag > -tol_im)
|
|
638
|
+
|
|
639
|
+
def _unary_cost(w, iz, t):
|
|
640
|
+
# penalize wrong sign heavily
|
|
641
|
+
if not _herglotz_ok(w, z_list[iz]):
|
|
642
|
+
return 1.0e30
|
|
643
|
+
|
|
644
|
+
# time continuity
|
|
645
|
+
wt = complex(w_prev[iz])
|
|
646
|
+
c = lam_time * (abs(w - wt) ** 2)
|
|
647
|
+
|
|
648
|
+
# discourage tiny-imag traps (safe substitute for any global Im-reward)
|
|
649
|
+
im = abs(w.imag)
|
|
650
|
+
c += lam_im / max(im, 1e-16)
|
|
651
|
+
|
|
652
|
+
# asymptotic anchor only near ends
|
|
653
|
+
if edge_k > 0 and (iz < edge_k or iz >= nz - edge_k):
|
|
654
|
+
z = complex(z_list[iz])
|
|
655
|
+
c += lam_asym * abs(z * w + 1.0)
|
|
656
|
+
|
|
657
|
+
return c
|
|
658
|
+
|
|
659
|
+
def _trans_cost(w_left, w_right):
|
|
660
|
+
return lam_space * (abs(w_right - w_left) ** 2)
|
|
661
|
+
|
|
662
|
+
# -------------
|
|
663
|
+
# time evolution
|
|
664
|
+
# -------------
|
|
665
|
+
|
|
666
|
+
for it in range(1, nt):
|
|
667
|
+
t = float(t_grid[it])
|
|
668
|
+
|
|
669
|
+
# build candidates list per spatial index
|
|
670
|
+
C = []
|
|
671
|
+
for iz in range(nz):
|
|
672
|
+
C.append(_candidates(iz, t))
|
|
673
|
+
|
|
674
|
+
# DP tables with variable state sizes
|
|
675
|
+
dp = []
|
|
676
|
+
prev_idx = []
|
|
677
|
+
|
|
678
|
+
# init
|
|
679
|
+
c0 = C[0]
|
|
680
|
+
dp0 = numpy.array([_unary_cost(w, 0, t) for w in c0], dtype=float)
|
|
681
|
+
dp.append(dp0)
|
|
682
|
+
prev_idx.append(numpy.full(dp0.size, -1, dtype=int))
|
|
683
|
+
|
|
684
|
+
# forward pass
|
|
685
|
+
for iz in range(1, nz):
|
|
686
|
+
ci = C[iz]
|
|
687
|
+
dp_i = numpy.full(len(ci), numpy.inf, dtype=float)
|
|
688
|
+
prev_i = numpy.full(len(ci), -1, dtype=int)
|
|
689
|
+
|
|
690
|
+
dp_prev = dp[iz - 1]
|
|
691
|
+
c_prev = C[iz - 1]
|
|
692
|
+
|
|
693
|
+
for j, wj in enumerate(ci):
|
|
694
|
+
u = _unary_cost(wj, iz, t)
|
|
695
|
+
|
|
696
|
+
best = numpy.inf
|
|
697
|
+
best_k = -1
|
|
698
|
+
for k, wk in enumerate(c_prev):
|
|
699
|
+
val = dp_prev[k] + _trans_cost(wk, wj)
|
|
700
|
+
if val < best:
|
|
701
|
+
best = val
|
|
702
|
+
best_k = k
|
|
703
|
+
|
|
704
|
+
dp_i[j] = u + best
|
|
705
|
+
prev_i[j] = best_k
|
|
706
|
+
|
|
707
|
+
dp.append(dp_i)
|
|
708
|
+
prev_idx.append(prev_i)
|
|
709
|
+
|
|
710
|
+
# backtrack
|
|
711
|
+
w_row = numpy.empty(nz, dtype=complex)
|
|
712
|
+
ok_row = numpy.zeros(nz, dtype=bool)
|
|
713
|
+
|
|
714
|
+
j = int(numpy.argmin(dp[-1]))
|
|
715
|
+
w_row[-1] = complex(C[-1][j])
|
|
716
|
+
|
|
717
|
+
for iz in range(nz - 1, 0, -1):
|
|
718
|
+
j = int(prev_idx[iz][j])
|
|
719
|
+
if j < 0:
|
|
720
|
+
j = 0
|
|
721
|
+
w_row[iz - 1] = complex(C[iz - 1][j])
|
|
722
|
+
|
|
723
|
+
# optional Newton refinement on chosen path
|
|
724
|
+
if refine_newton:
|
|
725
|
+
for iz in range(nz):
|
|
726
|
+
w_sol, success = fd_solve_w(
|
|
727
|
+
z_list[iz], t, a_coeffs, w_row[iz],
|
|
728
|
+
max_iter=max_iter, tol=tol, armijo=armijo,
|
|
729
|
+
min_lam=min_lam, w_min=w_min)
|
|
730
|
+
w_row[iz] = w_sol
|
|
731
|
+
ok_row[iz] = success
|
|
732
|
+
else:
|
|
733
|
+
ok_row[:] = True
|
|
734
|
+
|
|
735
|
+
W[it, :] = w_row
|
|
736
|
+
ok[it, :] = ok_row
|
|
737
|
+
w_prev = w_row
|
|
738
|
+
|
|
739
|
+
return W, ok
|