<reponame>J-Massey/postproc
# -*- coding: utf-8 -*-
"""
@author: <NAME>
@description: Module to compute the wavenumber spectra of a n-dimensional field.
@contact: <EMAIL>
"""

# Imports
import numpy as np
import scipy.signal as signal
from tqdm import tqdm

# Functions
def ke_spectra(u_i, x_i, **kwargs):
	"""
	Computes the kinetic energy (KE) of a n-dimensional velocity vector field such as u_i = (u, v, w) for 3D.
	:param u_i: n-dimensional velocity vector field
	:param x_i: n-dimensional spatial vector field
	:param kwargs: k_res for the resolution of k_mod_line which defines the bandwitdh dk.
	:return: k_mod 1D array and ke_integral 1D array.
	"""
	if len(u_i) > 3 or len(u_i) < 1 or len(x_i) > 3 or len(x_i) < 1:
		raise ValueError('Invalid field dimensions')
	# Wavenumbers
	k_i = _wavenumbers(*x_i) # k_i = (kx, ky, kz)
	# FFT to compute KE
	ke = 0
	for u in u_i:
		u = _window_ndim(u, signal.hanning) # Windowing
		uk = np.fft.fftn(u)/u.size # FFT
		ke += uk*uk.conjugate() # KE
	ke = 0.5*ke
	# Calc spectra
	return _pair_integrate_fast(ke, *k_i, **kwargs)


def scalar_spectra(a, *args, **kwargs):
	"""
	Computes the wavenumber spectra of a n-dimensional scalar field by means of nFFT. The spectra is for the wavenumber
	modulus of ak. Hence it performs and spherical integration for all the components of k_i (the wavenumber vector)
	:param a: The scalar field
	:param args: The spatial vector x_i = (x, Y, z)
	:param kwargs: k_res for the resolution of k_mod_line which defines the bandwitdh dk.
	:return: k_mod 1D array and ak_integral 1D array.
	"""
	if a.ndim > 3 or a.ndim < 1:
		raise ValueError('Invalid field dimensions')
	elif a.ndim != len(args):
		raise ValueError('Field dimensions must much the spatial vector components passed to the function')

	k_i = _wavenumbers(*args) # k_i = (kx, ky, kz)
	a = _window_ndim(a, signal.hanning) # Windowing
	ak = np.fft.fftn(a)/a.size # FFT add power spectrum??
	return _pair_integrate_fast(ak, *k_i, **kwargs) # Calc spectra


def _pair_integrate_fast(ak, *args, **kwargs):
	"""
	Internal function which computes the wavenumber modulus k_mod for each fft coefficient of the ak ndarray and
	integrates the components contributing to the same k_mod with a certain bandwidth dk
	:param ak: The nFFT of (a)
	:param args: the wavenumber vector *k_i
	:param kwargs: k_res for the resolution of k_mod_line which defines the bandwitdh dk.
	:return: k_mod 1D array and ak_integral 1D array.
	"""
	k_res = kwargs.get('k_res', 200)
	k2_sum_max = 0
	k2_sum_min = 0
	for k in args:
		k2_sum_max += np.max(k**2)
		k2_sum_min += np.min(k**2)
	k_min, k_max = np.sqrt(k2_sum_min), np.sqrt(k2_sum_max)
	dk = (k_max - k_min) / k_res

	ak_integral = np.zeros(k_res)
	k_mod_line = np.linspace(0, k_res - 1, k_res) * dk + dk / 2  # k values at half of each bandwidth

	print('	Find ak(k_i) with its k modulus and integrate it in ak(k)')
	with tqdm(total=ak.size) as pbar:
		for index in np.ndindex(ak.shape):
			ak_p = ak[index]
			k2_sum = 0
			for i, k in enumerate(args):
				k2_sum += k[index[i]] ** 2
			k_mod = np.sqrt(k2_sum)

			kint = int(k_mod / dk)
			if kint >= k_res:
				ak_integral[-1] += np.abs(ak_p)
			else:
				ak_integral[kint] += np.abs(ak_p)
			pbar.update(1)

	return k_mod_line, ak_integral

def _wavenumbers(*args):
	"""
	:param args: the wavenumber vector *k_i
	:return: the wavenumber vector for a position vector (1D, 2D or 3D) such as: x, Y, z of a uniform grid.
	"""
	k_i = () # wavenumber vector; k_i = (kx, ky, kz). Each component is a 1D array type.
	for arg in args:
		N = arg.size # number of points in the spatial vector component
		alpha = 2*np.pi/(np.max(arg)-np.min(arg)) # basic wavenumber
		index = np.fft.fftfreq(N, d=1/N) # index per wavenumber direction: eg x: 0,1,2,...,N/2-1,-N/2,-N/2+1,...,-1

		k = np.zeros(N) # wavenumber vector component k_i
		for i in range(0, N):
			k[i] = alpha*index[i]
		k_i = k_i + (k,)
	return k_i


