syncem 0.6.0__py3-none-any.whl

syncem/__init__.py

@@ -0,0 +1,98 @@
+"""syncem: moment expansion of synchrotron emission (JAX + Equinox).
+
+Note: importing this package enables ``jax_enable_x64`` (float64) as a global
+side effect, because the Bessel quadrature and the likelihood-style sums need
+float64 for accuracy. Do this once at the top of any driver script, before any
+array is created.
+
+Modules
+-------
+bessel           : differentiable J_n, J_n', K_nu (integral representations)
+stokes           : exact Schott harmonic Stokes (I, Q, V) with correct
+                   normalisation
+derivatives      : derivative spectra dS/dp, d^2S/dp^2 via autodiff
+expansion         : eqx.Module for quadratic Taylor averaging (mean/covariance)
+cumulants        : finite Bell moment contractions from supplied cumulants
+ultrarel         : ultra-relativistic F(x), G(x) fast path
+rm               : rotation measure (RM) constants + Burn depolarisation
+faraday          : joint emission-depth averages and finite responses
+transfer         : Mueller matrix + numerical slab/LOS matrix exponentials
+solutions        : analytic limiting solutions (thin/rotation/self-absorbed)
+kirchhoff        : emissivity/absorption moment expansion + Kirchhoff closure
+conversion       : cold-plasma Faraday rotation/conversion coefficients + moments
+magnus           : Magnus expansion (Omega1, Omega2) of the LOS transfer
+los_moments      : explicit reduced-units and physical-CGS moment-driven slabs
+sed              : frequency-domain SED (spectral index, curvature, abs. emissivity)
+model            : finite joint response of channel Stokes spectra and spectral
+                   fits (subpackage: index, errors, channels, phase, moments,
+                   bounds, assumptions, kernels, harmonic,
+                   high_order_harmonic, screens, basis, high_order_basis,
+                   predict, adapters, fit.observation, fit.linear,
+                   fit.nonlinear, fit.diagnostics, fit.result, fit.layout,
+                   fit.reduction, fit.combinations, fit.combination_result);
+                   import its modules by full path
+"""
+
+from __future__ import annotations
+
+__version__ = "0.6.0"
+
+import jax
+
+# Enable float64 for accuracy of Bessel quadrature and likelihood-style sums.
+jax.config.update("jax_enable_x64", True)
+
+from . import (  # noqa: E402
+    bessel,
+    conversion,
+    cumulants,
+    derivatives,
+    kirchhoff,
+    los_moments,
+    magnus,
+    expansion,
+    rm,
+    faraday,
+    sed,
+    solutions,
+    stokes,
+    transfer,
+    ultrarel,
+)  # noqa: E402,F401
+from .stokes import larmor_power, stokes_harmonic  # noqa: E402,F401
+from .derivatives import derivative_spectra  # noqa: E402,F401
+from .expansion import (  # noqa: E402,F401
+    CumulantExpansion,
+    QuadraticTaylorExpansion,
+    build_expansion,
+    mixed_moments,
+)
+from .ultrarel import F, G  # noqa: E402,F401
+
+__all__ = [
+    "__version__",
+    "bessel",
+    "stokes",
+    "derivatives",
+    "expansion",
+    "cumulants",
+    "ultrarel",
+    "rm",
+    "faraday",
+    "transfer",
+    "solutions",
+    "kirchhoff",
+    "conversion",
+    "magnus",
+    "los_moments",
+    "sed",
+    "stokes_harmonic",
+    "larmor_power",
+    "derivative_spectra",
+    "CumulantExpansion",
+    "QuadraticTaylorExpansion",
+    "build_expansion",
+    "mixed_moments",
+    "F",
+    "G",
+]

syncem/_expm.py

