FFTLog in Python, using only numpy and scipy.special.loggamma

fftlognt.py

import numpy as np
from scipy.fft import rfft, irfft
from scipy.special import loggamma


# constants
LN_2 = np.log(2)
LN_10 = np.log(10)


def lnkrgood(lnkr, dlnr, mu, q=0.0):

    # find low-ringing kr
    xp = (mu+1+q)/2
    xm = (mu+1-q)/2
    y = np.pi/(2*dlnr)
    zp = loggamma(xp + 1j*y)
    zm = loggamma(xm + 1j*y)
    arg = (LN_2 - lnkr)/dlnr + (zp.imag + zm.imag)/np.pi
    return lnkr + (arg - np.round(arg))*dlnr


def fhtq(f, dlnr, lnkr, mu, q=0.0):
    '''compute a discrete version of the biased Hankel transform

    Parameters
    ----------
    f : array_like (..., N)
        Function values. Can be multidimensional, the transform is performed
        over the last axis.
    dlnr : float
        Log-spacing of abscissae.
    lnkr : float
        Relation between k and r arrays, `k = exp(lnkr)/r[::-1]`.
    mu : float
        Order of the transform.
    q : float, optional
        Exponent of power law bias.
    '''

    assert np.ndim(f) >= 1, 'f must be array'
    assert np.isscalar(dlnr), 'dlnr must be scalar'
    assert np.isscalar(lnkr), 'kr must be scalar'
    assert np.isscalar(mu), 'mu must be scalar'

    # size of transform
    n = np.shape(f)[-1]

    # compute Hankel transform coefficients
    a = (mu+1+q)/2
    b = (mu+1-q)/2
    c = np.arange(0, n//2+1, dtype=float)
    c *= np.pi/(n*dlnr)
    u = np.empty(len(c), dtype=complex)
    v = np.empty(len(c), dtype=complex)
    u.real[:] = a
    v.real[:] = b
    u.imag[:] = c
    v.imag[:] = c
    loggamma(u, out=u)
    loggamma(v, out=v)
    c *= 2*(lnkr - LN_2)
    u.real -= v.real
    u.real += LN_2*q
    u.imag += v.imag
    u.imag -= c
    np.exp(u, out=u)
    if np.isnan(u[0]):
        if a > 0:
            u[0] = 0
        else:
            raise ValueError(f'singular transform')
    del(v)

    # Hankel transform via real FFT
    g = rfft(f, axis=-1)
    g *= u
    g = irfft(g, n, axis=-1)
    g[..., :] = g[..., ::-1]

    return g