freealg 0.7.11__py3-none-any.whl → 0.7.14__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/_edge.py

@@ -13,7 +13,7 @@
 
 import numpy
 from ._continuation_algebraic import eval_roots
-from ._decompress import eval_P_partials
+from ._decompress_util import eval_P_partials
 
 __all__ = ['evolve_edges', 'merge_edges']
 
@@ -129,65 +129,28 @@ def _init_edge_point_from_support(x_edge, a_coeffs, eta=1e-3):
 # evolve edges
 # ============
 
-def evolve_edges(t_grid, a_coeffs, support=None, eta=1e-3, dt_max=0.1,
-                 max_iter=30, tol=1e-12):
+def evolve_edges(
+        t_grid,
+        a_coeffs,
+        support=None,
+        eta=1e-3,
+        dt_max=0.1,
+        max_iter=30,
+        tol=1e-12,
+        return_preimage=False):
     """
-    Evolve spectral edges under free decompression using the fitted
-    polynomial P.
-
-    At time t, edges are computed as critical values of the FD map restricted
-    to the spectral curve P(zeta,y)=0. We solve for (zeta(t), y(t)):
+    Evolve spectral edges under free decompression using the fitted polynomial
+    P.
 
+    Solves for (zeta(t), y(t)) on the spectral curve:
         P(zeta,y) = 0,
         y^2 * Py(zeta,y) - (exp(t)-1) * Pzeta(zeta,y) = 0,
 
-    then map to the physical coordinate:
+    then maps to physical coordinate:
         z_edge(t) = zeta - (exp(t)-1)/y.
 
-    Parameters
-    ----------
-    t_grid : array_like of float
-        Strictly increasing time grid.
-    a_coeffs : ndarray
-        Coefficients defining P(zeta,y).
-    support : list of (float, float), optional
-        List of intervals [(a1,b1),...,(ak,bk)] at t=0. If provided, these
-        endpoints are used as labels/initial guesses and all are tracked.
-        If omitted, this function currently raises ValueError (auto-detection
-        is intentionally not implemented here to avoid fragile heuristics).
-    eta : float, optional
-        Small imaginary part used only to pick an initial physical root near
-        each endpoint at t=0.
-    dt_max : float, optional
-        Maximum internal time step used for substepping in t.
-    max_iter : int, optional
-        Newton iterations per time step.
-    tol : float, optional
-        Tolerance for the 2x2 Newton solve.
-
-    Returns
-    -------
-    edges : ndarray, shape (len(t_grid), 2*k)
-        Tracked edges in the order [a1,b1,a2,b2,...] for each time.
-    ok : ndarray of bool, same shape as edges
-        Flags indicating whether each edge solve succeeded.
-
-    Notes
-    -----
-    The solve is done by continuation in time. If two edges merge, the Newton
-    system may become ill-conditioned near the merge time.
-
-    Examples
-    --------
-    .. code-block:: python
-
-        t_grid = numpy.linspace(0.0, 3.0, 61)
-        support = [(a1,b1)]
-        edges, ok = fd_evolve_edges(t_grid, a_coeffs, support=support,
-                                    eta=1e-3)
-
-        a_t = edges[:, 0]
-        b_t = edges[:, 1]
+    If return_preimage=True, also returns zeta_hist and y_hist of shape
+    (nt, 2k).
     """
 
     t_grid = numpy.asarray(t_grid, dtype=float).ravel()
