simplified_syrinx_jax.ipynb

{
 "cells": [
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Implementing and fitting a simple syrinx model\n",
    "\n",
    "Model based on [Mindlin et al., 2003](https://journals.aps.org/pre/abstract/10.1103/PhysRevE.68.041908). "
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "We want to fit the system of first-order nonlinear equations:\n",
    "\n",
    "$$\n",
    "\\dot{x} = y \\\\\n",
    "\\dot{y} = -\\epsilon x - C x^2 y + By - D_0\n",
    "$$\n",
    "where\n",
    "$$\n",
    "\\epsilon = \\epsilon_1 + \\epsilon_2 K(t) \\\\\n",
    "B = \\beta_1 + \\beta_2 P(t) \\\\\n",
    "D_0 = \\delta D(t)\n",
    "$$\n",
    "and $K(t)$, $D(t)$, and $P(t)$ are the (linear envelopes of) tension in the ventral syringeal muscle (vS), the tracheobronchialis dorsalis (dTB), and sub-syringeal air pressure, respectively."
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "From the original paper, we take parameter values\n",
    "$$\n",
    "\\epsilon_1 = 1.25 \\times 10^8 \\, \\mathrm{s}^{-2} \\\\\n",
    "\\epsilon_2 = 7.5 \\times 10^9 \\, \\mathrm{V}^{-1}\\cdot \\mathrm{s}^{-2} \\\\\n",
    "C = 2 \\times 10^8 \\, \\mathrm{cm}^{-2} \\cdot \\mathrm{s}^{-1} \\\\\n",
    "\\beta_1 = -2 \\times 10^3 \\, \\mathrm{s}^{-1} \\\\\n",
    "\\beta_2 = 5.3 \\times 10^4 \\, \\mathrm{V}^{-1}\\cdot \\mathrm{s}^{-1} \\\\\n",
    "\\delta = 15 \\times 10^6 \\, \\mathrm{cm}\\cdot\\mathrm{V}^{-1} \\cdot \\mathrm{s}^{-2}\n",
    "$$"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "eps1 = 1.25e8\n",
    "eps2 = 7.5e9\n",
    "beta1 = -2e3\n",
    "beta2 = 5.3e5  # NOTE: 10x higher than in paper!\n",
    "C = 2e8\n",
    "delta = 15e6"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "But whereas the original paper used measured values for $K$, $P$, and $D$, we use simpler time series constructed to have the same shape:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import jax\n",
    "import jax.numpy as jnp\n",
    "import jax.random as jr\n",
    "from jax import lax\n",
    "import matplotlib.pyplot as plt\n",
    "\n",
    "%config InlineBackend.figure_format = 'retina'\n",
    "plt.rcParams.update({\n",
    "    \"text.usetex\": True,\n",
    "    \"font.family\": \"Helvetica\"\n",
    "})"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "sr = 4e4  # data sampling rate (Hz)\n",
    "T = 1.2\n",
    "t_axis = jnp.arange(0, T, 1/sr)\n",
    "params_true = jnp.array([eps1/1e8, eps2/1e8, beta1/1e3, beta2/1e3, C/1e8, delta/1e7])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def make_pulse_sequence(pulse_fun, arg_list):\n",
    "    def applicator(carry, args):\n",
    "       t = carry\n",
    "       return t, pulse_fun(t, *args)\n",
    "    \n",
    "    def scan_and_sum(t):\n",
    "        _, out = lax.scan(applicator, t, arg_list)\n",
    "        return jnp.sum(out)\n",
    "\n",
    "    return scan_and_sum"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "pfreq = 2 \n",
    "plocs = jnp.array([0.3, 0.9])\n",
    "pA = 0.025  # pressure amplitude (Volts)\n",
    "p0 = -0.005 # pressure DC offset (Volts)\n",
    "pwid = 0.08\n",
    "\n",
    "ppulse = lambda t, loc: pA * jnp.exp(-0.5 * (t - loc)**2/pwid**2) + p0\n",
    "\n",
    "P = jax.vmap(make_pulse_sequence(ppulse, (plocs,)))\n",
    "\n",
    "plt.plot(t_axis, P(t_axis))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def make_tension_pulse_fn(shape=1, scale=1, peak=1):\n",
    "    norm = peak * jnp.exp(shape - shape * jnp.log(shape) - shape * jnp.log(scale)) \n",
    "\n",
    "    fn = lambda t, loc: norm * jnp.exp((t - loc)/scale) * jnp.maximum((loc - t), 0)**shape \n",
    "\n",
    "    return fn\n",
    "    \n",
    "        \n",
    "kshape = 5  # shape parameter of gamma function\n",
    "kscale = 0.025  # rate parameter of gamma function (s)\n",
    "kpeak = 0.06  # peak value (Volts)\n",
    "klocs = jnp.array([0, 0.5, 1, 1.5])\n",
    "\n",
    "pulse = make_tension_pulse_fn(shape=kshape, scale=kscale, peak=kpeak)\n",
    "K = jax.vmap(make_pulse_sequence(pulse, (klocs,)))\n",
    "\n",
    "plt.plot(t_axis, K(t_axis))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "dfreq = 2 \n",
    "dlocs1 = jnp.arange(0.1, t_axis[-1], 1/dfreq)  # pressure pulse frequency (Hz)\n",
    "dlocs2 = jnp.arange(0.45, t_axis[-1], 1/dfreq)  # pressure pulse frequency (Hz)\n",
    "dlocs = jnp.sort(jnp.concatenate([dlocs1, dlocs2]))\n",
    "dA = jnp.array([0.05, 0.02, 0.01, 0.05, 0.03,])\n",
    "dwid = 0.01\n",
    "\n",
    "gpulse = lambda t, loc, amp: amp * jnp.exp(-0.5 * (t - loc)**2/dwid**2)\n",
    "\n",
    "D = jax.vmap(make_pulse_sequence(gpulse, (dlocs, dA)))\n",
    "\n",
    "plt.plot(t_axis, D(t_axis))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "plt.plot(t_axis, K(t_axis), label='vS tension')\n",
    "plt.plot(t_axis, D(t_axis), label='dTB tension')\n",
    "plt.plot(t_axis, P(t_axis), label='Pressure')\n",
    "plt.xlabel(\"time (s)\")\n",
    "plt.ylabel(\"envelope (V)\")\n",
    "plt.legend();"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Now let's define and integrate the ODE:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import diffrax"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def gradfun(t, y, args):\n",
    "    # params: eps1, eps2, beta1, beta2, C, delta\n",
    "    params, extra_args = args\n",
    "    K, D, P = extra_args\n",
    "    t_arr = jnp.array([t])\n",
    "    eps = (params[0] + params[1] * K(t_arr)) * 1e8\n",
    "    B = (params[2] + params[3] * P(t_arr)) * 1e3\n",
    "    C = params[4] * 1e8\n",
    "    D0 = params[5] * D(t_arr) * 1e7\n",
    "\n",
    "    xdot = y[1] \n",
    "    ydot = -eps * y[0] - C * y[0]**2 * y[1] + B * y[1] - D0\n",
    "\n",
    "    return jnp.array((xdot, ydot[0]))\n",
    "\n",
    "\n",
    "term = diffrax.ODETerm(gradfun)\n",
    "solver = diffrax.Dopri5()\n",
    "saveat = diffrax.SaveAt(ts=jnp.linspace(0, 1.2, int(sr)))\n",
    "stepsize_controller = diffrax.PIDController(rtol=1e-5, atol=1e-5)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "soln = diffrax.diffeqsolve(term, solver, t0=0, t1=T, dt0=0.5/sr, y0=jnp.array((0, 0)), saveat=saveat,\n",
    "                  stepsize_controller=stepsize_controller, args=(params_true, (K, D, P)), max_steps=int(1e6))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "plt.plot(soln.ts, soln.ys[:, 0])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "audio = jnp.interp(t_axis, soln.ts, soln.ys[:, 0])"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "If we wanted to save this generated data, we could do:\n",
    "```python\n",
    "import scipy.io as sio\n",
    "import numpy as np\n",
    "\n",
    "sio.wavfile.write('test.wav', int(sr), np.array(audio))\n",
    "```"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from scipy.signal import stft\n",
    "\n",
    "freqs, times, spec = stft(audio, fs=sr, \n",
    "                  nperseg=512, \n",
    "                  noverlap=480)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "fig, axs = plt.subplots(1, 2, figsize=(10, 4))\n",
    "axs[0].plot(t_axis, K(t_axis), label='vS tension')\n",
    "axs[0].plot(t_axis, D(t_axis), label='dTB tension')\n",
    "axs[0].plot(t_axis, P(t_axis), label='Pressure')\n",
    "axs[0].set_xlabel(\"time (s)\")\n",
    "axs[0].set_ylabel(\"envelope (V)\")\n",
    "axs[0].set_xlim([0, 1.2])\n",
    "axs[0].legend();\n",
    "\n",
    "axs[1].imshow(jnp.abs(spec), extent=[times[0], times[-1], freqs[0], freqs[-1]], aspect='auto')\n",
    "axs[1].set_xlabel('time (s)')\n",
    "axs[1].set_ylabel('frequency (Hz)')\n",
    "axs[1].set_ylim([12500, 19000])\n",
    "axs[1].set_xlim([0, 1.2])\n",
    "\n",
    "plt.tight_layout()"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Fit an ODE model to data\n",
    "\n",
    "Now we'll use data generated from the model above and see if we can recover the parameters of the model."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# split audio into short chunks and stack to make a dataset of snippets\n",
    "ys_full = jnp.stack([jnp.interp(t_axis, soln.ts, soln.ys[:, idx]) for idx in range(soln.ys.shape[-1])]).T\n",
    "\n",
    "# take derivative\n",
    "diff_order = 1\n",
    "ys_full_grads = jnp.diff(ys_full, diff_order, axis=0, prepend=jnp.zeros((diff_order, ys_full.shape[1]))) * sr\n",
    "T_snippet = 100  # samples\n",
    "ys = ys_full.reshape((-1, T_snippet, soln.ys.shape[-1]))\n",
    "ys_grads = ys_full_grads.reshape((-1, T_snippet, soln.ys.shape[-1]))\n",
    "ts = t_axis.reshape((-1, T_snippet))\n",
    "\n",
    "print(ys.shape, ts.shape)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "plt.plot(t_axis, ys_full_grads[:, 1])\n",
    "plt.xlim([0.3, 0.304])\n",
    "\n",
    "true_grads = jax.vmap(gradfun, in_axes=(0, 0, None))(t_axis, ys_full, (params_true, (K, D, P)))\n",
    "\n",
    "plt.figure()\n",
    "plt.plot(t_axis, true_grads[:, 1])\n",
    "plt.xlim([0.3, 0.304])\n",
    "\n",
    "plt.figure()\n",
    "plt.plot(t_axis, true_grads[:, 1] - ys_full_grads[:, 1])\n",
    "plt.xlim([0.3, 0.304])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# taken from https://docs.kidger.site/diffrax/examples/neural_ode/\n",
    "def dataloader(arrays, batch_size, *, key):\n",
    "    dataset_size = arrays[0].shape[0]\n",
    "    assert all(array.shape[0] == dataset_size for array in arrays)\n",
    "    indices = jnp.arange(dataset_size)\n",
    "    while True:\n",
    "        perm = jr.permutation(key, indices)\n",
    "        (key,) = jr.split(key, 1)\n",
    "        start = 0\n",
    "        end = batch_size\n",
    "        while end < dataset_size:\n",
    "            batch_perm = perm[start:end]\n",
    "            yield tuple(array[batch_perm] for array in arrays)\n",
    "            start = end\n",
    "            end = start + batch_size"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import equinox as eqx\n",
    "import optax\n",
    "import time\n",
    "from typing import Callable, Any"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "class NeuralODE(eqx.Module):\n",
    "    func: Callable\n",
    "    params: jax.Array\n",
    "    extra_args: Any\n",
    "\n",
    "    def __init__(self, gradfun, params, *args, **kwargs):\n",
    "        super().__init__(**kwargs)\n",
    "        self.func = gradfun\n",
    "        self.extra_args = args\n",
    "        self.params = params\n",
    "\n",
    "    def __call__(self, ts, y0):\n",
    "        solution = diffrax.diffeqsolve(\n",
    "            diffrax.ODETerm(self.func),\n",
    "            diffrax.Tsit5(),\n",
    "            t0=ts[0],\n",
    "            t1=ts[-1],\n",
    "            dt0=(ts[1] - ts[0]),\n",
    "            y0=y0,\n",
    "            stepsize_controller=diffrax.PIDController(rtol=1e-5, atol=1e-5),\n",
    "            saveat=diffrax.SaveAt(ts=ts),\n",
    "            args=(self.params, self.extra_args),\n",
    "            max_steps=int(1e6)\n",
    "        )\n",
    "        return solution.ys"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "class GradPredictor(eqx.Module):\n",
    "    func: Callable\n",
    "    params: jax.Array\n",
    "    extra_args: Any\n",
    "\n",
    "    def __init__(self, gradfun, params, *args, **kwargs):\n",
    "        super().__init__(**kwargs)\n",
    "        self.func = gradfun\n",
    "        self.extra_args = args\n",
    "        self.params = params\n",
    "    \n",
    "\n",
    "    def __call__(self, ts, ys):\n",
    "        grad_preds = jax.vmap(self.func, in_axes=(0, 0, None))(ts, ys, (self.params, self.extra_args))\n",
    "        return grad_preds"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from tensorboardX import SummaryWriter\n",
    "import uuid\n",
    "\n",
    "def main(\n",
    "    init,\n",
    "    batch_size=64,\n",
    "    lr_strategy=(3e-3, 3e-3),\n",
    "    steps_strategy=(500, 2500), \n",
    "    length_strategy=(0.1, 1),\n",
    "    seed=5678,\n",
    "    plot=True,\n",
    "    print_every=25,\n",
    "    train_grad=False,\n",
    "    runnum=None\n",
    "):\n",
    "    if not runnum:\n",
    "        runnum = uuid.uuid4()\n",
    "\n",
    "    key = jr.PRNGKey(seed)\n",
    "    data_key, model_key, loader_key = jr.split(key, 3)\n",
    "\n",
    "    _, length_size, data_size = ys.shape\n",
    "    \n",
    "    # eps1, eps2, beta1, beta2, C, delta\n",
    "    if train_grad:\n",
    "        model = GradPredictor(gradfun, init, K, D, P)\n",
    "    else:\n",
    "        model = NeuralODE(gradfun, init, K, D, P)\n",
    "\n",
    "    # Training loop like normal.\n",
    "    #\n",
    "    # Only thing to notice is that up until step 500 we train on only the first 10% of\n",
    "    # each time series. This is a standard trick to avoid getting caught in a local\n",
    "    # minimum.\n",
    "\n",
    "    @eqx.filter_value_and_grad\n",
    "    def grad_loss(model, ti, yi, dyi):\n",
    "        if train_grad:\n",
    "            y_pred = jax.vmap(model)(ti, yi)\n",
    "            return jnp.mean((dyi - y_pred) ** 2/sr**2)\n",
    "        else:\n",
    "            y_pred = jax.vmap(model)(ti, yi[:, 0])\n",
    "            return jnp.mean((yi[:, :, 0] - y_pred[:, :, 0]) ** 2)\n",
    "\n",
    "    @eqx.filter_jit\n",
    "    def make_step(ti, yi, dyi, model, opt_state):\n",
    "        loss, grads = grad_loss(model, ti, yi, dyi)\n",
    "        updates, opt_state = optim.update(grads, opt_state)\n",
    "        model = eqx.apply_updates(model, updates)\n",
    "        return loss, model, opt_state\n",
    "\n",
    "    writer = SummaryWriter(f\"logs/run{runnum}\", flush_secs=1)\n",
    "\n",
    "    globstep = 0\n",
    "    for lr, steps, length in zip(lr_strategy, steps_strategy, length_strategy):\n",
    "        optim = optax.adam(lr)\n",
    "        opt_state = optim.init(eqx.filter(model, eqx.is_inexact_array))\n",
    "        _ts = ts[:, : int(length_size * length)]\n",
    "        _ys = ys[:, : int(length_size * length)]\n",
    "        _dys = ys_grads[:, : int(length_size * length)]\n",
    "        for step, (ti, yi, dyi) in zip(\n",
    "            range(steps), dataloader((_ts, _ys, _dys), batch_size, key=loader_key)\n",
    "        ):\n",
    "            start = time.time()\n",
    "            loss, model, opt_state = make_step(ti, yi, dyi, model, opt_state)\n",
    "            end = time.time()\n",
    "            globstep += 1\n",
    "            if (step % print_every) == 0 or step == steps - 1:\n",
    "                print(f\"Step: {step}, Loss: {loss}, Computation time: {end - start}\")\n",
    "                writer.add_scalar('loss', loss, globstep)\n",
    "                writer.add_scalars('parameters', {\n",
    "                    'eps1': model.params[0]/params_true[0],\n",
    "                    'eps2': model.params[1]/params_true[1],\n",
    "                    'beta1': model.params[2]/params_true[2],\n",
    "                    'beta2': model.params[3]/params_true[3],\n",
    "                    'C': model.params[4]/params_true[4],\n",
    "                    'delta': model.params[5]/params_true[5],\n",
    "                }, globstep, end)\n",
    "\n",
    "    writer.close()\n",
    "    if plot:\n",
    "        plt.plot(ts[0], ys[0, :, 0], c=\"dodgerblue\", label=\"Real\")\n",
    "        plt.plot(ts[0], ys[0, :, 1], c=\"dodgerblue\")\n",
    "        if train_grad:\n",
    "            model_y = model(ts[0], ys[0])\n",
    "        else:\n",
    "            model_y = model(ts[0], ys[0, 0])\n",
    "        plt.plot(ts[0], model_y[:, 0], c=\"crimson\", label=\"Model\")\n",
    "        plt.plot(ts[0], model_y[:, 1], c=\"crimson\")\n",
    "        plt.legend()\n",
    "        plt.tight_layout()\n",
    "        plt.savefig(\"figs/neural_ode.png\")\n",
    "        plt.show()\n",
    "\n",
    "    return ts, ys, model"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "key = jr.PRNGKey(12345)\n",
    "params0 = jnp.array([1., 100., 1, 1000., 1, 1])\n",
    "ts, ys, model = main(init=params0, train_grad=False, steps_strategy=(500, 40000), \n",
    "                     lr_strategy=(1e-3, 1e-3), batch_size=64, print_every=100)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "params_true, params0, model.params"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "plt.plot(t_axis, ys_full[:, 0])\n",
    "\n",
    "plt.figure()\n",
    "y_pred = model(t_axis, ys_full[0])\n",
    "plt.plot(t_axis, y_pred[:, 0])\n",
    "\n",
    "plt.figure()\n",
    "y_pred = model(t_axis, ys_full[0])\n",
    "plt.plot(t_axis, y_pred[:, 0] - 1 * ys_full[:, 0])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "freqs, times, spec_pred = stft(y_pred[:, 0], fs=sr, \n",
    "                  nperseg=512, \n",
    "                  noverlap=480)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "fig, axs = plt.subplots(3, 1, figsize=(5, 12))\n",
    "axs[0].plot(t_axis, K(t_axis), label='vS tension')\n",
    "axs[0].plot(t_axis, D(t_axis), label='dTB tension')\n",
    "axs[0].plot(t_axis, P(t_axis), label='Pressure')\n",
    "axs[0].set_xlabel(\"time (s)\")\n",
    "axs[0].set_ylabel(\"envelope (V)\")\n",
    "axs[0].set_xlim([0, 1.2])\n",
    "axs[0].legend();\n",
    "\n",
    "minval = jnp.real(jnp.minimum(\n",
    "    jnp.min(jnp.abs(spec)**2), \n",
    "    jnp.min(jnp.abs(spec_pred)**2)))\n",
    "maxval = jnp.real(jnp.maximum(\n",
    "    jnp.max(jnp.abs(spec)**2), \n",
    "    jnp.max(jnp.abs(spec_pred)**2)))\n",
    "\n",
    "axs[1].imshow(jnp.abs(spec)**2, extent=[times[0], times[-1], freqs[0], freqs[-1]], \n",
    "              aspect='auto',\n",
    "              vmin=minval,\n",
    "              vmax=maxval)\n",
    "axs[1].set_xlabel('time (s)')\n",
    "axs[1].set_ylabel('frequency (Hz)')\n",
    "axs[1].set_ylim([12500, 19000])\n",
    "axs[1].set_xlim([0, 1.2])\n",
    "curr_fontsize = plt.rcParams['font.size']\n",
    "axs[1].annotate('Original', (0.025, 0.9), \n",
    "                xycoords='axes fraction', \n",
    "                color='white',\n",
    "                fontsize=2 * curr_fontsize)\n",
    "\n",
    "axs[2].imshow(jnp.abs(spec_pred)**2, extent=[times[0], times[-1], freqs[0], freqs[-1]], \n",
    "              aspect='auto',\n",
    "              vmin=minval,\n",
    "              vmax=maxval)\n",
    "axs[2].set_xlabel('time (s)')\n",
    "axs[2].set_ylabel('frequency (Hz)')\n",
    "axs[2].set_ylim([12500, 19000])\n",
    "axs[2].set_xlim([0, 1.2])\n",
    "axs[2].annotate('Fitted', (0.025, 0.9), \n",
    "                xycoords='axes fraction', \n",
    "                color='white',\n",
    "                fontsize=2 * curr_fontsize)\n",
    "\n",
    "plt.tight_layout()\n",
    "plt.savefig('figs/syrinx_recon.svg')"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "syrinx",
   "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.10.12"
  },
  "orig_nbformat": 4
 },
 "nbformat": 4,
 "nbformat_minor": 2
}