def _window_ndim(a, wfunction):
	"""
	Performs an in-place windowing on N-dimensional spatial-domain dat.
	This is done to mitigate boundary effects in the FFT.
	:param a: n-dimensional array input dat to be windowed, modified in place.
	:param wfunction: 1D window generation function. Function should accept one argument: the window length.
		   Example: scipy.signal.hamming
	:return: windowed n-dimensional array a
	"""
	if a.ndim == 0:
		raise ValueError('Input dat to be windowed cannot be scalar')
	for axis, axis_size in enumerate(a.shape):
		window = wfunction(axis_size)
		for i in range(len(a.shape)):
			if i == axis:
				continue
			else:
				window = np.stack([window] * a.shape[i], axis=i)
		a *= window
	return a
# -*- coding: utf-8 -*-
"""
@author: <NAME>
@description: Module to compute the wavenumber spectra of a n-dimensional field.
@contact: <EMAIL>
"""

# Imports
import numpy as np
import scipy.signal as signal
from tqdm import tqdm

# Functions
def ke_spectra(u_i, x_i, **kwargs):
	"""
	Computes the kinetic energy (KE) of a n-dimensional velocity vector field such as u_i = (u, v, w) for 3D.
	:param u_i: n-dimensional velocity vector field
	:param x_i: n-dimensional spatial vector field
	:param kwargs: k_res for the resolution of k_mod_line which defines the bandwitdh dk.
	:return: k_mod 1D array and ke_integral 1D array.
	"""
	if len(u_i) > 3 or len(u_i) < 1 or len(x_i) > 3 or len(x_i) < 1:
		raise ValueError('Invalid field dimensions')
	# Wavenumbers
	k_i = _wavenumbers(*x_i) # k_i = (kx, ky, kz)
	# FFT to compute KE
	ke = 0
	for u in u_i:
		u = _window_ndim(u, signal.hanning) # Windowing
		uk = np.fft.fftn(u)/u.size # FFT
		ke += uk*uk.conjugate() # KE
	ke = 0.5*ke
	# Calc spectra
	return _pair_integrate_fast(ke, *k_i, **kwargs)


def scalar_spectra(a, *args, **kwargs):
	"""
	Computes the wavenumber spectra of a n-dimensional scalar field by means of nFFT. The spectra is for the wavenumber
	modulus of ak. Hence it performs and spherical integration for all the components of k_i (the wavenumber vector)
	:param a: The scalar field
	:param args: The spatial vector x_i = (x, Y, z)
	:param kwargs: k_res for the resolution of k_mod_line which defines the bandwitdh dk.
	:return: k_mod 1D array and ak_integral 1D array.
	"""
	if a.ndim > 3 or a.ndim < 1:
		raise ValueError('Invalid field dimensions')
	elif a.ndim != len(args):
		raise ValueError('Field dimensions must much the spatial vector components passed to the function')

	k_i = _wavenumbers(*args) # k_i = (kx, ky, kz)
	a = _window_ndim(a, signal.hanning) # Windowing
	ak = np.fft.fftn(a)/a.size # FFT add power spectrum??
	return _pair_integrate_fast(ak, *k_i, **kwargs) # Calc spectra


def _pair_integrate_fast(ak, *args, **kwargs):
	"""
	Internal function which computes the wavenumber modulus k_mod for each fft coefficient of the ak ndarray and
	integrates the components contributing to the same k_mod with a certain bandwidth dk
	:param ak: The nFFT of (a)
	:param args: the wavenumber vector *k_i
	:param kwargs: k_res for the resolution of k_mod_line which defines the bandwitdh dk.
	:return: k_mod 1D array and ak_integral 1D array.
	"""
	k_res = kwargs.get('k_res', 200)
	k2_sum_max = 0
	k2_sum_min = 0
	for k in args:
		k2_sum_max += np.max(k**2)
		k2_sum_min += np.min(k**2)
	k_min, k_max = np.sqrt(k2_sum_min), np.sqrt(k2_sum_max)
	dk = (k_max - k_min) / k_res

	ak_integral = np.zeros(k_res)
	k_mod_line = np.linspace(0, k_res - 1, k_res) * dk + dk / 2  # k values at half of each bandwidth

	print('	Find ak(k_i) with its k modulus and integrate it in ak(k)')
	with tqdm(total=ak.size) as pbar:
		for index in np.ndindex(ak.shape):
			ak_p = ak[index]
			k2_sum = 0
			for i, k in enumerate(args):
				k2_sum += k[index[i]] ** 2
			k_mod = np.sqrt(k2_sum)

			kint = int(k_mod / dk)
			if kint >= k_res:
				ak_integral[-1] += np.abs(ak_p)
			else:
				ak_integral[kint] += np.abs(ak_p)
			pbar.update(1)

	return k_mod_line, ak_integral

