freealg 0.7.12__py3-none-any.whl → 0.7.15__py3-none-any.whl

freealg/_algebraic_form/_decompress_util.py

@@ -0,0 +1,172 @@
+# 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
+from ._continuation_algebraic import powers
+
+__all__ = ['build_time_grid', 'eval_P_partials']
+
+
+# ===============
+# build time grid
+# ===============
+
+def build_time_grid(sizes, n0, min_n_times=0):
+    """
+    sizes: list/array of requested matrix sizes (e.g. [2000,3000,4000,8000])
+    n0:    initial size (self.n)
+    min_n_times: minimum number of time points to run Newton sweep on
+
+    Returns
+    -------
+    t_all: sorted time grid to run solver on
+    idx_req: indices of requested times inside t_all (same order as sizes)
+    """
+
+    sizes = numpy.asarray(sizes, dtype=float)
+    alpha = sizes / float(n0)
+    t_req = numpy.log(alpha)
+
+    # Always include t=0 and T=max(t_req)
+    T = float(numpy.max(t_req)) if t_req.size else 0.0
+    base = numpy.unique(numpy.r_[0.0, t_req, T])
+    t_all = numpy.sort(base)
+
+    # Add points only if needed: split largest gaps
+    N = int(min_n_times) if min_n_times is not None else 0
+    while t_all.size < N and t_all.size >= 2:
+        gaps = numpy.diff(t_all)
+        k = int(numpy.argmax(gaps))
+        mid = 0.5 * (t_all[k] + t_all[k+1])
+        t_all = numpy.sort(numpy.unique(numpy.r_[t_all, mid]))
+
+    # Map each requested time to an index in t_all (stable, no float drama)
+    # (t_req values came from same construction, so they should match exactly;
+    # still: use searchsorted + assert)
+    idx_req = numpy.searchsorted(t_all, t_req)
+    # optional sanity:
+    # assert numpy.allclose(t_all[idx_req], t_req, rtol=0, atol=0)
+
+    return t_all, idx_req
+
+
+# ===============
+# eval P partials
+# ===============
+
+def eval_P_partials(z, m, a_coeffs):
+    """
+    Evaluate P(z,m) and its partial derivatives dP/dz and dP/dm.
+
+    This assumes P is represented by `a_coeffs` in the monomial basis
+
+        P(z, m) = sum_{j=0..s} a_j(z) * m^j,
+        a_j(z) = sum_{i=0..deg_z} a_coeffs[i, j] * z^i.
+
+    The function returns P, dP/dz, dP/dm with broadcasting over z and m.
+
+    Parameters
+    ----------
+    z : complex or array_like of complex
+        First argument to P.
+    m : complex or array_like of complex
+        Second argument to P. Must be broadcast-compatible with `z`.
+    a_coeffs : ndarray, shape (deg_z+1, s+1)
+        Coefficient matrix for P in the monomial basis.
+
+    Returns
+    -------
+    P : complex or ndarray of complex
+        Value P(z,m).
+    Pz : complex or ndarray of complex
+        Partial derivative dP/dz evaluated at (z,m).
+    Pm : complex or ndarray of complex
+        Partial derivative dP/dm evaluated at (z,m).
+
+    Notes
+    -----
+    For scalar (z,m), this uses Horner evaluation for a_j(z) and then Horner
+    in m. For array inputs, it uses precomputed power tables via `_powers` for
+    simplicity.
+
+    Examples
+    --------
+    .. code-block:: python
+
+        P, Pz, Pm = eval_P_partials(1.0 + 1j, 0.2 + 0.3j, a_coeffs)
+    """
+
+    z = numpy.asarray(z, dtype=complex)
+    m = numpy.asarray(m, dtype=complex)
+
+    deg_z = int(a_coeffs.shape[0] - 1)
+    s = int(a_coeffs.shape[1] - 1)
+
+    if (z.ndim == 0) and (m.ndim == 0):
+        zz = complex(z)
+        mm = complex(m)
+
+        a = numpy.empty(s + 1, dtype=complex)
+        ap = numpy.empty(s + 1, dtype=complex)
+
+        for j in range(s + 1):
+            c = a_coeffs[:, j]
+
+            val = 0.0 + 0.0j
+            for i in range(deg_z, -1, -1):
+                val = val * zz + c[i]
+            a[j] = val
+
+            dval = 0.0 + 0.0j
+            for i in range(deg_z, 0, -1):
+                dval = dval * zz + (i * c[i])
+            ap[j] = dval
+
+        p = a[s]
+        pm = 0.0 + 0.0j
+        for j in range(s - 1, -1, -1):
+            pm = pm * mm + p
+            p = p * mm + a[j]
+
+        pz = ap[s]
+        for j in range(s - 1, -1, -1):
+            pz = pz * mm + ap[j]
+
+        return p, pz, pm
+
+    shp = numpy.broadcast(z, m).shape
+    zz = numpy.broadcast_to(z, shp).ravel()
+    mm = numpy.broadcast_to(m, shp).ravel()
+
+    zp = powers(zz, deg_z)
+    mp = powers(mm, s)
+
+    dzp = numpy.zeros_like(zp)
+    for i in range(1, deg_z + 1):
+        dzp[:, i] = i * zp[:, i - 1]
+
+    P = numpy.zeros(zz.size, dtype=complex)
+    Pz = numpy.zeros(zz.size, dtype=complex)
+    Pm = numpy.zeros(zz.size, dtype=complex)
+
+    for j in range(s + 1):
+        aj = zp @ a_coeffs[:, j]
+        P += aj * mp[:, j]
+
+        ajp = dzp @ a_coeffs[:, j]
+        Pz += ajp * mp[:, j]
+
+        if j >= 1:
+            Pm += (j * aj) * mp[:, j - 1]
+
+    return P.reshape(shp), Pz.reshape(shp), Pm.reshape(shp)

