Python .WAV generator

Question

The following is a .wav file generator that I wrote briefly while (re)learning about audio sampling:

import math

# Generates the .wav file header for a given set of samples and specs
def genHeader(sampleRate, bitsPerSample, channels, samples):
    datasize = len(samples) * channels * bitsPerSample // 8
    o = bytes("RIFF",'ascii')                                               # (4byte) Marks file as RIFF
    o += (datasize + 36).to_bytes(4,'little')                               # (4byte) File size in bytes excluding this and RIFF marker
    o += bytes("WAVE",'ascii')                                              # (4byte) File type
    o += bytes("fmt ",'ascii')                                              # (4byte) Format Chunk Marker
    o += (16).to_bytes(4,'little')                                          # (4byte) Length of above format data
    o += (1).to_bytes(2,'little')                                           # (2byte) Format type (1 - PCM)
    o += (channels).to_bytes(2,'little')                                    # (2byte)
    o += (sampleRate).to_bytes(4,'little')                                  # (4byte)
    o += (sampleRate * channels * bitsPerSample // 8).to_bytes(4,'little')  # (4byte)
    o += (channels * bitsPerSample // 8).to_bytes(2,'little')               # (2byte)
    o += (bitsPerSample).to_bytes(2,'little')                               # (2byte)
    o += bytes("data",'ascii')                                              # (4byte) Data Chunk Marker
    o += (datasize).to_bytes(4,'little')                                    # (4byte) Data size in bytes
    return o

# Expects samples in the form [[channel1sample0,channel2sample0,...],[channel1sample1,channel2sample1,...]]
# Samples are values between -1 and 1 then upscaled to the correct amplitude (-2**15 -> 2**15 roughly) for 16 bit
def genFile(sampleRate, bitsPerSample, channels, samples):
    with open("test.wav","wb") as f:
        f.write(genHeader(sampleRate, bitsPerSample, channels, samples))
        m = 2**(bitsPerSample-1)-1
        for i in range(len(samples)):
            for j in range(channels):
                f.write(int(samples[i][j] * m).to_bytes(bitsPerSample//8,'little',signed=True))

# Generates samples for a given frequency, sample rate and duration
# c is the current amplitude of the wave between -1 and 1
def genSamples(freq, sampleRate, duration, wave, c):
    if wave == 'square':
        c, samples = genSamples(freq,sampleRate,duration,'sin')
        for i in range(len(samples)):
            if(samples[i] > 0):
                samples[i] = 1
            elif(samples[i] < 0):
                samples[i] = -1
            else:
                samples[i] = 0
    elif wave == 'sin':
        sampleCount = int(sampleRate * duration)
        waves = duration * freq
        inc = waves/sampleCount * 360
        samples = []
        for i in range(sampleCount):
            samples += [math.sin(math.radians(c))]
            c += inc
    else:
        print("NOT IMPLEMENTED")
        exit()
    return c, samples

# This is just a temporary method until multiple channels is implemented
def prepSamples(samples):
    samplesOut = []
    for i in range(len(samples)):
        samplesOut += [[samples[i]]]
    return samplesOut

# Generate samples from 20Hz to 20kHz
# Use c to have waves of 1 frequency continue from where previous
# frequency left of
samples = []
c = 0
for i in range(20,20001):
    c, r = genSamples(i, 44100, 0.003003003003003, "sin", c)
    samples += r
samples = prepSamples(samples)
genFile(44100, 16, 1, samples)

From this code review I'm most interested in comments w.r.t the genSamples function, the program as a whole seems to work well for frequencies between 20Hz and 18000Hz although there are some strange distortions after that as my sin wave becomes distorted (normally I wouldn't be able to hear 18kHz+ but as the the distortions are clearly audible). Any suggestions on how to improve the wave generation functionality would be greatly appreciated.

Note

I am aware of PEP8 and I am less interested in the styling of the code as this was just a quick write up to test what I'd been reading and to try out a couple of things, the code will be tidied when I more fully understand what I'm doing and how I want to continue. Until then I'd appreciate it if just the functionality of the code were to be critiqued.

Also note that the current setup of the code takes some time to run and any performance gains that are possible would also be appreciated.

Answer 1

Docstrings

Your comments about each function are essentially docstrings. They should be properly made as such, by putting them in triple quotes just after the function definitions.

def genSamples(freq, sampleRate, duration, wave, c):
    """Generates samples for a given frequency, sample rate and duration
    c is the current amplitude of the wave between -1 and 1"""

This means they can appear by typing help(function_name). IDEs or other tools may present them in various contexts.

Iteration

Prefer for sample in samples: over for i in range(len(samples)):

Consider your API design

We should strive for organizing code into independent, logical units (functions, classes, ...). Your functions are often interdependent in ways that make the code harder to grasp.

genSamples returns a tuple, which is rather unexpected at first. The value c seems intimately connected to the frequency sweep functionality. It would be clearer to let genSamples return only samples and place the sweep generation in another function which uses it.

Then we have genFile, which expects data in a list nested per-sample, and prepSamples is interlinked with this. Why this format? Probably because of how the Wav file is laid out on disk. It would make more sense to me to handle any multichannel data as one list per channel, and let genFile be responsible for transforming it to the file format. (If you need multichannel at all; otherwise this might be a good place to apply the YAGNI principle.)

Several functions take a lot of arguments which means they take more effort to use. One way to help matters is to put the "input" parameter first and provide default values for (some of) the rest (e.g. 16 bits, 44100 rate).

genHeader doesn't need to take the samples as a parameter since it depends only on the number of them. Also, it doesn't make much sense to use on its own. You might want to nest it in genFile or even merge it.

Use list comprehensions to build lists

The use of list comprehensions (and/or generator expressions) to build lists usually leads to more readable and faster code, at least when you're used to them. Consider this version of the sample generation function.

Since the comprehensions clearly have no side effects and don't depend on previous iterations of the internal loop, I think it's easier to see what is going on here. Such as what phase will be at each point.

def sine_signal(freq, rate, duration, phi=0):
    """Generates samples for a given frequency, sample rate and duration"""
    sample_count = int(rate * duration)
    k = 2*math.pi * freq * duration / sample_count
    phase = (i * k + phi for i in range(sample_count))
    signal = [math.sin(p) for p in phase]
    new_phase = 2*math.pi * freq * duration + phi
    return new_phase, signal

More efficient file writing

The list nesting done by prepSamples and the many calls to f.write when writing the file is significantly time-consuming. Let us try to format the data first and then write it all "at once". The struct module looks like the best way to go here. Here is a version of the file writing function modified to use this approach. (It's only for mono but shouldn't be too hard to extend to more channels.)

def genFile(sampleRate, bitsPerSample, channels, samples):
    """Expects samples in the form [sample0, sample1,...]
    Samples are values between -1 and 1 then upscaled to the correct 
    amplitude (-2**15 -> 2**15 roughly) for 16 bit"""

    header = gen_header(sampleRate, bitsPerSample, channels, samples)
    m = 2**(bitsPerSample - 1) - 1
    int_samples = [int(m * sample) for sample in samples]
    data = struct.pack('<{}h'.format(len(samples)), *int_samples)
    with open("test.wav", "wb") as f:
        f.write(header)
        f.write(data) 

Place top-level code in a function

It's good practice to put the "main" code into its own function, and using

if __name__ == "__main__":
    main()

This makes the file importable. And it helped me to profile the code.

Or, use Numpy instead

Until now I've focused on changes that can be made without any additional dependencies, but things become a lot easier and faster when we bring in Numpy. Numpy is a Python library for working with arrays, of which audio data is a typical example. Then, the frequency sweep and file-writing can be written as follows. (Additional logic may be needed to handle different byte orders.)

import numpy as np

def genFile(sampleRate, samples):
    dtype=np.int16
    m = np.iinfo(dtype).max
    samples *= m
    int_signal = samples.astype(dtype)    
    header = gen_header(sampleRate, 16, 1, samples)
    with open("test.wav", "wb") as f:
        f.write(header)
        int_signal.tofile(f)

def main():
    """Generate samples from 20Hz to 20kHz"""
    rate = 44100
    duration = 60
    freq = np.geomspace(20, 20000, num=rate*duration, dtype=np.float32)
    phase = 2*np.pi / rate * np.cumsum(freq)
    signal = 0.1 * np.sin(phase)
    genFile(rate, samples=signal)

Note the use of geomspace which makes the frequencies more perceptually evenly distributed. Use linspace instead to get the same effect as before.

Phase calculation from this answer.