soothe2-re/scripts/fit_vlaw_by_group.py

#!/usr/bin/env python3
"""
fit_vlaw_by_group.py — Fit VLAW parameters (α, β, c, Δ) per configuration group.

VLAW model (framed_model.cpp:205-208):
  cs = α * log1p(lvl / β) + c + (delta ? Δ : 0)
  applied_gain = 10^(-cs / 20)  [gamma0=1 already absorbed into α,c,Δ]

Need to fit these for each (fc, q, sens) configuration group:
  t1kq: fc=800..1200, q=1.0, sens=12 (input tone1kq)
  t1k: fc=500..2000, q=1.0, sens=12 (input tone1k)
  al: fc=1000, q=1.0, sens=3..24 (input lvl_tone_lvX)
  res: fc=300..700, q=1.0, sens=12 (input resonant)
  dual: fc=500, q=0.1..10.0, sens=12 (input dual)
"""

import numpy as np
import os
import sys
import subprocess
import json

sys.path.insert(0, '/home/m/re-tools/scripts')
import corpus

corpus.RB = '/home/m/re-tools/dsp/build/render48k'

# Reference errors from baseline_bridge.json (target)
with open('scripts/baseline_bridge.json') as f:
    REF_ERRORS = json.load(f)

def structural_cases():
    out = []
    for name, inp, args, ref, f in corpus.build_cases():
        joined = [','.join(args)] if len(args) == 3 else args
        out.append((name, inp, joined, ref, f))
    return out

def run_vlaw(inp, args, alpha, beta, c, delta):
    """Run render48k with VLAW parameters and return output path."""
    out = f'/tmp/vlaw_fit_{alpha}_{beta}_{c}_{delta}_{os.path.basename(inp)}.wav'
    env = {
        **os.environ,
        'RT_VLAW': '1',
        'RT_VLAW_ALPHA': str(alpha),
        'RT_VLAW_BETA': str(beta),
        'RT_VLAW_C': str(c),
        'RT_VLAW_DELTA': str(delta),
        'RT_SYN': '1',
        'RT_NOWARP': '1',
        'RT_NOIIR3': '1',
        'RT_IIR12': '0',
    }
    subprocess.run(
        [corpus.RB, inp, out] + args,
        capture_output=True, text=True, env=env,
        cwd='/home/m/re-tools'
    )
    return out

def eval_error(out, ref, f):
    """Evaluate error in dB between output and reference at frequency f."""
    if not os.path.exists(out) or os.path.getsize(out) == 0:
        return None
    ref_sig = corpus.load_mono(ref)
    out_sig = corpus.load_mono(out)
    min_len = min(len(ref_sig), len(out_sig))
    ref_sig = ref_sig[-min_len:]
    out_sig = out_sig[-min_len:]
    ref_ta = corpus.ta(ref_sig, f)
    out_ta = corpus.ta(out_sig, f)
    return corpus.db(out_ta / ref_ta)

def group_key(name):
    return name.split('_')[0]

def evaluate_params(alpha, beta, c, delta, cases_subset=None):
    """Evaluate VLAW params on all cases, return per-group mean abs error."""
    all_cases = structural_cases()
    if cases_subset:
        all_cases = [c for c in all_cases if group_key(c[0]) in cases_subset]
    
    errs = {}
    
    for name, inp, args, ref, f in all_cases:
        out = run_vlaw(inp, args, alpha, beta, c, delta)
        err = eval_error(out, ref, f)
        if err is not None:
            errs[name] = err
    
    # Group stats
    groups = {}
    for k, v in errs.items():
        g = group_key(k)
        groups.setdefault(g, []).append(v)
    
    out_stats = {g: float(np.mean(np.abs(v))) for g, v in groups.items()}
    out_stats['TOTAL'] = float(np.mean(np.abs(list(errs.values()))))
    return out_stats, errs

