complex-bessel 0.5.0

//! Y Bessel function upper interface.
//!
//! Translation of Fortran ZBESY from TOMS 644 (zbsubs.f lines 1177-1461).
//! Y(fnu, z) = i*CC*I(fnu, arg) - (2/pi)*conj(CC)*K(fnu, arg)
//! where CC = exp(-i*pi*fnu/2), arg = z*exp(-i*pi/2).

use num_complex::Complex;

use crate::algo::constants::HPI;
use crate::besi::zbesi;
use crate::besk::zbesk;
use crate::machine::BesselFloat;
use crate::types::{Accuracy, Error, Scaling};
use crate::utils::{Workspace, mul_i, mul_neg_i, with_workspace, zabs};

/// The K term of scaled Y when 2|Im z| >= ELIM: `zv * exp(-tay)`.
///
/// Fortran ZBESY drops it there (EY = 0) as an underflow, but for large fnu
/// K(fnu, zn) is large enough that this term dominates Y (e.g. fnu = 575,
/// z = 360i, where Fortran returns only the tiny I term). The term is kept
/// when it lies above 1e3 times the smallest normal number (Amos's ARM); each
/// half exp(-tay/2) is then a normal number too, so applying them in turn
/// loses no precision. Below that it is dropped as in Fortran.
/// Deliberate deviation from TOMS 644.
#[inline]
fn far_k_term<T: BesselFloat>(zv: Complex<T>, tay: T) -> Complex<T> {
    let arm = T::from_f64(1.0e3) * T::MACH_TINY;
    let m = zv.re.abs().max(zv.im.abs());
    if m == T::zero() || m.ln() - tay <= arm.ln() {
        return Complex::new(T::zero(), T::zero());
    }
    let half = (-tay * T::from_f64(0.5)).exp();
    zv * half * half
}

/// Compute Y_{fnu+j}(z) for j = 0, 1, ..., n-1.
///
/// Results are written into the `y` slice. `n` is derived from `y.len()`.
///
/// Uses the identity Y(v,z) = i*CC*I(v,arg) - (2/pi)*conj(CC)*K(v,arg)
/// where CC = exp(-i*pi*v/2) and arg = z*exp(-i*pi/2).
///
/// Equivalent to Fortran ZBESY in TOMS 644.
#[inline]
pub(crate) fn zbesy<T: BesselFloat>(
    z: Complex<T>,
    fnu: T,
    scaling: Scaling,
    y: &mut [Complex<T>],
) -> Result<(usize, Accuracy), Error> {
    zbesy_with(z, fnu, scaling, y, Workspace::Stack)
}

