import%20marimo%0A%0A__generated_with%20%3D%20%220.19.8%22%0Aapp%20%3D%20marimo.App()%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%20Helmholtz'%20equation%20in%20an%20annulus%20using%20Jaxfun%20with%20polar%20coordinates%0A%0A%20%20%20%20In%20this%20demo%20we%20will%20solve%20Helmholtz'%20equation%20in%20an%20annulus%20domain%20using%20polar%20coordinates%20and%20spectral%20accuracy.%20Helmholtz'%20equation%20reads%0A%0A%20%20%20%20%24%24%0A%20%20%20%20%5Cnabla%5E2%20u(x%2C%20y)%20%2B%20%5Calpha%20u(x%2C%20y)%20%3D%20f(x%2C%20y)%2C%5Cquad%20x%2C%20y%20%5Cin%20%5COmega%0A%20%20%20%20%24%24%0A%0A%20%20%20%20where%20%24%5COmega%20%3D%20%5C%7Bx%2C%20y%20%5C%2C%20%7C%20%5C%2C%20%20r_0%20%3C%20%5Csqrt%7Bx%5E2%20%2B%20y%5E2%7D%20%3C%20r_1%5C%7D%24%20and%20%24r_0%20%3C%20r_1%24%20are%20the%20inner%20and%20outer%20radii%20of%20the%20annulus.%20The%20parameter%20%24%5Calpha%24%20is%20a%20constant%20and%20%24f(x%2C%20y)%24%20is%20a%20real%20function.%20We%20use%20homogeneous%20Dirichlet%20boundary%20conditions%20on%20the%20entire%20boundary.%0A%0A%20%20%20%20We%20start%20by%20importing%20necessary%20classes%20and%20functions%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%20jax.numpy%20as%20jnp%0A%20%20%20%20import%20matplotlib.pyplot%20as%20plt%0A%20%20%20%20import%20sympy%20as%20sp%0A%0A%20%20%20%20from%20jaxfun.coordinates%20import%20get_CoordSys%0A%20%20%20%20from%20jaxfun.galerkin%20import%20(%0A%20%20%20%20%20%20%20%20Fourier%2C%0A%20%20%20%20%20%20%20%20FunctionSpace%2C%0A%20%20%20%20%20%20%20%20Legendre%2C%0A%20%20%20%20%20%20%20%20TensorProductSpace%2C%0A%20%20%20%20%20%20%20%20TestFunction%2C%0A%20%20%20%20%20%20%20%20TrialFunction%2C%0A%20%20%20%20%20%20%20%20inner%2C%0A%20%20%20%20)%0A%20%20%20%20from%20jaxfun.operators%20import%20Div%2C%20Grad%0A%20%20%20%20from%20jaxfun.utils%20import%20Domain%0A%0A%20%20%20%20return%20(%0A%20%20%20%20%20%20%20%20Div%2C%0A%20%20%20%20%20%20%20%20Domain%2C%0A%20%20%20%20%20%20%20%20Fourier%2C%0A%20%20%20%20%20%20%20%20FunctionSpace%2C%0A%20%20%20%20%20%20%20%20Grad%2C%0A%20%20%20%20%20%20%20%20Legendre%2C%0A%20%20%20%20%20%20%20%20TensorProductSpace%2C%0A%20%20%20%20%20%20%20%20TestFunction%2C%0A%20%20%20%20%20%20%20%20TrialFunction%2C%0A%20%20%20%20%20%20%20%20get_CoordSys%2C%0A%20%20%20%20%20%20%20%20inner%2C%0A%20%20%20%20%20%20%20%20jnp%2C%0A%20%20%20%20%20%20%20%20plt%2C%0A%20%20%20%20%20%20%20%20sp%2C%0A%20%20%20%20)%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%20We%20then%20create%20a%20polar%20coordinate%20system%0A%20%20%20%20%22%22%22)%0A%20%20%20%20return%0A%0A%0A%40app.cell%0Adef%20_(get_CoordSys%2C%20sp)%3A%0A%20%20%20%20_r%2C%20_theta%20%3D%20sp.symbols(%22r%2Ctheta%22%2C%20real%3DTrue%2C%20positive%3DTrue)%0A%20%20%20%20C%20%3D%20get_CoordSys(%0A%20%20%20%20%20%20%20%20%22C%22%2C%20sp.Lambda((_r%2C%20_theta)%2C%20(_r%20*%20sp.cos(_theta)%2C%20_r%20*%20sp.sin(_theta)))%0A%20%20%20%20)%0A%20%20%20%20return%20(C%2C)%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%20This%20coordinate%20system%20is%20further%20used%20to%20create%20a%20tensor%20product%20space%2C%20where%20the%20radial%20direction%20uses%20composite%20Legendre%20polynomials%20(basis%20functions%20%24%5Cpsi_i%3DP_i-P_%7Bi%2B2%7D%24)%20and%20the%20angular%20direction%20uses%20Fourier%20exponentials.%0A%20%20%20%20%22%22%22)%0A%20%20%20%20return%0A%0A%0A%40app.cell%0Adef%20_(C%2C%20Domain%2C%20Fourier%2C%20FunctionSpace%2C%20Legendre%2C%20TensorProductSpace%2C%20sp)%3A%0A%20%20%20%20r0%2C%20r1%20%3D%20sp.S.Half%2C%201%0A%20%20%20%20R%20%3D%20FunctionSpace(%0A%20%20%20%20%20%20%20%2020%2C%0A%20%20%20%20%20%20%20%20Legendre.Legendre%2C%0A%20%20%20%20%20%20%20%20bcs%3D%7B%22left%22%3A%20%7B%22D%22%3A%200%7D%2C%20%22right%22%3A%20%7B%22D%22%3A%200%7D%7D%2C%0A%20%20%20%20%20%20%20%20domain%3DDomain(r0%2C%20r1)%2C%0A%20%20%20%20%20%20%20%20name%3D%22R%22%2C%0A%20%20%20%20%20%20%20%20fun_str%3D%22phi%22%2C%0A%20%20%20%20)%0A%20%20%20%20F%20%3D%20FunctionSpace(20%2C%20Fourier.Fourier%2C%20name%3D%22F%22%2C%20fun_str%3D%22E%22)%0A%20%20%20%20P%20%3D%20TensorProductSpace((R%2C%20F)%2C%20system%3DC%2C%20name%3D%22P%22)%0A%20%20%20%20return%20(P%2C)%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%20To%20implement%20Helmholtz'%20equation%20we%20need%20test%20and%20trial%20functions%20for%20the%20tensor%20product%20space%0A%20%20%20%20%22%22%22)%0A%20%20%20%20return%0A%0A%0A%40app.cell%0Adef%20_(P%2C%20TestFunction%2C%20TrialFunction)%3A%0A%20%20%20%20u%20%3D%20TrialFunction(P%2C%20name%3D%22u%22)%0A%20%20%20%20v%20%3D%20TestFunction(P%2C%20name%3D%22v%22)%0A%20%20%20%20return%20u%2C%20v%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%20The%20test%20and%20trial%20functions%20are%20subclasses%20of%20the%20Sympy%20%5BFunction%5D(https%3A%2F%2Fdocs.sympy.org%2Flatest%2Fmodules%2Fcore.html%23sympy.core.function.Function)%0A%20%20%20%20%22%22%22)%0A%20%20%20%20return%0A%0A%0A%40app.cell%0Adef%20_(u)%3A%0A%20%20%20%20u%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%20Here%20we%20see%20that%20the%20trial%20function%20%60u%60%20is%20a%20function%20of%20spatial%20coordinates%20%60x%60%20and%20%60y%60%2C%20and%20it%20is%20a%20trial%20function%20of%20the%20%60P%60%20function%20space.%20In%20computational%20coordinates%20%60u%60%20is%20a%20tensor%20product%20function%20using%20a%20tensor%20product%20between%20the%20trial%20functions%20from%20the%20one-dimensional%20%60R%60%20and%20%60F%60%20spaces.%20%60u%60%20is%20evaluated%20in%20computational%20coordinates%20using%20the%20%60doit()%60%20method%0A%20%20%20%20%22%22%22)%0A%20%20%20%20return%0A%0A%0A%40app.cell%0Adef%20_(u)%3A%0A%20%20%20%20u.doit()%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%20Notice%20that%20%24%5Cphi_k(r)%24%20is%20component%20%24k%24%20of%20the%20trial%20functions%20of%20the%20radial%20coordinate%20and%20%24E_l(%5Ctheta)%3D%5Cexp(%5Cimath%20%5Cunderline%7Bl%7D%20%5Ctheta)%24%20is%20component%20%24l%24%20of%20the%20complex%20exponentials.%20(We%20defined%20the%20function%20names%20%22phi%22%20and%20%22E%22%20when%20creating%20the%20spaces.)%20The%20trial%20function%20is%20actually%20an%20expansion%0A%0A%20%20%20%20%24%24%0A%20%20%20%20u(x%2C%20y)%20%3D%20U(r%2C%20%5Ctheta)%20%3D%20%5Csum_%7Bk%3D0%7D%5E%7BN-3%7D%5Csum_%7Bl%3D0%7D%5E%7BN-1%7D%20%5Chat%7Bu%7D_%7Bkl%7D%20%5Cphi_%7Bk%7D(r)%20%5Cexp(%5Cimath%20%5Cunderline%7Bl%7D%20%5Ctheta)%0A%20%20%20%20%24%24%0A%0A%20%20%20%20where%0A%0A%20%20%20%20%24%24%0A%20%20%20%20%5Cunderline%7Bl%7D%20%3D%20%5Cbegin%7Bcases%7D%20l%2C%20%5Cquad%20%26%5Ctext%7Bif%7D%20%5C%2C%20l%20%3C%20N%2F2%20%5C%5C%0A%20%20%20%20-(N-l)%20%5Cquad%20%26%5Ctext%7Bif%7D%20%5C%2C%20l%20%5Cge%20N%2F2%0A%20%20%20%20%5Cend%7Bcases%7D%0A%20%20%20%20%24%24%0A%0A%20%20%20%20However%2C%20we%20do%20not%20print%20the%20expansion%20factors%20%24%5C%7B%5Chat%7Bu%7D_%7Bkl%7D%5C%7D%24%2C%20only%20the%20basis%20functions.%20The%20expansion%20factors%20are%20the%20unknown%20that%20we%20will%20compute%20in%20the%20end.%0A%0A%20%20%20%20Helmholtz'%20equation%20is%20now%0A%20%20%20%20%22%22%22)%0A%20%20%20%20return%0A%0A%0A%40app.cell%0Adef%20_(Div%2C%20Grad%2C%20u)%3A%0A%20%20%20%20alpha%20%3D%201%0A%20%20%20%20eq%20%3D%20Div(Grad(u))%20%2B%20alpha%20*%20u%0A%20%20%20%20eq%0A%20%20%20%20return%20alpha%2C%20eq%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%20Like%20the%20trial%20function%20we%20can%20also%20evaluate%20the%20equation%20in%20computational%20coordinates%20using%20%60doit()%60%0A%20%20%20%20%22%22%22)%0A%20%20%20%20return%0A%0A%0A%40app.cell%0Adef%20_(eq)%3A%0A%20%20%20%20eq.doit()%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%20Notice%20that%20we%20get%20the%20Helmholtz%20equation%20in%20polar%20computational%20coordinates%2C%20and%20the%20trial%20functions%20are%20expanded%20into%20tensor%20products.%20It%20looks%20a%20bit%20better%20if%20we%20multiply%20through%20with%20the%20radius%20%24r%5E2%24%3A%0A%20%20%20%20%22%22%22)%0A%20%20%20%20return%0A%0A%0A%40app.cell%0Adef%20_(C%2C%20eq)%3A%0A%20%20%20%20r%2C%20theta%20%3D%20C.base_scalars()%0A%20%20%20%20C.simplify((r**2%20*%20eq).doit())%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%20If%20you%20want%20to%20see%20what%20the%20Helmholtz%20equation%20looks%20like%20with%20a%20regular%20(non-tensor-product)%20function%2C%20then%20you%20can%20create%20a%20%60ScalarFunction%60%20and%20use%20that%20instead%20of%20the%20trial%20function%0A%20%20%20%20%22%22%22)%0A%20%20%20%20return%0A%0A%0A%40app.cell%0Adef%20_(C)%3A%0A%20%20%20%20from%20jaxfun.galerkin.arguments%20import%20ScalarFunction%0A%0A%20%20%20%20h%20%3D%20ScalarFunction(%22h%22%2C%20C)%0A%20%20%20%20h%0A%20%20%20%20return%20(h%2C)%0A%0A%0A%40app.cell%0Adef%20_(C%2C%20Div%2C%20Grad%2C%20alpha%2C%20h)%3A%0A%20%20%20%20C.simplify((Div(Grad(h))%20%2B%20alpha%20*%20h).doit())%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%20Notice%20that%20the%20scalar%20function%20use%20a%20lowercase%20letter%20in%20physical%20space%20and%20an%20upper%20case%20letter%20in%20computational%20space%0A%0A%20%20%20%20%24%24%0A%20%20%20%20h(x%2C%20y)%20%3D%20H(r%2C%20%5Ctheta).%0A%20%20%20%20%24%24%0A%20%20%20%20%22%22%22)%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%20In%20order%20to%20solve%20the%20Helmholtz%20equation%20we%20need%20to%20define%20%24f(x%2C%20y)%24%20and%20assemble%20the%20inner%20products%0A%0A%20%20%20%20%24%24%0A%20%20%20%20(%5Cnabla%5E2u%20%2B%20%5Calpha%20u%2C%20v)_%7BL%5E2(%5COmega)%7D%20%3D%20(f%2C%20v)_%7BL%5E2(%5COmega)%7D%0A%20%20%20%20%24%24%0A%0A%20%20%20%20Since%20the%20Fourier%20exponentials%20are%20complex%20functions%20we%20need%20to%20use%20a%20complex%20inner%20product%20defined%20as%0A%0A%20%20%20%20%24%24%0A%20%20%20%20(a%2C%20b)_%7BL%5E2(%5COmega)%7D%20%3D%20%5Cint_%7B%5COmega%7D%20%5C%2Ca%20%5C%2C%20%5Coverline%7Bb%7D%5C%2C%20d%5COmega%2C%0A%20%20%20%20%24%24%0A%0A%20%20%20%20where%20%24%5Coverline%7Bb%7D%24%20is%20the%20complex%20conjugate%20of%20%24b%24.%0A%0A%20%20%20%20We%20can%20assemble%20both%20sides%20of%20Helmholtz'%20equation%20in%20one%20call.%20Here%20we%20actually%20assemble%20all%20terms%20in%0A%0A%20%20%20%20%24%24%0A%20%20%20%20(%5Cnabla%5E2%20u%20%2B%20%5Calpha%20u%20-%20f%2C%20v)_%7BL%5E2(%5COmega)%7D%20%3D%200%0A%20%20%20%20%24%24%0A%20%20%20%20%22%22%22)%0A%20%20%20%20return%0A%0A%0A%40app.cell%0Adef%20_(eq%2C%20inner%2C%20sp%2C%20v)%3A%0A%20%20%20%20A%2C%20b%20%3D%20inner((eq%20-%202)%20*%20sp.conjugate(v))%0A%20%20%20%20return%20A%2C%20b%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%20Now%20all%20matrices%20will%20be%20collected%20in%20%60A%60%2C%20whereas%20the%20right%20hand%20side%20vector%20is%20collected%20in%20%60b%60%0A%20%20%20%20%22%22%22)%0A%20%20%20%20return%0A%0A%0A%40app.cell%0Adef%20_(A)%3A%0A%20%20%20%20A%0A%20%20%20%20return%0A%0A%0A%40app.cell%0Adef%20_(b)%3A%0A%20%20%20%20b.shape%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%20The%20tensor%20product%20matrices%20in%20A%20are%20tensors%20with%20four%20indices%2C%20and%20each%20item%20is%20the%20outer%20product%20of%20two%20(1D)%20matrices.%20We%20can%20see%20which%20matrices%20there%20are%20by%20inspecting%20the%20equation%20in%20computational%20space%0A%20%20%20%20%22%22%22)%0A%20%20%20%20return%0A%0A%0A%40app.cell%0Adef%20_(C%2C%20eq%2C%20sp%2C%20v)%3A%0A%20%20%20%20C.simplify((eq%20*%20sp.conjugate(v)).doit())%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%20Here%20the%20test%20function%20uses%20indices%20%24i%24%20and%20%24j%24%2C%20whereas%20the%20trial%20function%20uses%20%24k%24%20and%20%24l%24.%20Each%20term%20with%20four%20indices%20is%20one%20%60TPMatrix%60%20in%20%60A%60.%20As%20seen%20above%2C%20there%20are%204%20tensor%20product%20matrices%20(4%20items%20with%204%20indices).%20Note%20that%20the%20complex%20conjugate%20of%20the%20test%20function%20is%20evaluated%20automatically%20inside%20the%20%60inner%60%20function%20and%20it%20is%20not%20necessary%20to%20call%20it%20explicitly%20as%20above.%20Hence%20%60eq*v%60%20would%20produce%20the%20same%20result%2C%20but%20it%20would%20not%20print%20as%20nicely%20(and%20correctly)%20as%20above.%20Also%2C%20when%20evaluating%20the%20inner%20product%2C%20the%20equation%20is%20multiplied%20by%20an%20additional%20%60r%60%2C%20since%20the%20domain%20measure%20%24d%5COmega%20%3D%20rdrd%5Ctheta%24%0A%0A%20%20%20%20In%20the%20end%20we%20assemble%20a%20Kronecker%20product%20matrix%20from%20all%20the%20items%20in%20%60A%60%20and%20then%20solve%20the%20linear%20system%20of%20equations%3A%0A%20%20%20%20%22%22%22)%0A%20%20%20%20return%0A%0A%0A%40app.cell%0Adef%20_(A%2C%20b%2C%20jnp)%3A%0A%20%20%20%20M%20%3D%20jnp.sum(jnp.array(%5Bjnp.kron(*a.mats)%20for%20a%20in%20A%5D)%2C%20axis%3D0)%0A%20%20%20%20uh%20%3D%20jnp.linalg.solve(M%2C%20b.flatten()).reshape(b.shape)%0A%20%20%20%20return%20(uh%2C)%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%20We%20can%20plot%20the%20solution%20in%20Cartesian%20coordinates%20on%20a%20uniform%20grid%20with%20100%20by%20100%20points%3A%0A%20%20%20%20%22%22%22)%0A%20%20%20%20return%0A%0A%0A%40app.cell%0Adef%20_(P%2C%20plt%2C%20uh)%3A%0A%20%20%20%20xc%2C%20yc%20%3D%20P.cartesian_mesh(kind%3D%22uniform%22%2C%20N%3D(100%2C%20100))%0A%20%20%20%20uj%20%3D%20P.backward(uh%2C%20kind%3D%22uniform%22%2C%20N%3D(100%2C%20100))%0A%20%20%20%20fig%20%3D%20plt.figure(figsize%3D(8%2C%206))%0A%20%20%20%20ax%20%3D%20fig.gca()%0A%20%20%20%20ax.contourf(xc%2C%20yc%2C%20uj.real%2C%2050)%0A%20%20%20%20plt.axis(%22equal%22)%0A%20%20%20%20plt.show()%0A%20%20%20%20return%0A%0A%0A%40app.cell(hide_code%3DTrue)%0Adef%20_()%3A%0A%20%20%20%20import%20marimo%20as%20mo%0A%0A%20%20%20%20return%20(mo%2C)%0A%0A%0Aif%20__name__%20%3D%3D%20%22__main__%22%3A%0A%20%20%20%20app.run()%0A
72ea688b36b5275d108a6fbfb20322f0