#!/usr/bin/env python3 """ unmix: music source separation as a constrained linear inverse problem, solved with PDHG. Unknowns (time domain): v the lead voice, mono, constrained to the stereo centre h_L,h_R harmonic accompaniment (chords, bass), stereo p_L,p_R percussion, stereo Measurement: g_c = v + h_c + p_c (c = L, R) i.e. g = A f The mixing matrix is perfectly conditioned on its row space but has a (K-2)T-dimensional null space, so everything comes from the priors: min a_H Σ_c || D_t E F h_c ||_1 harmonic: phase-corrected spectra barely change over time + a_P Σ_c Σ_frames || (F p_c)_frame ||_2 percussion: few active frames, broadband + a_V || z ||_1 voice: v = F^T (M ⊙ z), sparse coefficients on the comb M (harmonics of the tracked pitch, human range only) || (F_c (g_c - v - h_c - p_c))_b ||_2 <= eps_b for every frequency band b (split data fidelity) F is a tight (Parseval) STFT frame, E removes each bin's phase advance using the mixture's instantaneous frequency (as in Masuyama et al. 2019), D_t is a forward difference over frames. Solved with Chambolle–Pock, one dual variable per term, per-block dual step sizes. """ from __future__ import annotations import numpy as np SR = 16000 # ───────────────────────────────────────────── tight STFT frame class Frame: """Circular sqrt-Hann STFT with ||F x|| = ||x|| exactly (Parseval). Full complex spectrum.""" def __init__(self, T, N=1024, hop=256): assert T % hop == 0 self.T, self.N, self.hop = T, N, hop self.nF = T // hop n = np.arange(N) self.w = np.sqrt(0.5 - 0.5 * np.cos(2 * np.pi * n / N)).astype(np.float32) self.c = float((self.w ** 2).sum() / hop) self.idx = (np.arange(self.nF)[:, None] * hop + n[None, :]) % T self.scale = 1.0 / np.sqrt(N * self.c) def fwd(self, x): return (np.fft.fft(x[self.idx] * self.w, axis=1) * self.scale).astype(np.complex64) def bin_hz(self): k = np.arange(self.N) return np.where(k <= self.N // 2, k, k - self.N) * SR / self.N # signed frequency per bin def frame_times(self): return (np.arange(self.nF) * self.hop + self.N / 2) / SR def fast_adj(F: Frame, Y): """Overlap-add adjoint without np.add.at (much faster): sum the hop-shifted frames.""" fr = (np.real(np.fft.ifft(Y, axis=1)) * F.N * F.w * F.scale).astype(np.float32) out = np.zeros(F.T + F.N, dtype=np.float32) R = F.N // F.hop for r in range(R): # frames r, r+R, r+2R ... don't overlap blk = fr[r::R].reshape(-1) s = r * F.hop out[s:s + blk.size] += blk out[:F.N] += out[F.T:F.T + F.N] # circular wrap return out[:F.T] Frame.adj = fast_adj # overlap-add adjoint # ───────────────────────────────────────────── operators def phase_correction(F: Frame, Gmix): """E = exp(-i Θ): removes each bin's accumulated phase advance (from the mixture's instantaneous frequency).""" N, hop = F.N, F.hop k = np.arange(N) omega = 2 * np.pi * np.where(k <= N // 2, k, k - N) / N # rad / sample, signed ph = np.angle(Gmix) dphi = np.diff(ph, axis=0, prepend=ph[:1]) - omega * hop dphi = (dphi + np.pi) % (2 * np.pi) - np.pi inst = omega + dphi / hop theta = np.cumsum(inst * hop, axis=0) return np.exp(-1j * theta).astype(np.complex64) def Dt(Z): # forward difference along frames (last row 0) out = np.zeros_like(Z); out[:-1] = Z[1:] - Z[:-1]; return out def DtT(Y): # adjoint of Dt out = np.zeros_like(Y); out[1:] += Y[:-1]; out[:-1] -= Y[:-1]; return out def comb_support(F: Frame, f0_frames, fmax=5000.0, tol_bins=1.5, tol_rel=0.03, dilate=1): """Boolean mask (frames × bins): True near h·f0 for voiced frames whose f0 is a plausible human pitch.""" hz = np.abs(F.bin_hz()) mask = np.zeros((F.nF, F.N), dtype=bool) for n, f0 in enumerate(f0_frames): if not (75.0 <= f0 <= 1100.0): # range check: not a human singing pitch → no voice here continue H = int(fmax // f0) for h in range(1, H + 1): fc = h * f0 mask[n] |= np.abs(hz - fc) <= max(tol_bins * SR / F.N, tol_rel * fc) if dilate: m2 = mask.copy() for d in range(1, dilate + 1): m2[d:] |= mask[:-d]; m2[:-d] |= mask[d:] mask = m2 return mask def band_index(F: Frame, edges_hz): """Band label per bin (by |frequency|).""" hz = np.abs(F.bin_hz()) return np.digitize(hz, edges_hz) - 1 # ───────────────────────────────────────────── solver def separate(g, f0_frames, a_H=1.0, a_P=1.0, a_V=1.0, eps_rel=0.0, bands=None, iters=300, N=1024, hop=256, stereo=True, gamma=None, rho=1.0, verbose=True, callback=None): """ g: (2, T) stereo mixture (T multiple of hop). f0_frames: voice pitch per STFT frame (Hz, 0 = none). eps_rel: per-band data tolerance as a fraction of the mixture's band energy (0 = exact equality). gamma: relative dual step weights per block {'H','P','D'}. The voice is in synthesis form, v = F^T (M ⊙ z): z lives on the comb M, so v is always feasible. (An analysis constraint "F v = 0 off the comb" is almost never satisfiable in a redundant frame: onsets and window side-lobes leak off the comb, and the feasible set collapses towards v = 0.) Returns dict(v, h, p, z, support, frame). """ C = 2 if stereo else 1 if not stereo: g = g.mean(axis=0, keepdims=True) T = g.shape[1] F = Frame(T, N, hop) Gmix = [F.fwd(g[c]) for c in range(C)] E = phase_correction(F, sum(Gmix) / C) M = comb_support(F, f0_frames) edges = bands if bands is not None else [0, 300, 1000, 3000, 1e9] bidx = band_index(F, edges) nb = len(edges) - 1 sel = [bidx == b for b in range(nb)] eps = [np.array([np.linalg.norm(Gmix[c][:, sel[b]]) for b in range(nb)]) * eps_rel for c in range(C)] # operator norms: ||D_t E F|| <= 2; ||F|| = 1; data block: stereo [[A,I,0,I,0],[A,0,I,0,I]] with ||A|| <= 1 → 2 L = {'H': 2.0, 'P': 1.0, 'D': 2.0 if C == 2 else np.sqrt(3.0)} gamma = gamma or {'H': 1.0, 'P': 1.0, 'D': 1.0} sig = {k: gamma[k] / L[k] ** 2 for k in L} tau = 0.99 / sum(gamma.values()) z = np.zeros((F.nF, N), np.complex64) h = np.zeros((C, T), np.float32) p = np.zeros((C, T), np.float32) zb, hb, pb = z.copy(), h.copy(), p.copy() yH = [np.zeros((F.nF, N), np.complex64) for _ in range(C)] yP = [np.zeros((F.nF, N), np.complex64) for _ in range(C)] yD = [np.zeros((F.nF, N), np.complex64) for _ in range(C)] voice = lambda zz: F.adj(M * zz) for it in range(iters): vb = voice(zb) # ── dual updates (prox of conjugates) ── for c in range(C): w = yH[c] + sig['H'] * Dt(E * F.fwd(hb[c])) yH[c] = w / np.maximum(1.0, np.abs(w) / a_H) # ℓ1 → clip magnitudes w = yP[c] + sig['P'] * F.fwd(pb[c]) yP[c] = w / np.maximum(1.0, np.linalg.norm(w, axis=1, keepdims=True) / a_P) # ℓ2,1 over frames w = yD[c] + sig['D'] * F.fwd(vb + hb[c] + pb[c]) u = w / sig['D'] # prox of σ ι_C^* = w - σ P_C(w/σ) Pu = u.copy() for b in range(nb): r = u[:, sel[b]] - Gmix[c][:, sel[b]] nr = np.linalg.norm(r) Pu[:, sel[b]] = Gmix[c][:, sel[b]] + (r * (eps[c][b] / nr) if nr > eps[c][b] and nr > 0 else r) yD[c] = w - sig['D'] * Pu # ── primal updates ── z_old, h_old, p_old = z, h, p aD = [F.adj(yD[c]) for c in range(C)] gz = M * F.fwd(sum(aD)) # adjoint of z ↦ F^T(M z) w = z - tau * gz mag = np.abs(w) z = (M * w * np.maximum(0.0, 1.0 - tau * a_V / np.maximum(mag, 1e-12))).astype(np.complex64) # soft-threshold on the comb h = np.stack([h[c] - tau * (F.adj(np.conj(E) * DtT(yH[c])) + aD[c]) for c in range(C)]) p = np.stack([p[c] - tau * (F.adj(yP[c]) + aD[c]) for c in range(C)]) if rho != 1.0: z = z_old + rho * (z - z_old); h = h_old + rho * (h - h_old); p = p_old + rho * (p - p_old) zb, hb, pb = 2 * z - z_old, 2 * h - h_old, 2 * p - p_old if callback and (it % 25 == 0 or it == iters - 1): callback(it, voice(z), h, p) if verbose and (it % 50 == 0 or it == iters - 1): v = voice(z) res = np.sqrt(sum(np.sum((g[c] - v - h[c] - p[c]) ** 2) for c in range(C)) / np.sum(g ** 2)) print(f" it {it:4d} relative residual {res:.4f}", flush=True) v = voice(z) if not stereo: h = np.repeat(h, 2, axis=0); p = np.repeat(p, 2, axis=0) return {"v": v, "h": h, "p": p, "z": z, "support": M, "frame": F} def refit(g, out, power=2.0): """Undo ℓ1 shrinkage: use the recovered stems as soft masks and re-split the measured mixture. m_k = |S_k|^p / Σ_j |S_j|^p per channel, and each stem becomes F^T(m_k ⊙ G_c): a Wiener-style re-split that keeps the separation's shape but restores the level ℓ1 shrank. The voice stays centred (average of its two channel estimates). Chosen on the tuning song (p = 2). """ F = out["frame"] V = F.fwd(out["v"]) v_new, h_new, p_new = np.zeros_like(out["v"]), np.zeros_like(out["h"]), np.zeros_like(out["p"]) for c in range(2): G = F.fwd(g[c]) Hc, Pc = F.fwd(out["h"][c]), F.fwd(out["p"][c]) a, b, d = np.abs(V) ** power, np.abs(Hc) ** power, np.abs(Pc) ** power tot = a + b + d + 1e-12 v_new += 0.5 * F.adj(a / tot * G) h_new[c] = F.adj(b / tot * G) p_new[c] = F.adj(d / tot * G) return {**out, "v": v_new, "h": h_new, "p": p_new} # ═════════════════════════════════════════════ v2: MP3 cells + pitch-following directional TV class MP3Data: """Data term from the MP3 bitstream for an excerpt of granules [g0, g1). Psi = Gram^{-1/2} S^T is the decoder's synthesis transposed and normalised (Psi Psi^T = I), so the constraint "the encoder's analysis of the mix lies in its quantiser cell" reads Gram^{1/2} lo <= Psi y_c <= Gram^{1/2} hi on reliable cells, jointly over the two coded channels (mid/side or left/right, per granule). unreliable="free" leaves the other coefficients unconstrained; "decoded" pins them to the decoded value (the hybrid: quantiser cells where the file is informative, the decoded value elsewhere).""" def __init__(self, bs_full, g0, g1, slack=0.5, edge=2, unreliable="free"): import mp3.mp3model as MM bs = {k: (v[g0:g1] if isinstance(v, np.ndarray) and v.ndim >= 1 and v.shape[0] == bs_full["G"] else v) for k, v in bs_full.items() if k != "pcm"} bs["G"] = g1 - g0 self.S = MM.Synthesis(bs) lo, hi = MM.cells(bs, slack=slack) ok = MM.reliable(bs) ok[:edge] = False; ok[-edge:] = False # excerpt edges: analysis is truncated there self.r = np.sqrt(self.S.gram_diag()) if unreliable == "decoded": # hybrid: cells where informative, decoded value elsewhere dec = bs["x"] * self.r self.lo = np.where(ok, lo * self.r, dec) self.hi = np.where(ok, hi * self.r, dec) self.lo[:edge] = -np.inf; self.hi[:edge] = np.inf; self.lo[-edge:] = -np.inf; self.hi[-edge:] = np.inf else: self.lo = np.where(ok, lo * self.r, -np.inf) self.hi = np.where(ok, hi * self.r, np.inf) self.ok = ok self.T = 576 * (g1 - g0) def fwd(self, y2): # stereo (2, T) -> normalised coded coefficients (G, 2, 576) return self.S.T(y2.astype(np.float64)) / self.r def adj(self, Y): # adjoint return self.S(Y / self.r) def separate2(g, f0_frames, a_H=1.0, a_P=1.0, a_V=1.0, a_T=0.0, voice_tv=None, f0_hz_cqt=None, mp3=None, iters=300, N=1024, hop=256, gamma=None, verbose=False, cqt_kw=None, phase_from=None): """ Like separate() (stereo), plus: voice_tv: None | 'fixed' | 'follow' directional TV on the voice in a constant-Q frame, weight a_T. 'follow' steps along the pitch slope with harmonic phase prediction; 'fixed' uses dk = 0 with the mixture's per-channel phase advance (Masuyama-style). mp3: None -> exact equality with the decoded mixture g; MP3Data -> quantiser-cell box. phase_from: signal whose constant-Q phase advance predicts the voice's phase (e.g. the voice estimate from a previous solve). Default: harmonics of the pitch ('follow') or the mixture ('fixed'). """ import cqt as CQ C = 2 T = g.shape[1] F = Frame(T, N, hop) Gmix = [F.fwd(g[c]) for c in range(C)] E = phase_correction(F, sum(Gmix) / C) Msk = comb_support(F, f0_frames) ft = F.frame_times() use_tv = voice_tv is not None and a_T > 0 if use_tv: fr = CQ.CQTFrame(T, **(cqt_kw or {})) f0c = f0_hz_cqt if f0_hz_cqt is not None else f0_frames ref = fr.fwd(np.asarray(phase_from, dtype=np.float64)) if phase_from is not None else None if voice_tv == "follow": D = CQ.DirectionalTV(fr, ft, f0c, follow=True, phase_ref=ref) else: D = CQ.DirectionalTV(fr, ft, f0c, follow=False, phase_ref=ref if ref is not None else fr.fwd(g.mean(axis=0))) L = {"H": 2.0, "P": 1.0, "D": 2.0, "T": 2.0} gamma = gamma or {"H": 1.0, "P": 1.0, "D": 1.0, "T": 1.0 if use_tv else 0.0} sig = {k: (gamma[k] / L[k] ** 2 if gamma[k] > 0 else 0.0) for k in L} tau = 0.99 / sum(v for v in gamma.values()) z = np.zeros((F.nF, N), np.complex64) h = np.zeros((C, T), np.float32); p = np.zeros((C, T), np.float32) zb, hb, pb = z.copy(), h.copy(), p.copy() yH = [np.zeros((F.nF, N), np.complex64) for _ in range(C)] yP = [np.zeros((F.nF, N), np.complex64) for _ in range(C)] yT = np.zeros((fr.K, fr.M - 1), np.complex128) if use_tv else None if mp3 is None: yD = np.zeros((C, T), np.float64) # time domain, equality with g else: yD = np.zeros((mp3.S.bs["G"], 2, 576)) voice = lambda zz: F.adj(Msk * zz) for it in range(iters): vb = voice(zb) for c in range(C): w = yH[c] + sig["H"] * Dt(E * F.fwd(hb[c])) yH[c] = w / np.maximum(1.0, np.abs(w) / a_H) w = yP[c] + sig["P"] * F.fwd(pb[c]) yP[c] = w / np.maximum(1.0, np.linalg.norm(w, axis=1, keepdims=True) / a_P) mixy = np.stack([vb + hb[0] + pb[0], vb + hb[1] + pb[1]]).astype(np.float64) if mp3 is None: # prox of sigma * iota_{=g}^* w = yD + sig["D"] * mixy yD = w - sig["D"] * g else: # prox of sigma * iota_box^* w = yD + sig["D"] * mp3.fwd(mixy) yD = w - sig["D"] * np.clip(w / sig["D"], mp3.lo, mp3.hi) if use_tv: w = yT + sig["T"] * D(fr.fwd(vb.astype(np.float64))) yT = w / np.maximum(1.0, np.abs(w) / a_T) # primal z_old, h_old, p_old = z, h, p aD = yD if mp3 is None else mp3.adj(yD) aD = aD.astype(np.float32) gv = aD[0] + aD[1] if use_tv: gv = gv + fr.adj(D.T(yT)).astype(np.float32) w = z - tau * (Msk * F.fwd(gv)) mag = np.abs(w) z = (Msk * w * np.maximum(0.0, 1.0 - tau * a_V / np.maximum(mag, 1e-12))).astype(np.complex64) h = np.stack([h[c] - tau * (F.adj(np.conj(E) * DtT(yH[c])) + aD[c]) for c in range(C)]) p = np.stack([p[c] - tau * (F.adj(yP[c]) + aD[c]) for c in range(C)]) zb, hb, pb = 2 * z - z_old, 2 * h - h_old, 2 * p - p_old if verbose and (it % 50 == 0 or it == iters - 1): v = voice(z) res = np.sqrt(sum(np.sum((g[c] - v - h[c] - p[c]) ** 2) for c in range(C)) / np.sum(g ** 2)) print(f" it {it:4d} residual vs decoded mix {res:.4f}", flush=True) return {"v": voice(z), "h": h, "p": p, "z": z, "support": Msk, "frame": F} def separate_outer(g, f0_frames, outer=2, voice_tv="follow", a_T=0.0, f0_from_voice=True, **kw): """Outer loop for the pitch-following TV: solve, then re-estimate the voice's phase (and, optionally, its pitch) from the recovered voice, and solve again with the TV predicting that phase.""" import voiceline as VL out = separate2(g, f0_frames, a_T=0.0, voice_tv=None, **kw) f0 = f0_frames for _ in range(outer - 1): v = out["v"].astype(np.float64) if f0_from_voice: fv = VL.track(np.stack([v, v]), out["frame"], centre_power=0, vibrato=0.0, voicing=0.03, fold=True) f0 = np.where(fv > 0, fv, f0) out = separate2(g, f0_frames, a_T=a_T, voice_tv=voice_tv, f0_hz_cqt=f0, phase_from=v, **kw) return out