example train colorless FDN

%23%20gallery_category%3A%20Optimization%0A%23%20gallery_title%3A%20Train%20a%20colorless%20FDN%0A%23%20gallery_description%3A%20Optimize%20an%20FDN%20for%20a%20flat%20lossless%20magnitude%20response%2C%20extract%20its%20build%2C%20and%20add%20decay%20for%20listening.%0A%0Aimport%20marimo%0A%0A__generated_with%20%3D%20%220.24.0%22%0Aapp%20%3D%20marimo.App()%0A%0A%0A%40app.cell%0Adef%20_()%3A%0A%20%20%20%20import%20marimo%20as%20mo%0A%0A%20%20%20%20return%20(mo%2C)%0A%0A%0A%40app.cell(hide_code%3DTrue)%0Adef%20_(mo%2C%20pyFDN)%3A%0A%20%20%20%20mo.md(f%22%22%22%0A%20%20%20%20%23%20Colorless%20FDN%2C%20trained%20in-notebook%0A%0A%20%20%20%20The%20companion%20to%20**Colorless%20FDN**%2C%20which%20loads%20pre-optimized%2C%20build-ready%20JSON%20presets.%20Here%20we%20run%20the%20optimization%20ourselves%20with%20%60pyFDN%60's%20training%20API%2C%20following%20*%7BpyFDN.paper_link(%22Differentiable_FDN_For_Colorless_Reverberation%22)%7D*%20(and%20its%20%22tiny%20colorless%20FDN%22%20follow-up)%3A%0A%0A%20%20%20%201.%20%60pyFDN.build_fdn%60%20--%20a%20standard%20FDN%20skeleton%20with%20random%20orthogonal%20feedback%20matrix.%0A%20%20%20%202.%20%60pyFDN.train_fdn(model%2C%20loss)%60%20--%20optimize%20the%20feedback%20matrix%20and%20gains%20toward%20an%20objective%20we%20write%20out%20in%20full%2C%20in%20place.%0A%20%20%20%203.%20%60pyFDN.extract_build%60%20--%20read%20both%20the%20initial%20and%20optimized%20FDNs%20back%20out.%0A%0A%20%20%20%20Delays%20stay%20fixed.%20We%20then%20add%20homogeneous%20decay%20so%20the%20result%20is%20audible.%0A%0A%20%20%20%20%23%23%20The%20objective%2C%20spelled%20out%0A%0A%20%20%20%20Every%20loss%20is%20a%20function%20of%20the%20model's%20**impulse%20response**%2C%20so%20an%20objective%20is%20a%20sum%20of%20named%20parts%20rather%20than%20a%20mode%20to%20select%3A%0A%0A%20%20%20%20%60%60%60python%0A%20%20%20%20loss%20%3D%20pyFDN.FlatMagnitude()%20%2B%200.2%20*%20pyFDN.Sparsity(pyFDN.param(model%2C%20%22feedback%22))%0A%20%20%20%20%60%60%60%0A%0A%20%20%20%20*%20%60pyFDN.FlatMagnitude()%60%20reads%20%60pyFDN.Response.magnitude%60%20--%20the%20%60rfft%60%20of%20the%20impulse%20response%20--%20and%20fits%20it%20to%20a%20constant.%20That%20is%20the%20colorless%20objective.%0A%20%20%20%20*%20%60pyFDN.Sparsity(...)%60%20is%20a%20cost%20on%20a%20*model%20parameter*%20rather%20than%20on%20the%20response%2C%20so%20it%20has%20to%20name%20one%3A%20%60pyFDN.param(model%2C%20%22feedback%22)%60%20resolves%20the%20feedback%20matrix%20in%20this%20model's%20graph%2C%20and%20%60pyFDN.params(model)%60%20lists%20everything%20else%20on%20offer.%20It%20rewards%20a%20dense%20matrix%2C%20i.e.%20good%20mixing.%0A%0A%20%20%20%20%60FlatMagnitude%60%20weighs%20a%20peak%20and%20a%20dip%20by%20whatever%20the%20squared%20error%20happens%20to%20make%20of%20them.%20The%20last%20section%20of%20this%20notebook%20swaps%20it%20for%20%60pyFDN.AsymmetricFlatMagnitude%60%2C%20which%20weighs%20them%20on%20purpose.%0A%0A%20%20%20%20%23%23%20Two%20details%20that%20make%20or%20break%20the%20fit%0A%0A%20%20%20%20*%20**Rendering.**%20A%20**lossless**%20FDN%20has%20every%20pole%20*exactly*%20on%20the%20unit%20circle%2C%20where%20the%20FFT-domain%20evaluation%20is%20near-singular%20and%20the%20impulse%20response%20comes%20out%20wrong.%20%60build_fdn(rt%3DNone)%60%20therefore%20defaults%20to%20%60alias_decay_db%3D%60%7BpyFDN.LOSSLESS_ALIAS_DECAY_DB%7D%3A%20flamo%20evaluates%20the%20system%20on%20a%20slightly%20smaller%20circle%2C%20and%20the%20%60%22time%22%60%20output%20layer%20removes%20that%20%24%5C%5Cgamma%5En%24%20envelope%20again.%20What%20a%20loss%20sees%20is%20the%20impulse%20response%20*itself*%2C%20accurate%20to%20%7BpyFDN.LOSSLESS_ALIAS_DECAY_DB%3A.0f%7D%20dB%20--%20which%20is%20all%20%60alias_decay_db%60%20means.%20It%20never%20reaches%20the%20extracted%20build.%0A%20%20%20%20*%20**Reading%20the%20result.**%20The%20numbers%20below%20are%20the%20loss%20terms%20themselves%3A%20%60TrainLog.loss_log%60%20keeps%20each%20term's%20history%20*unweighted*%2C%20so%20%60FlatMagnitude%60%20is%20exactly%20the%20mean%20squared%20deviation%20of%20%24%7CH%7C%24%20from%201%20that%20was%20optimized%2C%20and%20%60Sparsity%60%20the%20density%20of%20the%20feedback%20matrix.%20They%20compare%20across%20runs%20at%20the%20same%20%60nfft%60%20--%20but%20not%20across%20different%20%60nfft%60%2C%20since%20truncating%20to%20a%20longer%20window%20widens%20the%20peak-to-median%20range%20of%20%24%7CH%7C%24%20and%20the%20same%20FDN%20then%20scores%20a%20larger%20MSE.%0A%0A%20%20%20%20-%20pyFDN%20training%20pipeline%3A%20Jeremy%20B.%20Bai%2C%202026-06-19%0A%20%20%20%20%22%22%22)%0A%20%20%20%20return%0A%0A%0A%40app.cell%0Adef%20_()%3A%0A%20%20%20%20import%20numpy%20as%20np%0A%0A%20%20%20%20import%20pyFDN%0A%0A%20%20%20%20return%20np%2C%20pyFDN%0A%0A%0A%40app.cell%0Adef%20_(pyFDN)%3A%0A%20%20%20%20fs%20%3D%2048000%0A%20%20%20%20%23%20The%20objective's%20frequency%20resolution%20is%20this%20nfft%20and%20nothing%20else%3A%20it%20is%0A%20%20%20%20%23%20the%20rfft%20of%20an%20nfft-long%20impulse%20response.%20More%20of%20it%20resolves%20the%20modes%0A%20%20%20%20%23%20more%20finely%20and%20fits%20them%20better%2C%20at%20proportionally%20more%20time%20per%20step.%0A%20%20%20%20nfft%20%3D%202**14%0A%0A%20%20%20%20%23%201.%20build%20a%20small%20%22tiny%20colorless%22%20lossless%20skeleton.%20rt%3DNone%20also%20switches%0A%20%20%20%20%23%20%20%20%20on%20the%20default%20anti-aliasing%20decay%2C%20without%20which%20a%20lossless%20FDN's%0A%20%20%20%20%23%20%20%20%20impulse%20response%20cannot%20be%20rendered%20at%20all.%0A%20%20%20%20delays%20%3D%20pyFDN.sample_delay_lengths(%0A%20%20%20%20%20%20%20%2016%2C%20(200%2C%20600)%2C%20distribution%3D%22geometric%22%2C%20coprime%3DTrue%2C%20sort%3DTrue%2C%20rng%3D2%0A%20%20%20%20)%0A%20%20%20%20model%20%3D%20pyFDN.build_fdn(delays%3Ddelays%2C%20rt%3DNone%2C%20nfft%3Dnfft%2C%20device%3D%22cpu%22%2C%20rng%3D2)%0A%20%20%20%20init_build%20%3D%20pyFDN.extract_build(model)%20%20%23%20random%20init%2C%20before%20training%0A%0A%20%20%20%20%23%20%7CH%7C%20exactly%20as%20the%20loss%20sees%20it%3A%20the%20rfft%20of%20the%20model's%20impulse%20response.%0A%20%20%20%20mag_init%20%3D%20pyFDN.model_response(model).magnitude.detach().numpy().squeeze()%0A%0A%20%20%20%20%23%202.%20write%20the%20objective%20out%3A%20a%20flat%20magnitude%20response%2C%20plus%20a%20density%20reward%20on%20this%20model's%20feedback%20matrix.%0A%20%20%20%20loss%20%3D%20pyFDN.FlatMagnitude()%20%2B%200.2%20*%20pyFDN.Sparsity(pyFDN.param(model%2C%20%22feedback%22))%0A%20%20%20%20%23%20loss%20%3D%20pyFDN.AsymmetricFlatMagnitude(peak_power%3D8)%0A%0A%20%20%20%20%23%203.%20train%20in%20place%3B%20then%20extract.%0A%20%20%20%20log%20%3D%20pyFDN.train_fdn(%0A%20%20%20%20%20%20%20%20model%2C%0A%20%20%20%20%20%20%20%20loss%2C%0A%20%20%20%20%20%20%20%20optimizer%3D%22adam%22%2C%0A%20%20%20%20%20%20%20%20max_steps%3D200%2C%0A%20%20%20%20%20%20%20%20lr%3D8%20*%201e-2%2C%0A%20%20%20%20%20%20%20%20patience%3D50%2C%0A%20%20%20%20%20%20%20%20device%3D%22cpu%22%2C%0A%20%20%20%20%20%20%20%20rng%3D1%2C%0A%20%20%20%20)%0A%20%20%20%20opt_build%20%3D%20pyFDN.extract_build(model)%0A%20%20%20%20mag_opt%20%3D%20pyFDN.model_response(model).magnitude.detach().numpy().squeeze()%0A%0A%20%20%20%20%23%204.%20report%20the%20loss%20itself%20--%20the%20total%2C%20then%20each%20term%20unweighted.%0A%20%20%20%20print(%0A%20%20%20%20%20%20%20%20f%22ran%20%7Blog.steps_run%7D%20steps%2C%20total%20loss%20%22%0A%20%20%20%20%20%20%20%20f%22%7Blog.train_loss%5B0%5D%3A.4f%7D%20-%3E%20%7Blog.train_loss%5B-1%5D%3A.4f%7D%22%0A%20%20%20%20)%0A%20%20%20%20for%20_term%2C%20_history%20in%20log.loss_log.items()%3A%0A%20%20%20%20%20%20%20%20print(f%22%20%20%7B_term%3A18%7D%20%7B_history%5B0%5D%3A.4f%7D%20-%3E%20%7B_history%5B-1%5D%3A.4f%7D%22)%0A%20%20%20%20return%20fs%2C%20init_build%2C%20log%2C%20mag_init%2C%20mag_opt%2C%20nfft%2C%20opt_build%0A%0A%0A%40app.cell(hide_code%3DTrue)%0Adef%20_(mo)%3A%0A%20%20%20%20mo.md(r%22%22%22%0A%20%20%20%20%23%23%20Magnitude%20response%20and%20training%20loss%0A%0A%20%20%20%20The%20magnitude%20response%20the%20loss%20actually%20saw%20--%20%60pyFDN.Response.magnitude%60%2C%20the%20rfft%20of%20the%20model's%20%60nfft%60-long%20impulse%20response%20--%20before%20and%20after%20training.%20%60FlatMagnitude%60%20fits%20this%20curve%20to%20the%20dashed%20line%20at%200%20dB%2C%20and%20the%20peaks%20and%20notches%20pull%20in%20toward%20it.%0A%20%20%20%20%22%22%22)%0A%20%20%20%20return%0A%0A%0A%40app.cell%0Adef%20_(fs%2C%20log%2C%20mag_init%2C%20mag_opt%2C%20mo%2C%20nfft%2C%20np%2C%20pyFDN)%3A%0A%20%20%20%20import%20matplotlib.pyplot%20as%20plt%0A%0A%20%20%20%20_freqs%20%3D%20np.fft.rfftfreq(nfft%2C%201.0%20%2F%20fs)%0A%20%20%20%20_fig%2C%20_axes%20%3D%20plt.subplots(1%2C%202%2C%20figsize%3D(11%2C%203.2))%0A%20%20%20%20_axes%5B0%5D.plot(_freqs%2C%20pyFDN.lin_to_db(mag_init)%2C%20alpha%3D0.5%2C%20label%3D%22init%22)%0A%20%20%20%20_axes%5B0%5D.plot(_freqs%2C%20pyFDN.lin_to_db(mag_opt)%2C%20label%3D%22colorless%22)%0A%20%20%20%20_axes%5B0%5D.axhline(0.0%2C%20color%3D%22k%22%2C%20ls%3D%22--%22%2C%20lw%3D0.8%2C%20label%3D%22target%20%7CH%7C%20%3D%201%22)%0A%20%20%20%20_axes%5B0%5D.set(%0A%20%20%20%20%20%20%20%20xlabel%3D%22frequency%20%5BHz%5D%22%2C%0A%20%20%20%20%20%20%20%20ylabel%3D%22magnitude%20%5BdB%5D%22%2C%0A%20%20%20%20%20%20%20%20xscale%3D%22log%22%2C%0A%20%20%20%20%20%20%20%20title%3D%22Magnitude%20response%20as%20the%20loss%20sees%20it%22%2C%0A%20%20%20%20)%0A%20%20%20%20_axes%5B0%5D.legend(fontsize%3D8)%0A%20%20%20%20_axes%5B0%5D.grid(True%2C%20alpha%3D0.3)%0A%0A%20%20%20%20_axes%5B1%5D.plot(log.train_loss%2C%20lw%3D1)%0A%20%20%20%20_axes%5B1%5D.set(xlabel%3D%22step%22%2C%20ylabel%3D%22loss%22%2C%20yscale%3D%22log%22%2C%20title%3D%22Training%20loss%22)%0A%20%20%20%20_axes%5B1%5D.grid(True%2C%20alpha%3D0.3)%0A%20%20%20%20_fig.tight_layout()%0A%20%20%20%20mo.as_html(_fig)%0A%20%20%20%20return%0A%0A%0A%40app.cell(hide_code%3DTrue)%0Adef%20_(mo)%3A%0A%20%20%20%20mo.md(r%22%22%22%0A%20%20%20%20%23%23%20FDN%20parameters%3A%20random%20init%20vs%20colorless%0A%0A%20%20%20%20The%20stored%20%60FDNBuild%60%20parameters%20side%20by%20side%20--%20delay%20lengths%2C%20the%20feedback%20matrix%20%60A%60%2C%20and%20the%20input%2Foutput%2Fdirect%20gains%20%60b%60%2C%20%60c%60%2C%20%60d%60%2C%20on%20a%20shared%20color%20scale.%20Training%20reshapes%20%60A%60%20and%20the%20gains%20(the%20delays%20stay%20fixed)%20to%20flatten%20the%20magnitude%20response.%0A%20%20%20%20%22%22%22)%0A%20%20%20%20return%0A%0A%0A%40app.cell%0Adef%20_(init_build%2C%20mo%2C%20opt_build%2C%20pyFDN)%3A%0A%20%20%20%20_build_init%20%3D%20pyFDN.plot_FDN_build(init_build%2C%20title%3D%22Random%20init%22)%0A%20%20%20%20_build_opt%20%3D%20pyFDN.plot_FDN_build(opt_build%2C%20title%3D%22Colorless%22)%0A%20%20%20%20mo.hstack(%5Bmo.as_html(_build_init)%2C%20mo.as_html(_build_opt)%5D%2C%20gap%3D2)%0A%20%20%20%20return%0A%0A%0A%40app.cell(hide_code%3DTrue)%0Adef%20_(mo)%3A%0A%20%20%20%20mo.md(r%22%22%22%0A%20%20%20%20%23%23%20Listen%3A%20random%20init%20vs%20colorless%0A%0A%20%20%20%20Two%20renderings%20of%20each%20FDN%2C%20built%20end%20to%20end%20through%20the%20render%20API%20(%60pyFDN.build_set_decay%60%20-%3E%20%60pyFDN.build_to_impz%60)%2C%20peak-normalized%20with%20%60pyFDN.peak_normalize%60%20so%20the%20A%2FB%20compares%20*timbre*%2C%20not%20level%3A%0A%0A%20%20%20%20*%20**Long%20tail%20(hear%20the%20colour)**%20--%20a%20very%20long%20reverberation%20time%2C%20so%20the%20FDN%20rings%20with%20almost%20no%20decay%20and%20you%20hear%20its%20colour%20directly%3A%20the%20random%20init%20is%20tonal%2Fmetallic%20(sharp%20modal%20resonances)%2C%20the%20colorless%20one%20is%20noise-like%20(flat%20spectrum).%20%60build_to_impz%60%20renders%20in%20the%20time%20domain%20(%60pyFDN.process_fdn%60)%2C%20so%20the%20long%20ring%20is%20captured%20faithfully%20without%20the%20FFT%20wrap-around%20an%20nfft-length%20frequency-domain%20render%20would%20alias%20back%20onto%20the%20start.%0A%20%20%20%20*%20**Reverb%20tail**%20--%20a%20short%20homogeneous%20%24T_%7B60%7D%24%2C%20giving%20an%20audible%20decay.%0A%20%20%20%20%22%22%22)%0A%20%20%20%20return%0A%0A%0A%40app.cell%0Adef%20_(fs%2C%20init_build%2C%20opt_build%2C%20pyFDN)%3A%0A%20%20%20%20%23%20A%20long%20%22ring%22%20RT%20keeps%20the%20colour%20audible%20with%20little%20decay%3B%20a%20short%20RT%20gives%0A%20%20%20%20%23%20an%20audible%20reverb%20tail.%0A%20%20%20%20rt_ring%2C%20rt_rev%20%3D%2060.0%2C%202.0%0A%20%20%20%20n_samples%20%3D%20int(2.0%20*%20fs)%0A%0A%20%20%20%20def%20render(build%2C%20rt)%3A%0A%20%20%20%20%20%20%20%20%22%22%22build%20(with%20homogeneous%20decay)%20-%3E%20peak-normalized%201-D%20impulse%20response.%22%22%22%0A%20%20%20%20%20%20%20%20ir%20%3D%20pyFDN.build_to_impz(pyFDN.build_set_decay(build%2C%20rt)%2C%20n_samples).squeeze()%0A%20%20%20%20%20%20%20%20%23%20peak-normalize%20so%20the%20A%2FB%20compares%20timbre%20at%20matched%20level%2C%20not%20loudness.%0A%20%20%20%20%20%20%20%20return%20pyFDN.fade_out(pyFDN.peak_normalize(ir)%2C%202048)%0A%0A%20%20%20%20%23%20Long%20%22ring%22%3A%20the%20tail%20is%20still%20loud%20at%20the%20buffer%20end%2C%20so%20fade%20it%20out%20to%0A%20%20%20%20%23%20avoid%20a%20click%20on%20the%20abrupt%20cutoff.%0A%20%20%20%20init_noise%20%3D%20render(init_build%2C%20rt_ring)%0A%20%20%20%20opt_noise%20%3D%20render(opt_build%2C%20rt_ring)%0A%0A%20%20%20%20%23%20Reverb%3A%20a%20short%20homogeneous%20T60%20gives%20an%20audible%20decaying%20tail.%0A%20%20%20%20init_decay%20%3D%20render(init_build%2C%20rt_rev)%0A%20%20%20%20opt_decay%20%3D%20render(opt_build%2C%20rt_rev)%0A%20%20%20%20return%20init_decay%2C%20init_noise%2C%20opt_decay%2C%20opt_noise%0A%0A%0A%40app.cell%0Adef%20_(fs%2C%20init_decay%2C%20init_noise%2C%20mo%2C%20opt_decay%2C%20opt_noise%2C%20pyFDN)%3A%0A%20%20%20%20_plot%20%3D%20pyFDN.plot_impulse_response(%0A%20%20%20%20%20%20%20%20opt_decay%2C%20init_decay%2C%20fs%3Dfs%2C%20labels%3D%5B%22Colorless%22%2C%20%22Random%20init%22%5D%0A%20%20%20%20)%0A%20%20%20%20_audio%20%3D%20mo.hstack(%0A%20%20%20%20%20%20%20%20%5B%0A%20%20%20%20%20%20%20%20%20%20%20%20mo.vstack(%0A%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20%5B%0A%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20mo.Html(%22%3Cb%3ELong%20tail%20(hear%20the%20colour)%3C%2Fb%3E%22).style(%0A%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20%7B%22font-size%22%3A%20%221.2em%22%7D%0A%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20)%2C%0A%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20pyFDN.labeled_audio(%22Random%20init%22%2C%20init_noise%2C%20fs%3Dfs)%2C%0A%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20pyFDN.labeled_audio(%22Colorless%22%2C%20opt_noise%2C%20fs%3Dfs)%2C%0A%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20%5D%2C%0A%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20gap%3D1%2C%0A%20%20%20%20%20%20%20%20%20%20%20%20)%2C%0A%20%20%20%20%20%20%20%20%20%20%20%20mo.vstack(%0A%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20%5B%0A%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20mo.Html(%22%3Cb%3EReverb%20tail%3C%2Fb%3E%22).style(%7B%22font-size%22%3A%20%221.2em%22%7D)%2C%0A%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20pyFDN.labeled_audio(%22Random%20init%22%2C%20init_decay%2C%20fs%3Dfs)%2C%0A%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20pyFDN.labeled_audio(%22Colorless%22%2C%20opt_decay%2C%20fs%3Dfs)%2C%0A%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20%5D%2C%0A%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20gap%3D1%2C%0A%20%20%20%20%20%20%20%20%20%20%20%20)%2C%0A%20%20%20%20%20%20%20%20%5D%2C%0A%20%20%20%20%20%20%20%20gap%3D2%2C%0A%20%20%20%20)%0A%0A%20%20%20%20mo.vstack(%5B_plot%2C%20_audio%5D%2C%20gap%3D3)%0A%20%20%20%20return%0A%0A%0Aif%20__name__%20%3D%3D%20%22__main__%22%3A%0A%20%20%20%20app.run()%0A