def calculate_distance(r1: Union[np.ndarray, float], r2: Union[np.ndarray, float], 
                      z1: Union[np.ndarray, float], z2: Union[np.ndarray, float]) -> Union[np.ndarray, float]:
    r"""Euclidean distance between $(r_1, z_1)$ and $(r_2, z_2)$ in the poloidal plane.

    $$d = \sqrt{(r_2 - r_1)^2 + (z_2 - z_1)^2}$$

    Parameters
    ----------
    r1 : float or np.ndarray
        Radius of the first point(s) [m].
    r2 : float or np.ndarray
        Radius of the second point(s) [m].
    z1 : float or np.ndarray
        Height of the first point(s) [m].
    z2 : float or np.ndarray
        Height of the second point(s) [m].

    Returns
    -------
    float or np.ndarray
        Distance, broadcast over the inputs [m].

    Notes
    -----
    A distance in the $(R, Z)$ half-plane, not the 3-D distance between points
    on two current loops.
    """
    return np.sqrt((r2 - r1) ** 2 + (z2 - z1) ** 2)

def complete_elliptic_integral_e(m: Union[float, np.ndarray]) -> Union[float, np.ndarray]:
    r"""Complete elliptic integral of the second kind $E(m)$.

    $$E(m) = \int_0^{\pi/2}\sqrt{1 - m\sin^2\theta}\,d\theta$$

    Parameters
    ----------
    m : float or np.ndarray
        Parameter $m = k^2$, in $[0, 1]$ [-].

    Returns
    -------
    float or np.ndarray
        $E(m)$, from $\pi/2$ at $m = 0$ to 1 at $m = 1$ [-].

    Convention
    ----------
    Takes the parameter $m = k^2$, as ``scipy.special.ellipe`` does.

    References
    ----------
    .. [1] NIST Digital Library of Mathematical Functions, Sec. 19.2(ii),
           https://dlmf.nist.gov/19.2.
    .. [2] M. Abramowitz and I. A. Stegun, *Handbook of Mathematical Functions*,
           Dover (1972), Sec. 17.3.
    """
    return ellipe(m)

def complete_elliptic_integral_k(m: Union[float, np.ndarray]) -> Union[float, np.ndarray]:
    r"""Complete elliptic integral of the first kind $K(m)$.

    $$K(m) = \int_0^{\pi/2}\frac{d\theta}{\sqrt{1 - m\sin^2\theta}}$$

    Parameters
    ----------
    m : float or np.ndarray
        Parameter $m = k^2$, in $[0, 1)$ [-].

    Returns
    -------
    float or np.ndarray
        $K(m)$ [-].

    Convention
    ----------
    Takes the *parameter* $m = k^2$, as ``scipy.special.ellipk`` does, not the
    modulus $k$; $K \to \infty$ logarithmically as $m \to 1$.

    References
    ----------
    .. [1] NIST Digital Library of Mathematical Functions, Sec. 19.2(ii),
           https://dlmf.nist.gov/19.2.
    .. [2] M. Abramowitz and I. A. Stegun, *Handbook of Mathematical Functions*,
           Dover (1972), Sec. 17.3.
    """
    return ellipk(m)

def green_br_bz_exact(r_obs, z_obs, r_src, z_src) -> Tuple[np.ndarray, np.ndarray]:
    r"""$(B_r, B_z)$ per unit ring current, exact elliptic integrals, broadcasting.

    $$B_r = -\frac{\mu_0}{2\pi r}\,\frac{\partial G}{\partial z}, \qquad
      B_z = +\frac{\mu_0}{2\pi r}\,\frac{\partial G}{\partial r}$$

    Parameters
    ----------
    r_obs : array-like
        Observation major radius [m].
    z_obs : array-like
        Observation height [m].
    r_src : array-like
        Source major radius [m].
    z_src : array-like
        Source height [m].

    Returns
    -------
    br : np.ndarray
        Radial field per ampere [T/A].
    bz : np.ndarray
        Vertical field per ampere [T/A].

    Convention
    ----------
    Right-handed $(r, \varphi, z)$, current in $+\varphi$; consistent with
    :func:`green_psi_exact` through $B_z = (1/r)\,\partial\psi/\partial r
    \times 1/2\pi$ for a full-weber $\psi$.

    Limitations
    -----------
    Observation points on the geometric axis ($r_{obs} = 0$) return the analytic
    limits $B_r = 0$, $B_z = \mu_0 r_0^2/(2(r_0^2 + (z - z_0)^2)^{3/2})$; the
    coincident-point caveat of :func:`greens_function_exact` applies.

    References
    ----------
    .. [1] W. R. Smythe, *Static and Dynamic Electricity*, 3rd ed., McGraw-Hill
           (1968), Sec. 7.10.
    """
    r_obs, z_obs, r_src, z_src = np.broadcast_arrays(
        np.asarray(r_obs, dtype=float),
        np.asarray(z_obs, dtype=float),
        np.asarray(r_src, dtype=float),
        np.asarray(z_src, dtype=float),
    )
    on_axis = r_obs == 0.0
    r_safe = np.where(on_axis, 1.0, r_obs)
    coeff = MU0 / (2.0 * np.pi * r_safe)
    br = -coeff * greens_function_exact(r_obs, z_obs, r_src, z_src, "dpsi_dz")
    bz = coeff * greens_function_exact(r_obs, z_obs, r_src, z_src, "dpsi_dr")
    if np.any(on_axis):
        dz2 = (z_obs - z_src) ** 2
        bz_axis = MU0 * r_src**2 / (2.0 * (r_src**2 + dz2) ** 1.5)
        bz = np.where(on_axis, bz_axis, bz)
        br = np.where(on_axis, 0.0, br)
    return br, bz

