Poisson problem described as a potential
Imports
[ ]:
import jax
import jax.numpy as jnp
from autopdex import SimState, dae, spaces
jax.config.update("jax_enable_x64", True)
Weak form
[2]:
def poisson_model(ctx: SimState.ModelContext):
phi_fun = ctx.trial_ansatz["phi"]
phi = phi_fun(ctx.x_int, ctx.t)
x = ctx.trial_ansatz["physical coor"](ctx.x_int)
# ctx.x_int is the integration point coordinate in reference configuration
# "physical coor" is the corresponding coordinate in initial configuration
def source_term(x):
x2 = x - jnp.asarray([1.0, 0.5])
return 20.0 * (
jnp.sin(10.0 * x @ x)
- jnp.cos(10.0 * x2 @ x2)
)
if ctx.mode == "output":
return {
"phi": phi,
"source": source_term(x),
}
grad_phi = jax.jacfwd(phi_fun, 0)(ctx.x_int, ctx.t)
# Although evaluated at ctx.x_int, the gradients are w.r.t. the initial configuration
return 1/2 * grad_phi @ grad_phi - source_term(x) * phi
Solve with SimState
[3]:
vertices = [[0.0, 0.0], [1.0, 0.0], [1.0, 1.0], [0.0, 1.0]]
n_elements = (1000, 1000)
sim = SimState({"phi": spaces.H1(order=1, dim=2, field_dimension=1)})
sim.add_structured_mesh(n_elements, vertices, "quad")
sim.add_temporal_discretization({"phi": dae.NoTimeDerivative()})
sim.add_model("__all__", "potential", poisson_model)
def on_outer_boundary(x):
return jnp.logical_or(
jnp.logical_or(jnp.isclose(x[0], 0.0), jnp.isclose(x[0], 1.0)),
jnp.logical_or(jnp.isclose(x[1], 0.0), jnp.isclose(x[1], 1.0)),
)
sim.add_strong_bc("phi", on_outer_boundary, lambda x, t: 0.0)
sim.set_postprocessing_policy(dae.SaveAllPolicy(), result_folder_name="./short_example_hli.res")
sim.initialize(verbose=1)
sim.prepare()
sim = sim.run(dt0=1.0, time_span=1.0, num_time_steps=1)
Linear solver: Pardiso(lu).
Iteration 1, Residual norm: 4.949433423077815e-17
Progress: 100%, Time: 1.00e+00, dt: 1.00e+00, iterations: 1
[ ]: