encode_audio.py

#!/usr/bin/env python3
# Delta modulation audio encoder for playback via Uthernet II streaming.
#
# Simulates the Apple II speaker at 1MHz (i.e. cycle-level) resolution,
# by modeling it as a damped harmonic oscillator.
#
# On the Apple II side we use an audio player that is able to toggle the
# speaker with 1MHz precision, i.e. on any CPU clock cycle, although with a
# lower limit of 10 cycles between toggles (i.e. 102KHz maximum
# frequency).
#
# In order to reproduce a target audio waveform, we upscale it to 1MHz sample
# rate, i.e. to determine the desired speaker position at every CPU clock
# cycle, and compute the sequence of player operations on the Apple II side to
# best reproduce this waveform.
#
# This means that we are able to control the Apple II speaker with cycle-level
# precision, which results in high audio fidelity with low noise.
#
# To further optimize the audio quality we look ahead some defined number of
# cycles and choose a speaker trajectory that minimizes errors over this range.
# e.g. this allows us to anticipate large amplitude changes by pre-moving
# the speaker to better approximate them.
#
# This also needs to take into account scheduling the "end of frame" opcode
# every 2048 output bytes, where the Apple II will manage the TCP socket buffer
# while ticking the speaker at a regular (a, b) cadence to attempt to
# continue tracking the waveform as best we can.  Since we are stepping away
# from cycle-level management of the speaker during this period, it does
# introduce some quality degradation (manifesting as a slight background
# "crackle" to the audio)

import argparse
import collections
import contextlib
import functools
import librosa
import numpy
import soundfile as sf
from eta import ETA
from typing import Tuple

import lookahead
import opcodes
import opcodes_generated

# How many bytes to use per frame in the audio stream.  At the end of each
# frame we need to switch to special end-of-frame operations to cause the
# Apple II to manage the TCP socket buffers (ACK data received so far, and
# check we have at least another frame of data available)
#
# With an 8KB socket buffer this seems to be about the maximum we can get away
# with - it has to be page aligned, and 4KB causes stuttering even from a
# local playback source.
FRAME_SIZE = 2048


class Speaker:
    """Simulates the response of the Apple II speaker."""

    # TODO: move lookahead.evolve into Speaker method

    def __init__(self, sample_rate: float, freq: float, damping: float,
                 scale: float):
        """Initialize the Speaker object

        :arg sample_rate The sample rate of the simulated speaker (Hz)
        :arg freq The resonant frequency of the speaker
        :arg damping The exponential decay factor of the speaker response
        :arg scale Scale factor to normalize speaker position to desired range
        """
        self.sample_rate = sample_rate
        self.freq = freq
        self.damping = damping
        self.scale = numpy.float64(scale)  # TODO: analytic expression

        # See _Signal Processing in C_, C. Reid, T. Passin
        # https://archive.org/details/signalprocessing0000reid/
        dt = numpy.float64(1 / sample_rate)
        w = numpy.float64(freq * 2 * numpy.pi * dt)

        d = damping * dt
        e = numpy.exp(d)
        c1 = 2 * e * numpy.cos(w)
        c2 = e * e

        # Square wave impulse response parameters
        b2 = 0.0
        b1 = 1.0

        self.c1 = c1
        self.c2 = c2
        self.b1 = b1
        self.b2 = b2

def total_error(positions: numpy.ndarray, data: numpy.ndarray) -> numpy.ndarray:
    """Computes the total squared error for speaker position matrix vs data."""

    # Deal gracefully with the case where our speaker operation would slightly
    # run past the end of data
    min_len = min(len(positions), len(data))
    return numpy.sum(numpy.square(positions[:min_len] - data[:min_len]),
                     axis=-1)


@functools.lru_cache(None)
def frame_horizon(frame_offset: int, lookahead_steps: int):
    """Optimize frame_offset when more than lookahead_steps from end of frame.

    Candidate opcodes for all values of frame_offset are equal, until the
    end-of-frame opcode comes within our lookahead horizon.  This avoids
    needing to recompute many copies of the same candidate opcodes.

    TODO: this is a bit of a hack and we should be able to make the candidate
      opcode selection itself smarter about avoiding duplicate work
    """
    # TODO: This could be made tighter because a step is always at least 5
    #  cycles towards lookahead_steps.
    if frame_offset < (FRAME_SIZE - lookahead_steps):
        return 0
    return frame_offset


