1"""Fits simulation parameters to a target clip by CMA-ES over rendered batches.
3 python3 ml/fit.py TARGET.wav --out FIT_DIR [--init E=1.7e11,h=8e-4,...] [--generations 12] [--popsize 16]
5Each generation renders one batch of parameter sets with scripts/render-batch.ts
6(one simulation per worker in parallel on the GPU), scores every clip against
7the target, and updates the CMA-ES distribution. The free parameters, in a
8normalized log (or linear) space, are
10 E, h, Lx, Ly, sigma0, sigma1, strikeVelocity, strikeFracX, strikeFracY, bow sag (across the width)
12with rho fixed at the default (E and rho enter the sound only through E / rho
13at a fixed level, so one of them is redundant) and a single parabolic bow as
14the shape. The objective combines the log-mel distance of match.py (timbre and
15decay) with the fine-structure distances of the early and whole-clip spectra
16(the positions of the partials). Every generation's best clip is kept in
17FIT_DIR/best-GEN.wav and a listening page FIT_DIR/fit.html is written.
18"""
20import argparse
21import json
22import math
23import os
24import shutil
25import subprocess
26import sys
27import time
29import numpy as np
31sys.path.insert(0, os.path.dirname(os.path.abspath(__file__)))
32from match import read_wav_any, onset, segment, features # noqa: E402
33import linesim # noqa: E402
35ROOT = os.path.join(os.path.dirname(os.path.abspath(__file__)), "..")
36NODE = ["node", "--experimental-strip-types", os.path.join(ROOT, "scripts", "render-batch.ts")]
38# name, kind (log or lin), lower, upper, default
39SPACE = [
40 ("E", "log", 60e9, 300e9, 170e9),
41 ("h", "log", 0.3e-3, 2.0e-3, 0.79e-3),
42 ("Lx", "log", 6 * 0.0254, 30 * 0.0254, 12 * 0.0254),
43 ("Ly", "log", 6 * 0.0254, 30 * 0.0254, 18 * 0.0254),
44 ("sigma0", "log", 0.15, 5.0, 1.15),
45 ("sigma1", "log", 0.5e-4, 3e-3, 5e-4),
46 ("strikeVelocity", "log", 0.8, 6.0, 3.0),
47 ("strikeFracX", "lin", 0.5, 0.92, 0.7),
48 ("strikeFracY", "lin", 0.12, 0.55, 0.35),
49 ("bow", "lin", -6e-3, 6e-3, 3e-3),
50 ("bowY", "lin", -6e-3, 6e-3, 0.0),
51]
52FREE = [n for n, *_ in SPACE] # set from --free
55class SimpleCMA:
56 """Minimal (mu/mu_w, lambda)-CMA-ES (Hansen's tutorial, without the
57 rank-mu weights beyond the elite) for when the cma package is absent.
58 Candidates are clipped to the unit cube."""
60 def __init__(self, x0, sigma, popsize, seed=1):
61 self.n = len(x0)
62 self.mean = np.array(x0, float)
63 self.sigma = sigma
64 self.lam = popsize
65 self.mu = popsize // 2
66 w = np.log(self.mu + 0.5) - np.log(np.arange(1, self.mu + 1))
67 self.w = w / w.sum()
68 self.mueff = 1 / np.sum(self.w ** 2)
69 n = self.n
70 self.cc = (4 + self.mueff / n) / (n + 4 + 2 * self.mueff / n)
71 self.cs = (self.mueff + 2) / (n + self.mueff + 5)
72 self.c1 = 2 / ((n + 1.3) ** 2 + self.mueff)
73 self.cmu = min(1 - self.c1, 2 * (self.mueff - 2 + 1 / self.mueff) / ((n + 2) ** 2 + self.mueff))
74 self.damps = 1 + 2 * max(0, math.sqrt((self.mueff - 1) / (n + 1)) - 1) + self.cs
75 self.chiN = math.sqrt(n) * (1 - 1 / (4 * n) + 1 / (21 * n * n))
76 self.pc = np.zeros(n); self.ps = np.zeros(n); self.C = np.eye(n)
77 self.rng = np.random.default_rng(seed)
78 self.gen = 0
80 def ask(self):
81 self.B, self.D = None, None
82 vals, vecs = np.linalg.eigh(self.C)
83 self.D = np.sqrt(np.maximum(vals, 1e-20)); self.B = vecs
84 self.z = self.rng.standard_normal((self.lam, self.n))
85 self.y = (self.B @ (self.D[:, None] * self.z.T)).T
86 X = np.clip(self.mean + self.sigma * self.y, 0, 1)
87 return [x for x in X]
89 def tell(self, X, scores):
90 order = np.argsort(scores)
91 X = np.array(X)
92 old = self.mean.copy()
93 self.mean = self.w @ X[order[: self.mu]]
94 y_w = (self.mean - old) / self.sigma
95 C_inv_sqrt = self.B @ np.diag(1 / self.D) @ self.B.T
96 self.ps = (1 - self.cs) * self.ps + math.sqrt(self.cs * (2 - self.cs) * self.mueff) * (C_inv_sqrt @ y_w)
97 self.gen += 1
98 hsig = np.linalg.norm(self.ps) / math.sqrt(1 - (1 - self.cs) ** (2 * self.gen)) / self.chiN < 1.4 + 2 / (self.n + 1)
99 self.pc = (1 - self.cc) * self.pc + hsig * math.sqrt(self.cc * (2 - self.cc) * self.mueff) * y_w
100 ys = (X[order[: self.mu]] - old) / self.sigma
101 rank_mu = sum(wi * np.outer(yi, yi) for wi, yi in zip(self.w, ys))
102 self.C = (1 - self.c1 - self.cmu) * self.C + self.c1 * (np.outer(self.pc, self.pc) + (1 - hsig) * self.cc * (2 - self.cc) * self.C) + self.cmu * rank_mu
103 self.sigma *= math.exp((self.cs / self.damps) * (np.linalg.norm(self.ps) / self.chiN - 1))
106FIXED = {} # name -> value for the parameters that are not free
109def free_space():
110 return [e for e in SPACE if e[0] in FREE]
113def to_unit(vals):
114 """Physical values of the free parameters -> [0, 1] coordinates."""
115 u = []
116 for (name, kind, lo, hi, _), v in zip(free_space(), vals):
117 u.append((math.log(v / lo) / math.log(hi / lo)) if kind == "log" else (v - lo) / (hi - lo))
118 return np.array(u)
121def from_unit(u):
122 vals = []
123 for (name, kind, lo, hi, _), x in zip(free_space(), np.clip(u, 0, 1)):
124 vals.append(lo * (hi / lo) ** x if kind == "log" else lo + (hi - lo) * x)
125 return vals
128def params_of(vals):
129 """Full parameter dict (free values plus the fixed ones) with the shape as a feature string."""
130 d = {name: d0 for name, _, _, _, d0 in SPACE}
131 d.update(FIXED)
132 d.update({name: v for (name, *_), v in zip(free_space(), vals)})
133 bow, bow_y = d.pop("bow"), d.pop("bowY")
134 feats = [f"bow:x,{bow:.6g},0"] if abs(bow) > 1e-5 else []
135 if abs(bow_y) > 1e-5:
136 feats.append(f"bow:y,{bow_y:.6g},0")
137 d["features"] = ";".join(feats) if feats else "none"
138 return d
141def objective(ref, seg, tol=None):
142 """Default: the mixed log-mel and fine-structure objective. With tol set: the tolerant line measure of linesim.py."""
143 if tol is not None:
144 d, parts = linesim.distance(linesim.picture(ref, tol), linesim.picture(seg, tol))
145 return d, parts
146 fe = features(seg)
147 d_mel = float(np.mean(np.abs(ref[3] - fe[3])))
148 d_early = 1 - float(ref[0] @ fe[0])
149 d_whole = 1 - float(ref[1] @ fe[1])
150 return d_mel + 4 * d_early + 4 * d_whole, (d_mel, d_early, d_whole)
153def render(jobs, out, duration, workers):
154 os.makedirs(out, exist_ok=True)
155 jf = os.path.join(out, "jobs.json")
156 with open(jf, "w") as f:
157 json.dump({"jobs": jobs}, f)
158 env = dict(os.environ)
159 env.pop("DISPLAY", None)
160 subprocess.run(NODE + ["--batch", jf, "--out", out, "--workers", str(workers), "--duration", str(duration)], check=True, env=env, cwd=ROOT,
161 stdout=open(os.path.join(out, "render.log"), "a"), stderr=subprocess.STDOUT)
164def main():
165 ap = argparse.ArgumentParser()
166 ap.add_argument("target")
167 ap.add_argument("--out", required=True)
168 ap.add_argument("--init", default="", help="comma-separated name=value in SI for the starting point")
169 ap.add_argument("--sigma", type=float, default=0.15, help="initial CMA-ES step in the unit cube")
170 ap.add_argument("--generations", type=int, default=12)
171 ap.add_argument("--popsize", type=int, default=16)
172 ap.add_argument("--workers", type=int, default=16)
173 ap.add_argument("--duration", type=float, default=3.0)
174 ap.add_argument("--free", default=",".join(FREE), help="comma-separated names of the free parameters")
175 ap.add_argument("--set", default="", help="comma-separated name=value (SI) for fixed parameters that differ from the defaults")
176 ap.add_argument("--linesim", action="store_true", help="use the tolerant line measure as the objective")
177 ap.add_argument("--tol-start", type=float, default=0.05)
178 ap.add_argument("--tol-end", type=float, default=0.02)
179 args = ap.parse_args()
180 FREE[:] = [n.strip() for n in args.free.split(",") if n.strip()]
181 for kv in args.set.split(","):
182 if "=" in kv:
183 k, v = kv.split("="); FIXED[k.strip()] = float(v)
184 os.makedirs(args.out, exist_ok=True)
185 x = read_wav_any(args.target)
186 st = onset(x)
187 ref_seg = segment(x, st)
188 ref = ref_seg if args.linesim else features(ref_seg)
189 init = {name: d for name, _, _, _, d in SPACE}
190 init.update(FIXED)
191 for kv in args.init.split(","):
192 if "=" in kv:
193 k, v = kv.split("=")
194 init[k.strip()] = float(v)
195 x0 = to_unit([init[name] for name, *_ in free_space()])
196 print(f"target {args.target}: onset {st / 48000:.3f} s; start {params_of(from_unit(x0))}", flush=True)
198 try:
199 import cma
200 es = cma.CMAEvolutionStrategy(x0, args.sigma, {"popsize": args.popsize, "bounds": [0, 1], "seed": 1, "verbose": -9})
201 except ImportError:
202 es = SimpleCMA(x0, args.sigma, args.popsize)
203 log = open(os.path.join(args.out, "fit.jsonl"), "a")
204 best_overall = (float("inf"), None, None)
205 history = []
206 t0 = time.time()
207 for gen in range(args.generations):
208 tol = args.tol_start + (args.tol_end - args.tol_start) * gen / max(1, args.generations - 1) if args.linesim else None
209 X = es.ask()
210 jobs = [{"id": f"g{gen:02d}-{k:02d}", "params": params_of(from_unit(u))} for k, u in enumerate(X)]
211 gdir = os.path.join(args.out, f"gen{gen:02d}")
212 render(jobs, gdir, args.duration, args.workers)
213 scores = []
214 for j, u in zip(jobs, X):
215 p = os.path.join(gdir, "clips", j["id"] + ".wav")
216 if not os.path.exists(p):
217 scores.append(1e3); continue
218 s, parts = objective(ref, segment(read_wav_any(p), 0), tol)
219 scores.append(s)
220 log.write(json.dumps({"gen": gen, "id": j["id"], "score": s, "tol": tol, "parts": parts, "params": j["params"]}) + "\n")
221 log.flush()
222 es.tell(X, scores)
223 k = int(np.argmin(scores))
224 best = (scores[k], jobs[k]["params"], os.path.join(gdir, "clips", jobs[k]["id"] + ".wav"))
225 shutil.copy(best[2], os.path.join(args.out, f"best-{gen:02d}.wav"))
226 if best[0] < best_overall[0]:
227 best_overall = best
228 history.append({"gen": gen, "best": best[0], "mean": float(np.mean(scores)), "params": best[1]})
229 print(f"gen {gen}: best {best[0]:.3f} (mean {np.mean(scores):.3f}, overall {best_overall[0]:.3f}{f', tol {tol:.3f}' if tol else ''}) after {(time.time() - t0) / 60:.1f} min; "
230 + " ".join(f"{n}={v:.4g}" for n, v in best[1].items() if n != "features") + f" {best[1]['features']}", flush=True)
231 write_page(args, history, best_overall)
232 print("best:", json.dumps(best_overall[1]))
235def write_page(args, history, best_overall):
236 rows = "".join(
237 f"<tr><td>{h['gen']}</td><td>{h['best']:.3f}</td><td>{h['mean']:.3f}</td><td><audio controls preload='none' src='best-{h['gen']:02d}.wav'></audio></td>"
238 + "".join(f"<td>{h['params'][n]:.4g}</td>" for n, *_ in SPACE if n not in ("bow", "bowY")) + f"<td class='l'>{h['params']['features']}</td></tr>" for h in history)
239 html = f"""<!doctype html><meta charset="utf-8"><title>fit to {os.path.basename(args.target)}</title>
240<style>body{{font:14px system-ui,sans-serif;margin:16px}} table{{border-collapse:collapse}} td,th{{padding:3px 8px;border-bottom:1px solid #ddd;text-align:right;white-space:nowrap}} td.l{{text-align:left}} audio{{height:28px;width:200px}}</style>
241<h1>Simulation fitted to {os.path.basename(args.target)}</h1>
242<p>Target: <audio controls src="target.wav"></audio> CMA-ES, population {args.popsize}, objective = log-mel distance + 4 (early fine-structure distance + whole-clip fine-structure distance). Best so far: {best_overall[0]:.3f}.</p>
243<table><tr><th>gen</th><th>best</th><th>mean</th><th>best clip</th>{"".join(f"<th>{n}</th>" for n, *_ in SPACE if n not in ("bow", "bowY"))}<th class="l">shape</th></tr>{rows}</table>"""
244 with open(os.path.join(args.out, "fit.html"), "w") as f:
245 f.write(html)
248if __name__ == "__main__":
249 main()