@@ -0,0 +1,125 @@
+"""Matrix exponential whose branches survive ``vmap`` under reverse mode.
+
+``jax.scipy.linalg.expm`` picks its Pade degree with ``lax.switch`` and each
+squaring with ``lax.cond``, on values of the operand. Under ``vmap`` with a
+batched operand JAX turns both into selects that wrap every operand in
+``stop_gradient``. When the exponential is differentiated inside a
+``custom_jvp`` rule (``transfer._tangent``), those operands include the
+rule's tangents, and ``stop_gradient`` of a tangent cannot be transposed, so
+``grad`` of a sum over ``vmap`` in ``K`` would raise ``NotImplementedError``.
+Here the degree is fixed (Pade 13) and a squaring runs while any member of the
+batch needs it, selected per member.
+"""
+
+from __future__ import annotations
+
+from functools import partial
+
+import jax
+import jax.numpy as jnp
+
+# Higham (2005) Pade-13 coefficients b_0..b_13 and theta_13 (float64). The
+# approximant P/Q is unchanged by a common factor, so the coefficients are held
+# at the exact scale 2^-40: with b_0 = 6.5e16, V - U ~ b_0 I and the transpose
+# of the solve divided cotangents by b_0, which flushed float32 reverse-mode
+# derivatives for small sources (tests/test_transfer_scale_and_budget.py).
+# Dividing by b_0 fully overflows rev-over-fwd Hessians near 1e306.
+_PADE13 = tuple(
+    b * 2.0**-40
+    for b in (64764752532480000.0, 32382376266240000.0, 7771770303897600.0,
+              1187353796428800.0, 129060195264000.0, 10559470521600.0,
+              670442572800.0, 33522128640.0, 1323241920.0, 40840800.0,
+              960960.0, 16380.0, 182.0, 1.0)
+)  # fmt: skip
+_THETA13 = 5.371920351148152
+
+
+@partial(jax.jit, static_argnames="max_squarings")
+def expm_pade13(A, max_squarings):
+    """``exp(A)`` by Pade 13 on ``A 2^-n``, ``n = max(0, ceil(log2(|A|_1 / theta_13)))``
+    (Higham 2005: backward error below the unit roundoff).
+
+    ``jax.scipy.linalg.expm`` takes ``floor`` and so evaluates Pade 13 up to
+    ``2 theta_13``; for a rotation-dominated ``A`` just below ``2^(m+1) theta_13``
+    that costs relative errors from 1e-10 (m = 0) to 1e-1 (m = 31) where
+    ``ceil`` stays at the conditioning limit (``tests/test_transfer_expm_boundaries.py``).
+    Where ``ceil`` exceeds ``max_squarings`` but ``floor`` does not, ``floor``
+    is used, so the result is NaN exactly where ``expm``'s is in float64 (in
+    float32 ``expm`` used Pade 7, ``theta_7 = 3.93``, and a lower NaN bound).
+    Used for the slab value (5x5) and its propagators (8x8, ``propagators``).
+    Compiled, as ``expm`` is: op by op, XLA CPU rounded it an ulp differently."""
+    norm = jax.lax.stop_gradient(jnp.max(jnp.sum(jnp.abs(A), axis=0)))
+    ratio = jnp.log2(norm / _THETA13)
+    low, high = jnp.maximum(0, jnp.floor(ratio)), jnp.maximum(0, jnp.ceil(ratio))
+    return expm_squared(A, jnp.where(high > max_squarings, low, high), max_squarings)
+
+
+@partial(jax.jit, static_argnames="max_squarings")
+def expm_squared(A, n, max_squarings):
+    """Pade-13 of ``A / 2^n`` squared ``n`` times (NaN if ``n > max_squarings``).
+    A squaring step runs while any member needs it (``any_member``) and is
+    applied per member (``batched_only``)."""
+    dot = partial(jnp.matmul, precision=jax.lax.Precision.HIGHEST)
+    A = A / 2 ** n.astype(A.dtype)
+    b, ident = _PADE13, jnp.eye(A.shape[0], dtype=A.dtype)
+    A2 = dot(A, A)
+    A4 = dot(A2, A2)
+    A6 = dot(A4, A2)
+    U = dot(A6, b[13] * A6 + b[11] * A4 + b[9] * A2)
+    U = dot(A, U + b[7] * A6 + b[5] * A4 + b[3] * A2 + b[1] * ident)
+    V = dot(A6, b[12] * A6 + b[10] * A4 + b[8] * A2)
+    V = V + b[6] * A6 + b[4] * A4 + b[2] * A2 + b[0] * ident
+    R = jnp.linalg.solve(V - U, V + U)
+
+    def squaring(R, i):
+        def square(R):
+            return jnp.where(batched_only(i < n), dot(R, R), R)
+
+        return jax.lax.cond(any_member(i < n), square, lambda R: R, R), None
+
+    steps = jnp.arange(max_squarings, dtype=n.dtype)
+    R = jax.lax.scan(squaring, R, steps)[0]
+    return jnp.where(n > max_squarings, jnp.nan, R)
+
+
+def propagators(K, ds, max_squarings):
+    """``(Phi, G) = (e^{-K ds}, int_0^ds e^{-K t} dt)`` from one 8x8 exponential.
+
+    The identity block (not ``I ds``) keeps the norm at ``max(|K ds|, 1)``, at
+    most that of the slab value's 5x5 exponential (same ``expm_pade13``), so
+    they exceed ``max_squarings`` only where the value is refused. There is no
+    branch on a member's value, so a ``jax.jvp`` of it inside a ``custom_jvp``
+    rule transposes under a batched ``vmap``."""
+    A = jnp.zeros((8, 8), dtype=K.dtype)
+    A = A.at[:4, :4].set(-K * ds).at[:4, 4:].set(jnp.eye(4, dtype=K.dtype))
+    E = expm_pade13(A, max_squarings)
+    return E[:4, :4], E[:4, 4:] * ds
+
+
+@jax.custom_batching.custom_vmap
+def any_member(flag):
+    """``flag``; under ``vmap`` its ``any`` over the batch, unbatched, so that a
+    ``lax.cond`` on it stays a branch instead of evaluating both sides. The rule
+    re-enters ``any_member``, so an enclosing ``vmap`` reduces its axis too (a
+    plain ``jnp.any`` stayed batched there, and reverse mode through nested
+    ``vmap`` hit ``stop_gradient`` in the select that replaced the ``cond``)."""
+    return flag
+
+
+@any_member.def_vmap
+def _any_member_vmap(axis_size, in_batched, flag):
+    return any_member(jnp.any(flag, axis=0) if in_batched[0] else flag), False
+
+
+@jax.custom_batching.custom_vmap
+def batched_only(flag):
+    """``True``; under ``vmap`` with a batched ``flag``, ``flag``. In a branch
+    taken on ``any_member(flag)`` it selects the members per step, and folds
+    to nothing when unbatched (a ``where`` on ``flag`` cost 1.4x in reverse).
+    Unbatched at one level, it re-enters itself for an enclosing ``vmap``."""
+    return jnp.ones_like(flag)
+
+
+@batched_only.def_vmap
+def _batched_only_vmap(axis_size, in_batched, flag):
+    return (flag if in_batched[0] else batched_only(flag)), in_batched[0]

