use std::time::Duration;
use ndarray::s;
use rubato::audioadapter_buffers::direct::SequentialSlice;
use rubato::{
Fft, FixedSync, ResampleError, Resampler, ResamplerConstructionError, WindowFunction,
};
use crate::signal::math::db_to_gain;
use crate::signal::{SignalError, TimeSignal};
use crate::dsp::apply_hann_right;
#[derive(Debug, thiserror::Error)]
pub enum TimeError {
#[error("invalid trim range: start={start}, end={end}, signal_len={signal_len}")]
InvalidTrimRange {
start: usize,
end: usize,
signal_len: usize,
},
#[error("invalid analysis threshold: {threshold_db} dB")]
InvalidThreshold { threshold_db: f64 },
#[error(
"impulse response start could not be detected in channel {channel}: peak power {peak_power} below threshold {required_peak_power} from noise floor {noise_power}"
)]
LowSnr {
channel: usize,
peak_power: f64,
noise_power: f64,
required_peak_power: f64,
},
#[error("trimmed window duration must map to at least one sample, got {duration_seconds} s")]
InvalidWindowDuration { duration_seconds: f64 },
#[error("invalid sample rate for resampling: {sample_rate} Hz")]
InvalidResampleRate { sample_rate: f64 },
#[error("shift vector length {got} does not match channel count {expected}")]
InvalidShiftCount { expected: usize, got: usize },
#[error("linear shift magnitude {shift} exceeds signal length {signal_len}")]
ShiftExceedsSignal { shift: isize, signal_len: usize },
#[error(transparent)]
ResamplerConstruction(#[from] ResamplerConstructionError),
#[error(transparent)]
Resampler(#[from] ResampleError),
#[error(transparent)]
Signal(#[from] SignalError),
}
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub enum TimeShiftMode {
Linear,
Cyclic,
}
pub fn trim_samples(
signal: &TimeSignal,
start_sample: usize,
end_sample: i64,
) -> Result<TimeSignal, TimeError> {
let signal_len = signal.num_time_steps();
let end_sample = if end_sample < 0 {
signal_len.saturating_sub((-end_sample) as usize)
} else {
end_sample as usize
};
if start_sample >= end_sample || end_sample > signal_len {
return Err(TimeError::InvalidTrimRange {
start: start_sample,
end: end_sample,
signal_len,
});
}
Ok(TimeSignal::new(
signal
.time_data()
.slice(s![.., start_sample..end_sample])
.to_owned(),
signal.sample_rate(),
)?)
}
pub fn trim_duration(
signal: &TimeSignal,
start_duration: Duration,
end_duration: Duration,
) -> Result<TimeSignal, TimeError> {
let start_sample = (signal.sample_rate() * start_duration.as_secs_f64()).round() as usize;
let end_sample = (signal.sample_rate() * end_duration.as_secs_f64()).round() as i64;
trim_samples(signal, start_sample, end_sample)
}
pub fn apply_gain(signal: &TimeSignal, gain: f64) -> Result<TimeSignal, TimeError> {
Ok(TimeSignal::new(
signal.time_data().mapv(|sample| sample * gain),
signal.sample_rate(),
)?)
}
pub fn apply_gain_db(signal: &TimeSignal, gain_db: f64) -> Result<TimeSignal, TimeError> {
apply_gain(signal, db_to_gain(gain_db))
}
pub fn time_shift(
signal: &TimeSignal,
shift_samples: isize,
mode: TimeShiftMode,
pad_value: f64,
) -> Result<TimeSignal, TimeError> {
let shifts = vec![shift_samples; signal.num_channels()];
time_shift_per_channel(signal, &shifts, mode, pad_value)
}
pub fn time_shift_per_channel(
signal: &TimeSignal,
shift_samples: &[isize],
mode: TimeShiftMode,
pad_value: f64,
) -> Result<TimeSignal, TimeError> {
if shift_samples.len() != signal.num_channels() {
return Err(TimeError::InvalidShiftCount {
expected: signal.num_channels(),
got: shift_samples.len(),
});
}
let signal_len = signal.num_time_steps();
if let Some(shift) = shift_samples
.iter()
.copied()
.find(|&shift| matches!(mode, TimeShiftMode::Linear) && shift.unsigned_abs() > signal_len)
{
return Err(TimeError::ShiftExceedsSignal { shift, signal_len });
}
let mut shifted = signal.clone();
for (channel_index, mut channel) in shifted.channel_iter_mut().enumerate() {
let shift = shift_samples[channel_index];
rotate_channel(
channel.as_slice_mut().expect("channel view is contiguous"),
shift,
);
if matches!(mode, TimeShiftMode::Linear) {
if shift > 0 {
channel.slice_mut(s![..shift as usize]).fill(pad_value);
} else if shift < 0 {
let start = signal_len - shift.unsigned_abs();
channel.slice_mut(s![start..]).fill(pad_value);
}
}
}
Ok(shifted)
}
pub fn find_impulse_response_start(
impulse_response: &TimeSignal,
threshold_db: Option<f64>,
) -> Result<Vec<usize>, TimeError> {
let threshold_db = threshold_db.unwrap_or(20.0);
if !threshold_db.is_finite() || threshold_db < 0.0 {
return Err(TimeError::InvalidThreshold { threshold_db });
}
let num_time_steps = impulse_response.num_time_steps();
let mask_start = ((0.9 * num_time_steps as f64) as usize).min(num_time_steps);
let analysis_len = num_time_steps.saturating_sub(mask_start);
let snr_ratio = 10.0_f64.powf(threshold_db / 10.0);
let relative_threshold = 10.0_f64.powf(-threshold_db / 10.0);
let mut start_samples = Vec::with_capacity(impulse_response.num_channels());
for (channel_index, channel) in impulse_response.channel_iter().enumerate() {
let power: Vec<f64> = channel.iter().map(|sample| sample.abs().powi(2)).collect();
let noise_power = if analysis_len == 0 {
0.0
} else {
power.iter().skip(mask_start).sum::<f64>() / analysis_len as f64
};
let (peak_index, peak_power) = power
.iter()
.copied()
.enumerate()
.max_by(|(_, a), (_, b)| a.partial_cmp(b).unwrap_or(std::cmp::Ordering::Equal))
.unwrap_or((0, 0.0));
let required_peak_power = noise_power * snr_ratio;
if peak_power < required_peak_power || peak_index > mask_start {
return Err(TimeError::LowSnr {
channel: channel_index,
peak_power,
noise_power,
required_peak_power,
});
}
let first_above_threshold =
power
.iter()
.take(peak_index + 1)
.enumerate()
.find_map(|(index, value)| {
if peak_power > 0.0 && (*value / peak_power) >= relative_threshold {
Some(index)
} else {
None
}
});
let start_sample = first_above_threshold
.map(|index| index.saturating_sub(1))
.unwrap_or(0);
start_samples.push(start_sample);
}
Ok(start_samples)
}
pub fn window_and_trim(
impulse_response: &TimeSignal,
duration_seconds: f64,
) -> Result<(TimeSignal, Vec<usize>), TimeError> {
let len = (duration_seconds * impulse_response.sample_rate()).ceil() as usize;
if len == 0 {
return Err(TimeError::InvalidWindowDuration { duration_seconds });
}
let starts = find_impulse_response_start(impulse_response, None)?;
let mut windowed = TimeSignal::zeros(
impulse_response.num_channels(),
len,
impulse_response.sample_rate(),
)?;
for ((mut output_channel, input_channel), &start) in windowed
.channel_iter_mut()
.zip(impulse_response.channel_iter())
.zip(starts.iter())
{
output_channel
.iter_mut()
.zip(input_channel.iter().skip(start).take(len))
.for_each(|(dst, src)| *dst = *src);
if let Some(slice) = output_channel.as_slice_mut() {
apply_hann_right(slice, 0, len);
}
}
Ok((windowed, starts))
}
pub fn resample(signal: &TimeSignal, new_sample_rate: f64) -> Result<TimeSignal, TimeError> {
if (new_sample_rate - signal.sample_rate()).abs() < 1e-9 {
return Ok(signal.clone());
}
let sample_rate_in = validated_resample_rate(signal.sample_rate())?;
let sample_rate_out = validated_resample_rate(new_sample_rate)?;
let num_channels = signal.num_channels();
let num_input_frames = signal.num_time_steps();
let mut resampler = Fft::<f64>::new_custom(
sample_rate_in,
sample_rate_out,
1024,
2,
num_channels,
WindowFunction::BlackmanHarris2,
FixedSync::Both,
)?;
let input_sequential: Vec<f64> = signal.time_data().iter().copied().collect();
let buffer_in = SequentialSlice::new(&input_sequential, num_channels, num_input_frames)
.expect("sequential input length matches channel/frame dimensions");
let num_output_frames = resampler.process_all_needed_output_len(num_input_frames);
let mut output_sequential = vec![0.0; num_output_frames * num_channels];
let mut buffer_out =
SequentialSlice::new_mut(&mut output_sequential, num_channels, num_output_frames)
.expect("allocated output length matches channel/frame dimensions");
let (_, actual_output_frames) =
resampler.process_all_into_buffer(&buffer_in, &mut buffer_out, num_input_frames, None)?;
output_sequential.truncate(actual_output_frames * num_channels);
let data =
ndarray::Array2::from_shape_vec((num_channels, actual_output_frames), output_sequential)
.expect("resampler output length matches channel/frame dimensions");
TimeSignal::new(data, new_sample_rate).map_err(TimeError::from)
}
fn validated_resample_rate(sample_rate: f64) -> Result<usize, TimeError> {
if !sample_rate.is_finite() || sample_rate <= 0.0 {
return Err(TimeError::InvalidResampleRate { sample_rate });
}
let rounded = sample_rate.round();
if (sample_rate - rounded).abs() > 1e-9 {
return Err(TimeError::InvalidResampleRate { sample_rate });
}
Ok(rounded as usize)
}
fn rotate_channel(channel: &mut [f64], shift: isize) {
if channel.is_empty() {
return;
}
let normalized_shift = shift.rem_euclid(channel.len() as isize) as usize;
channel.rotate_right(normalized_shift);
}
#[cfg(test)]
mod tests {
use std::time::Duration;
use crate::signal::{SineConfig, generate_sine};
use approx::assert_abs_diff_eq;
use ndarray::arr2;
use super::*;
#[test]
fn trim_samples_accepts_positive_and_negative_end() {
let signal = TimeSignal::new(arr2(&[[0.0, 1.0, 2.0, 3.0, 4.0]]), 10.0).unwrap();
let trimmed = trim_samples(&signal, 1, 4).unwrap();
assert_eq!(trimmed.num_time_steps(), 3);
assert_abs_diff_eq!(trimmed.channel(0)[0], 1.0, epsilon = 1e-12);
let trimmed = trim_samples(&signal, 1, -1).unwrap();
assert_eq!(trimmed.num_time_steps(), 3);
assert_abs_diff_eq!(trimmed.channel(0)[2], 3.0, epsilon = 1e-12);
}
#[test]
fn trim_duration_maps_to_sample_positions() {
let signal = TimeSignal::new(arr2(&[[0.0; 20]]), 10.0).unwrap();
let trimmed = trim_duration(
&signal,
Duration::from_secs_f64(0.2),
Duration::from_secs_f64(1.1),
)
.unwrap();
assert_eq!(trimmed.num_time_steps(), 9);
}
#[test]
fn trim_rejects_invalid_ranges() {
let signal = TimeSignal::new(arr2(&[[0.0, 1.0, 2.0]]), 1.0).unwrap();
let result = trim_samples(&signal, 2, 2);
assert!(matches!(result, Err(TimeError::InvalidTrimRange { .. })));
}
#[test]
fn apply_gain_scales_samples_linearly() {
let signal = TimeSignal::new(arr2(&[[0.25, -0.5, 1.0]]), 48_000.0).unwrap();
let gained = apply_gain(&signal, 2.0).unwrap();
assert_abs_diff_eq!(gained.channel(0)[0], 0.5, epsilon = 1e-12);
assert_abs_diff_eq!(gained.channel(0)[1], -1.0, epsilon = 1e-12);
assert_abs_diff_eq!(gained.channel(0)[2], 2.0, epsilon = 1e-12);
}
#[test]
fn apply_gain_db_converts_db_to_linear_gain() {
let signal = TimeSignal::new(arr2(&[[0.25, -0.5, 1.0]]), 48_000.0).unwrap();
let gained = apply_gain_db(&signal, -6.0).unwrap();
let expected_gain = db_to_gain(-6.0);
assert_abs_diff_eq!(gained.channel(0)[0], 0.25 * expected_gain, epsilon = 1e-12);
assert_abs_diff_eq!(gained.channel(0)[1], -0.5 * expected_gain, epsilon = 1e-12);
assert_abs_diff_eq!(gained.channel(0)[2], expected_gain, epsilon = 1e-12);
}
#[test]
fn time_shift_cyclic_wraps_signal() {
let signal = TimeSignal::new(arr2(&[[1.0, 2.0, 3.0, 4.0]]), 48_000.0).unwrap();
let shifted = time_shift(&signal, 1, TimeShiftMode::Cyclic, 0.0).unwrap();
assert_abs_diff_eq!(
shifted.channel(0),
ndarray::arr1(&[4.0, 1.0, 2.0, 3.0]),
epsilon = 1e-12
);
}
#[test]
fn time_shift_linear_pads_exposed_samples() {
let signal = TimeSignal::new(arr2(&[[1.0, 2.0, 3.0, 4.0]]), 48_000.0).unwrap();
let shifted = time_shift(&signal, -2, TimeShiftMode::Linear, 0.0).unwrap();
assert_abs_diff_eq!(
shifted.channel(0),
ndarray::arr1(&[3.0, 4.0, 0.0, 0.0]),
epsilon = 1e-12
);
}
#[test]
fn time_shift_accepts_per_channel_shifts() {
let signal =
TimeSignal::new(arr2(&[[1.0, 2.0, 3.0], [10.0, 20.0, 30.0]]), 48_000.0).unwrap();
let shifted =
time_shift_per_channel(&signal, &[1, -1], TimeShiftMode::Linear, -1.0).unwrap();
assert_abs_diff_eq!(
shifted.channel(0),
ndarray::arr1(&[-1.0, 1.0, 2.0]),
epsilon = 1e-12
);
assert_abs_diff_eq!(
shifted.channel(1),
ndarray::arr1(&[20.0, 30.0, -1.0]),
epsilon = 1e-12
);
}
#[test]
fn find_impulse_response_start_detects_onset_per_channel() {
let signal = TimeSignal::new(
arr2(&[
[0.0, 0.0, 0.1, 0.4, 1.0, 0.5, 0.0, 0.0, 0.0, 0.0],
[0.0, 0.0, 0.0, 0.2, 0.6, 1.0, 0.2, 0.0, 0.0, 0.0],
]),
10.0,
)
.unwrap();
let starts = find_impulse_response_start(&signal, Some(20.0)).unwrap();
assert_eq!(starts, vec![1, 2]);
}
#[test]
fn find_impulse_response_start_rejects_low_snr() {
let signal = TimeSignal::new(
arr2(&[[0.0, 0.0, 0.01, 0.0, 0.0, 0.01, 0.01, 0.01, 0.01, 0.01]]),
10.0,
)
.unwrap();
let result = find_impulse_response_start(&signal, Some(20.0));
assert!(matches!(result, Err(TimeError::LowSnr { channel: 0, .. })));
}
#[test]
fn window_and_trim_uses_detected_onset_and_applies_fade_out() {
let signal = TimeSignal::new(
arr2(&[[0.0, 0.0, 0.1, 0.4, 1.0, 0.5, 0.25, 0.0, 0.0, 0.0]]),
10.0,
)
.unwrap();
let (trimmed, starts) = window_and_trim(&signal, 0.5).unwrap();
assert_eq!(starts, vec![1]);
assert_eq!(trimmed.num_time_steps(), 5);
assert_abs_diff_eq!(trimmed.channel(0)[0], 0.0, epsilon = 1e-12);
assert!(trimmed.channel(0)[1] > 0.0);
assert!(trimmed.channel(0)[4].abs() < trimmed.channel(0)[3].abs());
}
#[test]
fn window_and_trim_rejects_zero_duration() {
let signal = TimeSignal::new(arr2(&[[1.0, 0.0, 0.0]]), 10.0).unwrap();
let result = window_and_trim(&signal, 0.0);
assert!(matches!(
result,
Err(TimeError::InvalidWindowDuration { .. })
));
}
#[test]
fn resample_preserves_duration_and_frequency() {
let signal = generate_sine(
4_410,
440.0,
SineConfig {
amplitude: 0.5,
sample_rate: 44_100.0,
num_channels: 1,
},
)
.unwrap();
let resampled = resample(&signal, 48_000.0).unwrap();
assert_eq!(resampled.sample_rate(), 48_000.0);
assert_abs_diff_eq!(
resampled.length_in_seconds(),
signal.length_in_seconds(),
epsilon = 1e-3
);
let (peak_freq, _) = resampled.into_freq().to_magnitude().max_per_channel()[0];
assert_abs_diff_eq!(peak_freq, 440.0, epsilon = 1.0);
}
#[test]
fn resample_rejects_non_integer_target_rate() {
let signal = TimeSignal::new(arr2(&[[0.0, 1.0, 0.0]]), 48_000.0).unwrap();
let result = resample(&signal, 48_000.5);
assert!(matches!(
result,
Err(TimeError::InvalidResampleRate {
sample_rate: 48_000.5
})
));
}
#[test]
fn resample_short_circuits_when_sample_rate_is_unchanged() {
let signal = TimeSignal::new(arr2(&[[0.0, 1.0, 0.0, -1.0]]), 48_000.0).unwrap();
let resampled = resample(&signal, 48_000.0).unwrap();
assert_eq!(resampled, signal);
}
}