@@ -197,32 +160,43 @@ def evolve_edges(t_grid, a_coeffs, support=None, eta=1e-3, dt_max=0.1,
         raise ValueError("t_grid must be strictly increasing.")
 
     if support is None:
-        raise ValueError(
-            "support must be provided (auto-detection not implemented).")
+        raise ValueError("support must be provided (auto-detection not " +
+                         "implemented).")
 
-    # Flatten endpoints in the order [a1,b1,a2,b2,...]
+    # Flatten endpoints in fixed order [a1,b1,a2,b2,...]
     endpoints0 = []
     for a, b in support:
         endpoints0.append(float(a))
         endpoints0.append(float(b))
 
     m = len(endpoints0)
-    edges = numpy.empty((t_grid.size, m), dtype=float)
+    complex_edges = numpy.empty((t_grid.size, m), dtype=numpy.complex128)
     ok = numpy.zeros((t_grid.size, m), dtype=bool)
 
-    # Initialize spectral points (zeta,y) at t=0 for each endpoint
-    zeta = numpy.empty(m, dtype=complex)
-    y = numpy.empty(m, dtype=complex)
+    if return_preimage:
+        zeta_hist = numpy.empty((t_grid.size, m), dtype=numpy.complex128)
+        y_hist = numpy.empty((t_grid.size, m), dtype=numpy.complex128)
+    else:
+        zeta_hist = None
+        y_hist = None
+
+    # Initialize (zeta,y) at t=0 from support endpoints
+    zeta = numpy.empty(m, dtype=numpy.complex128)
+    y = numpy.empty(m, dtype=numpy.complex128)
 
     for j in range(m):
         z0, y0, ok0 = _init_edge_point_from_support(endpoints0[j], a_coeffs,
                                                     eta=eta)
         zeta[j] = z0
         y[j] = y0
-        edges[0, j] = float(numpy.real(z0))  # at t=0, z_edge = zeta
         ok[0, j] = ok0
+        complex_edges[0, j] = z0  # at t=0, tau-1 = 0 => z_edge = zeta
+
+    if return_preimage:
+        zeta_hist[0, :] = zeta
+        y_hist[0, :] = y
 
-    # Time continuation
+    # Time stepping
     for it in range(1, t_grid.size):
         t0 = float(t_grid[it - 1])
         t1 = float(t_grid[it])
@@ -236,19 +210,22 @@ def evolve_edges(t_grid, a_coeffs, support=None, eta=1e-3, dt_max=0.1,
             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], a_coeffs, max_iter=max_iter, tol=tol
                 )
                 ok[it, j] = okj
 
         tau = float(numpy.exp(t1))
         c = tau - 1.0
-        z_edge = zeta - c / y
+        complex_edges[it, :] = zeta - c / y
 
-        edges[it, :] = numpy.real(z_edge)
-        # ok[it,:] already set in last substep loop
+        if return_preimage:
+            zeta_hist[it, :] = zeta
+            y_hist[it, :] = y
 
-    return edges, ok
+    if return_preimage:
+        return complex_edges, ok, zeta_hist, y_hist
+
+    return complex_edges, ok
 
 
 # ===========
@@ -282,6 +259,7 @@ def merge_edges(edges, tol=0.0):
     active_k : ndarray, shape (nt,)
         Number of remaining bulks (connected components) at each time.
     """
+
     edges = numpy.asarray(edges, dtype=float)
     nt, m = edges.shape
     if m % 2 != 0:

freealg/_algebraic_form/_homotopy.py

@@ -5,6 +5,7 @@
 import numpy
 from ._moments import AlgebraicStieltjesMoments
 from tqdm import tqdm
+from math import comb
 
 __all__ = ['StieltjesPoly']
 
@@ -86,6 +87,8 @@ class StieltjesPoly(object):
         Coefficient matrix defining P(z, m) in the monomial basis. For fixed
         z, the coefficients of the polynomial in m are assembled from powers
         of z.
+    mom : callable, optional
+        A callable providing raw moments ``m_k = mom(k)``
     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|).
@@ -114,7 +117,7 @@ class StieltjesPoly(object):
     None, eps is chosen per element as 1e-8 * max(1, |z|).
     """
 
-    def __init__(self, a, eps=None, height=2.0, steps=100, order=15):
+    def __init__(self, a, mom=None, eps=None, height=2.0, steps=100, order=15):
         a = numpy.asarray(a)
         if a.ndim != 2:
             raise ValueError("a must be a 2D array.")
@@ -125,13 +128,28 @@ class StieltjesPoly(object):
         self.height = height
         self.steps = steps
         self.order = order
+        if order < 3:
+            raise RuntimeError("order is too small, choose a larger value.")
 
-        self.mom = AlgebraicStieltjesMoments(a)
-        self.rad = 1.0 + self.height * self.mom.radius(self.order)
+        if mom is None:
+            self.mom = AlgebraicStieltjesMoments(a)
+        else:
+            self.mom = mom
+        self.mu = numpy.array([self.mom(j) for j in range(self.order+1)])
+        self.rad = max([numpy.abs(self.mu[j] / self.mu[j-1])
+                        for j in range(2, self.order+1)])
+        self.rad = 1.0 + self.height * self.rad
         self.z0_p = 1j * self.rad
