modwave wavetable surface plots pdf generator

ENV.yml

name: plots
channels:
  - conda-forge
  - defaults
dependencies:
  - python=3.11.0
  - tqdm
  - numpy
  - matplotlib
  - colorcet
  - librosa
  - fpdf2

plot_wavetables.py

"""
To run

$ conda create -n plots -f ENV.yml
$ conda activate plots
$ python plot_cqt_wavetables.py

"""

from pathlib import Path

import colorcet as cc
from fpdf import FPDF
from matplotlib import cm
import matplotlib.pyplot as plt
from matplotlib.colors import LightSource, ListedColormap
import numpy as np
import librosa
from tqdm import tqdm


SAMPLES_PER_WAVE = 2048

# Image layout params
IMAGE_WIDTH = 40
X_INCREMENT = 40
X_LIM = X_INCREMENT * 5
X_LEFT = 5
Y_INCREMENT = 32
Y_LIM = Y_INCREMENT * 9
Y_TOP = 5

# Choose image type to embed in pdf:
# IMAGE_TYPE = '.png'
IMAGE_TYPE = ".jpg"


def all_data_len(filepath):
    raw32 = np.fromfile(filepath, dtype=np.float32, offset=0x82)
    return len(raw32)


def wavetable_data(filepath):
    waves = int(all_data_len(filepath) / SAMPLES_PER_WAVE)
    data = np.fromfile(filepath, dtype=np.float32, offset=0x82)
    data = data[: SAMPLES_PER_WAVE * waves]
    data.shape = (-1, SAMPLES_PER_WAVE)
    data = np.fliplr(data)
    return data


class PDF(FPDF):
    def __init__(self):
        super().__init__()
        self.xx = X_LEFT
        self.yy = Y_TOP

    def add_image(self, filename):
        self.image(f"{filename}{IMAGE_TYPE}", self.xx, self.yy, IMAGE_WIDTH)
        self.xx += X_INCREMENT
        if self.xx >= X_LIM:
            self.xx = X_LEFT
            self.yy += Y_INCREMENT
            if self.yy >= Y_LIM:
                self.yy = Y_TOP
                self.add_page()


def waterfall(title, wavetable):
    """Creates and saves a png thumbnail of the wavetable

    Args:
        title: object with png filename
        wavetable (ndarray): n x 64 float32 ndarray of wavetable data

    """
    fig = plt.figure()
    ax = fig.add_subplot(projection="3d")
    ax.set(
        box_aspect=(1, 1, 0.1),
        yticklabels=[],
        zticklabels=[],
        yticks=[],
        zticks=[],
        xmargin=0,
        ymargin=0,
        frame_on=False,
    )
    ax.azim = 200
    ax.elev = 35
    ax.autoscale(tight=True)
    ax.grid(visible=False)
    plt.title(title, y=0.1, fontsize=30)

    wave_count, wavelength = wavetable.shape
    X, Y = np.mgrid[:wave_count, :wavelength]

    ax.plot_wireframe(X, Y, wavetable, cstride=wavelength, lw=0.5, color="k")

    fig.subplots_adjust(top=1.22, bottom=-0.1, left=-0.1, right=1.1)
    plt.savefig(f"{title}{IMAGE_TYPE}")
    plt.close()


def cqt(title, wavetable):
    """spectrogram using librosa's constant-q transform
    Creates and saves a png thumbnail of the wavetable

    Args:
        title: object with png filename
        wavetable (ndarray): n x 64 float32 ndarray of wavetable data

    """
    fig, ax = plt.subplots()

    cmap_f = cc.m_gouldian
    cmap = ListedColormap(cmap_f(np.linspace(0.1, 0.8, 256)))

    sr = 48000 * 2
    _wave_count, _wavelength = wavetable.shape
    fw = np.tile(wavetable, (512,)).flatten()
    cq = np.abs(librosa.cqt(fw, sr=sr))
    librosa.display.specshow(
        librosa.amplitude_to_db(cq, ref=np.max), sr=sr, ax=ax, cmap=cmap
    )
    fig.subplots_adjust(top=0.95, bottom=0.15, left=0.07, right=0.93)
    plt.title(title, y=-0.15, fontsize=30)
    plt.savefig(f"{title}{IMAGE_TYPE}")
    plt.close()


def terrain(title, wavetable):
    """Creates and saves a png thumbnail of the wavetable

    Args:
        title: object with png filename
        wavetable (ndarray): n x 64 float32 ndarray of wavetable data

    """
    fig = plt.figure()
    ax = fig.add_subplot(projection="3d")
    ax.set(
        box_aspect=(1, 1, 0.1),
        yticklabels=[],
        zticklabels=[],
        yticks=[],
        zticks=[],
        xmargin=0,
        ymargin=0,
        frame_on=False,
    )
    ax.azim = 200
    ax.elev = 35
    ax.autoscale(tight=True)
    ax.grid(visible=False)
    plt.title(title, y=0.1, fontsize=30)

    wave_count, wavelength = wavetable.shape
    if wave_count == 1:
        wave_count = 2
        wavetable = np.tile(wavetable, (2, 1))
    X, Y = np.mgrid[:wave_count, :wavelength]

    # See https://matplotlib.org/stable/gallery/mplot3d/custom_shaded_3d_surface.html
    # ls = LightSource(270, 65)
    ls = LightSource(210, 65)
    # To use a custom hillshading mode, override the built-in shading and pass
    # in the rgb colors of the shaded surface calculated from "shade".
    x = X
    y = Y
    z = wavetable
    # cm.cividis cm.autumn cm.copper cm.YlOrBr_r
    # cc.m_gouldian cc.m_bgy cc.m_bky cc.m_CET_I1
    cmap_f = cc.m_gouldian
    cmap = ListedColormap(cmap_f(np.linspace(0.1, 0.8, 256)))
    rgb = ls.shade(z, cmap=cmap, vert_exag=5, blend_mode="soft")
    surf = ax.plot_surface(
        x,
        y,
        z,
        rstride=1,
        cstride=1,
        facecolors=rgb,
        linewidth=0,
        antialiased=False,
        shade=False,
    )

    fig.subplots_adjust(top=1.22, bottom=-0.1, left=-0.1, right=1.1)
    plt.savefig(f"{title}{IMAGE_TYPE}")
    plt.close()


if __name__ == "__main__":
    pdf = PDF()
    pdf.alias_nb_pages()
    pdf.add_page()

    p = Path(r".")
    total = len(list(p.glob("*.mwwavetable")))
    for i, f in enumerate(pbar := tqdm(p.glob("*.mwwavetable"), total=total)):
        pbar.set_description(f.stem)
        wave_data = wavetable_data(f)
        cqt(f.stem, wave_data)
        # terrain(f.stem, wave_data)
        # waterfall(f.stem, wave_data)
        pdf.add_image(f.stem)
    pdf.output("factory_wavetables.pdf", "F")