Research-Stack/5-Applications/tools-scripts/audio/pipewire_dsp_workloads.py

273 lines
9.5 KiB
Python

#!/usr/bin/env python3
# ==============================================================================
# COPYRIGHT NO ONE EVERYWHERE LLC (WYOMING HOLDING COMPANY)
# PROJECT: SOVEREIGN STACK
# This artifact is entirely proprietary and cryptographically proven.
# Open-Source usage requires explicit permission from Brandon Scott Schneider.
# ==============================================================================
"""
DSP-like front-end workloads for PipeWire waveprobe experiments.
These are bounded host-side transforms intended to approximate the shape of a
front-end DSP lane without claiming to replace a later PipeWire filter node,
custom DSP block, or HDL path.
"""
from __future__ import annotations
from typing import Dict, Tuple
import sys
import os
sys.path.insert(0, os.path.abspath(os.path.join(os.path.dirname(__file__), "..")))
from math_harness_compat import xp, AnyArray
WORKLOAD_RAW = "raw"
WORKLOAD_SPECTRAL_FOCUS = "spectral_focus"
WORKLOAD_TRANSIENT_EDGE = "transient_edge"
WORKLOAD_HYBRID = "hybrid"
AVAILABLE_WORKLOADS = (
WORKLOAD_RAW,
WORKLOAD_SPECTRAL_FOCUS,
WORKLOAD_TRANSIENT_EDGE,
WORKLOAD_HYBRID,
)
def available_workloads() -> Tuple[str, ...]:
return AVAILABLE_WORKLOADS
def _decode_pcm_mono(
chunk_bytes: bytes,
sample_width_bytes: int,
channels: int,
) -> Tuple[AnyArray, int]:
frame_width = max(1, sample_width_bytes * max(1, channels))
usable = len(chunk_bytes) - (len(chunk_bytes) % frame_width)
if usable <= 0:
return xp.zeros(0, dtype=xp.float32), 0
trimmed = chunk_bytes[:usable]
if sample_width_bytes == 1:
arr = xp.frombuffer(trimmed, dtype=xp.uint8).astype(xp.float32)
arr = (arr - 128.0) / 128.0
elif sample_width_bytes == 2:
arr = xp.frombuffer(trimmed, dtype="<i2").astype(xp.float32)
arr = arr / 32768.0
elif sample_width_bytes == 4:
arr = xp.frombuffer(trimmed, dtype="<i4").astype(xp.float32)
arr = arr / 2147483648.0
else:
raise ValueError(f"Unsupported sample width: {sample_width_bytes}")
arr = arr.reshape(-1, max(1, channels))
mono = arr.mean(axis=1).astype(xp.float32, copy=False)
return mono, usable
def _encode_pcm_mono(
samples: AnyArray,
sample_width_bytes: int,
channels: int,
) -> bytes:
if samples.size == 0:
return b""
clipped = xp.clip(samples.astype(xp.float32, copy=False), -1.0, 1.0)
if channels > 1:
clipped = xp.repeat(clipped[:, None], channels, axis=1).reshape(-1)
if sample_width_bytes == 1:
out = xp.clip(xp.round(clipped * 127.0 + 128.0), 0, 255).astype(xp.uint8)
elif sample_width_bytes == 2:
out = xp.clip(xp.round(clipped * 32767.0), -32768, 32767).astype("<i2")
elif sample_width_bytes == 4:
out = xp.clip(
xp.round(clipped * 2147483647.0),
-2147483648,
2147483647,
).astype("<i4")
else:
raise ValueError(f"Unsupported sample width: {sample_width_bytes}")
return out.tobytes()
def _match_rms(processed: AnyArray, reference: AnyArray) -> AnyArray:
if processed.size == 0:
return processed.astype(xp.float32, copy=False)
ref_rms = float(xp.sqrt(xp.mean(reference * reference))) if reference.size else 0.0
proc_rms = float(xp.sqrt(xp.mean(processed * processed)))
out = processed.astype(xp.float32, copy=True)
if ref_rms > 1e-9 and proc_rms > 1e-9:
out *= ref_rms / proc_rms
peak = float(xp.max(xp.abs(out))) if out.size else 0.0
if peak > 0.999:
out *= 0.999 / peak
return out
def _spectral_focus(samples: AnyArray) -> AnyArray:
if samples.size < 8:
return samples.astype(xp.float32, copy=True)
window = xp.hanning(samples.size).astype(xp.float32)
spec = xp.fft.rfft(samples * window)
mag = xp.abs(spec)
if mag.size <= 1:
return samples.astype(xp.float32, copy=True)
max_mag = float(xp.max(mag[1:])) if mag.size > 1 else float(xp.max(mag))
if max_mag <= 1e-12:
return xp.zeros_like(samples, dtype=xp.float32)
weights = 0.15 + 0.85 * xp.sqrt(mag / max_mag)
weights[0] *= 0.35
focused = xp.fft.irfft(spec * weights, n=samples.size).real.astype(xp.float32)
blended = (0.65 * samples + 0.35 * focused).astype(xp.float32, copy=False)
return _match_rms(blended, samples)
def _transient_edge(samples: AnyArray) -> AnyArray:
if samples.size < 4:
return samples.astype(xp.float32, copy=True)
diff = xp.diff(samples, prepend=samples[0]).astype(xp.float32, copy=False)
kernel = xp.array([0.25, 0.5, 0.25], dtype=xp.float32)
smoothed = xp.convolve(diff, kernel, mode="same")
edged = xp.tanh(2.5 * smoothed).astype(xp.float32, copy=False)
blended = (0.55 * samples + 0.45 * edged).astype(xp.float32, copy=False)
return _match_rms(blended, samples)
def _hybrid(samples: AnyArray) -> AnyArray:
focused = _spectral_focus(samples)
edged = _transient_edge(samples)
mixed = (0.5 * samples + 0.3 * focused + 0.2 * edged).astype(xp.float32, copy=False)
return _match_rms(mixed, samples)
def _compute_metrics(samples: AnyArray, sample_rate_hz: int) -> Dict[str, float]:
if samples.size == 0:
return {
"rms": 0.0,
"zero_crossing_rate": 0.0,
"spectral_centroid_hz": 0.0,
"spectral_flatness": 0.0,
"dominant_freq_hz": 0.0,
"transient_ratio": 0.0,
"band_energy_low": 0.0,
"band_energy_mid": 0.0,
"band_energy_high": 0.0,
}
rms = float(xp.sqrt(xp.mean(samples * samples)))
if samples.size >= 2:
zc = float(xp.mean((samples[:-1] * samples[1:]) < 0.0))
diff = xp.diff(samples, prepend=samples[0])
transient_ratio = float(xp.mean(xp.abs(diff)) / max(rms, 1e-9))
else:
zc = 0.0
transient_ratio = 0.0
if samples.size < 8 or sample_rate_hz <= 0:
return {
"rms": rms,
"zero_crossing_rate": zc,
"spectral_centroid_hz": 0.0,
"spectral_flatness": 0.0,
"dominant_freq_hz": 0.0,
"transient_ratio": transient_ratio,
"band_energy_low": 0.0,
"band_energy_mid": 0.0,
"band_energy_high": 0.0,
}
window = xp.hanning(samples.size).astype(xp.float32)
spec = xp.fft.rfft(samples * window)
power = xp.abs(spec) ** 2 + 1e-12
freqs = xp.fft.rfftfreq(samples.size, d=1.0 / sample_rate_hz)
total_power = float(xp.sum(power))
centroid = float(xp.sum(freqs * power) / total_power)
dom_idx = int(xp.argmax(power[1:]) + 1) if power.size > 1 else 0
dominant = float(freqs[dom_idx]) if freqs.size > dom_idx else 0.0
flatness = float(xp.exp(xp.mean(xp.log(power))) / max(xp.mean(power), 1e-12))
low_mask = freqs < 1000.0
mid_mask = (freqs >= 1000.0) & (freqs < 4000.0)
high_mask = freqs >= 4000.0
low = float(xp.sum(power[low_mask]) / total_power)
mid = float(xp.sum(power[mid_mask]) / total_power)
high = float(xp.sum(power[high_mask]) / total_power)
return {
"rms": rms,
"zero_crossing_rate": zc,
"spectral_centroid_hz": centroid,
"spectral_flatness": flatness,
"dominant_freq_hz": dominant,
"transient_ratio": transient_ratio,
"band_energy_low": low,
"band_energy_mid": mid,
"band_energy_high": high,
}
def apply_dsp_workload(
chunk_bytes: bytes,
sample_width_bytes: int,
channels: int,
sample_rate_hz: int,
workload: str = WORKLOAD_RAW,
) -> Tuple[bytes, Dict[str, float]]:
if workload not in AVAILABLE_WORKLOADS:
raise ValueError(
f"Unknown DSP workload {workload!r}; expected one of {AVAILABLE_WORKLOADS}"
)
try:
samples, usable = _decode_pcm_mono(
chunk_bytes,
sample_width_bytes=sample_width_bytes,
channels=channels,
)
except ValueError:
return chunk_bytes, {"workload": workload, "fallback_raw": 1.0}
if usable <= 0:
return b"", {"workload": workload, "fallback_raw": 1.0}
if workload == WORKLOAD_RAW:
processed = samples
elif workload == WORKLOAD_SPECTRAL_FOCUS:
processed = _spectral_focus(samples)
elif workload == WORKLOAD_TRANSIENT_EDGE:
processed = _transient_edge(samples)
else:
processed = _hybrid(samples)
metrics_in = _compute_metrics(samples, sample_rate_hz=sample_rate_hz)
metrics_out = _compute_metrics(processed, sample_rate_hz=sample_rate_hz)
processed_bytes = _encode_pcm_mono(
processed,
sample_width_bytes=sample_width_bytes,
channels=channels,
)
metrics = {
"workload": workload,
"input_rms": metrics_in["rms"],
"output_rms": metrics_out["rms"],
"rms_ratio": metrics_out["rms"] / max(metrics_in["rms"], 1e-9),
"zero_crossing_rate": metrics_out["zero_crossing_rate"],
"spectral_centroid_hz": metrics_out["spectral_centroid_hz"],
"spectral_flatness": metrics_out["spectral_flatness"],
"dominant_freq_hz": metrics_out["dominant_freq_hz"],
"transient_ratio": metrics_out["transient_ratio"],
"band_energy_low": metrics_out["band_energy_low"],
"band_energy_mid": metrics_out["band_energy_mid"],
"band_energy_high": metrics_out["band_energy_high"],
"centroid_shift_hz": metrics_out["spectral_centroid_hz"]
- metrics_in["spectral_centroid_hz"],
"transient_shift": metrics_out["transient_ratio"]
- metrics_in["transient_ratio"],
}
return processed_bytes, metrics