import%20marimo%0A%0A__generated_with%20%3D%20%220.19.11%22%0Aapp%20%3D%20marimo.App(width%3D%22medium%22)%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%20Poisson's%20equation%20in%202D%0A%0A%20%20%20%20Poisson's%20equation%20on%20the%20unit%20square%20with%20homogeneous%20Dirichlet%20boundary%20conditions%20(the%20solution%20is%20zero%20on%20the%20entire%20domain)%20and%20a%20constant%20right%20hand%20side%20is%20given%20as%0A%0A%20%20%20%20%24%24%0A%20%20%20%20%5Cbegin%7Balign%7D%0A%20%20%20%20%5Cnabla%5E2%20u(x%2C%20y)%20%26%3D%202%2C%20%5Cquad%20(x%2C%20y)%20%5Cin%20%5COmega%20%3D%20(0%2C%201)%5E2%20%5Ctag%7B1%7D%5C%5C%0A%20%20%20%20u(x%2C%20y)%20%26%3D%200%2C%20%5Cquad%20(x%2C%20y)%20%5Cin%20%5Cpartial%20%5COmega%20%5Ctag%7B2%7D%0A%20%20%20%20%5Cend%7Balign%7D%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_(mo)%3A%0A%20%20%20%20mo.md(r%22%22%22%0A%20%20%20%20In%20order%20to%20solve%20this%20problem%20with%20the%20Galerkin%20method%20we%20choose%20a%20functionspace%20for%20our%20solution%20%24T%24%20and%20attempt%20to%20find%20%24u_%7BN%7D%20%5Cin%20T%24%20such%20that%0A%0A%20%20%20%20%24%24%0A%20%20%20%20(%5Cnabla%5E2%20u_N%2C%20v)_%7BL%5E2(%5COmega)%7D%20%3D%20(2%2C%20v)_%7BL%5E2(%5COmega)%7D%2C%20%5Cquad%20%5Cforall%20%5C%2C%20v%20%5Cin%20T%2C%20%5Ctag%7B3%7D%0A%20%20%20%20%24%24%0A%0A%20%20%20%20where%20the%20real%20multidimensional%20%24L%5E2(%5COmega)%24%20inner%20product%0A%0A%20%20%20%20%24%24%0A%20%20%20%20(u%2C%20v)_%7BL%5E2(%5COmega)%7D%20%3D%20%5Cint_%7B%5COmega%7D%20u%20%7Bv%7D%20d%5COmega.%20%5Ctag%7B4%7D%0A%20%20%20%20%24%24%0A%0A%20%20%20%20That's%20the%20full%20extent%20of%20the%20Galerkin%20method!%20Now%20all%20we%20have%20to%20do%20is%20be%20specific%20about%20%24T%24%20and%20everything%20will%20fall%20out%20from%20the%20variational%20equation.%20If%20we%20want%20to%20we%20can%20also%20use%20Green's%20first%20identity%20to%20arrive%20at%20another%20formulation%20that%20is%20very%20popular%20with%20the%20finite%20element%20method%20because%20the%20regularity%20of%20the%20solution%20%24u_N%24%20can%20be%20lower%20(%24u_N%24%20only%20need%20to%20be%20differentiable%20once%2C%20not%20twice)%0A%0A%20%20%20%20%24%24%0A%20%20%20%20-(%5Cnabla%20u_N%2C%20%5Cnabla%20v)_%7BL%5E2(%5COmega)%7D%20%3D%20(2%2C%20v)_%7BL%5E2(%5COmega)%7D%2C%20%5Cquad%20%5Cforall%20%5C%2C%20v%20%5Cin%20T.%20%5Ctag%7B5%7D%0A%20%20%20%20%24%24%0A%0A%20%20%20%20For%20the%20spectral%20Galerkin%20method%20using%20basis%20functions%20of%20very%20high%20order%20and%20regularity%2C%20there%20is%20practically%20no%20difference%20between%20the%20two%20approaches.%0A%0A%20%20%20%20Lets%20solve%20this%20problem%20using%20a%20composition%20of%20Legendre%20polynomials%20%24P_j%24%2C%20such%20that%20each%20basis%20function%20satisfies%20the%20homogeneous%20boundary%20conditions%0A%0A%20%20%20%20%24%24%0A%20%20%20%20%5Cphi_j%20%3D%20P_j%20-%20P_%7Bj%2B2%7D%2C%0A%20%20%20%20%24%24%0A%0A%20%20%20%20and%20%24V%20%3D%20%5Ctext%7Bspan%7D%5C%7B%5Cphi_j%5C%7D_%7Bj%3D0%7D%5E%7BN-1%7D%24%20and%20%24T%20%3D%20V%20%5Cotimes%20V%24.%20Hence%20we%20have%20test%20functions%20%24v%20%5Cin%20T%24%20such%20that%20each%20basis%20function%20is%0A%0A%20%20%20%20%24%24%0A%20%20%20%20v(x%2C%20y)%20%3D%20%5Cphi_i(x)%20%5Cphi_j(y).%0A%20%20%20%20%24%24%0A%0A%20%20%20%20We%20also%20create%20a%20trial%20function%20%24u_N%24%20with%20unknown%20expansion%20coefficients%20%24%5Chat%7Bu%7D_%7Bkl%7D%24%0A%0A%20%20%20%20%24%24%0A%20%20%20%20u_N(x%2C%20y)%20%3D%20%5Csum_%7Bk%7D%5Csum_l%20%5Chat%7Bu%7D_%7Bkl%7D%20%5Cphi_k(x)%20%5Cphi_l(y).%0A%20%20%20%20%24%24%0A%0A%20%20%20%20The%20unknown%20%24U%3D%5C%7B%5Chat%7Bu%7D_%7Bkl%7D%5C%7D_%7Bk%2Cl%3D0%7D%5E%7BN-1%7D%24%20is%20what%20we%20need%20to%20find%20through%20the%20Galerkin%20method.%0A%0A%20%20%20%20Inserting%20for%20test%20and%20trial%20functions%20in%20Eq.%20(3)%20we%20get%0A%0A%20%20%20%20%24%24%0A%20%20%20%20%5Csum_%7Bk%3D0%7D%5E%7BN-1%7D%5Csum_%7Bl%3D0%7D%5E%7BN-1%7D%20%5Cint_%7B%5COmega%7D%20%5Cleft(%20%5Cphi''_k(x)%20%5Cphi_i(x)%20%5Cphi_l(y)%5Cphi_j(y)%20%2B%20%5Cphi_k(x)%20%5Cphi_i(x)%20%5Cphi''_l(y)%5Cphi_j(y)%20%5Cright)%5C%2C%20d%20%5COmega%20%5C%2C%20%5Chat%7Bu%7D_%7Bkl%7D%20%3D%5Cint_%7B%5COmega%7D%202%20%5Cphi_i(x)%20%5Cphi_j(y)%20d%5COmega%2C%0A%20%20%20%20%24%24%0A%0A%20%20%20%20where%20the%20prime%20represents%20a%20derivative%20with%20respect%20to%20the%20given%20functions%20variable.%20That%20is%2C%20%24%5Cphi'_j(x)%20%3D%20%5Cfrac%7B%5Cpartial%20%5Cphi_j(x)%7D%7B%5Cpartial%20x%7D%24%20and%20%24%5Cphi'_j(y)%20%3D%20%5Cfrac%7B%5Cpartial%20%5Cphi_j(y)%7D%7B%5Cpartial%20y%7D%24.%0A%0A%20%20%20%20The%20tensor%20on%20the%20left%20hand%20side%20contains%20no%20less%20than%204%20indices%20and%20is%20a%20tensor%20of%20rank%204.%20However%2C%20since%20the%20basis%20functions%20are%20separable%20(they%20are%20all%20functions%20of%20one%20variable)%2C%20we%20can%20split%20the%20integral%20into%0A%0A%20%20%20%20%24%24%0A%20%20%20%20%5Csum_%7Bk%3D0%7D%5E%7BN-1%7D%5Csum_%7Bl%3D0%7D%5E%7BN-1%7D%20%5Cleft(%20%5Cint_%7B0%7D%5E1%20%5Cphi''_k(x)%20%5Cphi_i(x)%20dx%20%5Cint_%7B0%7D%5E1%20%5Cphi_l(y)%5Cphi_j(y)%20dy%20%2B%20%5Cint_%7B0%7D%5E1%20%5Cphi_k(x)%20%5Cphi_i(x)%20dx%20%5Cint_%7B0%7D%5E1%20%5Cphi''_l(y)%5Cphi_j(y)%5C%2C%20dy%20%5Cright)%20%5C%2C%20%5Chat%7Bu%7D_%7Bkl%7D%20%3D%5Cint_%7B%5COmega%7D%202%20%5Cphi_i(x)%20%5Cphi_j(y)%20d%5COmega.%0A%20%20%20%20%24%24%0A%0A%20%20%20%20Now%20introduce%20%24s_%7Bik%7D%20%3D%20%5Cint_0%5E1%5Cphi''_k(x)%20%5Cphi_i(x)%20dx%20%3D%20%5Cint_0%5E1%5Cphi''_k(y)%20%5Cphi_i(y)%20dy%24%20and%20%24a_%7Bjl%7D%20%3D%20%5Cint_0%5E1%20%5Cphi_j(x)%20%5Cphi_l(x)%20dx%20%3D%20%5Cint_0%5E1%20%5Cphi_j(y)%20%5Cphi_l(y)%20dy%24%2C%20where%20we%20get%20the%20same%20matrix%20for%20%24x%24%20and%20%24y%24%20directions%20since%20we%20use%20the%20same%20basis%20functions%20in%20both%20%24x%24%20and%20%24y%24.%20The%20equation%20becomes%0A%0A%20%20%20%20%24%24%0A%20%20%20%20%5Csum_%7Bk%3D0%7D%5E%7BN-1%7D%5Csum_%7Bl%3D0%7D%5E%7BN-1%7D%20%5Cleft(%20s_%7Bik%7D%20a_%7Bjl%7D%20%2B%20a_%7Bik%7D%20s_%7Bjl%7D%5Cright)%20%5Chat%7Bu%7D_%7Bkl%7D%20%3D%20f_%7Bij%7D%2C%0A%20%20%20%20%24%24%0A%0A%20%20%20%20where%20%24f_%7Bij%7D%20%3D%20%5Cint_%7B%5COmega%7D%202%20%5Cphi_i(x)%20%5Cphi_j(y)%20d%5COmega%24.%20In%20matrix%20form%2C%20with%20%24S%3D%5C%7Bs_%7Bik%7D%5C%7D_%7Bi%2Ck%3D0%7D%5E%7BN-1%7D%24%2C%20%24A%3D%5C%7Ba_%7Bik%7D%5C%7D_%7Bi%2Ck%3D0%7D%5E%7BN-1%7D%24%20and%20%24F%3D%5C%7Bf_%7Bij%7D%5C%7D_%7Bi%2Cj%3D0%7D%5E%7BN-1%7D%24%20we%20get%20the%20linear%20algebra%20problem%0A%0A%20%20%20%20%24%24%0A%20%20%20%20SUA%20%2B%20AUS%20%3D%20F%0A%20%20%20%20%24%24%0A%0A%20%20%20%20In%20order%20to%20solve%20for%20%24U%24%20we%20can%20use%20the%20%5BKronecker%20product%5D(https%3A%2F%2Fen.wikipedia.org%2Fwiki%2FKronecker_product)%20rule%20%24%5Ctext%7Bvec%7D(SUA)%20%3D%20(S%20%5Cotimes%20A%5ET)%20%5Ctext%7Bvec%7D(U)%24%2C%20where%20%24%5Cotimes%24%20represent%20a%20Kronecker%20product%20and%20%24%5Ctext%7Bvec%7D%24%20represents%20a%20vectorization%20(flattening)%20(see%20also%20%5Blecture%20notes%20in%20MAT-MEK4270%5D(https%3A%2F%2Fmatmek-4270.github.io%2Fmatmek4270-book%2Flecture6.html%23vectorization-the-vec-trick)).%20The%20linear%20algebra%20problem%20becomes%0A%0A%20%20%20%20%24%24%0A%20%20%20%20(S%20%5Cotimes%20A%5ET%20%2B%20A%20%5Cotimes%20S%5ET)%20%5Ctext%7Bvec%7D(U)%20%3D%20%5Ctext%7Bvec%7D(F)%0A%20%20%20%20%24%24%0A%0A%20%20%20%20which%20can%20be%20solved%20for%20the%20flattened%20%24%5Ctext%7Bvec%7D(U)%24.%0A%0A%20%20%20%20We%20are%20now%20ready%20to%20solve%20the%20Poisson%20problem%20using%20Jaxfun.%20With%20Jaxfun%20we%20choose%20functionspaces%20first%20and%20the%20assemble%20the%20matrices%20%24S%2C%20A%24%20and%20%24F%24%20from%20above.%20The%20entire%20problem%20is%20solved%20in%20the%20function%20%60poisson()%60%20below%20and%20should%20be%20pretty%20self-explanatory.%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%20matplotlib.pyplot%20as%20plt%0A%20%20%20%20import%20jax.numpy%20as%20jnp%0A%20%20%20%20from%20jax%20import%20Array%0A%20%20%20%20from%20jaxfun.operators%20import%20Div%2C%20Grad%0A%20%20%20%20from%20jaxfun.galerkin%20import%20(%0A%20%20%20%20%20%20%20%20inner%2C%0A%20%20%20%20%20%20%20%20FunctionSpace%2C%0A%20%20%20%20%20%20%20%20JAXFunction%2C%0A%20%20%20%20%20%20%20%20Legendre%2C%0A%20%20%20%20%20%20%20%20TestFunction%2C%0A%20%20%20%20%20%20%20%20TensorProduct%2C%0A%20%20%20%20%20%20%20%20TrialFunction%2C%0A%20%20%20%20)%0A%20%20%20%20from%20jaxfun.galerkin.tensorproductspace%20import%20vec%0A%20%20%20%20from%20scipy%20import%20sparse%20as%20scipy_sparse%0A%0A%20%20%20%20def%20poisson(N%3A%20int)%20-%3E%20Array%3A%0A%20%20%20%20%20%20%20%20V%20%3D%20FunctionSpace(%0A%20%20%20%20%20%20%20%20%20%20%20%20N%20%2B%202%2C%0A%20%20%20%20%20%20%20%20%20%20%20%20Legendre.Legendre%2C%0A%20%20%20%20%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%20%20%20%20%20domain%3D(0%2C%201)%2C%0A%20%20%20%20%20%20%20%20%20%20%20%20name%3D%22V%22%2C%0A%20%20%20%20%20%20%20%20%20%20%20%20fun_str%3D%22phi%22%2C%0A%20%20%20%20%20%20%20%20)%0A%20%20%20%20%20%20%20%20T%20%3D%20TensorProduct(V%2C%20V%2C%20name%3D%22T%22)%0A%20%20%20%20%20%20%20%20u%20%3D%20TrialFunction(T%2C%20name%3D%22u%22)%0A%20%20%20%20%20%20%20%20v%20%3D%20TestFunction(T%2C%20name%3D%22v%22)%0A%20%20%20%20%20%20%20%20A%20%3D%20inner(Div(Grad(u))%20*%20v)%0A%20%20%20%20%20%20%20%20F%20%3D%20inner(2%20*%20v)%0A%20%20%20%20%20%20%20%20A0%20%3D%20vec(A)%0A%20%20%20%20%20%20%20%20h%20%3D%20jnp.array(scipy_sparse.linalg.spsolve(A0%2C%20F.flatten()).reshape(F.shape))%0A%20%20%20%20%20%20%20%20xj%20%3D%20T.mesh(kind%3D%22uniform%22%2C%20N%3D(50%2C%2050)%2C%20broadcast%3DFalse)%0A%20%20%20%20%20%20%20%20plt.contourf(xj%5B0%5D%2C%20xj%5B1%5D%2C%20T.backward(h%2C%20kind%3D%22uniform%22%2C%20N%3D(50%2C%2050)))%0A%20%20%20%20%20%20%20%20plt.colorbar()%0A%20%20%20%20%20%20%20%20plt.show()%0A%20%20%20%20%20%20%20%20return%20JAXFunction(h%2C%20T%2C%20name%3D%22h%22)%0A%0A%20%20%20%20h%20%3D%20poisson(20)%0A%20%20%20%20return%20h%2C%20jnp%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%20returned%20%60JAXFunction%60%20%60h%60%20represents%0A%0A%20%20%20%20%24%24%0A%20%20%20%20h(x%2C%20y)%20%3D%20%5Csum_%7Bi%3D0%7D%5E%7BN-1%7D%5Csum_%7Bj%3D0%7D%5E%7BN-1%7D%20%5Chat%7Bh%7D_%7Bij%7D%20%5Cphi_i(x)%20%5Cphi_j(y)%2C%0A%20%20%20%20%24%24%0A%0A%20%20%20%20and%20can%20be%20evaluated%20for%20any%20point%20in%20the%20domain%20%24%5B0%2C%201%5D%5E2%24.%20We%20can%20display%20%60h%60%20both%20in%20unevaluated%20and%20evaluated%20state%3A%0A%20%20%20%20%22%22%22)%0A%20%20%20%20return%0A%0A%0A%40app.cell%0Adef%20_(h)%3A%0A%20%20%20%20from%20IPython.display%20import%20display%0A%0A%20%20%20%20display(h)%0A%20%20%20%20display(h.doit(linear%3DTrue))%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%20Note%20the%20repeated%20indices%20on%20the%20last%20expression%2C%20which%20imply%20summation.%0A%0A%20%20%20%20Evaluate%20%60h%60%20at%20a%20single%20point%20%24h(x%3D0.5%2C%20y%3D0.5)%24%20in%20space%3A%0A%20%20%20%20%22%22%22)%0A%20%20%20%20return%0A%0A%0A%40app.cell%0Adef%20_(h%2C%20jnp)%3A%0A%20%20%20%20h(jnp.array(%5B0.5%2C%200.5%5D))%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%20Evaluate%20%60h%60%20on%20all%20quadrature%20points%20in%20the%20mesh%3A%0A%20%20%20%20%22%22%22)%0A%20%20%20%20return%0A%0A%0A%40app.cell%0Adef%20_(h)%3A%0A%20%20%20%20xj%20%3D%20h.functionspace.mesh()%0A%20%20%20%20hj%20%3D%20h.evaluate_mesh(xj)%0A%20%20%20%20print(hj.shape)%0A%20%20%20%20return%20(hj%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%20operation%20is%20normally%20called%20a%20backward%20transform%2C%20and%20it%20can%20also%20be%20computed%20as%0A%20%20%20%20%22%22%22)%0A%20%20%20%20return%0A%0A%0A%40app.cell%0Adef%20_(h%2C%20hj%2C%20jnp)%3A%0A%20%20%20%20hj2%20%3D%20h.backward()%0A%20%20%20%20assert%20jnp.allclose(hj%2C%20hj2)%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%20Note%20that%20each%20functionspace%20V%20contains%20%24N%24%20unknowns%2C%20but%20%24N%2B2%24%20quadrature%20points%20since%20there%20are%20two%20restrictions%20on%20the%20boundary%20conditions.%0A%0A%20%20%20%20Alternatively%2C%20use%20a%20flattened%20mesh%20of%20shape%20%24((N%2B2)%5E2%2C%202)%24%3A%0A%20%20%20%20%22%22%22)%0A%20%20%20%20return%0A%0A%0A%40app.cell%0Adef%20_(h)%3A%0A%20%20%20%20xf%20%3D%20h.functionspace.flatmesh()%0A%20%20%20%20hf%20%3D%20h(xf)%0A%20%20%20%20print(hf.shape%2C%20xf.shape)%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
7d0800b3ee65dc311429c41fb9038854