408 lines
18 KiB
C++
408 lines
18 KiB
C++
#include "doa_algorithms.hpp"
|
|
|
|
#include <algorithm>
|
|
#include <cmath>
|
|
#include <complex>
|
|
#include <limits>
|
|
#include <stdexcept>
|
|
#include <utility>
|
|
|
|
namespace esp32_doa {
|
|
namespace {
|
|
|
|
constexpr double kPi = 3.141592653589793238462643383279502884;
|
|
constexpr double kTolerance = 1e-14;
|
|
|
|
std::size_t index(const Stft &s, int snapshot, int frequency, int channel) {
|
|
return (static_cast<std::size_t>(snapshot) * s.bins + frequency) *
|
|
s.channels + channel;
|
|
}
|
|
|
|
void fft(std::vector<Complex> &x, bool inverse) {
|
|
const std::size_t n = x.size();
|
|
if (n == 0 || (n & (n - 1)) != 0) {
|
|
throw std::invalid_argument("nfft must be a power of two");
|
|
}
|
|
for (std::size_t i = 1, j = 0; i < n; ++i) {
|
|
std::size_t bit = n >> 1;
|
|
for (; j & bit; bit >>= 1) j ^= bit;
|
|
j ^= bit;
|
|
if (i < j) std::swap(x[i], x[j]);
|
|
}
|
|
for (std::size_t len = 2; len <= n; len <<= 1) {
|
|
const double sign = inverse ? 1.0 : -1.0;
|
|
const Complex root = std::polar(1.0, sign * 2.0 * kPi / len);
|
|
for (std::size_t start = 0; start < n; start += len) {
|
|
Complex w(1.0, 0.0);
|
|
const std::size_t half = len >> 1;
|
|
for (std::size_t j = 0; j < half; ++j) {
|
|
const Complex u = x[start + j];
|
|
const Complex v = x[start + j + half] * w;
|
|
x[start + j] = u + v;
|
|
x[start + j + half] = u - v;
|
|
w *= root;
|
|
}
|
|
}
|
|
}
|
|
if (inverse) {
|
|
for (Complex &v : x) v /= static_cast<double>(n);
|
|
}
|
|
}
|
|
|
|
double hann(int n, int length) {
|
|
if (length <= 1) return 1.0;
|
|
return 0.5 - 0.5 * std::cos(2.0 * kPi * n / (length - 1));
|
|
}
|
|
|
|
std::vector<int> frequency_bins(const Config &c) {
|
|
const int first = static_cast<int>(std::floor(
|
|
c.min_frequency_hz * c.nfft / c.sample_rate + 0.5));
|
|
const int last = static_cast<int>(std::floor(
|
|
c.max_frequency_hz * c.nfft / c.sample_rate + 0.5));
|
|
std::vector<int> bins;
|
|
for (int k = first; k < last && k <= c.nfft / 2; ++k)
|
|
if (k >= 0) bins.push_back(k);
|
|
return bins;
|
|
}
|
|
|
|
void validate(const Config &c) {
|
|
if (c.sample_rate <= 0 || c.channels <= 0 || c.nfft <= 0 ||
|
|
(c.nfft & (c.nfft - 1)) != 0 || c.hop <= 0 || c.hop > c.nfft)
|
|
throw std::invalid_argument("invalid STFT configuration");
|
|
if (static_cast<int>(c.microphone_azimuths_deg.size()) != c.channels)
|
|
throw std::invalid_argument("one microphone azimuth is required per channel");
|
|
if (!(c.min_frequency_hz >= 0.0 && c.min_frequency_hz < c.max_frequency_hz &&
|
|
c.max_frequency_hz < c.sample_rate / 2.0))
|
|
throw std::invalid_argument("invalid frequency range");
|
|
}
|
|
|
|
std::vector<std::pair<double, double>> microphone_positions(const Config &c) {
|
|
std::vector<std::pair<double, double>> p(c.channels);
|
|
for (int m = 0; m < c.channels; ++m) {
|
|
const double a = c.microphone_azimuths_deg[m] * kPi / 180.0;
|
|
p[m] = {c.radius_m * std::cos(a), c.radius_m * std::sin(a)};
|
|
}
|
|
return p;
|
|
}
|
|
|
|
double wrap_degrees(double angle) {
|
|
angle = std::fmod(angle, 360.0);
|
|
return angle < 0.0 ? angle + 360.0 : angle;
|
|
}
|
|
|
|
std::vector<Complex> mode_vector(const Config &c, int frequency,
|
|
double angle_degrees) {
|
|
const auto positions = microphone_positions(c);
|
|
const double a = angle_degrees * kPi / 180.0;
|
|
const double dx = std::cos(a), dy = std::sin(a);
|
|
const double omega = 2.0 * kPi * c.sample_rate * frequency / c.nfft;
|
|
std::vector<Complex> mode(c.channels);
|
|
for (int m = 0; m < c.channels; ++m) {
|
|
const double tau = (dx * positions[m].first + dy * positions[m].second) /
|
|
c.sound_speed_m_s;
|
|
mode[m] = std::polar(1.0, omega * tau);
|
|
}
|
|
return mode;
|
|
}
|
|
|
|
std::vector<Complex> solve_linear(std::vector<Complex> a, std::vector<Complex> b,
|
|
int n) {
|
|
for (int col = 0; col < n; ++col) {
|
|
int pivot = col;
|
|
double best = std::abs(a[col * n + col]);
|
|
for (int row = col + 1; row < n; ++row) {
|
|
const double value = std::abs(a[row * n + col]);
|
|
if (value > best) best = value, pivot = row;
|
|
}
|
|
if (best < 1e-20) throw std::runtime_error("singular covariance matrix");
|
|
if (pivot != col) {
|
|
for (int j = col; j < n; ++j)
|
|
std::swap(a[col * n + j], a[pivot * n + j]);
|
|
std::swap(b[col], b[pivot]);
|
|
}
|
|
const Complex diagonal = a[col * n + col];
|
|
for (int j = col; j < n; ++j) a[col * n + j] /= diagonal;
|
|
b[col] /= diagonal;
|
|
for (int row = 0; row < n; ++row) {
|
|
if (row == col) continue;
|
|
const Complex factor = a[row * n + col];
|
|
if (std::abs(factor) == 0.0) continue;
|
|
for (int j = col; j < n; ++j) a[row * n + j] -= factor * a[col * n + j];
|
|
b[row] -= factor * b[col];
|
|
}
|
|
}
|
|
return b;
|
|
}
|
|
|
|
struct EigenResult {
|
|
std::vector<double> values;
|
|
std::vector<Complex> vectors; // column-major eigenvectors
|
|
};
|
|
|
|
// Complex Hermitian Jacobi eigensolver. This small implementation is enough
|
|
// for the fixed 8-microphone covariance matrices and avoids BLAS/LAPACK.
|
|
EigenResult hermitian_eigen(const std::vector<Complex> &input, int n) {
|
|
std::vector<Complex> a = input;
|
|
std::vector<Complex> v(static_cast<std::size_t>(n) * n, Complex(0.0));
|
|
for (int i = 0; i < n; ++i) v[i * n + i] = 1.0;
|
|
for (int iteration = 0; iteration < 80 * n * n; ++iteration) {
|
|
int p = 0, q = 1;
|
|
double largest = 0.0;
|
|
for (int i = 0; i < n; ++i)
|
|
for (int j = i + 1; j < n; ++j)
|
|
if (std::abs(a[i * n + j]) > largest)
|
|
largest = std::abs(a[i * n + j]), p = i, q = j;
|
|
if (largest < 1e-12) break;
|
|
const Complex b = a[p * n + q];
|
|
const double phase = std::arg(b);
|
|
const double magnitude = std::abs(b);
|
|
const double theta = 0.5 * std::atan2(2.0 * magnitude,
|
|
a[q * n + q].real() - a[p * n + p].real());
|
|
const double cs = std::cos(theta), sn = std::sin(theta);
|
|
const Complex ph = std::polar(1.0, phase);
|
|
const double app = a[p * n + p].real(), aqq = a[q * n + q].real();
|
|
for (int k = 0; k < n; ++k) {
|
|
if (k == p || k == q) continue;
|
|
const Complex akp = a[k * n + p], akq = a[k * n + q];
|
|
a[k * n + p] = akp * ph * cs - akq * sn;
|
|
a[p * n + k] = std::conj(a[k * n + p]);
|
|
a[k * n + q] = akp * ph * sn + akq * cs;
|
|
a[q * n + k] = std::conj(a[k * n + q]);
|
|
}
|
|
a[p * n + p] = app * cs * cs + aqq * sn * sn - 2.0 * magnitude * cs * sn;
|
|
a[q * n + q] = app * sn * sn + aqq * cs * cs + 2.0 * magnitude * cs * sn;
|
|
a[p * n + q] = a[q * n + p] = Complex(0.0);
|
|
for (int k = 0; k < n; ++k) {
|
|
const Complex vkp = v[k * n + p], vkq = v[k * n + q];
|
|
v[k * n + p] = vkp * ph * cs - vkq * sn;
|
|
v[k * n + q] = vkp * ph * sn + vkq * cs;
|
|
}
|
|
}
|
|
std::vector<int> order(n);
|
|
for (int i = 0; i < n; ++i) order[i] = i;
|
|
std::sort(order.begin(), order.end(), [&](int x, int y) {
|
|
return a[x * n + x].real() < a[y * n + y].real();
|
|
});
|
|
EigenResult result;
|
|
result.values.resize(n);
|
|
result.vectors.resize(static_cast<std::size_t>(n) * n);
|
|
for (int col = 0; col < n; ++col) {
|
|
result.values[col] = a[order[col] * n + order[col]].real();
|
|
for (int row = 0; row < n; ++row)
|
|
result.vectors[row * n + col] = v[row * n + order[col]];
|
|
}
|
|
return result;
|
|
}
|
|
|
|
double spectrum_sharpness(const std::vector<double> &spectrum) {
|
|
if (spectrum.empty()) return 0.0;
|
|
std::vector<double> sorted = spectrum;
|
|
std::sort(sorted.begin(), sorted.end());
|
|
const double peak = sorted.back();
|
|
const double floor = sorted[sorted.size() / 2];
|
|
const double spread = peak - sorted.front();
|
|
return std::clamp((peak - floor) / (spread + 1e-12), 0.0, 1.0);
|
|
}
|
|
|
|
DoaResult finish_spectrum(std::vector<double> spectrum) {
|
|
auto peak = std::max_element(spectrum.begin(), spectrum.end());
|
|
const int index = static_cast<int>(peak - spectrum.begin());
|
|
DoaResult result;
|
|
result.angle_degrees = wrap_degrees(360.0 * index / spectrum.size());
|
|
result.peak_sharpness = spectrum_sharpness(spectrum);
|
|
result.spectrum = std::move(spectrum);
|
|
return result;
|
|
}
|
|
|
|
} // namespace
|
|
|
|
Complex &Stft::at(int snapshot, int frequency, int channel) {
|
|
return spectra[index(*this, snapshot, frequency, channel)];
|
|
}
|
|
const Complex &Stft::at(int snapshot, int frequency, int channel) const {
|
|
return spectra[index(*this, snapshot, frequency, channel)];
|
|
}
|
|
|
|
Stft analyze_stft(const std::vector<double> &audio, int sample_count,
|
|
const Config &config) {
|
|
validate(config);
|
|
if (sample_count < config.nfft ||
|
|
static_cast<std::size_t>(sample_count) * config.channels > audio.size())
|
|
throw std::invalid_argument("audio does not contain enough samples");
|
|
Stft result;
|
|
result.sample_count = sample_count;
|
|
result.snapshots = (sample_count - config.nfft) / config.hop + 1;
|
|
result.bins = config.nfft / 2 + 1;
|
|
result.channels = config.channels;
|
|
result.spectra.assign(static_cast<std::size_t>(result.snapshots) * result.bins *
|
|
result.channels,
|
|
Complex(0.0));
|
|
std::vector<double> means(config.channels, 0.0);
|
|
for (int n = 0; n < sample_count; ++n)
|
|
for (int m = 0; m < config.channels; ++m)
|
|
means[m] += audio[static_cast<std::size_t>(n) * config.channels + m] /
|
|
sample_count;
|
|
for (int s = 0; s < result.snapshots; ++s) {
|
|
const int start = s * config.hop;
|
|
for (int m = 0; m < config.channels; ++m) {
|
|
std::vector<Complex> frame(config.nfft);
|
|
for (int n = 0; n < config.nfft; ++n)
|
|
frame[n] = (audio[static_cast<std::size_t>(start + n) * config.channels + m] -
|
|
means[m]) * hann(n, config.nfft);
|
|
fft(frame, false);
|
|
for (int k = 0; k < result.bins; ++k) result.at(s, k, m) = frame[k];
|
|
}
|
|
}
|
|
return result;
|
|
}
|
|
|
|
DoaResult srp_phat(const Stft &stft, const Config &config) {
|
|
validate(config);
|
|
const std::vector<int> bins = frequency_bins(config);
|
|
const int pairs = config.channels * (config.channels - 1) / 2;
|
|
std::vector<double> response(config.grid_points, 0.0);
|
|
const auto positions = microphone_positions(config);
|
|
for (int g = 0; g < config.grid_points; ++g) {
|
|
const double angle = 360.0 * g / config.grid_points * kPi / 180.0;
|
|
const double dx = std::cos(angle), dy = std::sin(angle);
|
|
double score = 0.0;
|
|
for (int k : bins) {
|
|
std::vector<std::vector<Complex>> cc(config.channels,
|
|
std::vector<Complex>(config.channels));
|
|
for (int i = 0; i < config.channels; ++i)
|
|
for (int j = 0; j < config.channels; ++j)
|
|
for (int s = 0; s < stft.snapshots; ++s) {
|
|
Complex xi = stft.at(s, k, i), xj = stft.at(s, k, j);
|
|
const double ai = std::max(std::abs(xi), kTolerance);
|
|
const double aj = std::max(std::abs(xj), kTolerance);
|
|
cc[i][j] += (xi / ai) * std::conj(xj / aj);
|
|
}
|
|
const double omega = 2.0 * kPi * config.sample_rate * k / config.nfft;
|
|
for (int i = 0; i < config.channels; ++i) {
|
|
const double ti = (dx * positions[i].first + dy * positions[i].second) /
|
|
config.sound_speed_m_s;
|
|
for (int j = i + 1; j < config.channels; ++j) {
|
|
const double tj = (dx * positions[j].first + dy * positions[j].second) /
|
|
config.sound_speed_m_s;
|
|
const Complex mode = std::polar(1.0, omega * (tj - ti));
|
|
score += 2.0 * std::real(cc[i][j] * mode);
|
|
}
|
|
}
|
|
}
|
|
const double dc = static_cast<double>(stft.snapshots * config.channels * bins.size());
|
|
response[g] = (score + dc) /
|
|
(stft.snapshots * bins.size() * static_cast<double>(pairs));
|
|
}
|
|
return finish_spectrum(std::move(response));
|
|
}
|
|
|
|
DoaResult norm_music(const Stft &stft, const Config &config) {
|
|
validate(config);
|
|
const std::vector<int> bins = frequency_bins(config);
|
|
std::vector<double> per_frequency(static_cast<std::size_t>(bins.size()) *
|
|
config.grid_points, 0.0);
|
|
for (std::size_t frequency_index = 0; frequency_index < bins.size(); ++frequency_index) {
|
|
const int k = bins[frequency_index];
|
|
std::vector<Complex> covariance(static_cast<std::size_t>(config.channels) * config.channels,
|
|
Complex(0.0));
|
|
for (int i = 0; i < config.channels; ++i)
|
|
for (int j = 0; j < config.channels; ++j)
|
|
for (int s = 0; s < stft.snapshots; ++s)
|
|
covariance[i * config.channels + j] +=
|
|
stft.at(s, k, i) * std::conj(stft.at(s, k, j)) / stft.snapshots;
|
|
const EigenResult eig = hermitian_eigen(covariance, config.channels);
|
|
for (int g = 0; g < config.grid_points; ++g) {
|
|
const auto mode = mode_vector(config, k, 360.0 * g / config.grid_points);
|
|
Complex denom(0.0);
|
|
for (int i = 0; i < config.channels; ++i) {
|
|
for (int j = 0; j < config.channels; ++j) {
|
|
Complex cross = (i == j ? Complex(1.0) : Complex(0.0));
|
|
cross -= eig.vectors[i * config.channels + config.channels - 1] *
|
|
std::conj(eig.vectors[j * config.channels + config.channels - 1]);
|
|
denom += std::conj(mode[i]) * cross * mode[j];
|
|
}
|
|
}
|
|
per_frequency[frequency_index * config.grid_points + g] =
|
|
1.0 / std::max(std::abs(denom), 1e-20);
|
|
}
|
|
}
|
|
std::vector<double> response(config.grid_points, 0.0);
|
|
// Pyroomacoustics' NormMUSIC normalizes along the frequency axis for
|
|
// every candidate direction (Pssl / max(Pssl, axis=0)).
|
|
for (int g = 0; g < config.grid_points; ++g) {
|
|
double maximum = 0.0;
|
|
for (std::size_t f = 0; f < bins.size(); ++f)
|
|
maximum = std::max(maximum, per_frequency[f * config.grid_points + g]);
|
|
for (std::size_t f = 0; f < bins.size(); ++f)
|
|
response[g] += per_frequency[f * config.grid_points + g] /
|
|
std::max(maximum, 1e-20);
|
|
response[g] /= bins.size();
|
|
}
|
|
return finish_spectrum(std::move(response));
|
|
}
|
|
|
|
MvdrResult mvdr(const Stft &stft, const Config &config, double angle_degrees) {
|
|
validate(config);
|
|
const std::vector<int> bins = frequency_bins(config);
|
|
const int n = config.channels;
|
|
std::vector<Complex> weights(static_cast<std::size_t>(stft.bins) * n);
|
|
for (int k = 0; k < stft.bins; ++k) {
|
|
const auto steering = mode_vector(config, k, angle_degrees);
|
|
for (int m = 0; m < n; ++m) weights[k * n + m] = steering[m] / n;
|
|
}
|
|
for (int k : bins) {
|
|
std::vector<Complex> covariance(static_cast<std::size_t>(n) * n, Complex(0.0));
|
|
for (int i = 0; i < n; ++i)
|
|
for (int j = 0; j < n; ++j)
|
|
for (int s = 0; s < stft.snapshots; ++s)
|
|
covariance[i * n + j] += stft.at(s, k, i) * std::conj(stft.at(s, k, j)) /
|
|
stft.snapshots;
|
|
double power = 0.0;
|
|
for (int i = 0; i < n; ++i) power += covariance[i * n + i].real();
|
|
const double loading = config.diagonal_loading * std::max(power / n, 1e-12);
|
|
for (int i = 0; i < n; ++i) covariance[i * n + i] += loading;
|
|
std::vector<Complex> solved = solve_linear(covariance, steering, n);
|
|
Complex denominator(0.0);
|
|
for (int m = 0; m < n; ++m) denominator += std::conj(steering[m]) * solved[m];
|
|
if (std::abs(denominator) < 1e-12) denominator = Complex(1e-12);
|
|
for (int m = 0; m < n; ++m) weights[k * n + m] = solved[m] / denominator;
|
|
}
|
|
std::vector<double> output(stft.sample_count, 0.0), normalization(stft.sample_count, 0.0);
|
|
for (int s = 0; s < stft.snapshots; ++s) {
|
|
std::vector<Complex> spectrum(config.nfft, Complex(0.0));
|
|
for (int k = 0; k < stft.bins; ++k) {
|
|
Complex value(0.0);
|
|
for (int m = 0; m < n; ++m) value += std::conj(weights[k * n + m]) * stft.at(s, k, m);
|
|
spectrum[k] = value;
|
|
if (k > 0 && k < config.nfft / 2) spectrum[config.nfft - k] = std::conj(value);
|
|
}
|
|
fft(spectrum, true);
|
|
const int start = s * config.hop;
|
|
for (int i = 0; i < config.nfft; ++i) {
|
|
const double w = hann(i, config.nfft);
|
|
output[start + i] += spectrum[i].real() * w;
|
|
normalization[start + i] += w * w;
|
|
}
|
|
}
|
|
for (int i = 0; i < stft.sample_count; ++i)
|
|
output[i] = normalization[i] > 1e-12 ? output[i] / normalization[i] : 0.0;
|
|
return {std::move(output)};
|
|
}
|
|
|
|
DoaResult estimate(const std::vector<double> &audio, int sample_count,
|
|
const Config &config, const std::string &algorithm) {
|
|
const Stft stft = analyze_stft(audio, sample_count, config);
|
|
if (algorithm == "SRP-PHAT" || algorithm == "SRP" || algorithm == "SRP-PATH")
|
|
return srp_phat(stft, config);
|
|
if (algorithm == "NormMUSIC" || algorithm == "MUSIC" || algorithm == "NORMMUSIC")
|
|
return norm_music(stft, config);
|
|
throw std::invalid_argument("unknown algorithm: " + algorithm);
|
|
}
|
|
|
|
MvdrResult beamform(const std::vector<double> &audio, int sample_count,
|
|
const Config &config, double angle_degrees) {
|
|
return mvdr(analyze_stft(audio, sample_count, config), config, angle_degrees);
|
|
}
|
|
|
|
} // namespace esp32_doa
|