freealg 0.6.2__py3-none-any.whl → 0.7.0__py3-none-any.whl

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

freealg/_geometric_form/_torus_maps.py

@@ -0,0 +1,118 @@
+# 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
+from ._continuation_genus1 import _poly_eval
+
+__all__ = ['make_torus_grid', 'u_from_angles', 'eval_fitted_m_on_torus']
+
+
+# ===============
+# make torus grid
+# ===============
+
+def make_torus_grid(n_theta=101, n_phi=101, R=1.0, r=0.35):
+    """
+    Returns an embedded torus mesh (X,Y,Z) with angles (TH,PH).
+    TH: around major circle, PH: around tube.
+    """
+
+    if n_phi % 2 == 0:
+        raise ValueError('n_phi should be odd number to avoid rendering ' +
+                         'issue at phi=0.')
+
+    # fundamental angles (no endpoint duplicates)
+    theta = numpy.linspace(0.0, 2.0*numpy.pi, int(n_theta), endpoint=False)
+    phi = numpy.linspace(0.0, 2.0*numpy.pi, int(n_phi),   endpoint=False)
+
+    TH, PH = numpy.meshgrid(theta, phi)  # shapes (n_phi, n_theta)
+
+    # --- wrap/close the grid by appending first row/col ---
+    TH = numpy.vstack([TH, TH[0:1, :]])     # add phi seam row
+    PH = numpy.vstack([PH, PH[0:1, :]])
+
+    TH = numpy.hstack([TH, TH[:, 0:1]])     # add theta seam col
+    PH = numpy.hstack([PH, PH[:, 0:1]])
+
+    # torus embedding
+    X = (R + r*numpy.cos(PH)) * numpy.cos(TH)
+    Y = (R + r*numpy.cos(PH)) * numpy.sin(TH)
+    Z = r * numpy.sin(PH)
+
+    return X, Y, Z, TH, PH
+
+
+# =============
+# u from angles
+# =============
+
+def u_from_angles(TH, PH, lam, center=(0.0, 0.0)):
+    """
+    Map angles (TH,PH) in [0,2pi)^2 to the elliptic u-plane fundamental cell.
+
+    u = u0 + (omega1/2pi)*TH + (omega2/2pi)*PH,
+    omega1 = 2K(m), omega2 = 2 i K(1-m).
+    """
+
+    m = float(lam)
+    K = scipy.special.ellipk(m)
+    Kp = scipy.special.ellipk(1.0 - m)
+    omega1 = 2.0 * K
+    omega2 = 2.0j * Kp
+
+    u0 = complex(center[0], center[1])  # shift inside the fundamental domain
+    u = u0 + (omega1/(2.0*numpy.pi))*TH + (omega2/(2.0*numpy.pi))*PH
+
+    return u
+
+
+# ==========================
+# evaluate fitted m on torus
+# ==========================
+
+def eval_fitted_m_on_torus(u, a1, b1, a2, b2, p0, p1, q, lam):
+    """
+    Evaluate the *uniformized* branch m(u) = (p0(X)+Y p1(X))/q(X)
+    with X = lam * sn(u)^2 and Y chosen to match the elliptic derivative sign.
+
+    Requires jacobi_ellipj(u, m) that supports complex arrays and returns
+    (sn,cn,dn).
+    """
+
+    sn, cn, dn, _ = ellipj(u, lam)
+
+    X = lam * (sn * sn)
+
+    # canonical algebraic Y from curve: Y^2 = X(1-X)(X-lam)
+    D = X * (1.0 - X) * (X - lam)
+
+    Y = numpy.sqrt(D)
+
+    # enforce "positive imag" convention first
+    Y = numpy.where(numpy.imag(Y) < 0.0, -Y, Y)
+
+    # now align sign with elliptic reference:
+    # for Legendre normalization: dX/du = 2 i Y_ref, where Y_ref is equal to
+    # i*lam*sn*cn*dn (up to consistent conventions)
+    Y_ref = 1j * lam * (sn * cn * dn)
+    flip = numpy.real(Y * numpy.conjugate(Y_ref)) < 0.0
+    Y = numpy.where(flip, -Y, Y)
+
+    num0 = _poly_eval(X, p0)
+    num1 = _poly_eval(X, p1)
+    den = _poly_eval(X, q)
+
+    m = (num0 + Y * num1) / den
+    return m