Added a streaming spectrogram function.

cpp/zimt/zimtohrli.cc

@@ -266,6 +266,55 @@ Distance Zimtohrli::Distance(
   }
 }
 
+hwy::AlignedNDArray<float, 2> Zimtohrli::StreamingSpectrogram(
+    hwy::Span<const float> signal) {
+  hwy::AlignedNDArray<float, 2> chunk_energy_db_buffer(
+      {1, cam_filterbank->filter.Size()});
+  hwy::AlignedNDArray<float, 2> chunk_partial_energy_db_buffer(
+      {1, cam_filterbank->filter.Size()});
+  hwy::AlignedNDArray<float, 2> chunk_spectrogram_buffer(
+      {1, cam_filterbank->filter.Size()});
+  const size_t samples_per_chunk =
+      static_cast<size_t>(cam_filterbank->sample_rate / perceptual_sample_rate);
+  hwy::AlignedNDArray<float, 2> chunk_channels_buffer(
+      {samples_per_chunk, cam_filterbank->filter.Size()});
+  hwy::AlignedNDArray<float, 2> spectrogram(
+      {static_cast<size_t>(signal.size() / samples_per_chunk),
+       cam_filterbank->filter.Size()});
+  FilterbankState filter_state = cam_filterbank->filter.NewState();
+  for (size_t step = 0; (step + 1) * samples_per_chunk < signal.size();
+       ++step) {
+    cam_filterbank->filter.Filter(
+        hwy::Span<const float>(signal.data() + step * samples_per_chunk,
+                               samples_per_chunk),
+        filter_state, chunk_channels_buffer);
+    ComputeEnergy(chunk_channels_buffer, chunk_energy_db_buffer);
+    ToDb(chunk_energy_db_buffer, full_scale_sine_db, epsilon,
+         chunk_energy_db_buffer);
+    if (apply_masking) {
+      masking.CutFullyMasked(chunk_energy_db_buffer, cam_filterbank->cam_delta,
+                             chunk_partial_energy_db_buffer);
+    } else {
+      hwy::CopyBytes(chunk_energy_db_buffer.data(),
+                     chunk_partial_energy_db_buffer.data(),
+                     chunk_energy_db_buffer.memory_size() * sizeof(float));
+    }
+    if (apply_loudness) {
+      loudness.PhonsFromSPL(chunk_partial_energy_db_buffer,
+                            cam_filterbank->thresholds_hz,
+                            chunk_spectrogram_buffer);
+    } else {
+      hwy::CopyBytes(
+          chunk_partial_energy_db_buffer.data(),
+          chunk_spectrogram_buffer.data(),
+          chunk_partial_energy_db_buffer.memory_size() * sizeof(float));
+    }
+    hwy::CopyBytes(chunk_spectrogram_buffer.data(), spectrogram[{step}].data(),
+                   chunk_spectrogram_buffer.memory_size() * sizeof(float));
+  }
+  return spectrogram;
+}
+
 void Zimtohrli::Spectrogram(
     hwy::Span<const float> signal, FilterbankState& state,
     hwy::AlignedNDArray<float, 2>& channels,

cpp/zimt/zimtohrli.h

@@ -240,6 +240,16 @@ struct Zimtohrli {
       bool verbose, const hwy::AlignedNDArray<float, 2>& spectrogram_a,
       const hwy::AlignedNDArray<float, 2>& spectrogram_b) const;
 
+  // Memory thrifty spectrogram calculation that computes one chunk (samples per
+  // perceptual sample rate period) at a time.
+  //
+  // signal is a span of audio samples between -1 and 1.
+  //
+  // Returns a (num_downscaled_samples, num_channels)-shaped array of Phons
+  // values reprecenting the perceptual intensity of each channel.
+  hwy::AlignedNDArray<float, 2> StreamingSpectrogram(
+      hwy::Span<const float> signal);
+
   // Convenience method to analyze a signal.
   //
   // Allocates an Analysis instance, and executes Spectrogram on it along with

cpp/zimt/zimtohrli_test.cc

@@ -494,6 +494,60 @@ void BM_SpectrogramDistanceVsSeconds(benchmark::State& state) {
 }
 BENCHMARK_RANGE(BM_SpectrogramDistanceVsSeconds, 1, 8);
 
+void BM_StreamingSpectrogram(benchmark::State& state) {
+  const size_t sample_rate = 48000;
+  Zimtohrli z{.cam_filterbank = Cam{}.CreateFilterbank(sample_rate)};
+  hwy::AlignedNDArray<float, 1> signal(
+      {static_cast<size_t>(sample_rate * state.range(0))});
+  for (auto s : state) {
+    z.StreamingSpectrogram(signal[{}]);
+  }
+  state.SetItemsProcessed(signal.size() * state.iterations());
+}
+BENCHMARK_RANGE(BM_StreamingSpectrogram, 1, 32);
+
+void BM_PreallocSpectrogram(benchmark::State& state) {
+  const size_t sample_rate = 48000;
+  Zimtohrli z{.cam_filterbank = Cam{}.CreateFilterbank(sample_rate)};
+  hwy::AlignedNDArray<float, 1> signal(
+      {static_cast<size_t>(sample_rate * state.range(0))});
+  hwy::AlignedNDArray<float, 2> channels(
+      {signal.shape()[1], z.cam_filterbank->filter.Size()});
+  hwy::AlignedNDArray<float, 2> energy_channels_db(
+      {static_cast<size_t>(100 * state.range(0)), channels.shape()[1]});
+  hwy::AlignedNDArray<float, 2> partial_energy_channels_db(
+      {static_cast<size_t>(100 * state.range(0)), channels.shape()[1]});
+  hwy::AlignedNDArray<float, 2> spectrogram(
+      {energy_channels_db.shape()[0], energy_channels_db.shape()[1]});
+  for (auto s : state) {
+    z.Spectrogram(signal[{}], channels, energy_channels_db,
+                  partial_energy_channels_db, spectrogram);
+  }
+  state.SetItemsProcessed(signal.size() * state.iterations());
+}
+BENCHMARK_RANGE(BM_PreallocSpectrogram, 1, 32);
+
+void BM_RepeatedAllocSpectrogram(benchmark::State& state) {
+  const size_t sample_rate = 48000;
+  Zimtohrli z{.cam_filterbank = Cam{}.CreateFilterbank(sample_rate)};
+  hwy::AlignedNDArray<float, 1> signal(
+      {static_cast<size_t>(sample_rate * state.range(0))});
+  for (auto s : state) {
+    hwy::AlignedNDArray<float, 2> channels(
+        {signal.shape()[1], z.cam_filterbank->filter.Size()});
+    hwy::AlignedNDArray<float, 2> energy_channels_db(
+        {static_cast<size_t>(100 * state.range(0)), channels.shape()[1]});
+    hwy::AlignedNDArray<float, 2> partial_energy_channels_db(
+        {static_cast<size_t>(100 * state.range(0)), channels.shape()[1]});
+    hwy::AlignedNDArray<float, 2> spectrogram(
+        {energy_channels_db.shape()[0], energy_channels_db.shape()[1]});
+    z.Spectrogram(signal[{}], channels, energy_channels_db,
+                  partial_energy_channels_db, spectrogram);
+  }
+  state.SetItemsProcessed(signal.size() * state.iterations());
+}
+BENCHMARK_RANGE(BM_RepeatedAllocSpectrogram, 1, 32);
+
 void BM_SpectrogramDistanceVsResolution(benchmark::State& state) {
   const size_t sample_rate = 48000;
   const float seconds_of_audio = 1;