def audio_bytestream(data: numpy.ndarray, step: int, lookahead_steps: int,
                     sample_rate: int):
    """Computes optimal sequence of player opcodes to reproduce audio data."""
    # At resonance freq the scale is about 22400 but we can only access about 7%
    # of it across the frequency range.  This is also the equilibrium speaker
    # position when voltage is held constant. Normalize to this working
    # range for convenience.
    inv_scale = 22400 * 0.07759626164027278  # XXX

    # Speaker response parameters can be fitted in several ways.  First, by
    # recording audio of
    # - a single speaker click (impulse response)
    # - a frequency sweep over the entire audible spectrum
    #
    # Resonant frequency can be read off from the frequency spectrum.  For my
    # speaker there were two primary frequencies, at ~3875 and ~480Hz.
    # Looking at the click waveform the higher frequency mode dominates at
    # short time scales, and the lower frequency mode dominates at late
    # times.  Since we're interested in short timescale speaker control we
    # can ignore the latter.
    #
    # Damping factor can be estimated by fitting against the speaker click
    # waveform or the spectrum, e.g. by simulating a speaker click and then
    # computing its spectrum.
    #
    # TODO: other Apple II speakers almost certainly have different response
    #  characteristics, but hopefully not too widely varying.
    sp = Speaker(sample_rate, freq=3875, damping=-1210, scale=1 / inv_scale)

    # Starting speaker applied voltage.
    voltage1 = voltage2 = 1.0

    # last 2 speaker positions.
    # XXX 0.0?
    y1 = y2 = 1.0

    # Leave enough padding at the end to look ahead from the last data value,
    # and in case we schedule an end-of-frame opcode towards the end.
    # TODO: avoid temporarily doubling memory footprint to concatenate
    data = numpy.ascontiguousarray(numpy.concatenate(
        [data, numpy.zeros(max(lookahead_steps, opcodes.cycle_length(
            opcodes_generated.PlayerOps.TICK_00)), dtype=numpy.float32)]))

    # index within input audio data
    i = 0
    # Position in 2048-byte TCP frame
    frame_offset = 0

    # Total squared error of audio output
    total_err = 0.0

    # Keep track of how many large deviations we see from the target waveform
    # as another measure of audio quality
    clicks = 0
    # total squared error threshold to consider an operation as producing a
    # click
    click_threshold = 0.3

    # Keep track of how many opcodes we schedule, so we can print summary
    # statistics at the end
    opcode_counts = collections.defaultdict(int)

    # Always look ahead at least this many cycles.  We pick a higher value
    # when approaching end of frame, so we can maximize the
    min_lookahead_steps = lookahead_steps

    # Progress tracker
    # TODO: find a more adaptive ETA tracker, this one doesn't estimate well
    #   if the processing rate changes (which it does for us, since first few
    #   steps do extra work to precompute values that are cached for later
    #   steps)
    eta = ETA(total=1000, min_ms_between_updates=0)
    # Value of i at which we should next update eta
    next_eta_tick = 0

    while i < len(data):
        if i >= next_eta_tick:
            eta.print_status()
            next_eta_tick = int(eta.i * len(data) / 1000)

        if frame_offset >= (FRAME_SIZE - 5):  # XXX
            lookahead_steps = min_lookahead_steps + 130  # XXX parametrize
        else:
            lookahead_steps = min_lookahead_steps

        # The EOF opcodes need to act as a matched pair, so if we're emitting
        # the second one we need to pass in its predecessor
        last_opcode = opcode if frame_offset == FRAME_SIZE - 1 else None

        # Compute all possible opcode sequences we could emit starting at this
        # frame offset
        next_candidate_opcodes, voltages, lookahead_steps = \
            opcodes.candidate_opcodes(
                frame_horizon(frame_offset, lookahead_steps),
                lookahead_steps, last_opcode)

        # Simulate speaker trajectory for all of these  candidate opcode
        # sequences and pick the one that minimizes total error
        opcode_idx = lookahead.evolve_return_best(
            sp, y1, y2, voltage1, voltage2, voltage1 * voltages,
            data[i:i + lookahead_steps])
        opcode = next_candidate_opcodes[opcode_idx]
        opcode_length = opcodes.cycle_length(opcode)
        opcode_counts[opcode] += 1

        # Apply this opcode to evolve the speaker position
        opcode_voltages = (voltage1 * opcode.voltages).reshape((1, -1))
        all_positions = lookahead.evolve(
            sp, y1, y2, voltage1, voltage2, opcode_voltages)

        assert all_positions.shape[0] == 1
        assert all_positions.shape[1] == opcode_length

        # Update to new speaker state
        voltage1 = opcode_voltages[0, -1]
        voltage2 = opcode_voltages[0, -2]
        y1 = all_positions[0, -1]
        y2 = all_positions[0, -2]

        # Track accumulated error between desired and actual speaker trajectory
        new_error = total_error(
            all_positions[0] * sp.scale, data[i:i + opcode_length]).item()
        total_err += new_error
        if new_error > click_threshold:
            clicks += 1
            # XXX
            print(frame_offset, i / sample_rate, opcode, new_error,
                  numpy.mean(data[i:i + opcode_length]))

        # Emit chosen operation and simulated audio samples for recording
        yield opcode, numpy.array(
            all_positions * sp.scale, dtype=numpy.float32).reshape(-1)

        # Update input and output stream positions
        i += opcode_length
        frame_offset = (frame_offset + 1) % FRAME_SIZE

    # Make sure we have at least 2k left in stream so the player will do a
    # complete read of the last frame.
    # for _ in range(frame_offset % 2048, 2048):
    #     yield opcodes.Opcode.EXIT
    eta.done()

    # Print summary statistics
    print("Total error %f" % total_err)
    print("%d clicks" % clicks)
    print("Opcodes used:")
    for v, k in sorted(list(opcode_counts.items()), key=lambda kv: kv[1],
                       reverse=True):
        print("%s: %d" % (v, k))


