EAPSK63 waterfall plotter

waterfally.py

#!/usr/bin/env python3

import numpy as np
from scipy.fftpack import fft, fftshift
from scipy.signal import blackman
import matplotlib.pyplot as plt

import sys

# Parameters

# Good to crop a 1920x1080 image of whole contest
#N = 4096
#averaging = 1077

# Banner 1920x1080 (to crop)
N = 1024
averaging = 1077*4

# High resolution
#N = 16384
#averaging = 29 # 5.12s

samples = np.memmap('/home/daniel/psk63.wav', dtype=np.int16, mode='r', offset=0x28)

total_transforms = (len(samples)//2)//(N//2) - 1
lines = total_transforms//averaging

window = blackman(N)

try:
    data = np.load('data.npz')
    waterfall = data['waterfall']
    lines_done = data['lines_done']
except FileNotFoundError:
    waterfall = np.zeros((lines, N), dtype=np.float32)
    lines_done = 0

done = lines_done == lines
print('Loaded {}/{} lines already computed'.format(lines_done, lines))

try:
    for line in range(lines_done, lines):
        print('Computing line {}/{}'.format(line + 1, lines))
        transforms = np.zeros((averaging, N), dtype=np.float32)
        for transform in range(averaging):
            start = (line * averaging + transform)*N
            x = (np.float32(samples[start:start+2*N:2]) + 1j*np.float32(samples[start + 1:start+1+2*N:2]))/np.iinfo(np.int16).max
            f = fftshift(fft(x * window))
            transforms[transform,:] = np.real(f)**2 + np.imag(f)**2
        waterfall[line,:] = 10*np.log10(np.average(transforms, axis=0)) - 20*np.log10(N)
        lines_done += 1
    if not done:
        np.savez('data.npz', waterfall=waterfall, lines_done=lines_done)
except KeyboardInterrupt:
    if not done:
        print('Keyboard Interrupt. Saving work')
        np.savez('data.npz', waterfall=waterfall, lines_done=lines_done)
    sys.exit()

print('Plotting PNG images')
chunk = 5000
q, r = divmod(waterfall.T.shape[1], chunk)
chunks = q + int(bool(r))
for i in range(chunks):
    print('Plotting image {}/{}'.format(i+1, chunks))
    plt.imsave('waterfall{}.png'.format(i), waterfall.T[:,i*chunk:(i+1)*chunk], vmin = -130, vmax=-80, cmap='viridis', origin='bottom')