freealg 0.7.12__tar.gz → 0.7.15__tar.gz

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.
Files changed (81) hide show
  1. {freealg-0.7.12 → freealg-0.7.15}/MANIFEST.in +0 -1
  2. {freealg-0.7.12 → freealg-0.7.15}/PKG-INFO +1 -1
  3. freealg-0.7.15/freealg/__version__.py +1 -0
  4. freealg-0.7.15/freealg/_algebraic_form/_cusp.py +357 -0
  5. freealg-0.7.15/freealg/_algebraic_form/_cusp_wrap.py +268 -0
  6. {freealg-0.7.12 → freealg-0.7.15}/freealg/_algebraic_form/_decompress2.py +2 -0
  7. freealg-0.7.15/freealg/_algebraic_form/_decompress4.py +739 -0
  8. freealg-0.7.15/freealg/_algebraic_form/_decompress5.py +738 -0
  9. freealg-0.7.15/freealg/_algebraic_form/_decompress6.py +492 -0
  10. freealg-0.7.15/freealg/_algebraic_form/_decompress7.py +355 -0
  11. freealg-0.7.15/freealg/_algebraic_form/_decompress8.py +369 -0
  12. freealg-0.7.15/freealg/_algebraic_form/_decompress9.py +363 -0
  13. freealg-0.7.15/freealg/_algebraic_form/_decompress_new.py +431 -0
  14. freealg-0.7.15/freealg/_algebraic_form/_decompress_new_2.py +1631 -0
  15. freealg-0.7.15/freealg/_algebraic_form/_decompress_util.py +172 -0
  16. freealg-0.7.15/freealg/_algebraic_form/_homotopy2.py +289 -0
  17. freealg-0.7.15/freealg/_algebraic_form/_homotopy3.py +215 -0
  18. freealg-0.7.15/freealg/_algebraic_form/_homotopy4.py +320 -0
  19. freealg-0.7.15/freealg/_algebraic_form/_homotopy5.py +185 -0
  20. {freealg-0.7.12 → freealg-0.7.15}/freealg/_algebraic_form/_moments.py +0 -1
  21. freealg-0.7.15/freealg/_algebraic_form/_support.py +264 -0
  22. {freealg-0.7.12 → freealg-0.7.15}/freealg/_algebraic_form/algebraic_form.py +21 -2
  23. freealg-0.7.15/freealg/distributions/_compound_poisson.py +481 -0
  24. {freealg-0.7.12 → freealg-0.7.15}/freealg/distributions/_deformed_marchenko_pastur.py +6 -7
  25. {freealg-0.7.12 → freealg-0.7.15}/freealg.egg-info/PKG-INFO +1 -1
  26. {freealg-0.7.12 → freealg-0.7.15}/freealg.egg-info/SOURCES.txt +16 -0
  27. freealg-0.7.12/freealg/__version__.py +0 -1
  28. freealg-0.7.12/freealg/_algebraic_form/_support.py +0 -309
  29. {freealg-0.7.12 → freealg-0.7.15}/AUTHORS.txt +0 -0
  30. {freealg-0.7.12 → freealg-0.7.15}/CHANGELOG.rst +0 -0
  31. {freealg-0.7.12 → freealg-0.7.15}/LICENSE.txt +0 -0
  32. {freealg-0.7.12 → freealg-0.7.15}/README.rst +0 -0
  33. {freealg-0.7.12 → freealg-0.7.15}/freealg/__init__.py +0 -0
  34. {freealg-0.7.12 → freealg-0.7.15}/freealg/_algebraic_form/__init__.py +0 -0
  35. {freealg-0.7.12 → freealg-0.7.15}/freealg/_algebraic_form/_branch_points.py +0 -0
  36. {freealg-0.7.12 → freealg-0.7.15}/freealg/_algebraic_form/_constraints.py +0 -0
  37. {freealg-0.7.12 → freealg-0.7.15}/freealg/_algebraic_form/_continuation_algebraic.py +0 -0
  38. {freealg-0.7.12 → freealg-0.7.15}/freealg/_algebraic_form/_decompress.py +0 -0
  39. {freealg-0.7.12 → freealg-0.7.15}/freealg/_algebraic_form/_edge.py +0 -0
  40. {freealg-0.7.12 → freealg-0.7.15}/freealg/_algebraic_form/_homotopy.py +0 -0
  41. {freealg-0.7.12 → freealg-0.7.15}/freealg/_algebraic_form/_sheets_util.py +0 -0
  42. {freealg-0.7.12 → freealg-0.7.15}/freealg/_free_form/__init__.py +0 -0
  43. {freealg-0.7.12 → freealg-0.7.15}/freealg/_free_form/_chebyshev.py +0 -0
  44. {freealg-0.7.12 → freealg-0.7.15}/freealg/_free_form/_damp.py +0 -0
  45. {freealg-0.7.12 → freealg-0.7.15}/freealg/_free_form/_decompress.py +0 -0
  46. {freealg-0.7.12 → freealg-0.7.15}/freealg/_free_form/_density_util.py +0 -0
  47. {freealg-0.7.12 → freealg-0.7.15}/freealg/_free_form/_jacobi.py +0 -0
  48. {freealg-0.7.12 → freealg-0.7.15}/freealg/_free_form/_linalg.py +0 -0
  49. {freealg-0.7.12 → freealg-0.7.15}/freealg/_free_form/_pade.py +0 -0
  50. {freealg-0.7.12 → freealg-0.7.15}/freealg/_free_form/_plot_util.py +0 -0
  51. {freealg-0.7.12 → freealg-0.7.15}/freealg/_free_form/_sample.py +0 -0
  52. {freealg-0.7.12 → freealg-0.7.15}/freealg/_free_form/_series.py +0 -0
  53. {freealg-0.7.12 → freealg-0.7.15}/freealg/_free_form/_support.py +0 -0
  54. {freealg-0.7.12 → freealg-0.7.15}/freealg/_free_form/free_form.py +0 -0
  55. {freealg-0.7.12 → freealg-0.7.15}/freealg/_geometric_form/__init__.py +0 -0
  56. {freealg-0.7.12 → freealg-0.7.15}/freealg/_geometric_form/_continuation_genus0.py +0 -0
  57. {freealg-0.7.12 → freealg-0.7.15}/freealg/_geometric_form/_continuation_genus1.py +0 -0
  58. {freealg-0.7.12 → freealg-0.7.15}/freealg/_geometric_form/_elliptic_functions.py +0 -0
  59. {freealg-0.7.12 → freealg-0.7.15}/freealg/_geometric_form/_sphere_maps.py +0 -0
  60. {freealg-0.7.12 → freealg-0.7.15}/freealg/_geometric_form/_torus_maps.py +0 -0
  61. {freealg-0.7.12 → freealg-0.7.15}/freealg/_geometric_form/geometric_form.py +0 -0
  62. {freealg-0.7.12 → freealg-0.7.15}/freealg/_util.py +0 -0
  63. {freealg-0.7.12 → freealg-0.7.15}/freealg/distributions/__init__.py +0 -0
  64. {freealg-0.7.12 → freealg-0.7.15}/freealg/distributions/_chiral_block.py +0 -0
  65. {freealg-0.7.12 → freealg-0.7.15}/freealg/distributions/_deformed_wigner.py +0 -0
  66. {freealg-0.7.12 → freealg-0.7.15}/freealg/distributions/_kesten_mckay.py +0 -0
  67. {freealg-0.7.12 → freealg-0.7.15}/freealg/distributions/_marchenko_pastur.py +0 -0
  68. {freealg-0.7.12 → freealg-0.7.15}/freealg/distributions/_meixner.py +0 -0
  69. {freealg-0.7.12 → freealg-0.7.15}/freealg/distributions/_wachter.py +0 -0
  70. {freealg-0.7.12 → freealg-0.7.15}/freealg/distributions/_wigner.py +0 -0
  71. {freealg-0.7.12 → freealg-0.7.15}/freealg/visualization/__init__.py +0 -0
  72. {freealg-0.7.12 → freealg-0.7.15}/freealg/visualization/_glue_util.py +0 -0
  73. {freealg-0.7.12 → freealg-0.7.15}/freealg/visualization/_rgb_hsv.py +0 -0
  74. {freealg-0.7.12 → freealg-0.7.15}/freealg.egg-info/dependency_links.txt +0 -0
  75. {freealg-0.7.12 → freealg-0.7.15}/freealg.egg-info/not-zip-safe +0 -0
  76. {freealg-0.7.12 → freealg-0.7.15}/freealg.egg-info/requires.txt +0 -0
  77. {freealg-0.7.12 → freealg-0.7.15}/freealg.egg-info/top_level.txt +0 -0
  78. {freealg-0.7.12 → freealg-0.7.15}/pyproject.toml +0 -0
  79. {freealg-0.7.12 → freealg-0.7.15}/requirements.txt +0 -0
  80. {freealg-0.7.12 → freealg-0.7.15}/setup.cfg +0 -0
  81. {freealg-0.7.12 → freealg-0.7.15}/setup.py +0 -0
