o X۷i@sdZddlmZddlZddlmZddlZddlmZddl m Z m Z z ddl m ZejZ[WneyCGdd d ZeZYnwd d gZed Zeej ddd Z eejddd ZdddZdddZdS)z{Fast Hankel transforms using the FFTLog algorithm. The implementation closely follows the Fortran code of Hamilton (2000). ) annotationsN)warn)_fft)loggammapoch)fhtc@seZdZddZdS) _DummyModulecCsdS)N)selfnamer r M/home/ubuntu/vllm_env/lib/python3.10/site-packages/cupyx/scipy/fft/_fftlog.py __getattr__sz_DummyModule.__getattr__N)__name__ __module__ __qualname__r r r r r rs rrifhtc Cs|jd}|dkr"|dd}t|}|t| |||}t|||||d}t||} |dkrD| t| ||||9} | S)aCompute the fast Hankel transform. Computes the discrete Hankel transform of a logarithmically spaced periodic sequence using the FFTLog algorithm [1]_, [2]_. Parameters ---------- a : cupy.ndarray (..., n) Real periodic input array, uniformly logarithmically spaced. For multidimensional input, the transform is performed over the last axis. dln : float Uniform logarithmic spacing of the input array. mu : float Order of the Hankel transform, any positive or negative real number. offset : float, optional Offset of the uniform logarithmic spacing of the output array. bias : float, optional Exponent of power law bias, any positive or negative real number. Returns ------- A : cupy.ndarray (..., n) The transformed output array, which is real, periodic, uniformly logarithmically spaced, and of the same shape as the input array. See Also -------- :func:`scipy.special.fht` :func:`scipy.special.fhtoffset` : Return an optimal offset for `fht`. References ---------- .. [1] Talman J. D., 1978, J. Comp. Phys., 29, 35 .. [2] Hamilton A. J. S., 2000, MNRAS, 312, 257 (astro-ph/9905191) rroffsetbiasshapecupyarangeexpfhtcoeff_fhtq) adlnmurrnj_cjuAr r r r"s (    c Cs|jd}|dkr#|dd}t|}|t|||||}t|||||d}t||dd} |dkrE| t| |||} | S)aPCompute the inverse fast Hankel transform. Computes the discrete inverse Hankel transform of a logarithmically spaced periodic sequence. This is the inverse operation to `fht`. Parameters ---------- A : cupy.ndarray (..., n) Real periodic input array, uniformly logarithmically spaced. For multidimensional input, the transform is performed over the last axis. dln : float Uniform logarithmic spacing of the input array. mu : float Order of the Hankel transform, any positive or negative real number. offset : float, optional Offset of the uniform logarithmic spacing of the output array. bias : float, optional Exponent of power law bias, any positive or negative real number. Returns ------- a : cupy.ndarray (..., n) The transformed output array, which is real, periodic, uniformly logarithmically spaced, and of the same shape as the input array. See Also -------- :func:`scipy.special.ifht` :func:`scipy.special.fhtoffset` : Return an optimal offset for `fht`. rrrrrT)inverser) r'r!r"rrr#r$r%r&r r r r ras #  c Cs`||}}|d|d}|d|d}tdtj|d|||dd} tj|ddtd} tj|ddtd} | | jdd<|| jdd<t| | d|| jdd<t| | d| dt |9} | j| j8_| jt |7_| j| j7_| j| 7_tj | | d|ddkrd| jd<t | dsd|t |||| d<| S)z?Compute the coefficient array for a fast Hankel transform. rrr)dtypeN)outr) rlinspacemathpiemptycompleximagrealrLN_2risfiniter) r#r!r"rrlnkrqxpxmyr&vr r r rs, (    rFcCs|jd}t|dr|std|}d|d<n|ddkr0|r0td|}tj|d<tj|dd}|s>||9}n||}tj ||dd}|ddddf}|S)zUCompute the biased fast Hankel transform. This is the basic FFTLog routine. rrz.singular transform; consider changing the biasz6singular inverse transform; consider changing the bias)axis.N) rrisinfrcopyinfrrfftconjirfft)r r&r(r#r'r r r rs     r)rr)F)__doc__ __future__rr,warningsrrcupyx.scipy.fftrcupyx.scipy.specialrr scipy.fftr_fht _scipy_fft ImportErrorr__all__logr2 _implementsrrrr r r r s.          >  9)