示例#1
0
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]))
示例#2
0
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)
示例#3
0
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
示例#4
0
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
示例#5
0
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
示例#6
0
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
示例#7
0
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
示例#8
0
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)
示例#9
0
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
示例#10
0
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