def phase_correlation(signal, ref):
fft_s = fft(signal)
fft_r = fft(ref)
prod_of_ffts = (fft_s * np.conjugate(fft_r))
pc = prod_of_ffts / np.abs(prod_of_ffts)
return abs(ifft(pc))
def cryo_conv_vol(x, kernel_f):
n = x.shape[0]
n_ker = kernel_f.shape[0]
if np.any(np.array(x.shape) != n):
raise ValueError('Volume in `x` must be cubic')
if np.any(np.array(kernel_f.shape) != n_ker):
raise ValueError('Convolution kernel in `kernel_f` must be cubic')
is_singleton = len(x.shape) == 3
shifted_kernel_f = np.fft.ifftshift(np.fft.ifftshift(np.fft.ifftshift(kernel_f, 0), 1), 2)
if is_singleton:
x = numpy_fft.fftn(x, [n_ker] * 3)
else:
x = numpy_fft.fft(x, n=n_ker, axis=0)
x = numpy_fft.fft(x, n=n_ker, axis=1)
x = numpy_fft.fft(x, n=n_ker, axis=2)
x *= shifted_kernel_f
if is_singleton:
x = numpy_fft.ifftn(x)
x = x[:n, :n, :n]
else:
x = numpy_fft.ifft(x, axis=0)
x = numpy_fft.ifft(x, axis=1)
x = numpy_fft.ifft(x, axis=2)
x = x.real
return x
def time_to_frequency(ad, ed, td, wd, del_t, q, compensate_window=True):
"""
Convert time domain data to discrete Fourier transformed data,
with normalization by del_t x q x n / k1
Parameters
----------
ad
ed
td
wd
del_t
q
Returns
-------
"""
# ==================================================================================================================
# Now we get extract the data and transform it to frequency domain
# ==================================================================================================================
if compensate_window:
resc = ad.shape[0] / np.sum(wd)
else:
resc = 1.0
a_df = fft(wd * ad) * del_t * q * resc
e_df = fft(wd * ed) * del_t * q * resc
t_df = fft(wd * td) * del_t * q * resc
return a_df, e_df, t_df
def ft1D(x,y,*,axis=0,zero_DC=False):
"""Takes in x and y = y(x), and returns k and the Fourier transform
of y(x) -> f(k) along a single (1D) axis
Handles all of the annoyances of fftshift and ifftshift, and gets the
normalization right
Args:
x (np.ndarray) : independent variable, must be 1D
y (np.ndarray) : dependent variable, can be nD
Kwargs:
axis (int) : which axis to perform FFT
zero_DC (bool) : if true, sets f(0) = 0
"""
dx = x[1]-x[0]
k = fftshift(fftfreq(x.size,d=dx))*2*np.pi
fft_norm = dx
shifted_x = ifftshift(x)
if np.isclose(shifted_x[0],0):
f = fft(ifftshift(y,axes=(axis)),axis=axis)*fft_norm
else:
f = fft(y,axis=axis)*fft_norm
if zero_DC:
nd_slice = [slice(None) for i in range(len(f.shape))]
nd_slice[axis] = slice(0,1,1)
nd_slice = tuple(nd_slice)
f[nd_slice] = 0
f = fftshift(f,axes=(axis))
return k, f
def toeplitz_multiplication(v, first_row, first_column):
"""
Performs the matrix-vector product T * v where
T is a Toeplitz matrix, using FFT
Parameters
----------
v : array_like
input vector of size n_data
first_row : array_like
first row of the Toepltiz matrix (size n_data)
first_column : array_like
first column of the Toepltiz matrix (size n_data)
Returns
-------
y : numpy array
vector such that y = T * v
"""
n = first_row.shape[0]
a_2n_fft = fft(np.concatenate((first_row,[0],first_column[1:][::-1])))
return np.real(ifft(a_2n_fft*fft(v, 2*n))[0:n])
def toeplitz_matvec(col0, vec, n):
"""
col0: first column of the Toeplitz matrix with shape (n,)
vec: vector with shape (n,)
"""
p = (ifft(
fft(np.r_[col0, col0[-2:0:-1]]) *
fft(np.r_[vec, np.zeros(n - 2)])).real)[:n]
return p
def toeplitz_matvec_block(col0, vec, n):
"""
col0: first column of the Toeplitz matrix with shape (n,)
vec: vector with shape (m, n)
"""
m, n = vec.shape
padded_vec = np.zeros((m, n * 2 - 2))
padded_vec[:, :n] = vec
p = ifft(fft(np.r_[col0, col0[-2:0:-1]]) * fft(padded_vec)).real[:, :n]
return p
def csr_convolution(a, b):
P = len(a)
Q = len(b)
L = P + Q - 1
K = 2**nextpow2(L)
a_pad = np.pad(a, (0, K - P), 'constant', constant_values=(0))
b_pad = np.pad(b, (0, K - Q), 'constant', constant_values=(0))
c = ifft(fft(a_pad) * fft(b_pad))
c = c[0:L - 1].real
return c
def multiple_toepltiz_inverse(c_mat, lambda_n, a):
"""
Efficiently compute several Toepltiz systems with the same Toepltiz
matrix
T xj = cj
which is in matrix form
T x_mat = c_mat
Where T is a n_data x n_data Toeplitz matrix and c_mat is a n_data x n_knots matrix
"""
n = c_mat.shape[0]
#zero_vect = np.zeros(n_data)
# PRECOMPUTATIONS
# Cf. Step 2 of Ref. [1]
#ae_2n = np.concatenate(([1],a,zero_vect))
#ae_2n_fft = fft(ae_2n)
ae_2n_fft = fft(np.concatenate(([1],a)),2*n)
# using hermitian and real property of covariance matrices:
# be_2n_fft = fft(be_2n)
be_2n_fft = ae_2n_fft.conj() #np.real(ae_2n_fft) - 1j*np.imag(ae_2n_fft)
signs = (-1)**(np.arange(2*n)+1)
x_mat = np.empty(np.shape(c_mat))
print("shape of c_mat is " + str(c_mat.shape))
for j in range(c_mat.shape[1]):
#ce_2n = np.zeros(2*n_data)
#ce_2n[0:n_data] = c_mat[:,j]
#ce_2n = np.concatenate((c_mat[:,j],zero_vect))
#ce_2n_fft = fft(ce_2n)
ce_2n_fft = fft(c_mat[:,j],2*n)
u_2n = ifft( ae_2n_fft*ce_2n_fft )
v_2n = ifft( be_2n_fft*ce_2n_fft )
#pe_2n_fft = fft( np.concatenate((v_2n[0:n_data],zero_vect)) )
#qe_2n_fft = fft( np.concatenate((u_2n[n_data:],zero_vect)) )
pe_2n_fft = fft( v_2n[0:n] , 2*n )
qe_2n_fft = fft( u_2n[n:] , 2*n )
we_2n = ifft( ae_2n_fft*pe_2n_fft + signs*be_2n_fft*qe_2n_fft )
x_mat[:,j] = np.real(we_2n[0:n]/lambda_n)
return x_mat
def update_missing_data(self, y_gw_list):
"""
Update the value of missing data
Parameters
----------
y_gw_list : list
list of waveform in the time domain for each channel
Returns
-------
Nothing, just update the attribute self.data_dft
"""
# Update Gaussian process mean
self.model.update_mean(y_gw_list)
# Impute missing data
y_imp = self.model.impute(self.data)
# Transform back to Fourier domain, applying the windowing for complete
# time series, with re-scaling
resc = self.del_t * self.resc_full
data_dft = [
fft(y_imp0 * self.wd_full)[self.inds] * resc for y_imp0 in y_imp
]
# self.data_dft = data_dft[:]
return y_imp, data_dft
def _encoder(self, x, decimate, axis=-1, planner_effort='FFTW_ESTIMATE'):
"""
Transform the time series data into a multiple sub-banded demodulated signal in frequency domain.
"""
Nch, Nsamp = x.shape
Nsamp_dec = int(Nsamp / decimate)
X = fft.fft(x, axis=axis, planner_effort=planner_effort)
if Nsamp % 2 == 0:
X[:, 1:Nsamp // 2] *= 2
X[:, Nsamp // 2:] = 0
else:
X[:, 1:(Nsamp + 1) // 2] *= 2
X[:, (Nsamp + 1) // 2] = 0
X_ = fft.fftshift(X, axes=-1)[:,
int((Nsamp - Nsamp_dec) // 2):int(
(Nsamp + Nsamp_dec) // 2)] / decimate
if self.n_processes > 1:
func = self.pfunc.result
else:
func = self.multiply
return func(X_[:, self.encoder_rule], np.atleast_2d(self.Hwin))
def gls_covariance_time(mat, psd):
"""
Compute the covariance of the generalized least-square estimator for
a stationary noise.
Parameters
----------
mat : ndarray
Model design matrix of size n_data x n_params
psd : ndarray
Noise power spectrum = PSD * fs / 2 where fs is the sampling frequency
and PSD is the one-sided noise power spectral density.
Returns
-------
type
Description of returned object.
"""
# if type(wind) == str:
# wd = np.hanning(mat.shape[0])
# elif type(wind) == np.ndarray:
# wd = wind[:]
# Apply time-windowing
# mat_wind = mat * np.array([wd]).T
# Fourier transform the model
# mat_dft = fft(mat_wind, axis=0)
return gsl_covariance_freq(fft(mat, axis=0), psd)
def check_efield_resolution(self, efield, *, plot_fields=False):
efield_tail = np.max(np.abs([efield[0], efield[-1]]))
if efield_tail > np.max(np.abs(efield)) / 100:
warnings.warn(
'Consider using larger time interval, pulse does not decay to less than 1% of maximum value in time domain'
)
efield_fft = fftshift(fft(ifftshift(efield))) * self.dt
efield_fft_tail = np.max(np.abs([efield_fft[0], efield_fft[-1]]))
if efield_fft_tail > np.max(np.abs(efield_fft)) / 100:
warnings.warn(
'''Consider using smaller value of dt, pulse does not decay to less than 1% of maximum value in frequency domain'''
)
if plot_fields:
fig, axes = plt.subplots(1, 2)
l1, l2, = axes[0].plot(self.efield_t, np.real(efield),
self.efield_t, np.imag(efield))
plt.legend([l1, l2], ['Real', 'Imag'])
axes[1].plot(self.efield_w, np.real(efield_fft), self.efield_w,
np.imag(efield_fft))
axes[0].set_ylabel('Electric field Amp')
axes[0].set_xlabel('Time ($\omega_0^{-1})$')
axes[1].set_xlabel('Frequency ($\omega_0$)')
fig.suptitle(
'Check that efield is well-resolved in time and frequency')
plt.show()
def apply_gap_convolution(self, y_gw_fft_pos):
"""
Transform the frequency-domain waveform to account for gaps.
Parameters
----------
y_gw_fft_pos : list[ndarray]
list of frequency-domain waveforms for all channels.
Returns
-------
y_gw_masked_fft : list[ndarray]
list of distorted frequency-domain waveforms for all channels.
"""
# Convert waveform to time domain
y_gw_list = [
self.frequency_to_time(y_gw_fft_pos[i])
for i in range(len(y_gw_fft_pos))
]
# Apply mask window and Fourier transform back
y_gw_masked_fft = [
fft(self.wd * dat)[self.inds] * self.del_t * self.resc
for dat in y_gw_list
]
return y_gw_masked_fft
def track(self, im, pos, base_target_sz, current_scale_factor):
"""
track the scale using the scale filter
"""
# get scale filter features
scales = current_scale_factor * self.scale_size_factors
xs = self._extract_scale_sample(im, pos, base_target_sz, scales, self.scale_model_sz)
# project
xs = self.basis.dot(xs) * self.window
# get scores
xsf = fft(xs, axis=1)
scale_responsef = np.sum(self.sf_num * xsf, 0) / (self.sf_den + config.lamBda)
interp_scale_response = np.real(ifft(resize_dft(scale_responsef, config.number_of_interp_scales)))
recovered_scale_index = np.argmax(interp_scale_response)
if config.do_poly_interp:
# fit a quadratic polynomial to get a refined scale estimate
id1 = (recovered_scale_index - 1) % config.number_of_interp_scales
id2 = (recovered_scale_index + 1) % config.number_of_interp_scales
poly_x = np.array([self.interp_scale_factors[id1], self.interp_scale_factors[recovered_scale_index], self.interp_scale_factors[id2]])
poly_y = np.array([interp_scale_response[id1], interp_scale_response[recovered_scale_index], interp_scale_response[id2]])
poly_A = np.array([[poly_x[0]**2, poly_x[0], 1],
[poly_x[1]**2, poly_x[1], 1],
[poly_x[2]**2, poly_x[2], 1]], dtype=np.float32)
poly = np.linalg.inv(poly_A).dot(poly_y.T)
scale_change_factor = - poly[1] / (2 * poly[0])
else:
scale_change_factor = self.interp_scale_factors[recovered_scale_index]
return scale_change_factor
def toepltiz_mat_vect_prod(y, s_2n):
"""
Linear operator that calculate T y_in assuming that we can write:
Com = F* Lambda F
where Lambda is a P x P diagonal matrix and F is the P x n_data Discrete Fourier
Transform matrix.
Parameters
----------
y : numpy array
input data vector of size n_data
S_2N : numpy array (size P >= 2N)
PSD vector
Returns
-------
y_out : numpy array
y_out = T * y_in transformed output vector of size N_out
"""
return np.real(ifft(s_2n * fft(y, len(s_2n)))[0:len(y)])
def create_filter(order,
cutoff,
nyquist,
N,
ftype='fir',
output='freq',
shift=True):
"""
Create a prototype filter.
"""
h = firwin(order, cutoff, nyq=nyquist)
if output == 'freq':
w = fft.fftfreq(N)
w *= (nyquist * 2)
H = fft.fft(h, n=N, axis=-1, planner_effort='FFTW_ESTIMATE')
if shift:
return fft.fftshift(w), fft.fftshift(H)
else:
return w, H
else:
return h
def delay(x, time, fs, axis=-1, keeplength=False, pad=1):
extra_pad = 200 # add 200 samples to prevent wrapping
samps = int(np.floor(time * fs))
s = list(x.shape)
sz_pre = np.copy(s)
sz_post = np.copy(s)
sz_fft = np.copy(s)
sz_pre[axis] = samps
sz_post[axis] = pad + extra_pad
x = np.concatenate((np.zeros(sz_pre), x, np.zeros(sz_post)), axis)
sz_fft[axis] = int(np.round(2 ** np.ceil(np.log2(x.shape[axis])) -
x.shape[axis]))
x = np.concatenate((x, np.zeros(sz_fft)), axis)
new_len = sz_pre[axis] + s[axis] + sz_post[axis] + sz_fft[axis]
# x[n-k] <--> X(jw)e^(-jwk) where w in [0, 2pi)
if type(time) is not int:
theta = (-np.arange(new_len).astype(float) * fs * 2 * np.pi / new_len *
(time - np.float(samps) / fs))
theta[-(new_len // 2) + 1:] = -theta[(new_len // 2):1:-1]
st = [1 for _ in range(x.ndim)]
st[axis] = new_len
x = np.real(fft.ifft(fft.fft(x, axis=axis) *
np.exp(1j * theta.reshape(st))))
if keeplength:
x = np.take(x, range(s[axis]), axis)
else:
x = np.take(x, range(s[axis] + samps + pad), axis)
inds = tuple([slice(si) for si in sz_pre])
x[inds] = 0
return x
def detrend(y_data, detrend_method='poly', max_order=1, n_knots=10, psd=None):
"""
Remove linear or polynomial trend of order max_order
Parameter
----------
y_data : 1d array
input data of size n
max_order : int
maximum order of the polynomial to fit
psd : 1d array
noise PSD at Fourier frequencies (size n)
Returns
-------
y_detrend : 1d array
output detrended data (size n)
"""
t_norm = np.arange(0, y_data.shape[0])
if detrend_method == 'poly':
mat_linear = np.hstack(
[np.array([t_norm**k]).T for k in range(0, max_order + 1)])
if psd is not None:
amp = regression.generalized_least_squares(fft(mat_linear, axis=0),
fft(y_data), psd)
else:
amp = regression.least_squares(mat_linear, y_data)
trend = np.real(np.dot(mat_linear, amp))
y_detrend = y_data - trend
elif detrend_method == 'spline':
# Detrending using splines
n_seg = y_data.shape[0] // n_knots
t_knots = np.linspace(t_norm[n_seg], t_norm[-n_seg], n_knots)
spl = interpolate.LSQUnivariateSpline(t_norm,
y_data,
t_knots,
k=3,
ext="const")
trend = spl(t_norm)
y_detrend = y_data - trend
return y_detrend, trend
def plot2d_fft(self,
*,
delay_time_start=1,
create_figure=True,
color_range='auto',
subtract_DC=True,
draw_colorbar=True,
frequency_range=[-1000, 1000],
normalize=False,
phase=False,
save_fig=True,
wT_frequency_range='auto'):
w_ind = np.where((self.w > frequency_range[0])
& (self.w < frequency_range[1]))[0]
w = self.w[w_ind]
sig = self.signal_vs_delay_times[w_ind, :]
delay_time_indices = np.where(self.delay_times > delay_time_start)[0]
delay_times = self.delay_times[delay_time_indices]
sig = sig[:, delay_time_indices]
if normalize:
sig /= np.dot(self.dipoles, self.dipoles)**2
wT = fftshift(
fftfreq(delay_times.size,
d=(delay_times[1] - delay_times[0]))) * 2 * np.pi
sig_fft = fft(sig, axis=1)
if subtract_DC:
sig_fft[:, 0] = 0
sig_fft = fftshift(sig_fft, axes=(1))
ww, wTwT = np.meshgrid(wT, w)
if create_figure:
plt.figure()
if phase:
plt.title('Phase')
plot_sig = np.arctan2(np.imag(sig_fft), np.real(sig_fft))
else:
plt.title('Magnitude')
plot_sig = np.abs(sig_fft)
if color_range == 'auto':
plt.pcolormesh(ww, wTwT, plot_sig)
else:
plt.pcolormesh(ww,
wTwT,
plot_sig,
vmin=color_range[0],
vmax=color_range[1])
if draw_colorbar:
plt.colorbar()
plt.xlabel('$\omega_T$ ($\omega_0$)', fontsize=16)
plt.ylabel('Detection Frequency ($\omega_0$)', fontsize=16)
if wT_frequency_range == 'auto':
plt.xlim([0, np.max(wT)])
else:
plt.xlim(wT_frequency_range)
if save_fig:
plt.savefig(self.base_path + 'TA_spectra_fft')
def plotTA_units(self,
*,
frequency_range=[-1000, 1000],
subtract_DC=True,
create_figure=True,
color_range='auto',
draw_colorbar=True,
save_fig=True,
omega_0=1):
"""Plots the transient absorption spectra with detection frequency on the
y-axis and delay time on the x-axis.
Args:
frequency_range (list): sets the min (list[0]) and max (list[1]) detection frequency for y-axis
subtract_DC (bool): if True subtract the DC component of the TA
color_range (list): sets the min (list[0]) and max (list[1]) value for the colorbar
draw_colorbar (bool): if True add a colorbar to the plot
save_fig (bool): if True save the figure that is produced
omega_0 (float): convert from unitless variables, omega_0 should be provided in wavenumbers
"""
self.load_eigen_params()
f0_thz = omega_0 * 3E10 / 1.0E12 # omega_0 in wavenumbers
T_ps = self.delay_times / f0_thz / (2 * np.pi)
# Cut out unwanted detection frequency points
self.w += self.ground_to_excited_transition + self.center
self.w *= omega_0
w_ind = np.where((self.w > frequency_range[0])
& (self.w < frequency_range[1]))[0]
w = self.w[w_ind]
sig = self.signal[w_ind, :]
if omega_0 == 1:
xlab = r'Delay time ($\omega_0^{-1}$)'
ylab = r'Detection Frequency ($\omega_0$)'
else:
xlab = 'Delay time (ps)'
ylab = r'Detection Frequency (cm$^{-1}$)'
if subtract_DC:
sig_fft = fft(sig, axis=1)
sig_fft[:, 0] = 0
sig = np.real(ifft(sig_fft))
ww, tt = np.meshgrid(T_ps, w)
if create_figure:
plt.figure()
if color_range == 'auto':
plt.pcolormesh(ww, tt, sig)
else:
plt.pcolormesh(ww,
tt,
sig,
vmin=color_range[0],
vmax=color_range[1])
if draw_colorbar:
plt.colorbar()
plt.xlabel(xlab, fontsize=16)
plt.ylabel(ylab, fontsize=16)
if save_fig:
plt.savefig(self.base_path + 'TA_spectra_iso_ave')
def teopltiz_precompute(R,p=10,Nit = 1000,tol = 1e-4):
"""
Solve the system T y = e1 where T is symmetric Toepltiz.
where e1 = [1 0 0 0 0 0].T to compute the vector a and lambda_n for
further fast Toepltiz inversions.
Compute the prefactor lambda_n and vector a_{n-1} such that the inverse of
T writes
T^{-1} = (1/lambda_n) * [ 1 a*_{n-1}
a_{n-1} S_{n-1} ]
Parameters
----------
R : array_like
autocovariance function (first row of the Toepltiz matrix)
Returns
-------
lambda_n : scalar float
prefactor of the inverse of T
a : numpy array
vector of size N-1 involved in the computation of the inverse of T
p : scalar integer
maximum number of lags for the preconditionning
Nit : scalar integer
maximum number of iterations for PCG
tol : scalar float
relative error convergence criterium for the PCG algorithm
References
----------
[1] Jain, Fast Inversion of Banded Toeplitz Matrices by Circular
Decompositions, 1978
"""
N = len(R)
# First basis vector (of orthonormal cartesian basis)
e1 = np.concatenate(([1],np.zeros(N-1)))
# Compute spectrum
S_2N = fft(np.concatenate((R,[0],R[1:][::-1])))
# Linear operator correponding to the Toeplitz matrix
T_op = toepltizLinearOp(N,S_2N)
# Preconditionner to approximate T^{-1}
Psolver = computeToepltizPrecond(R,p=p)
# Build the associated linear operator
P_op = matrixalgebra.precondLinearOp(Psolver,N,N)
# Initial guess
z,info = sparse.linalg.bicgstab(T_op, e1, x0=np.zeros(N),tol=tol,
maxiter=Nit,M=P_op,callback=None)
matrixalgebra.printPCGstatus(info)
lambda_n = 1/z[0]
a = lambda_n * z[1:]
return lambda_n,a
def phaseran(signal):
"""
Performs phase randomization for coherence matrices.
NOTE: Signal input has to be nTimepoints x nElectrodes (there is a check)
"""
# check that it is in right orientation
if signal.shape[1] > signal.shape[0]:
signal = signal.T
# Get parameters
nTimepoints = signal.shape[0]
nElectrodes = signal.shape[1]
# Check to make sure that it is an odd number of samples
if nTimepoints % 2 == 0:
signal = signal[:-1, :]
nTimepoints = nTimepoints - 1
nTimepoints = signal.shape[0]
len_ser = (nTimepoints - 1) / 2
interv1 = np.arange(1, len_ser + 1)
interv2 = np.arange(len_ser + 1, nTimepoints)
# fft_A = pfft.builders.fft(signal, axis=0) # FFT of original data
try:
fft_A = pfft.fft(signal, axis=0, threads=15)
except:
fft_A = pfft.fft(signal, axis=0)
# Create the random phases for all the time series
ph_rnd = np.random.rand(len_ser, nElectrodes)
ph_interv1 = np.exp(2 * np.pi * 1j * ph_rnd)
ph_interv2 = np.conj(np.flipud(ph_interv1))
# Randomize all time series simultaneously
fft_recblk_surr = fft_A
fft_recblk_surr[interv1, :] = fft_A[interv1, :] * ph_interv1
fft_recblk_surr[interv2, :] = fft_A[interv2, :] * ph_interv2
surrblk = np.float32(pfft.ifft(fft_recblk_surr, axis=0)).T
return surrblk
def fourier_transform(array,axis,new_axis,inverse=False):
import numpy as np
from numpy.fft import fftshift,ifftshift
from expresso.pycas import pi
from ..coordinate_ndarray import CoordinateNDArray
try:
from pyfftw.interfaces.numpy_fft import fft,ifft
except ImportError:
from numpy.fft import fft,ifft
axi = array.axis.index(axis)
if len(array.axis) == 1:
if not inverse:
new_data = fftshift(fft(array.data,axis=axi),axes=[axi])
new_data *= 1/np.sqrt(2*np.pi)
else:
new_data = ifft(ifftshift(array.data,axes=[axi]),axis=axi)
new_data *= np.sqrt(2*np.pi)
else:
from pypropagate.progressbar import ProgressBar
axt = (axi + 1) % len(array.axis)
axi = axi if axt > axi else axi - 1
transposed_data = np.rollaxis(array.data,axt,start=0)
new_data = np.zeros(array.data.shape,dtype=complex)
transposed_new_data = np.rollaxis(new_data,axt,start=0)
for i in ProgressBar(range(transposed_data.shape[0])):
if not inverse:
transposed_new_data[i] = fftshift(fft(transposed_data[i],axis=axi),axes=[axi])
transposed_new_data[i] *= 1/np.sqrt(2*np.pi)
else:
transposed_new_data[i] = ifft(ifftshift(transposed_data[i],axes=[axi]),axis=axi)
transposed_new_data[i] *= np.sqrt(2*np.pi)
sw = array.bounds[axi][1] - array.bounds[axi][0]
tmin,tmax = array.evaluate((-(pi*array.shape[axi])/sw,
(pi*array.shape[axi])/sw))
new_bounds = [(b[0],b[1]) if i!=axi else (tmin,tmax) for i,b in enumerate(array.bounds)]
new_axis = [a if i!=axi else new_axis for i,a in enumerate(array.axis)]
return CoordinateNDArray(new_data,new_bounds,new_axis,array.evaluate)
def dft(self, wind='tukey', n_wind=500, normalized=True):
self.w = self.compute_window(wind=wind, n_wind=n_wind)
if normalized:
norm = np.sum(self.w) / (self.del_t * 2)
else:
norm = 1.0
return fft(self * self.w) / norm
def integrated_ft(self,delay_time_start = 1,delay_time_stop = 300):
delay_time_indices = np.where((self.delay_times > delay_time_start) & (self.delay_times < delay_time_stop))[0]
delay_times = self.delay_times[delay_time_indices]
sig = self.signal_vs_delay_times[:,delay_time_indices]
integrated = np.trapz(sig,x=self.TA.w,axis=0)
w_T = fftshift(fftfreq(delay_times.size,d=(delay_times[1] - delay_times[0])))*2*np.pi
integrated_fft = fft(integrated)
integrated_fft[0] = 0
integrated_fft = fftshift(integrated_fft)
return w_T, integrated_ft
def update(self, im, pos, base_target_sz, current_scale_factor):
"""
update the scale filter
"""
# get scale filter features
scales = current_scale_factor * self.scale_size_factors
xs = self._extract_scale_sample(im, pos, base_target_sz, scales,
self.scale_model_sz)
first_frame = not hasattr(self, 's_num')
if first_frame:
self.s_num = xs
else:
self.s_num = (1 - config.scale_learning_rate
) * self.s_num + config.scale_learning_rate * xs
# compute projection basis
if self.max_scale_dim:
self.basis, _ = scipy.linalg.qr(self.s_num, mode='economic')
scale_basis_den, _ = scipy.linalg.qr(xs, mode='economic')
else:
U, _, _ = np.linalg.svd(self.s_num)
self.basis = U[:, :self.s_num_compressed_dim]
self.basis = self.basis.T
# compute numerator
feat_proj = self.basis.dot(self.s_num) * self.window
sf_proj = fft(feat_proj, axis=1)
self.sf_num = self.yf * np.conj(sf_proj)
# update denominator
xs = scale_basis_den.T.dot(xs) * self.window
xsf = fft(xs, axis=1)
new_sf_den = np.sum(np.real(xsf * np.conj(xsf)), 0)
if first_frame:
self.sf_den = new_sf_den
else:
self.sf_den = (
1 - config.scale_learning_rate
) * self.sf_den + config.scale_learning_rate * new_sf_den
def _psd2invntt(psd, bandwidth, ncorr, fftw_flag='FFTW_MEASURE'):
"""
Compute the first row of the inverse of the noise time-time correlation
matrix from two-sided PSD with frequency increment fsamp/psd.size.
"""
fftsize = psd.shape[-1]
fsamp = bandwidth * fftsize
psd = (1 / fsamp) / psd
nthreads = multiprocessing.cpu_count()
out = fft.fft(psd, axis=-1, planner_effort=fftw_flag, threads=nthreads)
invntt = out[..., :ncorr+1].real / fftsize
return invntt
def compute_periodogram(self, x):
"""
Compute the windowed periodogram from time series x,
along with Fourier frequency grid
"""
# If size of analysed data is the same as the presets
if (x.shape[0] == self.N):
if self.wind == 'hanning':
# Compute periodogram
per = np.abs(fft(x * self.w))**2 / self.s2
elif self.wind == 'ones':
per = np.abs(fft(x))**2 / self.s2
else:
if self.wind == 'hanning':
w = np.hanning(x.shape[0])
s2 = np.sum(w**2)
per = np.abs(fft(x * w))**2 / s2
elif self.wind == 'ones':
s2 = x.shape[0]
per = np.abs(fft(x))**2 / s2
return per
def _psd2invntt(psd, bandwidth, ncorr, fftw_flag='FFTW_MEASURE'):
"""
Compute the first row of the inverse of the noise time-time correlation
matrix from two-sided PSD with frequency increment fsamp/psd.size.
"""
fftsize = psd.shape[-1]
fsamp = bandwidth * fftsize
psd = (1 / fsamp) / psd
nthreads = multiprocessing.cpu_count()
out = fft.fft(psd, axis=-1, planner_effort=fftw_flag, threads=nthreads)
invntt = out[..., :ncorr + 1].real / fftsize
return invntt
def add_gaussian_linewidth(self, sigma):
self.old_signal = self.signal.copy()
sig_tau_t = fftshift(fft(ifftshift(self.old_signal, axes=(-1)),
axis=-1),
axes=(-1))
sig_tau_t = sig_tau_t * (
np.exp(-self.t**2 / (2 * sigma**2))[np.newaxis, np.newaxis, :] *
np.exp(-self.t21_array**2 /
(2 * sigma**2))[:, np.newaxis, np.newaxis])
sig_tau_w = fftshift(ifft(ifftshift(sig_tau_t, axes=(-1)), axis=-1),
axes=(-1))
self.signal = sig_tau_w
def __init__(
self,
target_sz,
):
init_target_sz = target_sz
num_scales = config.number_of_scales_filter
scale_step = config.scale_step_filter
scale_sigma = config.number_of_interp_scales * config.scale_sigma_factor
scale_exp = np.arange(
-np.floor(num_scales - 1) / 2,
np.ceil(num_scales - 1) / 2 + 1,
dtype=np.float32) * config.number_of_interp_scales / num_scales
scale_exp_shift = np.roll(scale_exp,
(0, -int(np.floor((num_scales - 1) / 2))))
interp_scale_exp = np.arange(
-np.floor((config.number_of_interp_scales - 1) / 2),
np.ceil((config.number_of_interp_scales - 1) / 2) + 1,
dtype=np.float32)
interp_scale_exp_shift = np.roll(
interp_scale_exp,
[0, -int(np.floor(config.number_of_interp_scales - 1) / 2)])
self.scale_size_factors = scale_step**scale_exp
self.interp_scale_factors = scale_step**interp_scale_exp_shift
ys = np.exp(-0.5 * (scale_exp_shift**2) / (scale_sigma**2))
self.yf = np.real(fft(ys))[np.newaxis, :]
self.window = signal.hann(ys.shape[0])[np.newaxis, :].astype(
np.float32)
# make sure the scale model is not to large, to save computation time
if config.scale_model_factor**2 * np.prod(
init_target_sz) > config.scale_model_max_area:
scale_model_factor = np.sqrt(config.scale_model_max_area /
np.prod(init_target_sz))
else:
scale_model_factor = config.scale_model_factor
# set the scale model size
self.scale_model_sz = np.maximum(
np.floor(init_target_sz * scale_model_factor), np.array([8, 8]))
self.max_scale_dim = config.s_num_compressed_dim == 'MAX'
if self.max_scale_dim:
self.s_num_compressed_dim = len(self.scale_size_factors)
self.num_scales = num_scales
self.scale_step = scale_step
self.scale_factors = np.array([1])
def _sampling2psd(sampling, sampling_frequency, fftw_flag='FFTW_MEASURE'):
"""
Return folded PSD without binning.
"""
sampling = np.asarray(sampling)
n = sampling.size
nthreads = multiprocessing.cpu_count()
psd = np.abs(fft.fft(sampling, axis=-1, planner_effort=fftw_flag,
threads=nthreads))**2
freq = np.fft.fftfreq(n, d=1/sampling_frequency)[:n//2+1]
freq[-1] += sampling_frequency
psd = np.concatenate([[0.], 2 * psd[1:n//2], [psd[n//2]]]) / \
(n * sampling_frequency)
return freq, psd
def _gaussian_sample(nsamples, sampling_frequency, psd, twosided=False,
out=None, fftw_flag='FFTW_MEASURE'):
"""
Generate a gaussian N-sample sampled at fs from a one- or two-sided
Power Spectrum Density sampled at fs/N.
Parameter
---------
nsamples : int
Number of time samples.
sampling_frequency : float
Sampling frequency [Hz].
psd : array-like
One- or two-sided Power Spectrum Density [signal unit**2/Hz].
twosided : boolean, optional
Whether or not the input psd is one-sided (only positive frequencies)
or two-sided (positive and negative frequencies).
out : ndarray
Placeholder for the output buffer.
"""
psd = np.asarray(psd)
if out is None:
out = empty(psd.shape[:-1] + (nsamples,))
if not twosided:
psd = _unfold_psd(psd)
shape = psd.shape[:-1] + (nsamples,)
gauss = np.random.randn(*shape)
nthreads = multiprocessing.cpu_count()
ftgauss = fft.fft(gauss, planner_effort=fftw_flag, threads=nthreads)
ftgauss[..., 0] = 0
spec = ftgauss * np.sqrt(psd)
out[...] = fft.ifft(spec, planner_effort=fftw_flag,
threads=nthreads).real * np.sqrt(sampling_frequency)
return out
# For pre-channelized data, data are always complex,
# and should have shape (ntint, nchan, npol).
# For baseband data, we wish to get to the same shape for
# incoherent or by_channel, or just to fully channelized for coherent.
=======
if fh.telescope in ('arochime','arochime-raw'):
# take complex conjugate to ensure by-channel de-dispersion is applied correctly
vals = np.conj(vals)
>>>>>>> e6d25c1bfcc5a3426bc62becc3b3075fbea23ebc
if fh.nchan == 1:
# If we need coherent dedispersion, do FT of whole thing,
# otherwise to output channels, mimicking pre-channelized data.
if raw.dtype.kind == 'c': # complex data
nsamp = len(vals) if dedisperse == 'coherent' else nchan
vals = fft(vals.reshape(-1, nsamp, npol), axis=1,
**_fftargs)
else: # real data
nsamp = len(vals) if dedisperse == 'coherent' else nchan * 2
vals = rfft(vals.reshape(-1, nsamp, npol), axis=1,
**_rfftargs)
# Sadly, the way data are stored depends on what FFT routine
# one is using. We cannot deal with scipy's.
if vals.dtype.kind == 'f':
raise TypeError("Can no longer deal with scipy's format "
"for storing FTs of real data.")
if fedge_at_top:
# take complex conjugate to ensure by-channel de-dispersion is
# applied correctly.
# This needs to be done for ARO data, since we are in 2nd Nyquist
# zone; not clear it is needed for other telescopes.
def muenchetal(im, args):
"""Process a sinogram image with the Munch et al. de-striping algorithm.
Parameters
----------
im : array_like
Image data as numpy array.
wlevel : int
Levels of the wavelet decomposition.
sigma : float
Smoothing effect.
(Parameters wlevel and sigma have to passed as a string separated by ;)
Example (using tiffile.py)
--------------------------
>>> im = imread('sino_orig.tif')
>>> im = munchetal(im, '4;1.0')
>>> imsave('sino_flt.tif', im)
References
----------
B. Munch, P. Trtik, F. Marone, M. Stampanoni, Stripe and ring artifact removal with
combined wavelet-Fourier filtering, Optics Express 17(10):8567-8591, 2009.
"""
# Disable a warning:
simplefilter("ignore", ComplexWarning)
# Get args:
wlevel, sigma = args.split(";")
wlevel = int(wlevel)
sigma = float(sigma)
# The wavelet transform to use : {'haar', 'db1'-'db20', 'sym2'-'sym20', 'coif1'-'coif5', 'dmey'}
wname = "db5"
# Wavelet decomposition:
coeffs = wavedec2(im.astype(float32), wname, level=wlevel)
coeffsFlt = [coeffs[0]]
# FFT transform of horizontal frequency bands:
for i in range(1, wlevel + 1):
# Padding and windowing of input signal:
n_byte_align(coeffs[i][1], simd_alignment)
siz = coeffs[i][1].shape
tmp = pad(coeffs[i][1], pad_width=((coeffs[i][1].shape[0] / 2, coeffs[i][1].shape[0] / 2), (0,0)), mode='constant') # or 'constant' for zero padding
tmp = pad(tmp, pad_width=((0,0) ,(coeffs[i][1].shape[1] / 2, coeffs[i][1].shape[1] / 2)), mode='constant') # or 'constant' for zero padding
tmp = _windowing_lr(tmp, siz[1])
tmp = _windowing_lr(tmp.T, siz[0]).T
# FFT:
fcV = fftshift(fft(tmp, axis=0, threads=2))
my, mx = fcV.shape
# Damping of vertical stripes:
damp = 1 - npexp(-(arange(-floor(my / 2.),-floor(my / 2.) + my) ** 2) / (2 * (sigma ** 2)))
dampprime = kron(ones((1,mx)), damp.reshape((damp.shape[0],1)))
fcV = fcV * dampprime
# Inverse FFT:
fcV = ifftshift(fcV)
n_byte_align(fcV, simd_alignment)
fcVflt = ifft(fcV, axis=0, threads=2)
## Crop image:
tmp = fcVflt[fcVflt.shape[0] / 4:(fcVflt.shape[0] / 4 + siz[0]), fcVflt.shape[1] / 4:(fcVflt.shape[1] / 4 + siz[1])]
# Dump back coefficients:
cVHDtup = (coeffs[i][0], tmp, coeffs[i][2])
coeffsFlt.append(cVHDtup)
# Get wavelet reconstruction:
im_f = real(waverec2(coeffsFlt, wname))
# Return filtered image (an additional row and/or column might be present):
return im_f[0:im.shape[0],0:im.shape[1]].astype(float32)
def fold(fh, comm, samplerate, fedge, fedge_at_top, nchan,
nt, ntint, ngate, ntbin, ntw, dm, fref, phasepol,
dedisperse='incoherent',
do_waterfall=True, do_foldspec=True, verbose=True,
progress_interval=100, rfi_filter_raw=None, rfi_filter_power=None,
return_fits=False):
"""
FFT data, fold by phase/time and make a waterfall series
Folding is done from the position the file is currently in
Parameters
----------
fh : file handle
handle to file holding voltage timeseries
comm: MPI communicator or None
will use size, rank attributes
samplerate : Quantity
rate at which samples were originally taken and thus double the
band width (frequency units)
fedge : float
edge of the frequency band (frequency units)
fedge_at_top: bool
whether edge is at top (True) or bottom (False)
nchan : int
number of frequency channels for FFT
nt, ntint : int
total number nt of sets, each containing ntint samples in each file
hence, total # of samples is nt*ntint, with each sample containing
a single polarisation
ngate, ntbin : int
number of phase and time bins to use for folded spectrum
ntbin should be an integer fraction of nt
ntw : int
number of time samples to combine for waterfall (does not have to be
integer fraction of nt)
dm : float
dispersion measure of pulsar, used to correct for ism delay
(column number density)
fref: float
reference frequency for dispersion measure
phasepol : callable
function that returns the pulsar phase for time in seconds relative to
start of the file that is read.
dedisperse : None or string (default: incoherent).
None, 'incoherent', 'coherent', 'by-channel'.
Note: None really does nothing
do_waterfall, do_foldspec : bool
whether to construct waterfall, folded spectrum (default: True)
verbose : bool or int
whether to give some progress information (default: True)
progress_interval : int
Ping every progress_interval sets
return_fits : bool (default: False)
return a subint fits table for rank == 0 (None otherwise)
"""
assert dedisperse in (None, 'incoherent', 'by-channel', 'coherent')
need_fine_channels = dedisperse in ['by-channel', 'coherent']
assert nchan % fh.nchan == 0
if dedisperse in ['incoherent', 'by-channel'] and fh.nchan > 1:
oversample = nchan // fh.nchan
assert ntint % oversample == 0
else:
oversample = 1
if dedisperse == 'coherent' and fh.nchan > 1:
raise ValueError("Cannot coherently dedisperse channelized data.")
if comm is None:
mpi_rank = 0
mpi_size = 1
else:
mpi_rank = comm.rank
mpi_size = comm.size
npol = getattr(fh, 'npol', 1)
assert npol == 1 or npol == 2
if verbose > 1 and mpi_rank == 0:
print("Number of polarisations={}".format(npol))
# initialize folded spectrum and waterfall
# TODO: use estimated number of points to set dtype
if do_foldspec:
foldspec = np.zeros((ntbin, nchan, ngate, npol**2), dtype=np.float32)
icount = np.zeros((ntbin, nchan, ngate), dtype=np.int32)
else:
foldspec = None
icount = None
if do_waterfall:
nwsize = nt*ntint//ntw//oversample
waterfall = np.zeros((nwsize, nchan, npol**2), dtype=np.float64)
else:
waterfall = None
if verbose and mpi_rank == 0:
print('Reading from {}'.format(fh))
nskip = fh.tell()/fh.blocksize
if nskip > 0:
if verbose and mpi_rank == 0:
print('Starting {0} blocks = {1} bytes out from start.'
.format(nskip, nskip*fh.blocksize))
dt1 = (1./samplerate).to(u.s)
# need 2*nchan real-valued samples for each FFT
if fh.telescope == 'lofar':
dtsample = fh.dtsample
else:
dtsample = nchan // oversample * 2 * dt1
tstart = dtsample * ntint * nskip
# pre-calculate time delay due to dispersion in coarse channels
# for channelized data, frequencies are known
tb = -1. if fedge_at_top else +1.
if fh.nchan == 1:
if getattr(fh, 'data_is_complex', False):
# for complex data, really each complex sample consists of
# 2 real ones, so multiply dt1 by 2.
freq = fedge + tb * fftfreq(nchan, 2.*dt1)
if dedisperse == 'coherent':
fcoh = fedge + tb * fftfreq(nchan*ntint, 2.*dt1)
fcoh.shape = (-1, 1)
elif dedisperse == 'by-channel':
fcoh = freq + tb * fftfreq(ntint, dtsample)[:, np.newaxis]
else: # real data
freq = fedge + tb * rfftfreq(nchan*2, dt1)
if dedisperse == 'coherent':
fcoh = fedge + tb * rfftfreq(ntint*nchan*2, dt1)
fcoh.shape = (-1, 1)
elif dedisperse == 'by-channel':
fcoh = freq + tb * fftfreq(ntint, dtsample)[:, np.newaxis]
freq_in = freq
else:
# Input frequencies may not be the ones going out.
freq_in = fh.frequencies
if oversample == 1:
freq = freq_in
else:
freq = freq_in[:, np.newaxis] + tb * fftfreq(oversample, dtsample)
fcoh = freq_in + tb * fftfreq(ntint, dtsample)[:, np.newaxis]
# print('fedge_at_top={0}, tb={1}'.format(fedge_at_top, tb))
# By taking only up to nchan, we remove the top channel at the Nyquist
# frequency for real, unchannelized data.
ifreq = freq[:nchan].ravel().argsort()
# pre-calculate time offsets in (input) channelized streams
dt = dispersion_delay_constant * dm * (1./freq_in**2 - 1./fref**2)
if need_fine_channels:
# pre-calculate required turns due to dispersion.
#
# set frequency relative to which dispersion is coherently corrected
if dedisperse == 'coherent':
_fref = fref
else:
_fref = freq_in[np.newaxis, :]
# (check via eq. 5.21 and following in
# Lorimer & Kramer, Handbook of Pulsar Astronomy
dang = (dispersion_delay_constant * dm * fcoh *
(1./_fref-1./fcoh)**2) * u.cycle
with u.set_enabled_equivalencies(u.dimensionless_angles()):
dd_coh = np.exp(dang * 1j).conj().astype(np.complex64)
# add dimension for polarisation
dd_coh = dd_coh[..., np.newaxis]
# Calculate the part of the whole file this node should handle.
size_per_node = (nt-1)//mpi_size + 1
start_block = mpi_rank*size_per_node
end_block = min((mpi_rank+1)*size_per_node, nt)
for j in range(start_block, end_block):
if verbose and j % progress_interval == 0:
print('#{:4d}/{:4d} is doing {:6d}/{:6d} [={:6d}/{:6d}]; '
'time={:18.12f}'
.format(mpi_rank, mpi_size, j+1, nt,
j-start_block+1, end_block-start_block,
(tstart+dtsample*j*ntint).value)) # time since start
# Just in case numbers were set wrong -- break if file ends;
# better keep at least the work done.
try:
raw = fh.seek_record_read(int((nskip+j)*fh.blocksize),
fh.blocksize)
except(EOFError, IOError) as exc:
print("Hit {0!r}; writing data collected.".format(exc))
break
if verbose >= 2:
print("#{:4d}/{:4d} read {} items"
.format(mpi_rank, mpi_size, raw.size), end="")
if npol == 2 and raw.dtype.fields is not None:
raw = raw.view(raw.dtype.fields.values()[0][0])
if fh.nchan == 1: # raw.shape=(ntint*npol)
raw = raw.reshape(-1, npol)
else: # raw.shape=(ntint, nchan*npol)
raw = raw.reshape(-1, fh.nchan, npol)
if dedisperse == 'incoherent' and oversample > 1:
raw = ifft(raw, axis=1, **_fftargs).reshape(-1, nchan, npol)
raw = fft(raw, axis=1, **_fftargs)
if rfi_filter_raw is not None:
raw, ok = rfi_filter_raw(raw)
if verbose >= 2:
print("... raw RFI (zap {0}/{1})"
.format(np.count_nonzero(~ok), ok.size), end="")
if np.can_cast(raw.dtype, np.float32):
vals = raw.astype(np.float32)
else:
assert raw.dtype.kind == 'c'
vals = raw
# For pre-channelized data, data are always complex,
# and should have shape (ntint, nchan, npol).
# For baseband data, we wish to get to the same shape for
# incoherent or by_channel, or just to fully channelized for coherent.
if fh.nchan == 1:
# If we need coherent dedispersion, do FT of whole thing,
# otherwise to output channels, mimicking pre-channelized data.
if raw.dtype.kind == 'c': # complex data
nsamp = len(vals) if dedisperse == 'coherent' else nchan
vals = fft(vals.reshape(-1, nsamp, npol), axis=1,
**_fftargs)
else: # real data
nsamp = len(vals) if dedisperse == 'coherent' else nchan * 2
vals = rfft(vals.reshape(-1, nsamp, npol), axis=1,
**_rfftargs)
# Sadly, the way data are stored depends on what FFT routine
# one is using. We cannot deal with scipy's.
if vals.dtype.kind == 'f':
raise TypeError("Can no longer deal with scipy's format "
"for storing FTs of real data.")
if fedge_at_top:
# take complex conjugate to ensure by-channel de-dispersion is
# applied correctly.
# This needs to be done for ARO data, since we are in 2nd Nyquist
# zone; not clear it is needed for other telescopes.
np.conj(vals, out=vals)
# Now we coherently dedisperse, either all of it or by channel.
if need_fine_channels:
# for by_channel, we have vals.shape=(ntint, nchan, npol),
# and want to FT over ntint to get fine channels;
if vals.shape[0] > 1:
fine = fft(vals, axis=0, **_fftargs)
else:
# for coherent, we just reshape:
# (1, ntint*nchan, npol) -> (ntint*nchan, 1, npol)
fine = vals.reshape(-1, 1, npol)
# Dedisperse.
fine *= dd_coh
# Still have fine.shape=(ntint, nchan, npol),
# w/ nchan=1 for coherent.
if fine.shape[1] > 1 or raw.dtype.kind == 'c':
vals = ifft(fine, axis=0, **_fftargs)
else:
vals = irfft(fine, axis=0, **_rfftargs)
if fine.shape[1] == 1 and nchan > 1:
# final FT to get requested channels
if vals.dtype.kind == 'f':
vals = vals.reshape(-1, nchan*2, npol)
vals = rfft(vals, axis=1, **_rfftargs)
else:
vals = vals.reshape(-1, nchan, npol)
vals = fft(vals, axis=1, **_fftargs)
elif dedisperse == 'by-channel' and oversample > 1:
vals = vals.reshape(-1, oversample, fh.nchan, npol)
vals = fft(vals, axis=1, **_fftargs)
vals = vals.transpose(0, 2, 1, 3).reshape(-1, nchan, npol)
# vals[time, chan, pol]
if verbose >= 2:
print("... dedispersed", end="")
if npol == 1:
power = vals.real**2 + vals.imag**2
else:
p0 = vals[..., 0]
p1 = vals[..., 1]
power = np.empty(vals.shape[:-1] + (4,), np.float32)
power[..., 0] = p0.real**2 + p0.imag**2
power[..., 1] = p0.real*p1.real + p0.imag*p1.imag
power[..., 2] = p0.imag*p1.real - p0.real*p1.imag
power[..., 3] = p1.real**2 + p1.imag**2
if verbose >= 2:
print("... power", end="")
# current sample positions and corresponding time in stream
isr = j*(ntint // oversample) + np.arange(ntint // oversample)
tsr = (isr*dtsample*oversample)[:, np.newaxis]
if rfi_filter_power is not None:
power = rfi_filter_power(power, tsr.squeeze())
print("... power RFI", end="")
# correct for delay if needed
if dedisperse in ['incoherent', 'by-channel']:
# tsample.shape=(ntint/oversample, nchan_in)
tsr = tsr - dt
if do_waterfall:
# # loop over corresponding positions in waterfall
# for iw in xrange(isr[0]//ntw, isr[-1]//ntw + 1):
# if iw < nwsize: # add sum of corresponding samples
# waterfall[iw, :] += np.sum(power[isr//ntw == iw],
# axis=0)[ifreq]
iw = np.round((tsr / dtsample / oversample).to(1)
.value / ntw).astype(int)
for k, kfreq in enumerate(ifreq): # sort in frequency while at it
iwk = iw[:, (0 if iw.shape[1] == 1 else kfreq // oversample)]
iwk = np.clip(iwk, 0, nwsize-1, out=iwk)
iwkmin = iwk.min()
iwkmax = iwk.max()+1
for ipow in range(npol**2):
waterfall[iwkmin:iwkmax, k, ipow] += np.bincount(
iwk-iwkmin, power[:, kfreq, ipow], iwkmax-iwkmin)
if verbose >= 2:
print("... waterfall", end="")
if do_foldspec:
ibin = (j*ntbin) // nt # bin in the time series: 0..ntbin-1
# times and cycles since start time of observation.
tsample = tstart + tsr
phase = (phasepol(tsample.to(u.s).value.ravel())
.reshape(tsample.shape))
# corresponding PSR phases
iphase = np.remainder(phase*ngate, ngate).astype(np.int)
for k, kfreq in enumerate(ifreq): # sort in frequency while at it
iph = iphase[:, (0 if iphase.shape[1] == 1
else kfreq // oversample)]
# sum and count samples by phase bin
for ipow in range(npol**2):
foldspec[ibin, k, :, ipow] += np.bincount(
iph, power[:, kfreq, ipow], ngate)
icount[ibin, k, :] += np.bincount(
iph, power[:, kfreq, 0] != 0., ngate).astype(np.int32)
if verbose >= 2:
print("... folded", end="")
if verbose >= 2:
print("... done")
#Commented out as workaround, this was causing "Referenced before assignment" errors with JB data
#if verbose >= 2 or verbose and mpi_rank == 0:
# print('#{:4d}/{:4d} read {:6d} out of {:6d}'
# .format(mpi_rank, mpi_size, j+1, nt))
if npol == 1:
if do_foldspec:
foldspec = foldspec.reshape(foldspec.shape[:-1])
if do_waterfall:
waterfall = waterfall.reshape(waterfall.shape[:-1])
return foldspec, icount, waterfall