def _wavenumbers(*args):
	"""
	:param args: the wavenumber vector *k_i
	:return: the wavenumber vector for a position vector (1D, 2D or 3D) such as: x, Y, z of a uniform grid.
	"""
	k_i = () # wavenumber vector; k_i = (kx, ky, kz). Each component is a 1D array type.
	for arg in args:
		N = arg.size # number of points in the spatial vector component
		alpha = 2*np.pi/(np.max(arg)-np.min(arg)) # basic wavenumber
		index = np.fft.fftfreq(N, d=1/N) # index per wavenumber direction: eg x: 0,1,2,...,N/2-1,-N/2,-N/2+1,...,-1

		k = np.zeros(N) # wavenumber vector component k_i
		for i in range(0, N):
			k[i] = alpha*index[i]
		k_i = k_i + (k,)
	return k_i


def _window_ndim(a, wfunction):
	"""
	Performs an in-place windowing on N-dimensional spatial-domain dat.
	This is done to mitigate boundary effects in the FFT.
	:param a: n-dimensional array input dat to be windowed, modified in place.
	:param wfunction: 1D window generation function. Function should accept one argument: the window length.
		   Example: scipy.signal.hamming
	:return: windowed n-dimensional array a
	"""
	if a.ndim == 0:
		raise ValueError('Input dat to be windowed cannot be scalar')
	for axis, axis_size in enumerate(a.shape):
		window = wfunction(axis_size)
		for i in range(len(a.shape)):
			if i == axis:
				continue
			else:
				window = np.stack([window] * a.shape[i], axis=i)
		a *= window
	return a
# -*- coding: utf-8 -*- @author: <NAME> @description: Module to compute the wavenumber spectra of a n-dimensional field. @contact: <EMAIL> # Imports # Functions Computes the kinetic energy (KE) of a n-dimensional velocity vector field such as u_i = (u, v, w) for 3D. 	:param u_i: n-dimensional velocity vector field 	:param x_i: n-dimensional spatial vector field 	:param kwargs: k_res for the resolution of k_mod_line which defines the bandwitdh dk. 	:return: k_mod 1D array and ke_integral 1D array. # Wavenumbers # k_i = (kx, ky, kz) # FFT to compute KE # Windowing # FFT # KE # Calc spectra Computes the wavenumber spectra of a n-dimensional scalar field by means of nFFT. The spectra is for the wavenumber 	modulus of ak. Hence it performs and spherical integration for all the components of k_i (the wavenumber vector) 	:param a: The scalar field 	:param args: The spatial vector x_i = (x, Y, z) 	:param kwargs: k_res for the resolution of k_mod_line which defines the bandwitdh dk. 	:return: k_mod 1D array and ak_integral 1D array. # k_i = (kx, ky, kz) # Windowing # FFT add power spectrum?? # Calc spectra Internal function which computes the wavenumber modulus k_mod for each fft coefficient of the ak ndarray and 	integrates the components contributing to the same k_mod with a certain bandwidth dk 	:param ak: The nFFT of (a) 	:param args: the wavenumber vector *k_i 	:param kwargs: k_res for the resolution of k_mod_line which defines the bandwitdh dk. 	:return: k_mod 1D array and ak_integral 1D array. # k values at half of each bandwidth :param args: the wavenumber vector *k_i 	:return: the wavenumber vector for a position vector (1D, 2D or 3D) such as: x, Y, z of a uniform grid. # wavenumber vector; k_i = (kx, ky, kz). Each component is a 1D array type. # number of points in the spatial vector component # basic wavenumber # index per wavenumber direction: eg x: 0,1,2,...,N/2-1,-N/2,-N/2+1,...,-1 # wavenumber vector component k_i Performs an in-place windowing on N-dimensional spatial-domain dat. 	This is done to mitigate boundary effects in the FFT. 	:param a: n-dimensional array input dat to be windowed, modified in place. 	:param wfunction: 1D window generation function. Function should accept one argument: the window length. 		   Example: scipy.signal.hamming 	:return: windowed n-dimensional array a