""" Methods to transform the waveform """
import numpy as np
# from numba import jit, njit
# @njit(parallel=True)
[docs]def compute_fft(udata_x, delta_x):
"""Find the FFT of the samples in time-space, and return with the frequencies.
Parameters
----------
udata_x: 1d array
The samples in time-space.
delta_x: float
The stepping delta_x
Returns
-------
freqs: 1d array
The frequency axis, shifted approriately.
utilde: 1d array
The samples in frequency space, with conventions applied.
"""
# import necessary libraries.
from numpy.fft import fft
# FFT
utilde_orig = fft(udata_x)
# Apply conventions.
utilde = set_fft_conven(utilde_orig)
# Get frequency axes.
Nlen = len(utilde)
# print(Nlen)
# Naxis = np.arange(Nlen)
# freq_orig = fftfreq(Nlen)
# freq_axis = fftshift(freq_orig)*Nlen
# delta_x = xdata[1] - xdata[0]
# Naxis = np.arange(Nlen)
freq_axis = np.linspace(-0.5 / delta_x, 0.5 / delta_x, Nlen)
return freq_axis, utilde
# @njit(parallel=True)
[docs]def compute_ifft(utilde, delta_f):
"""Find the inverse FFT of the samples in frequency-space, and return with the time axis.
Parameters
----------
utilde : 1d array
The samples in frequency-space.
delta_f: float
The frequency stepping
Returns
-------
time_axis: 1d array
The time axis.
udata_time: 1d array
The samples in time domain.
"""
# import necessary libraries.
from numpy.fft import ifft
# FFT
utilde_orig = unset_fft_conven(utilde)
# Inverse transform
udata_time = ifft(utilde_orig)
# Get frequency axes.
Nlen = len(udata_time)
# print(Nlen)
# Naxis = np.arange(Nlen)
# freq_orig = fftfreq(Nlen)
# freq_axis = fftshift(freq_orig)*Nlen
# delta_x = xdata[1] - xdata[0]
# Naxis = np.arange(Nlen)
delta_t = 1.0 / (delta_f * Nlen)
# Dt = Nlen * delta_f/2
time_axis = np.linspace(0, delta_t * Nlen, Nlen)
return time_axis, udata_time
# @njit(parallel=True)
[docs]def set_fft_conven(utilde_orig):
"""Make a numppy fft consistent with the chosen conventions.
This takes care of the zero mode factor and array position.
Also, it shifts the negative frequencies using numpy's fftshift.
Parameters
----------
utilde_orig: 1d array
The result of a numpy fft.
Returns
-------
utilde_conven: 1d array
The fft with set conventions.
"""
# Multiply by 2, take conjugate.
utilde_conven = 2 * np.conj(utilde_orig) / len(utilde_orig)
# Restore the zero mode.
utilde_conven[0] = utilde_conven[0] / 2
# Shift the frequency axis.
utilde_conven = np.fft.fftshift(utilde_conven)
return utilde_conven
# @njit(parallel=True)
[docs]def unset_fft_conven(utilde_conven):
"""Make an actual conventional fft consistent with numpy's conventions.
The inverse of set_conv.
Parameters
----------
utilde_conven: 1d array
The conventional fft data vector.
Returns
-------
utilde_np
"""
utilde_np = np.fft.ifftshift(utilde_conven)
utilde_np = len(utilde_np) * np.conj(utilde_np) / 2
# print(utilde_original[0])
utilde_np[0] *= 2
# print(utilde_original[0])
return utilde_np
[docs]def Yslm(spin_weight, ell, emm, theta, phi):
"""Spin-weighted spherical harmonics data defined as a function of zeta and phi, for qlm data decomposition.
Inputs
-----------
spin_weight : int
The Spin weight.
ell : int
The mode number :math:`\\ell'.
emm : int
The azimuthal mode number :math:`m'.
theta : float
The polar angle :math:`\\theta` in radians,
phi : float
The aximuthal angle :math:`\\phi' in radians.
Returns
--------
Yslm : float
The value of Yslm at :math:`\\theta, phi'.
"""
import sympy as sp
# theta, phi = sp.symbols('theta phi')
fact = np.math.factorial
#fact = sp.factorial
Sum = 0
for aar in range(ell - spin_weight + 1):
if (aar + spin_weight - emm) < 0 or (ell - aar - spin_weight) < 0:
# print('Continuing')
continue
else:
# print('r, l, s, m', r, l, s, m)
# a1 = sp.binomial(ell - spin_weight, aar)
# print(a1)
# a2 = sp.binomial(ell + spin_weight, aar + spin_weight - emm)
# print(a2)
# a3 = np.exp(1j * emm * phi)
# print(a3)
# a4 = np.tan(theta / 2)
# print(a4)
Sum += (
sp.binomial(ell - spin_weight, aar)
* sp.binomial(ell + spin_weight, aar + spin_weight - emm)
* np.power((-1), (ell - aar - spin_weight))
* np.exp(1j * emm * phi)
/ np.power(np.tan(theta / 2), (2 * aar + spin_weight - emm))
)
Sum = complex(Sum)
# print(type(m))
# print((-1)**int(m))
# print(np.sin(th/2)**(2*l))
Yslm = (-1) ** emm * (
np.sqrt(
fact(ell + emm)
* fact(ell - emm)
* (2 * ell + 1)
/ (4 * np.pi * fact(ell + spin_weight) * fact(ell - spin_weight))
)
* np.sin(theta / 2) ** (2 * ell)
* Sum
)
return Yslm
[docs]def Yslm_vec(spin_weight, ell, emm, theta_grid, phi_grid):
"""Spin-weighted spherical harmonics data defined as a function of zeta and phi, for qlm data decomposition.
Inputs
-----------
spin_weight : int
The Spin weight.
ell : int
The mode number :math:`\\ell'.
emm : int
The azimuthal mode number :math:`m'.
theta : float
The polar angle :math:`\\theta` in radians,
phi : float
The aximuthal angle :math:`\\phi' in radians.
Returns
--------
Yslm : float
The value of Yslm at :math:`\\theta, phi'.
"""
# spin_weight = abs(spin_weight)
# theta, phi = sp.symbols('theta phi')
from math import comb
# import sympy as sp
fact = np.math.factorial
theta_grid = np.array(theta_grid)
phi_grid = np.array(phi_grid)
Sum = 0 + 1j * 0
for aar in range(0, ell - spin_weight + 1):
# print('aar', aar)
subterm = 0
if (aar + spin_weight - emm) < 0 or (ell - aar - spin_weight) < 0:
# print('Continuing')
continue
else:
# print(aar + spin_weight-emm)
term1 = comb(ell - spin_weight, aar)
term2 = comb(ell + spin_weight, aar + spin_weight - emm)
term3 = np.power(float(-1), (ell - aar - spin_weight))
# term3 = (-1)**(ell-aar-spin_weight)
term4 = np.exp(1j * emm * phi_grid)
term5 = np.power(np.tan(theta_grid / 2), (-2 * aar - spin_weight + emm))
subterm = term1 * term2 * term3 * term4 * term5
# print(term1, term2, term3, term4)
Sum += subterm
# print('arr, subterm', aar, subterm)
# print(ell+emm, ell+spin_weight, ell-spin_weight)
Yslmv = float(-1) ** emm * (
np.sqrt(
fact(ell + emm)
* fact(ell - emm)
* (2 * ell + 1)
/ (4 * np.pi * fact(ell + spin_weight) * fact(ell - spin_weight))
)
* np.sin(theta_grid / 2) ** (2 * ell)
* Sum
)
return Yslmv
[docs]def Yslm_pres(spin_weight, ell, emm, theta, phi, pres=16):
''' Spin-weighted spherical harmonics data defined as a function of zeta and phi, for qlm data decomposition.
Inputs
-----------
spin_weight : int
The Spin weight.
ell : int
The mode number :math:`\\ell'.
emm : int
The azimuthal mode number :math:`m'.
theta : float
The polar angle :math:`\\theta` in radians,
phi : float
The aximuthal angle :math:`\\phi' in radians.
pres : int, optional
The precision i.e. number of digits to compute
upto. Default value is 16.
Returns
--------
Yslm : float
The value of Yslm at :math:`\\theta, phi'.
'''
import sympy as sp
tv, pv = theta, phi
theta, phi = sp.symbols('theta phi')
#fact = np.math.factorial
fact = sp.factorial
Sum = 0
for aar in range(ell-spin_weight + 1):
if (aar + spin_weight-emm)<0 or (ell - aar - spin_weight)<0:
#print('Continuing')
continue
else:
#print('r, l, s, m', r, l, s, m)
a1 = sp.binomial(ell-spin_weight, aar)
#print(a1)
a2 = sp.binomial(ell+spin_weight, aar+spin_weight-emm)
#print(a2)
a3 = sp.exp(1j*emm*phi)
#print(a3)
a4 = sp.tan(theta/2)
#print(a4)
Sum += sp.binomial(ell - spin_weight, aar)*sp.binomial(ell + spin_weight, aar + spin_weight-emm)*sp.Pow((-1), (ell - aar - spin_weight))*sp.exp(1j* emm * phi)/sp.Pow(sp.tan(theta / 2), (2 * aar + spin_weight - emm))
Yslm_expr = sp.Pow(-1, emm) * (sp.sqrt(fact(ell + emm) * fact(ell - emm) * (2*ell + 1)/(4 * sp.pi * fact(ell + spin_weight) * fact(ell-spin_weight)))*sp.Pow(sp.sin(theta / 2), (2 * ell)) * Sum)
Yslm_expr = sp.simplify(Yslm_expr)
return Yslm_expr.evalf(pres, subs={theta: sp.Float(f'{tv}'), phi : sp.Float(f'{pv}')})