freealg/_algebraic_form/_homotopy2.py

@@ -0,0 +1,289 @@
+# =======
+# Imports
+# =======
+
+import numpy
+from ._moments import AlgebraicStieltjesMoments
+from tqdm import tqdm
+
+__all__ = ['StieltjesPoly']
+
+
+# ===========
+# select root
+# ===========
+
+# def select_root(roots, z, target):
+#     """
+#     Select the root among Herglotz candidates at a given z closest to a
+#     given target.
+
+#     Parameters
+#     ----------
+#     roots : array_like of complex
+#         Candidate roots for m at the given z.
+#     z : complex
+#         Evaluation point. The Stieltjes/Herglotz branch satisfies
+#         sign(Im(m)) = sign(Im(z)) away from the real axis.
+#     target : complex
+#         Previous continuation value used to enforce continuity, or
+#         target value.
+
+#     Returns
+#     -------
+#     w : complex
+#         Selected root corresponding to the Stieltjes branch.
+#     """
+
+def select_root(roots, z, target, tol_im=1e-10, tiny_im=1e-6, ratio=50.0):
+    z = complex(z)
+    roots = numpy.asarray(roots, dtype=numpy.complex128).ravel()
+
+    s = numpy.sign(z.imag)
+    if s == 0.0:
+        s = 1.0
+
+    # candidates in the correct half-plane (soft)
+    im = numpy.imag(roots) * s
+    cand = roots[im > -tol_im]
+    if cand.size == 0:
+        cand = roots
+        im = numpy.imag(cand) * s
+    else:
+        im = numpy.imag(cand) * s
+
+    # if there is a "big-imag" candidate and a "tiny-imag" candidate, pick big-imag
+    im_max = float(im.max())
+    im_min = float(im.min())
+
+    if im_max > tiny_im and im_min < tiny_im and im_max / max(im_min, 1e-16) > ratio:
+        return cand[int(numpy.argmax(im))]
+
+    # otherwise fall back to continuity
+    t = complex(target)
+    return cand[int(numpy.argmin(numpy.abs(cand - t)))]
+
+# def select_root(roots, z, target, tol_im=1e-10, tiny_im=1e-6, ratio=50.0):
+#     z = complex(z)
+#     roots = numpy.asarray(roots, dtype=numpy.complex128).ravel()
+
+#     s = numpy.sign(z.imag)
+#     if s == 0.0:
+#         s = 1.0
+
+#     im = numpy.imag(roots) * s
+#     cand = roots[im > -tol_im]
+#     if cand.size == 0:
+#         cand = roots
+#         im = numpy.imag(cand) * s
+#     else:
+#         im = numpy.imag(cand) * s
+
+#     im_max = float(im.max())
+#     im_min = float(im.min())
+
+#     # If there's a clear "big-Im" vs "tiny-Im" split, take big-Im
+#     if im_max > tiny_im and im_min < tiny_im and im_max / max(im_min, 1e-16) > ratio:
+#         return cand[int(numpy.argmax(im))]
+
+#     # Otherwise continuity
+#     t = complex(target)
+#     return cand[int(numpy.argmin(numpy.abs(cand - t)))]
+
+
+
+
+
+# ==============
+# stieltjes poly
+# ==============
+
+class StieltjesPoly(object):
+    """
+    Stieltjes-branch evaluator for an algebraic equation P(z, m) = 0.
+
+    Parameters
+    ----------
+    a : ndarray, shape (L, K)
+        Coefficient matrix defining P(z, m) in the monomial basis.
+    eps : float or None, optional
+        If Im(z) == 0, use z + i*eps as the boundary evaluation point.
+        If None and Im(z) == 0, eps is set to 1e-8 * max(1, |z|).
+    height : float, default = 2.0
+        Imaginary height factor used to build a safe start radius.
+    steps : int, default = 100
+        Number of continuation steps along the vertical leg.
+    order : int, default = 15
+        Number of moments in Stieltjes estimate.
+    reanchor : int, default = 1
+        During 1D line-sweeps, every ``reanchor`` points we re-run a full
+        evaluate() at the current point. Setting this to 1 is the most robust.
+    """
+
+    def __init__(self, a, eps=None, height=2.0, steps=100, order=15,
+                 reanchor=1):
+        a = numpy.asarray(a)
+        if a.ndim != 2:
+            raise ValueError("a must be a 2D array.")
+
+        self.a = a
+        self.a_l, _ = a.shape
+        self.eps = eps
+        self.height = float(height)
+        self.steps = int(steps)
+        self.order = int(order)
+        self.reanchor = reanchor
+
+        self.mom = AlgebraicStieltjesMoments(a)
+
+        # Start point far enough away (imag direction)
+        self.rad = 1.0 + self.height * self.mom.radius(self.order)
+        self.z0_p = 1j * self.rad
+        self.z0_m = -1j * self.rad
+
+        # Moment anchors at z0
+        self.m0_p = self.mom.stieltjes(self.z0_p, self.order)
+        self.m0_m = self.mom.stieltjes(self.z0_m, self.order)
+
+    def _poly_coeffs_m(self, z_val):
+        z_powers = z_val ** numpy.arange(self.a_l)
+        return (z_powers @ self.a)[::-1]
+
+    def _poly_roots(self, z_val):
+        coeffs = numpy.asarray(self._poly_coeffs_m(z_val),
+                               dtype=numpy.complex128)
+        return numpy.roots(coeffs)
+
+    def evaluate(self, z):
+        """
+        Evaluate the Stieltjes-branch solution m(z) at a single point.
+
+        Robust strategy for multi-bulk cubics:
+        1) Move to z_mid = x + i*sign(y)*rad (high above the real axis).
+        2) At z_mid, pick the root closest to the moment-based estimate m_mom(z_mid).
+           (This is the crucial re-anchoring that prevents choosing the wrong sheet.)
+        3) Continue vertically from z_mid down to z_eval = x + i*y (or x + i*eps).
+        """
+
+        z = complex(z)
+
+        if self.steps < 1:
+            raise ValueError("steps must be >= 1.")
+
+        # Boundary-value interpretation on the real axis
+        if z.imag == 0.0:
+            if self.eps is None:
+                eps_loc = 1e-8 * max(1.0, abs(z))
+            else:
+                eps_loc = float(self.eps)
+            z_eval = z + 1j * eps_loc
+        else:
+            z_eval = z
+
+        half_sign = numpy.sign(z_eval.imag)
+        if half_sign == 0.0:
+            half_sign = 1.0
+
+        # High-imag anchor at same real part
+        z_mid = complex(z_eval.real, half_sign * self.rad)
+
+        r_mid = self._poly_roots(z_mid)
+        w_prev = r_mid[int(numpy.argmax(numpy.imag(r_mid) * numpy.sign(z_mid.imag)))]
+
+
+        # Moment-based target at z_mid (THIS fixes the wrong-bulk issue)
+        # m_mid_target = self.mom.stieltjes(z_mid, self.order)
+        m_mid_target = -1.0 / z_mid
+
+        # Select correct branch at z_mid using the moment target
+        w_prev = select_root(self._poly_roots(z_mid), z_mid, m_mid_target)
+
+        # Vertical continuation: z_mid -> z_eval
+        for tau in numpy.linspace(0.0, 1.0, int(self.steps) + 1)[1:]:
+            z_tau = z_mid + tau * (z_eval - z_mid)
+            w_prev = select_root(self._poly_roots(z_tau), z_tau, w_prev)
+
+        return w_prev
+
+    def _is_flat_imag_line_1d(self, z_flat):
+        if z_flat.ndim != 1 or z_flat.size < 2:
+            return False
+
+        y = numpy.imag(z_flat)
+        if not numpy.all(numpy.isfinite(y)):
+            return False
+
+        y0 = float(y[0])
+        if y0 == 0.0:
+            return False
+
+        # all in same half-plane
+        if not numpy.all(numpy.sign(y) == numpy.sign(y0)):
+            return False
+
+        # nearly constant imaginary part (relative)
+        tol = 1e-14 * max(1.0, abs(y0))
+        return numpy.max(numpy.abs(y - y0)) <= tol
+
+    def _sweep_line(self, z_sorted, progress=False):
+        n = z_sorted.size
+        out = numpy.empty_like(z_sorted, dtype=numpy.complex128)
+
+        # Always anchor the first point robustly
+        m_prev = self.evaluate(z_sorted[0])
+        out[0] = m_prev
+
+        it = range(1, n)
+        if progress:
+            it = tqdm(it, total=n - 1)
+
+        reanchor = self.reanchor
+        do_reanchor = (reanchor is not None) and (int(reanchor) > 0)
+        reanchor = int(reanchor) if do_reanchor else 0
+
+        for k in it:
+            zk = z_sorted[k]
+
+            # With default reanchor=1, this becomes "evaluate every point"
+            if do_reanchor and (k % reanchor == 0):
+                m_prev = self.evaluate(zk)
+            else:
+                rk = self._poly_roots(zk)
+                m_prev = select_root(rk, zk, m_prev)
+
+            out[k] = m_prev
+
+        return out
+
+    def __call__(self, z, progress=False):
+        # Scalar fast-path
+        if numpy.isscalar(z):
+            return self.evaluate(z)
+
+        z_arr = numpy.asarray(z, dtype=numpy.complex128)
+        out = numpy.empty(z_arr.shape, dtype=numpy.complex128)
+
+        # 1D horizontal line sweep (density-style queries)
+        if z_arr.ndim == 1:
+            z_flat = z_arr
+            if self._is_flat_imag_line_1d(z_flat):
+                order = numpy.argsort(numpy.real(z_flat))
+                inv = numpy.empty_like(order)
+                inv[order] = numpy.arange(order.size)
+
+                z_sorted = z_flat[order]
+                out_sorted = self._sweep_line(z_sorted, progress=progress)
+
+                out[:] = out_sorted[inv]
+                return out
+
+        # Fallback: elementwise
+        if progress:
+            indices = tqdm(numpy.ndindex(z_arr.shape), total=z_arr.size)
+        else:
+            indices = numpy.ndindex(z_arr.shape)
+
+        for idx in indices:
+            out[idx] = self.evaluate(z_arr[idx])
+
+        return out

