freealg 0.7.17__py3-none-any.whl → 0.7.18__py3-none-any.whl

freealg/_algebraic_form/_decompressible.py

@@ -0,0 +1,292 @@
+# SPDX-FileCopyrightText: Copyright 2026, Siavash Ameli <sameli@berkeley.edu>
+# SPDX-License-Identifier: BSD-3-Clause
+# SPDX-FileType: SOURCE
+#
+# This program is free software: you can redistribute it and/or modify it under
+# the terms of the license found in the LICENSE.txt file in the root directory
+# of this source tree.
+
+
+# =======
+# Imports
+# =======
+
+import numpy
+
+__all__ = ['precheck_laurent']
+
+
+# =====
+# L add
+# =====
+
+def L_add(A, B):
+
+    C = dict(A)
+    for p, c in B.items():
+        C[p] = C.get(p, 0.0) + c
+        if abs(C[p]) == 0:
+            C.pop(p, None)
+    return C
+
+
+# =======
+# L scale
+# =======
+
+def L_scale(A, s):
+    return {p: s*c for p, c in A.items()}
+
+
+# =====
+# L mul
+# =====
+
+def L_mul(A, B, pmin, pmax):
+
+    C = {}
+    for pa, ca in A.items():
+        for pb, cb in B.items():
+            p = pa + pb
+            if p < pmin or p > pmax:
+                continue
+            C[p] = C.get(p, 0.0) + ca*cb
+
+    return C
+
+
+# =====
+# L pow
+# =====
+
+def L_pow(A, k, pmin, pmax):
+
+    if k == 0:
+        return {0: 1.0+0.0j}
+    if k == 1:
+        return dict(A)
+    # fast exponentiation
+    out = {0: 1.0+0.0j}
+    base = dict(A)
+    e = k
+    while e > 0:
+        if e & 1:
+            out = L_mul(out, base, pmin, pmax)
+        e >>= 1
+        if e:
+            base = L_mul(base, base, pmin, pmax)
+    return out
+
+
+# ============
+# L inv series
+# ============
+
+def L_inv_series(A, pmin, pmax):
+    """
+    Invert a *power series* around power 0: requires A[0] != 0 and no negative
+    powers.
+
+    Returns power series B with nonnegative powers only.
+    """
+
+    if any(p < 0 for p in A.keys()):
+        raise ValueError("L_inv_series expects no negative powers")
+
+    a0 = A.get(0, 0.0)
+    if abs(a0) < 1e-18:
+        raise ValueError("Need nonzero constant term to invert")
+
+    # compute up to pmax (nonnegative)
+    B = {0: 1.0/a0}
+    for n in range(1, pmax+1):
+        s = 0.0+0.0j
+        for k in range(1, n+1):
+            s += A.get(k, 0.0) * B.get(n-k, 0.0)
+        B[n] = -s/a0
+
+    # clip
+    B = {p: c for p, c in B.items() if 0 <= p <= pmax}
+    return B
+
+
+# ==================
+# Ptau Laurent alpha
+# ==================
+
+def Ptau_Laurent_alpha(a, tau, alpha, mu, pmin, pmax):
+    """
+    Build Laurent series of P_tau(z,m) at z=1/w,
+    m(w)=-(alpha w + mu1 w^2 + ...).
+
+    Returns dict power->coeff for powers in [pmin,pmax].
+    """
+
+    # m(w)
+    m = {1: -(alpha+0.0j)}
+    for k, muk in enumerate(mu, start=1):
+        m[k+1] = -(muk+0.0j)
+
+    # y = tau m
+    y = L_scale(m, tau)
+
+    # m = w*m1(w) with m1(w) = -(alpha + mu1 w + mu2 w^2 + ...)
+    m1 = {0: -(alpha+0.0j)}
+    for k, muk in enumerate(mu, start=1):
+        m1[k] = -(muk+0.0j)
+
+    inv_m1 = L_inv_series(m1, pmin=0, pmax=max(0, pmax))
+
+    # 1/m = w^{-1} * inv_m1
+    inv_m = {p-1: c for p, c in inv_m1.items() if (p-1) >= pmin}
+
+    # z = 1/w
+    z = {-1: 1.0+0.0j}
+
+    c = (1.0 - 1.0/tau)
+    zeta = L_add(z, L_scale(inv_m, c))  # z + c*(1/m)
+
+    deg_z = a.shape[0]-1
+    deg_m = a.shape[1]-1
+
+    out = {}
+    zeta_pows = [L_pow(zeta, i, pmin, pmax) for i in range(deg_z+1)]
+    y_pows = [L_pow(y,    j, pmin, pmax) for j in range(deg_m+1)]
+
+    for i in range(deg_z+1):
+        for j in range(deg_m+1):
+            coeff = a[i, j]
+            if abs(coeff) < 1e-18:
+                continue
+            term = L_mul(zeta_pows[i], y_pows[j], pmin, pmax)
+            out = L_add(out, L_scale(term, coeff))
+
+    return out
+
+
+# ======================
+# solve laurent alpha ls
+# ======================
+
+def solve_laurent_alpha_ls(a, tau, K=8, L=None, max_iter=40, tol_step=1e-12):
+    """
+    Solve for x = [alpha, mu1..muK] so that Laurent coeffs vanish
+    for powers p in [-L, ..., K]. Uses LS on both Re and Im parts (robust).
+    """
+
+    deg_z = a.shape[0]-1
+    if L is None:
+        L = deg_z + 2
+
+    pmin, pmax = -L, K
+    powers = list(range(pmin, pmax+1))
+
+    # unknowns: alpha + mu1..muK
+    x = numpy.zeros(K+1, dtype=numpy.float64)
+    x[0] = 1.0  # alpha init
+
+    def build_out(xx):
+        alpha = float(xx[0])
+        mu = xx[1:]
+        return Ptau_Laurent_alpha(a, tau, alpha, mu, pmin, pmax)
+
+    # LS Newton / Gauss-Newton
+    for it in range(max_iter):
+        out = build_out(x)
+
+        # complex residual vector r_p = coeff(p)
+        r = numpy.array([out.get(p, 0.0+0.0j) for p in powers],
+                        dtype=numpy.complex128)
+
+        # stack real+imag (this is the key fix)
+        F = numpy.concatenate([r.real, r.imag], axis=0)
+        nrm = numpy.linalg.norm(F)
+
+        if nrm < 1e-12:
+            return x, True, out, powers
+
+        # Jacobian by FD
+        eps = 1e-6
+        J = numpy.zeros((F.size, x.size), dtype=numpy.float64)
+        for k in range(x.size):
+            x2 = x.copy()
+            x2[k] += eps
+            out2 = build_out(x2)
+            r2 = numpy.array([out2.get(p, 0.0+0.0j) for p in powers],
+                             dtype=numpy.complex128)
+            F2 = numpy.concatenate([r2.real, r2.imag], axis=0)
+            J[:, k] = (F2 - F) / eps
+
+        # LS step
+        dx, *_ = numpy.linalg.lstsq(J, -F, rcond=None)
+
+        if numpy.linalg.norm(dx) < tol_step:
+            return x, False, out, powers
+
+        x += dx
+
+    return x, False, out, powers
+
+
+# ================
+# precheck laurent
+# ================
+
+def precheck_laurent(a, tau, K_list=(6, 8, 10), L=3, tol=1e-8, verbose=True):
+    """
+    For fixed tau, try several K and pick the best (smallest max |coeff|).
+    Returns dict with bestK, alpha, max_abs, ok, and per-K details.
+    """
+
+    deg_z = a.shape[0]-1
+    if L is None:
+        L = deg_z + 2
+
+    best = None
+    perK = []
+
+    for K in K_list:
+        x, ok_solve, out, powers = solve_laurent_alpha_ls(a, tau, K=K, L=L)
+        alpha = x[0]
+        coeffs = numpy.array([out.get(p, 0.0+0.0j) for p in powers],
+                             dtype=numpy.complex128)
+        max_abs = float(numpy.max(numpy.abs(coeffs)))
+        worst_p = powers[int(numpy.argmax(numpy.abs(coeffs)))]
+        perK.append((K, alpha, max_abs, worst_p, ok_solve))
+
+        if best is None or max_abs < best["max_abs"]:
+            best = {
+                "K": K,
+                "alpha": alpha,
+                "max_abs": max_abs,
+                "worst_p": worst_p,
+                "ok_solve": ok_solve
+            }
+
+    alphas = numpy.array([row[1] for row in perK], dtype=float)  # alpha per K
+    alpha_std = float(numpy.std(alphas))
+    alpha_span = float(numpy.max(alphas) - numpy.min(alphas))
+
+    ok = (best["max_abs"] <= tol) and best["ok_solve"]
+    ok = ok and (alpha_std < 1e-3)   # or use alpha_span < 3e-3
+
+    if verbose:
+        print(f"--- tau={tau} --- ok={ok} bestK={best['K']} "
+              f"max_abs={best['max_abs']:.3e} "
+              f"alpha={best['alpha']:.12g} worst_p={best['worst_p']}")
+
+        print(f"  alpha_std={alpha_std:.3e}, alpha_span={alpha_span:.3e}, "
+              f"alphas={alphas}")
+
+        for (K, alpha, max_abs, worst_p, ok_solve) in perK:
+            print(f"  K={K:2d} max_abs={max_abs:.3e} worst_p={worst_p:2d} "
+                  f"alpha={alpha:.12g} solve_ok={ok_solve}")
+
+    res = {
+        "best": best,
+        "perK": perK,
+        "alpha_std": alpha_std,
+        "alpha_span": alpha_span
+    }
+
+    return ok, res

freealg/_algebraic_form/_edge.py

@@ -22,7 +22,7 @@ __all__ = ['evolve_edges', 'merge_edges']
 # edge newton step
 # ================
 
-def _edge_newton_step(t, zeta, y, a_coeffs, max_iter=30, tol=1e-12):
+def _edge_newton_step(t, zeta, y, coeffs, max_iter=30, tol=1e-12):
     """
     """
 
@@ -30,7 +30,7 @@ def _edge_newton_step(t, zeta, y, a_coeffs, max_iter=30, tol=1e-12):
     c = tau - 1.0
 
     for _ in range(max_iter):
-        P, Pz, Py = eval_P_partials(zeta, y, a_coeffs)
+        P, Pz, Py = eval_P_partials(zeta, y, coeffs)
 
         # F1 = P(zeta,y)
         F1 = complex(P)
@@ -45,11 +45,11 @@ def _edge_newton_step(t, zeta, y, a_coeffs, max_iter=30, tol=1e-12):
         eps_z = 1e-8 * (1.0 + abs(zeta))
         eps_y = 1e-8 * (1.0 + abs(y))
 
-        Pp, Pzp, Pyp = eval_P_partials(zeta + eps_z, y, a_coeffs)
+        Pp, Pzp, Pyp = eval_P_partials(zeta + eps_z, y, coeffs)
         F1_zp = (complex(Pp) - F1) / eps_z
         F2_zp = (complex((y * y) * Pyp - c * Pzp) - F2) / eps_z
 
