{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "11018e82-afc7-4937-8914-a23d5286f656",
   "metadata": {},
   "outputs": [],
   "source": [
    "!pip3 install matplotlib --break-system-packages"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "005e6c38-f394-45b7-b890-ed1b72bf3995",
   "metadata": {},
   "outputs": [],
   "source": [
    "import os\n",
    "import time\n",
    "\n",
    "import matplotlib.pyplot as plt\n",
    "import numpy as np\n",
    "import scipy.io.wavfile as wav\n",
    "import scipy.signal\n",
    "\n",
    "from IPython.display import Audio\n"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "ae1158c3-aa77-4c23-a93d-f0a4854ac6c2",
   "metadata": {},
   "source": [
    "# Utility functions"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "7e97b1b2-d8b4-475d-9638-f3e239c9b414",
   "metadata": {},
   "outputs": [],
   "source": [
    "def lin_to_db(l):\n",
    "    return 20 * np.log10(l)\n",
    "\n",
    "def db_to_lin(d):\n",
    "    return np.pow(10.0, d / 20.0)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "12559a5a-067c-485a-8f36-61b3cc5c4fb5",
   "metadata": {},
   "outputs": [],
   "source": [
    "print(np.reshape(np.tile(np.array([[1,2,3,4],[5,6,7,8]])[:,:, np.newaxis], [1, 1, 2]), shape=(2, 8)))\n",
    "\n",
    "# I think replicate_last is not currently used.\n",
    "\n",
    "def replicate_last(x, num_reps):\n",
    "    \"\"\"Replicate the last dimension of an array num_reps times.\"\"\"\n",
    "    shape_in = x.shape\n",
    "    shape_out = np.array(shape_in)\n",
    "    shape_out[-1] *= num_reps\n",
    "    return np.reshape(np.tile(x[..., np.newaxis], list(np.ones_like(shape_in)) + [num_reps]), shape=shape_out)\n",
    "\n",
    "print(replicate_last(np.array([[1,2,3,4],[5,6,7,8]]), 5))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "292a19ac-6fce-477e-8a69-ca8d40fa238d",
   "metadata": {},
   "outputs": [],
   "source": [
    "w = np.hanning(256)\n",
    "x = np.random.rand(1000)\n",
    "print(np.convolve(w, x, 'same').shape)\n",
    "print(scipy.signal.convolve(x, w, 'same').shape)\n",
    "# \"same\" means \"same as 1st arg\" for scipy, but \"same as longer arg\" for numpy.\n",
    "#convolve = np.convolve\n",
    "convolve = scipy.signal.convolve"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "954bd4ca-1290-4282-90e4-31ed7c7b3530",
   "metadata": {},
   "outputs": [],
   "source": [
    "def old_slinterp(x, interp, truncate=False):\n",
    "    \"\"\"Linear interpolation of vector.  Ramps down to implicit zero at end unless truncate.\"\"\"\n",
    "    x_shape = list(x.shape)\n",
    "    x_out_shape = list(x.shape)\n",
    "    x_out_shape[-1] *= interp\n",
    "    x_in_flat = np.reshape(x, (np.prod([1] + x_shape[:-1]), x_shape[-1]))\n",
    "    n_rows, n_cols = x_in_flat.shape\n",
    "    x_out_flat = np.zeros((n_rows, interp * n_cols))\n",
    "    interpolation_window = np.hstack([np.linspace(0, 1, 1 + interp)[1:], np.linspace(1, 0, 1 + interp)[1:]])\n",
    "    for i in range(x_in_flat.shape[0]):\n",
    "        x_atrous = np.reshape(np.hstack([x_in_flat[[i], :].T, np.zeros((n_cols, interp - 1))]), (n_cols * interp,))\n",
    "        x_out_flat[i] = convolve(x_atrous, interpolation_window, 'same')\n",
    "    if truncate:\n",
    "        if interp > 1:  # Skip for degenerate interp == 1; indexing by [:, :-(1 - 1)] deletes the entire array.\n",
    "            x_out_flat = x_out_flat[:, :-(interp - 1)]\n",
    "        x_out_shape[-1] -= (interp - 1)\n",
    "    #print(f'{x_out_flat.shape=}, {x_out_shape=}')\n",
    "    return np.reshape(x_out_flat, x_out_shape)\n",
    "\n",
    "def slinterp(x, interp, truncate=False):\n",
    "    \"\"\"Linear interpolation of vector.  Ramps down to implicit zero at end unless truncate. CONV-FREE VERSION.\"\"\"\n",
    "    x_shape = list(x.shape)\n",
    "    x_out_shape = list(x.shape)\n",
    "    x_out_shape[-1] *= interp\n",
    "    x_in_flat = np.reshape(x, (np.prod([1] + x_shape[:-1]), x_shape[-1]))\n",
    "    n_rows, n_cols = x_in_flat.shape\n",
    "    x_out_flat = np.zeros((n_rows, interp * (n_cols + 1)))\n",
    "    interpolation_window = np.expand_dims(np.hstack([np.linspace(0, 1, 1 + interp)[1:], np.linspace(1, 0, 1 + interp)[1:]]), 0)\n",
    "    for i in range(x_in_flat.shape[-1]):\n",
    "        #x_atrous = np.reshape(np.hstack([x_in_flat[[i], :].T, np.zeros((n_cols, interp - 1))]), (n_cols * interp,))\n",
    "        #x_out_flat[i] = convolve(x_atrous, interpolation_window, 'same')\n",
    "        x_out_flat[..., i * interp : (i + 2) * interp] += x_in_flat[..., i] * interpolation_window\n",
    "    x_out_flat = x_out_flat[..., (interp - 1):-1]\n",
    "    if truncate:\n",
    "        if interp > 1:  # Skip for degenerate interp == 1; indexing by [:, :-(1 - 1)] deletes the entire array.\n",
    "            x_out_flat = x_out_flat[:, :-(interp - 1)]\n",
    "        x_out_shape[-1] -= (interp - 1)\n",
    "    #print(f'{x_out_flat.shape=}, {x_out_shape=}')\n",
    "    return np.reshape(x_out_flat, x_out_shape)\n",
    "\n",
    "print(slinterp(np.arange(4), 4, truncate=True).shape)\n",
    "print(slinterp(np.arange(4), 4, truncate=False).shape)\n",
    "print(slinterp(np.arange(4), 1, truncate=True).shape)\n",
    "print(slinterp(np.arange(4), 1, truncate=False).shape)\n",
    "\n",
    "plt.plot(np.arange(25), slinterp(np.array([[1, 1, 0, 2, 2]]), 5).T, '.',\n",
    "         np.arange(21), slinterp(np.array([[1, 1, 0, 2, 2]]), 5, truncate=True).T, '.')\n",
    "plt.grid()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "bb0499fb-3721-409a-ac1e-7f56e785b4a1",
   "metadata": {},
   "outputs": [],
   "source": [
    "from scipy import fft\n",
    "\n",
    "def spectrogram(x, fft_len=4096, win_len=None, hop_len=None, window_fn=np.hanning, sr=1, f_max=None, title=None):\n",
    "    if win_len is None:\n",
    "        win_len = fft_len\n",
    "    if hop_len is None:\n",
    "        hop_len = win_len // 2\n",
    "    # Window.\n",
    "    prepad_len = (fft_len - win_len) // 2\n",
    "    win = np.hstack([np.zeros(prepad_len), window_fn(win_len), np.zeros(fft_len - win_len - prepad_len)])\n",
    "    # Frame.\n",
    "    frame_indices = np.arange(0, len(x) - win_len, hop_len)[:, np.newaxis] + np.arange(win_len)[np.newaxis, :]\n",
    "    x_chunks_windowed = x[frame_indices] * win[np.newaxis, :]\n",
    "    # Transform.\n",
    "    stft_mag_db = 20 * np.log10(np.abs(fft.fft(x_chunks_windowed, n=fft_len)[:, :(fft_len // 2 + 1)]))\n",
    "    num_frames, num_bins = stft_mag_db.shape\n",
    "    t_base = np.arange(num_frames) * hop_len / sr\n",
    "    f_base = np.arange(num_bins) * sr / fft_len\n",
    "    plt.imshow(stft_mag_db.T, aspect='auto', origin='lower', extent=[t_base[0], t_base[-1], f_base[0], f_base[-1]])\n",
    "    plt.clim(np.max(stft_mag_db) + [-80, 0])\n",
    "    #plt.ylim([0, f_max / sr * fft_len])\n",
    "    if f_max:\n",
    "        plt.ylim([0, f_max])\n",
    "    #x_times = np.arange(0, len(x) / sr, 0.25)\n",
    "    #y_freqs = np.arange(0, 10000, 1000)\n",
    "    #plt.xticks(x_times * sr / hop_len, x_times)\n",
    "    #plt.yticks(y_freqs / sr * fft_len, y_freqs)\n",
    "    plt.colorbar()\n",
    "    if title:\n",
    "        plt.title(title)\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "e1b105f4-c55c-49e1-8f29-2d30a32a787b",
   "metadata": {},
   "outputs": [],
   "source": [
    "def synth_partial(mag, freq, frame_rate, sr):\n",
    "    \"\"\"Reconstruct audio waveform at sr from mag and freq 1-row arrays at frame_rate.\"\"\"\n",
    "    # Linear interpolation of both amplitude and frequency.\n",
    "    samps_per_frame = int(round(sr / frame_rate))\n",
    "    num_partials, num_frames = mag.shape\n",
    "    output_samples = num_frames * samps_per_frame\n",
    "    t = np.arange(output_samples) / sr\n",
    "    dt = 1 / sr\n",
    "    phase = np.cumsum(slinterp(2 * np.pi * freq * dt, samps_per_frame))\n",
    "    carrier = np.cos(phase)\n",
    "    return slinterp(mag, samps_per_frame) * carrier\n",
    "\n",
    "x = synth_partial(np.array([[1, 2, 3, 2]]), np.array([[40, 80, 120, 20]]), 10, 1000)\n",
    "plt.figure(figsize=(18, 4))\n",
    "plt.plot(x.T)\n",
    "plt.grid()\n",
    "print(x.shape)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "48b4585d-36ac-4633-a27a-b0fec2da930f",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Synthesize from (possibly downsampled) mag and freq contours.\n",
    "\n",
    "def synth_harmonic(freq, mag, interp_factor=1, sample_rate=44100, mag_is_db=True, truncate=True):\n",
    "    \"\"\"Convert a single pair of frequency and magnitude vectors to an audio waveform.\"\"\"\n",
    "    # As a special case, freq can be a scalar indicating a constant frequency.\n",
    "    carrier = None\n",
    "    if isinstance(freq, (list, np.ndarray)):\n",
    "        if len(freq) == 1:\n",
    "            freq = freq[0]\n",
    "        else:\n",
    "            assert len(freq) == len(mag)\n",
    "            # Ignore the final freq value, which tends to go crazy.  Replicate the preceding one.\n",
    "            #freq[-1] = freq[-2]\n",
    "            if interp_factor > 1:\n",
    "                freq = slinterp(freq, interp_factor, truncate=truncate)\n",
    "            carrier = np.cos(np.cumsum(2 * np.pi / sample_rate * freq))\n",
    "    if carrier is None:\n",
    "        # freq must be a scalar\n",
    "        #freq = freq * np.ones(1 + (len(mag) - 1) * interp_factor)\n",
    "        carrier = np.cos(2 * np.pi * freq / sample_rate * np.arange(truncate + (len(mag) - truncate) * interp_factor))\n",
    "\n",
    "    # Can now expand mag to match.\n",
    "    if interp_factor > 1:\n",
    "        mag = slinterp(mag, interp_factor, truncate=truncate)\n",
    "    # freq is cycles per sec.  We want radians per sample\n",
    "    if mag_is_db:\n",
    "        mag = db_to_lin(mag)\n",
    "    return mag * carrier\n",
    "\n",
    "x1 = synth_harmonic(40, np.array([1, 2, 3, 2]), 100, 1000, truncate=False, mag_is_db=False)\n",
    "x2 = synth_harmonic(np.array([40, 80, 120, 20]), np.array([1, 2, 3, 2]), 100, 1000, truncate=False, mag_is_db=False)\n",
    "plt.figure(figsize=(18, 4))\n",
    "plt.plot(x2.T)\n",
    "plt.plot(x.T)\n",
    "plt.plot(x1.T)\n",
    "plt.grid()\n",
    "print(x2.shape)\n",
    "\n",
    "def synth_harmonics(freqs, mags, interp_factor=1, sample_rate=44100):\n",
    "    \"\"\"Convert arrays where each row is a freq or mag vector to a waveform.\"\"\"\n",
    "    n_harms, n_frames = mags.shape\n",
    "    n_samps = 1 + interp_factor * (n_frames - 1)\n",
    "    result = np.zeros(n_samps)\n",
    "    for freq, mag in zip(freqs, mags):\n",
    "        result += synth_harmonic(freq, mag, interp_factor, sample_rate)\n",
    "    return result\n",
    "    "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "3bff32a8-278f-4d9b-9068-435506d12b93",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Helpers from piano-partials.ipynb\n",
    "\n",
    "def lin_fit(x):\n",
    "    \"\"\"Least-squares linear fit to points.\"\"\"\n",
    "    # xhat = a + b.index\n",
    "    # err = x[i] - (a + b.i)\n",
    "    # err^2 = x[i]^2 - 2 x[i] (a + b.i) + (a^2 + 2 a b i +b^2 i^2)\n",
    "    # Define axes around midpoint.\n",
    "    lenx = len(x)\n",
    "    if not lenx:\n",
    "        return []\n",
    "    index = np.arange(lenx) - (lenx - 1) / 2\n",
    "    # Make x zero mean\n",
    "    a = np.mean(x)\n",
    "    x_z = x - a\n",
    "    b = 0\n",
    "    if lenx > 1:\n",
    "        # Linear fit is normalized inner product\n",
    "        b = np.sum(x_z * index) / np.sum(index * index)\n",
    "    return a + b * index\n",
    "\n",
    "def expand_params(params):\n",
    "    \"\"\"Convert [val, npts, val, npts, val] sequence to points.\"\"\"\n",
    "    val0 = params[0]\n",
    "    outpts = [[val0]]\n",
    "    npts_val_pairs = np.vstack([params[1::2], params[2::2]])\n",
    "    for npts, val in npts_val_pairs.transpose():\n",
    "        outpts.append(np.linspace(val0, val, int(npts + 1))[1:])\n",
    "        val0 = val\n",
    "    return np.concatenate(outpts)\n",
    "\n",
    "plt.plot(expand_params([0, 1, 1, 2, 0, 4, 1]), '.-')\n"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "08a7c732-6d39-4a9a-ae17-b123bf5fa580",
   "metadata": {},
   "source": [
    "# Read waveforms"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "25c4d418-462a-4300-ae5c-9879c9402647",
   "metadata": {},
   "outputs": [],
   "source": [
    "def wavread(filename):\n",
    "  \"\"\"Read in audio data from a wav file.  Return d, sr.\"\"\"\n",
    "  # Read in wav file.\n",
    "  file_handle = open(filename, 'rb')\n",
    "  samplerate, wave_data = wav.read(file_handle)\n",
    "  # Normalize short ints to floats in range [-1..1).\n",
    "  data = np.asarray(wave_data, dtype=np.float32) / 32768.0\n",
    "  return data, samplerate\n",
    "\n",
    "def wavwrite(data, samplerate, filename):\n",
    "  \"\"\"Write a waveform to a WAV file.\"\"\"\n",
    "  wav.write(filename, samplerate, (32768.0 * data).astype(np.int16))\n",
    "\n",
    "def read_and_trim(filename, duration=2.0, rel_threshold=0.005, abs_threshold=0, sr=44100, channel=None, do_plot=False, pre_time=0.010):\n",
    "    d, file_sr = wavread(filename)\n",
    "    #print(\"d.shape=\", d.shape)\n",
    "    if len(d.shape) > 1:  # Stereo file\n",
    "        if channel is None:\n",
    "            d = np.mean(d, axis=-1)\n",
    "        else:\n",
    "            d = d[:, channel]\n",
    "    assert file_sr == sr\n",
    "    d = d - np.mean(d)\n",
    "    threshold = np.maximum(np.max(np.abs(d)) * rel_threshold, abs_threshold)\n",
    "    # Tight window at start?\n",
    "    drop_initial_samps = np.maximum(0, -int(round(pre_time * sr)) + np.min(np.flatnonzero(np.abs(d) > threshold)))\n",
    "    d_return = d[int(round(drop_initial_samps)) + np.arange(int(round(duration * sr)))]\n",
    "    if do_plot:\n",
    "        t = np.arange(len(d)) / sr\n",
    "        plt.figure(figsize=(12, 4))\n",
    "        plt.subplot(121)\n",
    "        plt.plot(t[:50000], d[:50000])\n",
    "        plt.ylim(threshold * np.array([-1, 1]))\n",
    "        plt.subplot(122)\n",
    "        plt.plot(t[:len(d_return)], d_return)\n",
    "        plt.xlim([0, 0.1])\n",
    "        print(\"drop initial\", drop_initial_samps / sr)\n",
    "    return d_return, sr\n",
    "\n",
    "def filename_for_note(note, octave, strike='ff', directory=None):\n",
    "    filename = '.'.join(['Piano', strike, note + str(octave), 'wav'])\n",
    "    if directory:\n",
    "        filename = os.path.join(directory, filename)\n",
    "    return filename\n",
    "\n",
    "#filename = 'Piano.ff.D4.wav'\n",
    "directory = '/Users/dpwe/Downloads/uiowa-piano'\n",
    "#note = 'C'\n",
    "#octave = 3\n",
    "note = 'Ab'\n",
    "octave = 1\n",
    "strike = 'pp'  # ['pp', 'mf', 'ff']:\n",
    "filename = filename_for_note(note, octave, strike, directory)\n",
    "print(filename)\n",
    "waveform, sr = read_and_trim(filename, channel=0, duration=5.0, rel_threshold=0.015, abs_threshold=0.0007, do_plot=True, pre_time=0.002)\n",
    "\n",
    "Audio(data=waveform.T, rate=sr)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "b613b2b0-a4f5-452f-877d-3ce01fa1ae2b",
   "metadata": {},
   "source": [
    "# Fundamental estimation from autocorrelation"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "1a051178-a322-4102-baf8-cf86932d1564",
   "metadata": {},
   "outputs": [],
   "source": [
    "def freq_from_autoco(d, sr=44100, max_duration=1.0, drop_initial_sec=0.05, do_plot=False, prop_max_thresh=0.75):\n",
    "    # Estimate fundamental by autocorrelation.\n",
    "    min_period = int(round(sr / 4000.0))\n",
    "    max_period = int(round(sr / 20.0))\n",
    "    drop_initial_samples = int(round(drop_initial_sec * sr))\n",
    "    d = d[drop_initial_samples : drop_initial_samples + int(round(sr * max_duration))]\n",
    "    #acr = np.correlate(d, d, 'full')[len(d) - 1:]\n",
    "    # Using FFT makes it 1000x faster, seconds to milliseconds.  Identical result.\n",
    "    acr = np.real(np.fft.ifft(np.abs(np.fft.fft(d, len(d) * 2)**2)))[:len(d)]\n",
    "    acr = acr[:10 * max_period]\n",
    "    # Find Nth peak, divide by N\n",
    "    maxima = scipy.signal.argrelextrema(acr, np.greater)[0]\n",
    "    # Keep only peaks at least 75% of largest peak below 4kHz.\n",
    "    maxima = maxima[np.flatnonzero(maxima >= min_period)]\n",
    "    max_acr_peak = np.max(acr[maxima])\n",
    "    maxima = maxima[np.flatnonzero(acr[maxima] > prop_max_thresh * max_acr_peak)]\n",
    "    npeaks = np.minimum(30, len(maxima))\n",
    "    # Assume first peak corresponds to fundamental, keep peaks that are close to integer multiples.\n",
    "    est_period = maxima[0]\n",
    "    #print(sr / est_period)\n",
    "    multiples = maxima / est_period\n",
    "    #print(multiples)\n",
    "    #print(np.abs((multiples - np.round(multiples)) / multiples))\n",
    "    # Max 2% relative error\n",
    "    good_multiples = np.flatnonzero(np.abs((multiples - np.round(multiples)) / multiples) < 0.02)\n",
    "    if do_plot:\n",
    "        plt.plot(acr)\n",
    "        plt.plot(maxima[:npeaks], acr[maxima[:npeaks]], 'or')\n",
    "        plt.plot([0, len(acr)], prop_max_thresh * max_acr_peak * np.array([1, 1]), '--')\n",
    "        plt.plot(maxima[good_multiples], acr[maxima[good_multiples]], 'ob')\n",
    "        plt.xlim([0, 10*est_period])\n",
    "    #freqs = sr * np.arange(1, npeaks + 1) / maxima[:npeaks]\n",
    "    freqs = (sr / (maxima / np.round(multiples)))[good_multiples]\n",
    "    #print(freqs)\n",
    "    freq = np.mean(freqs[-40:])\n",
    "    return freq\n",
    "\n",
    "\n",
    "%time freq = freq_from_autoco(waveform[8000:32000], do_plot=True, prop_max_thresh=0.81)\n",
    "print(freq)\n",
    "# Nominal: D5= 587.3295\n",
    "# C2 = 65.89466739451501\n",
    "# Db6 = 1108.73 Hz (we get 1113.37 for ff, 1114.50 for pp)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "5ea80ccb-d182-40bb-a673-5616fa1331fd",
   "metadata": {},
   "source": [
    "# Find frequency peaks in long Fourier transform"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "1a4ff398-98d8-4e3d-b80e-de85482c09ac",
   "metadata": {},
   "outputs": [],
   "source": [
    "from scipy import fft\n",
    "\n",
    "def localmax(x, npts):\n",
    "    \"\"\"Return indices of x where value is largest among nearest npts.\"\"\"\n",
    "    maxima = []\n",
    "    padval = np.array([np.min(x) - 1])\n",
    "    padded_x = np.hstack([padval, x, padval])\n",
    "    localmaxes = np.flatnonzero(np.logical_and(padded_x[1:-1]> padded_x[:-2], padded_x[1:-1] >= padded_x[2:]))\n",
    "    #for i in np.arange(1, len(x)):\n",
    "    for i in localmaxes:  \n",
    "        locality = x[np.maximum(0, i - npts // 2) : np.minimum(len(x), i + npts // 2)]\n",
    "        if x[i] == np.max(locality) and x[i] > x[i - 1]:\n",
    "            maxima.append(i)\n",
    "    return np.array(maxima)\n",
    "\n",
    "np.random.seed(57730)\n",
    "print(localmax(np.random.rand(32), npts=6))\n",
    "\n",
    "\n",
    "def peak_freqs(x, n_fft=65536, f_max=10000.0, depth_db=20.0, local_smooth_points=2047, local_max_points=100, sr=44100, est_f0=None, do_plot=False):\n",
    "    \"\"\"Return frequencies of local maxima from a long FFT.\"\"\"\n",
    "    # Provided f0 estimate is solely used to set the exclusion window for local max picking.\n",
    "    if est_f0:\n",
    "        # Use expected f_0 to set local_max_points.\n",
    "        f0_in_bins = est_f0 / (sr / n_fft)\n",
    "        local_max_points = int(round(f0_in_bins / 1))\n",
    "        if do_plot:\n",
    "            print(\"local_max_points=\", local_max_points)\n",
    "    long_fft = fft.fft(x[:n_fft])\n",
    "    long_fft = long_fft[:(n_fft // 2)]\n",
    "    freqs = np.arange(n_fft // 2) * sr / n_fft\n",
    "    n_bins = np.sum(freqs < f_max)\n",
    "    X_spec = 20 * np.log10(np.abs(long_fft[:n_bins]))\n",
    "    # Local average\n",
    "    smoo_X_spec = convolve(X_spec, np.ones(local_smooth_points) / local_smooth_points, 'same')\n",
    "    unsmoo_X_spec = X_spec - smoo_X_spec\n",
    "\n",
    "    #peaks = scipy.signal.argrelextrema(np.maximum(np.max(unsmoo_X_spec) - 30, unsmoo_X_spec), np.greater)[0]\n",
    "    peaks = localmax(np.maximum(np.max(unsmoo_X_spec) - depth_db, unsmoo_X_spec), local_max_points)\n",
    "\n",
    "    # Remove peaks more than 1 octave below expected F0.\n",
    "    peaks = peaks[peaks > f0_in_bins / 2.5]\n",
    "\n",
    "    if do_plot:\n",
    "        plt.figure(figsize=(20, 6))\n",
    "        plt.plot(freqs[:n_bins], unsmoo_X_spec)\n",
    "        plt.plot(freqs[:n_bins], smoo_X_spec)\n",
    "        plt.plot(freqs[peaks], unsmoo_X_spec[peaks], 'or')\n",
    "        plt.xlim((0, 4000))\n",
    "    \n",
    "    return peaks * sr / n_fft\n",
    "\n",
    "est_f0 = freq_from_autoco(waveform, drop_initial_sec=0.1)\n",
    "print('f0_from_autoco=', est_f0)\n",
    "#pk_frqs = peak_freqs(waveform[8000:], n_fft=32768, depth_db=30, do_plot=True, est_f0=est_f0)\n",
    "\n",
    "pk_frqs = peak_freqs(waveform[2000:80000], n_fft=32768, depth_db=40, do_plot=True, est_f0=est_f0)\n",
    "\n",
    "plt.grid()\n",
    "plt.xlabel('freq / Hz')\n",
    "plt.ylabel('level / dB')\n",
    "plt.title('Spectral peaks for ' + os.path.basename(filename))\n",
    "\n",
    "print(pk_frqs[:6])\n",
    "print(pk_frqs[:6] / est_f0)\n",
    "\n"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "0e8da7ea-cc11-4452-a6a6-bfd8f22bd6be",
   "metadata": {},
   "source": [
    "# Inharmonicity estimation"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "f27957df-3219-475d-a966-9009561564a1",
   "metadata": {
    "scrolled": true
   },
   "outputs": [],
   "source": [
    "# Manual search for lowest-error f0 and beta\n",
    "# Error is (measured_freq / pred_harm_freq) - 1.\n",
    "avg_freqs = pk_frqs\n",
    "hn = 1 + np.arange(len(avg_freqs))\n",
    "beta = 0.0003128\n",
    "f0 = est_f0 # 262.2118\n",
    "print(f0, avg_freqs)\n",
    "prop_err = avg_freqs / (f0 * hn * np.sqrt(1 + beta * hn * hn)) - 1\n",
    "#plt.plot(hn, avg_freqs / hn, hn, f0 * np.sqrt(1 + beta * hn * hn))\n",
    "plt.plot(hn, prop_err)\n",
    "print(\"rms prop err =\", np.sqrt(np.mean(prop_err ** 2)))\n",
    "f0= 253.68557562939364\n",
    "beta= 0.0008292878905710716\n",
    "#plt.plot(hn, f0 * np.sqrt(1 + beta * hn * hn))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "3e1a0bac-26b2-4310-b180-ea2a086d6493",
   "metadata": {},
   "outputs": [],
   "source": [
    "def est_inharmonicity(freqs, f0=None, beta=0, tolerance=0.2, verbose=False):\n",
    "    \"\"\"Estimate an inharmonicity beta for a series of harmonic frequencies.\"\"\"\n",
    "    freq_order = np.argsort(freqs)\n",
    "    if not f0:\n",
    "        f0 = freqs[freq_order[0]]\n",
    "    # Find harmonic numbers sequentially.\n",
    "    h_nums = np.zeros_like(freqs, dtype=int)\n",
    "    valid = np.zeros_like(freqs, dtype=bool)\n",
    "    f0s = np.zeros_like(freqs)\n",
    "    freq_errs = np.zeros_like(freqs)\n",
    "    last_freq = 0\n",
    "    last_hnum = 0\n",
    "    # Try every harmonic up to the largest candidate.\n",
    "    last_harm_freq = 0\n",
    "    for cand_hnum in 1 + np.arange(np.ceil(np.max(freqs) / f0)):\n",
    "        stretched_df = f0 * np.sqrt(1 + beta * cand_hnum * cand_hnum)\n",
    "        expected_freq = last_harm_freq + stretched_df\n",
    "        # Find actual freq closest to this.\n",
    "        best_freq_index = np.argmin(np.abs(freqs - expected_freq))\n",
    "        found_freq = freqs[best_freq_index]\n",
    "        delta_freq = expected_freq - found_freq\n",
    "        if np.abs(delta_freq / f0) > 0.1:\n",
    "            if verbose:\n",
    "                print(\"Closest harmonic to %d (%.1f) is %.1f Hz - reject\" % (cand_hnum, expected_freq, freqs[best_freq_index]))\n",
    "            last_harm_freq = expected_freq\n",
    "        else:\n",
    "            if verbose:\n",
    "                print(\"Found harmonic %d expected %.1f Hz got %.1f Hz\" % (cand_hnum, expected_freq, found_freq))\n",
    "            last_harm_freq = found_freq\n",
    "            h_nums[best_freq_index] = cand_hnum\n",
    "            valid[best_freq_index] = True\n",
    "    if verbose and np.sum(valid == False):\n",
    "        print(np.sum(valid == False), \"values invalid\");\n",
    "    #print(\"h_nums=\", h_nums)\n",
    "    #print(\"freq_errs=\", freq_err)\n",
    "    # We expect (f / h.f0) = sqrt(1 + beta h^2)\n",
    "    betas = (((freqs / (f0 * h_nums + (h_nums==0))) ** 2) - 1) / (h_nums ** 2 + (h_nums == 0))\n",
    "    #print(\"betas=\", betas)\n",
    "    weights = h_nums * h_nums * valid * (betas > 0) # Increase weighting in proportion to harmonic num for valid harmonics\n",
    "    #print(\"weights=\", weights)\n",
    "    # Weights can be all zero if all the calculated betas are <= 0.\n",
    "    if np.sum(weights):\n",
    "        beta = np.sum(weights * betas) / np.sum(weights)\n",
    "    # Now figure implied f0s\n",
    "    valid_ixs = np.flatnonzero(valid)\n",
    "    f0s = freqs / (h_nums * np.sqrt(1 + beta * (h_nums ** 2)) + (h_nums==0))\n",
    "    #print(\"f0s=\", f0s)\n",
    "    weights = valid / ((h_nums ** 2) + (valid == False))  # Downweight high harmonics for f0.\n",
    "    if np.sum(weights):\n",
    "        f0 = np.sum(weights * f0s) / np.sum(weights)\n",
    "    #f0 = np.median(f0s[valid])\n",
    "\n",
    "    if verbose:\n",
    "      plt.subplot(211)\n",
    "      plt.plot(h_nums, betas, '.')\n",
    "      plt.subplot(212)\n",
    "      plt.plot(h_nums, f0s, '.')\n",
    "    \n",
    "    return f0, beta, h_nums * valid\n",
    "    \n",
    "    # We are fitting freqs[h_num] = f0 * h_num * sqrt(1 + beta * (h_num ^ 2))\n",
    "    # Given beta, what is the (non-integral) h such that \n",
    "    # freq = f0 * h * sqrt(1 + beta * h^2)\n",
    "    # => freq/f0 = h sqrt(1 + beta h^2)\n",
    "    # => (f/f0)^2 = h^2 (1 + beta h^2)   => (f/(h.f0))^2 = 1 + beta h^2 => beta = ((f/(h.f0))^2 - 1) / h^2\n",
    "    # => (f/f0)^2 = h^2 + beta h^4\n",
    "    # => beta h^4 + h^2 - (f/f0)^2 = 0\n",
    "    # => h^2 = (-1 +/- sqrt(1 - 4 * beta * (-(f/f0)^2))) / (2 * beta)\n",
    "    # => h = sqrt( 1/(2 beta) * ( - 1))\n",
    "    # => freq^2 = f0^2 * h^2 * (1 + beta h^2) => beta h^4  + h^2 - (freq / f0)^2 = 0\n",
    "    # => h^2 = [-1 +/- sqrt(1 + 4 beta ((freq/f0) ^ 2) )] / 2 beta\n",
    "    # => h = sqrt( 1/(2 beta) sqrt(1 + 4 beta ((freq/f0) ^ 2)) )\n",
    "    #h_nums = np.sqrt( (1 / (2 * beta)) * (np.sqrt(1 + 4 * beta * ((freqs / f0) ** 2)) - 1))\n",
    "    #h_nums_q = np.round(h_nums)\n",
    "    #valid_wts = (np.abs(h_nums - h_nums_q) < 0.2) * (0 + np.arange(len(h_nums)))\n",
    "    #betas = (((freqs / (f0 * h_nums_q)) ** 2) - 1) / (h_nums_q ** 2)\n",
    "    #new_f0s = freqs / (h_nums_q * np.sqrt(1 + beta * (h_nums_q ** 2)))\n",
    "\n",
    "est_f0 = freq_from_autoco(waveform, sr, drop_initial_sec=0.1)\n",
    "# Trim transient based on fundamental frequency duration\n",
    "skip_initial_cycles = 10\n",
    "skip_initial_time = skip_initial_cycles / est_f0\n",
    "skip_initial_samples = int(round(skip_initial_time *sr))\n",
    "\n",
    "f0 = est_f0  # avg_freqs[0]\n",
    "beta = 0\n",
    "#f0 = 262.2118\n",
    "#beta = 0.0003128\n",
    "#freqs = avg_freqs[:25]\n",
    "#freqs = peak_freqs(d[8000:80000], n_fft=32768, depth_db=32, est_f0=est_f0)[:25]\n",
    "freqs = peak_freqs(waveform[skip_initial_samples:], n_fft=32768, depth_db=40, est_f0=est_f0)[:25]\n",
    "print(freqs)\n",
    "#del freqs[13]\n",
    "for _ in range(5):\n",
    "  print(f0, beta)\n",
    "  f0, beta, h_nums = est_inharmonicity(np.array(freqs), f0=f0, beta=beta, verbose=0)\n",
    "#for h_num, freq in zip(h_nums, avg_freqs):\n",
    "#    print(freq, h_num)\n",
    "print(h_nums)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "abe281f7-a747-4b40-998d-546073f9b5aa",
   "metadata": {},
   "outputs": [],
   "source": [
    "#f0 = 32.4\n",
    "#beta = 0.00025\n",
    "valid = np.flatnonzero(h_nums > 0)\n",
    "plt.plot(h_nums[valid], freqs[valid] / h_nums[valid], '.')\n",
    "\n",
    "_ = plt.plot(h_nums, f0 * np.sqrt(1 + beta * h_nums * h_nums), '.')\n",
    "#print(freqs[12])\n",
    "\n",
    "plt.legend(['Measured equivalent fundamentals of harmonics', 'Fundamentals for  inharmonicity parameter %f' % beta])\n",
    "plt.xlabel('harmonic number')\n",
    "plt.ylabel('fundamental frequency / Hz for ' + os.path.basename(filename))\n",
    "plt.grid()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "8dd7d682-63ca-4267-9393-7180bf00159c",
   "metadata": {},
   "source": [
    "# Calculate fundamental and inharmonicity across whole range"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "590179f0-f908-4048-a0d4-6cbca93fe627",
   "metadata": {},
   "outputs": [],
   "source": [
    "def expected_f0_hz(note, octave):\n",
    "    semis = {'C': 0, 'C#': 1, 'Db': 1, 'D': 2, 'D#': 3, 'Eb': 3, 'E': 4, 'F': 5, 'F#': 6, 'Gb': 6, 'G': 7, \n",
    "             'G#': 8, 'Ab': 8, 'A': 9, 'A#': 10, 'Bb': 10, 'B': 11}[note[0].upper() + note[1:]]\n",
    "    semis += (octave + 1) * 12  # semis is midi note number, C4 = 60.\n",
    "    return 440 * (2 ** ((semis - 69) / 12))\n",
    "\n",
    "print(expected_f0_hz('C', 4))\n",
    "print(expected_f0_hz('A', 4))\n",
    "print(expected_f0_hz('C', 5))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "4be2a0c6-f4cd-48ac-8a55-42d7c7eb3cb2",
   "metadata": {},
   "outputs": [],
   "source": [
    "directory = '/Users/dpwe/Downloads/uiowa-piano'\n",
    "strike = 'ff'  # ['pp', 'mf', 'ff']:\n",
    "\n",
    "filenames = []\n",
    "est_f0s = []\n",
    "f0s = []\n",
    "betas = []\n",
    "\n",
    "for octave in [1, 2, 3, 4, 5, 6]:\n",
    "  #for note in ['C', 'D', 'E', 'Gb', 'Ab', 'Bb']:\n",
    "  for note in ['C', 'Db', 'D', 'Eb', 'E', 'F', 'Gb', 'G', 'Ab', 'A', 'Bb', 'B']:\n",
    "      for strike in ['pp', 'mf', 'ff']:\n",
    "        filename =  filename_for_note(note, octave, strike)\n",
    "        #print(filename)\n",
    "        #d, sr = read_and_trim(os.path.join(directory, filename), channel=0, duration=5.0, \n",
    "        #                      rel_threshold=0.010, abs_threshold=0.0001, pre_time=0.01)\n",
    "        d, sr = read_and_trim(os.path.join(directory, filename), channel=0, duration=5.0, \n",
    "                              rel_threshold=0.015, abs_threshold=0.0007, pre_time=0.002)\n",
    "        drop_initial_samps = 8000  # int(round(drop_initial_sec * sr))\n",
    "        est_f0 = freq_from_autoco(d[drop_initial_samps : 32000], sr, prop_max_thresh=0.81)\n",
    "        est_f0s.append(est_f0)\n",
    "        filenames.append(filename)\n",
    "        freqs = peak_freqs(d[drop_initial_samps : 80000], n_fft=32768, depth_db=32, est_f0=est_f0)[:25]\n",
    "        f0 = est_f0\n",
    "        beta = 0\n",
    "        for _ in range(10):\n",
    "          f0, beta, h_nums = est_inharmonicity(np.array(freqs), f0=f0, beta=beta)\n",
    "        f0s.append(f0)\n",
    "        betas.append(beta)\n",
    "\n",
    "plt.plot(1200 * (np.log2(np.array(est_f0s) / (261.626 * 2)) - replicate_last(np.linspace(-4, 2, 72, endpoint=False), 3)), '.')\n",
    "plt.ylabel('tuning error / cents')\n",
    "_ = plt.xlabel('key number')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "c004668c-2064-4f43-8d89-80ba2e7f0eba",
   "metadata": {},
   "outputs": [],
   "source": [
    "for octave in [1, 2, 3, 4, 5, 6]:\n",
    "  #for note in ['C', 'D', 'E', 'Gb', 'Ab', 'Bb']:\n",
    "  for note in ['C', 'Db', 'D', 'Eb', 'E', 'F', 'Gb', 'G', 'Ab', 'A', 'Bb', 'B']:\n",
    "      for strike in ['pp', 'mf', 'ff']:\n",
    "        filename =  filename_for_note(note, octave, strike)\n",
    "        #print(filename)\n",
    "        d, sr = read_and_trim(os.path.join(directory, filename), channel=0, duration=5.0, \n",
    "                              rel_threshold=0.015, abs_threshold=0.0007, pre_time=0.002)\n",
    "        est_f0 = freq_from_autoco(d[8000:32000], sr, prop_max_thresh=0.81)\n",
    "        if np.abs(np.log2(est_f0 / expected_f0_hz(note, octave))) > 0.5/12:\n",
    "            print(filename, 'expected f0 %.1f, observed %.1f (%.1f cents)' % (\n",
    "                expected_f0_hz(note, octave), est_f0, 1200 * np.log2(est_f0 / expected_f0_hz(note, octave))))\n",
    "\n",
    "# Piano.pp.C1.wav was the same recording as Piano.pp.B0.wav, it's a data error.  \n",
    "# The error is present on the UIowa site, but I can't find any mention.\n",
    "# I resampled it up 100 cents with sox and rewrote it."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "a33d8b44-bdf5-4dc5-82ce-535cadef3b53",
   "metadata": {},
   "outputs": [],
   "source": [
    "#plt.plot(np.log2(np.array(est_f0s) / 440), '.')\n",
    "#print(len(est_f0s))\n",
    "#print(est_f0s)\n",
    "plt.plot(betas, '.-')\n",
    "print(filenames[5], est_f0s[5], betas[5])\n",
    "print(filenames[6], est_f0s[6], betas[6])"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "656b5fbc-c091-4d2e-888d-468d28d1912d",
   "metadata": {},
   "source": [
    "# Extract magnitude and frequency contours for a single harmonic\n",
    "\n",
    "`heterodyne_extract` extracts sample-by-sample magnitude and frequency contours for a single harmonic in the neighborhood of a specified input frequency (specified by a fundamental and a real-valued harmonic number, so the actual center frequency is simply the product of the two).  This is done by 'heterodyne extraction' - frequency-shifting the signal by multiplying it by the conjugate complex sinusoid at the center frequency (which shifts the center frequency down to 0 Hz), then low-pass filtering with a hanning window whose length is `smooth_cycles` times the fundamental period corresponding to `fundamental_freq`.  This window has zeros at the fundamental, so modulation at the fundamental rate is completely canceled.  This gives smooth contours for true harmonics of a fundamental."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "5278f393-a5b0-48ee-8b23-74293225a815",
   "metadata": {},
   "outputs": [],
   "source": [
    "#convolve = scipy.signal.convolve\n",
    "\n",
    "def heterodyne_extract(waveform, fundamental_freq, harmonic_number=1.0, sr=44100, smooth_cycles=2.0):\n",
    "    \"\"\"Return a full-sample-rate amplitude and frequency envelope near the specified frequency.\"\"\"\n",
    "    t = np.arange(len(waveform)) / sr\n",
    "    h_freq = harmonic_number * fundamental_freq\n",
    "    # Conjugate complex exponential at center freq.\n",
    "    complex_exp = np.exp(-1j * 2 * np.pi * h_freq * t)\n",
    "    f_period = sr / fundamental_freq\n",
    "    # Smoothing over e.g. 2 cycles of fundamental cancels periodic components, leaving only the DC shifted down from h_freq.\n",
    "    smooth_len = int(round(smooth_cycles * f_period))\n",
    "    #print(\"f_period=\", f_period, \"sm_len=\", smooth_len)\n",
    "    smooth_win = np.hanning(smooth_len)\n",
    "    smooth_win = smooth_win / np.sum(smooth_win)\n",
    "    smoothed_heterodyne_waveform = convolve(waveform * complex_exp, smooth_win, 'same')\n",
    "    heterodyne_magnitude = np.abs(smoothed_heterodyne_waveform)\n",
    "    # Unwrap the phase of the near-DC component, then also smooth at ~fundamental to get freq deviation.\n",
    "    unwrapped_phase = np.hstack([[0], np.unwrap(np.angle(smoothed_heterodyne_waveform))])\n",
    "    smoothed_dphi_dt = np.diff(convolve(unwrapped_phase, smooth_win, 'same'))\n",
    "    inst_freq = h_freq + smoothed_dphi_dt / (2 * np.pi) * sr\n",
    "    #plt.plot(unwrapped_phase[-10000:])\n",
    "    #plt.plot(np.angle((waveform * complex_exp)[-1000:]))\n",
    "    #plt.plot(np.imag(smoothed_heterodyne_waveform)[-1000:])\n",
    "    # If it moves more than 0.5 rad/sample, it's not tracking well.\n",
    "    dphi_dt_valid = (np.abs(smoothed_dphi_dt) < 0.5)\n",
    "    dphi_dt_weighting = heterodyne_magnitude * dphi_dt_valid\n",
    "    weighted_mean_dphi_dt = np.sum(dphi_dt_weighting * smoothed_dphi_dt) / np.sum(dphi_dt_weighting)   # radians / sample\n",
    "    avg_freq = h_freq + weighted_mean_dphi_dt / (2 * np.pi) * sr   # cycles / second\n",
    "    reconstructed_waveform = 2 * np.real(smoothed_heterodyne_waveform * np.conj(complex_exp))  # Smoothed complex waveform remodulated.\n",
    "    return lin_to_db(heterodyne_magnitude), inst_freq, avg_freq, reconstructed_waveform\n",
    "\n",
    "#f_freq = freq_from_autoco(waveform)  # 589.232\n",
    "#harmonic_num = 1\n",
    "#sm_cyc = 2.0\n",
    "#mag, freq, avg_freq, recons_waveform = heterodyne_extract(waveform, f_freq, harmonic_num, sr, sm_cyc)\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "d4fd5f4e-8901-4a5b-8a52-0888dcd1dfd1",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Heterodyne\n",
    "# h_num  delta_f\n",
    "# 1      -0.63\n",
    "# 2      -6.5\n",
    "# 3       4.1\n",
    "# 4      18.7\n",
    "# 5      35.3\n",
    "# 6      61.4\n",
    "# 7     104     4229\n",
    "# 8.2    31.4   4863\n",
    "# 9.3    97.1   5518\n",
    "# 10.2          6003\n",
    "# 10.8          6222\n",
    "# 11.2          6593\n",
    "# 11.7          6890\n",
    "# 13.3          7830\n",
    "# 14.15         8343\n",
    "# 15.45         9097\n",
    "# 16.7          9845\n",
    "\n",
    "f_freq = freq_from_autoco(waveform, drop_initial_sec=0.18)  # 589.232\n",
    "harmonic_num = 1\n",
    "sm_cyc = 2.0\n",
    "\n",
    "mag, freq, avg_freq, recons_waveform = heterodyne_extract(waveform, f_freq, harmonic_num, sr, sm_cyc)\n",
    "\n",
    "t = np.arange(len(waveform)) / sr\n",
    "ii = np.arange(22000) # \n",
    "#ii = np.arange(len(mag)-1000,len(mag))\n",
    "\n",
    "print(mag.shape)\n",
    "print(freq.shape)\n",
    "\n",
    "plt.subplot(211)\n",
    "#plt.plot(t[ii], db_to_lin(mag[ii]), '.')\n",
    "plt.plot(t[ii], mag[ii], '.')\n",
    "plt.grid()\n",
    "#plt.ylim([-60, -20])\n",
    "plt.title(f'harmonic {harmonic_num} - mag / dB')\n",
    "\n",
    "plt.subplot(212)\n",
    "plt.plot(t[ii], freq[ii], '.')\n",
    "plt.ylim([0.9 * avg_freq, 1.1 * avg_freq])\n",
    "plt.grid()\n",
    "plt.title(f'harmonic {harmonic_num} - freq / Hz')\n",
    "\n",
    "print(\"pre_freq=%.2f\" % (f_freq * harmonic_num), \"post_freq=%.2f\" % avg_freq, \"diff=%.2f\" % (avg_freq - f_freq * harmonic_num))\n",
    "\n",
    "#Audio(data=synth_partial(np.array([db_to_lin(mag)]), np.array([freq]), sr, sr), rate=sr)\n",
    "#Audio(data=synth_partial(np.array([db_to_lin(mag)]), avg_freq * np.ones((1, len(freq))), sr, sr), rate=sr)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "2617c691-9aa2-4776-ae0b-e27df473f486",
   "metadata": {},
   "outputs": [],
   "source": [
    "interp_factor = 44 * 10  # 10 ms sample period for parameters\n",
    "f_down = freq[::interp_factor]\n",
    "m_down = mag[::interp_factor]\n",
    "print(f'{f_down.shape=} {m_down.shape=}')\n",
    "h_recon = synth_harmonic(f_down, m_down, interp_factor=interp_factor, sample_rate=sr)\n",
    "plt.subplot(311)\n",
    "plt.plot(f_down)\n",
    "plt.subplot(312)\n",
    "plt.plot(m_down)\n",
    "plt.subplot(313)\n",
    "spectrogram(h_recon, sr=sr, f_max=4 * avg_freq)  #, fft_len=fft_len, f_max=f_max, title='original')\n",
    "#plt.plot(h_recon[-1000:])\n",
    "Audio(data=(h_recon).T, rate=sr)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "407b7b6e-e3f8-4ee4-9b47-c58b2154f17b",
   "metadata": {},
   "source": [
    "# Parent function extracts all harmonic envelopes for a note\n",
    "\n",
    "`extract_harmonics` goes through multiple steps to model a fixed-frequency note as a set of harmonics.\n",
    "* It makes an initial estimate of the fundamental from the autocorrelation of the whole note (with `freq_from_autoco`)\n",
    "* It finds all the most prominent peaks in the Fourier transform of the entire note.  The estimated fundamental is used here to set a local-dominance region so only the largest peak within one harmonics-spacing is returned.\n",
    "* It finds a best-fit for the inharmonicity of the upper harmonics by repeated application of `est_inharmonicity`.\n",
    "* It calls `heterodyne_extract` for up to the lowest `num_harms` harmonics, calculating a magnitude and (time-varying) frequency estimate for each one, then, for harmonics above `lowest_retry_harmonic`, extracts again based on the average of the extracted frequency.\n",
    "* The estimated harmonic is subtracted from the residual then goes on to extract the next harmonic.  The estimated harmonics are also accumulated to return a sinusoidal-model estimate of the input.\n",
    "* The final residual (aspects of the waveform not captured by the harmonics) is also returned."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "4f2d892b-2ec0-4437-b3b4-1867b64b3a5d",
   "metadata": {},
   "outputs": [],
   "source": [
    "def extract_harmonics(d, sr=44100, num_harms=8, skip_initial_cycles=10, verbose=False, lowest_retry_harmonic=0, est_f0=None):\n",
    "    \"\"\"Model a waveform with a set of harmonics.\"\"\"\n",
    "\n",
    "    recons_d = np.zeros_like(d)\n",
    "    recons_d_flat_f = np.zeros_like(d)\n",
    "\n",
    "    sm_cyc = 2.0\n",
    "\n",
    "    d_remaining = np.array(d)\n",
    "\n",
    "    mags = []\n",
    "    freqs = []\n",
    "    avg_freqs = []\n",
    "\n",
    "    # Initial F0 estimation from autocorrelation unless one was passed in.\n",
    "    if est_f0 is None:\n",
    "        est_f0 = freq_from_autoco(d, sr)\n",
    "    if verbose:\n",
    "        print('est_f0=', est_f0)\n",
    "    # Trim transient based on fundamental frequency duration\n",
    "    skip_initial_time = skip_initial_cycles / est_f0\n",
    "    skip_initial_samples = int(round(skip_initial_time * sr))\n",
    "    # Find all the prominent spectral peaks.  est_f0 is just for local max exclusion region.\n",
    "    f_peaks = peak_freqs(d[skip_initial_samples:], n_fft=32768, depth_db=32, est_f0=est_f0)[:num_harms]\n",
    "    if verbose:\n",
    "        print(\"f_peaks=\", f_peaks)\n",
    "\n",
    "    # We will extract tracks for each peak, but we will reference them to the independent f0 estimate.\n",
    "    #h_num_list = f_peaks / est_f0\n",
    "    max_inharm_peaks = 25  # Too many peaks messes up inharmonicity estimation for lower freqs.\n",
    "    f0 = est_f0\n",
    "    beta = 0\n",
    "    for _ in range(10):\n",
    "      # Iterate estimaing f0 and beta multiple times for covergence.\n",
    "      #print(f0, beta)\n",
    "      f0, beta, h_num_list = est_inharmonicity(f_peaks[:max_inharm_peaks], f0=f0, beta=beta)\n",
    "    # Use the fundamental extraction from the inharmonicity estimation.\n",
    "    est_f0 = f0\n",
    "    if verbose:\n",
    "        print('inh est_f0=', est_f0, 'beta=', beta)\n",
    "    \n",
    "    # Clip num_harmonics at 0.8 F_nyq\n",
    "    h_nums = (1 + np.arange(num_harms))\n",
    "    h_freqs = est_f0 * h_nums * np.sqrt(1 + beta * h_nums * h_nums)\n",
    "    num_harmonics_below_nyq = np.sum(h_freqs < 0.8 * sr / 2)\n",
    "    if num_harms > num_harmonics_below_nyq:\n",
    "        if verbose:\n",
    "            print('Clamping num_harms from', num_harms, 'to', num_harmonics_below_nyq)\n",
    "        num_harms = num_harmonics_below_nyq\n",
    "        h_nums = h_nums[:num_harms]\n",
    "    \n",
    "    for h_num in h_nums:\n",
    "        # Extract one harmonic from the residual so far.\n",
    "        stretched_hnum = h_num * np.sqrt(1 + beta * h_num * h_num)\n",
    "        # Oblige the search to be at the predicted frequency, incorporating inharmonic stretching\n",
    "        mag, freq, avg_freq, recons_waveform = heterodyne_extract(d_remaining, est_f0, stretched_hnum, sr, sm_cyc)\n",
    "        if verbose:\n",
    "            print(' h_num={:.3f}, imp_f0={:.1f}, f={:.1f} err_f={:.1f}, max_mag={:.1f} dB'.format(\n",
    "                h_num, avg_freq / stretched_hnum, avg_freq, avg_freq - est_f0 * stretched_hnum, max(mag)))\n",
    "\n",
    "        # Repeat with estimated frequency for higher harmonics.\n",
    "        if h_num > lowest_retry_harmonic:\n",
    "            mag, freq, avg_freq, recons_waveform = heterodyne_extract(d_remaining, avg_freq / h_num, h_num, sr, sm_cyc)\n",
    "            if verbose:\n",
    "                print('*h_num={:.3f}, imp_f0={:.1f}, f={:.1f} err_f={:.1f}, max_mag={:.1f} dB'.format(\n",
    "                    h_num, avg_freq / h_num, avg_freq, avg_freq - est_f0 * h_num, max(mag)))\n",
    "\n",
    "        recons_d += recons_waveform\n",
    "        d_remaining -= recons_waveform\n",
    "        #recons_d_flat_f += synth_partial(np.array([db_to_lin(mag)]), avg_freq * np.ones((1, len(freq))), sr, sr)[0]\n",
    "        recons_d_flat_f += synth_harmonic(avg_freq, np.array([db_to_lin(mag)]), 1, sr)[0]\n",
    "        mags.append(mag)\n",
    "        freqs.append(freq)\n",
    "        avg_freqs.append(avg_freq)\n",
    "\n",
    "    mags = np.array(mags)\n",
    "    freqs = np.array(freqs)\n",
    "    if verbose:\n",
    "        print(mags.shape)\n",
    "\n",
    "    return mags, recons_d, freqs, avg_freqs, est_f0\n",
    "\n",
    "#filename = filename_for_note('D', 4, 'ff')\n",
    "filename = filename_for_note('Ab', 1, 'pp')\n",
    "waveform, sr = read_and_trim(os.path.join(directory, filename), channel=0, duration=5.0, rel_threshold=0.015, abs_threshold=0.0007, pre_time=0.002)\n",
    "\n",
    "d = waveform\n",
    "\n",
    "mags, recons_d, freqs, avg_freqs, est_f0 = extract_harmonics(d, num_harms=40, lowest_retry_harmonic=0, verbose=1)\n",
    "#                                          extract_harmonics(d, num_harms=num_harms, est_f0=est_f0)\n",
    "print(f'{mags.shape=}, {freqs.shape=}')\n",
    "print(f'{avg_freqs=}')\n",
    "resid = d - recons_d\n",
    "\n",
    "t = np.arange(len(d)) / sr\n",
    "ii = np.arange(14000, 14200)\n",
    "plt.plot(t[ii], d[ii], '.', t[ii], recons_d[ii], '.', t[ii], resid[ii])\n",
    "#Audio(data=d.T, rate=sr)\n",
    "#Audio(data=recons_d.T, rate=sr)\n",
    "Audio(data=(resid).T, rate=sr)\n",
    "#Audio(data=recons_d_flat_f.T, rate=sr)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "0be16166-d18b-457d-8f6c-89eeea6f502f",
   "metadata": {},
   "outputs": [],
   "source": [
    "plt.figure(figsize=(20, 5))\n",
    "plt.imshow(mags, aspect='auto', origin='lower', interpolation='nearest')\n",
    "plt.colorbar()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "de0bf77e-8d15-4707-88f4-d9293865f7b6",
   "metadata": {},
   "outputs": [],
   "source": [
    "wavwrite(resid, sr, 'resid.wav')\n",
    "wavwrite(recons_d, sr, 'recons.wav')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "1a7ee247-10c1-47c2-8a5d-2cb00f711002",
   "metadata": {},
   "outputs": [],
   "source": [
    "f_max = 2000\n",
    "fft_len = 2 * 4096\n",
    "\n",
    "plt.figure(figsize=(18, 12))\n",
    "plt.subplot(311)\n",
    "spectrogram(d, sr=sr, fft_len=fft_len, f_max=f_max, title='original')\n",
    "plt.subplot(312)\n",
    "spectrogram(recons_d, sr=sr, fft_len=fft_len, f_max=f_max, title='partials')\n",
    "plt.subplot(313)\n",
    "spectrogram(resid, sr=sr, fft_len=fft_len, f_max=f_max, title='residual')\n",
    "#spectrogram(recons_d_flat_f, sr=sr, f_max=10000, title='flat_f partials')\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "12bfb193-1ba6-4f84-9bc8-897a7afa9bdd",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Resynthesize downsampled mags and freqs\n",
    "interp_factor = 44 * 10  # 10 ms sample period for parameters\n",
    "fs_down = freqs[:, ::interp_factor]\n",
    "ms_down = mags[:, ::interp_factor]\n",
    "print(f'{ms_down.shape=}')\n",
    "# Variable frequencies\n",
    "#print('frame-rate frequencies')\n",
    "#audio = synth_harmonics(fs_down, ms_down, interp_factor)\n",
    "# Fixed frequency per harmonic\n",
    "print('Constant frequencies')\n",
    "audio = synth_harmonics(avg_freqs, ms_down, interp_factor)\n",
    "\n",
    "spectrogram(audio, sr=sr, f_max=8000)  #, fft_len=fft_len, f_max=f_max, title='original')\n",
    "#plt.plot(h_recon[-1000:])\n",
    "Audio(data=(audio).T, rate=sr)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "89829014-3f7e-4f39-8242-7d5d2a31b043",
   "metadata": {},
   "outputs": [],
   "source": [
    "f = avg_freqs[0]\n",
    "ms = mags[0]\n",
    "%timeit a = synth_harmonic(f, ms, 1)\n",
    "plt.plot(synth_harmonic(f, ms[::2], 2))"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "cf340f1a-993c-4447-b993-ef65bdf74627",
   "metadata": {},
   "source": [
    "# Fixed log-time sampling\n",
    "\n",
    "We sample the magnitude contours at 21 log-spaced times (in milliseconds):\n",
    "```\n",
    "0, 4, 8, 12, 16, 24, 32, 48, 64, 96, 128, 192, 256, 384, 512, 768, 1024, 1536, 2048, 3072, 4096\n",
    "```"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "03fd6ceb-e399-4d03-be42-9ea078f054c9",
   "metadata": {},
   "outputs": [],
   "source": [
    "sample_times_ms = np.array([0, 4, 8, 12, 16, 24, 32, 48, 64, 96, 128, 192, 256, 384, 512, 768, 1024, 1536, 2048, 3072, 4096])\n"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "d8c09ace-f8fd-44cb-9a73-f43f21e781d2",
   "metadata": {},
   "source": [
    "# Adaptive magnitude sampling as convex hull"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "426cec34-6786-4bc2-b07f-f55ff5bf0636",
   "metadata": {},
   "outputs": [],
   "source": [
    "# New approach to linseg fitting - take all extrema, then filter down.\n",
    "filt_len = 31  # for D5\n",
    "#filt_len = 127  # for D4\n",
    "#filt_len = 1023\n",
    "filt_win = np.hanning(filt_len)/np.sum(np.hanning(filt_len))\n",
    "\n",
    "def lin_hull(x, ends=None):\n",
    "    \"\"\"Return a linear interpolation between the first and last values of x (or ends[0], ends[1]), of the same length as x.\"\"\"\n",
    "    if ends is None:\n",
    "        ends = (x[0], x[-1])\n",
    "    npts = len(x)\n",
    "    return np.linspace(ends[0], ends[1], npts)\n",
    "\n",
    "def max_prominence(x, ends=None, polarity=1):\n",
    "    \"\"\"Return index of point in x that is furthest above (below if polarity=-1) a line connecting the end points of x.\"\"\"\n",
    "    prominence = polarity * (x - lin_hull(x, ends))\n",
    "    argmax_prominence = np.argmax(prominence)\n",
    "    if np.max(prominence) > 0:\n",
    "        return argmax_prominence\n",
    "    return None\n",
    "\n",
    "def err_from_line(x, ends=None, abs_fn=np.abs):\n",
    "    \"\"\"Return sum(abs(error)) between x and line seg defined by ends.\"\"\"\n",
    "    return np.sum(abs_fn(x - lin_hull(x, ends)))\n",
    "\n",
    "def convex_hull(x, points=None, polarity=1):\n",
    "    \"\"\"Attempt to insert one new convex hull point in every interval defined by points.\"\"\"\n",
    "    if points is None:\n",
    "        points = [0, len(x) - 1]\n",
    "    new_points = [0]\n",
    "    for start in range(len(points) - 1):\n",
    "        new_point = max_prominence(x[points[start]:points[start + 1] + 1], polarity)\n",
    "        if new_point:\n",
    "            new_points.append(points[start] + int(new_point))\n",
    "        new_points.append(points[start + 1])\n",
    "    return new_points\n",
    "\n",
    "def hull(x, npoints):\n",
    "    \"\"\"Return npoints indices into x which attempt to provide a convex hull.\"\"\"\n",
    "    points = [0, len(x) - 1]\n",
    "    while len(points) < npoints:\n",
    "        max_err = 0\n",
    "        for i in range(len(points) - 1):\n",
    "            start = points[i]\n",
    "            end = points[i + 1]\n",
    "            err = err_from_line(x[start : end]) / np.sqrt((end - start))\n",
    "            #print(points, start, end, err, max_err)\n",
    "            if abs(err) > max_err:\n",
    "                max_err = abs(err)\n",
    "                worst_seg_ix = i\n",
    "        worst_start = points[worst_seg_ix]\n",
    "        worst_end = points[worst_seg_ix + 1]\n",
    "        polarity = 1 if err_from_line(x[worst_start : worst_end], abs_fn=np.asarray) > 0 else -1\n",
    "        points.insert(worst_seg_ix + 1, int(worst_start + max_prominence(x[worst_start : worst_end], polarity=polarity)))\n",
    "    return np.array(points)\n",
    "\n",
    "#h = convex_hull(mag_trim)\n",
    "#print(h)\n",
    "#h = convex_hull(mag_trim, h, -1)\n",
    "#print(h)\n",
    "#h = convex_hull(mag_trim, h)\n",
    "#print(h)\n",
    "#h = convex_hull(mag_trim, h, -1)\n",
    "#print(h)\n",
    "#h = convex_hull(mag_trim, h)\n",
    "#print(h)\n",
    "\n",
    "trim_start_samps = 128\n",
    "trim_end_samps = 1024\n",
    "\n",
    "hnum = 1\n",
    "mag_filt = lin_to_db(convolve(db_to_lin(mags[hnum]), filt_win, 'same'))\n",
    "mag_filt_trim = mag_filt[trim_start_samps:-trim_end_samps]\n",
    "max_times = scipy.signal.argrelmax(mag_filt_trim)[0]\n",
    "max_vals = mag_filt_trim[max_times]\n",
    "t = np.arange(len(mag_filt_trim)) / sr\n",
    "mmag = np.maximum(-130.0, mag_filt_trim)\n",
    "\n",
    "# Adaptive timing\n",
    "#h = hull(mmag, 12)\n",
    "# Fixed log-spaced timing\n",
    "h = np.round(sample_times_ms * 44.1).astype(np.int32)\n",
    "\n",
    "plt.figure(figsize=(16, 5))\n",
    "plt.subplot(121)\n",
    "#plt.plot(t, mmag, t[max_times], max_vals, 'o-')\n",
    "plt.plot(1000 * t, mmag)\n",
    "plt.plot(1000 * h / sr, mmag[h], 'o-')\n",
    "ylims = [-80, -20]\n",
    "t_lim = 0.15\n",
    "plt.xlim([0, 1000 * t_lim])\n",
    "#plt.legend(['mag', 'max_vals', 'hull'])\n",
    "plt.legend(['mag', 'hull'])\n",
    "plt.xlabel('time / ms')\n",
    "plt.ylabel('level / dB')\n",
    "plt.title('Envelope for harmonic %d of %s (%.1f Hz)' % (hnum, filename, est_f0))\n",
    "plt.ylim(ylims)\n",
    "plt.grid()\n",
    "plt.subplot(122)\n",
    "plt.plot(t, mmag)\n",
    "#plt.plot(t[max_times], max_vals, 'o-')\n",
    "plt.plot(h / sr, mmag[h], 'o-')\n",
    "# Line showing limit of left plot\n",
    "plt.plot([t_lim, t_lim], ylims, ':k')\n",
    "plt.xlim([0, 4.1])\n",
    "plt.ylim(ylims)\n",
    "plt.grid()\n",
    "plt.xlabel('time / s')\n",
    "print(h / 44.1)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "758f0649-c7d5-4c2a-89e4-80a11689ff65",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Convert sinusoids into bp sets\n",
    "\n",
    "# harms_params is a list of harmonic descriptions\n",
    "# each harmonic description is [const_freq, [(delta_time_ms, lin_mag), (delta_time_ms, lin_mag), ...]]\n",
    "\n",
    "def make_harms_params(mags, avg_freqs, max_harmonic=None, sr=44100, mag_floor=-100.0, terminal_slope=20.0, \n",
    "                      filt_len=31, adaptive_timing=True, verbose=False):\n",
    "  if not max_harmonic:\n",
    "      max_harmonic = mags.shape[0]\n",
    "    \n",
    "  dt = 1 / sr\n",
    "  sm_cyc = 2.0\n",
    "\n",
    "  #trim_ends_sec = 0.014\n",
    "  #trim_ends_samps = int(round(trim_ends_sec * sr))\n",
    "  trim_start_samps = 0 # 128\n",
    "  trim_end_samps = 1024\n",
    "\n",
    "  #filt_len = 31\n",
    "  filt_win = np.hanning(filt_len)/np.sum(np.hanning(filt_len))\n",
    "\n",
    "  harms_params = []\n",
    "\n",
    "  for hnum in range(max_harmonic):\n",
    "    mag_filt = lin_to_db(convolve(db_to_lin(mags[hnum]), filt_win, 'same'))\n",
    "    mag_filt_trim = mag_filt[trim_start_samps : -trim_end_samps]\n",
    "    #mag_filt_trim = mag_filt\n",
    "    mmag = np.maximum(mag_floor, mag_filt_trim)\n",
    "    if adaptive_timing:\n",
    "      # Adaptive timing\n",
    "      anchor_times = hull(mmag, 12)\n",
    "    else:\n",
    "      # Fixed log-spaced timing\n",
    "      sample_times_ms = np.array([0, 4, 8, 12, 16, 24, 32, 48, 64, 96, 128, 192, 256, 384, 512, 768, 1024, 1536, 2048, 3072, 4096])\n",
    "      anchor_times = np.round(sample_times_ms * 44.1).astype(np.int32)\n",
    "    anchor_vals = mmag[anchor_times]\n",
    "    nvals = list(zip([0] + list(np.diff(anchor_times)), anchor_vals))\n",
    "    bp_list = [(int(round(n * dt * 1000)), float(db_to_lin(val))) \n",
    "               for n, val in nvals]\n",
    "    last_time, last_mag = bp_list[-1]\n",
    "    last_mag_db = lin_to_db(last_mag)\n",
    "    final_mag_db = mag_floor\n",
    "    if adaptive_timing:\n",
    "      if last_mag_db > final_mag_db:\n",
    "        terminal_slope = terminal_slope  # -db/sec\n",
    "        final_dur = (last_mag_db - final_mag_db) / terminal_slope\n",
    "        bp_list.append((int(round(1000 * final_dur)), float(db_to_lin(final_mag_db))))\n",
    "      if adaptive_timing and bp_list[-1][1] == bp_list[-2][1]:  # Final points have same mag\n",
    "        bp_list = bp_list[:-1]\n",
    "    bp_list_pair = [float(avg_freqs[hnum]), bp_list]\n",
    "    if verbose:\n",
    "      print(\"h\", hnum + 1, bp_list_pair)\n",
    "    harms_params.append(bp_list_pair)\n",
    "  return harms_params\n",
    "\n",
    "harms_params = make_harms_params(mags, avg_freqs, mag_floor=-130, verbose=False, adaptive_timing=False)\n",
    "for hp in harms_params:\n",
    "    print(hp[0], len(hp[1]))\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "54a807c3-0072-45b4-a05a-226e1e0ba54c",
   "metadata": {},
   "outputs": [],
   "source": [
    "mags_dict = {}\n",
    "harms_dict = {}"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "43c1c74e-80e7-46a8-a663-6b759b86e424",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Filename to parameters set\n",
    "#filenames = ['Piano.mf.C5.wav', 'Piano.ff.C5.wav', 'Piano.mf.D5.wav', 'Piano.ff.D5.wav']\n",
    "#filenames = ['Piano.pp.D4.wav', 'Piano.pp.D5.wav', 'Piano.mf.D4.wav', 'Piano.mf.D5.wav', 'Piano.ff.D4.wav', 'Piano.ff.D5.wav']\n",
    "\n",
    "# Original U Iowa piano samples from https://theremin.music.uiowa.edu/mispiano.html converted to wav format.\n",
    "# Note: The file called Piano.pp.C0.wav was actually just a copy of Piano.pp.B1.wav (with different cropping).\n",
    "# I fixed its pitch by resampling:\n",
    "# dpwe@dpwe-macbookpro:~/Downloads$ sox uiowa-piano/Piano.pp.C1.wav piano.wav speed 100c \n",
    "# dpwe@dpwe-macbookpro:~/Downloads$ cp piano.wav uiowa-piano/Piano.pp.C1.wav \n",
    "\n",
    "directory = '/Users/dpwe/Downloads/uiowa-piano'\n",
    "\n",
    "num_harms = 40  # was 20\n",
    "\n",
    "filenames = []\n",
    "\n",
    "for octave in range(1, 8):\n",
    "    for note in ['C', 'E', 'Ab']:\n",
    "        est_f0 = None\n",
    "        for strike in ['ff', 'mf', 'pp']:\n",
    "            filename = filename = filename_for_note(note, octave, strike)\n",
    "            filenames.append(filename)\n",
    "            d, sr = read_and_trim(os.path.join(directory, filename), \n",
    "                                  channel=0, rel_threshold=0.015, abs_threshold=0.001, duration=5.0, pre_time=0.002)\n",
    "            #print(d.shape)\n",
    "            mags_dict[filename], recons_d, freqs, avg_freqs, est_f0 = extract_harmonics(d, num_harms=num_harms, est_f0=est_f0)\n",
    "            print(time.ctime(), filename, '%.1f' % est_f0, ['%.1f' % float(v) for v in avg_freqs[:4]])\n",
    "            #resid = d - recons_d\n",
    "            harms_dict[filename] = make_harms_params(mags_dict[filename], avg_freqs, mag_floor=-130, adaptive_timing=False)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "336ae992-222f-4fe9-9016-b9704d1734ec",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Describe note as initial amp + dB slope for single linear fit.\n",
    "filt_len = 31\n",
    "filt_win = np.hanning(filt_len)/np.sum(np.hanning(filt_len))\n",
    "\n",
    "plt.figure(figsize=(12, 6))\n",
    "colors = plt.cm.rainbow(np.linspace(0, 1, 21))\n",
    "avg_h_decay = np.zeros(len(filenames))\n",
    "avg_h_mag = np.zeros(len(filenames))\n",
    "for i, filename in enumerate(filenames):\n",
    "    #print(filename)\n",
    "    #plt.subplot(1, 3, (i % 3) + 1)\n",
    "    mags = mags_dict[filename]\n",
    "\n",
    "    num_harmonics = min(12, mags.shape[0])\n",
    "    h_init_amp = np.zeros(num_harmonics)\n",
    "    h_decay = np.zeros(num_harmonics)\n",
    "    h_nums = range(num_harmonics)\n",
    "    for h in h_nums:\n",
    "        mag_filt = lin_to_db(convolve(db_to_lin(mags[h]), filt_win, 'same'))\n",
    "        mag_filt_trim = mag_filt[trim_end_samps:-trim_end_samps]\n",
    "        t = np.arange(len(mag_filt_trim)) / sr\n",
    "        linfit = lin_fit(mag_filt_trim)\n",
    "        h_init_amp[h] = linfit[0]\n",
    "        h_decay[h] = (linfit[-1] - linfit[0]) / (t[-1] - t[0])\n",
    "    plt.plot(h_nums, h_decay, color=colors[i % len(colors)])\n",
    "    plt.ylim([-20, 10])\n",
    "    #plt.plot(h_nums, h_init_amp, color=colors[i])\n",
    "    #\n",
    "    #plt.plot(h_nums, lin_fit(h_decay), '--', h_nums, lin_fit(h_init_amp), '--', color=colors[i])\n",
    "    #plt.plot(h_nums, h_decay, color=colors[i])\n",
    "    #plt.plot(h_nums, lin_fit(h_decay), '--', color=colors[i])\n",
    "    #plt.legend(['decay dB/sec', 'init amp/dB'])\n",
    "    if i < 3:\n",
    "        plt.grid()\n",
    "        _ = plt.title(filename + ' avg dB/sec decay lin_fit harmonics')\n",
    "    avg_h_decay[i] = np.mean(h_decay)\n",
    "    avg_h_mag[i] = np.mean(h_init_amp)\n",
    "plt.legend(filenames)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "48d1ba1b-4cac-449c-a6db-dc10186342d8",
   "metadata": {},
   "outputs": [],
   "source": [
    "note_indices = np.arange(len(avg_h_decay))\n",
    "plt.figure(figsize=(16, 4))\n",
    "plt.plot(note_indices, avg_h_decay, '.-')\n",
    "plt.plot(note_indices, avg_h_mag, '.-')\n",
    "plt.legend(['avg_h_decay', 'avg_h_mag'])\n",
    "plt.grid()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "913bdbce-c0ed-4750-8f06-0366dc56ff05",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Mapping here is 1/8 shrunk after the first 250 ms.\n",
    "def time_mapper(t, boundary=0.25, magnification=8.0):\n",
    "        \"\"\"Convert time in secs to a nonlinear projection.\"\"\"\n",
    "        #return np.log10(0.01 + t)\n",
    "        t = np.asarray(t)\n",
    "        magnified = (t < boundary)\n",
    "        return magnified * t + (1 - magnified) * (boundary + (t - boundary) / magnification)\n",
    "\n",
    "def plot_mags(mags, filename, nharms=None, hop_len=64, sr=44100, mag_floor=-130.0, mag_ceil=-20.0, show_lin_fit=False):\n",
    "    if not nharms:\n",
    "        nharms = mags.shape[0]\n",
    "    nfrm = mags[:nharms, ::hop_len].shape[-1]\n",
    "    frmTime = np.arange(nfrm) * hop_len / sr\n",
    "    colors = plt.cm.rainbow(np.linspace(0, 1, max(12, nharms)))\n",
    "\n",
    "    boundary = 0.25\n",
    "    magnification = 8.0\n",
    "\n",
    "    for i in range(nharms):\n",
    "        plt.plot(time_mapper(frmTime[:nfrm]), mags[i, ::hop_len].T, '.', color=colors[i])\n",
    "        if show_lin_fit:\n",
    "            lin_fit_mag = lin_fit(mags[i, ::hop_len])\n",
    "            plt.plot(time_mapper(frmTime[:nfrm]), lin_fit_mag, '--', color=colors[i])\n",
    "\n",
    "    plt.plot(time_mapper([boundary, boundary]), [mag_floor, mag_ceil], '--k')\n",
    "    plt.ylim([mag_floor, mag_ceil])\n",
    "    xtimes = np.hstack([np.arange(0, 0.25, 0.0625), np.arange(0.5, 5.0, 0.5)])\n",
    "    plt.xlim([0, time_mapper(5.0)])\n",
    "    plt.xticks(time_mapper(xtimes), xtimes)\n",
    "    plt.legend(1 + np.arange(nharms), loc='upper right')\n",
    "    _ = plt.title((filename + ' - ' + str(nharms) + ' harmonics'))\n",
    "    \n",
    "\n",
    "def plot_harms(harms_params, filename, nharms=None, hop_len=64, sr=44100, mag_floor=-130.0, mag_ceil=-20.0):\n",
    "    if not nharms:\n",
    "        nharms = len(harms_params)\n",
    "    nfrm = mags[:nharms, ::hop_len].shape[-1]\n",
    "    frmTime = np.arange(nfrm) * hop_len / sr\n",
    "    colors = plt.cm.rainbow(np.linspace(0, 1, max(12, nharms)))\n",
    "\n",
    "    # Mapping here is 1/8 shrunk after the first 250 ms.\n",
    "    boundary = 0.25\n",
    "    magnification = 8.0\n",
    "\n",
    "    for i, hp in enumerate(harms_params):\n",
    "        times, vals = zip(*np.array(hp[1]))\n",
    "        times = np.cumsum(times) / 1000\n",
    "        vals = np.array(vals)\n",
    "        plt.fill_between(time_mapper(times), mag_floor, lin_to_db(vals), alpha=0.8, color=colors[i])\n",
    "    plt.plot(time_mapper([boundary, boundary]), [mag_floor, -mag_ceil], '--k')\n",
    "    plt.ylim([mag_floor, mag_ceil])\n",
    "    plt.legend(1 + np.arange(nharms))\n",
    "    xtimes = np.hstack([np.arange(0, 0.25, 0.0625), np.arange(0.5, 5.0, 0.5)])\n",
    "    plt.xticks(time_mapper(xtimes), xtimes)\n",
    "    plt.xlim([0, time_mapper(5.0)])\n",
    "    _ = plt.title((filename + ' - ' + str(nharms) + ' harmonics'))\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "bbd3fe33-19d5-4712-b652-76800dd98ba2",
   "metadata": {
    "scrolled": true
   },
   "outputs": [],
   "source": [
    "plt.figure(figsize=(15, 15))\n",
    "for i, filename in enumerate(filenames[:9]):\n",
    "    plt.subplot(3, 3, i + 1)\n",
    "    plot_mags(mags_dict[filename], filename, show_lin_fit=True)\n",
    "    #plot_harms(harms_dict[filename], filename)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "babe00ef-9fd2-4e74-aace-5beeec878f25",
   "metadata": {},
   "outputs": [],
   "source": [
    "plt.figure(figsize=(15, 15))\n",
    "for i, filename in enumerate(filenames[:9]):\n",
    "    plt.subplot(3, 3, i + 1)\n",
    "    #plot_mags(mags_dict[filename], filename)\n",
    "    plot_harms(harms_dict[filename], filename)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "5ae6a2d7-3fbd-4a87-8b10-7ba588892260",
   "metadata": {},
   "outputs": [],
   "source": [
    "def time_mapper(t, boundary=0.25, magnification=8.0):\n",
    "    \"\"\"Convert time in secs to a nonlinear projection.\"\"\"\n",
    "    return np.log10(0.01 + t)\n",
    "    #t = np.asarray(t)\n",
    "    #magnified = (t < boundary)\n",
    "    #return magnified * t + (1 - magnified) * (boundary + (t - boundary) / magnification)\n",
    "\n",
    "def inv_time_mapper(mapped_t, boundary=0.25, magnification=8.0):\n",
    "    magnified = mapped_t < boundary\n",
    "    return magnified * mapped_t + (1 - magnified) * (boundary + magnification * (mapped_t - boundary))\n",
    "\n",
    "from matplotlib.collections import PolyCollection\n",
    "\n",
    "def polygon_under_graph(x, y):\n",
    "    return [(x[0], 0.), *zip(x, y), (x[-1], 0.)]\n",
    "\n",
    "#filename = filename_for_note(note='C', octave=4, strike='ff')\n",
    "\n",
    "def harmonics_waterfall_plot(filename, axes=None):\n",
    "    global harms_dict\n",
    "    if axes == None:\n",
    "        axes = fig.add_subplot(projection='3d')\n",
    "    verts = []\n",
    "    mag_floor = -130\n",
    "    for i, hp in list(enumerate(harms_dict[filename])):\n",
    "            times, vals = zip(*np.array(hp[1]))\n",
    "            times = np.cumsum(times) / 1000\n",
    "            vals = np.array(vals)\n",
    "            times = np.hstack([[times[0]], times])\n",
    "            vals = np.hstack([[db_to_lin(mag_floor)], vals])\n",
    "            #plt.plot(np.log10(0.01 + times), lin_to_db(vals), zs=-i, zdir='y', alpha=0.8, color=colors[i])\n",
    "            #poly = axes.add_collection3d(plt.fill_between(2 + np.log10(0.01 + times), -100, lin_to_db(vals), alpha=0.8, color=colors[i]), \n",
    "            #                           zs=(i + 1), zdir='y')\n",
    "            verts.append(polygon_under_graph(time_mapper(times), -mag_floor + lin_to_db(vals)))\n",
    "    facecolors = plt.colormaps['rainbow'](np.linspace(0, 1, len(verts)))\n",
    "    poly = PolyCollection(verts, facecolors=facecolors, alpha=.7)\n",
    "    axes.add_collection3d(poly, zs=np.arange(len(harms_dict[filename]), 0, -1), zdir='y')\n",
    "    axes.set_zlim(0, 120)\n",
    "    axes.set_zlabel('level / dB')\n",
    "    \n",
    "    plt.title(filename)\n",
    "    plt.xlabel('time / s')\n",
    "    #tick_times = np.array([0, 0.1, 0.2, 1.0, 2.0, 3.0, 4.0])\n",
    "    tick_times = np.array([0, 0.02, 0.05, 0.1, 0.2, 0.5, 1.0, 2.0, 5.0])\n",
    "    plt.xticks(time_mapper(tick_times), tick_times)\n",
    "    #plt.xlim([-0.05, 0.8])\n",
    "    plt.ylabel('harm #')\n",
    "    plt.yticks(np.arange(2, 22, 2), np.arange(19, -1, -2))\n",
    "    #plt.zlabel('level / dB')\n",
    "\n",
    "    #plt.zlim([-100, -20])\n",
    "    #plt.legend(1 + np.arange(nharms))\n",
    "    #plt.xticks(np.log10(0.01 + xtimes), xtimes)\n",
    "    #plt.xlim([0, 3000])\n",
    "    #_ = plt.title((filename + ' - ' + str(nharms) + ' harmonics'))\n",
    "    axes.view_init(elev=20., azim=-115, roll=0)\n",
    "\n",
    "#fig = plt.figure(figsize=(8, 6))\n",
    "\n",
    "plt.figure(figsize=(18, 6))\n",
    "axes = plt.subplot(131, projection='3d')\n",
    "note = 'C'\n",
    "octave = 4\n",
    "harmonics_waterfall_plot(filename_for_note(note, octave, 'pp'), axes=axes)\n",
    "axes = plt.subplot(132, projection='3d')\n",
    "harmonics_waterfall_plot(filename_for_note(note, octave, 'mf'), axes=axes)\n",
    "axes = plt.subplot(133, projection='3d')\n",
    "harmonics_waterfall_plot(filename_for_note(note, octave, 'ff'), axes=axes)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "a332f17f-5863-4a99-8399-3dc7c5661f9d",
   "metadata": {},
   "outputs": [],
   "source": [
    "import json\n",
    "print(filename.split('.'))\n",
    "params_file = '.'.join(filename.split('.', -1)[:-1]) + '.json'\n",
    "print(params_file)\n",
    "with open(params_file, 'w') as f:\n",
    "    f.write(json.dumps(harms_params))\n",
    "print('saved to', params_file)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "63de5760-ec44-4246-8ede-17050a2ede1d",
   "metadata": {},
   "outputs": [],
   "source": [
    "import itertools\n",
    "\n",
    "hps = harms_params[0]\n",
    "bps = np.array(hps[1])\n",
    "print(bps[:, 0])\n",
    "plt.plot(expand_params(list(itertools.chain.from_iterable(bps))[1:]))\n",
    "plt.plot(np.cumsum(bps[:, 0]), bps[:, 1], '.')\n",
    "_ = plt.xlim(0, 10000)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "e346232b-bddf-45cc-929a-c9d90d859efc",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Synthesize various versions.\n",
    "# Original magnitudes, fixed frequencies\n",
    "\n",
    "def synth_harm_param(harm_param, sr=44100):\n",
    "    \"\"\"Synthesize a single harmonic described by a harm_param structure.\"\"\"\n",
    "    # harm_param is [freq, [(dtime_ms, mag), (dtime_ms, mag), ...]]\n",
    "    freq, breakpoints = harm_param\n",
    "    breakpoints = np.array(breakpoints)\n",
    "    breakpoints[:, 1] = lin_to_db(breakpoints[:, 1])\n",
    "    # magnitudes interpolated on millisecond scale in log domain.\n",
    "    # First value is assumed to be zero, so dropped.\n",
    "    mags = expand_params(list(itertools.chain.from_iterable(breakpoints))[1:])\n",
    "    return synth_harmonic(freq, mags, interp_factor=int(round(sr / 1000)), sample_rate=sr, mag_is_db=True)\n",
    "\n",
    "\n",
    "def synth_params(harms_params, verbose=False):\n",
    "    \"\"\"Synthesize a note described by a list of harm_param structures.\"\"\"\n",
    "    upper_bound_len_samples = 44100 * 10\n",
    "    result = np.zeros(upper_bound_len_samples)\n",
    "    max_harm_len_samples = 0\n",
    "    for hp in harms_params:\n",
    "        harm_audio = synth_harm_param(hp)\n",
    "        harm_len_samples = len(harm_audio)\n",
    "        result[:harm_len_samples] += harm_audio\n",
    "        max_harm_len_samples = np.maximum(max_harm_len_samples, harm_len_samples)\n",
    "    return result[:max_harm_len_samples], harm_audio\n",
    "\n",
    "\n",
    "def synth_params_orig(harms_params, verbose=False):\n",
    "    # Older version does not use the common synth_harmonic function.\n",
    "    result = np.zeros(0)\n",
    "    for f0, hp in harms_params:\n",
    "        # ignore first time, start list with first value.\n",
    "        params = [lin_to_db(hp[0][1])]\n",
    "        # Step through remaining pairs, converting times in ms to samples\n",
    "        for bpp in hp[1:]:\n",
    "            params.extend([int(round(bpp[0] * sr/1000)), lin_to_db(bpp[1])])\n",
    "        # Expand (val, npoints, val, npoints..., val) list\n",
    "        ep = expand_params(params)\n",
    "        carrier = np.cos(2 * np.pi * f0 * np.arange(len(ep)) / sr)\n",
    "        ep_len = len(ep)\n",
    "        if ep_len > len(result):\n",
    "            result = np.hstack([result, np.zeros(ep_len - len(result))])\n",
    "        result[:ep_len] += db_to_lin(ep) * carrier\n",
    "        if verbose:\n",
    "            print(f0)\n",
    "    return result, ep\n",
    "\n",
    "#audio, ep = synth_params([harms_params[i] for i in 1 + np.arange(10)])\n",
    "#filename = 'Piano.ff.C4.wav'\n",
    "filename = 'Piano.ff.C3.wav'\n",
    "#hd = harms_dict[filename][:40]\n",
    "# How does it sound if we skip some of the higher harmonics?\n",
    "# In particular, above the 16th harmonic, skip harmonics that are 2.n.f0 and 3.n.f0.  This lets us get higher, for brighter sounds, \n",
    "# and the missing harmonics aren't particularly noticeable in those regions.\n",
    "hd = [harms_dict[filename][i] for i in [0, 1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 12, 13, 14, 15, 16, 18, 22, 24, 28, 30, 34, 36]]\n",
    "h_audio, ep = synth_params(hd)\n",
    "#h_audio, ep = synth_params(harms_params)\n",
    "\n",
    "#f0 = harms_params[0][0]\n",
    "#carrier = np.cos(2 * np.pi * f0 * np.arange(len(ep)) / sr)\n",
    "\n",
    "#ii = np.arange(20000)\n",
    "ii = np.arange(4000)\n",
    "#plt.plot(t[ii], db_to_lin(mags[h_num - 1][ii]))\n",
    "#plt.plot(t[ii], db_to_lin(ep[ii]))\n",
    "#plt.title(filename)\n",
    "npts = len(resid)\n",
    "audio = h_audio[:npts]  # + resid\n",
    "Audio(data=audio, rate=sr)\n",
    "print(audio.shape, resid.shape, sr)\n",
    "#old_audio = audio\n",
    "start_pt = 950\n",
    "npts = 100\n",
    "plt.plot((start_pt + np.arange(npts)) / sr, audio[start_pt : start_pt + npts])\n",
    "#         (start_pt + np.arange(npts)) / sr, old_audio[start_pt : start_pt + npts])\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "ad4e2638-a2e3-48ac-a0af-e26085bb9646",
   "metadata": {},
   "outputs": [],
   "source": [
    "print(filename, len(audio), sr)\n",
    "Audio(data=audio, rate=sr)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "a32e7895-40b6-4ce3-ba32-cadfa79d2eec",
   "metadata": {},
   "outputs": [],
   "source": [
    "print(filename, len(audio), sr)\n",
    "Audio(data=audio, rate=sr)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "ff874fcf-dc6a-49e0-b144-74580dc34f28",
   "metadata": {},
   "outputs": [],
   "source": [
    "waveform, sr = read_and_trim(os.path.join(directory, filename), channel=0, duration=5.0, rel_threshold=0.010, abs_threshold=0.001, do_plot=False, pre_time=0.01)\n",
    "Audio(data=waveform, rate=sr)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "147eba90-e290-4f3f-b62a-7c082249da3b",
   "metadata": {},
   "outputs": [],
   "source": [
    "def interp_harm_param(hp0, hp1, alpha):\n",
    "    \"\"\"Return harm_param list that is alpha of the way to hp1 from hp0.\"\"\"\n",
    "    hp = []\n",
    "    for h0, h1 in zip(hp0, hp1):\n",
    "        f_a = np.exp(np.log(h0[0]) + alpha * (np.log(h1[0]) - np.log(h0[0])))\n",
    "        #if (abs(h0[0] - h1[0]) / (0.5 * (h0[0] + h1[0]))) > 0.1:\n",
    "        #    print(\"warn: interpolating\", h0[0], \"and\", h1[0])\n",
    "        # Don't warn since we're interpolating between notes too.\n",
    "        bps = []\n",
    "        eps = 1e-9\n",
    "        for bp0, bp1 in zip(h0[1], h1[1]):\n",
    "            bp = (bp0[0] + alpha * (bp1[0] - bp0[0]), np.exp(-eps + np.log(bp0[1] + eps) + alpha * (np.log(bp1[1] + eps) - np.log(bp0[1] + eps))))\n",
    "            bps.append(bp)\n",
    "        hp.append((f_a, bps))\n",
    "    return hp\n",
    "\n",
    "def print_hp(harm_params, max_harms=8):\n",
    "    print('***')\n",
    "    for harm_param in harm_params[:max_harms]:\n",
    "        s = '%.1f ' % harm_param[0]\n",
    "        for bp in harm_param[1][:8]:\n",
    "            s += '%d, %.1f / ' % (bp[0], 1000 * bp[1])\n",
    "        print(s)\n",
    "\n",
    "alpha = 0.9\n",
    "#ns = ['Piano.mf.E3.wav', 'Piano.mf.Ab3.wav', 'Piano.mf.C4.wav']\n",
    "ns = ['Piano.pp.C4.wav', 'Piano.mf.C4.wav', 'Piano.ff.C4.wav']\n",
    "\n",
    "#print_hp(harms_dict[n1])\n",
    "#print_hp(harms_dict[n2])\n",
    "#print_hp(interp_harm_param(harms_dict[n1], harms_dict[n2], alpha))\n",
    "#Audio(data=synth_params(interp_harm_param(harms_dict[n1], harms_dict[n2], alpha))[0], rate=sr)\n",
    "w = []\n",
    "for start in [0, 1]:\n",
    "    n1 = ns[start]\n",
    "    n2 = ns[start + 1]\n",
    "    for alpha in np.linspace(0, 1, 4, endpoint=False):\n",
    "        w.append(synth_params(interp_harm_param(harms_dict[n1], harms_dict[n2], alpha))[0][:sr])\n",
    "w = np.concat(w)\n",
    "Audio(data=w, rate=sr)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "06082d70-e4f1-4471-ba16-b20af0fc7121",
   "metadata": {},
   "outputs": [],
   "source": [
    "wavwrite(w, sr, '/tmp/tmp.wav')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "38a1588a-90b0-4e59-ad33-b38455c83db6",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Compare to pure scaling\n",
    "audio = synth_params(harms_dict[ns[1]])[0][:sr]\n",
    "ww = []\n",
    "for alpha in np.linspace(0, np.log(100/7), 10):\n",
    "    ww.append(np.exp(alpha) * 0.2 * audio)\n",
    "ww = np.concat(ww)\n",
    "wavwrite(ww, sr, '/tmp/tmp2.wav')"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "0d1e2e2e-70d3-4447-aee5-c67b0fdb94e7",
   "metadata": {},
   "source": [
    "# New plan:\n",
    "- Each note is described by a conformal set of parameters\n",
    "  - exactly one param set per harmonic (so harmonics of different notes are always in the same place)\n",
    "  - harmonic envelopes are sampled at a fixed set of log-spaced times (so we can interpolate notes without changing timing)\n",
    "We write the analysis out as a big blob:\n",
    " - A single vector of harmonic envelope time intervals (odd-numbered values of BP lists, e.g. 20 of them)\n",
    " - An array giving the base note, base velocity, and # harmonics for all the described notes.  For octaves 1-7 with 3 notes/oct ('C', 'E', 'Ab') and 3 velocities ('pp', 'mf', 'ff'), that's 63 notes.  At up to 20 harmonics per ote, that's 1260 harmonics.\n",
    " - A big array giving the 20 target values per harmonic per note.  Maybe these should be in dB, and 8 bit integers?  With 1260 harmonics, that's a total of 25,200 values (100 kB in float32)."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "14e41472-4bdc-45e1-9544-7e7785669527",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Construct the piano-params.json file.\n",
    "# Contents are:\n",
    "#   sample_times_ms - single vector of fixed log-spaced envelope sample times (in int16 integer ms)\n",
    "#   notes - Array of MIDI notes\n",
    "#   velocities - Array of strike velocities available for each note, the same for all notes\n",
    "#   num_harmonics - Array of (num_notes * num_velocities) counts of how many harmonics are defined for each note.\n",
    "#   harmonics_freq - Vector of (total_num_harmonics) int16s giving freq for each harmonic in cents 6900 = 440 Hz.\n",
    "#   harmonics_mags - Array of (total_num_harmonics, 20) uint8s giving 20 envelope samples for each harmonic.  In dB re: -130.\n",
    "\n",
    "# sample_times_ms is defined upscript.\n",
    "#print(f'{len(sample_times_ms)=}')\n",
    "#print(sample_times_ms)\n",
    "\n",
    "def filename_of_note_vel(note, vel):\n",
    "    return ('Piano.' + \n",
    "            {40: 'pp', 80: 'mf', 120: 'ff'}[velocity] + \n",
    "            '.' +\n",
    "            {0: 'C', 4: 'E', 8: 'Ab'}[note % 12] + str((note - 12) // 12) + \n",
    "            '.wav')\n",
    "\n",
    "def note_vel_of_filename(filename):\n",
    "    _, vel, note, _ = filename.split('.')\n",
    "    vel_val = {'pp': 40, 'mf': 80, 'ff': 120}[vel]\n",
    "    note_name = note[:-1]\n",
    "    note_oct = int(note[-1])\n",
    "    note_val = {'C': 0, 'E': 4, 'Ab': 8}[note_name] + 12 * note_oct + 12  # C4 is midi note 60.\n",
    "    return note_val, vel_val\n",
    "\n",
    "# Extract notes from filenames.\n",
    "assert set(filenames) == set(harms_dict.keys())\n",
    "velocities = set()\n",
    "notes = set()\n",
    "for filename in filenames:\n",
    "    note_val, strike_val = note_vel_of_filename(filename)\n",
    "    velocities.add(strike_val)\n",
    "    notes.add(note_val)\n",
    "\n",
    "notes = sorted(notes)\n",
    "velocities = sorted(velocities)\n",
    "#print(notes)\n",
    "\n",
    "num_harmonics = []\n",
    "\n",
    "for note in notes:\n",
    "    for velocity in velocities:\n",
    "        filename = filename_of_note_vel(note, velocity)\n",
    "        num_harmonics.append(len(harms_dict[filename]))\n",
    "\n",
    "total_num_harmonics = sum(num_harmonics)\n",
    "\n",
    "harmonics_freq = np.zeros(total_num_harmonics, dtype=int)\n",
    "harmonics_mags = np.zeros((total_num_harmonics, len(sample_times_ms) - 1), dtype=int)\n",
    "harmonic_index = 0\n",
    "\n",
    "def hz_to_cents(freq_hz):\n",
    "    \"\"\"Convert freq in hz to cents: 6000 = C4 = 261.63, 6900 = A4 = 440.0.\"\"\"\n",
    "    return np.round(1200 * np.log2(freq_hz / 440.0) + 6900).astype(int)\n",
    "\n",
    "def cents_to_hz(cents):\n",
    "    \"\"\"Convert freq cents back to Hzin hz to real-valued midi: 60 = C4 = 261.63, 69 = A4 = 440.0.\"\"\"\n",
    "    return 440 * np.exp2((cents - 6900) / 1200)\n",
    "\n",
    "\n",
    "print(hz_to_cents(440))\n",
    "print(hz_to_cents(444))\n",
    "print(cents_to_hz(hz_to_cents(444)))\n",
    "\n",
    "for note in notes:\n",
    "    for velocity in velocities:\n",
    "        filename = filename_of_note_vel(note, velocity)\n",
    "        harm_data = harms_dict[filename]\n",
    "        print(filename, \":\", len(harm_data), \"harmonics\")\n",
    "        for harm in harm_data:\n",
    "            # harm is [freq, [(time, mag), (time, mag), ...]]\n",
    "            # Round to 1 dp.\n",
    "            #harmonics_freq[harmonic_index] = np.round(10 * harm[0]) / 10.0\n",
    "            # Store freq in \"midi cents\" (a440 = 6900)\n",
    "            harmonics_freq[harmonic_index] = hz_to_cents(harm[0])\n",
    "            mag_info = np.array(harm[1])\n",
    "            # mag_info may include a final value which is not at a fixed time if it needed a \"run out\" to get down to -60dB.\n",
    "            len_mag_info = mag_info.shape[0]\n",
    "            if len_mag_info > len(sample_times_ms):\n",
    "                # print(f'{len_mag_info=} {len(sample_times_ms) + 1=}')\n",
    "                mag_info = mag_info[:-1]\n",
    "            #print(mag_info.T)\n",
    "            assert np.all(np.cumsum(mag_info[:, 0]) == sample_times_ms)\n",
    "            # Offset all harmonics mags as rounded to dB re: -130, the min from above.\n",
    "            harmonics_mags[harmonic_index] = 130 + np.round(lin_to_db(mag_info[1:, 1])).astype(int)\n",
    "            harmonic_index += 1\n",
    "\n",
    "# 3 notes per octave; 20 harmonics per note (for low notes), so C4 = 3 octaves in = 180 harmonics in.\n",
    "print(harmonics_freq[178:190])\n",
    "print(harmonics_mags[178:190])\n",
    "print(harmonic_index, \"total harmonics in\", len(notes) * len(velocities), \"notes\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "a012cd30-f232-4c00-a7d9-5d4437e351f5",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Save out the json file.\n",
    "data = {\n",
    "    # Drop the initial zero in sample_times_ms.  Initial sample is always 0.\n",
    "    'sample_times_ms': sample_times_ms[1:].tolist(),\n",
    "    'notes': notes,\n",
    "    'velocities': velocities,\n",
    "    'num_harmonics': num_harmonics,\n",
    "    'harmonics_freq': harmonics_freq.tolist(),\n",
    "    'harmonics_mags': harmonics_mags.tolist(),\n",
    "}\n",
    "json_file = 'piano-params.json'\n",
    "with open(json_file, 'w') as f:\n",
    "    f.write(json.dumps(data))\n",
    "print('saved to', json_file)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "e2a8b257-09fa-4b6d-b1f6-62e238138151",
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3 (ipykernel)",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.12.6"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}
