Poisson problem described as a potential

Download this notebook

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

[ ]: