polar

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%20Polar%20coordinates%20with%20jaxfun%0A%0A%20%20%20%20Just%20some%20tests%20to%20illustrate%20curvilinear%20coordinates.%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%20numpy%20as%20np%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%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)%0A%20%20%20%20from%20jaxfun.operators%20import%20Cross%2C%20Curl%2C%20Div%2C%20Dot%2C%20Grad%0A%20%20%20%20from%20jaxfun.utils.common%20import%20Domain%0A%0A%20%20%20%20system%20%3D%20%22cylindrical%22%0A%20%20%20%20_r%2C%20_theta%2C%20_z%20%3D%20sp.symbols(%22r%2Ctheta%2Cz%22%2C%20real%3DTrue%2C%20positive%3DTrue)%0A%20%20%20%20if%20system%20%3D%3D%20%22polar%22%3A%0A%20%20%20%20%20%20%20%20C%20%3D%20get_CoordSys(%0A%20%20%20%20%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%20%20%20%20)%0A%20%20%20%20%20%20%20%20R%20%3D%20FunctionSpace(%0A%20%20%20%20%20%20%20%20%20%20%20%2020%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%3DDomain(0.5%2C%201)%2C%0A%20%20%20%20%20%20%20%20%20%20%20%20name%3D%22R%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%20FunctionSpace(%0A%20%20%20%20%20%20%20%20%20%20%20%2020%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%3DDomain(0%2C%20np.pi)%2C%0A%20%20%20%20%20%20%20%20%20%20%20%20name%3D%22T%22%2C%0A%20%20%20%20%20%20%20%20%20%20%20%20fun_str%3D%22psi%22%2C%0A%20%20%20%20%20%20%20%20)%0A%20%20%20%20%20%20%20%20P%20%3D%20TensorProductSpace((R%2C%20T)%2C%20system%3DC%2C%20name%3D%22P%22)%0A%20%20%20%20elif%20system%20%3D%3D%20%22cylindrical%22%3A%0A%20%20%20%20%20%20%20%20%23%20system%20%3D%20'polar'%0A%20%20%20%20%20%20%20%20C%20%3D%20get_CoordSys(%0A%20%20%20%20%20%20%20%20%20%20%20%20%22C%22%2C%0A%20%20%20%20%20%20%20%20%20%20%20%20sp.Lambda((_r%2C%20_theta%2C%20_z)%2C%20(_r%20*%20sp.cos(_theta)%2C%20_r%20*%20sp.sin(_theta)%2C%20_z))%2C%0A%20%20%20%20%20%20%20%20)%0A%20%20%20%20%20%20%20%20R%20%3D%20FunctionSpace(%0A%20%20%20%20%20%20%20%20%20%20%20%2020%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%3DDomain(0.5%2C%201)%2C%0A%20%20%20%20%20%20%20%20%20%20%20%20name%3D%22R%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%20FunctionSpace(%0A%20%20%20%20%20%20%20%20%20%20%20%2020%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%3DDomain(0%2C%20np.pi)%2C%0A%20%20%20%20%20%20%20%20%20%20%20%20name%3D%22T%22%2C%0A%20%20%20%20%20%20%20%20%20%20%20%20fun_str%3D%22psi%22%2C%0A%20%20%20%20%20%20%20%20)%0A%20%20%20%20%20%20%20%20Z%20%3D%20FunctionSpace(%0A%20%20%20%20%20%20%20%20%20%20%20%2020%2C%20Legendre.Legendre%2C%20domain%3DDomain(0%2C%201)%2C%20name%3D%22Z%22%2C%20fun_str%3D%22L%22%0A%20%20%20%20%20%20%20%20)%0A%20%20%20%20%20%20%20%20P%20%3D%20TensorProductSpace((R%2C%20T%2C%20Z)%2C%20system%3DC%2C%20name%3D%22P%22)%0A%20%20%20%20return%20C%2C%20Cross%2C%20Curl%2C%20Div%2C%20Dot%2C%20Grad%2C%20P%2C%20TestFunction%2C%20TrialFunction%2C%20sp%0A%0A%0A%40app.cell%0Adef%20_(Div%2C%20Grad%2C%20P%2C%20TrialFunction)%3A%0A%20%20%20%20f%20%3D%20TrialFunction(P%2C%20name%3D%22f%22)%0A%20%20%20%20du%20%3D%20Div(Grad(f))%0A%20%20%20%20du%0A%20%20%20%20return%20du%2C%20f%0A%0A%0A%40app.cell%0Adef%20_(du)%3A%0A%20%20%20%20du.doit()%0A%20%20%20%20return%0A%0A%0A%40app.cell%0Adef%20_(du)%3A%0A%20%20%20%20from%20sympy%20import%20srepr%0A%0A%20%20%20%20srepr(du)%0A%20%20%20%20return%20(srepr%2C)%0A%0A%0A%40app.cell%0Adef%20_(C%2C%20Div%2C%20Grad%2C%20f%2C%20sp)%3A%0A%20%20%20%20dv%20%3D%20sp.Add.fromiter(C._parent.base_scalars())%20*%20Div(Grad(f))%0A%20%20%20%20dv%0A%20%20%20%20return%20(dv%2C)%0A%0A%0A%40app.cell%0Adef%20_(dv%2C%20srepr)%3A%0A%20%20%20%20srepr(dv)%0A%20%20%20%20return%0A%0A%0A%40app.cell%0Adef%20_(dv)%3A%0A%20%20%20%20dv.doit()%0A%20%20%20%20return%0A%0A%0A%40app.cell%0Adef%20_(Grad%2C%20f)%3A%0A%20%20%20%20Grad(f).doit()%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%20g%20%3D%20ScalarFunction(%22g%22%2C%20C)%0A%20%20%20%20g%0A%20%20%20%20return%20(g%2C)%0A%0A%0A%40app.cell%0Adef%20_(g)%3A%0A%20%20%20%20G%20%3D%20g.doit()%0A%20%20%20%20G%0A%20%20%20%20return%0A%0A%0A%40app.cell%0Adef%20_(Div%2C%20Grad%2C%20g)%3A%0A%20%20%20%20Div(Grad(g))%0A%20%20%20%20return%0A%0A%0A%40app.cell%0Adef%20_(Div%2C%20Grad%2C%20g)%3A%0A%20%20%20%20Div(Grad(g)).doit()%0A%20%20%20%20return%0A%0A%0A%40app.cell%0Adef%20_(P)%3A%0A%20%20%20%20P.tensorname%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%20%23%23%20Test%20operators%20on%20vectors%0A%20%20%20%20%22%22%22)%0A%20%20%20%20return%0A%0A%0A%40app.cell%0Adef%20_(P)%3A%0A%20%20%20%20from%20jaxfun.galerkin%20import%20VectorTensorProductSpace%0A%0A%20%20%20%20V%20%3D%20VectorTensorProductSpace(P%2C%20name%3D%22V%22)%0A%20%20%20%20V.tensorname%0A%20%20%20%20return%20(V%2C)%0A%0A%0A%40app.cell%0Adef%20_(V)%3A%0A%20%20%20%20V.name%0A%20%20%20%20return%0A%0A%0A%40app.cell%0Adef%20_(TestFunction%2C%20V)%3A%0A%20%20%20%20v%20%3D%20TestFunction(V%2C%20name%3D%22v%22)%0A%20%20%20%20v%0A%20%20%20%20return%20(v%2C)%0A%0A%0A%40app.cell%0Adef%20_(v)%3A%0A%20%20%20%20v.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%20Gradient%20of%20vector%3A%0A%20%20%20%20%22%22%22)%0A%20%20%20%20return%0A%0A%0A%40app.cell%0Adef%20_(Grad%2C%20u)%3A%0A%20%20%20%20Grad(u)%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%20the%20gradient%20of%20a%20vector%20%24%5Cboldsymbol%7Bu%7D%24%20is%20written%20as%20%24%5Cnabla%20%5Cboldsymbol%7Bu%7D%24%20since%20%24%5Cnabla%24%20is%20interpreted%20as%20an%20operator%20and%20not%20a%20vector.%20If%20%24%5Cnabla%24%20had%20instead%20been%20interpreted%20as%20a%20vector%2C%20then%20the%20gradient%20of%20%24%5Cboldsymbol%7Bu%7D%24%20would%20be%20written%20as%20%24(%5Cnabla%20%5Cboldsymbol%7Bu%7D)%5ET%24.%0A%20%20%20%20%22%22%22)%0A%20%20%20%20return%0A%0A%0A%40app.cell%0Adef%20_(Grad%2C%20u)%3A%0A%20%20%20%20Grad(u).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%20Dot%20product%3A%0A%20%20%20%20%22%22%22)%0A%20%20%20%20return%0A%0A%0A%40app.cell%0Adef%20_(Dot%2C%20v)%3A%0A%20%20%20%20Dot(v%2C%20v).doit()%0A%20%20%20%20return%0A%0A%0A%40app.cell%0Adef%20_(C%2C%20Cross%2C%20v)%3A%0A%20%20%20%20from%20IPython.display%20import%20display%0A%0A%20%20%20%20if%20C.dims%20%3D%3D%203%3A%0A%20%20%20%20%20%20%20%20y%20%3D%20Cross(v%2C%20v).doit()%0A%20%20%20%20%20%20%20%20display(y)%0A%20%20%20%20return%20(display%2C)%0A%0A%0A%40app.cell%0Adef%20_(C%2C%20Cross%2C%20display%2C%20v)%3A%0A%20%20%20%20if%20C.dims%20%3D%3D%203%3A%0A%20%20%20%20%20%20%20%20z%20%3D%20Cross(v%2C%20C.b_r)%0A%20%20%20%20%20%20%20%20display(z)%0A%20%20%20%20%20%20%20%20display(z.doit())%0A%20%20%20%20return%0A%0A%0A%40app.cell%0Adef%20_(C%2C%20Curl%2C%20display%2C%20v)%3A%0A%20%20%20%20if%20C.dims%20%3D%3D%203%3A%0A%20%20%20%20%20%20%20%20d%20%3D%20Curl(v)%0A%20%20%20%20%20%20%20%20F%20%3D%20C.to_cartesian(d.doit())%0A%20%20%20%20%20%20%20%20display(d)%0A%20%20%20%20%20%20%20%20display(d.doit())%0A%20%20%20%20%20%20%20%20display(F)%0A%20%20%20%20return%0A%0A%0A%40app.cell%0Adef%20_(C)%3A%0A%20%20%20%20from%20jaxfun.galerkin.arguments%20import%20VectorFunction%0A%0A%20%20%20%20h%20%3D%20VectorFunction(%22h%22%2C%20C)%0A%20%20%20%20h%0A%20%20%20%20return%20(h%2C)%0A%0A%0A%40app.cell%0Adef%20_(h)%3A%0A%20%20%20%20H%20%3D%20h.doit()%0A%20%20%20%20H%0A%20%20%20%20return%0A%0A%0A%40app.cell%0Adef%20_(Dot%2C%20TrialFunction%2C%20V%2C%20srepr%2C%20v)%3A%0A%20%20%20%20w%20%3D%20TrialFunction(V%2C%20name%3D%22u%22)%0A%20%20%20%20vw%20%3D%20Dot(w%2C%20v)%0A%20%20%20%20srepr(vw)%0A%20%20%20%20return%20(vw%2C)%0A%0A%0A%40app.cell%0Adef%20_(vw)%3A%0A%20%20%20%20VW%20%3D%20vw.doit()%0A%20%20%20%20VW%0A%20%20%20%20return%20(VW%2C)%0A%0A%0A%40app.cell%0Adef%20_(VW)%3A%0A%20%20%20%20VW.args%0A%20%20%20%20return%0A%0A%0A%40app.cell%0Adef%20_(C%2C%20VW%2C%20sp)%3A%0A%20%20%20%20a%20%3D%20sp.separatevars(VW.args%5B1%5D%2C%20dict%3DTrue%2C%20symbols%3DC._base_scalars)%0A%20%20%20%20return%20(a%2C)%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_(C%2C%20a)%3A%0A%20%20%20%20import%20contextlib%0A%0A%20%20%20%20for%20_r%20in%20C._base_scalars%3A%0A%20%20%20%20%20%20%20%20for%20j%20in%20a%5B_r%5D.args%3A%0A%20%20%20%20%20%20%20%20%20%20%20%20with%20contextlib.suppress(AttributeError)%3A%0A%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20%20print(j.local_index)%0A%20%20%20%20return%0A%0A%0A%40app.cell%0Adef%20_(C)%3A%0A%20%20%20%20from%20flax%20import%20nnx%0A%0A%20%20%20%20from%20jaxfun.operators%20import%20Identity%2C%20Outer%0A%20%20%20%20from%20jaxfun.pinns%20import%20FlaxFunction%2C%20MLPSpace%0A%0A%20%20%20%20W%20%3D%20MLPSpace(%5B8%5D%2C%20dims%3DC.dims%2C%20rank%3D1%2C%20system%3DC%2C%20name%3D%22V%22)%0A%20%20%20%20Q%20%3D%20MLPSpace(%5B8%5D%2C%20dims%3DC.dims%2C%20rank%3D0%2C%20system%3DC%2C%20name%3D%22Q%22)%0A%20%20%20%20u%20%3D%20FlaxFunction(W%2C%20%22u%22%2C%20rngs%3Dnnx.Rngs(2002))%20%20%23%20Vector%20space%20for%20velocity%0A%20%20%20%20p%20%3D%20FlaxFunction(Q%2C%20%22p%22%2C%20rngs%3Dnnx.Rngs(1001))%20%20%23%20Scalar%20space%20for%20pressure%0A%20%20%20%20return%20Identity%2C%20Outer%2C%20p%2C%20u%0A%0A%0A%40app.cell%0Adef%20_(C%2C%20Div%2C%20u)%3A%0A%20%20%20%20C.simplify(Div(u).doit())%0A%20%20%20%20return%0A%0A%0A%40app.cell%0Adef%20_(Div%2C%20Dot%2C%20Grad%2C%20p%2C%20u)%3A%0A%20%20%20%20R1%20%3D%20Dot(Grad(u)%2C%20u)%20-%20Div(Grad(u))%20%2B%20Grad(p)%0A%20%20%20%20R1%0A%20%20%20%20return%20(R1%2C)%0A%0A%0A%40app.cell%0Adef%20_(C%2C%20Div%2C%20Grad%2C%20Identity%2C%20Outer%2C%20p%2C%20sp%2C%20u)%3A%0A%20%20%20%20I%20%3D%20Identity(C)%0A%20%20%20%20R2%20%3D%20Div(Outer(u%2C%20u))%20-%20Div(%0A%20%20%20%20%20%20%20%20Grad(u)%20%2B%20Grad(u).T%20-%20sp.Rational(2%2C%203)%20*%20Div(u)%20*%20I%20%2B%20p%20*%20I%0A%20%20%20%20)%0A%20%20%20%20R2.doit()%0A%20%20%20%20return%0A%0A%0A%40app.cell%0Adef%20_(C%2C%20Cross%2C%20Curl%2C%20Div%2C%20Dot%2C%20Grad%2C%20display%2C%20p%2C%20sp%2C%20u)%3A%0A%20%20%20%20if%20C.dims%20%3D%3D%203%3A%0A%20%20%20%20%20%20%20%20R3%20%3D%20Cross(Curl(u)%2C%20u)%20%2B%20sp.S.Half%20*%20Grad(Dot(u%2C%20u))%20-%20Div(Grad(u))%20%2B%20Grad(p)%0A%20%20%20%20%20%20%20%20display(R3)%0A%20%20%20%20%20%20%20%20display(R3.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%20Note%20that%20the%20divergence%20of%20a%20dyadic%20%24A%24%20is%20defined%20as%20%24%5Ctext%7Bdiv%7D(A)%20%3D%20%5Cnabla%20%5Ccdot%20A%24%2C%20since%20we%20choose%20to%20interpret%20%24%5Cnabla%24%20as%20an%20operator%20and%20not%20a%20vector.%20This%20operator%20acts%20over%20the%20last%20axis%20of%20%24A%24.%20Some%20authors%20define%20%24%5Cnabla%24%20as%20a%20vector%20and%20in%20that%20case%20the%20definition%20of%20the%20divergence%20is%20%24%5Cnabla%20%5Ccdot%20A%5ET%24.%20This%20is%20perfectly%20fine%2C%20just%20a%20difference%20in%20notation.%20With%20index%20notation%20in%20Cartesian%20coordinates%20the%20divergence%20may%20also%20be%20written%20as%20%24%5Cfrac%7B%5Cpartial%20A_%7Bij%7D%7D%7B%5Cpartial%20x_j%7D%24%20or%20%24%5Cfrac%7B%5Cpartial%20A_%7Bij%7D%7D%7B%5Cpartial%20x_j%7D%20%5Cboldsymbol%7Bi%7D_i%24%2C%20where%20%24%5C%7B%5Cboldsymbol%7Bi%7D_j%5C%7D_%7Bj%3D1%7D%5E3%24%20are%20the%203%20Cartesian%20unit%20vectors.%0A%0A%20%20%20%20Hence%20we%20write%20Div(Grad(u))%20as%20%24%5Cnabla%20%5Ccdot%20%5Cnabla%20%5Cboldsymbol%7Bu%7D%24.%0A%20%20%20%20%22%22%22)%0A%20%20%20%20return%0A%0A%0A%40app.cell%0Adef%20_(C%2C%20Dot%2C%20sp%2C%20u)%3A%0A%20%20%20%20ut%20%3D%20Dot(u%2C%20C.get_contravariant_basis_vector(1)).doit()%0A%20%20%20%20sp.srepr(ut)%0A%20%20%20%20return%0A%0A%0A%40app.cell%0Adef%20_(Outer%2C%20u)%3A%0A%20%20%20%20uu%20%3D%20Outer(u%2C%20u).doit()%0A%20%20%20%20return%20(uu%2C)%0A%0A%0A%40app.cell%0Adef%20_(uu)%3A%0A%20%20%20%20uu%0A%20%20%20%20return%0A%0A%0A%40app.cell%0Adef%20_(C%2C%20Dot%2C%20R1)%3A%0A%20%20%20%20C.simplify(Dot(R1%2C%20C.get_contravariant_basis_vector(0)).doit())%0A%20%20%20%20return%0A%0A%0A%40app.cell%0Adef%20_(C%2C%20R1)%3A%0A%20%20%20%20C.simplify(R1.doit())%0A%20%20%20%20return%0A%0A%0A%40app.cell%0Adef%20_(C%2C%20Dot%2C%20R1)%3A%0A%20%20%20%20r1%20%3D%20Dot(R1%2C%20C.get_contravariant_basis_vector(0)).doit()%0A%20%20%20%20r1%0A%20%20%20%20return%20(r1%2C)%0A%0A%0A%40app.cell%0Adef%20_(r1)%3A%0A%20%20%20%20r1.expand()%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