CSES - FFT-spektri

Vaatii pygame ja pyaudio -kirjastot.

Hyödyllisiä Pythonin ominaisuuksia:

import cmath
imaginaariyksikkö = 1j
eulerin_kaava = cmath.exp(1j * 2 * cmath.pi * ...)
import pygame
import pyaudio
import math
import cmath
import struct

def fft(x):
    # TODO: toteuta fft
    pass

def isclose(a, b): return cmath.isclose(a, b, rel_tol=1e-4, abs_tol=1e-4)
def allclose(a, b): return all(isclose(x, y) for x, y in zip(a, b))

# Testejä taulukon koolle N=4
w = cmath.exp(2j * cmath.pi / 4)
tulos = fft([1, 1, 1, 1])
assert allclose(tulos, [4, 0, 0, 0]), tulos
tulos = fft([1, w, w**2, w**3])
assert allclose(tulos, [0, 4, 0, 0]), tulos
tulos = fft([1, w**2, w**4, w**6])
assert allclose(tulos, [0, 0, 4, 0]), tulos
tulos = fft([1, w**3, w**6, w**9])
assert allclose(tulos, [0, 0, 0, 4]), tulos

pygame.init()
screen_width = 800
screen_height = 600
screen = pygame.display.set_mode((screen_width, screen_height))
pygame.display.set_caption('FFT-spektri')
clock = pygame.time.Clock()

BLACK = (0, 0, 0)
WHITE = (255, 255, 255)

FORMAT = pyaudio.paInt16
CHANNELS = 1
RATE = 44100
CHUNK = 1024 # Kerrallaan analysoitavat näytteet

p = pyaudio.PyAudio()

stream = p.open(format=FORMAT,
                channels=CHANNELS,
                rate=RATE,
                input=True,
                frames_per_buffer=CHUNK)

hann_window = [0.5 - 0.5 * math.cos(2 * math.pi * n / CHUNK) for n in range(CHUNK)]

def draw_spectrum(spectrum, log_freq, screen_width, screen_height):
    screen.fill(BLACK)
    N = len(spectrum)
    for i, amp in enumerate(spectrum):
        # Muutetaan amplitudi desibeleiksi
        if amp <= 1e-8: dB = -100
        else: dB = 20. * math.log10(amp)
        # Skaalaus: Näytöllä näkyy -60dB..10dB
        y_pos = (dB + 60.) / 70. * screen_height

        if log_freq:
            x_pos = math.log10(i+1) / math.log10(N) * screen_width
        else:
            x_pos = i / N * screen_width

        pygame.draw.line(screen, WHITE, (x_pos, screen_height), (x_pos, screen_height - y_pos), 1)

log_freq = False
running = True
while running:
    for event in pygame.event.get():
        if event.type == pygame.QUIT:
            running = False
        elif event.type == pygame.KEYDOWN:
            if event.key == pygame.K_l:
                # Paina "L" niin taajuudet näytetään logaritmiasteikolla
                log_freq = not log_freq

    # Skipataan ylimääräinen data
    avail = stream.get_read_available()
    if avail > CHUNK:
        stream.read(avail-CHUNK)
    # Luetaan CHUNK verran näytteitä
    data = stream.read(CHUNK)
    samples = struct.unpack(str(CHUNK) + 'h', data)
    # Muutetaan välille -1..1
    samples = [sample / (1<<16) for sample in samples]
    # Ikkunointi (vältetään epäjatkuvuuskohta reunoilla)
    windowed_samples = [s * w for s, w in zip(samples, hann_window)]
    # Kutsutaan FFT-toteutusta ja käytetään vain "positiiviset" taajuudet
    fft_result = fft(windowed_samples)[:CHUNK//2]
    amplitude_spectrum = [abs(v) / math.sqrt(CHUNK) for v in fft_result]

    draw_spectrum(amplitude_spectrum, log_freq, screen_width, screen_height)
    pygame.display.flip()
    clock.tick(30)

pygame.quit()
stream.stop_stream()
stream.close()
p.terminate()