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

freealg/_geometric_form/__init__.py

@@ -0,0 +1,13 @@
+# 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.
+
+from .geometric_form import GeometricForm
+# from ._elliptic_functions import ellipj
+# from ._continuation_genus1 import mobius_z
+
+__all__ = ['GeometricForm']

freealg/_geometric_form/_continuation_genus0.py

@@ -0,0 +1,175 @@
+# 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__ = ['joukowski_z', 'joukowski_w', 'fit_pade', 'eval_pade',
+           'generate_pade']
+
+
+# ===========
+# joukowski z
+# ===========
+
+def joukowski_z(w, a, b):
+
+    c = 0.5 * (a + b)
+    d = 0.5 * (b - a)
+    z = c + 0.5 * d * (w + 1.0 / w)
+
+    return z
+
+
+# ===========
+# joukowski w
+# ===========
+
+def joukowski_w(z, a, b):
+
+    c = 0.5 * (a + b)
+    d = 0.5 * (b - a)
+    xi = (z - c) / d
+    s = numpy.sqrt(xi * xi - 1.0)
+
+    # Stabilize sqrt branch: make s have same sign as xi (helps continuity)
+    s = numpy.where(numpy.real(xi) < 0.0, -s, s)
+
+    w1 = xi + s
+    w2 = xi - s
+    w = numpy.where(numpy.abs(w1) >= 1.0, w1, w2)
+
+    return w
+
+
+# ========
+# fit pade
+# ========
+
+def fit_pade(w, m1, deg_p=12, deg_q=12, ridge_lambda=0.0):
+    """
+    Fit m1 on w-plane using Pade rational approximation.
+    """
+
+    n_samples = m1.size
+
+    wp = numpy.ones((n_samples, deg_p + 1), dtype=complex)
+    for k in range(1, deg_p + 1):
+        wp[:, k] = wp[:, k - 1] * w
+
+    wq = numpy.ones((n_samples, deg_q + 1), dtype=complex)
+    for k in range(1, deg_q + 1):
+        wq[:, k] = wq[:, k - 1] * w
+
+    A = numpy.hstack([wp, -m1[:, None] * wq[:, 1:]])
+    bvec = m1
+
+    # Scale columns for better conditioning on LS
+    s = numpy.linalg.norm(A, axis=0)
+    s[s == 0] = 1.0
+    As = A / s[None, :]
+
+    if ridge_lambda is None:
+        ridge_lambda = 0.0
+
+    if ridge_lambda > 0.0:
+        # Scale ridge by average diagonal magnitude of AhA
+        # Since columns of As have unit norm, this is typically ~1.
+        alpha = ridge_lambda
+
+        # Solving augmented least square
+        n_coef = As.shape[1]
+        A_aug = numpy.vstack([As, numpy.sqrt(alpha) * numpy.eye(
+            n_coef, dtype=complex)])
+        b_aug = numpy.concatenate([bvec, numpy.zeros(n_coef, dtype=complex)])
+
+        coef, _, _, _ = numpy.linalg.lstsq(A_aug, b_aug, rcond=None)
+    else:
+        coef, _, _, _ = numpy.linalg.lstsq(As, bvec, rcond=None)
+
+    coef = coef / s
+
+    p = coef[:deg_p + 1]
+    q = numpy.zeros(deg_q + 1, dtype=complex)
+    q[0] = 1.0
+    q[1:] = coef[deg_p + 1:]
+
+    return p, q
+
+
+# =========
+# eval pade
+# =========
+
+def eval_pade(w, p, q):
+
+    num = numpy.zeros_like(w, dtype=complex)
+    den = numpy.zeros_like(w, dtype=complex)
+
+    for k in range(len(p) - 1, -1, -1):
+        num = num * w + p[k]
+
+    for k in range(len(q) - 1, -1, -1):
+        den = den * w + q[k]
+
+    return num / den
+
+
+# =============
+# generate pade
+# =============
+
+def generate_pade(m1_fn, a, b, deg_p=12, deg_q=12, n_samples=4096, r=1.2,
+                  n_r=1, r_min=None, ridge_lambda=0.0):
+
+    if r_min is None:
+        r_min = 1.0 + 0.05 * (r - 1.0) if r > 1.0 else 1.0
+
+    if n_r is None or n_r < 1:
+        n_r = 1
+
+    if n_samples % 2 != 0:
+        raise ValueError('n_samples should be even.')
+
+    if n_r == 1:
+        rs = numpy.array([r], dtype=float)
+    else:
+        rs = numpy.linspace(r_min, r, n_r)
+
+    W_list = []
+    M_list = []
+
+    n_half = n_samples // 2
+
+    for r_i in rs:
+
+        # Generate sample points along theta
+        theta = numpy.pi * (numpy.arange(n_half) + 0.5) / n_half
+        w = r_i * numpy.exp(1j * theta)
+        z = joukowski_z(w, a, b)
+        m1 = m1_fn(z)
+
+        W_list.append(w)
+        M_list.append(m1)
+
+        # Add conjugate points which enforces Schwarz reflection
+        W_list.append(numpy.conjugate(w))
+        M_list.append(numpy.conjugate(m1))
+
+    w_all = numpy.concatenate(W_list)
+    m1_all = numpy.concatenate(M_list)
+
+    # Fit on the sample data from m1
+    p, q = fit_pade(w_all, m1_all, deg_p=deg_p, deg_q=deg_q,
+                    ridge_lambda=ridge_lambda)
+
+    return p, q