syncem/bessel.py

@@ -0,0 +1,268 @@
+"""Differentiable integer J and modified K functions with explicit quadrature.
+
+Integer J uses a periodic trapezoidal rule, with an exactly equivalent complex
+contour shift for ``0 < |x| < n`` to resolve exponentially small values without
+subtracting order-one oscillations. The contour is held fixed during autodiff:
+the integral is invariant under that shift, so this differentiates the same
+analytic function rather than the arbitrary contour-selection algorithm.
+
+For a concrete order the default grid resolves ``|x| <= n`` (the harmonic domain).
+For a traced order the default is 2048 nodes; choose a larger static ``n_nodes``
+when needed. Calls outside ``|n| + |x| <= n_nodes/2`` return NaN, rather than silently
+aliasing unresolved oscillations. This conservative resolution guard is not a
+roundoff or quadrature-error certificate. Compare increased resolution and an
+independent reference for a new domain, including the required derivatives.
+
+K uses one continuous cosh integral; there is no asymptotic switch at x=0.01.
+Its finite upper limit has the explicit bound returned by ``bessel_kn_tail_bound``.
+The quadrature error is separate. Finite independent tests cover orders 0, 2/3,
+5/3 over x=1e-12..100 and integer J through n=3200; they do not establish uniform
+accuracy outside those cases. Order, node count and tail cutoff are static;
+positive arguments are differentiated. SciPy is used only in independent tests.
+"""
+
+from __future__ import annotations
+
+from functools import lru_cache
+import math
+import numbers
+
+import jax
+import jax.numpy as jnp
+import numpy as np
+
+
+@lru_cache(maxsize=20)
+def _periodic_grid(n_nodes):
+    if n_nodes < 16:
+        raise ValueError("n_nodes must be at least 16")
+    t = 2.0 * np.pi * np.arange(n_nodes) / n_nodes
+    return np.sin(t), np.cos(t), t
+
+
+@lru_cache(maxsize=8)
+def _legendre_grid(n_nodes):
+    if n_nodes < 16:
+        raise ValueError("n_nodes must be at least 16")
+    return np.polynomial.legendre.leggauss(n_nodes)
+
+
+def _j_resolution(n, x, n_nodes):
+    if n_nodes is not None:
+        return int(n_nodes)
+    # No host materialisation of traced or device arrays in an online graph.
+    if isinstance(n, (numbers.Real, np.ndarray, list, tuple)):
+        extent = float(np.max(np.abs(np.asarray(n))))
+        if isinstance(x, (numbers.Real, np.ndarray, list, tuple)):
+            extent = max(extent, float(np.max(np.abs(np.asarray(x)))))
+        count = 2 ** math.ceil(math.log2(max(128, 4 * extent + 64)))
+        if count > 65536:
+            raise ValueError(
+                "automatic J quadrature exceeds 65536 nodes; choose a validated static n_nodes or a suitable reference approximation"
+            )
+        return count
+    return 2048
+
+
+def bessel_jn_and_prime(n, x, *, n_nodes=None):
+    """Return integer-order (J_n(x), J_n'(x)); broadcast n and x.
+
+    ``n >= 0`` must be integral. For JIT with a varying order, ``n_nodes`` is
+    a static resolution choice; unresolved inputs are NaN. At x=0 the analytic
+    function values are supplied; use a nonzero argument for autodiff.
+    """
+    count = _j_resolution(n, x, n_nodes)
+    sin_t, cos_t, t = _periodic_grid(count)
+    n, x = jnp.broadcast_arrays(jnp.asarray(n), jnp.asarray(x))
+    ax = jnp.abs(x)
+    safe_x = jnp.where(ax > 0, ax, 1.0)
+    shifted = (n > safe_x) & (n > 0)
+    # Fix the contour while differentiating. Scaled sinh/cosh avoid overflow
+    # when the shift is large but x*sinh(a) and x*cosh(a) remain modest.
+    d = jax.lax.stop_gradient(jnp.sqrt(jnp.maximum(n**2 - safe_x**2, 0.0)))
+    shift = jax.lax.stop_gradient(
+        jnp.where(shifted, jnp.log(jnp.maximum(n + d, 1.0)) - jnp.log(safe_x), 0.0)
+    )
+    reference_x = jax.lax.stop_gradient(safe_x)
+    ratio = safe_x / reference_x
+    xsinh = jnp.where(shifted, d * ratio, 0.0)
+    xcosh = jnp.where(shifted, n * ratio, safe_x)
+    phase = n[..., None] * t - xcosh[..., None] * sin_t
+    envelope = jnp.exp(xsinh[..., None] * (cos_t - 1.0))
+    pref = jnp.exp(xsinh - n * shift)
+    value = pref * jnp.mean(envelope * jnp.cos(phase), axis=-1)
+    prime = (
+        pref
+        * jnp.mean(
+            envelope
+            * (
+                xsinh[..., None] * cos_t * jnp.cos(phase)
+                + xcosh[..., None] * sin_t * jnp.sin(phase)
+            ),
+            axis=-1,
+        )
+        / safe_x
+    )
+    parity = jnp.where(n % 2 == 0, 1.0, -1.0)
+    value = jnp.where(x < 0, parity * value, value)
+    prime = jnp.where(x < 0, -parity * prime, prime)
+    value = jnp.where(ax == 0, jnp.where(n == 0, 1.0, 0.0), value)
+    prime = jnp.where(ax == 0, jnp.where(n == 1, 0.5, 0.0), prime)
+    valid = (n >= 0) & (n == jnp.floor(n)) & (n + ax <= count / 2)
+    return jnp.where(valid, value, jnp.nan), jnp.where(valid, prime, jnp.nan)
+
+
+def bessel_jn(n, x, *, n_nodes=None):
+    """Integer J_n; see ``bessel_jn_and_prime`` for the resolution contract."""
+    return bessel_jn_and_prime(n, x, n_nodes=n_nodes)[0]
+
+
+def bessel_jn_prime(n, x, *, n_nodes=None):
+    """Argument derivative of integer J_n, evaluated on the same contour."""
+    return bessel_jn_and_prime(n, x, n_nodes=n_nodes)[1]
+
+
+def bessel_jn_neighbours(n, x, *, n_nodes=None):
+    """``(J_{n-1}(x), J_{n+1}(x), J_n'(x))`` from one quadrature; broadcast n and x.
+
+    Dimensionless. ``n >= 1`` must be integral (traced floats allowed, as in
+    ``bessel_jn_and_prime``). One periodic trapezoidal rule on the contour
+    ``t + i a`` chosen for order ``n`` (``cosh a = n/|x|`` when ``n > |x|``,
+    else ``a = 0``) serves all three functions: with ``phi = n t - |x| cosh(a)
+    sin t`` and ``E = exp(|x| sinh(a) cos t)``, the order-``k`` integrand is
+    ``exp(-k a) E e^{i (phi + (k - n) t)}``, so ``J_{n+-1}`` use
+    ``cos(phi -+ t) = cos(phi) cos t +- sin(phi) sin t`` and ``J_n'`` uses the
+    argument derivative of the order-``n`` integrand, all from one set of
+    ``cos(phi)``, ``sin(phi)`` and ``E`` (three transcendental evaluations per
+    node instead of nine). The shift identity is exact for every integer
+    order (the integrand is entire and ``2 pi``-periodic), so no fallback to
+    separate contours is needed; the precision of the neighbours on the
+    order-``n`` contour is checked, not bounded, against SciPy over the
+    harmonic resolution classes ``n <= 1008``, ``x`` from ``1e-8`` through
+    the transition ``x ~ n`` to the resolution edge (relative differences
+    below ``1e-12`` in ``tests/model/test_memory_numerics.py``).
+
+    Resolution guard: ``n + 1 + |x| <= n_nodes/2`` (the default count
+    resolves ``|x| <= n + 1`` for a concrete order and is 2048 for a traced
+    one); unresolved inputs, ``n < 1`` and non-integral orders are NaN. This
+    guard is not an accuracy certificate. The contour is held fixed under
+    autodiff (see the module docstring); at ``x = 0`` the analytic values are
+    returned with a zero derivative. ``bessel_jn_and_prime`` is unchanged.
+
+    Autodiff accuracy: ``jax.jacfwd``/``jax.grad`` through this quadrature
+    lose accuracy at ``|x| << n``. On the shifted contour the neighbour
+    integrands carry ``exp(-+a)`` with ``a = arccosh(n/|x|)``, and roundoff
+    grows by about ``(2n/|x|)^k`` at derivative order ``k``. At ``n = 1``,
+    ``x = 1e-6`` (default count, 128 nodes) the third derivatives of
+    ``J_0`` and ``J_1'`` come out as -83 and 128 against SciPy's 3.8e-7
+    and 3.1e-7; the second derivative of ``J_0`` is off by 1.6e-4. For
+    derivatives use the order recurrence ``J_n' = (J_{n-1} - J_{n+1})/2``
+    on the same contour (``syncem.model._bessel_recurrence``:
+    ``neighbours_recurrence``, ``derivative_coefficients``), which
+    ``HarmonicKernel(derivatives="analytic")`` uses; at the same point it
+    agrees with SciPy to 1e-15 through the third derivative
+    (``tests/test_bessel_autodiff_accuracy.py``; checked, not bounded).
+    """
+    top = n
+    if n_nodes is None and isinstance(n, (numbers.Real, np.ndarray, list, tuple)):
+        top = np.abs(np.asarray(n, dtype=float)) + 1.0
+    count = _j_resolution(top, x, n_nodes)
+    sin_t, cos_t, t = _periodic_grid(count)
+    n, x = jnp.broadcast_arrays(jnp.asarray(n), jnp.asarray(x))
+    ax = jnp.abs(x)
+    safe_x = jnp.where(ax > 0, ax, 1.0)
+    shifted = (n > safe_x) & (n > 0)
+    d = jax.lax.stop_gradient(jnp.sqrt(jnp.maximum(n**2 - safe_x**2, 0.0)))
+    shift = jax.lax.stop_gradient(
+        jnp.where(shifted, jnp.log(jnp.maximum(n + d, 1.0)) - jnp.log(safe_x), 0.0)
+    )
+    ratio = safe_x / jax.lax.stop_gradient(safe_x)
+    xsinh = jnp.where(shifted, d * ratio, 0.0)
+    xcosh = jnp.where(shifted, n * ratio, safe_x)
+    phase = n[..., None] * t - xcosh[..., None] * sin_t
+    envelope = jnp.exp(xsinh[..., None] * (cos_t - 1.0))
+    c, s = envelope * jnp.cos(phase), envelope * jnp.sin(phase)
+    lower = jnp.exp(xsinh - (n - 1.0) * shift) * jnp.mean(
+        c * cos_t + s * sin_t, axis=-1
+    )
+    upper = jnp.exp(xsinh - (n + 1.0) * shift) * jnp.mean(
+        c * cos_t - s * sin_t, axis=-1
+    )
+    prime = (
+        jnp.exp(xsinh - n * shift)
+        * jnp.mean(xsinh[..., None] * cos_t * c + xcosh[..., None] * sin_t * s, axis=-1)
+        / safe_x
+    )
+    # J_{n+-1}(-x) = (-1)^(n+-1) J_{n+-1}(x) and J_n'(-x) = (-1)^(n+1) J_n'(x).
+    flip = jnp.where(x < 0, jnp.where(n % 2 == 0, -1.0, 1.0), 1.0)
+    lower = jnp.where(ax == 0, jnp.where(n == 1, 1.0, 0.0), flip * lower)
+    upper = jnp.where(ax == 0, 0.0, flip * upper)
+    prime = jnp.where(ax == 0, jnp.where(n == 1, 0.5, 0.0), flip * prime)
+    valid = (n >= 1) & (n == jnp.floor(n)) & (n + 1 + ax <= count / 2)
+    return tuple(jnp.where(valid, v, jnp.nan) for v in (lower, upper, prime))
+
+
+def _log_cosh(x):
+    magnitude = jnp.abs(x)
+    return magnitude + jnp.log1p(jnp.exp(-2.0 * magnitude)) - jnp.log(2.0)
+
+
+def _modified_tmax(nu, x, tail_cutoff):
+    cutoff = tail_cutoff + 2 * abs(nu)
+    # arccosh(1+cutoff/x), expressed without overflowing cutoff/x at tiny x.
+    return jnp.log(x + cutoff + jnp.sqrt(cutoff * (cutoff + 2 * x))) - jnp.log(x)
+
+
+def _modified_integral(nu, x, *, divide_cosh=False, n_nodes=128, tail_cutoff=50.0):
+    """Integral exp(-x cosh t) cosh(nu t) [ / cosh(t)] dt."""
+    nu = float(nu)
+    if not np.isfinite(nu) or tail_cutoff <= 0:
+        raise ValueError("nu must be finite and tail_cutoff must be positive")
+    nodes, weights = _legendre_grid(int(n_nodes))
+    x = jnp.asarray(x)
+    safe_x = jnp.where(x > 0, x, 1.0)
+    tmax = _modified_tmax(nu, safe_x, tail_cutoff)
+    t = 0.5 * tmax[..., None] * (nodes + 1.0)
+    log_cosh_t = _log_cosh(t)
+    log_integrand = -jnp.exp(jnp.log(safe_x)[..., None] + log_cosh_t) + _log_cosh(
+        nu * t
+    )
+    if divide_cosh:
+        log_integrand = log_integrand - log_cosh_t
+    integrand = jnp.exp(log_integrand)
+    result = 0.5 * tmax * jnp.sum(weights * integrand, axis=-1)
+    return jnp.where(x > 0, result, jnp.where(x == 0, jnp.inf, jnp.nan))
+
+
+def _modified_tail_bound(nu, x, *, divide_cosh=False, tail_cutoff=50.0):
+    """Analytic upper bound for the omitted positive-t integral, x > 0.
+
+    ``cosh(t) >= exp(t)/2`` and ``cosh(nu t) <= exp(|nu| t)`` reduce the tail to an
+    incomplete Gamma integral. Bounding log(u/z)<= (u-z)/z gives the elementary
+    bound below. It excludes quadrature and floating-point error.
+    """
+    nu = abs(float(nu))
+    x = jnp.asarray(x)
+    tmax = _modified_tmax(nu, x, tail_cutoff)
+    z = 0.5 * jnp.exp(jnp.log(x) + tmax)
+    power = nu - int(divide_cosh)
+    log_pref = power * jnp.log(2.0 / x) + int(divide_cosh) * jnp.log(2.0)
+    slope = max(power - 1.0, 0.0) / z
+    bound = jnp.exp(log_pref - z + (power - 1) * jnp.log(z)) / (1.0 - slope)
+    return jnp.where((x > 0) & (slope < 1), bound, jnp.inf)
+
+
+def bessel_kn(nu, x, *, n_nodes=128, tail_cutoff=50.0):
+    """K_nu(x), x > 0; nonnegative static nu is sufficient for current uses.
+
+    A continuous cosh integral with a parameter-dependent upper limit is used.
+    K_nu(0)=infinity and negative x returns NaN. Node count and tail_cutoff must
+    be static under JIT. Derivatives inherit both quadrature and finite-tail
+    errors, and must be checked separately from function values.
+    """
+    return _modified_integral(nu, x, n_nodes=n_nodes, tail_cutoff=tail_cutoff)
+
+
+def bessel_kn_tail_bound(nu, x, *, tail_cutoff=50.0):
+    """Upper bound on K's omitted integral, excluding discretisation error."""
+    return _modified_tail_bound(nu, x, tail_cutoff=tail_cutoff)

