def __init__(self,
file_path,
*,
num_conv_points=138,
dt=0.1,
center=0,
initial_state=0,
total_num_time_points=2000):
self.slicing_time = 0
self.interpolation_time = 0
self.expectation_time = 0
self.next_order_expectation_time = 0
self.convolution_time = 0
self.extend_time = 0
self.mask_time = 0
self.dipole_time = 0
self.base_path = file_path
self.undersample_factor = 1
self.set_homogeneous_linewidth(0.05)
self.set_inhomogeneous_linewidth(0)
self.load_eigenvalues()
self.load_mu()
self.efield_t = np.arange(-(num_conv_points // 2), num_conv_points // 2
+ num_conv_points % 2) * dt
self.efield_w = 2 * np.pi * fftshift(fftfreq(self.efield_t.size, d=dt))
# Code will not actually function until the following three empty lists are set by the user
self.efields = [] #initialize empty list of electric field shapes
self.polarization_sequence = [
] #initialize empty polarization sequence
self.pulse_times = [] #initialize empty list of pulse arrival times
HeavisideConvolve.__init__(self, num_conv_points)
# Initialize time array to be used for all desired delay times
self.t = np.arange(
-(total_num_time_points // 2),
total_num_time_points // 2 + total_num_time_points % 2) * dt
# The first pulse is assumed to arrive at t = 0, therefore shift array so that
# it includes only points where the signal will be nonzero (number of negative time points
# is essentially based upon width of the electric field, via the proxy of the size parameter
self.t += self.t[-(self.size // 2 + 1)]
self.dt = dt
# f = fftshift(fftfreq(self.t.size-self.t.size%2,d=self.dt))
f = fftshift(fftfreq(self.t.size, d=self.dt))
self.w = 2 * np.pi * f
self.initial_ground_state_index = initial_state
# Define the unitary operator for each manifold in the RWA given the rotating frequency center
self.recenter(new_center=center)
self.gamma_res = 6.91
def fillWithGaussianRandomField(self, ell, Cell, bufferFactor=1, threads=1):
"""
Generate a Gaussian random field from an input power spectrum specified
as ell, Cell.
Notes
-----
BufferFactor = 1 means the map will have periodic boundary function, while
BufferFactor > 1 means the map will be genrated on a patch bufferFactor
times larger in each dimension and then cut out so as to have
non-periodic boundary conditions.
Fills the data field of the map with the GRF realization.
"""
ft = fftTools.fftFromLiteMap(self, threads=threads)
Ny = self.Ny * bufferFactor
Nx = self.Nx * bufferFactor
bufferFactor = int(bufferFactor)
realPart = numpy.zeros([Ny, Nx])
imgPart = numpy.zeros([Ny, Nx])
ly = fftfreq(Ny, d=self.pixScaleY) * (2 * numpy.pi)
lx = fftfreq(Nx, d=self.pixScaleX) * (2 * numpy.pi)
# print ly
modLMap = numpy.zeros([Ny, Nx])
iy, ix = numpy.mgrid[0:Ny, 0:Nx]
modLMap[iy, ix] = numpy.sqrt(ly[iy] ** 2 + lx[ix] ** 2)
s = splrep(ell, Cell, k=3)
ll = numpy.ravel(modLMap)
kk = splev(ll, s)
id = numpy.where(ll > ell.max())
kk[id] = 0.0
# add a cosine ^2 falloff at the very end
# id2 = numpy.where( (ll> (ell.max()-500)) & (ll<ell.max()))
# lEnd = ll[id2]
# kk[id2] *= numpy.cos((lEnd-lEnd.min())/(lEnd.max() -lEnd.min())*numpy.pi/2)
# pylab.loglog(ll,kk)
area = Nx * Ny * self.pixScaleX * self.pixScaleY
p = numpy.reshape(kk, [Ny, Nx]) / area * (Nx * Ny) ** 2
realPart = numpy.sqrt(p) * numpy.random.randn(Ny, Nx)
imgPart = numpy.sqrt(p) * numpy.random.randn(Ny, Nx)
kMap = realPart + 1j * imgPart
if have_pyFFTW:
data = numpy.real(ifft2(kMap, threads=threads))
else:
data = numpy.real(ifft2(kMap))
b = bufferFactor
self.data = data[(b - 1) / 2 * self.Ny : (b + 1) / 2 * self.Ny, (b - 1) / 2 * self.Nx : (b + 1) / 2 * self.Nx]
def __init__(self,
file_path,
*,
num_conv_points=138,
dt=0.1,
center=0,
initial_state=0,
total_num_time_points=2000):
self.size = num_conv_points
# Initialize time array to be used for all desired delay times
self.t = np.arange(
-(total_num_time_points // 2),
total_num_time_points // 2 + total_num_time_points % 2) * dt
self.t += self.t[-(self.size // 2 + 1)]
self.dt = dt
parameter_file = os.path.join(file_path, 'params.yaml')
super().__init__(parameter_file, mask_by_occupation_num=True)
self.base_path = file_path
self.load_params()
self.set_diagrams_and_manifolds()
self.set_molecular_dipoles()
self.set_bottom_eigensystem()
self.set_H()
self.zero_hamiltonians()
self.rtol = 1E-6
self.atol = 1E-6
self.time_to_extend = 0
self.time_for_next_order = 0
############### Optical part
self.set_homogeneous_linewidth(0.05)
self.efield_t = np.arange(-(num_conv_points // 2), num_conv_points // 2
+ num_conv_points % 2) * dt
self.efield_w = 2 * np.pi * fftshift(fftfreq(self.efield_t.size, d=dt))
# Code will not actually function until the following three empty lists are set by the user
self.efields = [] #initialize empty list of electric field shapes
self.polarization_sequence = [
] #initialize empty polarization sequence
self.pulse_times = [] #initialize empty list of pulse arrival times
f = fftshift(fftfreq(self.t.size - self.t.size % 2, d=self.dt))
self.w = 2 * np.pi * f
# Define the unitary operator for each manifold in the RWA given the rotating frequency center
self.recenter(new_center=center)
def generate_signals(
shape: tuple,
specs: list[CompanionSpec],
template: np.ndarray,
derotation_angles: np.ndarray = None,
template_scale_factors: Union[np.ndarray, float, None] = None,
):
'''Inject signals for companions specified using
optional derotation angles to *counter*rotate the coordinates
before injection, such that rotation CCW (e.g. by `derotate_cube`)
aligns 0 deg PA with +Y
Parameters
----------
shape : tuple[int,int,int]
specs : list[CompanionSpec]
template : np.ndarray
derotation_angles : Optional[np.ndarray]
template_scale_factors : Union[np.ndarray,float,None]
Scale factor relative to 1.0 being the average brightness
of the primary over the observation, used to scale the
template image to reflect particularly sharp or poor
AO correction
Returns
-------
outcube : np.ndarray
'''
outcube = np.zeros(shape, dtype=template.dtype)
n_obs = shape[0]
template = improc.shift2(template, 0, 0, output_shape=shape[1:])
ft_template = fft.fft2(template)
xfreqs = fft.fftfreq(shape[2])
yfreqs = fft.fftfreq(shape[1])
if template_scale_factors is None:
template_scale_factors = np.ones(n_obs)
if np.isscalar(template_scale_factors):
template_scale_factors = np.repeat(np.array([template_scale_factors]),
n_obs)
if derotation_angles is None:
derotation_angles = np.zeros(n_obs)
for spec in specs:
theta = np.deg2rad(90 + spec.pa_deg - derotation_angles)
for i in range(n_obs):
dx = spec.r_px * np.cos(theta[i])
dy = spec.r_px * np.sin(theta[i])
shifter = np.exp(2j * np.pi * ((-dx * xfreqs[np.newaxis, :]) +
(-dy * yfreqs[:, np.newaxis])))
cube_contribution = fft.ifft2(ft_template * shifter).real
cube_contribution *= template_scale_factors[i] * spec.scale
outcube[i] += cube_contribution
return outcube
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 ift1D(k,f,*,axis=0,zero_DC=False):
"""Takes in k and f = f(k), and returns x and the discrete Fourier
transform of f(k) -> y(x).
Handles all of the annoyances of fftshift and ifftshift, and gets the
normalization right
Args:
x (np.ndarray): independent variable
y (np.ndarray): dependent variable
Kwargs:
axis (int) : which axis to perform FFT
"""
dk = k[1]-k[0]
x = fftshift(fftfreq(k.size,d=dk))*2*np.pi
ifft_norm = dk*k.size/(2*np.pi)
shifted_k = ifftshift(k)
if np.isclose(shifted_k[0],0):
y = ifft(ifftshift(f,axes=(axis)),axis=axis)*ifft_norm
else:
y = ifft(f,axis=axis)*ifft_norm
if zero_DC:
nd_slice = [slice(None) for i in range(len(y.shape))]
nd_slice[axis] = slice(0,1,1)
nd_slice = tuple(nd_slice)
y[nd_slice] = 0
y = fftshift(y,axes=(axis))
return x, y
def set_efields(self,times_list,efields_list,centers_list,phase_discrimination,*,reset_rhos = True,
plot_fields = False):
self.efield_times = times_list
self.efields = efields_list
self.centers = centers_list
self.set_phase_discrimination(phase_discrimination)
self.dts = []
self.efield_frequencies = []
if reset_rhos:
self.rhos = dict()
for t in times_list:
if t.size == 1:
dt = 1
w = np.array([0])
else:
dt = t[1] - t[0]
w = fftshift(fftfreq(t.size,d=dt))*2*np.pi
self.dts.append(dt)
self.efield_frequencies.append(w)
self.dt = self.dts[0]
if self.detection_type == 'polarization':
try:
self.local_oscillator = self.efields[-1].copy()
except:
self.local_oscillator = copy.deepcopy(self.efields[-1])
for field in self.efields:
if len(field) == 1:
# M = 1 is the impulsive limit
pass
else:
self.check_efield_resolution(field,plot_fields = plot_fields)
def logscale_normalization(spectra, srate=1., factor=20.):
t_bins, f_bins = spectra.shape
scale = np.linspace(0, 1, f_bins)**factor
scale *= (f_bins - 1) / scale.max()
scale = np.asarray(np.unique(np.round(scale)), dtype=np.int64)
# create spectrogram with new freq bins
spectra_ = np.zeros((t_bins, scale.size), dtype=np.complex128)
for i in range(0, scale.size):
if i < scale.size - 1:
spectra_[:, i] = np.sum(spectra[:, scale[i]:scale[i + 1]], axis=1)
else:
spectra_[:, i] = np.sum(spectra[:, scale[i]:], axis=1)
# list center freq of bins
allfreqs = np.abs(fft.fftfreq(f_bins * 2, 1. / srate)[:f_bins + 1])
freqs = []
for i in range(0, len(scale)):
if i == len(scale) - 1:
freqs += [np.mean(allfreqs[scale[i]:])]
else:
freqs += [np.mean(allfreqs[scale[i]:scale[i + 1]])]
return spectra_, np.asarray(freqs)
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 ft_shift2(image: np.ndarray,
dy: float,
dx: float,
flux_tol: Union[None, float] = 1e-15,
output_shape=None):
"""
Fast Fourier subpixel shifting
Parameters
----------
dy : float
Translation in +Y direction (i.e. a feature at (x, y) moves to (x, y + dy))
dx : float
Translation in +X direction (i.e. a feature at (x, y) moves to (x + dx, y))
flux_tol : float
Fractional flux change permissible
``(sum(output) - sum(image)) / sum(image) < flux_tol``
(default: 1e-15)
output_shape : tuple
shape of output array (default: same as input)
"""
if output_shape is None:
output_shape = image.shape
xfreqs = fft.fftfreq(output_shape[1])
yfreqs = fft.fftfreq(output_shape[0])
xform = fft.fft2(image, s=output_shape)
if output_shape is not None:
# compute center-to-center displacement such that
# supplying dx == dy == 0.0 will be a no-op (aside
# from changing shape)
orig_ctr_x, orig_ctr_y = (image.shape[1] - 1) / 2, (image.shape[0] -
1) / 2
new_ctr_x, new_ctr_y = (output_shape[1] - 1) / 2, (output_shape[0] -
1) / 2
base_dx, base_dy = new_ctr_x - orig_ctr_x, new_ctr_y - orig_ctr_y
else:
base_dx = base_dy = 0
modified_xform = xform * np.exp(2j * np.pi * (
(-(dx + base_dx) * xfreqs)[np.newaxis, :] +
(-(dy + base_dy) * yfreqs)[:, np.newaxis]))
new_image = fft.ifft2(modified_xform).real
frac_diff_flux = (np.sum(image) - np.sum(new_image)) / np.sum(image)
if flux_tol is not None and frac_diff_flux > flux_tol:
raise RuntimeError(
f"Flux conservation violated by {frac_diff_flux} fractional difference (more than {flux_tol})"
)
return new_image
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 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 _gen_kr(self):
"""Internal utiltiy to generate coordinate system and other internal
parameters"""
k = fftfreq(self.size, self.res)
kxx, kyy = np.meshgrid(k, k)
self._kr, self._phi = cart2pol(kyy, kxx)
# kmag is the radius of the spherical shell of the OTF
self._kmag = self.ni / self.wl
# because the OTF only exists on a spherical shell we can calculate
# a kz value for any pair of kx and ky values
self._kz = psqrt(self._kmag**2 - self._kr**2)
def set_t(self,optical_dephasing_rate,*,dt='auto'):
"""Sets the time grid upon which all frequency-detected signals will
be calculated on
"""
max_pos_t = int(self.gamma_res/optical_dephasing_rate)
if dt == 'auto':
dt = self.dts[-1] # signal detection bandwidth determined by local oscillator
self.t = np.arange(-max_pos_t,max_pos_t+dt/2,dt)
# if self.t.size % 2:
# self.t = self.t[:-1]
self.w = fftshift(fftfreq(self.t.size,d=dt)*2*np.pi)
def _update_fft_axes(axes, idx, shape, omitlast, ffunc):
for i in idx[:-1]:
__update_axes_label(axes, i)
axes[i].attrs['shift'] = axes[i].min # save lower bound. value of axes
axes[i].ax = fftmod.fftshift(fftmod.fftfreq(shape[i], d=axes[i].increment)) * 2 * np.pi
# the last dimension needs special care for rfft
i = idx[-1]
__update_axes_label(axes, i)
axes[i].attrs['shift'] = axes[i].min # save lower bound. value of axes
if omitlast:
axes[i].ax = ffunc(shape[i], d=axes[i].increment) * 2 * np.pi
else:
axes[i].ax = fftmod.fftshift(ffunc(shape[i], d=axes[i].increment)) * 2 * np.pi
def set_t(self, optical_dephasing_rate, *, dt='auto'):
"""Sets the time grid upon which all frequency-detected signals will
be calculated on
"""
max_pos_t = int(self.gamma_res / optical_dephasing_rate)
max_efield_t = max([np.max(u) for u in self.efield_times]) * 1.05
max_pos_t = max(max_pos_t, max_efield_t)
if dt == 'auto':
dt = self.dts[
-1] # signal detection bandwidth determined by local oscillator
n = int(max_pos_t / dt)
self.t = np.arange(-n, n + 1, 1) * dt
self.w = fftshift(fftfreq(self.t.size, d=dt) * 2 * np.pi)
def _gen_kr(self):
"""Internal utility function to generate internal state"""
# generate internal kspace coordinates
k = fftfreq(self.size, self.res)
kz = fftfreq(self.zsize, self.zres)
k_tot = np.meshgrid(kz, k, indexing="ij")
# calculate r
kr = norm(k_tot, axis=0)
# calculate the radius of the spherical shell in k-space
self.kmag = kmag = self.ni / self.wl
# determine k-space pixel size
dk, dkz = k[1] - k[0], kz[1] - kz[0]
# save output for user
self.dk, self.dkz = dk, dkz
# determine the min value for kz given the NA and wavelength
kz_min = np.sqrt(kmag**2 - (self.na / self.wl)**2)
# make sure we're not crazy
assert kz_min >= 0, "Something went horribly wrong"
# if the user gave us different z and x/y res we need to calculate
# the positional "error" in k-space to draw the spherical shell
if dk != dkz:
with np.errstate(invalid='ignore'):
dd = np.array((dkz, dk)).reshape(2, 1, 1)
dkr = norm(np.array(k_tot) * dd, axis=0) / kr
# we know the origin is zero so replace it
dkr[0, 0] = 0.0
else:
dkr = dk
if self.dual:
# if we want dual objectives we need two spherical shells
kzz = abs(k_tot[0])
else:
kzz = k_tot[0]
# calculate the points on the spherical shell, save them and the
# corresponding kz, ky and kx coordinates
self.valid_points = np.logical_and(
abs(kr - kmag) < dkr, kzz > kz_min + dkr)
self.kzz, self.krr = [k[self.valid_points] for k in k_tot]
def generate_white_noise(self):
"""Generate a white noise with the relevant power spectrum.
Returns
-------
white_noise : np.ndarray
"""
# Compute the k grid
d = self.Lbox / self.dimensions / (2 * np.pi)
all_k = [fft.fftfreq(self.dimensions, d=d)] * (self.Ndim - 1) + [
fft.rfftfreq(self.dimensions, d=d)
]
self.kgrid = kgrid = np.array(np.meshgrid(*all_k, indexing="ij"))
self.knorm = knorm = np.sqrt(np.sum(kgrid**2, axis=0))
# Compute Pk
Pk = np.zeros_like(knorm)
mask = knorm > 0
Pk[mask] = self.Pk(knorm[mask] * 2)
# Compute white noise (in Fourier space)
mu = np.random.standard_normal([self.dimensions] * self.Ndim)
muk = fft.rfftn(mu)
deltak = muk * np.sqrt(Pk)
# Compute field in real space
white_noise = fft.irfftn(deltak)
# Normalize variance
deltak_smoothed = deltak * self.filter.W(knorm)
field = fft.irfftn(deltak_smoothed)
std = field.std()
self.white_noise_fft = deltak * self.sigma8 / std
self.white_noise = white_noise * self.sigma8 / std
return self.white_noise
def _precompute(self):
# Computations that need only be done once per dataset.
self._nevts = self._v0.shape[0] # Number of events
# *NORMALIZE v0*
# Reshape to facilitate max_vz normalization using numpy broadcast rules.
v0 = np.moveaxis(self._v0, 0, -1)
# Normalize each event signal by the maximum z-component amplitude.
# We perform this succinctly using numpy multidimensional broadcasting rules.
self._max_vz = np.abs(v0[1, :, :]).max(axis=0)
self._max_vz.flags.writeable = False
v0 = v0 / self._max_vz
# Reshape back to original shape.
self._v0 = np.moveaxis(v0, -1, 0)
self._v0.flags.writeable = False
# Transform v0 to the spectral domain using real FFT
self._fv0 = fft(self._v0, axis=-1)
self._fv0.flags.writeable = False
# Compute discrete frequencies
self._w = 2 * np.pi * fftfreq(self._npts, self._dt)
self._w.flags.writeable = False
def set_efields(self,
times_list,
efields_list,
centers_list,
phase_discrimination,
*,
reset_psis=True,
plot_fields=False):
self.efield_times = times_list
self.efields = efields_list
self.centers = centers_list
self.set_phase_discrimination(phase_discrimination)
self.dts = []
self.efield_frequencies = []
if reset_psis:
self.psis = dict()
for t in times_list:
if t.size == 1:
dt = 1
w = np.array([0])
else:
dt = t[1] - t[0]
w = fftshift(fftfreq(t.size, d=dt)) * 2 * np.pi
self.dts.append(dt)
self.efield_frequencies.append(w)
self.heaviside_convolve_list.append(HeavisideConvolve(t.size))
self.dt = self.dts[0]
# Initialize unperturbed wavefunction
self.set_psi0(self.initial_state)
if self.detection_type == 'polarization' or 'integrated_polarization':
try:
self.local_oscillator = self.efields[-1].copy()
except:
self.local_oscillator = copy.deepcopy(self.efields[-1])
def fftFromLiteMap(liteMap, applySlepianTaper=False, nresForSlepian=3.0, threads=1):
"""
Create an fft2D object from a liteMap.
Parameters
----------
liteMap : liteMap.liteMap
The map object whose fft is being taken.
applySlepianTaper : bool, optional
If ``True``, apply the lowest order taper (to minimize edge-leakage).
Default is ``False``.
nresForSlepian : float, optional
If ``applySlepianTaper`` = ``True``, this specifies the resolution of
the taper to use. Default is 3.0.
threads : int, optional
Number of threads to use in pyFFTW calculations. Default is 1.
Returns
-------
ft : fftTools.fft2D
The fft2D object corresponding the input liteMap.
"""
ft = fft2D()
ft.Nx = liteMap.Nx
ft.Ny = liteMap.Ny
trace.issue("flipper.fftTools", 1, "Taking FFT of map with (Ny, Nx)= (%f, %f)" %(ft.Ny,ft.Nx))
ft.pixScaleX = liteMap.pixScaleX
ft.pixScaleY = liteMap.pixScaleY
lx = 2*numpy.pi*fftfreq(ft.Nx, d=ft.pixScaleX)
ly = 2*numpy.pi*fftfreq(ft.Ny, d=ft.pixScaleY)
ix = numpy.mod(numpy.arange(ft.Nx*ft.Ny), ft.Nx)
iy = numpy.arange(ft.Nx*ft.Ny)/ft.Nx
modLMap = numpy.zeros([ft.Ny, ft.Nx])
modLMap[iy,ix] = numpy.sqrt(lx[ix]**2 + ly[iy]**2)
ft.modLMap = modLMap
ft.lx = lx
ft.ly = ly
ft.ix = ix
ft.iy = iy
ft.thetaMap = numpy.zeros([ft.Ny, ft.Nx])
ft.thetaMap[iy[:], ix[:]] = numpy.arctan2(ly[iy[:]], lx[ix[:]])
ft.thetaMap *= 180./numpy.pi
mp = liteMap.data.copy()
taper = mp.copy()*0. + 1.0
if (applySlepianTaper):
try:
f = open(taperDir + os.path.sep + 'taper_Ny%d_Nx%d_Nres%3.1f' %(ft.Ny, ft.Nx, nresForSlepian))
taper = pickle.load(f)
f.close()
except:
taper = slepianTaper00(ft.Nx, ft.Ny,nresForSlepian)
f = open(taperDir + os.path.sep + 'taper_Ny%d_Nx%d_Nres%3.1f'%(ft.Ny, ft.Nx, nresForSlepian), mode="w")
pickle.dump(taper,f)
f.close()
if have_pyFFTW:
ft.kMap = fft2(mp*taper, threads=threads)
else:
ft.kMap = fft2(mp*taper)
del mp, modLMap, lx, ly
return ft
def create_filter(order,
cutoff,
nyquist,
N,
ftype='fir',
output='freq',
shift=True):
""" Create a lowpass FIR filter. This function is meant to create only the prototype filter,
where highpass, bandpass, or bandstop can all be transformed from the lowpass filter.
Parameters:
-----------
order: int
The filter order. The number of taps for an FIR filter.
cutoff: float
The cutoff frequency of the filter.
nyquist: int or float
This parameter is half the sampling rate.
ftype: str
Declare the filter to be an FIR ('fir') or an IIR ('iir').
output: str
Declare the return of the function to be in 'time' domain or 'freq' domain.
shift: bool
Declare if fftshift is applied to the FFT filter coeffiecients or not.
Returns:
--------
If output is 'freq'
w: ndarray
Frequencies corresponding to frequency components. The units is the
same as the ones chosen for the Nyquist rate.
H: ndarray, complex
The values of the FFT of the filter coefficients.
if output is 'time'
h: ndarray
The values of the filter coefficients
"""
if order > N:
raise ValueError(
"The order of the filter should not be longer than the length for FFT (binsize)."
)
if cutoff >= nyquist:
raise ValueError(
"The cutoff frequency must be at least 2 times smaller than the Nyquist rate."
)
if output not in ['time', 'freq']:
raise ValueError("'output' must be either 'time' or 'freq'!")
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
elif output == 'time':
return 1, h
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
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:
<<<<<<< HEAD
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()
=======
freq = (freq_in[:, np.newaxis] + tb * u.Hz *
rfftfreq(oversample*2, dtsample.value/2.)[::2])
# same as fine = rfftfreq(2*ntint, dtsample.value/2.)[::2]
#fcoh = freq_in[np.newaxis, :] + tb * u.Hz * rfftfreq(
# ntint*2, dtsample.value/2.)[::2, np.newaxis] #Old, incorrect by-channel
fcoh = freq_in - u.Hz * np.fft.fftfreq(ntint, dtsample.value)[:,np.newaxis]
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 == '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:
def fillWithGaussianRandomField(self,
ell,
Cell,
bufferFactor=1,
threads=1):
"""
Generate a Gaussian random field from an input power spectrum specified
as ell, Cell.
Notes
-----
BufferFactor = 1 means the map will have periodic boundary function, while
BufferFactor > 1 means the map will be genrated on a patch bufferFactor
times larger in each dimension and then cut out so as to have
non-periodic boundary conditions.
Fills the data field of the map with the GRF realization.
"""
ft = fftTools.fftFromLiteMap(self, threads=threads)
Ny = self.Ny * bufferFactor
Nx = self.Nx * bufferFactor
bufferFactor = int(bufferFactor)
realPart = numpy.zeros([Ny, Nx])
imgPart = numpy.zeros([Ny, Nx])
ly = fftfreq(Ny, d=self.pixScaleY) * (2 * numpy.pi)
lx = fftfreq(Nx, d=self.pixScaleX) * (2 * numpy.pi)
#print ly
modLMap = numpy.zeros([Ny, Nx])
iy, ix = numpy.mgrid[0:Ny, 0:Nx]
modLMap[iy, ix] = numpy.sqrt(ly[iy]**2 + lx[ix]**2)
s = splrep(ell, Cell, k=3)
ll = numpy.ravel(modLMap)
kk = splev(ll, s)
id = numpy.where(ll > ell.max())
kk[id] = 0.
#add a cosine ^2 falloff at the very end
# id2 = numpy.where( (ll> (ell.max()-500)) & (ll<ell.max()))
# lEnd = ll[id2]
# kk[id2] *= numpy.cos((lEnd-lEnd.min())/(lEnd.max() -lEnd.min())*numpy.pi/2)
#pylab.loglog(ll,kk)
area = Nx * Ny * self.pixScaleX * self.pixScaleY
p = numpy.reshape(kk, [Ny, Nx]) / area * (Nx * Ny)**2
realPart = numpy.sqrt(p) * numpy.random.randn(Ny, Nx)
imgPart = numpy.sqrt(p) * numpy.random.randn(Ny, Nx)
kMap = realPart + 1j * imgPart
if have_pyFFTW:
data = numpy.real(ifft2(kMap, threads=threads))
else:
data = numpy.real(ifft2(kMap))
b = bufferFactor
self.data = data[(b - 1) / 2 * self.Ny:(b + 1) / 2 * self.Ny,
(b - 1) / 2 * self.Nx:(b + 1) / 2 * self.Nx]
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 range(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
def fillWithGRFFromTemplate(self, twodPower, bufferFactor=1, threads=1):
"""
Generate a Gaussian random field from an input power spectrum
specified as a 2d powerMap
Notes
-----
BufferFactor = 1 means the map will have periodic boundary function, while
BufferFactor > 1 means the map will be genrated on a patch bufferFactor
times larger in each dimension and then cut out so as to have
non-periodic boundary conditions.
Fills the data field of the map with the GRF realization.
"""
ft = fftTools.fftFromLiteMap(self, threads=threads)
Ny = self.Ny * bufferFactor
Nx = self.Nx * bufferFactor
bufferFactor = int(bufferFactor)
assert (bufferFactor >= 1)
realPart = numpy.zeros([Ny, Nx])
imgPart = numpy.zeros([Ny, Nx])
ly = fftfreq(Ny, d=self.pixScaleY) * (2 * numpy.pi)
lx = fftfreq(Nx, d=self.pixScaleX) * (2 * numpy.pi)
#print ly
modLMap = numpy.zeros([Ny, Nx])
iy, ix = numpy.mgrid[0:Ny, 0:Nx]
modLMap[iy, ix] = numpy.sqrt(ly[iy]**2 + lx[ix]**2)
# divide out area factor
area = twodPower.Nx * twodPower.Ny * twodPower.pixScaleX * twodPower.pixScaleY
twodPower.powerMap *= (twodPower.Nx * twodPower.Ny)**2 / area
if bufferFactor > 1 or twodPower.Nx != Nx or twodPower.Ny != Ny:
lx_shifted = fftshift(twodPower.lx)
ly_shifted = fftshift(twodPower.ly)
twodPower_shifted = fftshift(twodPower.powerMap)
f_interp = interp2d(lx_shifted, ly_shifted, twodPower_shifted)
# ell = numpy.ravel(twodPower.modLMap)
# Cell = numpy.ravel(twodPower.powerMap)
# print ell
# print Cell
# s = splrep(ell,Cell,k=3)
#
#
# ll = numpy.ravel(modLMap)
# kk = splev(ll,s)
kk = f_interp(fftshift(lx), fftshift(ly))
kk = ifftshift(kk)
# id = numpy.where(modLMap > ell.max())
# kk[id] = 0.
# add a cosine ^2 falloff at the very end
# id2 = numpy.where( (ll> (ell.max()-500)) & (ll<ell.max()))
# lEnd = ll[id2]
# kk[id2] *= numpy.cos((lEnd-lEnd.min())/(lEnd.max() -lEnd.min())*numpy.pi/2)
# pylab.loglog(ll,kk)
area = Nx * Ny * self.pixScaleX * self.pixScaleY
#p = numpy.reshape(kk,[Ny,Nx]) /area * (Nx*Ny)**2
p = kk #/ area * (Nx*Ny)**2
else:
area = Nx * Ny * self.pixScaleX * self.pixScaleY
p = twodPower.powerMap #/area*(Nx*Ny)**2
realPart = numpy.sqrt(p) * numpy.random.randn(Ny, Nx)
imgPart = numpy.sqrt(p) * numpy.random.randn(Ny, Nx)
kMap = realPart + 1j * imgPart
if have_pyFFTW:
data = numpy.real(ifft2(kMap, threads=threads))
else:
data = numpy.real(ifft2(kMap))
b = bufferFactor
self.data = data[(b - 1) / 2 * self.Ny:(b + 1) / 2 * self.Ny,
(b - 1) / 2 * self.Nx:(b + 1) / 2 * self.Nx]
def fillWithGRFFromTemplate(self, twodPower, bufferFactor=1, threads=1):
"""
Generate a Gaussian random field from an input power spectrum
specified as a 2d powerMap
Notes
-----
BufferFactor = 1 means the map will have periodic boundary function, while
BufferFactor > 1 means the map will be genrated on a patch bufferFactor
times larger in each dimension and then cut out so as to have
non-periodic boundary conditions.
Fills the data field of the map with the GRF realization.
"""
ft = fftTools.fftFromLiteMap(self, threads=threads)
Ny = self.Ny * bufferFactor
Nx = self.Nx * bufferFactor
bufferFactor = int(bufferFactor)
assert bufferFactor >= 1
realPart = numpy.zeros([Ny, Nx])
imgPart = numpy.zeros([Ny, Nx])
ly = fftfreq(Ny, d=self.pixScaleY) * (2 * numpy.pi)
lx = fftfreq(Nx, d=self.pixScaleX) * (2 * numpy.pi)
# print ly
modLMap = numpy.zeros([Ny, Nx])
iy, ix = numpy.mgrid[0:Ny, 0:Nx]
modLMap[iy, ix] = numpy.sqrt(ly[iy] ** 2 + lx[ix] ** 2)
# divide out area factor
area = twodPower.Nx * twodPower.Ny * twodPower.pixScaleX * twodPower.pixScaleY
twodPower.powerMap *= (twodPower.Nx * twodPower.Ny) ** 2 / area
if bufferFactor > 1 or twodPower.Nx != Nx or twodPower.Ny != Ny:
lx_shifted = fftshift(twodPower.lx)
ly_shifted = fftshift(twodPower.ly)
twodPower_shifted = fftshift(twodPower.powerMap)
f_interp = interp2d(lx_shifted, ly_shifted, twodPower_shifted)
# ell = numpy.ravel(twodPower.modLMap)
# Cell = numpy.ravel(twodPower.powerMap)
# print ell
# print Cell
# s = splrep(ell,Cell,k=3)
#
#
# ll = numpy.ravel(modLMap)
# kk = splev(ll,s)
kk = f_interp(fftshift(lx), fftshift(ly))
kk = ifftshift(kk)
# id = numpy.where(modLMap > ell.max())
# kk[id] = 0.
# add a cosine ^2 falloff at the very end
# id2 = numpy.where( (ll> (ell.max()-500)) & (ll<ell.max()))
# lEnd = ll[id2]
# kk[id2] *= numpy.cos((lEnd-lEnd.min())/(lEnd.max() -lEnd.min())*numpy.pi/2)
# pylab.loglog(ll,kk)
area = Nx * Ny * self.pixScaleX * self.pixScaleY
# p = numpy.reshape(kk,[Ny,Nx]) /area * (Nx*Ny)**2
p = kk # / area * (Nx*Ny)**2
else:
area = Nx * Ny * self.pixScaleX * self.pixScaleY
p = twodPower.powerMap # /area*(Nx*Ny)**2
realPart = numpy.sqrt(p) * numpy.random.randn(Ny, Nx)
imgPart = numpy.sqrt(p) * numpy.random.randn(Ny, Nx)
kMap = realPart + 1j * imgPart
if have_pyFFTW:
data = numpy.real(ifft2(kMap, threads=threads))
else:
data = numpy.real(ifft2(kMap))
b = bufferFactor
self.data = data[(b - 1) / 2 * self.Ny : (b + 1) / 2 * self.Ny, (b - 1) / 2 * self.Nx : (b + 1) / 2 * self.Nx]
def initialize_nonparam_2d_nested_filter(X, gridres=1.0, **kwargs):
"""Function to compute the local Fourier filters using a nested approach.
Parameters
----------
X : array-like
Two-dimensional array containing the input field. All values are required
to be finite and the domain must be square.
gridres : float
Grid resolution in km.
Other Parameters
----------------
max_level : int
Localization parameter. 0: global noise, >0: increasing degree of localization.
Default : 3
win_type : string ['hanning', 'flat-hanning']
Type of window used for localization.
Default : flat-hanning
war_thr : float [0;1]
Threshold for the minimum fraction of rain needed for computing the FFT.
Default : 0.1
Returns
-------
F : array-like
Four-dimensional array containing the 2d fourier filters distributed over
a 2d spatial grid.
"""
if len(X.shape) != 2:
raise ValueError("X must be a two-dimensional array")
if X.shape[0] != X.shape[1]:
raise ValueError("a square array expected, but the shape of X is (%d,%d)" % \
(X.shape[0], X.shape[1]))
if np.any(np.isnan(X)):
raise ValueError("X must not contain NaNs")
# defaults
max_level = kwargs.get('max_level', 3)
win_type = kwargs.get('win_type', 'flat-hanning')
war_thr = kwargs.get('war_thr', 0.1)
# make sure non-rainy pixels are set to zero
min_value = np.min(X)
X = X.copy()
X -= min_value
#
dim = X.shape
dim_x = dim[1]
dim_y = dim[0]
# Nested algorithm
# prepare indices
Idxi = np.array([[0, dim_y]])
Idxj = np.array([[0, dim_x]])
Idxipsd = np.array([[0, 2**max_level]])
Idxjpsd = np.array([[0, 2**max_level]])
# generate the FFT sample frequencies
freq = fft.fftfreq(dim_y, gridres)
fx, fy = np.meshgrid(freq, freq)
freq_grid = np.sqrt(fx**2 + fy**2)
# domain fourier filter
F0 = initialize_nonparam_2d_fft_filter(X, win_type=win_type, donorm=True)
# and allocate it to the final grid
F = np.zeros((2**max_level, 2**max_level, F0.shape[0], F0.shape[1]))
F += F0[np.newaxis, np.newaxis, :, :]
# now loop levels and build composite spectra
level = 0
while level < max_level:
for m in range(len(Idxi)):
# the indices of rainfall field
Idxinext, Idxjnext = _split_field(Idxi[m, :], Idxj[m, :], 2)
# the indices of the field of fourier filters
Idxipsdnext, Idxjpsdnext = _split_field(Idxipsd[m, :],
Idxjpsd[m, :], 2)
for n in range(len(Idxinext)):
mask = _get_mask(dim, Idxinext[n, :], Idxjnext[n, :], win_type)
war = np.sum((X * mask) > 0.01) / float(
(Idxinext[n, 1] - Idxinext[n, 0])**2)
if war > war_thr:
# the new filter
newfilter = initialize_nonparam_2d_fft_filter(
X * mask, win_type=None, donorm=True)
# compute logistic function to define weights as function of frequency
# k controls the shape of the weighting function
# TODO: optimize parameters
k = 0.05
x0 = (Idxinext[n, 1] - Idxinext[n, 0]) / 2.
merge_weights = 1 / (1 + np.exp(-k * (1 / freq_grid - x0)))
newfilter *= (1 - merge_weights)
# perform the weighted average of previous and new fourier filters
F[Idxipsdnext[n, 0]:Idxipsdnext[n, 1],
Idxjpsdnext[n, 0]:Idxjpsdnext[
n, 1], :, :] *= merge_weights[np.newaxis,
np.newaxis, :, :]
F[Idxipsdnext[n, 0]:Idxipsdnext[n, 1],
Idxjpsdnext[n, 0]:Idxjpsdnext[n, 1], :, :] += newfilter[
np.newaxis, np.newaxis, :, :]
# update indices
level += 1
Idxi, Idxj = _split_field((0, dim[0]), (0, dim[1]), 2**level)
Idxipsd, Idxjpsd = _split_field((0, 2**max_level), (0, 2**max_level),
2**level)
return F
def plot_fixed_t32(self,
t32_time,
*,
part='real',
signal='rephasing',
ft=True,
savefig=True,
omega_0=1):
""""""
self.load_eigen_params()
t32_index, t32_time = self.get_closest_index_and_value(
t32_time, self.t32_array)
dt21 = self.t21_array[1] - self.t21_array[0]
dt43 = self.t43_array[1] - self.t43_array[0]
if signal == 'rephasing':
sig = self.rephasing_signal[:, t32_index, :]
elif signal == 'nonrephasing':
sig = self.nonrephasing_signal[:, t32_index, :]
if ft:
w21 = fftshift(fftfreq(self.t21_array.size, d=dt21)) * 2 * np.pi
w21 += self.ground_to_excited_transition + self.center
w21 *= omega_0
w43 = fftshift(fftfreq(self.t43_array.size, d=dt43)) * 2 * np.pi
w43 += self.ground_to_excited_transition + self.center
w43 *= omega_0
X, Y = np.meshgrid(w21, w43, indexing='ij')
if signal == 'nonrephasing':
ifft_t21_norm = self.t21_array.size * dt21
ifft_t43_norm = self.t43_array.size * dt43
sig = fftshift(ifftn(sig, axes=(0, 1)),
axes=(0, 1)) * ifft_t21_norm * ifft_t43_norm
elif signal == 'rephasing':
fft_t21_norm = dt21
ifft_t43_norm = self.t43_array.size * dt43
sig = fftshift(ifft(sig, axis=1), axes=(1)) * ifft_t43_norm
sig = fftshift(fft(sig, axis=0), axes=(0)) * fft_t21_norm
if omega_0 == 1:
xlab = '$\omega_{21}$ ($\omega_0$)'
ylab = '$\omega_{43}$ ($\omega_0$)'
else:
xlab = '$\omega_{21}$ (cm$^{-1}$)'
ylab = '$\omega_{43}$ (cm$^{-1}$)'
else:
X = self.T21[:, 0, :]
Y = self.T43[:, 0, :]
xlab = r'$t_{21}$ ($\omega_0^{-1}$)'
ylab = r'$t_{43}$ ($\omega_0^{-1}$)'
if part == 'real':
sig = sig.real
if part == 'imag':
sig = sig.imag
plt.figure()
plt.contour(X, Y, sig, 12, colors='k')
plt.contourf(X, Y, sig, 12)
plt.title(part + ' ' + signal + r' at $t_{32}$' +
' = {}'.format(t32_time))
plt.xlabel(xlab)
plt.ylabel(ylab)
plt.xlim([17000, 21000])
plt.ylim([17000, 21000])
plt.colorbar()
if savefig:
fig_name = os.path.join(
self.base_path,
part + '_' + signal + '_t_32_{}.png'.format(t32_time))
plt.savefig(fig_name)
plt.show()
def accelsearch(times,
signal,
delta_z=1,
fmin=1,
fmax=1e32,
gti=None,
zmax=100,
candidate_file=None,
ref_time=0,
debug=False,
interbin=False,
nproc=4,
det_p_value=0.15,
fft_rescale=None):
"""Find pulsars with accelerated search.
The theory behind these methods is described in Ransom+02, AJ 124, 1788.
Parameters
----------
times : array of floats
An evenly spaced list of times
signal : array of floats
The light curve, in counts; same length as ``times``
Other parameters
----------------
delta_z : float
The spacing in ``z`` space (delta_z = 1 -> delta_fdot = 1/T**2)
fmin : float, default 1.
Minimum frequency to search
fmax : float, default 1e32
Maximum frequency to search
gti : ``[[gti00, gti01], [gti10, gti11], ...]``, default None
Good Time Intervals. If None, it assumes the full range
``[[time[0] - dt / 2 -- time[-1] + dt / 2]]``
zmax : int, default 100
Maximum frequency derivative to search (pos and neg), in bins.
It corresponds to ``fdot_max = zmax / T**2``, where ``T`` is the
length of the observation.
candidate_file : str, default None
Save the final candidate table to this file. If None, the table
is just returned and not saved.
ref_time : float, default 0
Reference time for the times
det_p_value : float, default 0.015
Detection p-value (tail probability of noise powers, corrected for the
number of trials)
fft_rescale : function
Any function to apply to the initial FFT, normalized by the number of
photons as FT * np.sqrt(2/nph) so that || FT ||^2 are Leahy powers.
For example, a filter to flatten the spectrum in the presence of strong
red noise.
Returns
-------
candidate_table: :class:`Table`
Table containing the candidate frequencies and frequency derivatives,
the spectral power in Leahy normalization, the detection probability,
the time and the observation length.
"""
if not isinstance(times, np.ndarray):
times = np.asarray(times)
if not isinstance(signal, np.ndarray):
signal = np.asarray(signal)
dt = times[1] - times[0]
if gti is not None:
gti = np.asarray(gti)
# Fill in the data with a constant outside GTIs
gti_mask = create_gti_mask(times, gti)
expo_fraction = np.count_nonzero(gti_mask) / len(gti_mask)
bti_mask = ~gti_mask
mean_ops = np.mean
if np.mean(signal) > 10:
mean_ops = np.median
signal[bti_mask] = mean_ops(signal[gti_mask])
else:
expo_fraction = 1
gti = np.array([[times[0] - dt / 2, times[-1] + dt / 2]])
n_photons = np.sum(signal)
spectr = fft(signal) * np.sqrt(2 / n_photons)
freq = fftfreq(len(spectr), dt)
if debug:
_good_f = freq > 0
import matplotlib.pyplot as plt
fig = plt.figure(figsize=(12, 8))
plt.plot(freq[_good_f], (spectr * spectr.conj()).real[_good_f],
label='initial PDS')
plt.xlabel("Frequency (Hz)")
plt.ylabel("Power (Leahy)")
plt.loglog()
if fft_rescale is not None:
log.info("Applying initial filters...")
spectr = fft_rescale(spectr)
if debug:
plt.plot(freq[_good_f], (spectr * spectr.conj()).real[_good_f],
label='PDS after filtering (if any)')
fname = candidate_file + '_initial_spec.png' \
if candidate_file else 'initial_spec.png'
plt.legend(loc=2)
del _good_f
plt.savefig(fname)
plt.close(fig)
T = times[-1] - times[0] + dt
freq_intv_to_search = (freq >= fmin) & (freq < fmax)
log.info("Starting search over full plane...")
start_z = -zmax
end_z = zmax
range_z = np.arange(start_z, end_z, delta_z)
log.info("min and max possible r_dot: {}--{}".format(
delta_z / T**2,
np.max(range_z) / T**2))
freqs_to_search = freq[freq_intv_to_search]
candidate_table = Table(names=[
'time', 'length', 'frac_exposure', 'power', 'prob', 'frequency',
'fdot', 'fddot', 'ntrial'
],
dtype=[float] * 8 + [int])
detlev = pds_detection_level(ntrial=freqs_to_search.size,
epsilon=det_p_value)
responses = _create_responses(range_z)
candidate_rs, candidate_js, candidate_powers = \
_calculate_all_convolutions(spectr, responses,
freq_intv_to_search, detlev,
debug=debug, interbin=interbin,
nproc=nproc)
for r, j, cand_power in zip(candidate_rs, candidate_js, candidate_powers):
z = range_z[j]
cand_freq = r / T
fdot = z / T**2
prob = pds_probability(cand_power, freqs_to_search.size)
candidate_table.add_row([
ref_time + gti[0, 0], T, expo_fraction, cand_power, prob,
cand_freq, fdot, 0, freqs_to_search.size
])
if candidate_file is not None:
candidate_table.write(candidate_file + '.csv', overwrite=True)
return candidate_table