/// [`zbesy`] with the workspace for the K values of a sequence (n > 1).
#[inline]
pub(crate) fn zbesy_with<T: BesselFloat>(
    z: Complex<T>,
    fnu: T,
    scaling: Scaling,
    y: &mut [Complex<T>],
    ws: Workspace<'_, T>,
) -> Result<(usize, Accuracy), Error> {
    let n = y.len();
    let zero = T::zero();
    let one = T::one();
    let czero = Complex::new(zero, zero);
    let hpi_t = T::from_f64(HPI);

    // Zero the output buffer
    y.fill(czero);

    // Input validation (Fortran lines 1342-1348)
    if n < 1 {
        return Err(Error::InvalidInput);
    }
    if fnu < zero {
        return Err(Error::InvalidInput);
    }
    if z == czero {
        return Err(Error::InvalidInput);
    }
    // Reject non-finite inputs: NaN passes every ordered comparison below,
    // then panics at to_i32() deep in the algorithm chain (e.g. zmlri).
    if !z.re.is_finite() || !z.im.is_finite() || !fnu.is_finite() {
        return Err(Error::InvalidInput);
    }

    // Range check, as in the other upper interfaces. The delegates zbesi/zbesk
    // repeat it, but fnu.to_i32() below must already be protected here.
    let az = zabs(z);
    let fn_val = fnu + T::from_f64((n - 1) as f64);
    let aa_tol = T::from_f64(0.5) / T::tol();
    let bb = T::from_f64(2147483647.0 * 0.5);
    let aa = aa_tol.min(bb);

    if az > aa || fn_val > aa {
        return Err(Error::TotalPrecisionLoss);
    }

    // Partial precision loss (Fortran IERR=3). The delegates zbesi/zbesk hit
    // the same condition (|zn| == |z|), but their statuses are not returned,
    // so compute it here like the other upper interfaces do.
    let aa_sqrt = aa.sqrt();
    let status = if az > aa_sqrt || fn_val > aa_sqrt {
        Accuracy::Reduced
    } else {
        Accuracy::Normal
    };

    // Rotate argument: zn = (zz.im, -zz.re) where zz = z with Im >= 0
    // (Fortran lines 1349-1353)
    let zn = Complex::new(z.im.abs(), -z.re);

    // Compute coefficients CC and CSPN (Fortran lines 1359-1373)
    // Safety: fnu is finite (checked above) and <= aa < i32::MAX per range check
    let ifnu = fnu.to_i32().unwrap();
    let ffnu = fnu - T::from_f64(ifnu as f64);
    let arg = hpi_t * ffnu;
    let mut csgn = Complex::new(arg.cos(), arg.sin());

    // Multiply by i^(ifnu mod 4) (Fortran lines 1364-1367)
    csgn = match (ifnu % 4) as usize {
        0 => csgn,
        1 => mul_i(csgn),
        2 => -csgn,
        _ => mul_neg_i(csgn),
    };

    let rhpi = one / hpi_t;
    let mut cspn = csgn.conj() * rhpi;
    // CSGN *= i (Fortran lines 1371-1373)
    csgn = mul_i(csgn);

    // For n == 1 (single values): one-element arrays, no workspace
    if n == 1 {
        let mut i_buf = [czero];
        let mut k_buf = [czero];

        // Compute I(fnu, zn) via ZBESI (Fortran line 1354)
        let (nz1, _) = zbesi(zn, fnu, scaling, &mut i_buf)?;
        // Compute K(fnu, zn) via ZBESK (Fortran line 1356)
        let (nz2, _) = zbesk(zn, fnu, scaling, &mut k_buf)?;
        let nz = nz1.min(nz2);

        if scaling == Scaling::Unscaled {
            // KODE=1: simple combination (Fortran DO 50, lines 1375-1394)
            let cy_val = csgn * i_buf[0] - cspn * k_buf[0];

            // Conjugate if original Im(z) < 0 (Fortran lines 1390-1394)
            y[0] = if z.im < zero { cy_val.conj() } else { cy_val };

            return Ok((nz, status));
        }

        // KODE=2: scaled version with underflow protection (Fortran lines 1396-1456)
        let tol = T::tol();
        let elim = T::elim();

        let exr = z.re.cos();
        let exi = z.re.sin();
        let tay = (z.im + z.im).abs();
        // Fortran sets EY = 0 when TAY >= ELIM, dropping the K term. Here CSPN
        // is left unscaled and `far_k_term` applies exp(-TAY) instead.
        let k_term_far = tay >= elim;
        let ey = if k_term_far { one } else { (-tay).exp() };

        // Apply exponential scaling to CSPN (Fortran lines 1411-1413)
        cspn = Complex::new(exr, exi) * cspn * ey;

        let rtol = one / tol;
        let ascle = T::MACH_TINY * rtol * T::from_f64(1.0e3);

        // Scale K values if near underflow (Fortran lines 1423-1433)
        let (zv_s, zv_atol) = if k_buf[0].re.abs().max(k_buf[0].im.abs()) <= ascle {
            (k_buf[0] * rtol, tol)
        } else {
            (k_buf[0], one)
        };
        let mut zv = zv_s * cspn * zv_atol;
        if k_term_far {
            zv = far_k_term(zv, tay);
        }

        // Scale I values if near underflow (Fortran lines 1434-1444)
        let (zu_s, zu_atol) = if i_buf[0].re.abs().max(i_buf[0].im.abs()) <= ascle {
            (i_buf[0] * rtol, tol)
        } else {
            (i_buf[0], one)
        };
        let zu = zu_s * csgn * zu_atol;

        // CY(I) = ZU - ZV (Fortran lines 1445-1449)
        let cy_val = zu - zv;

        y[0] = if z.im < zero { cy_val.conj() } else { cy_val };

        let nz_out = if cy_val == czero && k_term_far { 1 } else { 0 };

        return Ok((nz_out, status));
    }

    zbesy_seq(z, zn, fnu, scaling, y, csgn, cspn, status, ws)
}

