Simulate a lens

import%20marimo%0A%0A__generated_with%20%3D%20%220.24.2%22%0Aapp%20%3D%20marimo.App(width%3D%22medium%22%2C%20app_title%3D%22Simulate%20a%20lens%22)%0A%0A%0A%40app.cell%0Adef%20_()%3A%0A%20%20%20%20import%20base64%0A%20%20%20%20import%20io%0A%20%20%20%20import%20json%0A%20%20%20%20import%20sys%0A%0A%20%20%20%20import%20marimo%20as%20mo%0A%20%20%20%20import%20numpy%20as%20np%0A%0A%20%20%20%20return%20base64%2C%20io%2C%20json%2C%20mo%2C%20np%2C%20sys%0A%0A%0A%40app.cell%0Adef%20_(json%2C%20mo%2C%20sys)%3A%0A%20%20%20%20%23%20A%20precomputed%20grid%20of%20simulations%20(notebooks%2Fdata%2Fplayground_grid.json)%2C%20so%20the%20widget%20runs%0A%20%20%20%20%23%20entirely%20in%20the%20browser%3A%20nothing%20is%20simulated%20on%20a%20server%20when%20someone%20moves%20a%20slider.%0A%20%20%20%20def%20_load_text(name%3A%20str)%20-%3E%20str%3A%0A%20%20%20%20%20%20%20%20%23%20In%20the%20browser%20build%20the%20data%20lives%20at%20%2Fmarimo%2Fdata%2F%20on%20the%20site%20(shared%20by%20all%20the%0A%20%20%20%20%20%20%20%20%23%20notebooks%2C%20independent%20of%20where%20the%20export's%20assets%20are)%3B%20locally%20it%20is%20next%20to%20the%20file.%0A%20%20%20%20%20%20%20%20if%20%22pyodide%22%20in%20sys.modules%3A%0A%20%20%20%20%20%20%20%20%20%20%20%20from%20urllib.parse%20import%20urlsplit%0A%0A%20%20%20%20%20%20%20%20%20%20%20%20from%20pyodide.http%20import%20open_url%0A%0A%20%20%20%20%20%20%20%20%20%20%20%20u%20%3D%20urlsplit(str(mo.notebook_location()))%0A%20%20%20%20%20%20%20%20%20%20%20%20return%20open_url(f%22%7Bu.scheme%7D%3A%2F%2F%7Bu.netloc%7D%2Fmarimo%2Fdata%2F%7Bname%7D%22).read()%0A%20%20%20%20%20%20%20%20return%20(mo.notebook_location()%20%2F%20%22data%22%20%2F%20name).read_text()%0A%0A%20%20%20%20data%20%3D%20json.loads(_load_text(%22playground_grid.json%22))%0A%20%20%20%20MF%2C%20LA%2C%20SEEDS%2C%20GRID%20%3D%20data%5B%22mass_fractions%22%5D%2C%20data%5B%22log_axion%22%5D%2C%20data%5B%22seeds%22%5D%2C%20data%5B%22grid%22%5D%0A%20%20%20%20return%20GRID%2C%20LA%2C%20MF%2C%20SEEDS%0A%0A%0A%40app.cell%0Adef%20_(LA%2C%20MF%2C%20SEEDS%2C%20mo)%3A%0A%20%20%20%20substructure%20%3D%20mo.ui.dropdown(options%3D%5B%22vortex%22%2C%20%22subhalo%22%2C%20%22none%22%5D%2C%20value%3D%22vortex%22%2C%20label%3D%22substructure%22)%0A%20%20%20%20mass_fraction%20%3D%20mo.ui.slider(steps%3DMF%2C%20value%3D0.03%2C%20label%3D%22mass%20fraction%22)%0A%20%20%20%20log_axion%20%3D%20mo.ui.slider(steps%3DLA%2C%20value%3D-23.0%2C%20label%3D%22log10%20axion%20mass%20(eV)%22)%0A%20%20%20%20seed%20%3D%20mo.ui.slider(steps%3DSEEDS%2C%20value%3D3%2C%20label%3D%22seed%22)%0A%20%20%20%20return%20log_axion%2C%20mass_fraction%2C%20seed%2C%20substructure%0A%0A%0A%40app.cell%0Adef%20_(log_axion%2C%20mass_fraction%2C%20mo%2C%20seed%2C%20substructure)%3A%0A%20%20%20%20controls%20%3D%20%5Bsubstructure%2C%20seed%5D%0A%20%20%20%20if%20substructure.value%20!%3D%20%22none%22%3A%0A%20%20%20%20%20%20%20%20controls.append(mass_fraction)%0A%20%20%20%20if%20substructure.value%20%3D%3D%20%22vortex%22%3A%0A%20%20%20%20%20%20%20%20controls.append(log_axion)%0A%20%20%20%20mo.md(%0A%20%20%20%20%20%20%20%20f%22%22%22%0A%20%20%20%20The%20simulator%20the%20agent%20drives%2C%20with%20DeepLenseSim's%20defaults%3A%20a%2010%3Csup%3E12%3C%2Fsup%3E%20M%3Csub%3E%E2%98%89%3C%2Fsub%3E%20lens%20at%20z%20%3D%200.5%2C%0A%20%20%20%20a%20S%C3%A9rsic%20source%20at%20z%20%3D%201%2C%20Euclid-like%20pixels.%20Left%3A%20the%20observed%20image.%20Right%3A%20what%20the%20substructure%20adds%20to%20the%20arcs.%0A%0A%20%20%20%20%7Bmo.hstack(controls%2C%20gap%3D2%2C%20justify%3D%22start%22%2C%20wrap%3DTrue)%7D%0A%20%20%20%20%22%22%22%0A%20%20%20%20)%0A%20%20%20%20return%0A%0A%0A%40app.cell%0Adef%20_(GRID%2C%20base64%2C%20io%2C%20log_axion%2C%20mass_fraction%2C%20mo%2C%20np%2C%20seed%2C%20substructure)%3A%0A%20%20%20%20from%20PIL%20import%20Image%0A%0A%20%20%20%20cls%20%3D%20substructure.value%0A%20%20%20%20key%20%3D%20f%22%7Bcls%7D%7C%7B0%20if%20cls%20%3D%3D%20'none'%20else%20mass_fraction.value%7D%7C%7Blog_axion.value%20if%20cls%20%3D%3D%20'vortex'%20else%200%7D%7C%7Bseed.value%7D%22%0A%20%20%20%20entry%20%3D%20GRID%5Bkey%5D%0A%0A%20%20%20%20def%20_u8(b64%3A%20str)%20-%3E%20%22np.ndarray%22%3A%0A%20%20%20%20%20%20%20%20return%20np.frombuffer(base64.b64decode(b64)%2C%20dtype%3Dnp.uint8).reshape(64%2C%2064)%0A%0A%20%20%20%20def%20_png(rgb%3A%20%22np.ndarray%22)%20-%3E%20bytes%3A%0A%20%20%20%20%20%20%20%20buf%20%3D%20io.BytesIO()%0A%20%20%20%20%20%20%20%20Image.fromarray(rgb).resize((320%2C%20320)%2C%20Image.NEAREST).save(buf%2C%20format%3D%22PNG%22)%0A%20%20%20%20%20%20%20%20return%20buf.getvalue()%0A%0A%20%20%20%20gray%20%3D%20_u8(entry%5B%22img%22%5D)%0A%20%20%20%20observed%20%3D%20np.stack(%5Bgray%5D%20*%203%2C%20axis%3D-1)%0A%0A%20%20%20%20%23%20residual%3A%20128%20%3D%20zero%3B%20blue%20below%2C%20red%20above%2C%20white%20at%20zero%20(a%20small%20diverging%20map%2C%20no%20matplotlib)%0A%20%20%20%20r%20%3D%20(_u8(entry%5B%22res%22%5D).astype(float)%20-%20128.0)%20%2F%20127.0%0A%20%20%20%20t%20%3D%20np.abs(r)%5B...%2C%20None%5D%0A%20%20%20%20red%20%3D%20np.array(%5B178%2C%2024%2C%2043%5D%2C%20dtype%3Dfloat)%0A%20%20%20%20blue%20%3D%20np.array(%5B33%2C%20102%2C%20172%5D%2C%20dtype%3Dfloat)%0A%20%20%20%20white%20%3D%20np.array(%5B255%2C%20255%2C%20255%5D%2C%20dtype%3Dfloat)%0A%20%20%20%20col%20%3D%20np.where(r%5B...%2C%20None%5D%20%3E%3D%200%2C%20red%2C%20blue)%0A%20%20%20%20residual%20%3D%20(white%20*%20(1%20-%20t)%20%2B%20col%20*%20t).clip(0%2C%20255).astype(np.uint8)%0A%0A%20%20%20%20f%20%3D%20entry%5B%22facts%22%5D%0A%20%20%20%20facts%20%3D%20%5Bf%22Einstein%20radius%20%7Bf%5B'theta_E'%5D%3A.2f%7D%E2%80%B3%22%2C%20f%22field%20of%20view%20%7Bf%5B'fov'%5D%3A.1f%7D%E2%80%B3%22%5D%0A%20%20%20%20if%20cls%20%3D%3D%20%22vortex%22%3A%0A%20%20%20%20%20%20%20%20facts.append(f%22vortex%20half-length%20%7Bf%5B'vortex_half_length'%5D%3A.2f%7D%E2%80%B3%22)%0A%20%20%20%20if%20cls%20%3D%3D%20%22subhalo%22%3A%0A%20%20%20%20%20%20%20%20facts.append(f%22%7Bf%5B'n_subhalos'%5D%7D%20subhalos%22)%0A%20%20%20%20if%20cls%20!%3D%20%22none%22%3A%0A%20%20%20%20%20%20%20%20facts.append(f%22peak%20perturbation%20%7Bf%5B'peak_fraction'%5D%3A.1%25%7D%20of%20peak%20brightness%22)%0A%20%20%20%20else%3A%0A%20%20%20%20%20%20%20%20facts.append(%22no%20substructure%2C%20so%20the%20residual%20is%20zero%22)%0A%0A%20%20%20%20mo.vstack(%0A%20%20%20%20%20%20%20%20%5B%0A%20%20%20%20%20%20%20%20%20%20%20%20mo.hstack(%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.vstack(%5Bmo.md(%22**observed**%22)%2C%20mo.image(src%3D_png(observed)%2C%20width%3D320)%5D%2C%20gap%3D0.3)%2C%0A%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20mo.vstack(%5Bmo.md(%22**perturbation%20from%20substructure**%22)%2C%20mo.image(src%3D_png(residual)%2C%20width%3D320)%5D%2C%20gap%3D0.3)%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.5%2C%0A%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20justify%3D%22start%22%2C%0A%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20wrap%3DTrue%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.md(%22%20%C2%B7%20%22.join(facts))%2C%0A%20%20%20%20%20%20%20%20%5D%0A%20%20%20%20)%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