Deep dive: evaluate state using sparse matrices

In this notebook we solve the model via the finite element method using FEniCSx.

import warnings

import numpy as np
import plotly.graph_objects as go
import plotly.io as pio
import ufl
from basix.ufl import element
from dolfinx import default_scalar_type, fem, mesh
from mpi4py import MPI
from plotly.subplots import make_subplots
from scipy.sparse import SparseEfficiencyWarning, bmat
from scipy.sparse.linalg import spsolve

warnings.simplefilter("ignore", SparseEfficiencyWarning)

The following line allows the plot at the end

pio.renderers.default = "notebook_connected"

Introduction

As in the previous notebook, we fix some of the physical parameters arbitrarily.

T = 2.0  # final time

x_0 = 0.0  # left point of the 1D interval
x_L = 1.0  # right point of the 1D interval

def kappa_1(x):
    """First diffusion coefficient."""
    return np.ones_like(x[0])

def kappa_2(x):
    """Second diffusion coefficient."""
    return np.ones_like(x[0])

def y0(x):
    """Initial value."""
    return 5 * np.ones_like(x[0])

mu = (3, 3, 3, 3)  # control vector

def u_func(t):
    """Input function."""
    return -3 * (t <= 4 / 3) + 3 * (t > 4 / 3) * (t <= 2.0)

def f(y, q):
    """Nonlinear coupling function."""
    return np.sqrt(y) * np.sinh(q)

Let us write again a weak formulation of the model:

