# Source code for dynasor.post_processing.filon

"""
This module provides an implementation of Filon's integration formula.
For information about Filon's formula, see e.g.
Abramowitz and Stegun, Handbook of Mathematical Functions,
section 25 <http://mathworld.wolfram.com/FilonsIntegrationFormula.html>_ or
Allen and Tildesley, *Computer Simulation of Liquids*, Appendix D :cite:AllTil87.

Integration is performed along one dimension (default axis=0), e.g.,

.. code::

[F0[0]  F1[0] ..  FN[0] ]     [f0[0]  f1[0] ..  fN[0] ]
[   .      .         .  ]     [   .      .         .  ]
[F0[.]  F1[.] ..  FN[.] ] = I([f0[.]  f1[.] ..  fN[.] ], dx, [q[0] .. q[Nq]])
[   .      .         .  ]     [   .      .         .  ]
[F0[Nq] F1[Nq] .. FN[Nq]]     [f0[Nx] f1[Nx] .. fN[Nx]]

where q and Fj have end index Nq, and fj has end index Nx.
Nq is automatically set by the length of q. Due to the algorithm,
fj[Nx] must be of odd length (Nx must be an even number), and should
correspond to a linearly spaced set of data points (separated by dx along the
integration axis).

:func:sin_integral and :func:cos_integral allow for shifted integration
intervals by the optional argument x0.

"""

__all__ = ['fourier_cos', 'sin_integral', 'cos_integral']

from numpy import sin, cos, linspace, pi, arange
import numpy as np
import numba
from numpy.typing import NDArray
from typing import Tuple, Callable

[docs]def fourier_cos(f: NDArray[float],
dx: float,
q: NDArray[float] = None,
axis: int = 0) -> Tuple[NDArray[float], NDArray[float]]:
r"""Calculates the direct Fourier cosine transform :math:F(q) of a
function :math:f(x) using Filon's integration method.

The array values f[0]..f[2n] are expected to correspond to
:math:f(0.0)\ldots f(2n\Delta x). Hence, f should contain an odd
number of elements.

The transform is approximated by the integral
:math:F(q) = 2\int_{0}^{x_{max}} f(x) \cos(q x) dx,
where :math:x_{max} = 2n \Delta x.

Parameters
----------
f
function values; must contain an odd number of elements
dx
spacing of x-axis (:math:\Delta x)
q
values of reciprocal axis, at which to evaluate transform;
if q is not provided, linspace(0.0, 2*pi/dx, f.shape[axis]),
will be used.
axis
axis along which to carry out integration

Returns
-------
tuple of q and F values

Example
-------
A common use case is

.. code-block:: python

q, F = fourier_cos(f, dx)
"""

if q is None:
q = linspace(0.0, 2 * pi / dx, f.shape[axis])

return q, 2 * cos_integral(f, dx, q, x0=0.0, axis=axis)

def cos_integral(f: NDArray[float],
dx: float,
q: NDArray[float],
x0: float = 0.0,
axis: int = 0) -> NDArray[float]:
r"""Calculates the integral
:math:\int_{x_0}^{2n\Delta x} f(x) \cos(q x) dx.

Parameters
----------
f
function values; must contain an odd number of elements
dx
spacing of x-axis (:math:\Delta x)
q
values of reciprocal axis, at which to evaluate transform;
if q is not provided, linspace(0.0, 2*pi/dx, f.shape[axis]),
will be used.
x0
offset for integration interval
axis
axis along which to carry out integration

Returns
-------
Integral values
"""

return _gen_sc_int(f, dx, q, x0, axis, cos)

def sin_integral(f: NDArray[float],
dx: float,
q: NDArray[float],
x0: float = 0.0,
axis: int = 0) -> NDArray[float]:
r"""Calculates the integral
:math:\int_{x_0}^{2n\Delta x} f(x) \sin(q x) dx.

Parameters
----------
f
function values; must contain an odd number of elements
dx
spacing of x-axis (:math:\Delta x)
q
values of reciprocal axis, at which to evaluate transform;
if q is not provided, linspace(0.0, 2*pi/dx, f.shape[axis]),
will be used.
x0
offset for integration interval
axis
axis along which to carry out integration

Returns
-------
Integral values
"""
return _gen_sc_int(f, dx, q, x0, axis, sin)

def _gen_sc_int(f: NDArray[float],
dx: float,
q: NDArray[float],
x0: float,
axis: int,
sc: Callable) -> NDArray[float]:
""" General sin/cos integral

Due to numba sin or cos are sent in as strings
"""
if sc == np.sin:
sc = 'sin'
else:
assert sc == np.cos
sc = 'cos'

# Let numpy handle the axis juggling
integral = np.apply_along_axis(_gen_sc_int_1D, axis, f, dx, q, x0, sc)

return integral

@numba.njit(fastmath=False, parallel=True)
def _gen_sc_int_1D(f: NDArray, dx: float, k: NDArray, x0: float, sc: str) -> NDArray[float]:
"""Calculate filon for a 1D array

Parameters
----------
f
function values; must contain an odd number of elements
dx
spacing of x-axis
k
values of reciprocal axis, at which to evaluate transform
x0
offset for integration interval
sc
sin or cos

Returns
-------
integral
integral of f at points k
"""

# Keep it simple
assert f.ndim == 1
assert k.ndim == 1

assert sc in ['sin', 'cos']

Nk = len(k)
Nx = len(f)

assert Nx >= 3 and Nx % 2 == 1  # Nx odd and not one

x = x0 + dx * arange(Nx)

integral = np.empty(Nk)

for i in numba.prange(Nk):
integral[i] = filon_single_k(f, x, k[i], sc)

return integral

@numba.njit(fastmath=False)
def filon_single_k(f: NDArray, x: NDArray, k: float, sc: str) -> float:
"""performs the integral at a single k-point"""

dx = x[1] - x[0]

alpha, beta, gamma = _alpha_beta_gamma_single(dx * k)

if sc == 'sin':
sc_x = np.sin(k * x)
else:
sc_x = np.cos(k * x)

sc_x[0] *= 0.5
sc_x[-1] *= 0.5

Nx = len(x)
B, C = 0.0, 0.0
for i in range(0, Nx, 2):
B += sc_x[i] * f[i]
for i in range(1, Nx, 2):
C += sc_x[i] * f[i]

if sc == 'sin':
A = f[0] * np.cos(k * x[0]) - f[-1] * np.cos(k * x[-1])
else:
A = f[-1] * np.sin(k * x[-1]) - f[0] * np.sin(k * x[0])

integral = dx * (alpha * A + beta * B + gamma * C)

return integral

@numba.njit(fastmath=False)
def _alpha_beta_gamma_single(t: float):
# From theta, calculate alpha, beta, and gamma

if t == 0:
alpha, beta, gamma = 0.0, 2/3, 4/3
else:
alpha = (t**2 + t * np.sin(t) * np.cos(t) - 2 * np.sin(t)**2) / t**3
beta = 2 * (t * (1 + np.cos(t)**2) - 2 * np.sin(t) * np.cos(t)) / t**3
gamma = 4 * (np.sin(t) - t * np.cos(t)) / t**3

return alpha, beta, gamma