def preprocess_audio(
        filename: str, target_sample_rate: int, normalize: float,
        normalization_percentile: int) -> numpy.ndarray:
    """Upscale input audio to target sample rate and normalize signal."""

    data, _ = librosa.load(filename, sr=target_sample_rate, mono=True)

    max_value = numpy.percentile(data, normalization_percentile)
    data /= max_value
    data *= normalize

    return data


def downsample_audio(simulated_audio, original_audio, input_rate, output_rate,
                     noise_output=False):
    """Downscale the 1MHz simulated audio output suitable for writing as .wav

    :arg simulated_audio The simulated audio data to downsample
    :arg original_audio The original audio data that was simulated
    :arg input_rate Sample rate of input audio
    :arg output_rate Desired sample rate of output audio
    :arg noise_output Whether to also produce a noise waveform, i.e. difference
      between input and output audio

    :returns Tuple of downsampled audio and noise data (or None
    if noise_output==False)
    """
    downsampled_output = librosa.resample(
        numpy.array(simulated_audio, dtype=numpy.float32),
        orig_sr=input_rate,
        target_sr=output_rate)

    downsampled_noise = None
    if noise_output:
        noise_len = min(len(simulated_audio), len(original_audio))
        downsampled_noise = librosa.resample(
            numpy.array(
                simulated_audio[:noise_len] - original_audio[:noise_len]),
            orig_sr=input_rate,
            target_sr=output_rate)

    return downsampled_output, downsampled_noise


