freealg 0.6.2__tar.gz → 0.7.0__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 (59) hide show
  1. {freealg-0.6.2 → freealg-0.7.0}/PKG-INFO +9 -11
  2. {freealg-0.6.2 → freealg-0.7.0}/README.rst +8 -10
  3. freealg-0.7.0/freealg/__init__.py +20 -0
  4. freealg-0.7.0/freealg/__version__.py +1 -0
  5. freealg-0.7.0/freealg/_algebraic_form/__init__.py +11 -0
  6. freealg-0.7.0/freealg/_algebraic_form/_continuation_algebraic.py +503 -0
  7. freealg-0.7.0/freealg/_algebraic_form/_decompress.py +648 -0
  8. freealg-0.7.0/freealg/_algebraic_form/_edge.py +352 -0
  9. freealg-0.7.0/freealg/_algebraic_form/_sheets_util.py +145 -0
  10. freealg-0.7.0/freealg/_algebraic_form/algebraic_form.py +987 -0
  11. {freealg-0.6.2/freealg → freealg-0.7.0/freealg/_freeform}/__init__.py +3 -6
  12. {freealg-0.6.2/freealg → freealg-0.7.0/freealg/_freeform}/_decompress.py +0 -10
  13. freealg-0.6.2/freealg/_util.py → freealg-0.7.0/freealg/_freeform/_density_util.py +22 -67
  14. {freealg-0.6.2/freealg → freealg-0.7.0/freealg/_freeform}/_linalg.py +1 -1
  15. {freealg-0.6.2/freealg → freealg-0.7.0/freealg/_freeform}/_pade.py +0 -1
  16. {freealg-0.6.2/freealg → freealg-0.7.0/freealg/_freeform}/freeform.py +2 -31
  17. freealg-0.7.0/freealg/_geometric_form/__init__.py +13 -0
  18. freealg-0.7.0/freealg/_geometric_form/_continuation_genus0.py +175 -0
  19. freealg-0.7.0/freealg/_geometric_form/_continuation_genus1.py +275 -0
  20. freealg-0.7.0/freealg/_geometric_form/_elliptic_functions.py +174 -0
  21. freealg-0.7.0/freealg/_geometric_form/_sphere_maps.py +63 -0
  22. freealg-0.7.0/freealg/_geometric_form/_torus_maps.py +118 -0
  23. freealg-0.7.0/freealg/_geometric_form/geometric_form.py +1094 -0
  24. freealg-0.7.0/freealg/_util.py +72 -0
  25. {freealg-0.6.2 → freealg-0.7.0}/freealg/distributions/__init__.py +5 -1
  26. freealg-0.7.0/freealg/distributions/_chiral_block.py +440 -0
  27. freealg-0.7.0/freealg/distributions/_deformed_marchenko_pastur.py +617 -0
  28. freealg-0.7.0/freealg/distributions/_deformed_wigner.py +312 -0
  29. {freealg-0.6.2 → freealg-0.7.0}/freealg/distributions/_marchenko_pastur.py +197 -80
  30. freealg-0.7.0/freealg/visualization/__init__.py +12 -0
  31. freealg-0.7.0/freealg/visualization/_glue_util.py +32 -0
  32. freealg-0.7.0/freealg/visualization/_rgb_hsv.py +125 -0
  33. {freealg-0.6.2 → freealg-0.7.0}/freealg.egg-info/PKG-INFO +9 -11
  34. freealg-0.7.0/freealg.egg-info/SOURCES.txt +56 -0
  35. freealg-0.6.2/freealg/__version__.py +0 -1
  36. freealg-0.6.2/freealg.egg-info/SOURCES.txt +0 -35
  37. {freealg-0.6.2 → freealg-0.7.0}/AUTHORS.txt +0 -0
  38. {freealg-0.6.2 → freealg-0.7.0}/CHANGELOG.rst +0 -0
  39. {freealg-0.6.2 → freealg-0.7.0}/LICENSE.txt +0 -0
  40. {freealg-0.6.2 → freealg-0.7.0}/MANIFEST.in +0 -0
  41. {freealg-0.6.2/freealg → freealg-0.7.0/freealg/_freeform}/_chebyshev.py +0 -0
  42. {freealg-0.6.2/freealg → freealg-0.7.0/freealg/_freeform}/_damp.py +0 -0
  43. {freealg-0.6.2/freealg → freealg-0.7.0/freealg/_freeform}/_jacobi.py +0 -0
  44. {freealg-0.6.2/freealg → freealg-0.7.0/freealg/_freeform}/_plot_util.py +0 -0
  45. {freealg-0.6.2/freealg → freealg-0.7.0/freealg/_freeform}/_sample.py +0 -0
  46. {freealg-0.6.2/freealg → freealg-0.7.0/freealg/_freeform}/_series.py +0 -0
  47. {freealg-0.6.2/freealg → freealg-0.7.0/freealg/_freeform}/_support.py +0 -0
  48. {freealg-0.6.2 → freealg-0.7.0}/freealg/distributions/_kesten_mckay.py +0 -0
  49. {freealg-0.6.2 → freealg-0.7.0}/freealg/distributions/_meixner.py +0 -0
  50. {freealg-0.6.2 → freealg-0.7.0}/freealg/distributions/_wachter.py +0 -0
  51. {freealg-0.6.2 → freealg-0.7.0}/freealg/distributions/_wigner.py +0 -0
  52. {freealg-0.6.2 → freealg-0.7.0}/freealg.egg-info/dependency_links.txt +0 -0
  53. {freealg-0.6.2 → freealg-0.7.0}/freealg.egg-info/not-zip-safe +0 -0
  54. {freealg-0.6.2 → freealg-0.7.0}/freealg.egg-info/requires.txt +0 -0
  55. {freealg-0.6.2 → freealg-0.7.0}/freealg.egg-info/top_level.txt +0 -0
  56. {freealg-0.6.2 → freealg-0.7.0}/pyproject.toml +0 -0
  57. {freealg-0.6.2 → freealg-0.7.0}/requirements.txt +0 -0
  58. {freealg-0.6.2 → freealg-0.7.0}/setup.cfg +0 -0
  59. {freealg-0.6.2 → freealg-0.7.0}/setup.py +0 -0