syncem/constants.py

@@ -0,0 +1,23 @@
+"""Shared Gaussian-cgs constants adopted by the manuscript reference calculation.
+
+Electron mass is the CODATA2022 recommended value, converted from kg to g;
+elementary charge uses the manuscript's conventional Coulomb-to-statcoulomb
+conversion of the SI elementary charge. The charge denotes a positive
+magnitude (electron charge is -E_ESU). Values are centralised so emission and
+propagation cannot silently use different normalisations.
+
+Sources: https://physics.nist.gov/cuu/pdf/wall_2022.pdf and
+https://www.nist.gov/pml/special-publication-330/sp-330-section-2 .
+Mass uncertainty and any convention/conversion uncertainty are not randomised
+by this module; precision budgets requiring them should propagate them as
+shared physical-constant inputs rather than independent per-module errors.
+"""
+
+E_ESU = 4.803204712570263e-10  # charge magnitude [statC]
+M_E = 9.1093837139e-28  # electron mass [g]
+C_CGS = 2.99792458e10  # speed of light [cm/s]
+# SI speed of light [m/s]; only for the Faraday phase 2 (c/nu)^2 varphi with the
+# wavelength in metres (syncem.model.phase), never for the emission kernels.
+C_SI_M = 2.99792458e8
+
+__all__ = ["E_ESU", "M_E", "C_CGS", "C_SI_M"]