/// The n > 1 part of [`zbesy`], kept out of line so that a stack workspace
/// stays out of the frame of the single-value path.
#[inline(never)]
#[allow(clippy::too_many_arguments)]
fn zbesy_seq<T: BesselFloat>(
    z: Complex<T>,
    zn: Complex<T>,
    fnu: T,
    scaling: Scaling,
    y: &mut [Complex<T>],
    csgn: Complex<T>,
    cspn: Complex<T>,
    status: Accuracy,
    ws: Workspace<'_, T>,
) -> Result<(usize, Accuracy), Error> {
    let n = y.len();
    with_workspace(ws, n, |k_buf| {
        zbesy_combine(z, zn, fnu, scaling, y, csgn, cspn, status, k_buf)
    })
}

/// I into `y` in one call, K into the workspace `k_buf` (Fortran ZBESY's
/// CWRK), in chunks of `k_buf.len()` orders if it is shorter than `y`, and
/// the two combined into Y.
#[inline]
#[allow(clippy::too_many_arguments)]
fn zbesy_combine<T: BesselFloat>(
    z: Complex<T>,
    zn: Complex<T>,
    fnu: T,
    scaling: Scaling,
    y: &mut [Complex<T>],
    mut csgn: Complex<T>,
    mut cspn: Complex<T>,
    status: Accuracy,
    k_buf: &mut [Complex<T>],
) -> Result<(usize, Accuracy), Error> {
    let chunk = k_buf.len();
    assert!(chunk > 0, "empty K workspace");
    let zero = T::zero();
    let one = T::one();
    let czero = Complex::new(zero, zero);

    // Compute I(fnu, zn) via ZBESI (Fortran line 1354)
    let (nz1, _) = zbesi(zn, fnu, scaling, y)?;

    // KODE=2 setup (Fortran lines 1396-1413)
    let tol = T::tol();
    let rtol = one / tol;
    let ascle = T::MACH_TINY * rtol * T::from_f64(1.0e3);
    let tay = (z.im + z.im).abs();
    // Fortran sets EY = 0 when TAY >= ELIM, dropping the K term. Here CSPN
    // is left unscaled and `far_k_term` applies exp(-TAY) instead.
    let k_term_far = tay >= T::elim();
    if scaling == Scaling::Exponential {
        let ey = if k_term_far { one } else { (-tay).exp() };
        cspn = Complex::new(z.re.cos(), z.re.sin()) * cspn * ey;
    }

    let mut nz2: usize = 0;
    let mut nz_out: usize = 0;
    for (c, y_chunk) in y.chunks_mut(chunk).enumerate() {
        let k_chunk = &mut k_buf[..y_chunk.len()];
        // Compute K(fnu + c*chunk, zn) via ZBESK (Fortran line 1356)
        let start = if c == 0 {
            fnu
        } else {
            fnu + T::from_f64((c * chunk) as f64)
        };
        let (nz_k, _) = zbesk(zn, start, scaling, k_chunk)?;
        nz2 += nz_k;

        for (yi, &k) in y_chunk.iter_mut().zip(k_chunk.iter()) {
            let i = *yi;
            if scaling == Scaling::Unscaled {
                // KODE=1: CY(I) = CSGN*I(I) - CSPN*K(I) (Fortran DO 50, lines 1375-1394)
                *yi = csgn * i - cspn * k;
            } else {
                // KODE=2: scaled version with underflow protection
                // Scale K values if near underflow (Fortran lines 1423-1433)
                let (zv_s, zv_atol) = if k.re.abs().max(k.im.abs()) <= ascle {
                    (k * rtol, tol)
                } else {
                    (k, one)
                };
                let mut zv = zv_s * cspn * zv_atol;
                if k_term_far {
                    zv = far_k_term(zv, tay);
                }

                // Scale I values if near underflow (Fortran lines 1434-1444)
                let (zu_s, zu_atol) = if i.re.abs().max(i.im.abs()) <= ascle {
                    (i * rtol, tol)
                } else {
                    (i, one)
                };
                let zu = zu_s * csgn * zu_atol;

                // CY(I) = ZU - ZV (Fortran lines 1445-1449)
                let cy_val = zu - zv;
                *yi = if z.im < zero { cy_val.conj() } else { cy_val };
                if cy_val == czero && k_term_far {
                    nz_out += 1;
                }
            }

            // Advance CSGN *= i, CSPN *= -i (Fortran lines 1383-1388, 1450-1455)
            csgn = mul_i(csgn);
            cspn = mul_neg_i(cspn);
        }
    }

    if scaling == Scaling::Exponential {
        return Ok((nz_out, status));
    }

    // Conjugate if original Im(z) < 0 (Fortran lines 1390-1394)
    if z.im < zero {
        for yi in y.iter_mut() {
            *yi = yi.conj();
        }
    }
    Ok((nz1.min(nz2), status))
}

