| """Load-time grid-rate decimation for the grid-resolution ablation. |
| |
| The cache is rendered once at 250 Hz. Coarser conditioning rates are obtained by |
| decimating that rendering here, at __getitem__ time, so the ablation needs no |
| extra caches and leaves the audio/target side byte-identical across arms. |
| |
| Frames are partitioned into contiguous groups by target frame count, and each |
| channel is reduced with the rule that preserves its meaning: |
| |
| * state velocity -> max (a sustained state stays present) |
| * onset velocity -> max (a hit keeps its velocity, never averaged away) |
| * onset count -> sum (two hits merging into one frame is the effect |
| under study, so counts must accumulate) |
| * onset flags -> any |
| * class ids -> value at the loudest onset in the group, else the state |
| |
| Upsampling is refused: timing precision the 250 Hz raster discarded cannot be |
| recovered by interpolation, so a finer rate must be re-rendered from MIDI. |
| """ |
|
|
| from __future__ import annotations |
|
|
| import numpy as np |
| import torch |
|
|
| |
| |
| ROWS_PER_FAMILY = 3 |
| STATE_VEL_OFFSET = 0 |
| ONSET_VEL_OFFSET = 1 |
| ONSET_COUNT_OFFSET = 2 |
|
|
|
|
| def target_frame_count(*, grid_num_frames: int, duration_sec: float, target_rate_hz: float) -> int: |
| """Frame count for target_rate_hz, never exceeding what the cache holds.""" |
| frames = int(max(0, int(grid_num_frames))) |
| if frames <= 0: |
| return 0 |
| duration = float(max(0.0, float(duration_sec))) |
| rate = float(target_rate_hz) |
| if not rate > 0.0 or not duration > 0.0: |
| return frames |
| return int(min(int(frames), max(1, int(round(duration * rate))))) |
|
|
|
|
| def _group_bounds(num_frames: int, num_groups: int) -> np.ndarray: |
| """Contiguous group edges spreading num_frames over num_groups.""" |
| return np.linspace(0, int(num_frames), num=int(num_groups) + 1).round().astype(np.int64) |
|
|
|
|
| def downsample_grid_payload( |
| *, |
| grid_ft: torch.Tensor, |
| grid_ids_ft: torch.Tensor, |
| family_onsets_ft: torch.Tensor, |
| family_onset_count_ft: torch.Tensor, |
| grid_num_frames: int, |
| duration_sec: float, |
| target_rate_hz: float, |
| ) -> dict[str, object]: |
| """Decimate one cached example to target_rate_hz. |
| |
| Returns the replacement tensors plus the new frame count / effective rate. |
| A target rate matching the cached frame count is a no-op returning the |
| inputs unchanged. A genuinely higher rate raises instead of fabricating |
| timing precision by interpolation. |
| """ |
| frames = int(grid_num_frames) |
| requested_frames = int(max(1, round(float(duration_sec) * float(target_rate_hz)))) |
| if int(frames) > 0 and float(duration_sec) > 0.0 and int(requested_frames) > int(frames): |
| cached_rate = float(frames) / float(duration_sec) |
| raise ValueError( |
| "grid-rate upsampling is invalid: " |
| f"requested {float(target_rate_hz):g} Hz ({requested_frames} frames) from a " |
| f"{cached_rate:.3f} Hz cache ({frames} frames). Re-render the grid from MIDI." |
| ) |
| out_frames = target_frame_count( |
| grid_num_frames=frames, |
| duration_sec=duration_sec, |
| target_rate_hz=target_rate_hz, |
| ) |
| duration = float(max(1.0e-6, float(duration_sec))) |
|
|
| if int(out_frames) >= int(frames) or int(frames) <= 0: |
| |
| |
| return { |
| "grid_ft": grid_ft, |
| "grid_ids_ft": grid_ids_ft, |
| "family_onsets_ft": family_onsets_ft, |
| "family_onset_count_ft": family_onset_count_ft, |
| "grid_num_frames": int(frames), |
| "grid_frame_rate": float(frames) / duration, |
| } |
|
|
| bounds = _group_bounds(int(frames), int(out_frames)) |
| num_families = int(family_onsets_ft.shape[0]) |
|
|
| grid_out = torch.zeros((int(grid_ft.shape[0]), int(out_frames)), dtype=grid_ft.dtype) |
| ids_out = torch.zeros((int(grid_ids_ft.shape[0]), int(out_frames)), dtype=grid_ids_ft.dtype) |
| onsets_out = torch.zeros((num_families, int(out_frames)), dtype=family_onsets_ft.dtype) |
| counts_out = torch.zeros((num_families, int(out_frames)), dtype=torch.int64) |
|
|
| onset_count_i64 = family_onset_count_ft.to(dtype=torch.int64) |
|
|
| for out_idx in range(int(out_frames)): |
| start = int(bounds[int(out_idx)]) |
| stop = int(max(int(start) + 1, int(bounds[int(out_idx) + 1]))) |
|
|
| window_grid = grid_ft[:, int(start) : int(stop)] |
| |
| |
| grid_out[:, int(out_idx)] = window_grid.max(dim=1).values |
| for family_idx in range(int(num_families)): |
| count_row = int(ONSET_COUNT_OFFSET + ROWS_PER_FAMILY * int(family_idx)) |
| if int(count_row) < int(grid_ft.shape[0]): |
| grid_out[int(count_row), int(out_idx)] = window_grid[int(count_row)].sum() |
|
|
| onsets_out[:, int(out_idx)] = family_onsets_ft[:, int(start) : int(stop)].any(dim=1) |
| counts_out[:, int(out_idx)] = onset_count_i64[:, int(start) : int(stop)].sum(dim=1) |
|
|
| |
| |
| window_onset_vel = grid_ft[ONSET_VEL_OFFSET::ROWS_PER_FAMILY, int(start) : int(stop)] |
| for family_idx in range(int(ids_out.shape[0])): |
| if int(family_idx) < int(window_onset_vel.shape[0]): |
| vel_row = window_onset_vel[int(family_idx)] |
| if float(vel_row.max()) > 0.0: |
| pick = int(start) + int(torch.argmax(vel_row).item()) |
| else: |
| pick = int(start) |
| else: |
| pick = int(start) |
| ids_out[int(family_idx), int(out_idx)] = grid_ids_ft[int(family_idx), int(pick)] |
|
|
| count_dtype = family_onset_count_ft.dtype |
| if count_dtype == torch.uint8: |
| counts_out = counts_out.clamp_(0, 255) |
|
|
| return { |
| "grid_ft": grid_out.contiguous(), |
| "grid_ids_ft": ids_out.contiguous(), |
| "family_onsets_ft": onsets_out.contiguous(), |
| "family_onset_count_ft": counts_out.to(dtype=count_dtype).contiguous(), |
| "grid_num_frames": int(out_frames), |
| "grid_frame_rate": float(out_frames) / duration, |
| } |
|
|
|
|
| def uniform_grid_times(num_frames: int, duration_sec: float) -> torch.Tensor: |
| """Frame-centre times, matching _uniform_frame_times_from_duration (bpm geometry).""" |
| frames = int(max(0, int(num_frames))) |
| duration = float(max(0.0, float(duration_sec))) |
| if frames <= 0 or not duration > 0.0: |
| return torch.zeros((0,), dtype=torch.float32) |
| idx = torch.arange(int(frames), dtype=torch.float32) |
| return ((idx + 0.5) * (float(duration) / float(frames))).contiguous() |
|
|
