freealg 0.6.3__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.
- {freealg-0.6.3 → freealg-0.7.0}/PKG-INFO +1 -1
- freealg-0.7.0/freealg/__init__.py +20 -0
- freealg-0.7.0/freealg/__version__.py +1 -0
- freealg-0.7.0/freealg/_algebraic_form/__init__.py +11 -0
- freealg-0.7.0/freealg/_algebraic_form/_continuation_algebraic.py +503 -0
- freealg-0.7.0/freealg/_algebraic_form/_decompress.py +648 -0
- freealg-0.7.0/freealg/_algebraic_form/_edge.py +352 -0
- freealg-0.7.0/freealg/_algebraic_form/_sheets_util.py +145 -0
- freealg-0.7.0/freealg/_algebraic_form/algebraic_form.py +987 -0
- {freealg-0.6.3/freealg → freealg-0.7.0/freealg/_freeform}/__init__.py +3 -6
- freealg-0.6.3/freealg/_util.py → freealg-0.7.0/freealg/_freeform/_density_util.py +1 -57
- {freealg-0.6.3/freealg → freealg-0.7.0/freealg/_freeform}/_linalg.py +1 -1
- {freealg-0.6.3/freealg → freealg-0.7.0/freealg/_freeform}/freeform.py +2 -1
- freealg-0.7.0/freealg/_geometric_form/__init__.py +13 -0
- freealg-0.7.0/freealg/_geometric_form/_continuation_genus0.py +175 -0
- freealg-0.7.0/freealg/_geometric_form/_continuation_genus1.py +275 -0
- freealg-0.7.0/freealg/_geometric_form/_elliptic_functions.py +174 -0
- freealg-0.7.0/freealg/_geometric_form/_sphere_maps.py +63 -0
- freealg-0.7.0/freealg/_geometric_form/_torus_maps.py +118 -0
- freealg-0.7.0/freealg/_geometric_form/geometric_form.py +1094 -0
- freealg-0.7.0/freealg/_util.py +72 -0
- {freealg-0.6.3 → freealg-0.7.0}/freealg/distributions/__init__.py +5 -1
- freealg-0.7.0/freealg/distributions/_chiral_block.py +440 -0
- freealg-0.7.0/freealg/distributions/_deformed_marchenko_pastur.py +617 -0
- freealg-0.7.0/freealg/distributions/_deformed_wigner.py +312 -0
- {freealg-0.6.3 → freealg-0.7.0}/freealg/distributions/_marchenko_pastur.py +197 -80
- freealg-0.7.0/freealg/visualization/__init__.py +12 -0
- freealg-0.7.0/freealg/visualization/_glue_util.py +32 -0
- freealg-0.7.0/freealg/visualization/_rgb_hsv.py +125 -0
- {freealg-0.6.3 → freealg-0.7.0}/freealg.egg-info/PKG-INFO +1 -1
- freealg-0.7.0/freealg.egg-info/SOURCES.txt +56 -0
- freealg-0.6.3/freealg/__version__.py +0 -1
- freealg-0.6.3/freealg.egg-info/SOURCES.txt +0 -35
- {freealg-0.6.3 → freealg-0.7.0}/AUTHORS.txt +0 -0
- {freealg-0.6.3 → freealg-0.7.0}/CHANGELOG.rst +0 -0
- {freealg-0.6.3 → freealg-0.7.0}/LICENSE.txt +0 -0
- {freealg-0.6.3 → freealg-0.7.0}/MANIFEST.in +0 -0
- {freealg-0.6.3 → freealg-0.7.0}/README.rst +0 -0
- {freealg-0.6.3/freealg → freealg-0.7.0/freealg/_freeform}/_chebyshev.py +0 -0
- {freealg-0.6.3/freealg → freealg-0.7.0/freealg/_freeform}/_damp.py +0 -0
- {freealg-0.6.3/freealg → freealg-0.7.0/freealg/_freeform}/_decompress.py +0 -0
- {freealg-0.6.3/freealg → freealg-0.7.0/freealg/_freeform}/_jacobi.py +0 -0
- {freealg-0.6.3/freealg → freealg-0.7.0/freealg/_freeform}/_pade.py +0 -0
- {freealg-0.6.3/freealg → freealg-0.7.0/freealg/_freeform}/_plot_util.py +0 -0
- {freealg-0.6.3/freealg → freealg-0.7.0/freealg/_freeform}/_sample.py +0 -0
- {freealg-0.6.3/freealg → freealg-0.7.0/freealg/_freeform}/_series.py +0 -0
- {freealg-0.6.3/freealg → freealg-0.7.0/freealg/_freeform}/_support.py +0 -0
- {freealg-0.6.3 → freealg-0.7.0}/freealg/distributions/_kesten_mckay.py +0 -0
- {freealg-0.6.3 → freealg-0.7.0}/freealg/distributions/_meixner.py +0 -0
- {freealg-0.6.3 → freealg-0.7.0}/freealg/distributions/_wachter.py +0 -0
- {freealg-0.6.3 → freealg-0.7.0}/freealg/distributions/_wigner.py +0 -0
- {freealg-0.6.3 → freealg-0.7.0}/freealg.egg-info/dependency_links.txt +0 -0
- {freealg-0.6.3 → freealg-0.7.0}/freealg.egg-info/not-zip-safe +0 -0
- {freealg-0.6.3 → freealg-0.7.0}/freealg.egg-info/requires.txt +0 -0
- {freealg-0.6.3 → freealg-0.7.0}/freealg.egg-info/top_level.txt +0 -0
- {freealg-0.6.3 → freealg-0.7.0}/pyproject.toml +0 -0
- {freealg-0.6.3 → freealg-0.7.0}/requirements.txt +0 -0
- {freealg-0.6.3 → freealg-0.7.0}/setup.cfg +0 -0
- {freealg-0.6.3 → freealg-0.7.0}/setup.py +0 -0
|
@@ -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
|