def fft_2D(R, convert=3e-5, freq_bounds=None):
"""
Perform a 2D FFT to transform a 3rd order response function defined in the
rotating frame into a series of 2D spectra.
First argument must be a MetaArray object wth ticks and rw_freq defined.
Returns the FFT in another MetaArray object with ticks updated to the
calculated frequencies (converted using the convert argument which defaults
to cm to fs).
"""
# reverses frequency order to keep convention e^{+i \omega t}
R_2D = fftshift(fft2(ifftshift(R, axes=(0, 2)), axes=(0, 2)),
axes=(0, 2))[::-1, :, ::-1]
dt = [t[1] - t[0] for t in R.ticks]
freqs = [
fftshift(fftfreq(R.shape[axis], dt[axis] * convert)) + R.rw_freq
for axis in (0, 2)
]
if freq_bounds is not None:
bounds = [bound_indices(ticks, freq_bounds) for ticks in freqs]
freqs = [freq[bound[0]:bound[1]] for freq, bound in zip(freqs, bounds)]
R_2D = R_2D[bounds[0][0]:bounds[0][1], :, bounds[1][0]:bounds[1][1]]
return MetaArray(R_2D, ticks=(freqs[0], R.ticks[1], freqs[1]))
def trim_pulse_overlap(R, E_all=((1, ), (1, ), (1, ))):
# Not quite right?
start = len(E_all[0]) + len(E_all[1])
end = None
R_new = R[:, start:end, :]
ticks_new = (R.ticks[0], R.ticks[1][start:end], R.ticks[2])
return MetaArray(R_new, ticks=ticks_new, rw_freq=R.rw_freq)
def R3_to_P3(R, E_all=((1, ), (1, ), (1, )), trim=True, include_margin=True):
E3, E2, E1 = E_all
P3 = shift_time_indices(
convolve_third(convolve_second(convolve_first(R, E1), E2), E3), E_all,
trim, include_margin)
try:
return MetaArray(P3, ticks=expand_ticks(R, E_all), rw_freq=R.rw_freq)
except AttributeError:
return P3
def R3_to_P3_alt(R,
E_all=((1, ), (1, ), (1, )),
trim=True,
include_margin=True):
conv_mat, shapes, slice_maps = convolution_parameters(R.shape, E_all)
P3 = shift_time_indices(
reduce(loop_matvec, zip(conv_mat, shapes[1:], slice_maps), R), E_all,
trim, include_margin)
try:
return MetaArray(P3, ticks=expand_ticks(R, E_all), rw_freq=R.rw_freq)
except AttributeError:
return P3
def R3_resize(R3, E_all):
E3, E2, E1 = E_all
R3_new = np.zeros(resized_shape(R3.shape, E_all), dtype=R3.dtype)
R3_new[-(R3.shape[0] - E3.center):, :, :] = \
R3[:end(E3.center),
(E2.center + E3.center):end(E2.center),
E1.center:end(E1.center)]
try:
return MetaArray(R3_new,
ticks=resize_ticks(R3.ticks, E_all),
rw_freq=R3.rw_freq)
except AttributeError:
return R3_new
def R3_to_P3(R3, E_all, shortcut=False):
shapes, slice_maps = convolution_parameters(R3.shape, E_all)
S3 = reduce(loop_matvec, zip(shapes[1:], slice_maps), R3)
if shortcut:
return S3
P3 = P3_resize(shift_time_indices(S3, E_all), E_all)
try:
return MetaArray(P3,
ticks=resize_ticks(R3.ticks, E_all),
rw_freq=R3.rw_freq)
except AttributeError:
return P3
def R3_add_margin(R, E_all=((1, ), (1, ), (1, ))):
tau = np.array([int((len(E) - 1) / 2.0) for E in E_all])
tau3, tau2, tau1 = tau
L3, L2, L1 = R.shape
R_new = np.zeros(tuple(np.array(R.shape) + tau), dtype=R.dtype)
R_new[tau3:(L3 + tau3), tau2:(L2 + tau2), tau1:(L1 + tau1)] = R
try:
return MetaArray(R_new,
ticks=expand_ticks(R, E_all),
rw_freq=R.rw_freq)
except AttributeError:
return R_new
def combine_NR_PE(R_NR, R_PE):
x3, x2, x1 = R_NR.shape
y3, y2, y1 = R_PE.shape
if (x3 != y3) or (x2 != y2):
raise ValueError('Mismatched t3 or t2')
R_c = np.zeros((x3, x2, x1 + y1 - 1), dtype=R_NR.dtype)
R_c[:, :, :(x1 - 1)] = R_NR[:, :, :0:-1]
R_c[:, :, x1:] = R_PE[:, :, 1:]
R_c[:, :, (x1 - 1)] = (R_PE[:, :, 0] + R_NR[:, :, 0]) / 2.0
new_ticks = (R_NR.ticks[0], R_NR.ticks[1],
np.append(-R_NR.ticks[2][:0:-1], R_PE.ticks[2]))
return MetaArray(R_c, ticks=new_ticks, rw_freq=R_NR.rw_freq)
def P3_to_R3_opt_damp(P3,
E_all=((1, ), (1, ), (1, )),
damp0=None,
trim=True,
include_margin=True,
minimizer=scipy.optimize.minimize,
**min_kwargs):
conv_mat, shapes, slice_maps = convolution_parameters(P3.shape, E_all)
P3 = reduce(
invert_loop_matvec(minimizer, **min_kwargs),
zip(conv_mat[::-1], shapes[-2::-1],
[[(y, x) for (x, y) in sm] for sm in slice_maps[::-1]], damp0),
shift_time_indices_undo(P3, E_all, trim, include_margin))
try:
return MetaArray(P3,
ticks=unexpand_ticks(P3, E_all),
rw_freq=P3.rw_freq)
except AttributeError:
return P3
def P3_to_R3_fixed_damp(P3,
E_all=((1, ), (1, ), (1, )),
damp=None,
trim=True,
include_margin=True):
conv_mat, shapes, slice_maps = convolution_parameters(P3.shape, E_all)
deconv_mat = [
Ainv for (Ainv, _) in generalized_tikhonov_inverse(A, d)
for A, d in zip(conv_mat, damp)
]
P3 = reduce(
loop_matvec,
zip(deconv_mat[::-1], shapes[-2::-1],
[[(y, x) for (x, y) in slice_map]
for slice_map in slice_maps[::-1]]),
shift_time_indices_undo(P3, E_all, trim, include_margin))
try:
return MetaArray(P3,
ticks=unexpand_ticks(P3, E_all),
rw_freq=P3.rw_freq)
except AttributeError:
return P3