syncem/conversion.py

@@ -0,0 +1,50 @@
+"""Leading cold-plasma coefficients in CGS and the paper's Stokes convention.
+
+Inputs: frequency Hz, thermal density cm^-3, magnetic field Gauss. The
+``mueller_*`` functions return signed Stokes rates cm^-1 for dS/ds=eps-K S,
+Q=parallel-minus-perpendicular to projected B, and V=-2 Im(E_parallel E_perp*).
+In that natural basis conversion is rQ (U--V mixing); rU=0. At sky azimuth
+phi its components are rQ*cos(2phi), rQ*sin(2phi).
+
+``rotation_coefficient`` retains the position-angle rate for compatibility;
+``conversion_coefficient`` returns rQ itself. There is no additional factor
+of two on the latter. These leading high-frequency cold-dielectric terms
+require plasma and cyclotron frequencies small compared with nu. Hot and
+non-thermal corrections are distribution dependent and are not bounded here.
+"""
+
+from __future__ import annotations
+
+import numpy as np
+
+from .rm import E_ESU, M_E, C_CGS
+
+# rho_V = e^3 n_e B_par / (2 pi m_e^2 c^2 nu^2)
+C_ROT = E_ESU**3 / (2.0 * np.pi * M_E**2 * C_CGS**2)
+# rho_Q = e^4 n_e B_perp^2 / (4 pi^2 m_e^3 c^3 nu^3)
+C_CONV = E_ESU**4 / (4.0 * np.pi**2 * M_E**3 * C_CGS**3)
+
+
+def rotation_coefficient(nu, n_e, B_par):
+    """Physical Faraday rotation rate rho_V [rad/cm] (position angle)."""
+    return C_ROT * n_e * B_par / nu**2
+
+
+def conversion_coefficient(nu, n_e, B_perp):
+    """Natural-basis Stokes conversion rate rQ [cm^-1], signed."""
+    return -C_CONV * n_e * B_perp**2 / nu**3
+
+
+def conversion_rotation_ratio(nu, B_perp, B_par):
+    """Signed rQ/rV, independent of n_e; undefined when B_par=0."""
+    return -(E_ESU * B_perp**2) / (4.0 * np.pi * M_E * C_CGS * nu * B_par)
+
+
+def mueller_rotation(nu, n_e, B_par):
+    """Stokes-space rotation coefficient rV = 2 rho_V for syncem.transfer."""
+    return 2.0 * rotation_coefficient(nu, n_e, B_par)
+
+
+def mueller_conversion(nu, n_e, B_perp):
+    """Natural-basis Stokes conversion coefficient rQ for syncem.transfer."""
+    return conversion_coefficient(nu, n_e, B_perp)