def fit_single_case(name, inp, args, ref, f, init_params):
    """Grid search for best params on a single case."""
    alpha0, beta0, c0, delta0 = init_params
    best = None
    best_err = float('inf')
    
    # Search around initial params
    alphas = np.linspace(max(0.5, alpha0-1), alpha0+1, 9)
    betas = np.linspace(max(0.1, beta0-0.2), beta0+0.2, 9)
    cs = np.linspace(max(0.0, c0-0.5), c0+0.5, 9)
    deltas = np.linspace(max(0.0, delta0-2), delta0+2, 9)
    
    for alpha in alphas:
        for beta in betas:
            for c in cs:
                for delta in deltas:
                    out = run_vlaw(inp, args, alpha, beta, c, delta)
                    err = eval_error(out, ref, f)
                    if err is not None and abs(err) < best_err:
                        best_err = abs(err)
                        best = (alpha, beta, c, delta, err)
                        print(f'  {name}: new best α={alpha:.3f}, β={beta:.3f}, c={c:.3f}, Δ={delta:.3f} => err={err:.3f} dB')
    
    return best

def main():
    # Build case map by group
    all_cases = structural_cases()
    groups = {}
    for name, inp, args, ref, f in all_cases:
        g = group_key(name)
        groups.setdefault(g, []).append((name, inp, args, ref, f))
    
    print("Available groups:", list(groups.keys()))
    for g, cases in groups.items():
        print(f"  {g}: {len(cases)} cases")
    
    # Current calibrated params for dual(q=0.5)
    dual_params = (3.2193, 0.4927, 0.5423, 6.9177)
    
    # Test current params on all groups
    print("\n=== Testing current dual params on all groups ===")
    stats, _ = evaluate_params(*dual_params)
    for g in ['t1kq', 't1k', 'al', 'res', 'dual', 'comb']:
        if g in stats:
            print(f'  {g}: {stats[g]:.3f} dB')
    
    # For each group, pick a representative case and fit
    print("\n=== Fitting per group (representative case) ===")
    results = {}
    
    # For dual, use q=0.5 as reference (already calibrated)
    if 'dual' in groups:
        # Find q=0.5 case
        for name, inp, args, ref, f in groups['dual']:
            if '0.5' in name:
                best = fit_single_case(name, inp, args, ref, f, dual_params)
                if best:
                    results['dual'] = best[:4]
                break
    
    # For t1kq, use fc=1000
    if 't1kq' in groups:
        for name, inp, args, ref, f in groups['t1kq']:
            if '1000' in name:
                best = fit_single_case(name, inp, args, ref, f, dual_params)
                if best:
                    results['t1kq'] = best[:4]
                break
    
    # For t1k, use fc=1000
    if 't1k' in groups:
        for name, inp, args, ref, f in groups['t1k']:
            if '1000' in name:
                best = fit_single_case(name, inp, args, ref, f, dual_params)
                if best:
                    results['t1k'] = best[:4]
                break
    
    # For al, use sens=12
    if 'al' in groups:
        for name, inp, args, ref, f in groups['al']:
            if '12' in name:
                best = fit_single_case(name, inp, args, ref, f, dual_params)
                if best:
                    results['al'] = best[:4]
                break
    
    # For res, use fc=500
    if 'res' in groups:
        for name, inp, args, ref, f in groups['res']:
            if '500' in name:
                best = fit_single_case(name, inp, args, ref, f, dual_params)
                if best:
                    results['res'] = best[:4]
                break
    
    # Print results
    print("\n=== FITTED VLAW PARAMETERS BY GROUP ===")
    for g, (alpha, beta, c, delta) in results.items():
        print(f'{g}: α={alpha:.4f}, β={beta:.4f}, c={c:.4f}, Δ={delta:.4f}')
    
    # Save to JSON
    with open('/tmp/opencode/vlaw_params.json', 'w') as f:
        json.dump({g: {'alpha': a, 'beta': b, 'c': c, 'delta': d} 
                   for g, (a, b, c, d) in results.items()}, f, indent=2)
    print('\nSaved to /tmp/opencode/vlaw_params.json')

if __name__ == '__main__':
    main()