freealg/_geometric_form/_continuation_genus1.py

@@ -0,0 +1,275 @@
+# 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
+import scipy
+from ._elliptic_functions import ellipj
+
+__all__ = ['mobius_X', 'mobius_z', 'mobius_lambda', 'legendre_Y',
+           'fit_rational_xy', '_poly_eval', 'eval_rational_xy',
+           'eval_rational_xy_select', 'generate_rational_xy']
+
+
+# =============
+# sqrt pos imag
+# =============
+
+def sqrt_pos_imag(z):
+    """
+    Square root on a branch cut with always positive imaginary part.
+    """
+
+    sq = numpy.sqrt(z)
+    sq = numpy.where(sq.imag < 0, -sq, sq)
+
+    return sq
+
+
+# ========
+# mobius X
+# ========
+
+def mobius_X(z, a1, b1, a2, b2):
+
+    A = (a2 - b2) / (a2 - a1)
+    return A * (z - a1) / (z - b2)
+
+
+# ========
+# mobius z
+# ========
+
+def mobius_z(X, a1, b1, a2, b2):
+
+    A = (a2 - b2) / (a2 - a1)
+    return (X * b2 - A * a1) / (X - A)
+
+
+# =============
+# mobius lambda
+# =============
+
+def mobius_lambda(a1, b1, a2, b2):
+
+    num = (b1 - a1) * (b2 - a2)
+    den = (a2 - a1) * (b2 - b1)
+
+    return num / den
+
+
+# ==========
+# legendre Y
+# ==========
+
+def legendre_Y(X, lam):
+
+    D = X * (1.0 - X) * (X - lam)
+    Y = sqrt_pos_imag(D)
+
+    flip = numpy.real(X) > 1.0
+    Y = numpy.where(flip, -Y, Y)
+
+    return Y
+
+
+# ===============
+# fit rational xy
+# ===============
+
+def fit_rational_xy(X, Y, F, deg_p0=12, deg_p1=12, deg_q=12, ridge_lambda=0.0,
+                    weights=None):
+
+    n = F.size
+
+    max_deg = max(deg_p0, deg_p1, deg_q)
+    xp = numpy.ones((n, max_deg + 1), dtype=complex)
+    for k in range(1, max_deg + 1):
+        xp[:, k] = xp[:, k - 1] * X
+
+    A0 = xp[:, :deg_p0 + 1]
+    A1 = (Y[:, None] * xp[:, :deg_p1 + 1])
+    Aq = (-F[:, None] * xp[:, 1:deg_q + 1])
+
+    A = numpy.hstack([A0, A1, Aq])
+    bvec = F
+
+    if weights is not None:
+        w = numpy.sqrt(numpy.maximum(weights, 0.0))
+        A = A * w[:, None]
+        bvec = bvec * w
+
+    s = numpy.linalg.norm(A, axis=0)
+    s[s == 0] = 1.0
+    As = A / s[None, :]
+
+    if ridge_lambda is None:
+        ridge_lambda = 0.0
+
+    if ridge_lambda > 0.0:
+        alpha = ridge_lambda
+        n_coef = As.shape[1]
+        A_aug = numpy.vstack(
+            [As, numpy.sqrt(alpha) * numpy.eye(n_coef, dtype=complex)])
+        b_aug = numpy.concatenate([bvec, numpy.zeros(n_coef, dtype=complex)])
+        coef, _, _, _ = numpy.linalg.lstsq(A_aug, b_aug, rcond=None)
+    else:
+        coef, _, _, _ = numpy.linalg.lstsq(As, bvec, rcond=None)
+
+    coef = coef / s
+
+    i0 = deg_p0 + 1
+    i1 = i0 + (deg_p1 + 1)
+
+    p0 = coef[:i0]
+    p1 = coef[i0:i1]
+
+    q = numpy.zeros(deg_q + 1, dtype=complex)
+    q[0] = 1.0
+    q[1:] = coef[i1:]
+
+    return p0, p1, q
+
+
+# =========
+# poly eval
+# =========
+
+def _poly_eval(x, c):
+
+    y = numpy.zeros_like(x, dtype=complex)
+    for k in range(len(c) - 1, -1, -1):
+        y = y * x + c[k]
+
+    return y
+
+
+# ================
+# eval rational xy
+# ================
+
+def eval_rational_xy(X, Y, p0, p1, q, alt_branch):
+    num0 = _poly_eval(X, p0)
+    num1 = _poly_eval(X, p1)
+    den = _poly_eval(X, q)
+
+    m_plus = (num0 + Y * num1) / den
+    m_minus = (num0 - Y * num1) / den
+
+    if alt_branch:
+        return m_minus
+    return m_plus
+
+
+# =======================
+# eval rational xy select
+# =======================
+
+def eval_rational_xy_select(z, X, Y, p0, p1, q, alt_branch):
+
+    m_plus = eval_rational_xy(X, Y, p0, p1, q, alt_branch=False)
+    m_minus = eval_rational_xy(X, Y, p0, p1, q, alt_branch=True)
+
+    mask_p = numpy.imag(z) >= 0.0
+
+    pick_plus_p = (numpy.imag(m_plus) >= 0.0)
+    pick_plus_m = (numpy.imag(m_plus) <= 0.0)
+
+    pick_plus = numpy.where(mask_p, pick_plus_p, pick_plus_m)
+
+    m1 = numpy.where(pick_plus, m_plus, m_minus)
+    m2 = numpy.where(pick_plus, m_minus, m_plus)
+
+    if alt_branch:
+        return m2
+    return m1
+
+
+# ====================
+# generate rational xy
+# ====================
+
+def generate_rational_xy(m1_fn, a1, b1, a2, b2, deg_p0=12, deg_p1=12, deg_q=12,
+                         n_samples=4096, r=0.25, n_r=1, r_min=None,
+                         ridge_lambda=0.0):
+
+    if n_samples % 2 != 0:
+        raise ValueError('n_samples should be even.')
+
+    if n_r is None or n_r < 1:
+        n_r = 1
+
+    lam = mobius_lambda(a1, b1, a2, b2)
+
+    r = float(r)
+    if r > 0.6:
+        r = 0.30
+
+    if r_min is None:
+        r_min = 0.25 * r
+
+    if n_r == 1:
+        vs = numpy.array([r], dtype=float)
+    else:
+        vs = numpy.linspace(r_min, r, n_r)
+
+    K = float(scipy.special.ellipk(lam))
+    Kp = float(scipy.special.ellipk(1.0 - lam))
+
+    omega1 = 2.0 * K
+    shift2 = 1j * Kp
+
+    n_half = n_samples // 2
+    t = (numpy.arange(n_half) + 0.5) / n_half
+    u_re = omega1 * t
+
+    X_list = []
+    Y_list = []
+    F_list = []
+
+    for v in vs:
+        for sh in (0.0, shift2):
+
+            u = u_re + sh + 1j * v
+
+            sn, cn, dn, _ = ellipj(u, lam)
+
+            X = lam * (sn * sn)
+
+            Y = legendre_Y(X, lam)
+            Y_ref = 1j * lam * (sn * cn * dn)
+            if sh != 0.0:
+                Y_ref = -Y_ref
+
+            flip = numpy.real(Y * numpy.conjugate(Y_ref)) < 0.0
+            Y = numpy.where(flip, -Y, Y)
+
+            z = mobius_z(X, a1, b1, a2, b2)
+            f = m1_fn(z)
+
+            X_list.append(X)
+            Y_list.append(Y)
+            F_list.append(f)
+
+            X_list.append(numpy.conjugate(X))
+            Y_list.append(numpy.conjugate(Y))
+            F_list.append(numpy.conjugate(f))
+
+    X_all = numpy.concatenate(X_list)
+    Y_all = numpy.concatenate(Y_list)
+    f_all = numpy.concatenate(F_list)
+
+    p0, p1, q = fit_rational_xy(X_all, Y_all, f_all, deg_p0=deg_p0,
+                                deg_p1=deg_p1, deg_q=deg_q,
+                                ridge_lambda=ridge_lambda)
+
+    return p0, p1, q, lam

