interactive waveform generation and mixing using pygame and pyaudio

view_wave.py

import pygame
import pygame_gui
import math
import pyaudio
import struct

SCREEN_X = 800
SCREEN_Y = 600

SCOPE_DIMS = (SCREEN_X//4, SCREEN_Y // 4)

SCOPE_Y_SCALE = SCOPE_DIMS[1]
SCOPE_X_SCALE = 2

SAMPLE_RATE = 44100
BUFFER_LEN = 1024

def generate_wave(freq, sample_rate, length, offset = 0, scale = 0.1):
    #eliminate popping by offsetting the wave by a period
    offset = offset % (sample_rate / freq)
    for i in range(length):
        yield scale * math.sin(math.tau * freq * (i + offset) / sample_rate)
    
def mix_wave(wave_a, wave_b, ratio = 0.5):
    fac_a = 0 + ratio
    fac_b = 1 - ratio
    for i, sample in enumerate(wave_a):
        yield (sample * fac_a) + (wave_b[i] * fac_b)


def main():
    pygame.init()
    pygame.display.set_caption("Wave")
    clock = pygame.time.Clock()

    manager = pygame_gui.UIManager((SCREEN_X, SCREEN_Y))
    button_1 = pygame_gui.elements.UIButton(relative_rect=pygame.Rect((400, 280), (100, 50)), text='Stop',manager=manager)

    background = pygame.Surface((SCREEN_X, SCREEN_Y))
    background.fill(pygame.Color('#000000'))
    window_surface = pygame.display.set_mode((SCREEN_X, SCREEN_Y))

    wavsurfs = [pygame.Surface(SCOPE_DIMS), pygame.Surface(SCOPE_DIMS), pygame.Surface(SCOPE_DIMS)]

    mixfac = 0.5
    wavoff = [0]
    
    waves = [[], [], []]

    def gen_waves(offset = True):
        waves[0] = [*generate_wave(440, SAMPLE_RATE, BUFFER_LEN, offset=wavoff[0])]
        waves[1] = [*generate_wave(587.33 * 2, SAMPLE_RATE, BUFFER_LEN, offset=wavoff[0])]
        waves[2] = [*mix_wave(waves[0], waves[1], mixfac)]
        if offset:
            wavoff[0] = wavoff[0] + BUFFER_LEN

    gen_waves(offset=False)
    
    def audio_callback(in_data, frame_count, time_info, status):
        gen_waves()
        data = waves[2][:frame_count]
        data = struct.pack("%df" % frame_count, *data)
        return (data, pyaudio.paContinue)

    au = pyaudio.PyAudio()
    stream = au.open(format=pyaudio.paFloat32,
                    channels=1,
                    rate=SAMPLE_RATE,
                    frames_per_buffer=BUFFER_LEN,
                    stream_callback=audio_callback,
                    output=True,)

    while True:
        time_delta = clock.tick(60)/1000.0

        for event in pygame.event.get():
            if event.type == pygame.QUIT:
                stream.stop_stream()
                stream.close()
                au.terminate()
                pygame.quit()
                return
            
            if event.type == pygame.MOUSEWHEEL:
                if event.y > 0:
                    mixfac = min(1, mixfac + 0.05)
                    dirty = True
                elif event.y < 0:
                    mixfac = max(0, mixfac - 0.05)
                    dirty = True
            
            if event.type == pygame_gui.UI_BUTTON_PRESSED:
              if event.ui_element == button_1:
                    if stream.is_active():
                        stream.stop_stream()
                        button_1.set_text('Start')
                    else:
                        stream.start_stream()
                        button_1.set_text('Stop')

            manager.process_events(event)
        
        window_surface.blit(background, (0, 0))

        for surf in wavsurfs:
            surf.fill(pygame.Color('#AAAAAA'))

        pygame.draw.aalines(wavsurfs[0], (255, 50, 90), False, [(i/SCOPE_X_SCALE, SCOPE_Y_SCALE * (sample + 0.5)) for i, sample in enumerate(waves[0])])
        pygame.draw.aalines(wavsurfs[1], (30, 50, 255), False, [(i/SCOPE_X_SCALE, SCOPE_Y_SCALE * (sample + 0.5)) for i, sample in enumerate(waves[1])])
        pygame.draw.aalines(wavsurfs[2], (255, 50, 255), False, [(i/SCOPE_X_SCALE, SCOPE_Y_SCALE * (sample + 0.5)) for i, sample in enumerate(waves[2])])

        window_surface.blit(wavsurfs[0], (80, 20))
        window_surface.blit(wavsurfs[1], (80, 40 + SCOPE_DIMS[1]))
        window_surface.blit(wavsurfs[2], (100 + SCOPE_DIMS[0], 30 + int(SCOPE_DIMS[1]/2)))

        manager.update(time_delta)
        manager.draw_ui(window_surface)
        
        pygame.display.update()


if __name__ == "__main__":
    main()

Video Preview: gfycat.com/portlyvigilanteidolonhelvum