Generates some FFT figures for an SE question.

FFT Figure generation

"""
Generate a figure to demonstrate the effect of frequency mirroring in a given rotating frame.
"""

from matplotlib.pyplot import plot, figure, gca, text, tight_layout, savefig, subplot, show
from matplotlib.pyplot import xlim, ylim, xticks, yticks, tick_params, title
from matplotlib.patches import FancyArrowPatch
import numpy as np
from numpy import pi, exp, double
from numpy.fft import fft, rfft, fftshift, fftfreq
import re, os

def fftx(data, n=None, axis=-1):
    """
    Power preserving Fourier transform

    Scales the Fourier transform by 2/length_of_data, which preserves the spectrum power.
    """

    spec = fft(data, n, axis)
    nn = np.shape(data)[axis]

    return 2*spec/nn

def rfftx(data, n=None, axis=-1):
    """
    Power preserving Fourier transform

    Scales the Fourier transform by 2/length_of_data, which preserves the spectrum power.
    """

    spec = rfft(data, n, axis)
    nn = np.shape(data)[axis]

    return 2*spec/nn


def save_figs(loc, formats='pdf', **kwargs):
    """
    Saves the file at 'loc' with all the formats iterated in "formats"
    """
    # Check if the path exists and make it if it doesn't.
    basepath = '/'.join(re.split(r'[\\/]',loc)[0:-1])
    if len(basepath) > 0 and not os.path.exists(basepath):
        os.makedirs(basepath)

    # Can pass a string or a list/tuple of strings
    if isinstance(formats, str):
        formats = [formats]

    for cformat in formats:
        savefig(loc+'.'+cformat, format=cformat, **kwargs)

omega = 0.15
omega2 = omega*2*pi
num = 2**20
apod = 250      # Apodization factor

rc = '#D20202'
gc = '#007935'
pc = '#702AA2'
bc = '#2A48A2'
grc= '#A29F96'

t = np.linspace(0, 5000, num)
sig_pos = np.exp(1j*omega2*t)*np.exp(-t/apod)
sig_neg = np.exp(-1j*omega2*t)*np.exp(-t/apod)
np_fft = num

f = fftfreq(np_fft, double(np.diff(t[0:2])))
sr = 1/double(np.diff(t[0:2]))
f = np.linspace(-sr/2, sr/2, np_fft)
s_pos = fftshift(abs(fftx(sig_pos, np_fft)))
s_neg = fftshift(abs(fftx(sig_neg, np_fft)))

figure(1, figsize=(8, 4), dpi=120)

plot([omega, omega], [-1, 2], '-', color=grc, lw=0.8)
plot([-omega, -omega], [-1, 2], '--', color=grc, lw=0.8)
plot(np.where(f >= -omega/10, f, None), np.where(f >= -omega/10, s_pos, None), color=bc, lw=1.2)
plot(np.where(f <= omega/10, f, None), np.where(f <= omega/10, s_neg, None), color=gc, lw=1.2)

max_height = max(s_pos)*1.2

# Arrow
ah = np.mean([max(s_pos), max_height])
arrow_ms = 20;          arrow_lw = 1.5
arrow_kwargs = {'arrowstyle':'<|-|>',
                'fc':'k',
                'mutation_scale':arrow_ms,
                'transform':gca().transData,
                'zorder':-1,
                'lw':arrow_lw}

txtkwargs = {
    'fontsize':16,
    'va':'top',
    'ha':'center',
    'transform':gca().transData
}

del_arrow = FancyArrowPatch((-omega, ah), (omega, ah), **arrow_kwargs)
text(0, ah*0.98, '$2\omega_{0}$', **txtkwargs)
gca().add_patch(del_arrow)

xlim([x*1.25 for x in (-omega, omega)])
ylim([-0.0001, max_height])

xticks([-omega, 0, omega])
yticks([])
gca().xaxis.set_ticklabels(['$-\omega_{0}$', '0', '$\omega_{0}$'])

tick_params(top=False)

tight_layout()
save_figs('images/FrequencyMirror', formats='png', transparent=True)

figure(2, figsize=(8, 4), dpi=120)

s_1chan_pos = fftshift(abs(fftx(np.real(sig_pos), np_fft)))
s_1chan_neg = fftshift(abs(fftx(np.real(sig_neg), np_fft)))
max_height = max(s_1chan_pos)*1.05

subplot(2, 1, 1)
px, = plot(f, s_1chan_pos, color=bc, lw=1.2)

xlim([x*1.25 for x in (-omega, omega)])
ylim([-0.0001, max_height])

xticks([])
yticks([])
title('$\\omega_{0}$', fontsize=10)

subplot(2, 1, 2)
py, = plot(f, s_1chan_neg, color=gc, lw=1.2)
title('$-\\omega_{0}$', fontsize=10)


xlim([x*1.25 for x in (-omega, omega)])
ylim([-0.0001, max_height])

xticks([-omega, 0, omega])
yticks([])
gca().xaxis.set_ticklabels(['$-\omega_{0}$', '0', '$\omega_{0}$'])

tick_params(top=False)

save_figs('images/FrequencyMirrorOneChannelAcquisition', formats='png', transparent=True)

tight_layout()
show()