@@ -1,6 +1,6 @@
1
1
  Metadata-Version: 2.4
2
2
  Name: freealg
3
- Version: 0.6.2
3
+ Version: 0.7.0
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
@@ -163,19 +163,17 @@ code, we also welcome feature requests and bug reports.
163
163
  How to Cite
164
164
  ===========
165
165
 
166
- If you use this work, please cite our `arXiv paper <https://arxiv.org/abs/2506.11994>`__.
166
+ If you use this work, please cite our `paper <https://openreview.net/pdf?id=2CeGVUpOd7>`__.
167
167
 
168
168
  .. code::
169
169
 
170
- @article{spectral2025,
171
- title={Spectral Estimation with Free Decompression},
172
- author={Siavash Ameli and Chris van der Heide and Liam Hodgkinson and Michael W. Mahoney},
173
- year={2025},
174
- eprint={2506.11994},
175
- archivePrefix={arXiv},
176
- primaryClass={stat.ML},
177
- url={https://arxiv.org/abs/2506.11994},
178
- journal={arXiv preprint arXiv:2506.11994},
170
+ @inproceedings{
171
+ AMELI-2025,
172
+ title={Spectral Estimation with Free Decompression},
173
+ author={Siavash Ameli and Chris van der Heide and Liam Hodgkinson and Michael W. Mahoney},
174
+ booktitle={The Thirty-ninth Annual Conference on Neural Information Processing Systems},
175
+ year={2025},
176
+ url={https://openreview.net/forum?id=2CeGVUpOd7}
179
177
  }
180
178
 
181
179
 
@@ -87,19 +87,17 @@ code, we also welcome feature requests and bug reports.
87
87
  How to Cite
88
88
  ===========
89
89
 
90
- If you use this work, please cite our `arXiv paper <https://arxiv.org/abs/2506.11994>`__.
90
+ If you use this work, please cite our `paper <https://openreview.net/pdf?id=2CeGVUpOd7>`__.
91
91
 
92
92
  .. code::
93
93
 
94
- @article{spectral2025,
95
- title={Spectral Estimation with Free Decompression},
96
- author={Siavash Ameli and Chris van der Heide and Liam Hodgkinson and Michael W. Mahoney},
97
- year={2025},
98
- eprint={2506.11994},
99
- archivePrefix={arXiv},
100
- primaryClass={stat.ML},
101
- url={https://arxiv.org/abs/2506.11994},
102
- journal={arXiv preprint arXiv:2506.11994},
94
+ @inproceedings{
95
+ AMELI-2025,
96
+ title={Spectral Estimation with Free Decompression},
97
+ author={Siavash Ameli and Chris van der Heide and Liam Hodgkinson and Michael W. Mahoney},
98
+ booktitle={The Thirty-ninth Annual Conference on Neural Information Processing Systems},
99
+ year={2025},
100
+ url={https://openreview.net/forum?id=2CeGVUpOd7}
103
101
  }
104
102
 
105
103
 
@@ -0,0 +1,20 @@
1
+ # SPDX-FileCopyrightText: Copyright 2025, 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
6
+ # under the terms of the license found in the LICENSE.txt file in the root
7
+ # directory of this source tree.
8
+
9
+ from ._freeform import FreeForm, eigvalsh, cond, norm, trace, slogdet, supp, \
10
+ sample, kde
11
+ from ._algebraic_form import AlgebraicForm
12
+ from ._geometric_form import GeometricForm
13
+ from . import visualization
14
+ from . import distributions
15
+
16
+ __all__ = ['FreeForm', 'distributions', 'visualization', 'eigvalsh', 'cond',
17
+ 'norm', 'trace', 'slogdet', 'supp', 'sample', 'kde',
18
+ 'AlgebraicForm', 'GeometricForm']
19
+
20
+ from .__version__ import __version__ # noqa: F401 E402
@@ -0,0 +1 @@
1
+ __version__ = "0.7.0"
@@ -0,0 +1,11 @@
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
6
+ # under the terms of the license found in the LICENSE.txt file in the root
7
+ # directory of this source tree.
8
+
9
+ from .algebraic_form import AlgebraicForm
10
+
11
+ __all__ = ['AlgebraicForm']
@@ -0,0 +1,503 @@
1
+ # SPDX-FileCopyrightText: Copyright 2025, 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 .._geometric_form._continuation_genus0 import joukowski_z
16
+
17
+ __all__ = [
18
+ 'sample_z_joukowski', 'filter_z_away_from_cuts', 'powers',
19
+ 'fit_polynomial_relation', 'eval_P', 'eval_roots',
20
+ 'build_sheets_from_roots']
21
+
22
+
23
+ # ==================
24
+ # sample z joukowski
25
+ # ==================
26
+
27
+ def sample_z_joukowski(a, b, n_samples=4096, r=1.25, n_r=3, r_min=None):
28
+
29
+ if r_min is None:
30
+ r_min = 1.0 + 0.05 * (r - 1.0) if r > 1.0 else 1.0
31
+
32
+ if n_r is None or n_r < 1:
33
+ n_r = 1
34
+
35
+ if n_samples % 2 != 0:
36
+ raise ValueError('n_samples should be even.')
37
+
38
+ if n_r == 1:
39
+ rs = numpy.array([r], dtype=float)
40
+ else:
41
+ rs = numpy.linspace(r_min, r, n_r)
42
+
43
+ n_half = n_samples // 2
44
+ theta = numpy.pi * (numpy.arange(n_half) + 0.5) / n_half
45
+
46
+ z_list = []
47
+ for r_i in rs:
48
+ w = r_i * numpy.exp(1j * theta)
49
+ z = joukowski_z(w, a, b)
50
+ z_list.append(z)
51
+ z_list.append(numpy.conjugate(z))
52
+
53
+ return numpy.concatenate(z_list)
54
+
55
+
56
+ # =======================
57
+ # filter z away from cuts
58
+ # =======================
59
+
60
+ def filter_z_away_from_cuts(z, cuts, y_eps=1e-2, x_pad=0.0):
61
+
62
+ z = numpy.asarray(z, dtype=numpy.complex128).ravel()
63
+ x = numpy.real(z)
64
+ y = numpy.imag(z)
65
+
66
+ keep = numpy.ones(z.size, dtype=bool)
67
+ for a, b in cuts:
68
+ aa = a - x_pad
69
+ bb = b + x_pad
70
+ near_real_cut = (numpy.abs(y) <= y_eps) & (x >= aa) & (x <= bb)
71
+ keep &= ~near_real_cut
72
+
73
+ return z[keep]
74
+
75
+
76
+ # ======
77
+ # powers
78
+ # ======
79
+
80
+ def powers(x, deg):
81
+
82
+ n = x.size
83
+ xp = numpy.ones((n, deg + 1), dtype=complex)
84
+ for k in range(1, deg + 1):
85
+ xp[:, k] = xp[:, k - 1] * x
86
+ return xp
87
+
88
+
89
+ # =======================
90
+ # fit polynomial relation
91
+ # =======================
92
+
93
+ def fit_polynomial_relation(z, m, s, deg_z, ridge_lambda=0.0, weights=None,
94
+ triangular=None):
95
+
96
+ z = numpy.asarray(z, dtype=complex).ravel()
97
+ m = numpy.asarray(m, dtype=complex).ravel()
98
+
99
+ if z.size != m.size:
100
+ raise ValueError('z and m must have the same size.')
101
+ if s < 1:
102
+ raise ValueError('s must be >= 1.')
103
+ if deg_z < 0:
104
+ raise ValueError('deg_z must be >= 0.')
105
+
106
+ zp = powers(z, deg_z)
107
+ mp = powers(m, s)
108
+
109
+ if weights is None:
110
+ w = None
111
+ else:
112
+ w = numpy.asarray(weights, dtype=float).ravel()
113
+ if w.size != z.size:
114
+ raise ValueError('weights must have the same size as z.')
115
+ w = numpy.sqrt(numpy.maximum(w, 0.0))
116
+
117
+ tri = None
118
+ if triangular is not None:
119
+ tri = str(triangular).strip().lower()
120
+ if tri in ['none', '']:
121
+ tri = None
122
+
123
+ if tri is None:
124
+ pairs = [(i, j) for j in range(s + 1)
125
+ for i in range(deg_z + 1)]
126
+
127
+ elif tri in ['lower', 'l']:
128
+ pairs = [(i, j) for j in range(s + 1)
129
+ for i in range(deg_z + 1) if i >= j]
130
+
131
+ elif tri in ['upper', 'u']:
132
+ pairs = [(i, j) for j in range(s + 1)
133
+ for i in range(deg_z + 1) if i <= j]
134
+
135
+ elif tri in ['antidiag', 'anti', 'antidiagonal', 'ad']:
136
+ pairs = [(i, j) for j in range(s + 1)
137
+ for i in range(deg_z + 1) if (i + j) <= deg_z]
138
+
139
+ if len(pairs) == 0:
140
+ raise ValueError('antidiag constraint removed all coefficients.')
141
+ else:
142
+ raise ValueError("triangular must be None, 'lower', 'upper', or " +
143
+ "'antidiag'.")
144
+
145
+ n_coef = len(pairs)
146
+ A = numpy.empty((z.size, n_coef), dtype=complex)
147
+
148
+ for k, (i, j) in enumerate(pairs):
149
+ A[:, k] = zp[:, i] * mp[:, j]
150
+
151
+ if w is not None:
152
+ A = A * w[:, None]
153
+
154
+ s_col = numpy.max(numpy.abs(A), axis=0)
155
+ s_col[s_col == 0.0] = 1.0
156
+ As = A / s_col[None, :]
157
+
158
+ if ridge_lambda > 0.0:
159
+ L = numpy.sqrt(ridge_lambda) * numpy.eye(n_coef, dtype=complex)
160
+ As = numpy.vstack([As, L])
161
+
162
+ _, _, vh = numpy.linalg.svd(As, full_matrices=False)
163
+ coef_scaled = vh[-1, :]
164
+ coef = coef_scaled / s_col
165
+
166
+ full = numpy.zeros((deg_z + 1, s + 1), dtype=complex)
167
+ for k, (i, j) in enumerate(pairs):
168
+ full[i, j] = coef[k]
169
+
170
+ return full
171
+
172
+
173
+ # ======
174
+ # eval P
175
+ # ======
176
+
177
+ def eval_P(z, m, a_coeffs):
178
+
179
+ z = numpy.asarray(z, dtype=complex)
180
+ m = numpy.asarray(m, dtype=complex)
181
+ deg_z = int(a_coeffs.shape[0] - 1)
182
+ s = int(a_coeffs.shape[1] - 1)
183
+
184
+ shp = numpy.broadcast(z, m).shape
185
+ zz = numpy.broadcast_to(z, shp).ravel()
186
+ mm = numpy.broadcast_to(m, shp).ravel()
187
+
188
+ zp = powers(zz, deg_z)
189
+ mp = powers(mm, s)
190
+
191
+ P = numpy.zeros(zz.size, dtype=complex)
192
+ for j in range(s + 1):
193
+ aj = zp @ a_coeffs[:, j]
194
+ P = P + aj * mp[:, j]
195
+
196
+ return P.reshape(shp)
197
+
198
+
199
+ # ==============
200
+ # poly coef in m
201
+ # ==============
202
+
203
+ def _poly_coef_in_m(z, a_coeffs):
204
+
205
+ z = numpy.asarray(z, dtype=complex).ravel()
206
+ deg_z = int(a_coeffs.shape[0] - 1)
207
+ s = int(a_coeffs.shape[1] - 1)
208
+ zp = powers(z, deg_z)
209
+
210
+ c = numpy.empty((z.size, s + 1), dtype=complex)
211
+ for j in range(s + 1):
212
+ c[:, j] = zp @ a_coeffs[:, j]
213
+ return c
214
+
215
+
216
+ # ==============
217
+ # root quadratic
218
+ # ==============
219
+
220
+ def _roots_quadratic(c0, c1, c2):
221
+
222
+ disc = c1 * c1 - 4.0 * c2 * c0
223
+ sq = numpy.sqrt(disc)
224
+ den = 2.0 * c2
225
+
226
+ r1 = (-c1 + sq) / den
227
+ r2 = (-c1 - sq) / den
228
+ return numpy.stack([r1, r2], axis=1)
229
+
230
+
231
+ # ============
232
+ # cbrt complex
233
+ # ============
234
+
235
+ def _cbrt_complex(z):
236
+
237
+ z = numpy.asarray(z, dtype=complex)
238
+ r = numpy.abs(z)
239
+ th = numpy.angle(z)
240
+ return (r ** (1.0 / 3.0)) * numpy.exp(1j * th / 3.0)
241
+
242
+
243
+ # ==========
244
+ # root cubic
245
+ # ==========
246
+
247
+ def _roots_cubic(c0, c1, c2, c3):
248
+
249
+ c0 = numpy.asarray(c0, dtype=complex)
250
+ c1 = numpy.asarray(c1, dtype=complex)
251
+ c2 = numpy.asarray(c2, dtype=complex)
252
+ c3 = numpy.asarray(c3, dtype=complex)
253
+
254
+ a = c2 / c3
255
+ b = c1 / c3
256
+ c = c0 / c3
257
+
258
+ p = b - (a * a) / 3.0
259
+ q = (2.0 * a * a * a) / 27.0 - (a * b) / 3.0 + c
260
+
261
+ Delta = (q * q) / 4.0 + (p * p * p) / 27.0
262
+ sqrtD = numpy.sqrt(Delta)
263
+
264
+ A = -q / 2.0 + sqrtD
265
+ u = _cbrt_complex(A)
266
+
267
+ eps = 1e-30
268
+ small = numpy.abs(u) < eps
269
+ if numpy.any(small):
270
+ u2 = _cbrt_complex(-q / 2.0 - sqrtD)
271
+ u = numpy.where(small, u2, u)
272
+
273
+ small = numpy.abs(u) < eps
274
+ v = numpy.empty_like(u)
275
+ v[~small] = -p[~small] / (3.0 * u[~small])
276
+ v[small] = _cbrt_complex(-q[small])
277
+
278
+ y1 = u + v
279
+ w = complex(-0.5, numpy.sqrt(3.0) / 2.0)
280
+ y2 = w * u + numpy.conjugate(w) * v
281
+ y3 = numpy.conjugate(w) * u + w * v
282
+
283
+ x1 = y1 - a / 3.0
284
+ x2 = y2 - a / 3.0
285
+ x3 = y3 - a / 3.0
286
+
287
+ return numpy.stack([x1, x2, x3], axis=1)
288
+
289
+
290
+ # ==========
291
+ # eval roots
292
+ # ==========
293
+
294
+ def eval_roots(z, a_coeffs):
295
+
296
+ z = numpy.asarray(z, dtype=complex).ravel()
297
+ c = _poly_coef_in_m(z, a_coeffs)
298
+
299
+ s = int(c.shape[1] - 1)
300
+ if s == 1:
301
+ m = -c[:, 0] / c[:, 1]
302
+ return m[:, None]
303
+
304
+ if s == 2:
305
+ return _roots_quadratic(c[:, 0], c[:, 1], c[:, 2])
306
+
307
+ if s == 3:
308
+ return _roots_cubic(c[:, 0], c[:, 1], c[:, 2], c[:, 3])
309
+
310
+ roots = numpy.empty((z.size, s), dtype=complex)
311
+ for i in range(z.size):
312
+ roots[i, :] = numpy.roots(c[i, ::-1])
313
+ return roots
314
+
315
+
316
+ # =======================
317
+ # track one sheet on grid
318
+ # =======================
319
+
320
+ def track_one_sheet_on_grid(z, roots, sheet_seed, cuts=None, i0=None, j0=None):
321
+ z = numpy.asarray(z)
322
+ n_y, n_x = z.shape
323
+ s = roots.shape[1]
324
+ if s < 1:
325
+ raise ValueError("s must be >= 1.")
326
+
327
+ R = roots.reshape((n_y, n_x, s))
328
+
329
+ if i0 is None:
330
+ ycol = numpy.imag(z[:, 0])
331
+ pos = numpy.where(ycol > 0.0)[0]
332
+ i0 = int(pos[0]) if pos.size > 0 else (n_y // 2)
333
+
334
+ if j0 is None:
335
+ j0 = n_x // 2
336
+
337
+ seed_imag = float(numpy.imag(sheet_seed))
338
+ cand0 = R[i0, j0, :]
339
+ idx0 = int(numpy.argmin(numpy.abs(cand0 - sheet_seed)))
340
+
341
+ sheet = numpy.full((n_y, n_x), numpy.nan + 1j * numpy.nan, dtype=complex)
342
+ sheet[i0, j0] = cand0[idx0]
343
+
344
+ visited = numpy.zeros((n_y, n_x), dtype=bool)
345
+ q_i = numpy.empty(n_y * n_x, dtype=int)
346
+ q_j = numpy.empty(n_y * n_x, dtype=int)
347
+
348
+ head = 0
349
+ tail = 0
350
+ q_i[tail] = i0
351
+ q_j[tail] = j0
352
+ tail += 1
353
+ visited[i0, j0] = True
354
+
355
+ neighbors = [(-1, 0), (1, 0), (0, -1), (0, 1)]
356
+
357
+ y_unique = numpy.unique(numpy.imag(z[:, 0]))
358
+ if y_unique.size >= 2:
359
+ dy = float(numpy.min(numpy.diff(y_unique)))
360
+ y_eps = 0.49 * dy
361
+ else:
362
+ y_eps = 0.0
363
+
364
+ def crosses_cut(x_mid):
365
+ if cuts is None:
366
+ return False
367
+ for a, b in cuts:
368
+ if a <= x_mid <= b:
369
+ return True
370
+ return False
371
+
372
+ while head < tail:
373
+ i = int(q_i[head])
374
+ j = int(q_j[head])
375
+ head += 1
376
+
377
+ m_prev = sheet[i, j]
378
+ y1 = float(numpy.imag(z[i, j]))
379
+ x1 = float(numpy.real(z[i, j]))
380
+
381
+ for di, dj in neighbors:
382
+ i2 = i + di
383
+ j2 = j + dj
384
+ if i2 < 0 or i2 >= n_y or j2 < 0 or j2 >= n_x:
385
+ continue
386
+ if visited[i2, j2]:
387
+ continue
388
+
389
+ y2 = float(numpy.imag(z[i2, j2]))
390
+ x2 = float(numpy.real(z[i2, j2]))
391
+
392
+ if cuts is not None:
393
+ if (y1 > y_eps and y2 < -y_eps) or \
394
+ (y1 < -y_eps and y2 > y_eps):
395
+ x_mid = 0.5 * (x1 + x2)
396
+ if crosses_cut(x_mid):
397
+ continue
398
+
399
+ cand = R[i2, j2, :]
400
+ d = numpy.abs(cand - m_prev)
401
+ idx = int(numpy.argmin(d))
402
+
403
+ if seed_imag != 0.0:
404
+ y_sign = 1.0 if y2 >= 0.0 else -1.0
405
+ target = float(numpy.sign(seed_imag) * y_sign)
406
+ if target != 0.0:
407
+ sgn = numpy.sign(numpy.imag(cand))
408
+ ok = (sgn == numpy.sign(target)) | (sgn == 0.0)
409
+ if numpy.any(ok):
410
+ ok_idx = numpy.where(ok)[0]
411
+ idx = int(ok_idx[numpy.argmin(d[ok])])
412
+
413
+ sheet[i2, j2] = cand[idx]
414
+ visited[i2, j2] = True
415
+ q_i[tail] = i2
416
+ q_j[tail] = j2
417
+ tail += 1
418
+
419
+ return sheet
420
+
421
+
422
+ # =======================
423
+ # build sheets from roots
424
+ # =======================
425
+
426
+ def build_sheets_from_roots(z, roots, m1, cuts=None, i0=None, j0=None):
427
+ z = numpy.asarray(z)
428
+ m1 = numpy.asarray(m1)
429
+
430
+ n_y, n_x = z.shape
431
+ s = roots.shape[1]
432
+ if s < 1:
433
+ raise ValueError("s must be >= 1.")
434
+
435
+ if i0 is None:
436
+ ycol = numpy.imag(z[:, 0])
437
+ pos = numpy.where(ycol > 0.0)[0]
438
+ i0 = int(pos[0]) if pos.size > 0 else (n_y // 2)
439
+
440
+ if j0 is None:
441
+ j0 = n_x // 2
442
+
443
+ R0 = roots.reshape((n_y, n_x, s))[i0, j0, :]
444
+ idx_phys = int(numpy.argmin(numpy.abs(R0 - m1[i0, j0])))
445
+
446
+ idxs = list(range(s))
447
+ idxs.sort(key=lambda k: numpy.imag(R0[k]))
448
+
449
+ seeds = [R0[k] for k in idxs]
450
+ sheets = [track_one_sheet_on_grid(z, roots, seed, cuts=cuts, i0=i0, j0=j0)
451
+ for seed in seeds]
452
+
453
+ phys_pos = int(numpy.where(numpy.array(idxs, dtype=int) == idx_phys)[0][0])
454
+ if phys_pos != 0:
455
+ sheets[0], sheets[phys_pos] = sheets[phys_pos], sheets[0]
456
+ idxs[0], idxs[phys_pos] = idxs[phys_pos], idxs[0]
457
+
458
+ if cuts is not None:
459
+ y_unique = numpy.unique(numpy.imag(z[:, 0]))
460
+ if y_unique.size >= 2:
461
+ dy = float(numpy.min(numpy.diff(y_unique)))
462
+ eps_y = 0.49 * dy
463
+ else:
464
+ eps_y = 0.0
465
+
466
+ i_cut = numpy.where(numpy.abs(numpy.imag(z[:, 0])) <= eps_y)[0]
467
+ if i_cut.size > 0:
468
+ i_cut = int(i_cut[numpy.argmin(numpy.abs(
469
+ numpy.imag(z[i_cut, 0])))])
470
+
471
+ X = numpy.real(z[i_cut, :])
472
+ on_cut = numpy.zeros(n_x, dtype=bool)
473
+ for j in range(n_x):
474
+ xj = float(X[j])
475
+ for a, b in cuts:
476
+ if a <= xj <= b:
477
+ on_cut[j] = True
478
+ break
479
+
480
+ sheets[0][i_cut, on_cut] = m1[i_cut, on_cut]
481
+
482
+ ycol = numpy.imag(z[:, 0])
483
+ y_unique = numpy.unique(ycol)
484
+ if y_unique.size >= 2:
485
+ dy = float(numpy.min(numpy.diff(y_unique)))
486
+ eps_y = 1.1 * dy
487
+ else:
488
+ eps_y = 0.0
489
+
490
+ i_band = numpy.where(numpy.abs(ycol) <= eps_y)[0]
491
+ i_up = numpy.where(ycol > eps_y)[0]
492
+ i_dn = numpy.where(ycol < -eps_y)[0]
493
+ if (i_band.size > 0) and (i_up.size > 0) and (i_dn.size > 0):
494
+ i_up = int(i_up[0])
495
+ i_dn = int(i_dn[-1])
496
+ for r in range(1, len(sheets)):
497
+ for i in i_band:
498
+ if ycol[i] >= 0.0:
499
+ sheets[r][i, :] = sheets[r][i_up, :]
500
+ else:
501
+ sheets[r][i, :] = sheets[r][i_dn, :]
502
+
503
+ return sheets, idxs