freealg 0.7.11__py3-none-any.whl → 0.7.14__py3-none-any.whl

freealg/_algebraic_form/_decompress7.py

@@ -0,0 +1,355 @@
+# SPDX-FileCopyrightText: Copyright 2026, Siavash Ameli
+# SPDX-License-Identifier: BSD-3-Clause
+# SPDX-FileType: SOURCE
+
+"""Free Decompression (FD) Newton solver.
+
+This file defines `decompress_newton` used by AlgebraicForm.decompress(...,
+method='one').
+
+Implementation notes
+--------------------
+We solve, for each query point z and time t (tau = exp(t)), the 2x2 system
+in variables (zeta, y):
+
+    F1(zeta,y) = P(zeta, y) = 0
+    F2(zeta,y) = z - zeta + (tau-1)/y = 0
+
+Then w = m(t,z) = y/tau.
+
+Speed-critical change vs earlier variants: all polynomial evaluations
+(P, dP/dzeta, dP/dy) use scalar Horner without allocating power matrices.
+"""
+
+from __future__ import annotations
+
+import numpy
+
+__all__ = ['decompress_newton']
+
+
+# =====================
+# scalar poly evaluation
+# =====================
+
+def _eval_a_and_da(z: complex, a_coeffs: numpy.ndarray) -> tuple[numpy.ndarray, numpy.ndarray]:
+    """Evaluate a_j(z) and a'_j(z) for j=0..s where P(z,y)=sum_j a_j(z) y^j.
+
+    a_coeffs has shape (deg_z+1, s+1) storing coefficients in z ascending:
+        a_coeffs[i,j] = coeff of z^i in a_j(z).
+    """
+    deg_z = a_coeffs.shape[0] - 1
+    s = a_coeffs.shape[1] - 1
+
+    a = numpy.empty(s + 1, dtype=numpy.complex128)
+    da = numpy.empty(s + 1, dtype=numpy.complex128)
+
+    # Horner for each column j
+    for j in range(s + 1):
+        col = a_coeffs[:, j]
+        # a_j(z)
+        v = complex(col[deg_z])
+        for i in range(deg_z - 1, -1, -1):
+            v = v * z + col[i]
+        a[j] = v
+
+        # a'_j(z)
+        if deg_z == 0:
+            da[j] = 0.0 + 0.0j
+        else:
+            vp = complex(deg_z * col[deg_z])
+            for i in range(deg_z - 1, 0, -1):
+                vp = vp * z + (i * col[i])
+            da[j] = vp
+
+    return a, da
+
+
+def _eval_P_Pz_Py(z: complex, y: complex, a_coeffs: numpy.ndarray) -> tuple[complex, complex, complex]:
+    """Return P(z,y), Pz(z,y)=\partial_z P, Py(z,y)=\partial_y P (scalars)."""
+    a, da = _eval_a_and_da(z, a_coeffs)
+    s = a_coeffs.shape[1] - 1
+
+    # Build powers of y incrementally (cheap; s is small)
+    ypow = 1.0 + 0.0j
+    P = 0.0 + 0.0j
+    Pz = 0.0 + 0.0j
+
+    for j in range(s + 1):
+        P += a[j] * ypow
+        Pz += da[j] * ypow
+        ypow *= y
+
+    # Py
+    Py = 0.0 + 0.0j
+    if s >= 1:
+        ypow = 1.0 + 0.0j  # y^(j-1)
+        for j in range(1, s + 1):
+            Py += (j * a[j]) * ypow
+            ypow *= y
+
+    return P, Pz, Py
+
+
+# =========
+# 2x2 newton
+# =========
+
+def _newton_2x2(
+    z: complex,
+    tau: float,
+    zeta0: complex,
+    y0: complex,
+    a_coeffs: numpy.ndarray,
+    *,
+    max_iter: int,
+    tol: float,
+    damping: float,
+    step_clip: float,
+    w_min: float,
+    require_imw_pos: bool,
+    im_eps: float,
+) -> tuple[complex, complex, bool, int]:
+    """Solve for (zeta,y) at fixed (z,tau)."""
+    zeta = complex(zeta0)
+    y = complex(y0)
+    tau_m1 = tau - 1.0
+
+    # Avoid singular y
+    if abs(y) < w_min:
+        y = (w_min + 0.0j)
+
+    for it in range(max_iter):
+        P, Pz, Py = _eval_P_Pz_Py(zeta, y, a_coeffs)
+        F1 = P
+        F2 = z - zeta + tau_m1 / y
+
+        # Stop criterion on both equations
+        if (abs(F1) <= tol) and (abs(F2) <= tol):
+            w = y / tau
+            if (not require_imw_pos) or (z.imag <= 0.0) or (w.imag > im_eps):
+                return zeta, y, True, it
+
+        # Jacobian
+        # F1_zeta = Pz
+        # F1_y    = Py
+        # F2_zeta = -1
+        # F2_y    = -(tau-1)/y^2
+        inv_y2 = 1.0 / (y * y)
+        J11 = Pz
+        J12 = Py
+        J21 = -1.0 + 0.0j
+        J22 = -(tau_m1) * inv_y2
+
+        det = J11 * J22 - J12 * J21
+        if det == 0.0:
+            break
+
+        # d = -J^{-1}F
+        # d_zeta = (-F1*J22 + F2*J12)/det
+        # d_y    = ( J21*F1 - J11*F2)/det
+        d_zeta = (-F1 * J22 + F2 * J12) / det
+        d_y = (J21 * F1 - J11 * F2) / det
+
+        # Clip step
+        if step_clip is not None and step_clip > 0.0:
+            m = max(abs(d_zeta), abs(d_y))
+            if m > step_clip:
+                s = step_clip / m
+                d_zeta *= s
+                d_y *= s
+
+        lam = float(damping) if damping is not None else 1.0
+        if lam <= 0.0:
+            lam = 1.0
+
+        # Simple backtracking to enforce Im(w)>0 for z in C+
+        # (and avoid y crossing 0)
+        for _ in range(12):
+            zeta_try = zeta + lam * d_zeta
+            y_try = y + lam * d_y
+
+            if abs(y_try) < w_min:
+                lam *= 0.5
+                continue
+
+            if require_imw_pos and (z.imag > 0.0):
+                w_try = y_try / tau
+                if w_try.imag <= im_eps:
+                    lam *= 0.5
+                    continue
+
+            # Accept
+            zeta = zeta_try
+            y = y_try
+            break
+
+        else:
+            # could not find acceptable step
+            break
+
+    return zeta, y, False, max_iter
+
+
+# ===============
+# public interface
+# ===============
+
+def decompress_newton(
+    z_query,
+    t_all,
+    a_coeffs,
+    *,
+    w0_list=None,
+    max_iter: int = 40,
+    tol: float = 1e-13,
+    damping: float = 1.0,
+    step_clip: float = 5.0,
+    w_min: float = 1e-300,
+    max_split: int = 4,
+    require_imw_pos: bool = True,
+    im_eps: float = 1e-14,
+    sweep: bool = False,
+    verbose: bool = False,
+    debug: bool = False,
+    debug_idx=None,
+    eta_rescale: str | None = None,
+    **kwargs,
+):
+    """Compute w(t,z)=m(t,z) on a time grid using FD and Newton.
+
+    Parameters mirror earlier versions; extra kwargs are ignored.
+
+    eta_rescale:
+      - None (default): keep z_query fixed for all times.
+      - 'inv_tau': use z.imag/tau at time t. (Useful for mass checks.)
+
+    Returns
+    -------
+    W : (n_t, n_z) complex
+    ok : (n_t, n_z) bool
+    """
+
+    z_query = numpy.asarray(z_query, dtype=numpy.complex128).ravel()
+    t_all = numpy.asarray(t_all, dtype=float).ravel()
+    if t_all.size == 0:
+        raise ValueError('t_all is empty.')
+    if numpy.any(numpy.diff(t_all) < 0):
+        raise ValueError('t_all must be sorted increasing.')
+
+    n_z = z_query.size
+    n_t = t_all.size
+
+    if w0_list is None:
+        w0 = -1.0 / z_query
+    else:
+        w0 = numpy.asarray(w0_list, dtype=numpy.complex128).ravel()
+        if w0.size != n_z:
+            raise ValueError('w0_list must have same length as z_query.')
+
+    W = numpy.empty((n_t, n_z), dtype=numpy.complex128)
+    ok = numpy.zeros((n_t, n_z), dtype=bool)
+    W[0, :] = w0
+    ok[0, :] = numpy.isfinite(w0.real) & numpy.isfinite(w0.imag)
+
+    dbg_set = set()
+    if debug_idx is not None:
+        try:
+            dbg_set = set(int(i) for i in debug_idx)
+        except Exception:
+            dbg_set = set()
+
+    for it in range(1, n_t):
+        t_prev = float(t_all[it - 1])
+        t = float(t_all[it])
+        tau_prev = float(numpy.exp(t_prev))
+        tau = float(numpy.exp(t))
+
+        if eta_rescale == 'inv_tau':
+            z_t = z_query.real + 1j * (z_query.imag / tau)
+        else:
+            z_t = z_query
+
+        # warm start: previous w(t_prev,z) -> y seed
+        w_seed = W[it - 1, :].copy()
+        y_seed = tau_prev * w_seed
+
+        # build zeta seed that satisfies the constraint initially
+        y_safe = y_seed.copy()
+        tiny = numpy.abs(y_safe) < w_min
+        if numpy.any(tiny):
+            y_safe[tiny] = (w_min + 0.0j)
+        zeta_seed = z_t + (tau - 1.0) / y_safe
+
+        w_out = numpy.empty(n_z, dtype=numpy.complex128)
+        ok_out = numpy.zeros(n_z, dtype=bool)
+
+        # optional sweep within x at same time
+        y_last = None
+        for j in range(n_z):
+            z = z_t[j]
+
+            if sweep and (y_last is not None):
+                y0 = y_last
+                y_safe0 = y0 if abs(y0) > w_min else (w_min + 0.0j)
+                zeta0 = z + (tau - 1.0) / y_safe0
+            else:
+                y0 = y_seed[j]
+                zeta0 = zeta_seed[j]
+
+            zeta, y, okj, nit = _newton_2x2(
+                z, tau, zeta0, y0, a_coeffs,
+                max_iter=max_iter, tol=tol,
+                damping=damping, step_clip=step_clip,
+                w_min=w_min, require_imw_pos=require_imw_pos,
+                im_eps=im_eps,
+            )
+
+            if (not okj) and (max_split is not None) and (max_split > 0):
+                # try smaller pseudo-steps by tempering tau -> tau_mid
+                # (helps if time steps are too large)
+                for k in range(1, int(max_split) + 1):
+                    tau_mid = tau_prev + (tau - tau_prev) * (k / float(max_split))
+                    zeta_mid, y_mid, ok_mid, _ = _newton_2x2(
+                        z, tau_mid, zeta0, y0, a_coeffs,
+                        max_iter=max_iter, tol=tol,
+                        damping=damping, step_clip=step_clip,
+                        w_min=w_min, require_imw_pos=require_imw_pos,
+                        im_eps=im_eps,
+                    )
+                    if ok_mid:
+                        # now jump from (tau_mid) seed
+                        zeta0b = zeta_mid
+                        y0b = y_mid
+                        zeta, y, okj, nit = _newton_2x2(
+                            z, tau, zeta0b, y0b, a_coeffs,
+                            max_iter=max_iter, tol=tol,
+                            damping=damping, step_clip=step_clip,
+                            w_min=w_min, require_imw_pos=require_imw_pos,
+                            im_eps=im_eps,
+                        )
+                        if okj:
+                            break
+
+            wj = y / tau
+            w_out[j] = wj
+            ok_out[j] = bool(okj)
+
+            if okj:
+                y_last = y
+
+            if debug and (j in dbg_set):
+                print(
+                    f"[t={t:0.6f}] j={j} z={z.real:+.6f}{z.imag:+.2e}j "
+                    f"zeta={zeta.real:+.6f}{zeta.imag:+.2e}j "
+                    f"y={y.real:+.3e}{y.imag:+.3e}j "
+                    f"w={wj.real:+.3e}{wj.imag:+.3e}j ok={okj} it={nit}"
+                )
+
+        W[it, :] = w_out
+        ok[it, :] = ok_out
+
+        if verbose:
+            print(f'[t={t:0.6f}] ok={ok_out.mean():0.3f}')
+
+    return W, ok