@@ -18,7 +18,6 @@ exclude .coveragerc
18
18
  exclude tox.ini
19
19
  exclude TODO.rst
20
20
  exclude CONTRIBUTING.rst
21
- exclude notebook/experimental/clean_battle_20240814_public.json
22
21
 
23
22
  prune docs
24
23
  prune archive
@@ -1,6 +1,6 @@
1
1
  Metadata-Version: 2.4
2
2
  Name: freealg
3
- Version: 0.7.12
3
+ Version: 0.7.15
4
4
  Summary: Free probability for large matrices
5
5
  Home-page: https://github.com/ameli/freealg
6
6
  Download-URL: https://github.com/ameli/freealg/archive/main.zip
@@ -0,0 +1 @@
1
+ __version__ = "0.7.15"
@@ -0,0 +1,357 @@
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
+ import scipy.optimize
16
+ from ._decompress import eval_P_partials
17
+
18
+ __all__ = ["solve_cusp"]
19
+
20
+
21
+ # ==========
22
+ # newton 3x3
23
+ # ==========
24
+
25
+ def _newton_3x3(F, x0, max_iter=60, tol=1e-12, bounds=None, max_step=None):
26
+ x = numpy.array(x0, dtype=float)
27
+
28
+ # bounds: list/tuple of (lo, hi) per component (None means unbounded)
29
+ if bounds is not None:
30
+ b = []
31
+ for lo, hi in bounds:
32
+ b.append((None if lo is None else float(lo),
33
+ None if hi is None else float(hi)))
34
+ bounds = b
35
+
36
+ if max_step is not None:
37
+ max_step = numpy.asarray(max_step, dtype=float)
38
+ if max_step.shape != (3,):
39
+ raise ValueError("max_step must have shape (3,)")
40
+
41
+ def _apply_bounds(xv):
42
+ if bounds is None:
43
+ return xv
44
+ for i, (lo, hi) in enumerate(bounds):
45
+ if lo is not None and xv[i] < lo:
46
+ xv[i] = lo
47
+ if hi is not None and xv[i] > hi:
48
+ xv[i] = hi
49
+ return xv
50
+
51
+ x = _apply_bounds(x.copy())
52
+
53
+ fx = F(x)
54
+ if numpy.linalg.norm(fx) <= tol:
55
+ return x, True, fx
56
+
57
+ for _ in range(max_iter):
58
+ J = numpy.zeros((3, 3), dtype=float)
59
+ eps = 1e-6
60
+ for j in range(3):
61
+ xp = x.copy()
62
+ xp[j] += eps
63
+ xp = _apply_bounds(xp)
64
+ J[:, j] = (F(xp) - fx) / eps
65
+
66
+ try:
67
+ dx = numpy.linalg.solve(J, -fx)
68
+ except numpy.linalg.LinAlgError:
69
+ return x, False, fx
70
+
71
+ if max_step is not None:
72
+ dx = numpy.clip(dx, -max_step, max_step)
73
+
74
+ lam = 1.0
75
+ improved = False
76
+ for _ls in range(12):
77
+ x_try = x + lam * dx
78
+ x_try = _apply_bounds(x_try)
79
+ f_try = F(x_try)
80
+ if numpy.linalg.norm(f_try) < numpy.linalg.norm(fx):
81
+ x, fx = x_try, f_try
82
+ improved = True
83
+ break
84
+ lam *= 0.5
85
+
86
+ if not improved:
87
+ return x, False, fx
88
+
89
+ if numpy.linalg.norm(fx) <= tol:
90
+ return x, True, fx
91
+
92
+ return x, False, fx
93
+
94
+
95
+ __all__ = ["solve_cusp"]
96
+
97
+
98
+ def _second_partials_fd(zeta, y, a_coeffs, eps_z=None, eps_y=None):
99
+ zeta = float(zeta)
100
+ y = float(y)
101
+
102
+ if eps_z is None:
103
+ eps_z = 1e-7 * (1.0 + abs(zeta))
104
+ if eps_y is None:
105
+ eps_y = 1e-7 * (1.0 + abs(y))
106
+
107
+ _, Pz_p, Py_p = eval_P_partials(zeta + eps_z, y, a_coeffs)
108
+ _, Pz_m, Py_m = eval_P_partials(zeta - eps_z, y, a_coeffs)
109
+ Pzz = (Pz_p - Pz_m) / (2.0 * eps_z)
110
+ Pzy1 = (Py_p - Py_m) / (2.0 * eps_z)
111
+
112
+ _, Pz_p, Py_p = eval_P_partials(zeta, y + eps_y, a_coeffs)
113
+ _, Pz_m, Py_m = eval_P_partials(zeta, y - eps_y, a_coeffs)
114
+ Pzy2 = (Pz_p - Pz_m) / (2.0 * eps_y)
115
+ Pyy = (Py_p - Py_m) / (2.0 * eps_y)
116
+
117
+ Pzy = 0.5 * (Pzy1 + Pzy2)
118
+ return float(Pzz), float(Pzy), float(Pyy)
119
+
120
+
121
+ def _cusp_F_real(zeta, y, s, a_coeffs):
122
+ # tau = 1 + exp(s) => c = tau-1 = exp(s) > 0
123
+ c = float(numpy.exp(float(s)))
124
+
125
+ P, Pz, Py = eval_P_partials(float(zeta), float(y), a_coeffs)
126
+ P = float(numpy.real(P))
127
+ Pz = float(numpy.real(Pz))
128
+ Py = float(numpy.real(Py))
129
+
130
+ F1 = P
131
+ F2 = (y * y) * Py - c * Pz
132
+
133
+ Pzz, Pzy, Pyy = _second_partials_fd(zeta, y, a_coeffs)
134
+ F3 = y * (Pzz * (Py * Py) - 2.0 * Pzy * Pz * Py + Pyy * (Pz * Pz)) + \
135
+ 2.0 * (Pz * Pz) * Py
136
+
137
+ return numpy.array([F1, F2, F3], dtype=float)
138
+
139
+
140
+ # ================
141
+ # poly coeffs in y
142
+ # ================
143
+
144
+ def _poly_coeffs_in_y(a_coeffs, zeta):
145
+ a = numpy.asarray(a_coeffs)
146
+ deg_z = a.shape[0] - 1
147
+ deg_y = a.shape[1] - 1
148
+ z_pows = numpy.power(zeta, numpy.arange(deg_z + 1, dtype=numpy.int64))
149
+ c = numpy.empty((deg_y + 1,), dtype=numpy.complex128)
150
+ for j in range(deg_y + 1):
151
+ c[j] = numpy.dot(a[:, j], z_pows)
152
+ return c # ascending in y
153
+
154
+
155
+ # ===================
156
+ # pick realish root y
157
+ # ===================
158
+
159
+ def _pick_realish_root_y(a_coeffs, zeta):
160
+
161
+ c_asc = _poly_coeffs_in_y(a_coeffs, zeta)
162
+ c_desc = c_asc[::-1] # descending for numpy.roots
163
+
164
+ k = 0
165
+ while k < len(c_desc) and abs(c_desc[k]) == 0:
166
+ k += 1
167
+ c_desc = c_desc[k:] if k < len(c_desc) else c_desc
168
+
169
+ if len(c_desc) <= 1:
170
+ return 0.0
171
+
172
+ roots = numpy.roots(c_desc)
173
+ j = int(numpy.argmin(numpy.abs(numpy.imag(roots))))
174
+ return float(numpy.real(roots[j]))
175
+
176
+
177
+ # ==========
178
+ # solve cusp
179
+ # ==========
180
+
181
+ def solve_cusp(
182
+ a_coeffs,
183
+ t_init,
184
+ zeta_init,
185
+ y_init=None,
186
+ max_iter=80,
187
+ tol=1e-12,
188
+ t_bounds=None,
189
+ zeta_bounds=None):
190
+ """
191
+ Exact-derivative cusp solve for (zeta, y, t) with unknowns (zeta, y, s),
192
+ where tau = 1 + exp(s), t = log(tau), x = zeta - (tau-1)/y.
193
+
194
+ a_coeffs: array shape (deg_z+1, deg_y+1), P(zeta,y)=
195
+ sum_{i,j} a[i,j]*zeta^i*y^j
196
+ """
197
+
198
+ a = numpy.asarray(a_coeffs, dtype=numpy.complex128)
199
+ deg_z = a.shape[0] - 1
200
+ deg_y = a.shape[1] - 1
201
+
202
+ def _P_partials_all(zeta, y):
203
+ # returns (P, Pz, Py, Pzz, Pzy, Pyy) as complex
204
+ zeta = numpy.complex128(zeta)
205
+ y = numpy.complex128(y)
206
+
207
+ zi = numpy.power(zeta, numpy.arange(deg_z + 1, dtype=numpy.int64))
208
+ yj = numpy.power(y, numpy.arange(deg_y + 1, dtype=numpy.int64))
209
+
210
+ P = numpy.sum(a * zi[:, None] * yj[None, :])
211
+
212
+ # Pz
213
+ if deg_z >= 1:
214
+ iz = numpy.arange(1, deg_z + 1, dtype=numpy.int64)
215
+ zi_m1 = numpy.power(zeta, iz - 1)
216
+ Pz = numpy.sum(
217
+ (a[iz, :] * iz[:, None]) * zi_m1[:, None] * yj[None, :])
218
+ else:
219
+ Pz = 0.0 + 0.0j
220
+
221
+ # Py
222
+ if deg_y >= 1:
223
+ jy = numpy.arange(1, deg_y + 1, dtype=numpy.int64)
224
+ yj_m1 = numpy.power(y, jy - 1)
225
+ Py = numpy.sum(
226
+ (a[:, jy] * jy[None, :]) * zi[:, None] * yj_m1[None, :])
227
+ else:
228
+ Py = 0.0 + 0.0j
229
+
230
+ # Pzz
231
+ if deg_z >= 2:
232
+ iz = numpy.arange(2, deg_z + 1, dtype=numpy.int64)
233
+ zi_m2 = numpy.power(zeta, iz - 2)
234
+ Pzz = numpy.sum((a[iz, :] * (iz * (iz - 1))[:, None]) *
235
+ zi_m2[:, None] * yj[None, :])
236
+ else:
237
+ Pzz = 0.0 + 0.0j
238
+
239
+ # Pyy
240
+ if deg_y >= 2:
241
+ jy = numpy.arange(2, deg_y + 1, dtype=numpy.int64)
242
+ yj_m2 = numpy.power(y, jy - 2)
243
+ Pyy = numpy.sum((a[:, jy] * (jy * (jy - 1))[None, :]) *
244
+ zi[:, None] * yj_m2[None, :])
245
+ else:
246
+ Pyy = 0.0 + 0.0j
247
+
248
+ # Pzy
249
+ if (deg_z >= 1) and (deg_y >= 1):
250
+ iz = numpy.arange(1, deg_z + 1, dtype=numpy.int64)
251
+ jy = numpy.arange(1, deg_y + 1, dtype=numpy.int64)
252
+ zi_m1 = numpy.power(zeta, iz - 1)
253
+ yj_m1 = numpy.power(y, jy - 1)
254
+ coeff = a[numpy.ix_(iz, jy)] * (iz[:, None] * jy[None, :])
255
+ Pzy = numpy.sum(coeff * zi_m1[:, None] * yj_m1[None, :])
256
+ else:
257
+ Pzy = 0.0 + 0.0j
258
+
259
+ return P, Pz, Py, Pzz, Pzy, Pyy
260
+
261
+ def _F(vec):
262
+ zeta, y, s = float(vec[0]), float(vec[1]), float(vec[2])
263
+ c = float(numpy.exp(s)) # c = tau - 1 > 0
264
+ P, Pz, Py, Pzz, Pzy, Pyy = _P_partials_all(zeta, y)
265
+
266
+ # Work in reals: cusp lives on real zeta,y for real cusp
267
+ P = float(numpy.real(P))
268
+ Pz = float(numpy.real(Pz))
269
+ Py = float(numpy.real(Py))
270
+ Pzz = float(numpy.real(Pzz))
271
+ Pzy = float(numpy.real(Pzy))
272
+ Pyy = float(numpy.real(Pyy))
273
+
274
+ F1 = P
275
+ F2 = (y * y) * Py - c * Pz
276
+ F3 = y * (Pzz * (Py * Py) - 2.0 * Pzy * Pz * Py + Pyy * (Pz * Pz)) + \
277
+ 2.0 * (Pz * Pz) * Py
278
+ return numpy.array([F1, F2, F3], dtype=float)
279
+
280
+ z0 = float(zeta_init)
281
+
282
+ # seed y: keep your provided seed; else pick a real-ish root at z0
283
+ if y_init is None:
284
+ # build polynomial in y at fixed z0 and pick root with smallest imag
285
+ zi = numpy.power(z0, numpy.arange(deg_z + 1, dtype=numpy.int64))
286
+ c_asc = numpy.array([numpy.dot(a[:, j], zi) for j in range(deg_y + 1)],
287
+ dtype=numpy.complex128)
288
+ c_desc = c_asc[::-1]
289
+ kk = 0
290
+ while kk < len(c_desc) and abs(c_desc[kk]) == 0:
291
+ kk += 1
292
+ c_desc = c_desc[kk:] if kk < len(c_desc) else c_desc
293
+ roots = numpy.roots(c_desc) if len(c_desc) > 1 else numpy.array([0.0])
294
+ j = int(numpy.argmin(numpy.abs(numpy.imag(roots))))
295
+ y0 = float(numpy.real(roots[j]))
296
+ else:
297
+ y0 = float(y_init)
298
+
299
+ tau0 = float(numpy.exp(float(t_init)))
300
+ c0 = max(tau0 - 1.0, 1e-14)
301
+ s0 = float(numpy.log(c0))
302
+
303
+ # bounds for zeta, y, s
304
+ z_lo, z_hi = -numpy.inf, numpy.inf
305
+ if zeta_bounds is not None:
306
+ z_lo, z_hi = float(zeta_bounds[0]), float(zeta_bounds[1])
307
+ if z_hi < z_lo:
308
+ z_lo, z_hi = z_hi, z_lo
309
+
310
+ s_lo, s_hi = -numpy.inf, numpy.inf
311
+ if t_bounds is not None:
312
+ t_lo, t_hi = float(t_bounds[0]), float(t_bounds[1])
313
+ if t_hi < t_lo:
314
+ t_lo, t_hi = t_hi, t_lo
315
+ c_lo = max(float(numpy.expm1(t_lo)), 1e-14)
316
+ c_hi = max(float(numpy.expm1(t_hi)), 1e-14)
317
+ s_lo, s_hi = float(numpy.log(c_lo)), float(numpy.log(c_hi))
318
+
319
+ # keep y on the seeded sheet (this is crucial)
320
+ y_rad = 4.0 * (1.0 + abs(y0))
321
+ y_lo, y_hi = float(y0 - y_rad), float(y0 + y_rad)
322
+
323
+ lb = numpy.array([z_lo, y_lo, s_lo], dtype=float)
324
+ ub = numpy.array([z_hi, y_hi, s_hi], dtype=float)
325
+ x0 = numpy.array([z0, y0, s0], dtype=float)
326
+ x0 = numpy.minimum(numpy.maximum(x0, lb), ub)
327
+
328
+ res = scipy.optimize.least_squares(
329
+ _F,
330
+ x0,
331
+ bounds=(lb, ub),
332
+ method="trf",
333
+ max_nfev=int(max_iter) * 100,
334
+ ftol=tol,
335
+ xtol=tol,
336
+ gtol=tol,
337
+ x_scale="jac")
338
+
339
+ zeta, y, s = res.x
340
+ c = float(numpy.exp(float(s)))
341
+ tau = 1.0 + c
342
+ t = float(numpy.log(tau))
343
+ x = float(zeta - (tau - 1.0) / y)
344
+
345
+ F_final = _F(res.x)
346
+ ok = bool(res.success and
347
+ (numpy.max(numpy.abs(F_final)) <= max(1e-9, 50.0 * tol)))
348
+
349
+ return {
350
+ "ok": ok,
351
+ "t": t,
352
+ "tau": float(tau),
353
+ "zeta": float(zeta),
354
+ "y": float(y),
355
+ "x": x,
356
+ "F": F_final,
357
+ "success": bool(res.success)}
@@ -0,0 +1,268 @@
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
+ import scipy.optimize as opt
16
+
17
+
18
+ # ================
19
+ # poly coeffs in y
20
+ # ================
21
+
22
+ def _poly_coeffs_in_y(a_coeffs, zeta):
23
+ """
24
+ Build coefficients c_j(zeta) so that P(zeta, y) = sum_j c_j(zeta) y^j.
25
+
26
+ Assumes a_coeffs[i, j] multiplies z^i y^j (same layout as eval_P in
27
+ _continuation_algebraic). Returns coefficients in ascending powers of y.
28
+ """
29
+
30
+ a = numpy.asarray(a_coeffs)
31
+ deg_z = a.shape[0] - 1
32
+ deg_y = a.shape[1] - 1
33
+
34
+ # c_j(zeta) = sum_i a[i,j] zeta^i
35
+ z_pows = numpy.power(zeta, numpy.arange(deg_z + 1, dtype=numpy.int64))
36
+ c = numpy.empty((deg_y + 1,), dtype=numpy.complex128)
37
+ for j in range(deg_y + 1):
38
+ c[j] = numpy.dot(a[:, j], z_pows)
39
+
40
+ return c
41
+
42
+
43
+ # ===================
44
+ # pick realish root y
45
+ # ===================
46
+
47
+ def _pick_realish_root_y(a_coeffs, zeta):
48
+ """
49
+ Pick a reasonable real-ish root y of P(zeta, y)=0 to seed Newton.
50
+
51
+ Returns a float (real part of the selected root).
52
+ """
53
+
54
+ c_asc = _poly_coeffs_in_y(a_coeffs, zeta) # ascending in y
55
+ # numpy.roots wants descending order
56
+ c_desc = c_asc[::-1]
57
+ # strip leading ~0 coefficients
58
+ k = 0
59
+ while k < len(c_desc) and abs(c_desc[k]) == 0:
60
+ k += 1
61
+ c_desc = c_desc[k:] if k < len(c_desc) else c_desc
62
+
63
+ if len(c_desc) <= 1:
64
+ return 0.0
65
+
66
+ roots = numpy.roots(c_desc)
67
+ # choose the root closest to the real axis
68
+ j = int(numpy.argmin(numpy.abs(numpy.imag(roots))))
69
+ return float(numpy.real(roots[j]))
70
+
71
+
72
+ # =========
73
+ # cusp wrap
74
+ # =========
75
+
76
+ def cusp_wrap(self, t_grid, edge_kwargs=None, max_iter=80, tol=1e-12,
77
+ verbose=False):
78
+
79
+ if edge_kwargs is None:
80
+ edge_kwargs = {}
81
+
82
+ t_grid = numpy.asarray(t_grid, dtype=float).ravel()
83
+
84
+ # allow scalar / len-1 input
85
+ if t_grid.size == 1:
86
+ t0 = float(t_grid[0])
87
+ dt = 0.25
88
+ t_grid = numpy.linspace(max(0.0, t0 - dt), t0 + dt, 21)
89
+
90
+ if t_grid.size < 5:
91
+ raise ValueError("t_grid too small")
92
+
93
+ def gap_at(tt):
94
+ ce, _, _ = self.edge(numpy.array([float(tt)]), verbose=False,
95
+ **edge_kwargs)
96
+ return float(ce[0, 2].real - ce[0, 1].real)
97
+
98
+ # coarse grid gap
99
+ ce, _, _ = self.edge(t_grid, verbose=False, **edge_kwargs)
100
+ gap = ce[:, 2].real - ce[:, 1].real
101
+ m = numpy.isfinite(gap)
102
+
103
+ if numpy.count_nonzero(m) < 2:
104
+ return {"success": False, "reason": "gap is not finite on grid"}
105
+
106
+ tg = t_grid[m]
107
+ gg = gap[m]
108
+
109
+ # candidate bracket indices from coarse grid
110
+ s = numpy.sign(gg)
111
+ idx = numpy.where(s[:-1] * s[1:] < 0)[0]
112
+
113
+ bracketed = False
114
+ t_star = None
115
+
116
+ # robust: verify sign change using the true gap_at before calling brentq
117
+ if idx.size > 0:
118
+ for ii in idx[:5]: # try a few brackets
119
+ tL, tR = float(tg[ii]), float(tg[ii + 1])
120
+ gL = gap_at(tL)
121
+ gR = gap_at(tR)
122
+ if numpy.isfinite(gL) and numpy.isfinite(gR) and (gL * gR < 0.0):
123
+ t_star = float(opt.brentq(gap_at, tL, tR, xtol=1e-12,
124
+ rtol=1e-12, maxiter=200))
125
+ bracketed = True
126
+ break
127
+
128
+ # fallback: minimizer of |gap| on the coarse grid
129
+ if t_star is None:
130
+ i0 = int(numpy.argmin(numpy.abs(gg)))
131
+ t_star = float(tg[i0])
132
+ bracketed = False
133
+
134
+ # --- seed (zeta,y) correctly using zeta = x + (tau-1)/y ---
135
+ ce_star, _, _ = self.edge(numpy.array([t_star]), verbose=False,
136
+ **edge_kwargs)
137
+ x_seed = float(ce_star[0, 1].real) # inner edge b1
138
+ tau = float(numpy.exp(t_star))
139
+ c = tau - 1.0
140
+
141
+ a = numpy.asarray(self.a_coeffs, dtype=numpy.complex128)
142
+ deg_z = a.shape[0] - 1
143
+ deg_y = a.shape[1] - 1
144
+
145
+ def poly_in_y(zeta):
146
+ zi = numpy.power(zeta, numpy.arange(deg_z + 1, dtype=numpy.int64))
147
+ c_asc = numpy.array([numpy.dot(a[:, j], zi) for j in range(deg_y + 1)],
148
+ dtype=numpy.complex128)
149
+ return c_asc
150
+
151
+ zeta0 = float(x_seed)
152
+ c_asc = poly_in_y(zeta0)
153
+ roots = numpy.roots(c_asc[::-1])
154
+ j = int(numpy.argmin(numpy.abs(numpy.imag(roots))))
155
+ y0 = float(numpy.real(roots[j]))
156
+ if abs(y0) < 1e-12:
157
+ jj = numpy.argsort(numpy.abs(numpy.imag(roots)))
158
+ for k in jj:
159
+ if abs(numpy.real(roots[k])) > 1e-8:
160
+ y0 = float(numpy.real(roots[k]))
161
+ break
162
+
163
+ zeta_seed = float(x_seed + c / y0)
164
+ y_seed = float(y0)
165
+
166
+ def P_all(zeta, y):
167
+ zeta = numpy.complex128(zeta)
168
+ y = numpy.complex128(y)
169
+ zi = numpy.power(zeta, numpy.arange(deg_z + 1))
170
+ yj = numpy.power(y, numpy.arange(deg_y + 1))
171
+ P = numpy.sum(a * zi[:, None] * yj[None, :])
172
+
173
+ if deg_z >= 1:
174
+ iz = numpy.arange(1, deg_z + 1)
175
+ Pz = numpy.sum((a[iz, :] * iz[:, None]) *
176
+ numpy.power(zeta, iz - 1)[:, None] * yj[None, :])
177
+ else:
178
+ Pz = 0.0 + 0.0j
179
+
180
+ if deg_y >= 1:
181
+ jy = numpy.arange(1, deg_y + 1)
182
+ Py = numpy.sum((a[:, jy] * jy[None, :]) * zi[:, None] *
183
+ numpy.power(y, jy - 1)[None, :])
184
+ else:
185
+ Py = 0.0 + 0.0j
186
+
187
+ if deg_z >= 2:
188
+ iz = numpy.arange(2, deg_z + 1)
189
+ Pzz = numpy.sum((a[iz, :] * (iz * (iz - 1))[:, None]) *
190
+ numpy.power(zeta, iz - 2)[:, None] * yj[None, :])
191
+ else:
192
+ Pzz = 0.0 + 0.0j
193
+
194
+ if deg_y >= 2:
195
+ jy = numpy.arange(2, deg_y + 1)
196
+ Pyy = numpy.sum((a[:, jy] * (jy * (jy - 1))[None, :]) *
197
+ zi[:, None] * numpy.power(y, jy - 2)[None, :])
198
+ else:
199
+ Pyy = 0.0 + 0.0j
200
+
201
+ if (deg_z >= 1) and (deg_y >= 1):
202
+ iz = numpy.arange(1, deg_z + 1)
203
+ jy = numpy.arange(1, deg_y + 1)
204
+ coeff = a[numpy.ix_(iz, jy)] * (iz[:, None] * jy[None, :])
205
+ Pzy = numpy.sum(coeff * numpy.power(zeta, iz - 1)[:, None] *
206
+ numpy.power(y, jy - 1)[None, :])
207
+ else:
208
+ Pzy = 0.0 + 0.0j
209
+
210
+ return P, Pz, Py, Pzz, Pzy, Pyy
211
+
212
+ def G(v):
213
+ zeta, y = float(v[0]), float(v[1])
214
+ P, Pz, Py, _, _, _ = P_all(zeta, y)
215
+ P = float(numpy.real(P))
216
+ Pz = float(numpy.real(Pz))
217
+ Py = float(numpy.real(Py))
218
+ F2 = (y * y) * Py - c * Pz
219
+ return numpy.array([P, F2], dtype=float)
220
+
221
+ z_rad = 0.5
222
+ y_rad = 5.0 * (1.0 + abs(y_seed))
223
+ lb = numpy.array([zeta_seed - z_rad, y_seed - y_rad], dtype=float)
224
+ ub = numpy.array([zeta_seed + z_rad, y_seed + y_rad], dtype=float)
225
+
226
+ res = opt.least_squares(
227
+ G, numpy.array([zeta_seed, y_seed], dtype=float),
228
+ bounds=(lb, ub), method="trf",
229
+ max_nfev=8000, ftol=tol, xtol=tol, gtol=tol, x_scale="jac"
230
+ )
231
+
232
+ zeta_star = float(res.x[0])
233
+ y_star = float(res.x[1])
234
+ x_star = float(zeta_star - c / y_star)
235
+
236
+ P, Pz, Py, Pzz, Pzy, Pyy = P_all(zeta_star, y_star)
237
+ P = float(numpy.real(P))
238
+ Pz = float(numpy.real(Pz))
239
+ Py = float(numpy.real(Py))
240
+ Pzz = float(numpy.real(Pzz))
241
+ Pzy = float(numpy.real(Pzy))
242
+ Pyy = float(numpy.real(Pyy))
243
+
244
+ F2 = (y_star * y_star) * Py - c * Pz
245
+ F3 = y_star * (Pzz * (Py * Py) - 2.0 * Pzy * Pz * Py + Pyy * (Pz * Pz)) + \
246
+ 2.0 * (Pz * Pz) * Py
247
+ F = numpy.array([P, float(F2), float(F3)], dtype=float)
248
+
249
+ ok = bool(numpy.max(numpy.abs(F)) < 1e-8)
250
+
251
+ return {
252
+ "ok": ok,
253
+ "t": float(t_star),
254
+ "tau": float(tau),
255
+ "zeta": float(zeta_star),
256
+ "y": float(y_star),
257
+ "x": float(x_star),
258
+ "F": F,
259
+ "success": True,
260
+ "seed": {
261
+ "t": float(t_star),
262
+ "x": float(x_seed),
263
+ "zeta": float(zeta_seed),
264
+ "y": float(y_seed)
265
+ },
266
+ "merge": {"bracketed": bool(bracketed)},
267
+ "gap_at_t": float(gap_at(t_star)),
268
+ "lsq_success": bool(res.success)}
@@ -35,6 +35,7 @@ def decompress_coeffs(a, t, normalize=True):
35
35
  sum_{r=0..L} sum_{s=0..L+K} A[r, s](t) z^r m^s = 0,
36
36
  normalized by normalize_coefficients.
37
37
  """
38
+
38
39
  a = numpy.asarray(a)
39
40
  a[-1, 0] = 0.0
40
41
  if a.ndim != 2:
@@ -123,6 +124,7 @@ def plot_candidates(a, x, delta=1e-4, size=None, latex=False, verbose=False):
123
124
  ax : matplotlib.axes.Axes
124
125
  The axes the scatter plot was drawn on.
125
126
  """
127
+
126
128
  if not (isinstance(delta, (float, int)) and delta > 0):
127
129
  raise ValueError("delta must be a positive scalar.")
128
130