Spectrogram Plotting

Mix.install([
  {:nx_signal, "~> 0.2"},
  {:vega_lite, "~> 0.1.4"},
  {:kino_vega_lite, "~> 0.1.1"}
])

Generating the audio data

# You can load an audio file here.  For this example,
# we're producing 3 seconds of 220Hz, 440Hz, 1kHz and 3kHz sine waves

fs = 44.1e3
t_max = 3

full_n = ceil(fs * t_max)
half_n = div(full_n, 2)

# samples/sec * sec = samples
n = Nx.iota({half_n})

sin = fn freq, n ->
  Nx.sin(Nx.multiply(2 * :math.pi() * freq / fs, n))
end

sin220 = sin.(220, n)
sin440 = sin.(440, n)
sin1000 = sin.(1000, n)
sin3000 = sin.(3000, n)

data = Nx.concatenate([Nx.add(sin440, sin1000), Nx.add(sin220, sin3000)])
n = Nx.iota({full_n})

d = %{data: Nx.to_flat_list(data[[1000..1250]]), n: Nx.to_flat_list(n[[1000..1250]])}

VegaLite.new(width: 600, height: 600, title: "Audio Sample")
|> VegaLite.data_from_values(d, only: ["n", "data"])
|> VegaLite.mark(:line)
|> VegaLite.encode_field(:x, "n", type: :ordinal)
|> VegaLite.encode_field(:y, "data", type: :quantitative)

defmodule Spectrogram do
  alias VegaLite, as: Vl
  import Nx.Defn

  def calculate_stft_and_plot_spectrogram(
        input,
        fs,
        window_duration_ms,
        plot_cutoff_frequency \\ 4000
      ) do
    n_window = ceil(fs * window_duration_ms)
    {spectrogram, f, t, max_f} = stft(input, fs: fs, n_window: n_window)

    max_f = Nx.to_number(max_f)
    spectrogram = Nx.slice(spectrogram, [0, 0], [Nx.size(t), max_f])
    f = Nx.slice(f, [0], [max_f])

    spectrogram
    |> to_plot_data(f, t, plot_cutoff_frequency)
    |> plot()
  end

  defn stft(input, opts) do
    fs = opts[:fs]
    n_window = opts[:n_window]

    # ms to samples
    window = NxSignal.Windows.hann(n: n_window, is_periodic: true)

    # use the default overlap of 50%
    {s, t, f} = NxSignal.stft(input, window, sampling_rate: fs, fft_length: 1024)

    max_f =
      Nx.select(f >= fs / 2, Nx.iota(f.shape), Nx.size(f) + 1)
      |> Nx.argmin()

    spectrogram = Nx.abs(s)
    # to dBFS
    spectrogram = 20 * Nx.log(spectrogram / Nx.reduce_max(spectrogram)) / Nx.log(10)

    {spectrogram, f, t, max_f}
  end

  defp to_plot_data(s, f, t, plot_cutoff_frequency) do
    for t_idx <- 0..(Nx.size(t) - 1),
        f_idx <- 0..(Nx.size(f) - 1),
        Nx.to_number(f[[f_idx]]) <= plot_cutoff_frequency,
        reduce: %{"t" => [], "f" => [], "s" => []} do
      %{"t" => t_acc, "f" => f_acc, "s" => s_acc} ->
        %{
          "t" => [Nx.to_number(t[[t_idx]]) | t_acc],
          "f" => [Float.round(Nx.to_number(f[[f_idx]]), 3) | f_acc],
          "s" => [Nx.to_number(s[[t_idx, f_idx]]) | s_acc]
        }
    end
  end

  defp plot(dataset) do
    Vl.new(title: "Spectrogram", width: 500, height: 500)
    |> Vl.mark(:rect)
    |> Vl.data_from_values(dataset)
    |> Vl.encode_field(:x, "t",
      type: :quantitative,
      title: "Time (seconds)",
      axis: [tick_min_step: 0.1],
      grid: false
    )
    |> Vl.encode_field(:y, "f",
      type: :quantitative,
      sort: "-x",
      title: "Frequency (Hz)",
      axis: [tick_count: 25],
      grid: false
    )
    |> Vl.encode_field(:color, "s",
      aggregate: :max,
      type: :quantitative,
      scale: [scheme: "viridis"],
      legend: [title: "dBFS"]
    )
    |> Vl.config(view: [stroke: nil])
  end
end

Spectrogram.calculate_stft_and_plot_spectrogram(data, fs, 50.0e-3)

Notice how the first half of the spectrogram looks cleaner than the second one. This is due to the window length, that also interferes in how we can observe both our time and frequency resolutions.

Below we can see what happens if we use different window durations (150ms, 100ms and 25ms respectively).

Spectrogram.calculate_stft_and_plot_spectrogram(data, fs, 150.0e-3)

Spectrogram.calculate_stft_and_plot_spectrogram(data, fs, 100.0e-3)

Spectrogram.calculate_stft_and_plot_spectrogram(data, fs, 25.0e-3)