def main():
    parser = argparse.ArgumentParser()
    parser.add_argument("--clock", choices=['pal', 'ntsc'],
                        help="Whether target machine clock speed is PAL ("
                             "1015657Hz) or NTSC (1020484)",
                        required=True)
    # TODO: implement 6502 - JMP indirect takes 5 cycles instead of 6
    parser.add_argument("--step_size", type=int,
                        help="Delta encoding step size")
    # TODO: if we're not looking ahead beyond the longest (non-end-of-frame)
    #  opcode then this will reduce quality, e.g. two opcodes may truncate to
    #  the same prefix, but have different results when we apply them
    #  fully.
    parser.add_argument("--lookahead_cycles", type=int,
                        help="Number of clock cycles to look ahead in audio "
                             "stream.")
    parser.add_argument("--normalization", default=0.8, type=float,
                        help="Overall multiplier to rescale input audio "
                             "values.")
    parser.add_argument("--norm_percentile", default=100,
                        help="Normalize to specified percentile value of input "
                             "audio")
    parser.add_argument("--wav_output", type=str, help="output audio file")
    parser.add_argument("--noise_output", type=str, help="output audio file")
    parser.add_argument("input", type=str, help="input audio file to convert")
    parser.add_argument("output", type=str, help="output audio file")
    args = parser.parse_args()

    # Effective clock rate, including every-65 cycle "long cycle" that takes
    # 16/14 as long.
    cpu_clock_rate = 1015657 if args.clock == 'pal' else 1020484  # NTSC

    input_audio = preprocess_audio(
        args.input, cpu_clock_rate, args.normalization, args.norm_percentile)
    print("Done preprocessing audio")

    # Sample rate for output .wav files
    # TODO: flag
    output_rate = 44100

    # Buffers simulated audio output so we can downsample it in suitably
    # large chunks for writing to the output .wav file
    output_buffer = []

    # Python contexts for writing output files if requested
    opcode_context = open(args.output, "wb+")
    if args.wav_output:
        wav_context = sf.SoundFile(
            args.wav_output, "w", output_rate, channels=1, format='WAV')
    else:
        # We're not creating a file but still need a context
        # XXX does this work?
        wav_context = contextlib.nullcontext
    if args.noise_output:
        noise_context = sf.SoundFile(
            args.noise_output, "w", output_rate, channels=1,
            format='WAV')
    else:
        # We're not creating a file but still need a context
        noise_context = contextlib.nullcontext

    with wav_context as wav_f, noise_context as noise_f, opcode_context \
            as opcode_f:
        # Tracks current position in input audio waveform
        input_offset = 0

        # Process input audio, writing output to ][-Sound audio file
        # and (if requested) .wav files of simulated speaker audio and
        # noise (difference between original and simulated audio)
        for idx, sample_data in enumerate(audio_bytestream(
                input_audio, args.step_size, args.lookahead_cycles,
                cpu_clock_rate)):
            opcode, samples = sample_data
            opcode_f.write(bytes([opcode.byte]))

            output_buffer.extend(samples)
            input_offset += len(samples)
            # Keep accumulating as long as we have <1MB in the buffer, or are
            # within 1MB from the end.  This ensures we have enough samples to
            # downsample, including the last (partial) buffer.
            if (
                    len(output_buffer) < 1 * 1024 * 1024 or (
                    len(input_audio) - input_offset) < 1 * 1024 * 1024
            ):
                continue

            # TODO: don't bother computing if we're not writing wavs
            downsampled_audio, downsampled_noise = downsample_audio(
                output_buffer, input_audio[input_offset - len(output_buffer):],
                cpu_clock_rate, output_rate, bool(args.noise_output)
            )
            if args.wav_output:
                wav_f.write(downsampled_audio)
                wav_f.flush()
            if args.noise_output:
                noise_f.write(downsampled_noise)
                noise_f.flush()

            output_buffer = []

        # TODO: handle last buffer more cleanly than duplicating this code
        if output_buffer:
            downsampled_audio, downsampled_noise = downsample_audio(
                output_buffer, input_audio[input_offset - len(output_buffer):],
                cpu_clock_rate, output_rate, bool(args.noise_output)
            )
            if args.wav_output:
                wav_f.write(downsampled_audio)
            if args.noise_output:
                noise_f.write(downsampled_noise)


if __name__ == "__main__":
    main()