-        Pp, Pzp, Pyp = eval_P_partials(zeta, y + eps_y, a_coeffs)
+        Pp, Pzp, Pyp = eval_P_partials(zeta, y + eps_y, coeffs)
         F1_yp = (complex(Pp) - F1) / eps_y
         F2_yp = (complex(((y + eps_y) * (y + eps_y)) * Pyp - c * Pzp) - F2) / \
             eps_y
@@ -102,7 +102,7 @@ def _pick_physical_root(z, roots):
 # init edge point from support
 # ============================
 
-def _init_edge_point_from_support(x_edge, a_coeffs, eta=1e-3):
+def _init_edge_point_from_support(x_edge, coeffs, eta=1e-3):
     """
     Initialize (zeta,y) at t=0 for an edge near x_edge.
 
@@ -111,7 +111,7 @@ def _init_edge_point_from_support(x_edge, a_coeffs, eta=1e-3):
     """
 
     z = complex(x_edge + 1j * eta)
-    roots = eval_roots(numpy.array([z]), a_coeffs)[0]
+    roots = eval_roots(numpy.array([z]), coeffs)[0]
     y = _pick_physical_root(z, roots)
 
     # Move zeta to real axis as initial guess
@@ -119,7 +119,7 @@ def _init_edge_point_from_support(x_edge, a_coeffs, eta=1e-3):
 
     # Refine zeta,y to satisfy P=0 and Py=0 at t=0 (branch point)
     # This uses the same Newton system with c=0, i.e. F2 = y^2 Py.
