mwe_real_element.py

from mpi4py import MPI
import dolfinx
from dolfinx import log
from dolfinx.fem.petsc import NonlinearProblem
import numpy as np
from scifem import create_real_functionspace
import ufl


M = 8
mesh = dolfinx.mesh.create_unit_square(MPI.COMM_WORLD, M, M, dtype=np.float64)
V = dolfinx.fem.functionspace(mesh, ("Lagrange", 1))
R = create_real_functionspace(mesh)

W = ufl.MixedFunctionSpace(V, R)

tdim = mesh.topology.dim
fdim = tdim - 1
mesh.topology.create_connectivity(fdim, tdim)


def boundary_right(x):
    return np.isclose(x[0], 1)


def boundary_left(x):
    return np.isclose(x[0], 0)


dofs_right = dolfinx.fem.locate_dofs_geometrical(V, boundary_right)
dofs_other = dolfinx.fem.locate_dofs_geometrical(V, boundary_left)

right_facets = dolfinx.mesh.locate_entities_boundary(mesh, fdim, boundary_right)
right_values = np.full_like(right_facets, 1, dtype=np.int32)
other_facets = dolfinx.mesh.locate_entities_boundary(mesh, fdim, boundary_left)
other_values = np.full_like(other_facets, 2, dtype=np.int32)

facet_markers = dolfinx.mesh.meshtags(
    mesh,
    dim=fdim,
    entities=np.concatenate([right_facets, other_facets]),
    values=np.concatenate([right_values, other_values]),
)


u = dolfinx.fem.Function(V)
u_n = dolfinx.fem.Function(V)
pressure = dolfinx.fem.Function(R)
pressure_n = dolfinx.fem.Function(R)

dt = dolfinx.fem.Constant(mesh, 1.0)
K_S = dolfinx.fem.Constant(mesh, 2.0)
e = dolfinx.fem.Constant(mesh, 2.0)
volume = dolfinx.fem.Constant(mesh, 4.0)

v, pressure_v = ufl.TestFunctions(W)

pressure_ini_expr = dolfinx.fem.Expression(
    dolfinx.fem.Constant(mesh, 3.5), R.element.interpolation_points
)
pressure_n.interpolate(pressure_ini_expr)
pressure.x.array[:] = pressure_n.x.array[:]

bc_right_expr = dolfinx.fem.Expression(
    # K_S * pressure**0.5, V.element.interpolation_points
    pressure,
    V.element.interpolation_points,
)
u_bc_right = dolfinx.fem.Function(V)
u_bc_right.interpolate(bc_right_expr)

bc_right = dolfinx.fem.dirichletbc(u_bc_right, dofs_right)
bc_other = dolfinx.fem.dirichletbc(dolfinx.fem.Constant(mesh, 0.0), dofs_other, V)

n = ufl.FacetNormal(mesh)
flux = ufl.inner(ufl.grad(u), n)

ds = ufl.Measure("ds", domain=mesh, subdomain_data=facet_markers)

F = -(u - u_n) / dt * v * ufl.dx + ufl.dot(ufl.grad(u), ufl.grad(v)) * ufl.dx
F += (
    pressure - pressure_n
) / dt * pressure_v * ufl.dx - e / volume * flux * pressure_v * ds(2)

forms = ufl.extract_blocks(F)

solver = NonlinearProblem(
    forms,
    [u, pressure],
    bcs=[bc_other, bc_right],
    petsc_options_prefix="mwe",
    petsc_options={
        "ksp_type": "preonly",
        "pc_type": "lu",
        "pc_factor_mat_solver_type": "mumps",
        "snes_monitor": None,
    },
)

log.set_log_level(log.LogLevel.INFO)
t = 0
t_final = 10
times = []

all_pressures = []
while t < t_final:
    t += dt.value
    times.append(t)
    print(f"Solving at time {t:.2f}")

    # Solve the problem
    (u, pressure) = solver.solve()
    u_n.x.array[:] = u.x.array[:]
    pressure_n.x.array[:] = pressure.x.array[:]
    print(f"u = {u.x.array}")
    print(f"pressure = {pressure.x.array}")

    u_bc_right.interpolate(bc_right_expr)
    all_pressures.append(pressure.x.array[0].copy())

import matplotlib.pyplot as plt

plt.plot(times, all_pressures)
plt.xlabel("Time")
plt.ylabel("Pressure")
plt.show()

import pyvista
from dolfinx import plot


def make_ugrid(solution):
    topology, cell_types, geometry = plot.vtk_mesh(solution.function_space)
    u_grid = pyvista.UnstructuredGrid(topology, cell_types, geometry)
    u_grid.point_data["c"] = solution.x.array.real
    u_grid.set_active_scalars("c")
    return u_grid


pyvista.set_jupyter_backend("html")

u_plotter = pyvista.Plotter()
u_grid = make_ugrid(u)
u_plotter.add_mesh(u_grid, cmap="magma", show_edges=False)
u_plotter.view_xy()

if not pyvista.OFF_SCREEN:
    u_plotter.show()
else:
    figure = u_plotter.screenshot("concentration.png")