freealg/_geometric_form/_elliptic_functions.py

@@ -0,0 +1,174 @@
+# 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.
+
+
+"""
+The Jacobi elliptic functions sn, cn, and dn are implemented in:
+
+1. scipy.special: but they do not accept complex inputs
+2. mpmath: but they do not have vectorization (very slow)
+
+Because of these, we implemented sn, cn, and dn ourselves. There is no need to
+implement the elliptic K function, since scipy.special already has it (and
+vectorized) and we only need that for real inputs, not complex inputs. So, no
+need to re-implement the scipy.special.ellipk function.
+"""
+
+
+# =======
+# Imports
+# =======
+
+import numpy
+from scipy import special
+
+__all__ = ['ellipj']
+
+
+# ===========
+# jacobi nome
+# ===========
+
+def _jacobi_nome(m):
+    """
+    Compute the Jacobi nome q = exp(-pi K(1-m)/K(m)) and the complete elliptic
+    integrals K(m), K(1-m).
+    """
+
+    m = float(m)
+    K = special.ellipk(m)
+    Kp = special.ellipk(1.0 - m)
+    q = numpy.exp(-numpy.pi * Kp / K)
+
+    return q, K, Kp
+
+
+# =======
+# theta 1
+# =======
+
+def _theta1(v, q, n_terms):
+    """
+    Jacobi theta_1(v,q) via truncated Fourier series (vectorized for complex
+    v).
+    """
+
+    v = numpy.asarray(v, dtype=complex)
+    n = numpy.arange(n_terms, dtype=float)
+    a = 2.0 * n + 1.0
+    logq = numpy.log(q)
+    w = numpy.exp(((n + 0.5) ** 2) * logq) * ((-1.0) ** n)
+    s = numpy.sin(a[:, None] * v[None, :])
+
+    return 2.0 * (w[:, None] * s).sum(axis=0)
+
+
+# =======
+# theta 2
+# =======
+
+def _theta2(v, q, n_terms):
+    """
+    Jacobi theta_2(v,q) via truncated Fourier series (vectorized for complex
+    v).
+    """
+
+    v = numpy.asarray(v, dtype=complex)
+    n = numpy.arange(n_terms, dtype=float)
+    a = 2.0 * n + 1.0
+    logq = numpy.log(q)
+    w = numpy.exp(((n + 0.5) ** 2) * logq)
+    c = numpy.cos(a[:, None] * v[None, :])
+
+    return 2.0 * (w[:, None] * c).sum(axis=0)
+
+
+# =======
+# theta 3
+# =======
+
+def _theta3(v, q, n_terms):
+    """
+    Jacobi theta_3(v,q) via truncated Fourier series (vectorized for complex
+    v).
+    """
+
+    v = numpy.asarray(v, dtype=complex)
+    n = numpy.arange(1, n_terms + 1, dtype=float)
+    logq = numpy.log(q)
+    w = numpy.exp((n ** 2) * logq)
+    c = numpy.cos((2.0 * n)[:, None] * v[None, :])
+
+    return 1.0 + 2.0 * (w[:, None] * c).sum(axis=0)
+
+
+# =======
+# theta 4
+# =======
+
+def _theta4(v, q, n_terms):
+    """
+    Jacobi theta_4(v,q) via truncated Fourier series (vectorized for complex
+    v).
+    """
+
+    v = numpy.asarray(v, dtype=complex)
+    n = numpy.arange(1, n_terms + 1, dtype=float)
+    logq = numpy.log(q)
+    w = numpy.exp((n ** 2) * logq) * ((-1.0) ** n)
+    c = numpy.cos((2.0 * n)[:, None] * v[None, :])
+
+    return 1.0 + 2.0 * (w[:, None] * c).sum(axis=0)
+
+
+# ======
+# ellipj
+# ======
+
+def ellipj(u, m, n_terms=None):
+    """
+    Vectorized Jacobi elliptic sn, cn, dn for complex u and real m in (0,1),
+    computed from theta functions.
+    """
+
+    u = numpy.asarray(u, dtype=complex)
+    q, K, Kp = _jacobi_nome(m)
+
+    if n_terms is None:
+        if q < 1e-6:
+            n_terms = 16
+        elif q < 1e-3:
+            n_terms = 24
+        elif q < 1e-2:
+            n_terms = 32
+        elif q < 5e-2:
+            n_terms = 48
+        else:
+            n_terms = 80
+
+    v = (numpy.pi / (2.0 * K)) * u
+    v_flat = v.ravel()
+
+    th1 = _theta1(v_flat, q, n_terms)
+    th2 = _theta2(v_flat, q, n_terms)
+    th3 = _theta3(v_flat, q, n_terms)
+    th4 = _theta4(v_flat, q, n_terms)
+
+    th2_0 = _theta2(numpy.array([0.0], dtype=complex), q, n_terms)[0]
+    th3_0 = _theta3(numpy.array([0.0], dtype=complex), q, n_terms)[0]
+    th4_0 = _theta4(numpy.array([0.0], dtype=complex), q, n_terms)[0]
+
+    sn = (th3_0 * th1) / (th2_0 * th4)
+    cn = (th4_0 * th2) / (th2_0 * th4)
+    dn = (th4_0 * th3) / (th3_0 * th4)
+
+    sn = sn.reshape(v.shape)
+    cn = cn.reshape(v.shape)
+    dn = dn.reshape(v.shape)
+
+    return sn, cn, dn, q