-        self.m0_p = self.mom.stieltjes(self.z0_p, self.order)
+        self.m0_p = self._moment_est(self.z0_p)
         self.z0_m = -1j * self.rad
-        self.m0_m = self.mom.stieltjes(self.z0_m, self.order)
+        self.m0_m = self._moment_est(self.z0_m)
+
+    def _moment_est(self, z):
+        # Estimate Stieltjes transform (root) using moment
+        # expansion
+        z = numpy.asarray(z)
+        pows = z[..., numpy.newaxis]**(-numpy.arange(self.order+1)-1)
+        return -numpy.sum(pows * self.mu, axis=-1)
 
     def _poly_coeffs_m(self, z_val):
         z_powers = z_val ** numpy.arange(self.a_l)
@@ -142,7 +160,8 @@ class StieltjesPoly(object):
                                dtype=numpy.complex128)
         return numpy.roots(coeffs)
 
-    def evaluate(self, z, eps=None, height=2.0, steps=100, order=15):
+    def evaluate(self, z, eps=None, height=2.0, steps=100, order=15, extrap=2,
+                 num_angles=1):
         """
         Evaluate the Stieltjes-branch solution m(z) at a single point.
 
@@ -168,33 +187,46 @@ class StieltjesPoly(object):
         if half_sign == 0.0:
             half_sign = 1.0
 
-        # # If z is outside radius of convergence, no homotopy
-        # # necessary
-        # if numpy.abs(z) > self.rad:
-        #     target = self.mom.stieltjes(z, self.order)
-        #     return select_root(self._poly_roots(z), z, target)
-
-        if half_sign > 0.0:
-            z0 = self.z0_p
-            target = self.m0_p
-        else:
-            z0 = self.z0_m
-            target = self.m0_m
+        # If z is outside radius of convergence, no homotopy
+        # necessary
+        if numpy.abs(z) > self.rad:
+            target = self._moment_est(z)
+            return select_root(self._poly_roots(z), z, target)
+
+        # z0 = z.real
+        # z0 = z0 + 1j*numpy.sqrt(self.rad**2 - z0**2)
+        # target = self._moment_est(z0)
+        # if half_sign > 0.0:
+        #     z0 = self.z0_p
+        #     target = self.m0_p
+        # else:
+        #     z0 = self.z0_m
+        #     target = self.m0_m
 
         # Initialize at z0
-        w_prev = select_root(self._poly_roots(z0), z0, target)
-
-        # Straight-line homotopy continuation
-        for tau in numpy.linspace(0.0, 1.0, int(self.steps) + 1)[1:]:
-            z_tau = z0 + tau * (z_eval - z0)
-            w_prev = select_root(self._poly_roots(z_tau), z_tau, w_prev)
-
-        return w_prev
-
-    def __call__(self, z, progress=False):
+        res = 0
+        for theta in numpy.linspace(0, numpy.pi, num_angles+2)[1:-1]:
+            z0 = self.rad * numpy.exp(1j * theta) * half_sign
+            target = self._moment_est(z0)
+            coeffs = numpy.array([(-1)**k * comb(extrap, k + 1)
+                                for k in range(extrap)])
+            w_prev = numpy.ones(extrap) * \
+                select_root(self._poly_roots(z0), z0, target)
+
+            # Straight-line homotopy continuation
+            for tau in numpy.linspace(0.0, 1.0, int(self.steps) + 1)[1:]:
+                z_tau = z0 + tau * (z_eval - z0)
+                target = numpy.dot(coeffs, w_prev)
+                w_prev[1:] = w_prev[0:-1]
+                w_prev[0] = select_root(self._poly_roots(z_tau), z_tau, target)
+            res += w_prev[0]
+
+        return res / num_angles
+
+    def __call__(self, z, progress=False, num_angles=1):
         # Scalar fast-path
         if numpy.isscalar(z):
-            return self.evaluate(z)
+            return self.evaluate(z, num_angles=num_angles)
 
         # Array-like: evaluate elementwise, preserving shape
         z_arr = numpy.asarray(z)
@@ -206,7 +238,7 @@ class StieltjesPoly(object):
         else:
             indices = numpy.ndindex(z_arr.shape)
         for idx in indices:
-            out[idx] = self.evaluate(z_arr[idx])
+            out[idx] = self.evaluate(z_arr[idx], num_angles=num_angles)
 
         return out