-    zeta, y, ok = _edge_newton_step(0.0, zeta, y, a_coeffs, max_iter=50,
+    zeta, y, ok = _edge_newton_step(0.0, zeta, y, coeffs, max_iter=50,
                                     tol=1e-10)
 
     return zeta, y, ok
@@ -131,7 +131,7 @@ def _init_edge_point_from_support(x_edge, a_coeffs, eta=1e-3):
 
 def evolve_edges(
         t_grid,
-        a_coeffs,
+        coeffs,
         support=None,
         eta=1e-3,
         dt_max=0.1,
@@ -185,7 +185,7 @@ def evolve_edges(
     y = numpy.empty(m, dtype=numpy.complex128)
 
     for j in range(m):
-        z0, y0, ok0 = _init_edge_point_from_support(endpoints0[j], a_coeffs,
+        z0, y0, ok0 = _init_edge_point_from_support(endpoints0[j], coeffs,
                                                     eta=eta)
         zeta[j] = z0
         y[j] = y0
@@ -210,7 +210,7 @@ def evolve_edges(
             t = t0 + dt * (ks / float(n_sub))
             for j in range(m):
                 zeta[j], y[j], okj = _edge_newton_step(
-                    t, zeta[j], y[j], a_coeffs, max_iter=max_iter, tol=tol
+                    t, zeta[j], y[j], coeffs, max_iter=max_iter, tol=tol
                 )
                 ok[it, j] = okj
 

freealg/_algebraic_form/_homotopy4.py

@@ -23,13 +23,13 @@ __all__ = ["StieltjesPoly"]
 # Poly -> roots in m for z
 # =========================
 
-def _poly_m_coeffs(a_coeffs, z):
+def _poly_m_coeffs(coeffs, z):
     """Return coefficients b_j for \sum_j b_j m^j = 0 at fixed z.
 
-    a_coeffs[i,j] is coeff of z^i m^j.
+    coeffs[i,j] is coeff of z^i m^j.
     Returns b of length (deg_m+1) with b[j] = \sum_i a[i,j] z^i.
     """
-    a = numpy.asarray(a_coeffs, dtype=numpy.complex128)
+    a = numpy.asarray(coeffs, dtype=numpy.complex128)
     deg_z = a.shape[0] - 1
     deg_m = a.shape[1] - 1
 
@@ -47,9 +47,9 @@ def _poly_m_coeffs(a_coeffs, z):
     return b
 
 
-def _roots_m(a_coeffs, z):
+def _roots_m(coeffs, z):
     """All algebraic roots in m at fixed z."""
-    b = _poly_m_coeffs(a_coeffs, z)
+    b = _poly_m_coeffs(coeffs, z)
 
     # Drop leading zeros in highest power to keep numpy.roots stable
     # numpy.roots expects highest degree first.
@@ -239,9 +239,9 @@ def _viterbi_track(z_list, roots_list, mL, mR, *,
 class StieltjesPoly(object):
     """Callable m(z) for P(z,m)=0 using robust branch selection."""
 
-    def __init__(self, a_coeffs, *,
+    def __init__(self, coeffs, *,
                  viterbi_opt=None):
-        self.a_coeffs = numpy.asarray(a_coeffs, dtype=numpy.complex128)
+        self.coeffs = numpy.asarray(coeffs, dtype=numpy.complex128)
         self.viterbi_opt = dict(viterbi_opt or {})
 
     # ----------
@@ -249,7 +249,7 @@ class StieltjesPoly(object):
     # ----------
 
     def evaluate_scalar(self, z, target=None):
-        r = _roots_m(self.a_coeffs, z)
+        r = _roots_m(self.coeffs, z)
         return _pick_physical_root_scalar(z, r, target=target)
 
     # ---------------
@@ -267,7 +267,7 @@ class StieltjesPoly(object):
             z_list = z.ravel()
 
             # roots for each point
-            roots_list = [_roots_m(self.a_coeffs, zi) for zi in z_list]
+            roots_list = [_roots_m(self.coeffs, zi) for zi in z_list]
 
             # boundary anchors via scalar selection
             mL = self.evaluate_scalar(z_list[0])

freealg/_algebraic_form/_homotopy5.py

@@ -8,8 +8,8 @@
 import numpy
 
 
-def _poly_coeffs_in_m(a_coeffs, z):
-    a = a_coeffs
+def _poly_coeffs_in_m(coeffs, z):
+    a = coeffs
     dz = a.shape[0] - 1
     s = a.shape[1] - 1
     zp = numpy.array([z**i for i in range(dz + 1)], dtype=numpy.complex128)
@@ -19,8 +19,8 @@ def _poly_coeffs_in_m(a_coeffs, z):
     return coeff_m
 
 
-def _roots_m(a_coeffs, z):
-    coeff_m = _poly_coeffs_in_m(a_coeffs, z)
+def _roots_m(coeffs, z):
+    coeff_m = _poly_coeffs_in_m(coeffs, z)
     c = coeff_m[::-1]
     while c.size > 1 and numpy.abs(c[0]) == 0:
         c = c[1:]
@@ -96,12 +96,12 @@ def _viterbi_1d(z_list, roots_all, *, lam_space, lam_asym,
 class StieltjesPoly(object):
     """Callable m(z) for P(z,m)=0 using robust branch selection."""
 
-    def __init__(self, a_coeffs, *, viterbi_opt=None):
-        self.a_coeffs = numpy.asarray(a_coeffs, dtype=numpy.complex128)
+    def __init__(self, coeffs, *, viterbi_opt=None):
+        self.coeffs = numpy.asarray(coeffs, dtype=numpy.complex128)
         self.viterbi_opt = dict(viterbi_opt or {})
 
     def evaluate_scalar(self, z, target=None):
-        r = _roots_m(self.a_coeffs, z)
+        r = _roots_m(self.coeffs, z)
         if r.size == 0:
             return numpy.nan + 1j * numpy.nan
         tol_im = float(self.viterbi_opt.get("tol_im", 1e-14))
@@ -123,11 +123,11 @@ class StieltjesPoly(object):
 
         if z.ndim == 1 and z.size >= 2:
             z_list = z.ravel()
-            s = self.a_coeffs.shape[1] - 1
+            s = self.coeffs.shape[1] - 1
             roots_all = numpy.empty((z_list.size, s), dtype=numpy.complex128)
             ok_all = numpy.ones(z_list.size, dtype=bool)
             for k in range(z_list.size):
-                r = _roots_m(self.a_coeffs, z_list[k])
+                r = _roots_m(self.coeffs, z_list[k])
                 if r.size != s:
                     ok_all[k] = False
                     if r.size == 0:

freealg/_algebraic_form/_support.py

@@ -15,18 +15,18 @@ import numpy
 import numpy.polynomial.polynomial as poly
 from ._homotopy5 import StieltjesPoly
 
-__all__ = ['compute_support']
+__all__ = ['estimate_support']
 
 
 # =====================
 # poly coeffs in m at z
 # =====================
 
-def _poly_coeffs_in_m_at_z(a_coeffs, z):
-    s = a_coeffs.shape[1] - 1
+def _poly_coeffs_in_m_at_z(coeffs, z):
+    s = coeffs.shape[1] - 1
     a = numpy.empty(s + 1, dtype=numpy.complex128)
     for j in range(s + 1):
-        a[j] = poly.polyval(z, a_coeffs[:, j])
+        a[j] = poly.polyval(z, coeffs[:, j])
     return a
 
 
@@ -34,14 +34,14 @@ def _poly_coeffs_in_m_at_z(a_coeffs, z):
 # P and partials
 # ==============
 
-def _P_and_partials(a_coeffs, z, m):
-    s = a_coeffs.shape[1] - 1
+def _P_and_partials(coeffs, z, m):
+    s = coeffs.shape[1] - 1
 
     a = numpy.empty(s + 1, dtype=numpy.complex128)
     da = numpy.empty(s + 1, dtype=numpy.complex128)
     for j in range(s + 1):
-        a[j] = poly.polyval(z, a_coeffs[:, j])
-        da[j] = poly.polyval(z, poly.polyder(a_coeffs[:, j]))
+        a[j] = poly.polyval(z, coeffs[:, j])
+        da[j] = poly.polyval(z, poly.polyder(coeffs[:, j]))
 
     mpow = 1.0 + 0.0j
     P = 0.0 + 0.0j
@@ -67,13 +67,13 @@ def _P_and_partials(a_coeffs, z, m):
 # newton edge
 # ===========
 
-def _newton_edge(a_coeffs, x0, m0, tol=1e-12, max_iter=50):
+def _newton_edge(coeffs, x0, m0, tol=1e-12, max_iter=50):
     x = float(x0)
     m = float(m0)
 
     for _ in range(max_iter):
         z = x + 0.0j
-        P, Pz, Pm, Pzm, Pmm = _P_and_partials(a_coeffs, z, m)
+        P, Pz, Pm, Pzm, Pmm = _P_and_partials(coeffs, z, m)
 
         f0 = float(numpy.real(P))
         f1 = float(numpy.real(Pm))
@@ -157,12 +157,12 @@ def _bisect_edge(stieltjes_poly, x_lo, x_hi, eta, im_thr, max_iter=60):
 
 
 # ===============
-# compute support
+# estimate support
 # ===============
 
-def compute_support(a_coeffs, x_min, x_max, n_scan=4000, **kwargs):
+def estimate_support(coeffs, x_min, x_max, n_scan=4000, **kwargs):
 
-    a_coeffs = numpy.asarray(a_coeffs, dtype=numpy.complex128)
+    coeffs = numpy.asarray(coeffs, dtype=numpy.complex128)
 
     x_min = float(x_min)
     x_max = float(x_max)
@@ -182,7 +182,7 @@ def compute_support(a_coeffs, x_min, x_max, n_scan=4000, **kwargs):
         'tol_im': 1e-14,
     }
     vopt.update(kwargs.get('viterbi_opt', {}) or {})
-    stieltjes = StieltjesPoly(a_coeffs, viterbi_opt=vopt)
+    stieltjes = StieltjesPoly(coeffs, viterbi_opt=vopt)
 
     x_grid = numpy.linspace(x_min, x_max, n_scan)
     z_grid = x_grid + 1j * eta
@@ -234,18 +234,18 @@ def compute_support(a_coeffs, x_min, x_max, n_scan=4000, **kwargs):
         edges_ref = []
         for x0 in edges:
             m0 = float(numpy.real(stieltjes.evaluate_scalar(x0 + 1j * eta)))
-            xe, _, ok = _newton_edge(a_coeffs, x0, m0, tol=newton_tol)
+            xe, _, ok = _newton_edge(coeffs, x0, m0, tol=newton_tol)
             edges_ref.append(float(xe)
                              if ok and numpy.isfinite(xe) else float(x0))
         edges = _cluster_edges(edges_ref, edge_x_cluster_tol)
 
     edges.sort()
-    support = []
+    est_supp = []
     for k in range(0, edges.size - 1, 2):
         a = float(edges[k])
         b = float(edges[k + 1])
         if b > a:
-            support.append((a, b))
+            est_supp.append((a, b))
 
     info = {
         'x_grid': x_grid,
@@ -254,11 +254,11 @@ def compute_support(a_coeffs, x_min, x_max, n_scan=4000, **kwargs):
         'im_grid': im_grid,
         'im_thr': float(im_thr),
         'edges': edges,
-        'support': support,
+        'est_supp': est_supp,
         'x_min': x_min,
         'x_max': x_max,
         'n_scan': n_scan,
         'scale': scale,
     }
 
-    return support, info
+    return est_supp, info