/
/
/
1"""Tests for the pure numpy DBN postprocessor."""
2
3from __future__ import annotations
4
5import numpy as np
6import pytest
7
8from music_assistant.providers.smart_fades.dbn_postprocessor import DBNDownBeatTracker
9
10
11def test_state_space_sizes() -> None:
12 """Test that the bar state space has correct dimensions."""
13 positions, intervals = DBNDownBeatTracker._build_bar_state_space(
14 num_beats=4, min_interval=14, max_interval=55
15 )
16 # Each interval i contributes i states, repeated for each beat
17 expected_states = sum(range(14, 56)) * 4
18 assert len(positions) == expected_states
19 assert len(intervals) == expected_states
20 # Positions for a 4-beat bar range from [0, 4)
21 assert positions.min() >= 0.0
22 assert positions.max() < 4.0
23
24
25def test_state_space_small() -> None:
26 """Test state space with small interval range for easy verification."""
27 positions, _intervals = DBNDownBeatTracker._build_bar_state_space(
28 num_beats=2, min_interval=3, max_interval=4
29 )
30 # interval=3: 3 states, interval=4: 4 states, x 2 beats = 14 states
31 assert len(positions) == 14
32 # Beat 0 positions: [0/3, 1/3, 2/3, 0/4, 1/4, 2/4, 3/4]
33 # Beat 1 positions: [1+0/3, 1+1/3, 1+2/3, 1+0/4, 1+1/4, 1+2/4, 1+3/4]
34 assert positions[0] == pytest.approx(0.0)
35 assert positions[7] == pytest.approx(1.0)
36
37
38def test_transition_model_structure() -> None:
39 """Test that the transition model is a valid sparse matrix."""
40 positions, intervals = DBNDownBeatTracker._build_bar_state_space(
41 num_beats=4, min_interval=14, max_interval=55
42 )
43 tm_states, tm_pointers, _tm_log_probs = DBNDownBeatTracker._build_transition_model(
44 positions, intervals, num_beats=4, transition_lambda=100
45 )
46 num_states = len(positions)
47 # pointers has num_states + 1 entries (CSR format)
48 assert len(tm_pointers) == num_states + 1
49 # Every state must have at least one predecessor
50 for s in range(num_states):
51 assert tm_pointers[s + 1] > tm_pointers[s]
52 # All source state indices must be valid
53 assert tm_states.max() < num_states
54
55
56def test_dbn_tracker_constant_tempo() -> None:
57 """Test that the DBN produces regular beats for a constant-tempo signal."""
58 fps = 50
59 bpm = 120.0
60 duration = 10.0 # seconds
61 num_frames = int(duration * fps)
62
63 # Synthesize activations: strong peaks at expected beat positions
64 interval = 60.0 / bpm * fps # frames per beat
65 beat_act = np.full(num_frames, 0.05)
66 downbeat_act = np.full(num_frames, 0.02)
67 for i in range(int(duration * bpm / 60)):
68 frame = int(i * interval)
69 if frame < num_frames:
70 beat_act[frame] = 0.95
71 if i % 4 == 0:
72 downbeat_act[frame] = 0.95
73
74 combined = np.column_stack(
75 [
76 np.maximum(beat_act - downbeat_act, 1e-5),
77 downbeat_act,
78 ]
79 )
80
81 tracker = DBNDownBeatTracker(beats_per_bar=[4], min_bpm=55, max_bpm=215, fps=fps)
82 result, _num_beats = tracker(combined)
83
84 # Should detect ~20 beats in 10s at 120 BPM
85 beat_times = result[:, 0]
86 assert 18 <= len(beat_times) <= 22
87
88 # Inter-beat intervals should be close to 0.5s
89 ibis = np.diff(beat_times)
90 assert np.all(np.abs(ibis - 0.5) < 0.06)
91
92 # Should have downbeats (beat_position == 1)
93 downbeat_mask = result[:, 1] == 1
94 assert downbeat_mask.sum() >= 2
95
96
97def test_dbn_tracker_fills_intro_gaps() -> None:
98 """Test that the DBN fills in regular beats even when the intro has no peaks."""
99 fps = 50
100 bpm = 120.0
101 duration = 10.0
102 num_frames = int(duration * fps)
103
104 interval = 60.0 / bpm * fps
105 beat_act = np.full(num_frames, 0.05)
106 downbeat_act = np.full(num_frames, 0.02)
107
108 # Only place peaks after 5s (second half)
109 for i in range(int(5.0 * bpm / 60), int(duration * bpm / 60)):
110 frame = int(i * interval)
111 if frame < num_frames:
112 beat_act[frame] = 0.95
113 if i % 4 == 0:
114 downbeat_act[frame] = 0.95
115
116 combined = np.column_stack(
117 [
118 np.maximum(beat_act - downbeat_act, 1e-5),
119 downbeat_act,
120 ]
121 )
122
123 tracker = DBNDownBeatTracker(beats_per_bar=[4], min_bpm=55, max_bpm=215, fps=fps, threshold=0.0)
124 result, _num_beats = tracker(combined)
125
126 beat_times = result[:, 0]
127 # DBN should still find beats in the first 5s via tempo continuity
128 early_beats = beat_times[beat_times < 5.0]
129 assert len(early_beats) >= 5, f"Expected beats in intro region, got {len(early_beats)}"
130
131
132def test_winning_meter_is_returned() -> None:
133 """A 4/4 activation pattern reports beats_per_bar=4 alongside the beats."""
134 fps = 50
135 bpm = 120.0
136 duration = 20.0
137 num_frames = int(duration * fps)
138
139 interval = 60.0 / bpm * fps
140 beat_act = np.full(num_frames, 0.05)
141 downbeat_act = np.full(num_frames, 0.02)
142 for i in range(int(duration * bpm / 60)):
143 frame = int(i * interval)
144 if frame < num_frames:
145 beat_act[frame] = 0.95
146 if i % 4 == 0:
147 downbeat_act[frame] = 0.95
148
149 combined = np.column_stack(
150 [
151 np.maximum(beat_act - downbeat_act, 1e-5),
152 downbeat_act,
153 ]
154 )
155
156 tracker = DBNDownBeatTracker(beats_per_bar=[3, 4], min_bpm=55, max_bpm=215, fps=fps)
157 result, num_beats = tracker(combined)
158
159 assert num_beats == 4
160 assert len(result) > 0
161
162
163def test_dbn_tracker_output_format() -> None:
164 """Test that output matches the madmom interface: (M, 2) with [time, beat_pos]."""
165 fps = 50
166 num_frames = 500 # 10s
167 rng = np.random.default_rng(42)
168 combined = rng.uniform(0.01, 0.1, (num_frames, 2)).astype(np.float64)
169 # Place a few clear peaks
170 for i in range(0, num_frames, 25): # 120 BPM
171 combined[i, 0] = 0.9
172 if (i // 25) % 4 == 0:
173 combined[i, 1] = 0.9
174
175 tracker = DBNDownBeatTracker(beats_per_bar=[4], min_bpm=55, max_bpm=215, fps=fps)
176 result, _num_beats = tracker(combined)
177
178 assert result.ndim == 2
179 assert result.shape[1] == 2
180 # Column 0: times in seconds, should be positive and sorted
181 assert np.all(result[:, 0] >= 0)
182 assert np.all(np.diff(result[:, 0]) > 0)
183 # Column 1: beat positions, integers 1..num_beats
184 assert np.all(result[:, 1] >= 1)
185 assert np.all(result[:, 1] <= 4)
186