freealg/_algebraic_form/_decompress8.py

@@ -0,0 +1,369 @@
+# 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.
+
+"""
+Robust Newton solver for Free Decompression (FD) using the characteristic
+variables (zeta, y) on the spectral curve P(zeta, y)=0.
+
+This implementation avoids solving directly in w (where zeta = z + alpha / w
+introduces a pole at w=0). Instead, for each query z and time t (tau=e^t), we
+solve the 2x2 complex system:
+
+    F1(zeta, y) := P(zeta, y) = 0
+    F2(zeta, y) := zeta - (tau - 1)/y - z = 0
+
+Then m(t,z) = w = y/tau.
+
+The public API matches the existing decompress_newton used by AlgebraicForm.
+"""
+
+import numpy
+
+__all__ = ['decompress_newton']
+
+
+# =====================
+# Polynomial evaluation
+# =====================
+
+def _powers(x, deg):
+    """
+    Returns [1, x, x^2, ..., x^deg] for each element of x.
+    """
+    x = numpy.asarray(x, dtype=numpy.complex128)
+    xp = numpy.ones((x.size, deg + 1), dtype=numpy.complex128)
+    for k in range(1, deg + 1):
+        xp[:, k] = xp[:, k - 1] * x
+    return xp
+
+
+def _poly_coef_in_y(zeta, a_coeffs):
+    """
+    For each zeta, compute coefficients a_j(zeta) so that
+        P(zeta, y) = sum_{j=0}^s a_j(zeta) y^j
+    """
+    zeta = numpy.asarray(zeta, dtype=numpy.complex128).ravel()
+    deg_z = int(a_coeffs.shape[0] - 1)
+    s = int(a_coeffs.shape[1] - 1)
+
+    zp = _powers(zeta, deg_z)
+    a = numpy.empty((zeta.size, s + 1), dtype=numpy.complex128)
+    for j in range(s + 1):
+        a[:, j] = zp @ a_coeffs[:, j]
+    return a
+
+
+def _poly_coef_in_y_dzeta(zeta, a_coeffs):
+    """
+    For each zeta, compute coefficients da_j/dzeta(zeta) so that
+        d/dzeta P(zeta, y) = sum_{j=0}^s (da_j/dzeta)(zeta) y^j
+    """
+    zeta = numpy.asarray(zeta, dtype=numpy.complex128).ravel()
+    deg_z = int(a_coeffs.shape[0] - 1)
+    s = int(a_coeffs.shape[1] - 1)
+
+    if deg_z <= 0:
+        return numpy.zeros((zeta.size, s + 1), dtype=numpy.complex128)
+
+    # derivative powers: d/dzeta zeta^i = i zeta^(i-1)
+    zp = _powers(zeta, deg_z - 1)  # up to zeta^(deg_z-1)
+    da = numpy.empty((zeta.size, s + 1), dtype=numpy.complex128)
+    for j in range(s + 1):
+        col = a_coeffs[:, j]
+        # sum_{i=1..deg_z} i*c_{i,j} zeta^(i-1)
+        # build weighted coefficients for zp @ ...
+        w = numpy.arange(deg_z + 1, dtype=numpy.complex128) * col
+        da[:, j] = zp @ w[1:]
+    return da
+
+
+def _eval_P_and_partials(zeta, y, a_coeffs):
+    """
+    Evaluate P(zeta,y), P_zeta(zeta,y), P_y(zeta,y).
+
+    Returns:
+        P, Pz, Py  (arrays of shape (n,))
+    """
+    zeta = numpy.asarray(zeta, dtype=numpy.complex128).ravel()
+    y = numpy.asarray(y, dtype=numpy.complex128).ravel()
+
+    a = _poly_coef_in_y(zeta, a_coeffs)          # (n, s+1)
+    da = _poly_coef_in_y_dzeta(zeta, a_coeffs)   # (n, s+1)
+
+    s = int(a.shape[1] - 1)
+    # powers of y up to s
+    yp = _powers(y, s)  # (n, s+1)
+
+    P = numpy.sum(a * yp, axis=1)
+
+    # P_zeta = sum_j da_j(zeta) y^j
+    Pz = numpy.sum(da * yp, axis=1)
+
+    # P_y = sum_{j>=1} j*a_j(zeta) y^{j-1}
+    if s == 0:
+        Py = numpy.zeros_like(P)
+    else:
+        Py = numpy.zeros_like(P)
+        # yp[:, j-1] available
+        for j in range(1, s + 1):
+            Py += (j * a[:, j]) * yp[:, j - 1]
+
+    return P, Pz, Py
+
+
+# =======================
+# 2x2 complex Newton step
+# =======================
+
+def _newton_2x2(z, tau, zeta0, y0, a_coeffs, max_iter, tol,
+                armijo, min_lam, w_min, enforce_imag=True):
+    """
+    Solve for (zeta,y) at given (z,tau) using damped Newton on the 2x2 complex
+    system.
+
+    Returns:
+        (zeta, y, ok, iters)
+    """
+    zeta = numpy.complex128(zeta0)
+    y = numpy.complex128(y0)
+
+    # helper to compute residual norm
+    def F(zeta_, y_):
+        P, Pz, Py = _eval_P_and_partials(numpy.array([zeta_]),
+                                         numpy.array([y_]),
+                                         a_coeffs)
+        P = P[0]
+        Pz = Pz[0]
+        Py = Py[0]
+        F1 = P
+        F2 = (zeta_ - (tau - 1.0) / y_ - z)
+        # Jacobian entries
+        J11 = Pz
+        J12 = Py
+        J21 = 1.0 + 0.0j
+        J22 = (tau - 1.0) / (y_ * y_)  # d/dy of (-(tau-1)/y) is +(tau-1)/y^2
+        return F1, F2, J11, J12, J21, J22
+
+    # initial residual
+    F1, F2, J11, J12, J21, J22 = F(zeta, y)
+    r0 = max(abs(F1), abs(F2))
+
+    if not numpy.isfinite(r0):
+        return zeta, y, False, 0
+
+    for it in range(int(max_iter)):
+        r = max(abs(F1), abs(F2))
+        if r <= tol:
+            w = y / tau
+            if (abs(w) < w_min) or (not numpy.isfinite(w.real)) or (not numpy.isfinite(w.imag)):
+                return zeta, y, False, it
+            if enforce_imag and (z.imag > 0.0) and (w.imag <= 0.0):
+                return zeta, y, False, it
+            return zeta, y, True, it
+
+        # Solve 2x2 complex linear system J * d = -F
+        det = J11 * J22 - J12 * J21
+        if det == 0 or (not numpy.isfinite(det.real)) or (not numpy.isfinite(det.imag)):
+            return zeta, y, False, it
+
+        d_zeta = (-F1 * J22 - (-F2) * J12) / det
+        d_y = (J11 * (-F2) - J21 * (-F1)) / det
+
+        # Armijo damping on residual norm
+        lam = 1.0
+        if armijo is None or armijo <= 0.0:
+            lam = 1.0
+        else:
+            c = float(armijo)
+            r_curr = r
+            # Try steps until sufficient decrease or min_lam
+            while True:
+                zeta_try = zeta + lam * d_zeta
+                y_try = y + lam * d_y
+                F1t, F2t, J11t, J12t, J21t, J22t = F(zeta_try, y_try)
+                r_try = max(abs(F1t), abs(F2t))
+                if numpy.isfinite(r_try) and (r_try <= (1.0 - c * lam) * r_curr):
+                    # accept
+                    zeta, y = zeta_try, y_try
+                    F1, F2, J11, J12, J21, J22 = F1t, F2t, J11t, J12t, J21t, J22t
+                    break
+                lam *= 0.5
+                if lam < float(min_lam):
+                    # accept last trial if it improves, else fail
+                    if numpy.isfinite(r_try) and (r_try < r_curr):
+                        zeta, y = zeta_try, y_try
+                        F1, F2, J11, J12, J21, J22 = F1t, F2t, J11t, J12t, J21t, J22t
+                        break
+                    return zeta, y, False, it
+
+        # continue loop
+
+    # max_iter exceeded
+    return zeta, y, False, int(max_iter)
+
+
+# =================
+# Public entrypoint
+# =================
+
+def decompress_newton(z_query, t_all, a_coeffs, w0_list=None,
+                     max_iter=50, tol=1e-12,
+                     armijo=1e-4, min_lam=1e-6, w_min=1e-14,
+                     sweep=True, verbose=False, **kwargs):
+    """
+    Parameters
+    ----------
+    z_query : array_like (complex)
+        Query points z where m(t,z) should be evaluated (typically x + i*delta)
+
+    t_all : array_like (float)
+        Time grid including t=0, increasing.
+
+    a_coeffs : ndarray
+        Coefficient matrix for P(z,m) in monomial basis (deg_z+1, s+1)
+
+    w0_list : array_like (complex), optional
+        Initial condition m(0, z_query) on the physical branch. If None, this
+        function will approximate it as -1/z_query (may be poor near cuts).
+
+    Other parameters mirror the existing solver interface.
+
+    Returns
+    -------
+    W : ndarray (n_t, n_z) complex
+        Estimated m(t,z) on the tracked branch.
+
+    ok : ndarray (n_t, n_z) bool
+        Convergence flag for each point.
+    """
+
+    z_query = numpy.asarray(z_query, dtype=numpy.complex128).ravel()
+    t_all = numpy.asarray(t_all, dtype=float).ravel()
+
+    if t_all.size == 0:
+        raise ValueError('t_all is empty.')
+
+    # enforce sorted
+    if numpy.any(numpy.diff(t_all) < 0):
+        raise ValueError('t_all must be sorted increasing.')
+
+    n_z = z_query.size
+    n_t = t_all.size
+
+
+    # If the caller does not include t=0 as the first element, prepend it.
+    # The Newton march below assumes W[0] is the known initial condition at t=0
+    # (physical branch), and evolves forward for k=1..n_t-1.
+    drop_t0 = False
+    if n_t == 0:
+        raise ValueError('t_all must be non-empty.')
+    if abs(float(t_all[0])) > 0.0:
+        t_all = numpy.concatenate(([0.0], t_all))
+        n_t = t_all.size
+        drop_t0 = True
+    if w0_list is None:
+        w0 = -1.0 / z_query
+    else:
+        w0 = numpy.asarray(w0_list, dtype=numpy.complex128).ravel()
+        if w0.size != n_z:
+            raise ValueError('w0_list must have same length as z_query.')
+
+    # Output arrays
+    W = numpy.empty((n_t, n_z), dtype=numpy.complex128)
+    ok = numpy.zeros((n_t, n_z), dtype=bool)
+
+    # Initialize at t=0
+    W[0, :] = w0
+    ok[0, :] = numpy.isfinite(w0.real) & numpy.isfinite(w0.imag)
+
+    # For each time step, solve independently per z with continuation seeds
+    for it in range(1, n_t):
+        t_prev = float(t_all[it - 1])
+        t = float(t_all[it])
+        tau_prev = numpy.exp(t_prev)
+        tau = numpy.exp(t)
+
+        # seeds from previous time
+        w_seed = W[it - 1, :].copy()
+        # y seed from previous: y = tau_prev * w
+        y_seed = tau_prev * w_seed
+
+        # Optional sweep: use previous x point at the same time as init
+        zeta_seed = numpy.empty(n_z, dtype=numpy.complex128)
+
+        # Initialize zeta so that F2 is satisfied initially (good conditioning)
+        # zeta = z + (tau-1)/y
+        # Guard y=0
+        y_safe = y_seed.copy()
+        tiny = numpy.abs(y_safe) < 1e-300
+        if numpy.any(tiny):
+            y_safe[tiny] = (1e-300 + 0.0j)
+        zeta_seed[:] = z_query + (tau - 1.0) / y_safe
+
+        # Sweep order: left-to-right
+        if sweep:
+            order = range(n_z)
+        else:
+            order = range(n_z)
+
+        # Storage for this time
+        w_out = numpy.empty(n_z, dtype=numpy.complex128)
+        ok_out = numpy.zeros(n_z, dtype=bool)
+
+        prev_good_idx = None
+        for j in order:
+            z = z_query[j]
+            # choose initial guess
+            y0 = y_seed[j]
+            zeta0 = zeta_seed[j]
+
+            if sweep and (prev_good_idx is not None):
+                # use last successful (zeta,y) as initial, but adjust zeta to
+                # satisfy constraint using current z and that y
+                y0 = y_last
+                y_safe0 = y0 if abs(y0) > 1e-300 else (1e-300 + 0.0j)
+                zeta0 = z + (tau - 1.0) / y_safe0
+
+            # Solve
+            zeta, y, okj, _ = _newton_2x2(
+                z, tau, zeta0, y0, a_coeffs,
+                max_iter=max_iter, tol=tol,
+                armijo=armijo, min_lam=min_lam, w_min=w_min,
+                enforce_imag=True
+            )
+
+            if not okj:
+                # Fallback 1: asymptotic Stieltjes seed
+                w_asym = -1.0 / z
+                y0b = tau * w_asym
+                y_safe0b = y0b if abs(y0b) > 1e-300 else (1e-300 + 0.0j)
+                zeta0b = z + (tau - 1.0) / y_safe0b
+
+                zeta, y, okj, _ = _newton_2x2(
+                    z, tau, zeta0b, y0b, a_coeffs,
+                    max_iter=max_iter, tol=tol,
+                    armijo=armijo, min_lam=min_lam, w_min=w_min,
+                    enforce_imag=True
+                )
+
+            wj = y / tau
+            w_out[j] = wj
+            ok_out[j] = bool(okj)
+
+            if okj:
+                prev_good_idx = j
+                y_last = y
+
+        W[it, :] = w_out
+        ok[it, :] = ok_out
+
+        if verbose:
+            print(f'[t={t:0.6f}] success={ok_out.mean():0.3f}')
+    if drop_t0:
+        return W[1:, :], ok[1:, :]
+    return W, ok