import%20marimo%0A%0A__generated_with%20%3D%20%220.15.2%22%0Aapp%20%3D%20marimo.App(width%3D%22full%22%2C%20app_title%3D%221D%20Wave%20Equation%20in%20Phydrax%22)%0A%0A%0A%40app.cell%0Adef%20_()%3A%0A%20%20%20%20import%20time%0A%20%20%20%20from%20typing%20import%20Any%0A%0A%20%20%20%20import%20jax%0A%20%20%20%20import%20jax.numpy%20as%20jnp%0A%20%20%20%20import%20jax.random%20as%20jr%0A%20%20%20%20import%20marimo%20as%20mo%0A%20%20%20%20import%20matplotlib.pyplot%20as%20plt%0A%20%20%20%20import%20optax%0A%0A%20%20%20%20import%20phydrax%20as%20phx%0A%0A%20%20%20%20plt.style.use(%22seaborn-v0_8-whitegrid%22)%0A%20%20%20%20return%20jax%2C%20jnp%2C%20jr%2C%20mo%2C%20optax%2C%20phx%2C%20plt%2C%20time%0A%0A%0A%40app.cell(hide_code%3DTrue)%0Adef%20_(mo)%3A%0A%20%20%20%20mo.md(%0A%20%20%20%20%20%20%20%20r%22%22%22%0A%20%20%20%20%23%201D%20Wave%20Equation%20in%20Phydrax%0A%0A%20%20%20%20This%20notebook%20is%20a%20hands-on%20tutorial%20for%20a%20central%20Phydrax%20workflow%3A%0A%0A%20%20%20%201.%20Build%20a%20flexible%20neural%20field%20%5C(u_%5Ctheta%5C)%20from%20factorized%20latent%20branches.%0A%20%20%20%202.%20Enforce%20boundary%20and%20initial%20constraints%20**by%20construction**%20using%20Phydrax's%20%0A%20%20%20%20%20%20%20%20unique%20enforced%20constraint%20overlay%20pipeline.%0A%20%20%20%203.%20Train%20with%20a%20**single%20PDE%20residual%20objective**%20rather%20than%20juggling%20many%20soft%20%0A%20%20%20%20%20%20%20%20penalties%20using%20a%20**problem-optimal%20differentiation**%20configuration.%0A%0A%20%20%20%20---%0A%0A%20%20%20%20We%20solve%0A%0A%20%20%20%20%24%24%0A%20%20%20%20u_%7Btt%7D%20-%20c%5E2%20u_%7Bxx%7D%20%3D%200%2C%5Cqquad%20(x%2Ct)%5Cin%5B0%2C%5Cpi%5D%5Ctimes%5B0%2C2%5Cpi%5D%2C%0A%20%20%20%20%24%24%0A%0A%20%20%20%20with%20constraints%0A%0A%20%20%20%20%24%24%0A%20%20%20%20u(0%2Ct)%3Du(%5Cpi%2Ct)%3D0%2C%5Cqquad%0A%20%20%20%20u(x%2C0)%3D%5Csin%20x%2C%5Cqquad%0A%20%20%20%20u_t(x%2C0)%3D%5Csin%20x.%0A%20%20%20%20%24%24%0A%0A%20%20%20%20We%20also%20use%20the%20exact%20reference%20solution%0A%0A%20%20%20%20%24%24%0A%20%20%20%20u_%7B%5Ctext%7Bexact%7D%7D(x%2Ct)%3D%5Csin(x)%5Cbig(%5Csin(t)%2B%5Ccos(t)%5Cbig).%0A%20%20%20%20%24%24%0A%0A%20%20%20%20For%20reference%2C%20Nvidia%20also%20provides%20a%20%0A%20%20%20%20%5B**PhysicsNeMo%20implementation**%5D(https%3A%2F%2Fdocs.nvidia.com%2Fphysicsnemo%2Flatest%2Fphysicsnemo-sym%2Fuser_guide%2Ffoundational%2F1d_wave_equation.html)%20%0A%20%20%20%20of%20this%20problem.%20%0A%0A%20%20%20%20We%20demonstrate%20several%20optimizations%20developed%20at%20Phydra%20Labs%20that%20allow%20for%20%0A%20%20%20%20**orders-of-magnitude%20improvement%20across%20multiple%20axes**.%0A%0A%20%20%20%20---%0A%20%20%20%20%22%22%22%0A%20%20%20%20)%0A%20%20%20%20return%0A%0A%0A%40app.cell(hide_code%3DTrue)%0Adef%20_(mo)%3A%0A%20%20%20%20mo.md(%0A%20%20%20%20%20%20%20%20r%22%22%22%0A%20%20%20%20%23%23%20Why%20PCI%20overlays%20change%20the%20optimization%20problem%0A%0A%20%20%20%20A%20traditional%20PINN-style%20objective%20often%20looks%20like%0A%0A%20%20%20%20%24%24%0A%20%20%20%20J_%7B%5Ctext%7Bsoft%7D%7D(%5Ctheta)%3D%0A%20%20%20%20%5Clambda_r%5Cmathcal%7BL%7D_r%20%2B%0A%20%20%20%20%5Clambda_%7Bbc%7D%5Cmathcal%7BL%7D_%7Bbc%7D%20%2B%0A%20%20%20%20%5Clambda_%7Bic%2C0%7D%5Cmathcal%7BL%7D_%7Bic%2C0%7D%20%2B%20%0A%20%20%20%20%5Clambda_%7Bic%2C1%7D%5Cmathcal%7BL%7D_%7Bic%2C1%7D%2C%0A%20%20%20%20%24%24%0A%0A%20%20%20%20where%20tuning%20the%20%5C(%5Clambda%5C)'s%20becomes%20a%20balancing%20problem.%0A%0A%20%20%20%20Phydrax's%20PCI%20pipeline%20instead%20constructs%20an%20enforced%20ansatz%0A%0A%20%20%20%20%24%24%0A%20%20%20%20u_%5Ctheta%20%3D%20%5Cmathcal%7BE%7D_%7Bic%2C1%7D%5Ccirc%5Cmathcal%7BE%7D_%7Bic%2C0%7D%5Ccirc%5Cmathcal%7BE%7D_%7Bbc%7D%5Bu_%7B%5Ctheta%2C%5Ctext%7Braw%7D%7D%5D%2C%0A%20%20%20%20%24%24%0A%0A%20%20%20%20so%20BC%2FIC%20are%20pinned%20in%20the%20function%20form%20itself.%20Training%20is%20then%0A%0A%20%20%20%20%24%24%0A%20%20%20%20%5Cmin_%5Ctheta%5C%3B%5Cmathcal%7BL%7D_r(%5Ctheta)%2C%5Cqquad%0A%20%20%20%20%5Cmathcal%7BL%7D_r%3D%5Cmathbb%7BE%7D_%7B(x%2Ct)%7D%5Cleft%5B%0A%20%20%20%20%5Cleft(u_%7Btt%7D-c%5E2u_%7Bxx%7D%5Cright)%5E2%0A%20%20%20%20%5Cright%5D.%0A%20%20%20%20%24%24%0A%0A%20%20%20%20In%20short%3A%20PCI%20converts%20a%20multi-loss%20balancing%20act%20into%20a%20residual-focused%20solve%20on%20a%0A%20%20%20%20constrained%20function%20manifold.%20%0A%0A%20%20%20%20*A%20deeper%20mathematical%20dive%20into%20Phydrax's%20PCI%20pipeline%20can%20%0A%20%20%20%20be%20found%20%5Bhere%5D(https%3A%2F%2Fphydra-labs.github.io%2Fphydrax%2Fappendix%2Fphysics_constrained_interpolation%2F).*%0A%20%20%20%20%22%22%22%0A%20%20%20%20)%0A%20%20%20%20return%0A%0A%0A%40app.cell(hide_code%3DTrue)%0Adef%20_(mo)%3A%0A%20%20%20%20mo.callout(%0A%20%20%20%20%20%20%20%20mo.md(%0A%20%20%20%20%20%20%20%20%20%20%20%20r%22%22%22%0A%20%20%20%20%20%20%20%20%20%20%20%20%23%23%20Why%20double%20forward%20JVP%20is%20the%20right%20derivative%20engine%20here%0A%0A%20%20%20%20%20%20%20%20%20%20%20%20This%20example%20uses%20a%20latent-factorized%20architecture%20of%20the%20form%0A%0A%20%20%20%20%20%20%20%20%20%20%20%20%24%24%0A%20%20%20%20%20%20%20%20%20%20%20%20z(x%2Ct)%3Dz_x(x)%5Codot%20z_t(t)%5Cin%5Cmathbb%7BR%7D%5E%7BL%7D%2C%5Cqquad%0A%20%20%20%20%20%20%20%20%20%20%20%20u(x%2Ct)%3D%5Cmathcal%7BC%7D(z(x%2Ct))%2C%0A%20%20%20%20%20%20%20%20%20%20%20%20%24%24%0A%0A%20%20%20%20%20%20%20%20%20%20%20%20where%20%5C(z_x%2Cz_t%5C)%20map%20**low-dimensional%20coordinates**%20(1D%20%5C(x%5C)%2C%201D%20%5C(t%5C))%0A%20%20%20%20%20%20%20%20%20%20%20%20into%20a%20**high-dimensional%20latent**%20%5C(L%5C)%2C%20then%20%5C(%5Cmathcal%7BC%7D%5C)%20contracts%0A%20%20%20%20%20%20%20%20%20%20%20%20latent%20channels%20to%20a%20scalar%20field.%0A%0A%20%20%20%20%20%20%20%20%20%20%20%20---%0A%0A%20%20%20%20%20%20%20%20%20%20%20%20For%20wave%20residuals%20we%20need%20second%20derivatives%3A%0A%0A%20%20%20%20%20%20%20%20%20%20%20%20%24%24%0A%20%20%20%20%20%20%20%20%20%20%20%20%5Cmathcal%7BR%7D(x%2Ct)%3Du_%7Btt%7D(x%2Ct)-c%5E2u_%7Bxx%7D(x%2Ct).%0A%20%20%20%20%20%20%20%20%20%20%20%20%24%24%0A%0A%20%20%20%20%20%20%20%20%20%20%20%20In%20Phydrax%20with%20%60backend%3D%22ad%22%2C%20ad_engine%3D%22jvp%22%60%2C%20these%20are%20computed%20as%0A%20%20%20%20%20%20%20%20%20%20%20%20forward-over-forward%20directional%20derivatives%3A%0A%0A%20%20%20%20%20%20%20%20%20%20%20%20%24%24%0A%20%20%20%20%20%20%20%20%20%20%20%20u_%7Bxx%7D%3DD_%7Be_x%7D%5C!%5Cleft(D_%7Be_x%7Du%5Cright)%2C%5Cqquad%0A%20%20%20%20%20%20%20%20%20%20%20%20u_%7Btt%7D%3DD_%7Be_t%7D%5C!%5Cleft(D_%7Be_t%7Du%5Cright).%0A%20%20%20%20%20%20%20%20%20%20%20%20%24%24%0A%0A%20%20%20%20%20%20%20%20%20%20%20%20---%0A%0A%20%20%20%20%20%20%20%20%20%20%20%20Why%20this%20is%20ideal%20for%20this%20model%20class%3A%0A%0A%20%20%20%20%20%20%20%20%20%20%20%201.%20**Input%20dimension%20is%20tiny.**%20%20%0A%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20Forward-mode%20complexity%20is%20tied%20to%20tangent%20dimension%20(here%201%20per%20axis)%2C%0A%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20so%20JVP%20is%20naturally%20efficient%20for%20coordinate%20derivatives.%0A%0A%20%20%20%20%20%20%20%20%20%20%20%202.%20**Derivative%20flow%20matches%20factorization.**%20%20%0A%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20Differentiating%20w.r.t.%20%5C(x%5C)%20only%20perturbs%20the%20%5C(x%5C)-branch%20tangent%20while%0A%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20preserving%20multiplicative%20interaction%20with%20%5C(t%5C)-branch%20latents%2C%20and%20vice%20versa.%0A%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20This%20keeps%20derivative%20propagation%20aligned%20with%20the%20model%E2%80%99s%20separable%20algebra.%0A%0A%20%20%20%20%20%20%20%20%20%20%20%203.%20**No%20full%20Hessian%20materialization.**%20%20%0A%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20JVP-of-JVP%20gives%20the%20exact%20directional%20second%20derivative%20needed%20by%20PDE%20operators%0A%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20without%20constructing%20dense%20Hessian%20tensors.%0A%0A%20%20%20%20%20%20%20%20%20%20%20%204.%20**Great%20fit%20with%20coord-separable%20sampling.**%20%20%0A%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20In%20this%20notebook%20we%20sample%20%5C(x%5C)-axes%20and%20%5C(t%5C)-points%20in%20a%20structured%20way%2C%0A%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20so%20derivative%20computation%20stays%20close%20to%20the%20branchwise%20geometry%20instead%20of%0A%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20flattening%20everything%20into%20a%20monolithic%20dense-input%20path.%0A%0A%20%20%20%20%20%20%20%20%20%20%20%20Conceptually%3A%20this%20architecture%20is%20%22low-dimensional%20coordinate%20geometry%20%0A%20%20%20%20%20%20%20%20%20%20%20%20%2B%20large%20latent%20algebra%20%2B%20contraction.%22%0A%20%20%20%20%20%20%20%20%20%20%20%20Double%20forward%20JVP%20is%20the%20AD%20mechanism%20that%20most%20directly%20respects%20that%20geometry.%0A%20%20%20%20%20%20%20%20%20%20%20%20%22%22%22%0A%20%20%20%20%20%20%20%20)%2C%0A%20%20%20%20%20%20%20%20kind%3D%22success%22%2C%0A%20%20%20%20)%0A%20%20%20%20return%0A%0A%0A%40app.cell%0Adef%20_()%3A%0A%20%20%20%20%23%20-------------------------------------------------------------------------%0A%20%20%20%20%23%20Configurations%0A%20%20%20%20%23%20-------------------------------------------------------------------------%0A%20%20%20%20latent_size%20%3D%2032%0A%20%20%20%20width_size%20%3D%2020%0A%20%20%20%20depth%20%3D%206%0A%20%20%20%20num_iter%20%3D%20100%0A%20%20%20%20learning_rate%20%3D%207e-3%0A%20%20%20%20num_x_interior%20%3D%20400%0A%20%20%20%20num_t_interior%20%3D%20800%0A%20%20%20%20nx_plot%20%3D%20200%0A%20%20%20%20nt_plot%20%3D%20400%0A%20%20%20%20seed%20%3D%200%0A%20%20%20%20c%20%3D%201.0%0A%20%20%20%20x_min%20%3D%200.0%0A%20%20%20%20x_max%20%3D%203.141592653589793%0A%20%20%20%20t_min%20%3D%200.0%0A%20%20%20%20t_max%20%3D%206.283185307179586%0A%20%20%20%20return%20(%0A%20%20%20%20%20%20%20%20c%2C%0A%20%20%20%20%20%20%20%20depth%2C%0A%20%20%20%20%20%20%20%20latent_size%2C%0A%20%20%20%20%20%20%20%20learning_rate%2C%0A%20%20%20%20%20%20%20%20nt_plot%2C%0A%20%20%20%20%20%20%20%20num_iter%2C%0A%20%20%20%20%20%20%20%20num_t_interior%2C%0A%20%20%20%20%20%20%20%20num_x_interior%2C%0A%20%20%20%20%20%20%20%20nx_plot%2C%0A%20%20%20%20%20%20%20%20seed%2C%0A%20%20%20%20%20%20%20%20t_max%2C%0A%20%20%20%20%20%20%20%20t_min%2C%0A%20%20%20%20%20%20%20%20width_size%2C%0A%20%20%20%20%20%20%20%20x_max%2C%0A%20%20%20%20%20%20%20%20x_min%2C%0A%20%20%20%20)%0A%0A%0A%40app.cell%0Adef%20_(jax%2C%20jnp%2C%20jr%2C%20phx)%3A%0A%20%20%20%20%23%20-------------------------------------------------------------------------%0A%20%20%20%20%23%20Utility%20functions%0A%20%20%20%20%23%20-------------------------------------------------------------------------%0A%20%20%20%20def%20u_exact(x%3A%20jax.Array%2C%20t%3A%20jax.Array)%20-%3E%20jax.Array%3A%0A%20%20%20%20%20%20%20%20return%20jnp.sin(x)%20*%20(jnp.sin(t)%20%2B%20jnp.cos(t))%0A%0A%20%20%20%20def%20evaluate_on_grid(%0A%20%20%20%20%20%20%20%20u%3A%20phx.domain.DomainFunction%2C%0A%20%20%20%20%20%20%20%20*%2C%0A%20%20%20%20%20%20%20%20x_min%3A%20float%2C%0A%20%20%20%20%20%20%20%20x_max%3A%20float%2C%0A%20%20%20%20%20%20%20%20t_min%3A%20float%2C%0A%20%20%20%20%20%20%20%20t_max%3A%20float%2C%0A%20%20%20%20%20%20%20%20nx%3A%20int%2C%0A%20%20%20%20%20%20%20%20nt%3A%20int%2C%0A%20%20%20%20)%20-%3E%20tuple%5Bjax.Array%2C%20jax.Array%2C%20jax.Array%2C%20jax.Array%2C%20jax.Array%5D%3A%0A%20%20%20%20%20%20%20%20x%20%3D%20jnp.linspace(x_min%2C%20x_max%2C%20int(nx))%0A%20%20%20%20%20%20%20%20t%20%3D%20jnp.linspace(t_min%2C%20t_max%2C%20int(nt))%0A%20%20%20%20%20%20%20%20xx%2C%20tt%20%3D%20jnp.meshgrid(x%2C%20t%2C%20indexing%3D%22ij%22)%0A%0A%20%20%20%20%20%20%20%20x_flat%20%3D%20xx.reshape((-1%2C))%0A%20%20%20%20%20%20%20%20t_flat%20%3D%20tt.reshape((-1%2C))%0A%20%20%20%20%20%20%20%20point_eval%20%3D%20jax.jit(%0A%20%20%20%20%20%20%20%20%20%20%20%20jax.vmap(lambda%20x_i%2C%20t_i%3A%20u.func(jnp.asarray(%5Bx_i%5D%2C%20dtype%3Dfloat)%2C%20t_i))%0A%20%20%20%20%20%20%20%20)%0A%20%20%20%20%20%20%20%20u_pred%20%3D%20point_eval(x_flat%2C%20t_flat).reshape((int(nx)%2C%20int(nt)))%0A%20%20%20%20%20%20%20%20u_true%20%3D%20u_exact(xx%2C%20tt)%0A%20%20%20%20%20%20%20%20diff%20%3D%20u_pred%20-%20u_true%0A%20%20%20%20%20%20%20%20return%20x%2C%20t%2C%20u_pred%2C%20u_true%2C%20diff%0A%0A%20%20%20%20def%20constraint_errors(%0A%20%20%20%20%20%20%20%20u%3A%20phx.domain.DomainFunction%2C%0A%20%20%20%20%20%20%20%20*%2C%0A%20%20%20%20%20%20%20%20x_min%3A%20float%2C%0A%20%20%20%20%20%20%20%20x_max%3A%20float%2C%0A%20%20%20%20%20%20%20%20t_min%3A%20float%2C%0A%20%20%20%20%20%20%20%20t_max%3A%20float%2C%0A%20%20%20%20)%20-%3E%20dict%5Bstr%2C%20float%5D%3A%0A%20%20%20%20%20%20%20%20t_line%20%3D%20jnp.linspace(t_min%2C%20t_max%2C%20256)%0A%20%20%20%20%20%20%20%20x_line%20%3D%20jnp.linspace(x_min%2C%20x_max%2C%20256)%0A%0A%20%20%20%20%20%20%20%20left%20%3D%20jax.vmap(lambda%20t_i%3A%20u.func(jnp.asarray(%5Bx_min%5D%2C%20dtype%3Dfloat)%2C%20t_i))(%0A%20%20%20%20%20%20%20%20%20%20%20%20t_line%0A%20%20%20%20%20%20%20%20)%0A%20%20%20%20%20%20%20%20right%20%3D%20jax.vmap(lambda%20t_i%3A%20u.func(jnp.asarray(%5Bx_max%5D%2C%20dtype%3Dfloat)%2C%20t_i))(%0A%20%20%20%20%20%20%20%20%20%20%20%20t_line%0A%20%20%20%20%20%20%20%20)%0A%20%20%20%20%20%20%20%20u0%20%3D%20jax.vmap(lambda%20x_i%3A%20u.func(jnp.asarray(%5Bx_i%5D%2C%20dtype%3Dfloat)%2C%20t_min))(x_line)%0A%20%20%20%20%20%20%20%20ut%20%3D%20phx.operators.dt(u%2C%20var%3D%22t%22)%0A%20%20%20%20%20%20%20%20ut0%20%3D%20jax.vmap(lambda%20x_i%3A%20ut.func(jnp.asarray(%5Bx_i%5D%2C%20dtype%3Dfloat)%2C%20t_min))(%0A%20%20%20%20%20%20%20%20%20%20%20%20x_line%0A%20%20%20%20%20%20%20%20)%0A%0A%20%20%20%20%20%20%20%20return%20%7B%0A%20%20%20%20%20%20%20%20%20%20%20%20%22max_bc_error%22%3A%20float(%0A%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20jnp.maximum(jnp.max(jnp.abs(left))%2C%20jnp.max(jnp.abs(right)))%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%20%22max_u0_error%22%3A%20float(jnp.max(jnp.abs(u0%20-%20jnp.sin(x_line))))%2C%0A%20%20%20%20%20%20%20%20%20%20%20%20%22max_ut0_error%22%3A%20float(jnp.max(jnp.abs(ut0%20-%20jnp.sin(x_line))))%2C%0A%20%20%20%20%20%20%20%20%7D%0A%0A%20%20%20%20%23%20-------------------------------------------------------------------------%0A%20%20%20%20%23%20Main%20solver%20factory%0A%20%20%20%20%23%20-------------------------------------------------------------------------%0A%20%20%20%20def%20make_solver(%0A%20%20%20%20%20%20%20%20*%2C%0A%20%20%20%20%20%20%20%20c%3A%20float%2C%0A%20%20%20%20%20%20%20%20x_min%3A%20float%2C%0A%20%20%20%20%20%20%20%20x_max%3A%20float%2C%0A%20%20%20%20%20%20%20%20t_min%3A%20float%2C%0A%20%20%20%20%20%20%20%20t_max%3A%20float%2C%0A%20%20%20%20%20%20%20%20latent_size%3A%20int%2C%0A%20%20%20%20%20%20%20%20width_size%3A%20int%2C%0A%20%20%20%20%20%20%20%20depth%3A%20int%2C%0A%20%20%20%20%20%20%20%20num_x_interior%3A%20int%2C%0A%20%20%20%20%20%20%20%20num_t_interior%3A%20int%2C%0A%20%20%20%20%20%20%20%20key%3A%20jax.Array%2C%0A%20%20%20%20)%20-%3E%20phx.solver.FunctionalSolver%3A%0A%20%20%20%20%20%20%20%20geom%20%3D%20phx.domain.Interval1d(float(x_min)%2C%20float(x_max))%0A%20%20%20%20%20%20%20%20time_dom%20%3D%20phx.domain.TimeInterval(float(t_min)%2C%20float(t_max))%0A%20%20%20%20%20%20%20%20domain%20%3D%20geom%20%40%20time_dom%0A%0A%20%20%20%20%20%20%20%20key_x%2C%20key_t%2C%20key_bw%20%3D%20jr.split(key%2C%203)%0A%20%20%20%20%20%20%20%20model%20%3D%20phx.nn.LatentContractionModel(%0A%20%20%20%20%20%20%20%20%20%20%20%20latent_size%3Dint(latent_size)%2C%0A%20%20%20%20%20%20%20%20%20%20%20%20out_size%3D%22scalar%22%2C%0A%20%20%20%20%20%20%20%20%20%20%20%20execution_policy%3Dphx.nn.LatentExecutionPolicy(%0A%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20topology%3D%22best_effort_flat%22%2C%0A%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20layout%3D%22coord_separable%22%2C%0A%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20fallback%3D%22error%22%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%20x%3Dphx.nn.MLP(%0A%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20in_size%3D%22scalar%22%2C%0A%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20out_size%3Dint(latent_size)%2C%0A%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20width_size%3Dint(width_size)%2C%0A%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20depth%3Dint(depth)%2C%0A%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20activation%3Dphx.nn.Stan(int(width_size))%2C%0A%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20key%3Dkey_x%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%20t%3Dphx.nn.MLP(%0A%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20in_size%3D%22scalar%22%2C%0A%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20out_size%3Dint(latent_size)%2C%0A%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20width_size%3Dint(width_size)%2C%0A%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20depth%3Dint(depth)%2C%0A%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20activation%3Dphx.nn.Stan(int(width_size))%2C%0A%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20key%3Dkey_t%2C%0A%20%20%20%20%20%20%20%20%20%20%20%20)%2C%0A%20%20%20%20%20%20%20%20)%0A%20%20%20%20%20%20%20%20u_raw%20%3D%20domain.Model(%22x%22%2C%20%22t%22)(model)%0A%0A%20%20%20%20%20%20%20%20boundary%20%3D%20domain.component(%7B%22x%22%3A%20phx.domain.Boundary()%7D)%0A%20%20%20%20%20%20%20%20initial%20%3D%20domain.component(%7B%22t%22%3A%20phx.domain.FixedStart()%7D)%0A%0A%20%20%20%20%20%20%20%20%40domain.Function(%22x%22)%0A%20%20%20%20%20%20%20%20def%20u0_target(x)%3A%0A%20%20%20%20%20%20%20%20%20%20%20%20return%20jnp.sin(x%5B0%5D)%0A%0A%20%20%20%20%20%20%20%20%40domain.Function(%22x%22)%0A%20%20%20%20%20%20%20%20def%20ut0_target(x)%3A%0A%20%20%20%20%20%20%20%20%20%20%20%20return%20jnp.sin(x%5B0%5D)%0A%0A%20%20%20%20%20%20%20%20%23%20PCI%20overlays%3A%20exact%20BC%2FIC%20enforcement%20by%20construction.%0A%20%20%20%20%20%20%20%20overlays%20%3D%20%5B%0A%20%20%20%20%20%20%20%20%20%20%20%20phx.solver.SingleFieldEnforcedConstraint(%0A%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20%22u%22%2C%0A%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20boundary%2C%0A%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20lambda%20f%3A%20phx.constraints.enforce_dirichlet(%0A%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20f%2C%20boundary%2C%20var%3D%22x%22%2C%20target%3D0.0%0A%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)%2C%0A%20%20%20%20%20%20%20%20%20%20%20%20phx.solver.SingleFieldEnforcedConstraint(%0A%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20%22u%22%2C%0A%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20initial%2C%0A%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20lambda%20f%3A%20f%2C%0A%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20time_derivative_order%3D0%2C%0A%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20initial_target%3Du0_target%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%20phx.solver.SingleFieldEnforcedConstraint(%0A%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20%22u%22%2C%0A%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20initial%2C%0A%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20lambda%20f%3A%20f%2C%0A%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20time_derivative_order%3D1%2C%0A%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20initial_target%3Dut0_target%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%0A%0A%20%20%20%20%20%20%20%20%23%20Residual%20objective%20only%3A%20derivatives%20are%20computed%20with%20JVP%20engine.%0A%20%20%20%20%20%20%20%20structure_xt%20%3D%20phx.domain.ProductStructure(((%22x%22%2C)%2C%20(%22t%22%2C)))%0A%20%20%20%20%20%20%20%20residual%20%3D%20phx.constraints.ContinuousPointwiseInteriorConstraint(%0A%20%20%20%20%20%20%20%20%20%20%20%20%22u%22%2C%0A%20%20%20%20%20%20%20%20%20%20%20%20domain%2C%0A%20%20%20%20%20%20%20%20%20%20%20%20operator%3Dlambda%20f%3A%20(%0A%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20phx.operators.dt_n(f%2C%20var%3D%22t%22%2C%20order%3D2%2C%20ad_engine%3D%22jvp%22)%0A%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20-%20(float(c)%20**%202)%0A%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20*%20phx.operators.laplacian(f%2C%20var%3D%22x%22%2C%20ad_engine%3D%22jvp%22)%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%20num_points%3D%7B%22x%22%3A%20int(num_x_interior)%2C%20%22t%22%3A%20int(num_t_interior)%7D%2C%0A%20%20%20%20%20%20%20%20%20%20%20%20structure%3Dstructure_xt%2C%0A%20%20%20%20%20%20%20%20%20%20%20%20reduction%3D%22mean%22%2C%0A%20%20%20%20%20%20%20%20%20%20%20%20label%3D%22wave_residual%22%2C%0A%20%20%20%20%20%20%20%20)%0A%0A%20%20%20%20%20%20%20%20return%20phx.solver.FunctionalSolver(%0A%20%20%20%20%20%20%20%20%20%20%20%20functions%3D%7B%22u%22%3A%20u_raw%7D%2C%0A%20%20%20%20%20%20%20%20%20%20%20%20constraints%3D%5Bresidual%5D%2C%0A%20%20%20%20%20%20%20%20%20%20%20%20constraint_terms%3Doverlays%2C%0A%20%20%20%20%20%20%20%20%20%20%20%20boundary_weight_key%3Dkey_bw%2C%0A%20%20%20%20%20%20%20%20)%0A%20%20%20%20return%20constraint_errors%2C%20evaluate_on_grid%2C%20make_solver%0A%0A%0A%40app.cell%0Adef%20_(%0A%20%20%20%20c%2C%0A%20%20%20%20depth%2C%0A%20%20%20%20jr%2C%0A%20%20%20%20latent_size%2C%0A%20%20%20%20learning_rate%2C%0A%20%20%20%20make_solver%2C%0A%20%20%20%20mo%2C%0A%20%20%20%20num_iter%2C%0A%20%20%20%20num_t_interior%2C%0A%20%20%20%20num_x_interior%2C%0A%20%20%20%20optax%2C%0A%20%20%20%20seed%2C%0A%20%20%20%20t_max%2C%0A%20%20%20%20t_min%2C%0A%20%20%20%20time%2C%0A%20%20%20%20width_size%2C%0A%20%20%20%20x_max%2C%0A%20%20%20%20x_min%2C%0A)%3A%0A%20%20%20%20%23%20-------------------------------------------------------------------------%0A%20%20%20%20%23%20Main%20execution%20path%0A%20%20%20%20%23%20-------------------------------------------------------------------------%0A%20%20%20%20solver%20%3D%20make_solver(%0A%20%20%20%20%20%20%20%20c%3Dc%2C%0A%20%20%20%20%20%20%20%20x_min%3Dfloat(x_min)%2C%0A%20%20%20%20%20%20%20%20x_max%3Dfloat(x_max)%2C%0A%20%20%20%20%20%20%20%20t_min%3Dfloat(t_min)%2C%0A%20%20%20%20%20%20%20%20t_max%3Dfloat(t_max)%2C%0A%20%20%20%20%20%20%20%20latent_size%3Dint(latent_size)%2C%0A%20%20%20%20%20%20%20%20width_size%3Dint(width_size)%2C%0A%20%20%20%20%20%20%20%20depth%3Dint(depth)%2C%0A%20%20%20%20%20%20%20%20num_x_interior%3Dint(num_x_interior)%2C%0A%20%20%20%20%20%20%20%20num_t_interior%3Dint(num_t_interior)%2C%0A%20%20%20%20%20%20%20%20key%3Djr.key(int(seed))%2C%0A%20%20%20%20)%0A%20%20%20%20t0%20%3D%20time.perf_counter()%0A%20%20%20%20init_loss%20%3D%20float(solver.loss(key%3Djr.key(int(seed)%20%2B%201)))%0A%20%20%20%20trained_solver%20%3D%20solver.solve(%0A%20%20%20%20%20%20%20%20num_iter%3Dint(num_iter)%2C%0A%20%20%20%20%20%20%20%20optim%3Doptax.rprop(float(learning_rate))%2C%0A%20%20%20%20)%0A%20%20%20%20final_loss%20%3D%20float(trained_solver.loss(key%3Djr.key(int(seed)%20%2B%202)))%0A%20%20%20%20elapsed%20%3D%20float(time.perf_counter()%20-%20t0)%0A%20%20%20%20train_stats%20%3D%20%7B%0A%20%20%20%20%20%20%20%20%22init_loss%22%3A%20init_loss%2C%0A%20%20%20%20%20%20%20%20%22final_loss%22%3A%20final_loss%2C%0A%20%20%20%20%20%20%20%20%22elapsed_s%22%3A%20elapsed%2C%0A%20%20%20%20%20%20%20%20%22s_per_iter%22%3A%20elapsed%20%2F%20max(int(num_iter)%2C%201)%2C%0A%20%20%20%20%7D%0A%20%20%20%20train_status%20%3D%20mo.callout(%0A%20%20%20%20%20%20%20%20mo.md(%0A%20%20%20%20%20%20%20%20%20%20%20%20f%22%22%22%0A%20%20%20%20%20%20%20%20%20%20%20%20%E2%9C%85%20Training%20complete%0A%0A%20%20%20%20%20%20%20%20%20%20%20%20-%20initial%20loss%3A%20%60%7Binit_loss%3A.6e%7D%60%0A%20%20%20%20%20%20%20%20%20%20%20%20-%20final%20loss%3A%20%60%7Bfinal_loss%3A.6e%7D%60%0A%20%20%20%20%20%20%20%20%20%20%20%20-%20elapsed%3A%20%60%7Belapsed%3A.2f%7Ds%60%20(%60%7Btrain_stats%5B%22s_per_iter%22%5D%3A.4f%7Ds%2Fiter%60)%0A%20%20%20%20%20%20%20%20%20%20%20%20%22%22%22%0A%20%20%20%20%20%20%20%20)%2C%0A%20%20%20%20%20%20%20%20kind%3D%22success%22%2C%0A%20%20%20%20)%0A%0A%20%20%20%20train_status%0A%20%20%20%20return%20train_stats%2C%20trained_solver%0A%0A%0A%40app.cell%0Adef%20_(%0A%20%20%20%20constraint_errors%2C%0A%20%20%20%20evaluate_on_grid%2C%0A%20%20%20%20jnp%2C%0A%20%20%20%20nt_plot%2C%0A%20%20%20%20nx_plot%2C%0A%20%20%20%20t_max%2C%0A%20%20%20%20t_min%2C%0A%20%20%20%20trained_solver%2C%0A%20%20%20%20x_max%2C%0A%20%20%20%20x_min%2C%0A)%3A%0A%20%20%20%20u%20%3D%20trained_solver.ansatz_functions()%5B%22u%22%5D%0A%20%20%20%20x%2C%20t%2C%20u_pred%2C%20u_true%2C%20diff%20%3D%20evaluate_on_grid(%0A%20%20%20%20%20%20%20%20u%2C%0A%20%20%20%20%20%20%20%20x_min%3Dfloat(x_min)%2C%0A%20%20%20%20%20%20%20%20x_max%3Dfloat(x_max)%2C%0A%20%20%20%20%20%20%20%20t_min%3Dfloat(t_min)%2C%0A%20%20%20%20%20%20%20%20t_max%3Dfloat(t_max)%2C%0A%20%20%20%20%20%20%20%20nx%3Dint(nx_plot)%2C%0A%20%20%20%20%20%20%20%20nt%3Dint(nt_plot)%2C%0A%20%20%20%20)%0A%20%20%20%20errors%20%3D%20constraint_errors(%0A%20%20%20%20%20%20%20%20u%2C%0A%20%20%20%20%20%20%20%20x_min%3Dfloat(x_min)%2C%0A%20%20%20%20%20%20%20%20x_max%3Dfloat(x_max)%2C%0A%20%20%20%20%20%20%20%20t_min%3Dfloat(t_min)%2C%0A%20%20%20%20%20%20%20%20t_max%3Dfloat(t_max)%2C%0A%20%20%20%20)%0A%20%20%20%20diag_stats%20%3D%20%7B%0A%20%20%20%20%20%20%20%20%22l2_error%22%3A%20float(jnp.sqrt(jnp.mean((u_pred%20-%20u_true)%20**%202)))%2C%0A%20%20%20%20%20%20%20%20%22linf_error%22%3A%20float(jnp.max(jnp.abs(diff)))%2C%0A%20%20%20%20%20%20%20%20**errors%2C%0A%20%20%20%20%7D%0A%20%20%20%20plot_data%20%3D%20(x%2C%20t%2C%20u_pred%2C%20u_true%2C%20diff)%0A%20%20%20%20return%20diag_stats%2C%20plot_data%0A%0A%0A%40app.cell(hide_code%3DTrue)%0Adef%20_(diag_stats%2C%20mo%2C%20train_stats)%3A%0A%20%20%20%20diagnostics_panel%20%3D%20mo.callout(%22Diagnostics%20unavailable.%22%2C%20kind%3D%22warn%22)%0A%20%20%20%20if%20diag_stats%20is%20not%20None%3A%0A%20%20%20%20%20%20%20%20msg%20%3D%20f%22%22%22%0A%20%20%20%20%20%20%20%20%23%23%20Diagnostics%0A%0A%20%20%20%20%20%20%20%20-%20L2%20error%3A%20%60%7Bdiag_stats%5B%22l2_error%22%5D%3A.3e%7D%60%0A%20%20%20%20%20%20%20%20-%20Linf%20error%3A%20%60%7Bdiag_stats%5B%22linf_error%22%5D%3A.3e%7D%60%0A%20%20%20%20%20%20%20%20-%20max%20BC%20error%3A%20%60%7Bdiag_stats%5B%22max_bc_error%22%5D%3A.3e%7D%60%0A%20%20%20%20%20%20%20%20-%20max%20IC%20value%20error%3A%20%60%7Bdiag_stats%5B%22max_u0_error%22%5D%3A.3e%7D%60%0A%20%20%20%20%20%20%20%20-%20max%20IC%20slope%20error%3A%20%60%7Bdiag_stats%5B%22max_ut0_error%22%5D%3A.3e%7D%60%0A%20%20%20%20%20%20%20%20%22%22%22%0A%20%20%20%20%20%20%20%20if%20train_stats%20is%20not%20None%3A%0A%20%20%20%20%20%20%20%20%20%20%20%20msg%20%2B%3D%20(%0A%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20f%22%5Cn-%20loss%20(init%20%E2%86%92%20final)%3A%20%60%7Btrain_stats%5B'init_loss'%5D%3A.3e%7D%60%20%E2%86%92%20%22%0A%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20f%22%60%7Btrain_stats%5B'final_loss'%5D%3A.3e%7D%60%22%0A%20%20%20%20%20%20%20%20%20%20%20%20)%0A%20%20%20%20%20%20%20%20diagnostics_panel%20%3D%20mo.md(msg)%0A%20%20%20%20diagnostics_panel%0A%20%20%20%20return%0A%0A%0A%40app.cell(hide_code%3DTrue)%0Adef%20_(jnp%2C%20mo%2C%20plot_data%2C%20plt)%3A%0A%20%20%20%20plot_panel%20%3D%20mo.callout(%22No%20plot%20data%20available.%22%2C%20kind%3D%22warn%22)%0A%20%20%20%20if%20plot_data%20is%20not%20None%3A%0A%20%20%20%20%20%20%20%20x_plot%2C%20t_plot%2C%20u_pred_plot%2C%20u_true_plot%2C%20diff_plot%20%3D%20plot_data%0A%0A%20%20%20%20%20%20%20%20extent%20%3D%20(%0A%20%20%20%20%20%20%20%20%20%20%20%20float(x_plot%5B0%5D)%2C%0A%20%20%20%20%20%20%20%20%20%20%20%20float(x_plot%5B-1%5D)%2C%0A%20%20%20%20%20%20%20%20%20%20%20%20float(t_plot%5B0%5D)%2C%0A%20%20%20%20%20%20%20%20%20%20%20%20float(t_plot%5B-1%5D)%2C%0A%20%20%20%20%20%20%20%20)%0A%20%20%20%20%20%20%20%20field_min%20%3D%20float(jnp.minimum(jnp.min(u_pred_plot)%2C%20jnp.min(u_true_plot)))%0A%20%20%20%20%20%20%20%20field_max%20%3D%20float(jnp.maximum(jnp.max(u_pred_plot)%2C%20jnp.max(u_true_plot)))%0A%20%20%20%20%20%20%20%20diff_lim%20%3D%20float(jnp.max(jnp.abs(diff_plot)))%0A%20%20%20%20%20%20%20%20diff_lim%20%3D%20diff_lim%20if%20diff_lim%20%3E%200.0%20else%201e-12%0A%0A%20%20%20%20%20%20%20%20fig%2C%20axes%20%3D%20plt.subplots(1%2C%203%2C%20figsize%3D(10.4%2C%206.6)%2C%20constrained_layout%3DTrue)%0A%20%20%20%20%20%20%20%20ax0%2C%20ax1%2C%20ax2%20%3D%20axes%0A%0A%20%20%20%20%20%20%20%20im0%20%3D%20ax0.imshow(%0A%20%20%20%20%20%20%20%20%20%20%20%20jnp.asarray(u_pred_plot.T)%2C%0A%20%20%20%20%20%20%20%20%20%20%20%20origin%3D%22lower%22%2C%0A%20%20%20%20%20%20%20%20%20%20%20%20aspect%3D%22equal%22%2C%0A%20%20%20%20%20%20%20%20%20%20%20%20extent%3Dextent%2C%0A%20%20%20%20%20%20%20%20%20%20%20%20cmap%3D%22jet%22%2C%0A%20%20%20%20%20%20%20%20%20%20%20%20vmin%3Dfield_min%2C%0A%20%20%20%20%20%20%20%20%20%20%20%20vmax%3Dfield_max%2C%0A%20%20%20%20%20%20%20%20)%0A%20%20%20%20%20%20%20%20ax0.set_title(%22Phydrax%20Solution%22)%0A%20%20%20%20%20%20%20%20ax0.set_xlabel(%22x%22)%0A%20%20%20%20%20%20%20%20ax0.set_ylabel(%22t%22)%0A%20%20%20%20%20%20%20%20ax0.grid(False)%0A%20%20%20%20%20%20%20%20fig.colorbar(im0%2C%20ax%3Dax0%2C%20fraction%3D0.072%2C%20pad%3D0.042)%0A%0A%20%20%20%20%20%20%20%20im1%20%3D%20ax1.imshow(%0A%20%20%20%20%20%20%20%20%20%20%20%20jnp.asarray(u_true_plot.T)%2C%0A%20%20%20%20%20%20%20%20%20%20%20%20origin%3D%22lower%22%2C%0A%20%20%20%20%20%20%20%20%20%20%20%20aspect%3D%22equal%22%2C%0A%20%20%20%20%20%20%20%20%20%20%20%20extent%3Dextent%2C%0A%20%20%20%20%20%20%20%20%20%20%20%20cmap%3D%22jet%22%2C%0A%20%20%20%20%20%20%20%20%20%20%20%20vmin%3Dfield_min%2C%0A%20%20%20%20%20%20%20%20%20%20%20%20vmax%3Dfield_max%2C%0A%20%20%20%20%20%20%20%20)%0A%20%20%20%20%20%20%20%20ax1.set_title(%22Ground%20Truth%22)%0A%20%20%20%20%20%20%20%20ax1.set_xlabel(%22x%22)%0A%20%20%20%20%20%20%20%20ax1.set_ylabel(%22t%22)%0A%20%20%20%20%20%20%20%20ax1.grid(False)%0A%20%20%20%20%20%20%20%20fig.colorbar(im1%2C%20ax%3Dax1%2C%20fraction%3D0.072%2C%20pad%3D0.042)%0A%0A%20%20%20%20%20%20%20%20im2%20%3D%20ax2.imshow(%0A%20%20%20%20%20%20%20%20%20%20%20%20jnp.asarray(diff_plot.T)%2C%0A%20%20%20%20%20%20%20%20%20%20%20%20origin%3D%22lower%22%2C%0A%20%20%20%20%20%20%20%20%20%20%20%20aspect%3D%22equal%22%2C%0A%20%20%20%20%20%20%20%20%20%20%20%20extent%3Dextent%2C%0A%20%20%20%20%20%20%20%20%20%20%20%20cmap%3D%22jet%22%2C%0A%20%20%20%20%20%20%20%20%20%20%20%20vmin%3D-diff_lim%2C%0A%20%20%20%20%20%20%20%20%20%20%20%20vmax%3Ddiff_lim%2C%0A%20%20%20%20%20%20%20%20)%0A%20%20%20%20%20%20%20%20ax2.set_title(%22Error%20(pred%20-%20exact)%22)%0A%20%20%20%20%20%20%20%20ax2.set_xlabel(%22x%22)%0A%20%20%20%20%20%20%20%20ax2.set_ylabel(%22t%22)%0A%20%20%20%20%20%20%20%20ax2.grid(False)%0A%20%20%20%20%20%20%20%20fig.colorbar(im2%2C%20ax%3Dax2%2C%20fraction%3D0.072%2C%20pad%3D0.042)%0A%0A%20%20%20%20%20%20%20%20plot_panel%20%3D%20mo.hstack(%5Bmo.md(%22%22)%2C%20fig%2C%20mo.md(%22%22)%5D%2C%20widths%3D%5B1%2C%208%2C%201%5D)%0A%20%20%20%20plot_panel%0A%20%20%20%20return%0A%0A%0A%40app.cell(hide_code%3DTrue)%0Adef%20_(diag_stats%2C%20mo%2C%20num_iter%2C%20num_t_interior%2C%20num_x_interior%2C%20train_stats)%3A%0A%20%20%20%20physicsnemo_linf%20%3D%200.025%0A%20%20%20%20phydrax_reference_linf%20%3D%200.00186%0A%20%20%20%20physicsnemo_steps%20%3D%2010_000%0A%20%20%20%20physicsnemo_collocation%20%3D%203_200%0A%0A%20%20%20%20our_steps%20%3D%20int(num_iter)%0A%20%20%20%20our_collocation%20%3D%20int(num_x_interior)%20*%20int(num_t_interior)%0A%20%20%20%20step_reduction%20%3D%20100.0%20*%20((physicsnemo_steps%20-%20our_steps)%20%2F%20physicsnemo_steps)%0A%20%20%20%20collocation_ratio%20%3D%20our_collocation%20%2F%20physicsnemo_collocation%0A%20%20%20%20reference_linf_reduction%20%3D%20100.0%20*%20(%0A%20%20%20%20%20%20%20%20(physicsnemo_linf%20-%20phydrax_reference_linf)%20%2F%20physicsnemo_linf%0A%20%20%20%20)%0A%0A%20%20%20%20current_linf_line%20%3D%20%22%22%0A%20%20%20%20if%20diag_stats%20is%20not%20None%3A%0A%20%20%20%20%20%20%20%20current_linf%20%3D%20float(diag_stats%5B%22linf_error%22%5D)%0A%20%20%20%20%20%20%20%20current_reduction%20%3D%20100.0%20*%20((physicsnemo_linf%20-%20current_linf)%20%2F%20physicsnemo_linf)%0A%20%20%20%20%20%20%20%20current_linf_line%20%3D%20(%0A%20%20%20%20%20%20%20%20%20%20%20%20f%22%5Cn-%20Current%20notebook%20run%20Linf%3A%20%60%7Bcurrent_linf%3A.3e%7D%60%20%22%0A%20%20%20%20%20%20%20%20%20%20%20%20f%22(%60%7Bcurrent_reduction%3A.1f%7D%25%60%20lower%20than%20PhysicsNeMo%20baseline)%22%0A%20%20%20%20%20%20%20%20)%0A%0A%20%20%20%20speed_line%20%3D%20%22%22%0A%20%20%20%20if%20train_stats%20is%20not%20None%3A%0A%20%20%20%20%20%20%20%20speed_line%20%3D%20f%22%5Cn-%20This%20run%20time%20per%20iteration%3A%20%60%7Btrain_stats%5B's_per_iter'%5D%3A.3f%7Ds%2Fiter%60%22%0A%0A%20%20%20%20comparison_panel%20%3D%20mo.md(%0A%20%20%20%20%20%20%20%20f%22%22%22%0A%20%20%20%20%23%23%20Comparison%20with%20PhysicsNeMo%0A%0A%20%20%20%20From%20NVIDIA's%20published%20PhysicsNeMo%201D%20wave%20example%3A%0A%0A%20%20%20%20-%20PhysicsNeMo%20reported%20%24L_%5C%5Cinfty%24%20error%3A%20~%60%7Bphysicsnemo_linf%3A.3e%7D%60%20%20%0A%20%20%20%20%20%20Source%3A%20https%3A%2F%2Fdocs.nvidia.com%2Fphysicsnemo%2Flatest%2Fphysicsnemo-sym%2Fuser_guide%2Ffoundational%2F1d_wave_equation.html%23results%0A%20%20%20%20-%20Phydrax%20reference%20result%20(this%20setup%20family)%3A%20~%60%7Bphydrax_reference_linf%3A.3e%7D%60%20%20%0A%20%20%20%20%20%20%E2%86%92%20**%60%7Breference_linf_reduction%3A.1f%7D%25%60%20lower%20Linf%20error**%0A%20%20%20%20%7Bcurrent_linf_line%7D%0A%0A%20%20%20%20Training%2Fconfig%20context%3A%0A%0A%20%20%20%20-%20PhysicsNeMo%20config%20steps%3A%20%60%7Bphysicsnemo_steps%3A%2C%7D%60%0A%20%20%20%20-%20This%20notebook%20steps%3A%20%60%7Bour_steps%3A%2C%7D%60%20(**%60%7Bstep_reduction%3A.1f%7D%25%60%20fewer**)%0A%20%20%20%20-%20PhysicsNeMo%20collocation%20points%3A%20~%60%7Bphysicsnemo_collocation%3A%2C%7D%60%0A%20%20%20%20-%20This%20notebook%20collocation%20points%3A%20%60%7Bour_collocation%3A%2C%7D%60%20(**%60%7Bcollocation_ratio%3A.0f%7Dx%60%20more**)%0A%20%20%20%20%7Bspeed_line%7D%0A%0A%20%20%20%20PhysicsNeMo%20config%20reference%3A%0A%20%20%20%20https%3A%2F%2Fgithub.com%2FNVIDIA%2Fphysicsnemo-sym%2Fblob%2Fmain%2Fexamples%2Fwave_equation%2Fconf%2Fconfig.yaml%0A%0A%20%20%20%20Hardware%20context%20for%20these%20Phydrax%20runs%3A%20MacBook%20Pro%20M1%20Max%20laptop%20(no%20dedicated%20GPU).%0A%0A%0A%20%20%20%20**tl%3Bdr**%3A%20On%20a%20**laptop**%20with%20**no%20dedicated%20GPU**%2C%20we%20are%20able%20to%20**locally**%20solve%20this%20example%0A%20%20%20%20from%20the%20PhysicsNeMo%20documentation%20in%20a%20**couple%20of%20minutes**%2C%20with%20**100x**%20the%20batch%20size%2C%20%0A%20%20%20%20at%20**1%25**%20of%20the%20iterations%2C%0A%20%20%20%20and%20still%20get%20a%20**%3E90%25%20reduction%20in%20error**%2C%20while%20exactly%20satisfying%20***multiple%20exact%20initial%20and%20boundary%20conditions%20simultaneously***.%20All%20using%20out-of-the-box%20Phydrax%20components%2C%20highlighting%20our%20philosophy%20at%20Phydra%20Labs%20of%20%0A%20%20%20%20scalable%2C%20modular%2C%20composable%20high-performance%20SciML%20software%20tooling%2C%20where%20all%20the%20meticulous%20optimization%20and%20%0A%20%20%20%20implementation%20complexity%20is%20abstracted%20away%20from%20the%20user%20(human%20or%20agent!)%0A%0A%20%20%20%20*Interested%20in%20custom%20optimized%20software%20for%20your%20use%20case%3F%20Email%20us%20at%20partner%40phydra.ai*%0A%20%20%20%20%22%22%22%0A%20%20%20%20)%0A%20%20%20%20comparison_panel%0A%20%20%20%20return%0A%0A%0A%40app.cell(hide_code%3DTrue)%0Adef%20_(mo)%3A%0A%20%20%20%20mo.md(%0A%20%20%20%20%20%20%20%20%22%22%22%0A%20%20%20%20---%0A%20%20%20%20**Tips**%0A%0A%20%20%20%20-%20Run%20as%20a%20notebook%20editor%3A%20%60marimo%20edit%20examples%2Fwave1d.py%60%0A%20%20%20%20-%20Run%20as%20an%20app%3A%20%60python%20examples%2Fwave1d.py%60%0A%20%20%20%20%22%22%22%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
94d8d606b98d4bdc26714663107f2075fa5d6fb3c3f538e772755baecdd7fba2