freealg/_algebraic_form/_homotopy3.py

@@ -0,0 +1,215 @@
+# =======
+# Imports
+# =======
+
+import numpy
+from tqdm import tqdm
+from ._moments import AlgebraicStieltjesMoments
+
+__all__ = ['StieltjesPoly']
+
+
+# ==============
+# Stieltjes Poly
+# ==============
+
+class StieltjesPoly(object):
+    """
+    Stieltjes-branch evaluator for P(z, m)=0 with robust 1D tracking.
+
+    For 1D arrays on a horizontal line (z = x + i*delta), uses Viterbi
+    tracking across the whole line to avoid branch mis-selection.
+    Otherwise falls back to pointwise evaluate().
+
+    Parameters
+    ----------
+    a : ndarray
+        Polynomial coefficient matrix.
+    eps : float or None
+        Imaginary offset when Im(z)=0.
+    height : float
+        Radius factor for safe imaginary height.
+    steps : int
+        Vertical continuation steps in evaluate().
+    order : int
+        Moment order.
+    lam_im : float
+        Viterbi penalty strength for tiny |Im(m)|.
+    tol_im : float
+        Herglotz tolerance.
+    """
+
+    def __init__(self, a, eps=None, height=2.0, steps=100, order=15,
+                 lam_im=1.0e3, tol_im=1.0e-12):
+        a = numpy.asarray(a)
+        if a.ndim != 2:
+            raise ValueError("a must be a 2D array.")
+
+        self.a = a
+        self.a_l, _ = a.shape
+
+        self.eps = eps
+        self.height = float(height)
+        self.steps = int(steps)
+        self.order = int(order)
+
+        self.lam_im = float(lam_im)
+        self.tol_im = float(tol_im)
+
+        self.mom = AlgebraicStieltjesMoments(a)
+
+        self.rad = 1.0 + self.height * self.mom.radius(self.order)
+        self.z0_p = 1j * self.rad
+        self.z0_m = -1j * self.rad
+
+        self.m0_p = self.mom.stieltjes(self.z0_p, self.order)
+        self.m0_m = self.mom.stieltjes(self.z0_m, self.order)
+
+    # -----------
+    # poly roots
+    # -----------
+
+    def _poly_coeffs_m(self, z_val):
+        z_powers = z_val ** numpy.arange(self.a_l)
+        return (z_powers @ self.a)[::-1]
+
+    def _poly_roots(self, z_val):
+        coeffs = numpy.asarray(self._poly_coeffs_m(z_val),
+                               dtype=numpy.complex128)
+        return numpy.roots(coeffs)
+
+    # -------------
+    # point evaluate
+    # -------------
+
+    def _select_root_continuity(self, roots, z, target):
+        z = complex(z)
+        roots = numpy.asarray(roots, dtype=numpy.complex128).ravel()
+
+        s = numpy.sign(z.imag)
+        if s == 0.0:
+            s = 1.0
+
+        im = numpy.imag(roots) * s
+        cand = roots[im > -self.tol_im]
+        if cand.size == 0:
+            cand = roots
+
+        t = complex(target)
+        return cand[int(numpy.argmin(numpy.abs(cand - t)))]
+
+    def evaluate(self, z):
+        z = complex(z)
+
+        if z.imag == 0.0:
+            if self.eps is None:
+                eps_loc = 1e-8 * max(1.0, abs(z))
+            else:
+                eps_loc = float(self.eps)
+            z_eval = z + 1j * eps_loc
+        else:
+            z_eval = z
+
+        s = numpy.sign(z_eval.imag)
+        if s == 0.0:
+            s = 1.0
+
+        z_mid = complex(z_eval.real, s * self.rad)
+
+        # anchor using moment estimate at z_mid
+        target = self.mom.stieltjes(z_mid, self.order)
+        w_prev = self._select_root_continuity(self._poly_roots(z_mid), z_mid, target)
+
+        # vertical continuation
+        for tau in numpy.linspace(0.0, 1.0, int(self.steps) + 1)[1:]:
+            z_tau = z_mid + tau * (z_eval - z_mid)
+            w_prev = self._select_root_continuity(self._poly_roots(z_tau), z_tau, w_prev)
+
+        return w_prev
+
+    # -----------------
+    # viterbi utilities
+    # -----------------
+
+    def _is_flat_imag_line_1d(self, z):
+        if z.ndim != 1 or z.size < 2:
+            return False
+        y = numpy.imag(z)
+        if not numpy.all(numpy.isfinite(y)):
+            return False
+        y0 = float(y[0])
+        if y0 == 0.0:
+            return False
+        if not numpy.all(numpy.sign(y) == numpy.sign(y0)):
+            return False
+        tol = 1e-14 * max(1.0, abs(y0))
+        return numpy.max(numpy.abs(y - y0)) <= tol
+
+    def _herglotz_mask(self, roots, z):
+        s = numpy.sign(z.imag)
+        if s == 0.0:
+            s = 1.0
+        return (numpy.imag(roots) * s) > -self.tol_im
+
+    def _viterbi_track(self, z_sorted, roots_all):
+        z_sorted = numpy.asarray(z_sorted, dtype=numpy.complex128)
+        R = numpy.asarray(roots_all, dtype=numpy.complex128)
+        N, S = R.shape
+
+        unary = numpy.zeros((N, S), dtype=numpy.float64)
+        for k in range(N):
+            mask = self._herglotz_mask(R[k], z_sorted[k])
+            unary[k, ~mask] += 1e30
+            im = numpy.abs(numpy.imag(R[k]))
+            unary[k] += self.lam_im / numpy.maximum(im, 1e-16)
+
+        dp = numpy.full((N, S), numpy.inf, dtype=numpy.float64)
+        prev = numpy.full((N, S), -1, dtype=numpy.int64)
+
+        dp[0] = unary[0]
+
+        for k in range(1, N):
+            diff = R[k][None, :] - R[k-1][:, None]
+            cost = dp[k-1][:, None] + (numpy.abs(diff) ** 2)
+            i_star = numpy.argmin(cost, axis=0)
+            dp[k] = unary[k] + cost[i_star, numpy.arange(S)]
+            prev[k] = i_star
+
+        j = int(numpy.argmin(dp[-1]))
+        path = numpy.empty(N, dtype=numpy.int64)
+        path[-1] = j
+        for k in range(N - 1, 0, -1):
+            path[k - 1] = prev[k, path[k]]
+
+        return R[numpy.arange(N), path]
+
+    # -----
+    # call
+    # -----
+
+    def __call__(self, z, progress=False):
+        if numpy.isscalar(z):
+            return self.evaluate(z)
+
+        z_arr = numpy.asarray(z, dtype=numpy.complex128)
+
+        # Viterbi path only for 1D horizontal lines
+        if z_arr.ndim == 1 and self._is_flat_imag_line_1d(z_arr):
+            order = numpy.argsort(numpy.real(z_arr))
+            inv = numpy.empty_like(order)
+            inv[order] = numpy.arange(order.size)
+
+            z_sorted = z_arr[order]
+            roots_all = numpy.array([self._poly_roots(zk) for zk in z_sorted],
+                                    dtype=numpy.complex128)
+            m_sorted = self._viterbi_track(z_sorted, roots_all)
+            return m_sorted[inv]
+
+        # Fallback: elementwise
+        out = numpy.empty(z_arr.shape, dtype=numpy.complex128)
+        it = numpy.ndindex(z_arr.shape)
+        if progress:
+            it = tqdm(it, total=z_arr.size)
+        for idx in it:
+            out[idx] = self.evaluate(z_arr[idx])
+        return out