freealg/_geometric_form/_sphere_maps.py

@@ -0,0 +1,63 @@
+# 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_genus0 import joukowski_z
+
+__all__ = ['make_sphere_grid', 'eval_on_sphere']
+
+
+# ===========================
+# stereographic w from sphere
+# ===========================
+
+def stereographic_w_from_sphere(X, Y, Z):
+
+    den = 1.0 - Z
+    w = (X + 1j * Y) / den
+    return w
+
+
+# ================
+# make sphere grid
+# ================
+
+def make_sphere_grid(n_theta=100, n_phi=50):
+
+    theta = numpy.linspace(0.0, 2.0 * numpy.pi, n_theta, endpoint=False)
+    u = numpy.linspace(0.0, 1.0, n_phi)
+    phi = numpy.pi * (u * u)
+    TH, PH = numpy.meshgrid(theta, phi)
+    X = numpy.sin(PH) * numpy.cos(TH)
+    Y = numpy.sin(PH) * numpy.sin(TH)
+    Z = numpy.cos(PH)
+
+    return X, Y, Z
+
+
+# ==============
+# eval on sphere
+# ==============
+
+def eval_on_sphere(X, Y, Z, a, b, m1_fn, m2_fn, z_eps=1e-9):
+
+    Zc = numpy.minimum(Z, 1.0 - 1e-12)
+    w = stereographic_w_from_sphere(X, Y, Zc)
+    z = joukowski_z(w, a, b)
+    z = z + 1j * z_eps
+    m1 = m1_fn(z)
+    m2 = m2_fn(z)
+    out = numpy.array(m2, copy=True)
+    out[numpy.abs(w) <= 1.0] = m1[numpy.abs(w) <= 1.0]
+
+    return out