import%20marimo%0A%0A__generated_with%20%3D%20%220.19.8%22%0Aapp%20%3D%20marimo.App(width%3D%22medium%22)%0A%0A%0A%40app.cell%0Adef%20_(mo)%3A%0A%20%20%20%20mo.md(r%22%22%22%0A%20%20%20%20%23%20The%20form%20language%20in%20Jaxfun%20(WIP)%0A%0A%20%20%20%20Jaxfun%20consists%20of%20a%20rich%20form-language%20that%20is%20used%20to%20formulate%20problems.%20This%20language%20is%20built%20on%20top%20of%20Sympy's%20Function%20and%20Vector%20classes.%20The%20fundamental%20idea%20is%20that%20you%20choose%20a%20functionspace%20and%20then%20you%20create%20functions%20on%20that%20space.%20For%20the%20global%20Galerkin%20method%2C%20this%20means%20that%20you%20work%20with%20functions%20like%0A%0A%20%20%20%20%5Cbegin%7Bequation%7D%0A%20%20%20%20u(x)%20%3D%20%5Csum_%7Bi%3D0%7D%5EN%20%5Chat%7Bu%7D_i%20%5Cphi_i(x)%20%5Ctag%7B1%7D%0A%20%20%20%20%5Cend%7Bequation%7D%0A%0A%20%20%20%20where%20%24u(x)%24%20is%20your%20function%2C%20%24%5Cphi_i(x)%24%20are%20%24N%2B1%24%20basis%20functions%20and%20%24%5Chat%7Bu%7D_i%24%20are%20%24N%2B1%24%20unknown%20expansion%20coefficients.%20The%20functionspace%20corresponding%20to%20this%20function%20is%20%24V%3D%5Ctext%7Bspan%7D%5C%7B%5Cphi_i%5C%7D_%7Bi%3D0%7D%5EN%24.%20When%20we%20say%20that%20%24u(x)%20%5Cin%20V%24%2C%20this%20means%20that%20%24u%24%20can%20be%20expanded%20as%20shown%20in%20Eq.%20(1)%2C%20but%20we%20do%20not%20necessarily%20know%20the%20coefficients.%0A%0A%20%20%20%20A%20tensorproductspace%20in%202D%20Cartesian%20space%20can%20be%20defined%20as%20%24T%20%3D%20V%20%5Cotimes%20V%24.%20A%20scalar%20function%20%24f%20%5Cin%20T%24%20is%20then%20defined%20as%0A%0A%20%20%20%20%5Cbegin%7Bequation%7D%0A%20%20%20%20f(x%2C%20y)%20%3D%20%5Csum_%7Bi%3D0%7D%5EN%20%5Csum_%7Bj%3D0%7D%5EN%20%5Chat%7Bf%7D_%7Bij%7D%20%5Cphi_i(x)%20%5Cphi_j(y)%20%5Ctag%7B2%7D%0A%20%20%20%20%5Cend%7Bequation%7D%0A%0A%20%20%20%20with%20the%20unknown%20matrix%20%24F%3D%5C%7Bf_%7Bij%7D%5C%7D_%7Bi%2Cj%3D0%7D%5EN%20%5Cin%20%5Cmathbb%7BR%7D%5E%7B(N%2B1)%20%5Ctimes%20(N%2B1)%7D%24.%0A%0A%20%20%20%20A%20vector%20tensorproductspace%20in%202D%20can%20be%20defined%20as%20%24W%20%3D%20T%20%5Ctimes%20T%24.%20A%20function%20%24%5Cboldsymbol%7Bw%7D%20%5Cin%20W%24%20is%20then%20defined%20as%0A%0A%20%20%20%20%5Cbegin%7Bequation%7D%0A%20%20%20%20%5Cboldsymbol%7Bw%7D(x%2C%20y)%20%3D%20w_x(x%2C%20y)%20%5Cboldsymbol%7Bi%7D%20%2B%20w_y(x%2C%20y)%20%5Cboldsymbol%7Bj%7D%20%5Ctag%7B3%7D%0A%20%20%20%20%5Cend%7Bequation%7D%0A%0A%20%20%20%20where%20%24%5Cboldsymbol%7Bi%7D%24%20and%20%24%5Cboldsymbol%7Bj%7D%24%20are%20Cartesian%20unit%20vectors%20in%20%24x%24%20and%20%24y%24%20direction%2C%20respectively.%20In%20Eq.%20(3)%20both%20%24w_x%24%20and%20%24w_y%24%20can%20be%20expanded%20as%20shown%20in%20Eq%20(2).%0A%0A%20%20%20%20A%20multilayer%20perceptron%20neural%20network%20is%20defined%20as%20the%20function%0A%0A%20%20%20%20%24%24%0A%20%20%20%20M(%5Cboldsymbol%7Bz%7D%3B%20%5Ctheta)%20%3D%20%20W%5E%7BL%7D%20%5Csigma(%20W%5E%7BL-1%7D%20%5Cldots%20%5Csigma(%20W%5E1%20%5Cboldsymbol%7Bz%7D%20%2B%20%5Cboldsymbol%7Bb%7D%5E1)%20%5Cldots%20%2B%20%5Cboldsymbol%7Bb%7D%5E%7BL-1%7D)%20%20%2B%20%5Cboldsymbol%7Bb%7D%5EL%0A%20%20%20%20%24%24%0A%0A%20%20%20%20where%20%24%5Ctheta%20%3D%20%5C%7BW%5El%2C%20%5Cboldsymbol%7Bb%7D%5El%5C%7D_%7Bl%3D1%7D%5EL%24%20represents%20all%20the%20unknowns%20in%20the%20model%20and%20%24W%5El%2C%20%5Cboldsymbol%7Bb%7D%5El%24%20represents%20the%20weights%20and%20biases%20on%20level%20%24l%24.%20The%20model%20input%20%24%5Cboldsymbol%7Bz%7D%24%20can%20be%20anything.%20For%20example%2C%20%24%5Cboldsymbol%7Bz%7D%24%20can%20represent%20%24(x)%2C%20(x%2C%20y)%2C%20(x%2C%20y%2C%20z)%2C%20(x%2C%20t)%2C%20(x%2C%20y%2C%20t)%2C%20...%24.%20For%20%24Z%24%20number%20of%20inputs%20the%20scalar%20function%20%24M_%7B%5Ctheta%7D%3A%20%5Cmathbb%7BR%7D%5EZ%20%5Crightarrow%20%5Cmathbb%7BR%7D%24.%20Similarly%2C%20a%20vector%20function%20in%202D%20space%20is%20simply%20%24M_%7B%5Ctheta%7D%3A%20%5Cmathbb%7BR%7D%5EZ%20%5Crightarrow%20%5Cmathbb%7BR%7D%5E2%24.%0A%0A%20%20%20%20In%20Jaxfun%20you%20choose%20the%20functionspace%20and%20then%20you%20create%20functions.%20Let's%20create%20samples%20of%20all%20the%20three%20Galerkin%20functions%20defined%20above%20(%24u(x)%2C%20f(x%2C%20y)%2C%20%5Cboldsymbol%7Bw%7D(x%2C%20y)%24)%3A%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%20from%20IPython.display%20import%20display%0A%0A%20%20%20%20from%20jaxfun.galerkin%20import%20(%0A%20%20%20%20%20%20%20%20Legendre%2C%0A%20%20%20%20%20%20%20%20TensorProduct%2C%0A%20%20%20%20%20%20%20%20TestFunction%2C%0A%20%20%20%20%20%20%20%20VectorTensorProductSpace%2C%0A%20%20%20%20)%0A%0A%20%20%20%20N%20%3D%208%0A%20%20%20%20V%20%3D%20Legendre.Legendre(N%2C%20name%3D%22V%22%2C%20fun_str%3D%22phi%22)%0A%20%20%20%20T%20%3D%20TensorProduct(V%2C%20V%2C%20name%3D%22T%22)%0A%20%20%20%20W%20%3D%20VectorTensorProductSpace(T%2C%20name%3D%22W%22)%0A%20%20%20%20u%20%3D%20TestFunction(V%2C%20name%3D%22u%22)%0A%20%20%20%20f%20%3D%20TestFunction(T%2C%20name%3D%22f%22)%0A%20%20%20%20w%20%3D%20TestFunction(W%2C%20name%3D%22w%22)%0A%20%20%20%20display(u)%0A%20%20%20%20display(f)%0A%20%20%20%20display(w)%0A%20%20%20%20return%20N%2C%20T%2C%20V%2C%20W%2C%20display%2C%20f%2C%20u%2C%20w%0A%0A%0A%40app.cell%0Adef%20_(mo)%3A%0A%20%20%20%20mo.md(r%22%22%22%0A%20%20%20%20Note%20that%20the%20three%20displayed%20functions%20are%20all%20in%20unevaluated%20state.%20They%20appear%20as%20the%20left%20hand%20side%20of%20their%20definition.%20We%20see%20%24u(x%3B%20V)%24%2C%20which%20is%20the%20left%20hand%20side%20of%20Eq.%20(1)%2C%20with%20an%20additional%20%24V%24%20to%20show%20that%20this%20function%20has%20been%20created%20from%20the%20space%20named%20%22V%22.%0A%0A%20%20%20%20We%20can%20get%20the%20basisfunctions%20from%20the%20right%20hand%20side%20by%20evaluating%20the%20expressions%3A%0A%20%20%20%20%22%22%22)%0A%20%20%20%20return%0A%0A%0A%40app.cell%0Adef%20_(display%2C%20f%2C%20u%2C%20w)%3A%0A%20%20%20%20display(u.doit())%0A%20%20%20%20display(f.doit())%0A%20%20%20%20display(w.doit())%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%20Note%20that%20since%20TestFunctions%20do%20not%20contain%20any%20expansion%20coefficients%2C%20only%20the%20basis%20functions%20are%20shown.%0A%0A%20%20%20%20JAXFunctions%20are%20complete%20functions%2C%20containing%20also%20the%20expansion%20coefficients.%20When%20evaluated%2C%20the%20JAXFunction%20returns%20the%20product%20of%20the%20expansion%20coefficients%20and%20a%20basis%20function.%0A%20%20%20%20%22%22%22)%0A%20%20%20%20return%0A%0A%0A%40app.cell%0Adef%20_(N%2C%20V%2C%20display)%3A%0A%20%20%20%20import%20jax.numpy%20as%20jnp%0A%20%20%20%20from%20jaxfun.galerkin%20import%20JAXFunction%0A%0A%20%20%20%20a%20%3D%20JAXFunction(jnp.ones(N)%2C%20V%2C%20name%3D%22a%22)%0A%20%20%20%20display(a)%0A%20%20%20%20display(a.doit())%0A%20%20%20%20return%20JAXFunction%2C%20jnp%0A%0A%0A%40app.cell%0Adef%20_(JAXFunction%2C%20N%2C%20T%2C%20display%2C%20jnp)%3A%0A%20%20%20%20b%20%3D%20JAXFunction(jnp.ones((N%2C%20N))%2C%20T%2C%20name%3D%22b%22)%0A%20%20%20%20display(b)%0A%20%20%20%20display(b.doit())%0A%20%20%20%20return%0A%0A%0A%40app.cell%0Adef%20_(JAXFunction%2C%20N%2C%20W%2C%20display%2C%20jnp)%3A%0A%20%20%20%20c%20%3D%20JAXFunction(jnp.ones((2%2C%20N%2C%20N))%2C%20W%2C%20name%3D%22c%22)%0A%20%20%20%20display(c)%0A%20%20%20%20display(c.doit())%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
c5c75a2a215c310dad02a07a42a5dc5d