def _findCarrier(self, band0, band1, mask):
band = band0 * band1
ixf = abs(fft.fftshift(fft.fft2(fft.fftshift(band))))
if self.debug:
plt.figure()
plt.title('Find carrier')
plt.imshow((ixf - gaussian_filter(ixf, 20)) * mask)
pyc0, pxc0 = self._findPeak((ixf - gaussian_filter(ixf, 20)) * mask)
ixfz, Kx, Ky = self._zoomf(band, self.N, self._kx[pyc0, pxc0],
self._ky[pyc0, pxc0], 50, self._dk * self.N)
pyc, pxc = self._findPeak(abs(ixfz))
if self.debug:
plt.figure()
plt.title('Zoon Find carrier')
plt.imshow(abs(ixfz))
kx = Kx[pxc]
ky = Ky[pyc]
otf_exclude_min_radius = 0.5
otf_exclude_max_radius = 1.5
kr = sqrt(self._kx**2 + self._ky**2)
m = (kr < 2)
otf = fft.fftshift(self._tfm(kr, m) + (1 - m))
otf_mask = (kr > otf_exclude_min_radius) & (kr <
otf_exclude_max_radius)
otf_mask_for_band_common_freq = fft.fftshift(
otf_mask & scipy.ndimage.shift(otf_mask, (pyc0 -
(self.N // 2 + 1), pxc0 -
(self.N // 2 + 1)),
order=0))
band0_common = fft.ifft2(
fft.fft2(band0) / otf * otf_mask_for_band_common_freq)
xx = np.arange(-self.N / 2 * self._dx,
self.N / 2 * self._dx,
self._dx,
dtype=np.single)
phase_shift_to_xpeak = exp(-1j * kx * xx * 2 * pi * self.NA /
self.wavelength)
phase_shift_to_ypeak = exp(-1j * ky * xx * 2 * pi * self.NA /
self.wavelength)
band1_common = fft.ifft2(
fft.fft2(band1) / otf * otf_mask_for_band_common_freq) * np.outer(
phase_shift_to_ypeak, phase_shift_to_xpeak)
scaling = 1 / np.sum(band0_common * np.conjugate(band0_common))
cross_corr_result = np.sum(band0_common * band1_common) * scaling
ampl = np.abs(cross_corr_result) * 2
phase = np.angle(cross_corr_result)
return kx, ky, phase, ampl
def __CPUErosion__(self, num_erosions=1):
temp = np.absolute(fftshift(self._support)).astype(bool)
eroded = binary_erosion(temp,
structure=self.BEStruct,
iterations=num_erosions)
self._support = fftshift(eroded.astype(complex))
return
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 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 retrieve_phase_far_field(src_fname,
save_path,
output_fname=None,
pad_length=256,
n_epoch=100,
learning_rate=0.001):
# raw data is assumed to be centered at zero frequency
prj_np = dxchange.read_tiff(src_fname)
if output_fname is None:
output_fname = os.path.basename(
os.path.splitext(src_fname)[0]) + '_recon'
# take modulus and inverse shift
prj_np = ifftshift(np.sqrt(prj_np))
obj_init = np.random.normal(50, 10, list(prj_np.shape) + [2])
obj = tf.Variable(obj_init, dtype=tf.float32, name='obj')
prj = tf.constant(prj_np, name='prj')
obj_real = tf.cast(obj[:, :, 0], dtype=tf.complex64)
obj_imag = tf.cast(obj[:, :, 1], dtype=tf.complex64)
# obj_pad = tf.pad(obj, [[pad_length, pad_length], [pad_length, pad_length], [0, 0]], mode='SYMMETRIC')
det = tf.fft2d(obj_real + 1j * obj_imag, name='detector_plane')
loss = tf.reduce_mean(tf.squared_difference(tf.abs(det), prj, name='loss'))
sess = tf.Session()
optimizer = tf.train.AdamOptimizer(learning_rate=learning_rate)
optimizer = optimizer.minimize(loss)
sess.run(tf.global_variables_initializer())
for i_epoch in range(n_epoch):
t0 = time.time()
_, current_loss = sess.run([optimizer, loss])
print('Iteration {}: loss = {}, Δt = {} s.'.format(
i_epoch, current_loss,
time.time() - t0))
det_final = sess.run(det)
obj_final = sess.run(obj)
res = np.linalg.norm(obj_final, 2, axis=2)
dxchange.write_tiff(res,
os.path.join(save_path, output_fname),
dtype='float32',
overwrite=True)
dxchange.write_tiff(fftshift(np.angle(det_final)),
os.path.join(save_path, 'detector_phase'),
dtype='float32',
overwrite=True)
dxchange.write_tiff(fftshift(np.abs(det_final)**2),
os.path.join(save_path, 'detector_mag'),
dtype='float32',
overwrite=True)
return
def _decoder(self,
x_encoded,
result='full',
axis=-1,
planner_effort='FFTW_ESTIMATE'):
"""
Transform the encoded data back to time series.
result = 'full': return to original modulating frequencies for each sub-band.
result = 'auto': return to the lowest modulating frequencies for each sub-band.
"""
Nch, Nf, Nsamp = x_encoded.shape
if result == 'full':
out = np.zeros((Nch, self.n_freqs, self.n_samples),
dtype=np.complex64)
out[:, self.decoder_rule, self.encoder_rule] = x_encoded
out = fft.fftshift(out, axes=axis)
return fft.ifft(out, axis=axis, planner_effort=planner_effort)
elif result == 'auto':
out = x_encoded
out = fft.fftshift(out, axes=axis)
out = np.roll(
out,
int(self.bandwidth // 2 *
(self.n_samples // self.sample_rate)))
return fft.ifft(out,
n=self.n_samples,
axis=axis,
planner_effort=planner_effort)
def __init__(self, mag, phase, mse, pupil_diff, mse_diff, model):
"""The results of retrieving a pupil function's phase and magnitude
Paramters
---------
mag : ndarray (n, n)
Coefficients for the zernike decomposition of the magnitude
phase : ndarray (n, n)
Coefficients for the zernike decomposition of the phase
mse : ndarray (m, )
Mean squared error as a function of the number of iterations (m)
performed
pupil_diff : ndarray (m, )
The relative change in the retrieved pupil function as a function
of the number of iterations (m) performed
mse_diff : ndarray (m, )
The relative change in the mean squared error as a function of the
number of iterations (m) performed
model : HanserPSF object
the model used to retrieve the pupil function
"""
# update internals
self.mag = mag
self.phase = phase
self.mse = mse
self.pupil_diff = pupil_diff
self.mse_diff = mse_diff
self.model = model
# calculate coordinate system
model._gen_kr()
r, theta = model._kr, model._phi
self.r, self.theta = fftshift(r), fftshift(theta)
# pull specific model parameters
self.na, self.wl = model.na, model.wl
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 fillPowerFromTemplate(self, twodPower):
"""
Fill the power2D.powerMap with the input power array.
Parameters
----------
twodPower : array_like
The 2D data array specifying the template power to fill with.
"""
tdp = twodPower.copy()
# interpolate
if tdp.Nx != self.Nx or tdp.Ny != self.Ny:
# first divide out the area factor
area = tdp.Nx*tdp.Ny*tdp.pixScaleX*tdp.pixScaleY
tdp.powerMap *= (tdp.Nx*tdp.Ny)**2 / area
lx_shifted = fftshift(tdp.lx)
ly_shifted = fftshift(tdp.ly)
tdp_shifted = fftshift(tdp.powerMap)
f_interp = interp2d(lx_shifted, ly_shifted, tdp_shifted)
cl_new = f_interp(fftshift(self.lx), fftshift(self.ly))
cl_new = ifftshift(cl_new)
area = self.Nx*self.Ny*self.pixScaleX*self.pixScaleY
cl_new *= area / (self.Nx*self.Ny*1.)**2
self.powerMap[:] = cl_new[:]
else:
self.powerMap[:] = tdp.powerMap[:]
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 polarization_to_signal(self,P_of_t_in,*,return_polarization=False,
local_oscillator_number = -1):
"""This function generates a frequency-resolved signal from a polarization field
local_oscillator_number - usually the local oscillator will be the last pulse in the list self.efields"""
pulse_time = self.pulse_times[local_oscillator_number]
if self.gamma != 0:
exp_factor = np.exp(-self.gamma * (self.t-pulse_time))
P_of_t_in *= exp_factor
P_of_t = P_of_t_in
if return_polarization:
return P_of_t
if local_oscillator_number == 'impulsive':
efield = np.exp(1j*self.w*(pulse_time))
else:
pulse_time_ind = np.argmin(np.abs(self.t - pulse_time))
pulse_start_ind = pulse_time_ind - self.size//2
pulse_end_ind = pulse_time_ind + self.size//2 + self.size%2
t_slice = slice(pulse_start_ind, pulse_end_ind,None)
efield = np.zeros(self.t.size,dtype='complex')
efield[t_slice] = self.efields[local_oscillator_number]
efield = fftshift(ifft(ifftshift(efield)))*len(P_of_t)*(self.t[1]-self.t[0])/np.sqrt(2*np.pi)
if P_of_t.size%2:
P_of_t = P_of_t[:-1]
efield = efield[:len(P_of_t)]
P_of_w = fftshift(ifft(ifftshift(P_of_t)))*len(P_of_t)*(self.t[1]-self.t[0])/np.sqrt(2*np.pi)
signal = np.imag(P_of_w * np.conjugate(efield))
return signal
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 removePhaseRamps( img ): #... by centering the Bragg peak in the array
fimg = fftshift( fftn( fftshift( img ) ) )
intens = np.absolute( fimg )**2
maxHere = np.where( intens==intens.max() )
for n in [ 0, 1, 2 ]:
fimg = np.roll( fimg, fimg.shape[n]//2-maxHere[n], axis=n )
imgout = fftshift( ifftn( fftshift( fimg ) ) )
return imgout
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 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 polarization_to_signal(self,
P_of_t_in,
*,
return_polarization=False,
local_oscillator_number=-1,
undersample_factor=1):
"""This function generates a frequency-resolved signal from a polarization field
local_oscillator_number - usually the local oscillator will be the last pulse
in the list self.efields"""
undersample_slice = slice(None, None, undersample_factor)
P_of_t = P_of_t_in[undersample_slice]
t = self.t[undersample_slice]
dt = t[1] - t[0]
pulse_time = self.pulse_times[local_oscillator_number]
if self.gamma != 0:
exp_factor = np.exp(-self.gamma * np.abs(t - pulse_time))
P_of_t *= exp_factor
if self.sigma_I != 0:
inhomogeneous = np.exp(-(t - pulse_time)**2 * self.sigma_I**2 / 2)
P_of_t *= inhomogeneous
if return_polarization:
return P_of_t
pulse_time_ind = np.argmin(np.abs(self.t - pulse_time))
efield = np.zeros(self.t.size, dtype='complex')
if self.efield_t.size == 1:
# Impulsive limit
efield[pulse_time_ind] = self.efields[local_oscillator_number]
efield = fftshift(ifft(ifftshift(efield))) * efield.size / np.sqrt(
2 * np.pi)
else:
pulse_start_ind = pulse_time_ind - self.size // 2
pulse_end_ind = pulse_time_ind + self.size // 2 + self.size % 2
t_slice = slice(pulse_start_ind, pulse_end_ind, None)
efield[t_slice] = self.efields[local_oscillator_number]
efield = fftshift(ifft(ifftshift(efield))) * self.t.size * (
self.t[1] - self.t[0]) / np.sqrt(2 * np.pi)
# if P_of_t.size%2:
# P_of_t = P_of_t[:-1]
# t = t[:-1]
halfway = self.w.size // 2
pm = self.w.size // (2 * undersample_factor)
efield_min_ind = halfway - pm
efield_max_ind = halfway + pm + self.w.size % 2
efield = efield[efield_min_ind:efield_max_ind]
P_of_w = fftshift(ifft(
ifftshift(P_of_t))) * len(P_of_t) * dt / np.sqrt(2 * np.pi)
signal = np.imag(P_of_w * np.conjugate(efield))
return signal
def _refineCarrier(self, band0, band1, kx_in, ky_in):
pxc0 = np.int(np.round(kx_in / self._dk) + self.N / 2)
pyc0 = np.int(np.round(ky_in / self._dk) + self.N / 2)
otf_exclude_min_radius = self.eta / 2
otf_exclude_max_radius = 1.5
m = (self._kr < 2)
otf = fft.fftshift(self._tfm(self._kr, m) + (1 - m) * 0.0001)
otf_mask = (self._kr > otf_exclude_min_radius) & (
self._kr < otf_exclude_max_radius)
otf_mask_for_band_common_freq = fft.fftshift(
otf_mask & scipy.ndimage.shift(otf_mask, (pyc0 -
(self.N // 2), pxc0 -
(self.N // 2)),
order=0))
band0_common = fft.ifft2(
fft.fft2(band0) / otf * otf_mask_for_band_common_freq)
band1_common = fft.ifft2(
fft.fft2(band1) / otf * otf_mask_for_band_common_freq)
band = band0_common * band1_common
mag = 25 * self.N / 256
ixfz, Kx, Ky = self._zoomf(band, self.N, np.single(self._k[pxc0]),
np.single(self._k[pyc0]), mag,
self._dk * self.N)
pyc, pxc = self._findPeak(abs(ixfz))
if self.debug:
plt.figure()
plt.title('Zoom Find carrier')
plt.imshow(abs(ixfz))
kx = Kx[pxc]
ky = Ky[pyc]
xx = np.arange(-self.N / 2 * self._dx,
self.N / 2 * self._dx,
self._dx,
dtype=np.double)
phase_shift_to_xpeak = exp(-1j * kx * xx * 2 * pi * self.NA /
self.wavelength)
phase_shift_to_ypeak = exp(-1j * ky * xx * 2 * pi * self.NA /
self.wavelength)
scaling = 1 / np.sum(band0_common * np.conjugate(band0_common))
cross_corr_result = np.sum(band0_common * band1_common * np.outer(
phase_shift_to_ypeak, phase_shift_to_xpeak)) * scaling
ampl = np.abs(cross_corr_result) * 2
phase = np.angle(cross_corr_result)
return kx, ky, phase, ampl
def _initializeSupport(self, sigma=0.575):
temp = np.log10(np.absolute(fftshift(fftn(self._modulus))))
mask = (temp > sigma * temp.max()).astype(float)
labeled, features = label(mask)
support_label = list(
dict(
sorted(collections.Counter(labeled.ravel()).items(),
key=lambda item: -item[1])).keys())[1]
self._support = np.zeros(self._arraySize)
self._support[np.where(labeled == support_label)] = 1.
self._support = fftshift(self._support)
# self.BinaryErosion( 1 )
return
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 _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 __init__( self,
modulus,
support=None, # if None, use default support
beta=0.9,
binning=1, # for high-energy CDI. Set to 1 for regular phase retrieval.
gpu=False,
pcc=False,
pcc_params=None, # user-defined coherence function parameters
random_start=True
):
self.BEStruct = np.ones( ( 3, 3, 3 ) ) # default structuring element for 3D binary erosion
self.BinaryErosion = self.__CPUErosion__
self._modulus = fftshift( modulus )
self._arraySize = tuple( this*shp for this, shp in zip( [ binning, binning, 1 ], self._modulus.shape ) )
if support is None:
self._initializeSupport()
else:
self._support = fftshift( support )
self._support_comp = 1. - self._support
self._beta = beta
if random_start:
self._cImage = np.exp( 2.j * np.pi * np.random.random_sample( self._arraySize ) ) * self._support
else:
self._cImage = 1. * self._support
self._cachedImage = np.zeros( self._cImage.shape ).astype( complex )
self._cImage_fft_mod = np.absolute( fftn( self._cImage ) )
self._error = []
self._UpdateError()
self.generateAlgoDict()
if gpu==True:
gpack = self.generateGPUPackage( pcc=pcc, pcc_params=pcc_params )
self.gpusolver = accelerator.Solver( gpack )
#if pcc==True:
# self._pccSolver = PCSolver( np.absolute( self._modulus )**2, gpack )
# self._kernel_f = self._pccSolver.getBlurKernel()
# self._ModProject = self._ModProjectPC
return
def _rapsd(X):
"""Compute radially averaged PSD of input field X.
"""
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]))
L = X.shape[0]
if L % 2 == 1:
XC,YC = np.ogrid[-int(L/2):int(L/2)+1, -int(L/2):int(L/2)+1]
else:
XC,YC = np.ogrid[-int(L/2):int(L/2), -int(L/2):int(L/2)]
R = np.sqrt(XC*XC + YC*YC).astype(int)
F = fft.fftshift(fft.fft2(X, **fft_kwargs))
F = abs(F)**2
if L % 2 == 0:
r_range = np.arange(0, int(L/2)+1)
else:
r_range = np.arange(0, int(L/2))
result = []
for r in r_range:
MASK = R == r
F_vals = F[MASK]
result.append(np.mean(F_vals))
return np.array(result)
def __init__(
self,
modulus,
support,
beta=0.9,
binning=1, # for high-energy CDI. Set to 1 for regular phase retrieval.
gpu=False,
random_start=True):
self._modulus = fftshift(modulus)
self._support = support
self._beta = beta
# self._modulus_sum = modulus.sum()
self._support_comp = 1. - support
if random_start:
self._cImage = np.exp(2.j * np.pi * np.random.rand(
binning * self._modulus.shape[0],
binning * self._modulus.shape[1],
self._modulus.shape[2])) * self._support
else:
self._cImage = 1. * support
self._cachedImage = np.zeros(self._cImage.shape).astype(complex)
self._cImage_fft_mod = np.absolute(fftn(self._cImage))
self._error = []
self._UpdateError()
self.generateAlgoDict()
if gpu == True:
self.gpusolver = accelerator.Solver(self.generateGPUPackage())
return
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 _rapsd(X):
"""Compute radially averaged PSD of input field X.
"""
if len(X.shape) != 2:
raise ValueError("%i dimensions are found, but the number of dimensions should be 2" % \
len(X.shape))
M, N = X.shape
YC, XC = _compute_centred_coord_array(M, N)
R = np.sqrt(XC * XC + YC * YC).astype(int)
F = fft.fftshift(fft.fft2(X, **fft_kwargs))
F = abs(F)**2
L = max(X.shape[0], X.shape[1])
if L % 2 == 0:
r_range = np.arange(0, int(L / 2) + 1)
else:
r_range = np.arange(0, int(L / 2))
result = []
for r in r_range:
MASK = R == r
F_vals = F[MASK]
result.append(np.mean(F_vals))
return np.array(result)
def ft_interface(a, s, axes, norm, **kwargs):
# call fft and shift the result
shftax = axes[:-1] if omitlast else axes
if forward:
return fftmod.fftshift(ftfunc(a, s, axes, norm, **kwargs), shftax)
else:
return ftfunc(fftmod.ifftshift(a, shftax), s, axes, norm, **kwargs)
def from_recip(y):
"""
Converts Fourier frequencies to spatial coordinates.
Parameters
----------
y : `list` [`numpy.ndarray` [`float`]], of shape [(nx,), (ny,), ...]
List (or equivalent) of vectors which define a mesh in the dimension
equal to the length of `x`
Returns
-------
x : `list` [`numpy.ndarray` [`float`]], of shape [(nx,), (ny,), ...]
List of vectors defining a mesh such that for a function, `f`, defined on
the mesh given by `y`, ifft(f) is defined on the mesh given by `x`. 0 will be
in the middle of `x`.
"""
x = []
for Y in y:
if Y.size > 1:
x.append(fftfreq(Y.size, Y.item(1) - Y.item(0)) * (2 * pi))
else:
x.append(array([0]))
x[-1] = x[-1].astype(Y.dtype, copy=False)
return [fftshift(X) for X in x]
def ImportCore(self, varDict):
self._modulus = tf.constant(varDict['modulus'], dtype=tf.complex64)
self._support = tf.Variable(varDict['support'], dtype=tf.complex64)
self._support_comp = tf.Variable(1. - varDict['support'],
dtype=tf.complex64)
self._beta = tf.constant(varDict['beta'], dtype=tf.complex64)
self._cImage = tf.Variable(varDict['cImage'], dtype=tf.complex64)
self._cachedImage = tf.Variable(np.zeros(varDict['cImage'].shape),
dtype=tf.complex64)
self._modulus_sum = tf.reduce_sum(self._modulus)
self._cImage_fft_mod = tf.Variable(
tf.abs(tf.signal.fft3d(self._cImage)))
self.BinaryErosion = self.__GPUErosion__
self._error = []
self._UpdateError()
x, y, z = np.meshgrid(
*[np.arange(-n // 2., n // 2.) for n in varDict['support'].shape])
self._rsquared = tf.constant(ftools.reduce(
lambda a, b: a + b, [fftshift(this)**2 for this in [x, y, z]]),
dtype=tf.complex64)
# used for GPU shrinkwrap
return
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 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 decomposition_fft(X, filter, **kwargs):
"""Decompose a 2d input field into multiple spatial scales by using the Fast
Fourier Transform (FFT) and a bandpass filter.
Parameters
----------
X : array_like
Two-dimensional array containing the input field. All values are required
to be finite.
filter : dict
A filter returned by any method implemented in bandpass_filters.py.
Other Parameters
----------------
MASK : array_like
Optional mask to use for computing the statistics for the cascade levels.
Pixels with MASK==False are excluded from the computations.
Returns
-------
out : ndarray
A dictionary described in the module documentation. The parameter n is
determined from the filter (see bandpass_filters.py).
"""
MASK = kwargs.get("MASK", None)
if len(X.shape) != 2:
raise ValueError("the input is not two-dimensional array")
if MASK is not None and MASK.shape != X.shape:
raise ValueError("dimension mismatch between X and MASK: X.shape=%s, MASK.shape=%s" % \
(str(X.shape), str(MASK.shape)))
if X.shape != filter["weights_2d"].shape[1:3]:
raise ValueError(
"dimension mismatch between X and filter: X.shape=%s, filter['weights_2d'].shape[1:3]=%s"
% (str(X.shape), str(filter["weights_2d"].shape[1:3])))
if np.any(~np.isfinite(X)):
raise ValueError("X contains non-finite values")
result = {}
means = []
stds = []
F = fft.fftshift(fft.fft2(X, **fft_kwargs))
X_decomp = []
for k in range(len(filter["weights_1d"])):
W_k = filter["weights_2d"][k, :, :]
X_ = np.real(fft.ifft2(fft.ifftshift(F * W_k), **fft_kwargs))
X_decomp.append(X_)
if MASK is not None:
X_ = X_[MASK]
means.append(np.mean(X_))
stds.append(np.std(X_))
result["cascade_levels"] = np.stack(X_decomp)
result["means"] = means
result["stds"] = stds
return result
def plot(self, log=False, title='', show=False, zoomUptoL=None):
"""
Plot the fft2D object as two images, one for the real part and
another for the imaginary part.
Parameters
----------
log : bool, optional
If ``True``, use a log-scale. Default is ``False``.
title : str, optional
The title to put on the plots.
show : bool, optional
If ``True``, show the plots, or otherwise, create a pylab object
without showing. Default is ``False``.
zoomUptoL float, optional
The multipole number L to zoom to on the 2D fft sub-space, so that
we show [-L, L] X [-L, L]. Default is ``None``.
"""
pReal = fftshift(numpy.real(self.kMap.copy()))
pImag = fftshift(numpy.imag(self.kMap.copy()))
if log:
pReal = numpy.log(numpy.abs(pReal))
pImag = numpy.log(numpy.abs(pImag))
im = pylab.matshow(pReal,origin="down",extent=[numpy.min(self.lx),numpy.max(self.lx),\
numpy.min(self.ly),numpy.max(self.ly)])
pylab.xlabel(r'$\ell_x$',fontsize=15)
pylab.ylabel(r'$\ell_y$',fontsize=15)
pylab.colorbar()
pylab.title(title + '(Real Part)',fontsize=8)
im2 = pylab.matshow(pReal,origin="down",extent=[numpy.min(self.lx),numpy.max(self.lx),\
numpy.min(self.ly),numpy.max(self.ly)])
pylab.xlabel(r'$\ell_x$',fontsize=15)
pylab.ylabel(r'$\ell_y$',fontsize=15)
pylab.colorbar()
pylab.title(title + '(Imaginary Part)',fontsize=8)
if zoomUptoL!=None:
im.axes.set_xlim(-zoomUptoL,zoomUptoL)
im.axes.set_ylim(-zoomUptoL,zoomUptoL)
im2.axes.set_xlim(-zoomUptoL,zoomUptoL)
im2.axes.set_ylim(-zoomUptoL,zoomUptoL)
if show:
pylab.show()
def writeFits(self,file,overWrite=False):
"""
23-10-2009: added by JB Juin
12-02-2009: Complete re-write to add WCS info (Sudeep)
@brief Write a power2D as a Fits file
"""
h = pyfits.Header()
h.update("COMMENT","flipper.power2D")
idx = numpy.where(fftshift(self.lx == 0))
idy = numpy.where(fftshift(self.ly == 0))
h.update('CTYPE1','ANG-FREQ')
h.update('CTYPE2','ANG-FREQ')
h.update("CRPIX1",idx[0][0]+1)
h.update("CRPIX2",idy[0][0]+1)
h.update("CRVAL1",0.0)
h.update("CRVAL2",0.0)
h.update("CDELT1",numpy.abs(self.lx[0]-self.lx[1]))
h.update("CDELT2",numpy.abs(self.ly[0]-self.ly[1]))
pyfits.writeto(file,fftshift(self.powerMap),header=h,clobber=overWrite)
def createKspaceMask(self, verticalStripe=None,slantStripeLxLy=None,\
slantStripeLxLy2=None,smoothingRadius=None,\
apodizeWithGaussFWHM=None):
"""
@brief Creates a mask in L(K)-space, with Stripes set to zero. Vertical stripes
are given by [-lx,lx], while slantStripes are specified by the intercepts on the
X, and Y axes.
"""
mask = self.powerMap.copy()
mask[:,:] = 1.
if verticalStripe!=None:
idx = numpy.where((self.lx<verticalStripe[1]) & (self.lx > verticalStripe[0]))
#print idx
mask[:,idx] = 0.
if slantStripeLxLy != None:
Lx = slantStripeLxLy[0]
Ly = slantStripeLxLy[1]
phi = numpy.arctan(1.0*Ly/Lx)
#print phi
perp = Lx*numpy.sin(phi)
perpMap = self.modLMap.copy()*numpy.cos(self.thetaMap*numpy.pi/180.+phi-numpy.pi/2.)
#pylab.imshow(perpMap)
#pylab.show()
idxx =numpy. where(numpy.abs(perpMap) < numpy.abs(perp))
mask[idxx] = 0.
if slantStripeLxLy2 != None:
Lx = slantStripeLxLy2[0]
Ly = slantStripeLxLy2[1]
phi = numpy.arctan(1.0*Ly/Lx)
#print phi
perp = Lx*numpy.sin(phi)
perpMap = self.modLMap.copy()*numpy.cos(self.thetaMap*numpy.pi/180.+phi-numpy.pi/2.)
#pylab.imshow(perpMap)
#pylab.show()
idxxx =numpy. where(numpy.abs(perpMap) < numpy.abs(perp))
mask[idxxx] = 0.
if smoothingRadius!=None:
mask = fftshift(blur_image(fftshift(mask),smoothingRadius))
if apodizeWithGaussFWHM !=None:
mask *=numpy.exp(-self.modLMap**2*apodizeWithGaussFWHM**2/(8*numpy.log(2)))
self.kMask = mask
def _phase_correlation(im0, im1, callback=None, *args):
"""
Computes phase correlation between im0 and im1
Args:
im0
im1
callback (function): Process the cross-power spectrum (i.e. choose
coordinates of the best element, usually of the highest one).
Defaults to :func:`imreg_dft.utils.argmax2D`
Returns:
tuple: The translation vector (Y, X). Translation vector of (0, 0)
means that the two images match.
"""
if callback is None:
callback = utils._argmax2D
# TODO: Implement some form of high-pass filtering of PHASE correlation
f0, f1 = [fft.fft2(arr) for arr in (im0, im1)]
# spectrum can be filtered (already),
# so we have to take precaution against dividing by 0
eps = abs(f1).max() * 1e-15
# cps == cross-power spectrum of im0 and im1
cps = abs(fft.ifft2((f0 * f1.conjugate()) / (abs(f0) * abs(f1) + eps)))
# scps = shifted cps
scps = fft.fftshift(cps)
(t0, t1), success = callback(scps, *args)
ret = np.array((t0, t1))
# _compensate_fftshift is not appropriate here, this is OK.
t0 -= f0.shape[0] // 2
t1 -= f0.shape[1] // 2
ret -= np.array(f0.shape, int) // 2
return ret, success
def pattern_params(my_pat, size=2):
"""Find stuff"""
# REAL FFT!
# note the limited shifting, we don't want to shift the last axis
my_pat_fft = fftshift(rfftn(ifftshift(my_pat)),
axes=tuple(range(my_pat.ndim))[:-1])
my_abs_pat_fft = abs(my_pat_fft)
# find dc loc, center of FFT after shifting
sizeky, sizekx = my_abs_pat_fft.shape
# remember we didn't shift the last axis!
dc_loc = (sizeky // 2, 0)
# mask data and find next biggest peak
dc_power = my_abs_pat_fft[dc_loc]
my_abs_pat_fft[dc_loc] = 0
max_loc = np.unravel_index(my_abs_pat_fft.argmax(), my_abs_pat_fft.shape)
# pull the 3x3 region around the peak and fit
max_shift = localize_peak(my_abs_pat_fft[slice_maker(max_loc, 3)])
# calculate precise peak relative to dc
peak = np.array(max_loc) + np.array(max_shift) - np.array(dc_loc)
# correct location based on initial data shape
peak_corr = peak / np.array(my_pat.shape)
# calc angle
preciseangle = np.arctan2(*peak_corr)
# calc period
precise_period = 1 / norm(peak_corr)
# calc phase
phase = np.angle(my_pat_fft[max_loc[0], max_loc[1]])
# calc modulation depth
numerator = abs(my_pat_fft[slice_maker(max_loc, size)].sum())
mod = numerator / dc_power
return {"period": precise_period,
"angle": preciseangle,
"phase": phase,
"fft": my_pat_fft,
"mod": mod,
"max_loc": max_loc}
def writeFits(self, file, overWrite=False):
"""
Write a fftTools.fft2D objects as a FITS file
Parameters
----------
file : str
The name of the FITS file to write.
overWrite : bool, optional
Whether to overwrite any exisiting files with the desired name.
Default is ``False``.
"""
h = pyfits.Header()
h.update("COMMENT","flipper.fft2D")
idx = numpy.where(fftshift(self.lx == 0))
idy = numpy.where(fftshift(self.ly == 0))
h.update('CTYPE1','ANG-FREQ')
h.update('CTYPE2','ANG-FREQ')
h.update("CRPIX1",idx[0][0]+1)
h.update("CRPIX2",idy[0][0]+1)
h.update("CRVAL1",0.0)
h.update("CRVAL2",0.0)
h.update("CDELT1",numpy.abs(self.lx[0]-self.lx[1]))
h.update("CDELT2",numpy.abs(self.ly[0]-self.ly[1]))
realFile = file.split('.')[0]+'_real.fits'
pyfits.writeto(realFile,fftshift(numpy.real(self.kMap)),header=h,clobber=overWrite)
del h
h = pyfits.Header()
h.update("COMMENT","flipper.fft2D")
idx = numpy.where(fftshift(self.lx == 0))
idy = numpy.where(fftshift(self.ly == 0))
h.update('CTYPE1','ANG-FREQ')
h.update('CTYPE2','ANG-FREQ')
h.update("CRPIX1",idx[0][0]+1)
h.update("CRPIX2",idy[0][0]+1)
h.update("CRVAL1",0.0)
h.update("CRVAL2",0.0)
h.update("CDELT1",numpy.abs(self.lx[0]-self.lx[1]))
h.update("CDELT2",numpy.abs(self.ly[0]-self.ly[1]))
realFile = file.split('.')[0]+'_imag.fits'
pyfits.writeto(realFile,fftshift(numpy.imag(self.kMap)),header=h,clobber=overWrite)
def mapFromFFT(self, kFilter=None, kFilterFromList=None, showFilter=False,
setMeanToZero=False, returnFFT=False, threads=1):
"""
Perform the inverse fft (map from FFT) with an optional filter.
Parameters
----------
kFilter : array_like, optional
2D array specifying a k-space filter to apply. If applied, resulting
map is IFFT(fft*kFilter). Default is ``None``.
kFilterFromList : tuple, optional
Tuple of length 2 specifying 1D filter (ell, F_ell), which will be
interpolated into a 2D array. Default is ``None``.
showFilter : bool, optional
Whether to plot the filter. Default is ``False``.
setMeanToZero : bool, optional
Whether to set the ell = 0 pixel to zero (zeroes mean in real space).
Default is ``False``.
returnFFT : bool, optional
Whether to return the fftTools.fft2D class as well. Default is
``False``.
threads : int, optional
Number of threads to use in pyFFTW calculations. Default is 1.
Returns
-------
data : array_like
The (optinally filtered) 2D real space data array
ftMap : fftTools.fft2D, optional
The fft2D object, returned if returnFFT = ``True``.
"""
kMap = self.kMap.copy()
kFilter0 = numpy.real(kMap.copy())*0.+ 1.
if kFilter != None:
kFilter0 *= kFilter
if kFilterFromList != None:
kFilter = kMap.copy()*0.
l = kFilterFromList[0]
Fl = kFilterFromList[1]
FlSpline = splrep(l,Fl,k=3)
ll = numpy.ravel(self.modLMap)
kk = (splev(ll,FlSpline))
kFilter = numpy.reshape(kk,[self.Ny,self.Nx])
kFilter0 *= kFilter
if setMeanToZero:
id = numpy.where(self.modLMap == 0.)
kFilter0[id] = 0.
if showFilter:
pylab.semilogy(l,Fl,'r',ll,kk,'b.')
pylab.matshow(fftshift(kFilter0),origin="down",extent=[numpy.min(self.lx),\
numpy.max(self.lx),\
numpy.min(self.ly),\
numpy.max(self.ly)])
pylab.show()
kMap[:,:] *= kFilter0[:,:]
if have_pyFFTW:
data = numpy.real(ifft2(kMap, threads=threads))
else:
data = numpy.real(ifft2(kMap))
if returnFFT:
ftMap = self.copy()
ftMap.kMap = kMap.copy()
return data, ftMap
else:
return data
dy = domain / ny y_vals=np.linspace(1,ny,num=ny,endpoint=True)*domain/ny x_vals=np.linspace(1,nx,num=nx,endpoint=True)*domain/nx x_arr,y_arr=np.meshgrid(x_vals,y_vals) dt = 0.4 * 16.0 / nx # choose an initial dt. This will change # as the simulation progresses to maintain # numerical stability ## Spectral Domain dk = 2.0*pi/domain; k = np.arange(-n/2, n/2)*dk l = np.arange(-n/2, n/2)*dk kk, ll = [fftshift(q) for q in np.meshgrid(k, l)] # put in FFT order ksq = kk**2 + ll**2 ksq[ksq == 0] = 1.0 # avoid divide by zero - set ksq = 1 at zero wavenum rksq = 1.0 / ksq # reciprocal 1/(k^2 + l^2) ik = 1j*kk il = 1j*ll ## Dissipation & Anti-Aliasing nu = (((domain)/(np.floor(n/3)*2.0*pi))**4)/tau del4 = 1.0 / (1.0 + nu*ksq**2*dt) # dissipate at small scales k_max = AA_FAC*2*dk # anti-aliasing removes wavenumbers > k_max # from the non-linear term
def plot(self,log=False,colorbar=True,title='',powerOfL=0,pngFile=None,show=True,zoomUptoL=None,\
showMask=False, yrange = None, showBinsFromFile = None,drawCirclesAtL=None,\
drawVerticalLinesAtL = None, valueRange=None, colorbarLabel=None):
"""
@brief Display the power spectrum
"""
#modLMap = self.modLMap
#modLMap[numpy.where(modLMap ==0)] = 1.
p = self.powerMap.copy()
p[:] *= (self.modLMap[:]+1.)**powerOfL
p = fftshift(p)
if showBinsFromFile:
binLower,binUpper,binCenter= readBinningFile(showBinsFromFile)
theta = numpy.arange(0,2.*numpy.pi+0.05,0.05)
for i in xrange(len(binLower)):
x,y = binUpper[i]*numpy.cos(theta),binUpper[i]*numpy.sin(theta)
pylab.plot(x,y,'k')
if drawCirclesAtL !=None:
for ell in drawCirclesAtL:
theta = numpy.arange(0,2.*numpy.pi+0.05,0.05)
x,y = ell*numpy.cos(theta),ell*numpy.sin(theta)
pylab.plot(x,y,'k')
if len(drawCirclesAtL)<5:
pylab.text(ell*numpy.cos(numpy.pi/4.),ell*numpy.sin(numpy.pi/4.),\
'%d'%numpy.int(ell),rotation=-45,horizontalalignment = 'center',\
verticalalignment='bottom',fontsize=8)
if drawVerticalLinesAtL!=None:
for ell in drawVerticalLinesAtL:
pylab.axvline(ell)
if log:
p = numpy.log10(numpy.abs(p))
if yrange != None:
p[numpy.where(p < yrange[0])] = yrange[0]
p[numpy.where(p > yrange[1])] = yrange[1]
vmin = p.min()
vmax = p.max()
if valueRange != None:
vmin = valueRange[0]
vmax = valueRange[1]
im = pylab.imshow(p,origin="down",extent=[numpy.min(self.lx),numpy.max(self.lx),\
numpy.min(self.ly),numpy.max(self.ly)],aspect='equal',vmin=vmin,vmax=vmax, interpolation='nearest')
pylab.title(title,fontsize=8)
if colorbar:
cb=pylab.colorbar()
if colorbarLabel != None:
cb.set_label(colorbarLabel)
pylab.xlabel(r'$\ell_x$',fontsize=15)
pylab.ylabel(r'$\ell_y$',fontsize=15)
if showMask:
im2 = pylab.imshow(fftshift(self.kMask.copy()),\
origin="down",\
extent=[numpy.min(self.lx),\
numpy.max(self.lx),\
numpy.min(self.ly),\
numpy.max(self.ly)],\
aspect='equal')
pylab.xlabel(r'$\ell_x$',fontsize=15)
pylab.ylabel(r'$\ell_y$',fontsize=15)
if colorbar:
pylab.colorbar()
if zoomUptoL!=None:
im.axes.set_xlim(-zoomUptoL,zoomUptoL)
im.axes.set_ylim(-zoomUptoL,zoomUptoL)
if showMask:
im2.axes.set_xlim(-zoomUptoL,zoomUptoL)
im2.axes.set_ylim(-zoomUptoL,zoomUptoL)
if show:
pylab.show()
if pngFile!=None:
pylab.savefig(pngFile)
def upgradePixelPitch(m, N=1, threads=1):
"""
Go to finer pixels with fourier interpolation.
Parameters
----------
m : liteMap
The liteMap object holding the data to upgrade the pixel size of.
N : int, optional
Go to 2^N times smaller pixels. Default is 1.
threads : int, optional
Number of threads to use in pyFFTW calculations. Default is 1.
Returns
-------
mNew : liteMap
The map with smaller pixels.
"""
if N < 1:
return m.copy()
Ny = m.Ny * 2 ** N
Nx = m.Nx * 2 ** N
npix = Ny * Nx
if have_pyFFTW:
ft = fft2(m.data, threads=threads)
else:
ft = fft2(m.data)
ftShifted = fftshift(ft)
newFtShifted = numpy.zeros((Ny, Nx), dtype=numpy.complex128)
# From the numpy.fft.fftshift help:
# """
# Shift zero-frequency component to center of spectrum.
#
# This function swaps half-spaces for all axes listed (defaults to all).
# If len(x) is even then the Nyquist component is y[0].
# """
#
# So in the case that we have an odd dimension in our map, we want to put
# the extra zero at the beginning
if m.Nx % 2 != 0:
offsetX = (Nx - m.Nx) / 2 + 1
else:
offsetX = (Nx - m.Nx) / 2
if m.Ny % 2 != 0:
offsetY = (Ny - m.Ny) / 2 + 1
else:
offsetY = (Ny - m.Ny) / 2
newFtShifted[offsetY : offsetY + m.Ny, offsetX : offsetX + m.Nx] = ftShifted
del ftShifted
ftNew = ifftshift(newFtShifted)
del newFtShifted
# Finally, deconvolve by the pixel window
mPix = numpy.copy(numpy.real(ftNew))
mPix[:] = 0.0
mPix[
mPix.shape[0] / 2 - (2 ** (N - 1)) : mPix.shape[0] / 2 + (2 ** (N - 1)),
mPix.shape[1] / 2 - (2 ** (N - 1)) : mPix.shape[1] / 2 + (2 ** (N - 1)),
] = (1.0 / (2.0 ** N) ** 2)
if have_pyFFTW:
ftPix = fft2(mPix, threads=threads)
else:
ftPix = fft2(mPix)
del mPix
inds = numpy.where(ftNew != 0)
ftNew[inds] /= numpy.abs(ftPix[inds])
if have_pyFFTW:
newData = ifft2(ftNew, threads=threads) * (2 ** N) ** 2
else:
newData = ifft2(ftNew) * (2 ** N) ** 2
del ftNew
del ftPix
x0_new, y0_new = m.pixToSky(0, 0)
m = m.copy() # don't overwrite original
m.wcs.header.update("NAXIS1", 2 ** N * m.wcs.header["NAXIS1"])
m.wcs.header.update("NAXIS2", 2 ** N * m.wcs.header["NAXIS2"])
m.wcs.header.update("CDELT1", m.wcs.header["CDELT1"] / 2.0 ** N)
m.wcs.header.update("CDELT2", m.wcs.header["CDELT2"] / 2.0 ** N)
m.wcs.updateFromHeader()
p_x, p_y = m.skyToPix(x0_new, y0_new)
m.wcs.header.update("CRPIX1", m.wcs.header["CRPIX1"] - p_x)
m.wcs.header.update("CRPIX2", m.wcs.header["CRPIX2"] - p_y)
m.wcs.updateFromHeader()
mNew = liteMapFromDataAndWCS(numpy.real(newData), m.wcs)
mNew.data[:] = numpy.real(newData[:])
return mNew
def similarity(im0, im1):
"""Return similarity transformed image im1 and transformation parameters.
Transformation parameters are: isotropic scale factor, rotation angle (in
degrees), and translation vector.
A similarity transformation is an affine transformation with isotropic
scale and without shear.
Limitations:
Image shapes must be equal and square.
All image areas must have same scale, rotation, and shift.
Scale change must be less than 1.8.
No subpixel precision.
"""
if im0.shape != im1.shape:
raise ValueError("Images must have same shapes.")
elif len(im0.shape) != 2:
raise ValueError("Images must be 2 dimensional.")
f0 = fftshift(abs(fft2(im0)))
f1 = fftshift(abs(fft2(im1)))
h = highpass(f0.shape)
f0 *= h
f1 *= h
del h
f0, log_base = logpolar(f0)
f1, log_base = logpolar(f1)
f0 = fft2(f0)
f1 = fft2(f1)
r0 = abs(f0) * abs(f1)
ir = abs(ifft2((f0 * f1.conjugate()) / r0))
i0, i1 = np.unravel_index(np.argmax(ir), ir.shape)
angle = 180.0 * i0 / ir.shape[0]
scale = log_base ** i1
if scale > 1.8:
ir = abs(ifft2((f1 * f0.conjugate()) / r0))
i0, i1 = np.unravel_index(np.argmax(ir), ir.shape)
angle = -180.0 * i0 / ir.shape[0]
scale = 1.0 / (log_base ** i1)
if scale > 1.8:
raise ValueError("Images are not compatible. Scale change > 1.8")
if angle < -90.0:
angle += 180.0
elif angle > 90.0:
angle -= 180.0
im2 = ndii.zoom(im1, 1.0 / scale)
im2 = ndii.rotate(im2, angle)
if im2.shape < im0.shape:
t = np.zeros_like(im0)
t[:im2.shape[0], :im2.shape[1]] = im2
im2 = t
elif im2.shape > im0.shape:
im2 = im2[:im0.shape[0], :im0.shape[1]]
f0 = fft2(im0)
f1 = fft2(im2)
ir = abs(ifft2((f0 * f1.conjugate()) / (abs(f0) * abs(f1))))
t0, t1 = np.unravel_index(np.argmax(ir), ir.shape)
if t0 > f0.shape[0] // 2:
t0 -= f0.shape[0]
if t1 > f0.shape[1] // 2:
t1 -= f0.shape[1]
im2 = ndii.shift(im2, [t0, t1])
# correct parameters for ndimage's internal processing
if angle > 0.0:
d = int((int(im1.shape[1] / scale) * math.sin(math.radians(angle))))
t0, t1 = t1, d + t0
elif angle < 0.0:
d = int((int(im1.shape[0] / scale) * math.sin(math.radians(angle))))
t0, t1 = d + t1, d + t0
scale = (im1.shape[1] - 1) / (int(im1.shape[1] / scale) - 1)
return im2, scale, angle, [-t0, -t1]
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 process_frames(self, data):
sino = fft.fftshift(self.fft_object(data[0]))
sino[self.row1:self.row2] = \
sino[self.row1:self.row2] * self.filtercomplex
sino = fft.ifftshift(sino)
return self.ifft_object(sino).real
def _get_ang_scale(ims, bgval, exponent='inf', constraints=None, reports=None):
"""
Given two images, return their scale and angle difference.
Args:
ims (2-tuple-like of 2D ndarrays): The images
bgval: We also pad here in the :func:`map_coordinates`
exponent (float or 'inf'): The exponent stuff, see :func:`similarity`
constraints (dict, optional)
reports (optional)
Returns:
tuple: Scale, angle. Describes the relationship of
the subject image to the first one.
"""
assert len(ims) == 2, \
"Only two images are supported as input"
shape = ims[0].shape
ims_apod = [utils._apodize(im) for im in ims]
dfts = [fft.fftshift(fft.fft2(im)) for im in ims_apod]
filt = _logpolar_filter(shape)
dfts = [dft * filt for dft in dfts]
# High-pass filtering used to be here, but we have moved it to a higher
# level interface
pcorr_shape = _get_pcorr_shape(shape)
log_base = _get_log_base(shape, pcorr_shape[1])
stuffs = [_logpolar(np.abs(dft), pcorr_shape, log_base)
for dft in dfts]
(arg_ang, arg_rad), success = _phase_correlation(
stuffs[0], stuffs[1],
utils.argmax_angscale, log_base, exponent, constraints, reports)
angle = -np.pi * arg_ang / float(pcorr_shape[0])
angle = np.rad2deg(angle)
angle = utils.wrap_angle(angle, 360)
scale = log_base ** arg_rad
angle = - angle
scale = 1.0 / scale
if reports is not None:
reports["shape"] = filt.shape
reports["base"] = log_base
if reports.show("spectra"):
reports["dfts_filt"] = dfts
if reports.show("inputs"):
reports["ims_filt"] = [fft.ifft2(np.fft.ifftshift(dft))
for dft in dfts]
if reports.show("logpolar"):
reports["logpolars"] = stuffs
if reports.show("scale_angle"):
reports["amas-result-raw"] = (arg_ang, arg_rad)
reports["amas-result"] = (scale, angle)
reports["amas-success"] = success
extent_el = pcorr_shape[1] / 2.0
reports["amas-extent"] = (
log_base ** (-extent_el), log_base ** extent_el,
-90, 90
)
if not 0.5 < scale < 2:
raise ValueError(
"Images are not compatible. Scale change %g too big to be true."
% scale)
return scale, angle