def green_psi_exact(r_obs, z_obs, r_src, z_src) -> np.ndarray:
    r"""Poloidal flux per unit ring current, exact elliptic integrals.

    $$\psi = \mu_0\,G(r, z; r_0, z_0)$$

    with $G$ from :func:`greens_function_exact`.

    Parameters
    ----------
    r_obs : array-like
        Observation major radius [m].
    z_obs : array-like
        Observation height [m].
    r_src : array-like
        Source major radius [m].
    z_src : array-like
        Source height [m].

    Returns
    -------
    np.ndarray
        Flux per ampere, broadcast over observation and source arrays [Wb/A].

    Convention
    ----------
    Full weber, the IMAS Data Dictionary storage (COCOS 11-18), same as
    :func:`green_r`.  The per-radian flux of an EFIT g-file or of
    :func:`vaft.formula.equilibrium.poloidal_field_factor` with ``cocos=None``
    is this divided by $2\pi$; mixing the two silently mis-scales fields by
    $2\pi$.  Tracked in #354.

    References
    ----------
    .. [1] J. D. Jackson, *Classical Electrodynamics*, 3rd ed., Wiley (1999),
           Sec. 5.5.
    """
    return MU0 * greens_function_exact(r_obs, z_obs, r_src, z_src, "psi")

def greens_function_2d(R: np.ndarray,
                      Z: np.ndarray,
                      R0: float,
                      Z0: float) -> np.ndarray:
    r"""Axisymmetric kernel $\sqrt{RR_0}\,K(m)$ between a field point and a ring source.

    $$G_{2D} = \sqrt{R R_0}\;K(m), \qquad
      m = \frac{4RR_0}{(R + R_0)^2 + (Z - Z_0)^2}$$

    Parameters
    ----------
    R : np.ndarray
        Major radius of the field points [m].
    Z : np.ndarray
        Height of the field points [m].
    R0 : float
        Major radius of the ring source [m].
    Z0 : float
        Height of the ring source [m].

    Returns
    -------
    np.ndarray
        Kernel value at each field point [m].

    Convention
    ----------
    This is the leading term of the ring-current flux function, *not* the flux
    Green's function itself: the poloidal flux per unit current is
    $\mu_0\sqrt{RR_0}\,[(2 - m)K(m) - 2E(m)]/\sqrt{m}$
    (:func:`greens_function_exact`, :func:`green_psi_exact`).  Use those for
    any physical field; this kernel is kept for the legacy integrals below.

    Limitations
    -----------
    Diverges logarithmically at the source point; no guard.

    References
    ----------
    .. [1] J. D. Jackson, *Classical Electrodynamics*, 3rd ed., Wiley (1999),
           Sec. 5.5 (vector potential of a circular loop).
    """
    k2 = 4 * R * R0 / ((R + R0)**2 + (Z - Z0)**2)
    return np.sqrt(R * R0) * complete_elliptic_integral_k(k2)