#[cfg(test)]
mod tests {
    use super::*;
    use num_complex::Complex64;

    #[test]
    fn besy_z_zero_returns_error() {
        let z = Complex64::new(0.0, 0.0);
        let mut y = [Complex64::new(0.0, 0.0)];
        assert!(zbesy(z, 0.0, Scaling::Unscaled, &mut y).is_err());
    }

    #[test]
    fn besy_non_finite_inputs_return_invalid_input() {
        let nan = f64::NAN;
        let inf = f64::INFINITY;
        let mut y = [Complex64::new(0.0, 0.0)];
        for z in [
            Complex64::new(nan, 0.0),
            Complex64::new(0.0, nan),
            Complex64::new(nan, nan),
            Complex64::new(inf, 0.0),
        ] {
            assert!(matches!(
                zbesy(z, 1.0, Scaling::Unscaled, &mut y),
                Err(Error::InvalidInput)
            ));
        }
        let z = Complex64::new(1.0, 1.0);
        for fnu in [nan, inf] {
            assert!(matches!(
                zbesy(z, fnu, Scaling::Unscaled, &mut y),
                Err(Error::InvalidInput)
            ));
        }
    }

    #[test]
    fn besy_reports_reduced_accuracy_for_large_argument() {
        // |z| > sqrt(aa) ≈ 32767 (f64) means more than half of the significant
        // digits may be lost (Fortran IERR=3). zbesy used to discard the
        // delegate statuses and always report Normal.
        let mut y = [Complex64::new(0.0, 0.0)];
        let (_, status) =
            zbesy(Complex64::new(40000.0, 0.0), 0.0, Scaling::Unscaled, &mut y).unwrap();
        assert_eq!(status, Accuracy::Reduced);

        let (_, status) = zbesy(Complex64::new(1.0, 1.0), 0.0, Scaling::Unscaled, &mut y).unwrap();
        assert_eq!(status, Accuracy::Normal);
    }

    #[test]
    fn besy_large_finite_order_returns_total_precision_loss() {
        // fnu >= 2^31 used to panic at fnu.to_i32().unwrap() because zbesy
        // had no range check of its own before delegating to zbesi/zbesk.
        let z = Complex64::new(1.0, 1.0);
        let mut y = [Complex64::new(0.0, 0.0)];
        assert!(matches!(
            zbesy(z, 3.0e9, Scaling::Unscaled, &mut y),
            Err(Error::TotalPrecisionLoss)
        ));
    }
}