(1)\[\begin{split}\begin{aligned} \frac{\mathrm{d}}{\mathrm{d}t} \langle y (t), \varphi \rangle_H + \mu_1 a_1 (y(t), \varphi) - \mu_2 \langle f(y(t), q(t)), \varphi \rangle_{H} &= 0 \qquad \forall \varphi \in V, \quad t>0 \\ y(0) &= y_\circ \qquad \text{in } H, \\ \mu_3 a_2 (q(t), \psi) + \mu_4 \langle f(y(t), q(t)), \psi \rangle_{H} - \langle b(t), \psi \rangle_{V'_\circ, V_\circ} &= 0 \qquad \forall \psi \in V_\circ, \quad t>0. \end{aligned}\end{split}\]

Discretization

First, we discretize the space \(\Omega\) on 201 nodes between x_0 and x_L:

N_nodes_x = 201
Omega = mesh.create_interval(comm=MPI.COMM_WORLD, nx=N_nodes_x - 1, points=[x_0, x_L])

Given approximated spaces \(V^h = \mathrm{span} \, \{ \varphi_1, \dots, \varphi_{N_\mathsf{y}} \} \subset V\) and \(V_\circ^h = \mathrm{span} \, \{ \psi_1, \dots, \psi_{N_\mathsf{y}} \} \subset V_\circ\), the finite element (FE) method is based on the Galerkin projection, i.e. finding approximated solutions \(y^h\) of \(y\) and \(q^h\) of \(q\) of the form

\[\begin{equation*} y^h (t) = \sum_{i=1}^{N_\mathsf{y}} y^h_i (t) \varphi_i(x) \qquad \text{and} \qquad q^h (t) = \sum_{i=1}^{N_\mathsf{q}} q^h_i (t) \psi_i(x) \end{equation*}\]

solving the system (1) on the discretized spaces \(V^h\) and \(V_\circ^h\). This is:

(2)\[\begin{split}\begin{aligned} \frac{\mathrm{d}}{\mathrm{d}t} \langle y^h (t), \varphi_i \rangle_H + \mu_1 a_1 (y^h (t), \varphi_i) - \mu_2 \langle f(y^h(t), q^h(t)), \varphi_i \rangle_{H} &= 0 \qquad i=1,\dots,N_\mathsf{y}, \quad t>0, \\ \langle y^h(0) - y_\circ, \varphi_i \rangle_H &= 0 \qquad i=1,\dots,N_\mathsf{y}, \\ \mu_3 a_2 (q^h (t), \psi_i) + \mu_4 \langle f(y^h(t), q^h(t)), \psi_i \rangle_{H} - \langle b(t), \psi_i \rangle_{V'_\circ, V_\circ} &= 0 \qquad i=1,\dots,N_\mathsf{q}, \quad t>0. \end{aligned}\end{split}\]

We use in this notebook linear Lagrangian elements (P1).

FE_degree = 1
N_dofs = (N_nodes_x - 1) * FE_degree + 1
P1 = element("Lagrange", Omega.basix_cell(), FE_degree)
V_h = fem.functionspace(Omega, P1)

Then, we define the boundary conditions (that only apply to the state \(q\)).

left_boundary_dofs = fem.locate_dofs_geometrical(V_h, lambda x: np.isclose(x[0], x_0))
bc_left = fem.dirichletbc(default_scalar_type(x_0), left_boundary_dofs, V_h)
bcs = [bc_left]

In this formulation, the space V_h will be the same for \(y\) and \(q\). What differentiate them is only the boundary conditions, which we will apply for all the arrays appearing in the equation for \(q\).

We prepare the finite element (bilinear and linear) forms:

phi = ufl.TestFunction(V_h)
v = ufl.TrialFunction(V_h)

k_1 = fem.Function(V_h)
k_1.interpolate(kappa_1)
k_2 = fem.Function(V_h)
k_2.interpolate(kappa_2)

# stiffness matrices
a_1 = fem.form(k_1 * v.dx(0) * phi.dx(0) * ufl.dx)
a_2 = fem.form(k_2 * v.dx(0) * phi.dx(0) * ufl.dx)

# mass matrix
m = fem.form(v * phi * ufl.dx)

# initial value
y_circ_f = fem.Function(V_h)
y_circ_f.interpolate(y0)
y_circ_h_form = fem.form(y_circ_f * phi * ufl.dx)

# Neumann boundary
B = fem.form(phi * ufl.ds)
b_h = fem.assemble_vector(B).array
bc_left.set(b_h)  # we need to apply the Dirichlet BC to the RHS of q

We can define the coordinate vectors

\[\begin{equation*} \mathbf{y}^h (t) = \left[ \begin{array}{c} y^h_1 (t) \\ \vdots \\ y^h_{N_\mathsf{y}} (t) \end{array}\begin{array}{c} y^h_1 (t) \\ \vdots \\ y^h_{N_\mathsf{y}} (t) \end{array} \right] \in \mathbb{R}^{N_\mathsf{y}} \qquad \text{and} \qquad \mathbf{q}^h (t) = \left[ \begin{array}{c} q^h_1 (t) \\ \vdots \\ q^h_{N_\mathsf{q}} (t) \end{array}\begin{array}{c} q^h_1 (t) \\ \vdots \\ q^h_{N_\mathsf{q}} (t) \end{array} \right] \in \mathbb{R}^{N_\mathsf{q}} \end{equation*}\]

At this point, we can approximate (2) in matrix form as:

(3)\[\begin{split}\begin{aligned} \mathrm{M}_\mathsf{y} \dot{\mathbf{y}}^h (t) + \mu_1 \mathrm{A}_\mathsf{y} \mathbf{y}^h (t) - \mu_2 \mathrm{M}_\mathsf{y} f \left( \mathbf{y}^h (t), \mathbf{q}^h (t) \right) &= 0 \qquad \text{for } t>0, \\ \mathrm{M}_\mathsf{y} \mathbf{y}^h (0) &= \mathbf{y}^h_\circ, \\ \mu_3 \mathrm{A}_\mathsf{q} \mathbf{q}^h (t) + \mu_4 \mathrm{M}_\mathsf{q} f(\mathbf{y}^h (t), \mathbf{q}^h (t)) &= \mathbf{b}^h (t) \qquad \text{for } \quad t>0, \end{aligned}\end{split}\]

where:

\[\begin{split}& \left( \mathrm{M}_\mathsf{y} \right)_{i,j} = \langle \varphi_j, \varphi_i \rangle_H = \int_\Omega \varphi_j \varphi_i \mathrm{d}x, \\ & \left( \mathrm{A}_\mathsf{y} \right)_{i,j} = a^1_{\boldsymbol\mu} (\varphi_j, \varphi_i) = \int_\Omega \kappa_1 \varphi'_j \varphi'_i \mathrm{d}x, \\ & \left( \mathrm{A}_\mathsf{q} \right)_{i,j} = a^2_{\boldsymbol\mu} (\psi_j, \psi_i) = \int_\Omega \kappa_2 \psi'_j \psi'_i \mathrm{d}x, \\ & \left( \mathbf{y}^h_\circ \right)_i = \langle y_\circ, \varphi_i \rangle_H = \int_\Omega y_\circ \varphi_i \mathrm{d}x, \\ & \left( \mathbf{b}^h (t) \right)_i = \langle b(t), \psi_i \rangle_{V'_\circ, V_\circ},\end{split}\]

and the function \(f\) is evaluated component by component:

\[\left[ f \left( \mathbf{y}^h (t), \mathbf{q}^h (t) \right) \right]_i = f \left( y^h_i (t), q^h_i (t) \right) = \sqrt{y^h_i (t)}\sinh{q^h_i (t)}.\]

We assemble these FE matrices (only once!)

M_y = fem.assemble_matrix(m).to_scipy()
M_q = fem.assemble_matrix(m, bcs=bcs).to_scipy()
A_y = fem.assemble_matrix(a_1).to_scipy()
A_q = fem.assemble_matrix(a_2, bcs=bcs).to_scipy()
y_circ_h = fem.assemble_vector(y_circ_h_form).array
y_0_h = spsolve(M_y, y_circ_h)

To solve (3) we will next discretize in time and then solve the nonlinear problem using Newton’s method. We operate on a uniform time grid \(t_k = k \delta\) for \(k=0, \dots, K\), and \(\delta = T/K\), and we interpret \(\mathbf{y}^{h,k}\) the solution at time \(t_k\) of the fully approximated finite-dimensional system.

K = 200
delta = T / K
t_array = np.linspace(0, T, K + 1)
u = u_func(t_array)

In particular, implicit Euler (IE) discretization in time gives:

(4)\[\begin{split}\begin{aligned} \left( \mathrm{M}_\mathsf{y} + \mu_1 \delta \mathrm{A}_\mathsf{y} \right) \mathbf{y}^{h,k} - \mu_2 \delta \, \mathrm{M}_\mathsf{y} f \left( \mathbf{y}^{h,k}, \mathbf{q}^{h,k} \right) &= \mathrm{M}_\mathsf{y} \mathbf{y}^{h,k-1} \qquad \text{for } k = 1, \dots, K, \\ \mathrm{M}_\mathsf{y} \mathbf{y}^{h,0} &= \mathbf{y}^h_\circ, \\ \mu_3 \mathrm{A}_\mathsf{q} \mathbf{q}^{h,k} + \mu_4 \mathrm{M}_\mathsf{q} f(\mathbf{y}^{h,k}, \mathbf{q}^{h,k}) &= \mathbf{b}^{h,k} \qquad \text{for } k = 1, \dots, K, \end{aligned}\end{split}\]

Let us define the coupled variable

\[\begin{split}\mathbf{z}^{h,k} = \left( \begin{array}{c} \mathbf{y}^{h,k} \\ \mathbf{q}^{h,k} \end{array} \right) \in \mathbb{R}^{N_\mathsf{y} + N_\mathsf{q}}\end{split}\]

and the functional \(\mathrm{F}^k : \mathbb{R}^{N_\mathsf{y} + N_\mathsf{q}} \to \mathbb{R}^{N_\mathsf{y} + N_\mathsf{q}}\), defined as

\[\begin{split}\mathrm{F}^k \left( \mathbf{z}^{h,k} \right) = \left( \begin{array}{c} \mathrm{F}^k_\mathsf{y} \left( \mathbf{z}^{h,k} \right) \\ \mathrm{F}^k_\mathsf{q} \left( \mathbf{z}^{h,k} \right) \end{array} \right),\end{split}\]

with

\[\begin{split}& \mathrm{F}^k_\mathsf{y} \left( \mathbf{z}^{h,k} \right) = \left( \mathrm{M}_\mathsf{y} + \mu_1 \delta \mathrm{A}_\mathsf{y} \right) \mathbf{y}^{h,k} - \mu_2 \delta \, \mathrm{M}_\mathsf{y} f \left( \mathbf{y}^{h,k}, \mathbf{q}^{h,k} \right) - \mathrm{M}_\mathsf{y} \mathbf{y}^{h,k-1}, \\ & \mathrm{F}^k_\mathsf{q} \left( \mathbf{z}^{h,k} \right) = \mu_3 \mathrm{A}_\mathsf{q} \mathbf{q}^{h,k} + \mu_4 \mathrm{M}_\mathsf{q} f(\mathbf{y}^{h,k}, \mathbf{q}^{h,k}) - \mathbf{b}^{h,k}.\end{split}\]

We use Newton’s method for solving \(\mathrm{F}^k (\mathbf{z}^{h,k}) = 0\) at each time step \(k\). It works as follow:

  • Start from \(n=0\) and some initial value \(\mathbf{z}_n\).

  • Evaluate \(\mathbf{d}\) solution of \(\mathrm{J}_{\mathrm{F}^k} (\mathbf{z}_n) \mathbf{d} = - \mathrm{F}^k(\mathbf{z_n})\).

  • Generate new point \(\mathbf{z}_{n+1} = \mathbf{z}_n + \mathbf{d}\) and \(n=n+1\).

  • If \(\mathrm{F}^k(\mathbf{z}_{n})\) is smaller than some tolerance \(\tau\), then we set \(\mathbf{z}^{h,k} := \mathbf{z}_{n}\).

Additionally, at iteration \(k\) we set the previous iterate \(\mathbf{z}^{h,k-1}\) as the initial point \(\mathbf{z}_0\) and we use the variation “damped Newton method” to guarantee monotonicity in the residual decrease.

First, let us define some useful variables and functions:

mat = M_y + mu[0] * delta * A_y
i_bc = bc_left.dof_indices()[0] + N_dofs
Theta = []
for n in range(201):
    v2 = fem.Function(V_h)
    v2.x.array[n] = 1
    v2.x.scatter_forward()
    Tyy = fem.assemble_matrix(fem.form(v * v2 * phi * ufl.dx)).to_scipy().copy()
    Tyy.eliminate_zeros()
    Theta.append(Tyy)
lower_bound = np.concatenate((1e-18 * np.ones_like(y_0_h), -np.inf * np.ones_like(b_h)))


def Newton_F(z, rhs):
    """Function F to minimize."""
    y, q = np.split(z, 2)
    F_y = mat.dot(y) - mu[1] * delta * M_y.dot(f(y, q))
    F_q = mu[2] * A_q.dot(q) + mu[3] * M_q.dot(f(y, q))
    F = np.concatenate((F_y, F_q)) - rhs
    return F


def Newton_J_F(z, *argv, **kwargs):
    """Jacobian of function F."""
    y, q = np.split(z, 2)
    dyf = 0.5 * np.sinh(q) / np.sqrt(y)
    dqf = np.sqrt(y) * np.cosh(q)
    dy = np.tensordot(dyf, Theta, 1).item()
    dq = np.tensordot(dqf, Theta, 1).item()
    J_F = bmat(
        [
            [mat - mu[1] * delta * dy, -mu[1] * delta * dq],
            [mu[3] * dy, mu[2] * A_q + mu[3] * dq],
        ],
        format="csr",
    )
    J_F[i_bc, :] = 0
    J_F[:, i_bc] = 0
    J_F[i_bc, i_bc] = 1
    J_F.eliminate_zeros()
    return J_F

Next, we define the function to solve \(F^k=0\) using Newton’s method.

def solve_Newton(z_0, rhs_F, tol=1e-8):
    """Solve Newton's method for the root finding."""
    z = z_0.copy()
    F = Newton_F(z, rhs_F)
    J_F = Newton_J_F(z)
    err = F.dot(F)
    it = 0
    tau = 0.5
    while err > tol:
        it += 1
        p = spsolve(J_F, -F)
        derr = J_F.T.dot(F).dot(p)
        it_damped = 0
        alpha = 1
        for it_damped in range(20):
            z_new = z + alpha * p
            if np.all(z_new >= lower_bound):
                F_new = Newton_F(z_new, rhs_F)
                err_new = F_new.dot(F_new)
                if err_new - err <= +alpha * 1e-4 * derr:
                    err = err_new
                    F = F_new
                    z = z_new
                    # print(f"{it=}, {it_damped=}, {err=}")
                    break
            alpha *= tau
            it_damped += 1
        else:
            # print(f"{it=}, {it_damped=}, {err=}, {err_new=}")
            raise RuntimeError
    return z

The overall algorithm to find \(\left \{ \mathbf{y}^{h,k}, \mathbf{q}^{h,k} \right \}_{k=1, \dots, K}\) is:

  • Evaluate \(\mathbf{y}^{h,0}\) by projecting the intial value \(y_\circ\) onto the solution space.

  • Evaluate a consistent initial condition for \(\mathbf{q}^{h,0}\) (this step is optional).

  • Use Newton’s method to evaluate \(\mathbf{z}^{h,k}\) for \(k=1,\dots, K\).

We collect data into the matrices Y^h and Q^h:

\[\begin{equation*} \mathrm{Y}^h = \left[ \begin{array}{c|c|c} \mathbf{y}^{h,1} & \dots & \mathbf{y}^{h,K} \end{array} \right] \qquad \mathrm{Q}^h = \left[ \begin{array}{c|c|c} \mathbf{q}^{h,1} & \dots & \mathbf{q}^{h,K} \end{array} \right] \end{equation*}\]

These matrices will contain the columns of the coordinate arrays (in the FE basis).

%%time
Y_h = np.zeros((N_dofs, K))
Q_h = np.zeros((N_dofs, K))
y = y_0_h.copy()
q = spsolve(A_q, u[1] * b_h)
z = np.concatenate((y, q))
for k in range(1, K):
    rhs_F = np.concatenate((M_y.dot(y), u[k] * b_h))
    z_new = solve_Newton(z, rhs_F)
    y, q = np.split(z_new, 2)
    Y_h[:, k - 1] = y
    Q_h[:, k - 1] = q
    z = z_new

CPU times: user 5.61 s, sys: 5.2 ms, total: 5.62 s
Wall time: 5.66 s

Note

Note that our Newton algorithm is significantly faster than SciPy’s root optimizer (scipy.optimizer.root) while reaching almost identical accuracy. While the cell above takes 5.5 seconds, the following code executes in about 23 seconds.

Y_h = np.zeros((N_dofs, K))
Q_h = np.zeros((N_dofs, K))
y = y_0_h.copy()
q = spsolve(A_q, u[1] * b_h)
z = np.concatenate((y, q))
for k in range(1, K):
    rhs_F = np.concatenate((M_y.dot(y), u[k] * b_h))
    z_new = optimize.root(Newton_F, z, args=(rhs_F), jac=lambda z, _: Newton_J_F(z).toarray()).x
    y, q = np.split(z_new, 2)
    Y_h[:,k-1] = y
    Q_h[:,k-1] = q
    z = z_new

def residual(Y, Q):
    """Evaluate the residual of the solution."""
    res = 0
    w = delta * np.ones(K)
    w[0] = delta / 2
    w[-1] = delta / 2
    for k in range(1, K):
        y = Y[:, k - 1]
        q = Q[:, k - 1]
        z = np.concatenate((y, q))
        rhs_F = np.concatenate((M_y.dot(y), u[k] * b_h))
        res_k = Newton_F(z, rhs_F)
        res_k_y, res_k_q = np.split(res_k, 2)
        res += w[k - 1] * (
            res_k_y.dot(M_y.dot(res_k_y)) + res_k_q.dot(M_q.dot(res_k_q))
        )
    return res


residual(Y_h, Q_h)

np.float64(1.2955134623224107e-08)

fig = go.Figure()

# Array of x values
x = Omega.geometry.x[:, 0]

# Create subplots
fig = make_subplots(
    rows=1,
    cols=2,
    column_widths=[0.5, 0.5],
)

# Add all $y^k$ and $q^k$ functions to trace
for k, t in enumerate(t_array[1:]):
    fig.add_traces(
        [
            go.Scatter(
                x=x,
                y=Y_h[:, k],
                visible=False,
                name=f"y(t={t:0.3f})",
            ),
            go.Scatter(
                x=x,
                y=Q_h[:, k],
                visible=False,
                name=f"q(t={t:0.3f})",
            ),
        ],
        rows=[1, 1],
        cols=[1, 2],
    )

# Make t=0 visible
fig.data[0].visible = True
fig.data[1].visible = True

# Create and add slider
steps = []
for i in range(K - 1):
    step = dict(
        method="update",
        args=[
            {"visible": [False] * len(fig.data)},
            {"title": "Slider switched to step: " + str(i)},
        ],  # layout attribute
    )
    step["args"][0]["visible"][2 * i] = True  # Toggle i'th trace to "visible"
    step["args"][0]["visible"][2 * i + 1] = True  # Toggle i'th trace to "visible"
    steps.append(step)

sliders = [
    dict(active=0, currentvalue={"prefix": "Time: "}, pad={"t": 50}, steps=steps)
]

# Add slider to figure
fig.update_layout(
    sliders=sliders,
    yaxis={"range": [0, 7]},
    legend={"x": 0.5, "y": 1.3, "xanchor": "center", "orientation": "h"},
)

# Show final interactive plot
fig.show()