def greens_function_3d(R: np.ndarray,
                      Z: np.ndarray,
                      phi: np.ndarray,
                      R0: float,
                      Z0: float,
                      phi0: float) -> np.ndarray:
    r"""Toroidal-angle-resolved variant of the $\sqrt{RR_0}\,K(m)$ kernel.

    $$G_{3D} = \sqrt{R R_0}\;K(m), \qquad
      m = \frac{4RR_0}{(R + R_0)^2 + (Z - Z_0)^2 + 4RR_0\sin^2\!\big(\tfrac{\varphi - \varphi_0}{2}\big)}$$

    Parameters
    ----------
    R : np.ndarray
        Major radius of the field points [m].
    Z : np.ndarray
        Height of the field points [m].
    phi : np.ndarray
        Toroidal angle of the field points [rad].
    R0 : float
        Major radius of the source [m].
    Z0 : float
        Height of the source [m].
    phi0 : float
        Toroidal angle of the source [rad].

    Returns
    -------
    np.ndarray
        Kernel value at each field point [m].

    Limitations
    -----------
    Not the Biot-Savart kernel of a point source, and the difference is exact
    rather than vague. The chord between two points in a torus is

    $$d^2 = (R - R_0)^2 + (Z - Z_0)^2
            + 4RR_0\sin^2\!\big(\tfrac{\varphi - \varphi_0}{2}\big),$$

    while the ring kernel's denominator is
    $(R + R_0)^2 + (Z - Z_0)^2 = (R - R_0)^2 + (Z - Z_0)^2 + 4RR_0$. A point
    source therefore *replaces* the $4RR_0$ with
    $4RR_0\sin^2(\Delta\varphi/2)$; this expression **adds** the second term to
    the first, so the two coincide only at $\Delta\varphi = 0$, where the
    $\sin^2$ vanishes and the ring result is recovered.

    That $\Delta\varphi = 0$ reduction is the one property worth relying on.
    Away from it the expression decays with toroidal separation in a way that
    has no recorded derivation, so it should not be read as a field. Tracked in
    #356: it needs either a derivation or a replacement, and it has no
    production caller today.

    References
    ----------
    .. [1] J. D. Jackson, *Classical Electrodynamics*, 3rd ed., Wiley (1999),
           Sec. 5.5.
    """
    k2 = 4 * R * R0 / ((R + R0)**2 + (Z - Z0)**2 + 4 * R * R0 * np.sin((phi - phi0)/2)**2)
    return np.sqrt(R * R0) * complete_elliptic_integral_k(k2)

def greens_integral_2d(R: np.ndarray,
                      Z: np.ndarray,
                      R0: float,
                      Z0: float,
                      f: np.ndarray) -> float:
    r"""Line integral along $R$ of the 2-D kernel times a source density.

    $$I = \int G_{2D}(R, Z; R_0, Z_0)\,f(R)\,dR$$

    Parameters
    ----------
    R : np.ndarray
        Major radius of the samples, monotonic [m].
    Z : np.ndarray
        Height of the samples [m].
    R0 : float
        Major radius of the source [m].
    Z0 : float
        Height of the source [m].
    f : np.ndarray
        Source density at the samples [any].

    Returns
    -------
    float
        Integral in units of ``f`` times m^2 [any].

    Limitations
    -----------
    Integrates along the single supplied ``R`` array (a line, not an area) with
    the kernel of :func:`greens_function_2d`; not a flux or field.

    Numerical notes
    ---------------
    Trapezoidal rule.
    """
    G = greens_function_2d(R, Z, R0, Z0)
    return trapz_integral(R, G * f)

def greens_integral_3d(R: np.ndarray,
                      Z: np.ndarray,
                      phi: np.ndarray,
                      R0: float,
                      Z0: float,
                      phi0: float,
                      f: np.ndarray) -> float:
    r"""Line integral along $R$ of the 3-D kernel times a source density.

    $$I = \int G_{3D}(R, Z, \varphi; R_0, Z_0, \varphi_0)\,f(R)\,dR$$

    Parameters
    ----------
    R : np.ndarray
        Major radius of the samples, monotonic [m].
    Z : np.ndarray
        Height of the samples [m].
    phi : np.ndarray
        Toroidal angle of the samples [rad].
    R0 : float
        Major radius of the source [m].
    Z0 : float
        Height of the source [m].
    phi0 : float
        Toroidal angle of the source [rad].
    f : np.ndarray
        Source density at the samples [any].

    Returns
    -------
    float
        Integral in units of ``f`` times m^2 [any].

    Limitations
    -----------
    Same caveats as :func:`greens_integral_2d` and :func:`greens_function_3d`.

    Numerical notes
    ---------------
    Trapezoidal rule along ``R`` only.
    """
    G = greens_function_3d(R, Z, phi, R0, Z0, phi0)
    return trapz_integral(R, G * f)

def trapz_integral(x: np.ndarray, y: np.ndarray) -> float:
    r"""Trapezoidal integral $\int y\,dx$, local copy for the Green kernels.

    $$\int y\,dx \approx \sum_i \frac{y_i + y_{i+1}}{2}\,(x_{i+1} - x_i)$$

    Parameters
    ----------
    x : np.ndarray
        Sample abscissae, monotonic [any].
    y : np.ndarray
        Integrand at the samples [any].

    Returns
    -------
    float
        Integral over the sampled range, in units of ``y`` times ``x`` [any].

    Numerical notes
    ---------------
    Identical to :func:`vaft.formula.utils.trapz_integral` (``numpy.trapezoid``
    via :func:`vaft.compat.trapz_compat`); this copy is the one
    ``vaft.formula.trapz_integral`` resolves to, and it is kept because
    ``test_formula_lazy_namespace`` pins that shadowing. The reason it was
    written -- that importing ``utils`` dragged in scikit-learn -- is gone
    since #426.
    """
    return float(trapz_compat(y, x=x))