Possible reuse of LU decomposition when solving linear elasticity problem in same mesh with diferent bcs

General
1

Hi,

I would like to optimize a linear elasticity code to run faster by reusing the LU decomposition used by PETSc.

Below I have a version of my code that runs using the function LinearProblem(), and another version where I assemble the Lu=b system.

  1. What I noticed is that to assemble A using assemble_matrix() you need the boundary conditions, why is this? Is L modified? Once I can move the assembly of L outside the loop, how could I reuse the LU decomposition from PETSc?
  2. In the function solve_fea() there is also a step where I fix the displacements along some directions. From what I remember this means the solution is also Lu=b, but with the rows and columns corresponding to the degrees of freedom are removed. I think this also means I have to remove those rows and columns from the decomposition, and then map them back the full problem. Is this correct? How could I do this in PETSc?

A self contained running code is shown below.

Thank you,

Alex


import numpy as np
import time

from mpi4py import MPI
from petsc4py import PETSc

import ufl
from dolfinx import mesh, fem, io, default_scalar_type
from dolfinx.fem.petsc import assemble_matrix, assemble_vector, apply_lifting, set_bc, LinearProblem


r_loading = 2.0e-3





def in_sphere(x, point, r_range):
    
    return ((x[0]-point[0])**2 + (x[1]-point[1])**2 + (x[2]-point[2])**2) < r_range**2


def mark_boundaries(domain, facet_node,grasp_node,grasp_node_normal):
   
    dim = domain.topology.dim
    fdim = dim - 1
    domain.topology.create_connectivity(fdim, dim)

    def left_facets(xq):
        return in_sphere(xq, facet_node, r_loading)

    def right_facets(xq):
        return in_sphere(xq, grasp_node, r_loading)

    facets_left  = mesh.locate_entities(domain, fdim, left_facets)
    facets_right = mesh.locate_entities(domain, fdim, right_facets)

    facet_indices = np.concatenate([facets_left, facets_right])
    facet_markers = np.concatenate([np.full_like(facets_left,  1, dtype=np.int32),
                                    np.full_like(facets_right, 2, dtype=np.int32)])

    # Sort by facet index
    sort_perm = np.argsort(facet_indices)
    facet_indices = facet_indices[sort_perm]
    facet_markers = facet_markers[sort_perm]

    mt = mesh.meshtags(domain, fdim, facet_indices, facet_markers)
    return mt


def elasticity_parameters(E=10.0e4, nu=0.20):
    """Lamé parameters from E, nu."""
    lambda_ = E * nu  / ((1.0 + nu) * (1.0 - 2.0 * nu))
    mu      = E * 0.5 /  (1.0 + nu)
    
    return lambda_, mu


def epsilon(u):
    return ufl.sym(ufl.grad(u))


def sigma(u, lambda_, mu):
    return lambda_ * ufl.nabla_div(u) * ufl.Identity(len(u)) + 2.0 * mu * epsilon(u)


def build_primal_problem(domain, facet_tags, lambda_, mu, node_normal, lock_num,
                         p_traction=1.0/(np.pi*r_loading**2), body_force=(0.0, 0.0, -3941.0*9.8)):
    
    dim = domain.topology.dim
    fdim = dim - 1
    ds = ufl.Measure("ds", domain=domain, subdomain_data=facet_tags)

    V = fem.functionspace(domain, ("Lagrange", 1, (domain.geometry.dim,)))
    u = ufl.TrialFunction(V)
    v = ufl.TestFunction(V)

    f_vec = fem.Constant(domain, default_scalar_type(body_force))

    # Traction direction: push along +x (change sign if needed)
    T = fem.Constant(domain, -p_traction * default_scalar_type((node_normal[0], node_normal[1], node_normal[2])))

    a = ufl.inner(sigma(u, lambda_, mu), epsilon(v)) * ufl.dx
    L = ufl.dot(f_vec, v) * ufl.dx + ufl.dot(T, v) * ds(2)

    # Dirichlet BC on left jaw (tag 1)
    left_facets = facet_tags.find(1)
    domain.topology.create_connectivity(fdim, dim)
    left_dofs = fem.locate_dofs_topological(V, fdim, left_facets)

    u_D = np.array((0.0, 0.0, 0.0), dtype=default_scalar_type)
    bc_left = fem.dirichletbc(u_D, left_dofs, V)


    # # Dirichlet BC to stop deflection normal to x (lock_num == 1) or y (lock_num == 2) axis   
    # if lock_num == 1:
    #     right_facets = facet_tags.find(2)
    #     right_dofs1   = fem.locate_dofs_topological(V.sub(1), fdim, right_facets)
    #     bc_1 = fem.dirichletbc(0.0, right_dofs1, V.sub(1))
    #     right_dofs2   = fem.locate_dofs_topological(V.sub(2), fdim, right_facets)
    #     bc_2 = fem.dirichletbc(0.0, right_dofs2, V.sub(2))
    #     return V, a, L, [bc_left, bc_1, bc_2]
    # if lock_num == 2:
    #     right_facets = facet_tags.find(2)
    #     right_dofs0   = fem.locate_dofs_topological(V.sub(0), fdim, right_facets)
    #     bc_1 = fem.dirichletbc(0.0, right_dofs0, V.sub(0))
    #     right_dofs2   = fem.locate_dofs_topological(V.sub(2), fdim, right_facets)
    #     bc_2 = fem.dirichletbc(0.0, right_dofs2, V.sub(2))
    #     return V, a, L, [bc_left, bc_1, bc_2]
    
    return V, a, L, [bc_left]




def array_dist_squared(x, p):
    
    return (x[:,0]-p[0])**2 + (x[:,1]-p[1])**2 + (x[:,2]-p[2])**2



def fea_solve(domain, facet_node, grasp_node, grasp_node_normal, lock_num):
 
    lambda_, mu = elasticity_parameters()
    facet_tags  = mark_boundaries(domain, facet_node,grasp_node,grasp_node_normal)

    Vp, a_p, L_p, bcs_p = build_primal_problem(domain, facet_tags, lambda_, mu, grasp_node_normal, lock_num)
    
    problem = LinearProblem(a_p, L_p, bcs=bcs_p,
                            petsc_options_prefix="linear_problem",
                            petsc_options={"ksp_type": "preonly", "pc_type": "lu", "pc_factor_mat_solver_type": "mumps"})
    uh = problem.solve()

       
    
    # Get displacement at node where load is applied
    dof_coords = Vp.tabulate_dof_coordinates()
    distances  = np.linalg.norm(dof_coords - grasp_node, axis=1)
    dof_index  = np.argmin(distances)
    u_node     = np.sqrt(uh.x.array[3*dof_index]**2 + uh.x.array[3*dof_index+1]**2 + uh.x.array[3*dof_index+2]**2)
    
    return domain, uh, u_node



     
# Part FEA
##################################################################
length = 20.0e-3
width  = 10.0e-3

grasp_facet_nodes = np.array([[0.0,    width/2, width/2],
                              [0.0,    width/2, width/2 + 1.0e-3],
                              [0.0,    width/2, width/2 + 2.0e-3]])
  
grasp_nodes2      = np.array([[length, width/2, width/2],
                              [length, width/2, width/2 + 1.0e-3],
                              [length, width/2, width/2 + 2.0e-3]])

grasp_node_normal = np.array([1.0,    0.0,     0.0    ])

cell_size = 0.5e-3

n_l = int(length/cell_size)
n_w = int(width /cell_size)

lock_num = 1


domain = mesh.create_box(MPI.COMM_WORLD, ((0.0, 0.0, 0.0), (length,width,width)), (n_l,n_w,n_w), mesh.CellType.tetrahedron)


for n in range(grasp_facet_nodes.shape[0]):    
    
    t_start = time.perf_counter()
    domain_f, uh, u_node = fea_solve(domain, grasp_facet_nodes[n,:], grasp_nodes2[n,:], grasp_node_normal, lock_num)
    t_end   = time.perf_counter()
    print(f"FEA elapsed time: {t_end  - t_start:.6f} seconds")


    with io.XDMFFile(domain_f.comm, f"max_p_stress_part_loading_{n+1}.xdmf", "w") as xdmf:
        xdmf.write_mesh(domain_f)
        uh.name   = "u"
        xdmf.write_function(uh)
        
    print(u_node)


print('\n')





# Common assembly steps
t_start = time.perf_counter()

lambda_, mu = elasticity_parameters()

dim = domain.topology.dim
fdim = dim - 1

Vp = fem.functionspace(domain, ("Lagrange", 1, (domain.geometry.dim,)))
u  = ufl.TrialFunction(Vp)
v  = ufl.TestFunction(Vp)

body_force=(0.0, 0.0, -3941.0*9.8)
f_vec = fem.Constant(domain, default_scalar_type(body_force))

# Traction direction: push along +x (change sign if needed)
p_traction=1.0/(np.pi*r_loading**2)
T = fem.Constant(domain, -p_traction * default_scalar_type((grasp_node_normal[0], grasp_node_normal[1], grasp_node_normal[2])))

a_p = ufl.inner(sigma(u, lambda_, mu), epsilon(v)) * ufl.dx

domain.topology.create_connectivity(fdim, dim)

u_D   = np.array((0.0, 0.0, 0.0), dtype=default_scalar_type)


t_end   = time.perf_counter()
print(f"Common assembly time: {t_end  - t_start:.6f} seconds")





for n in range(grasp_facet_nodes.shape[0]):    
    
    t_start = time.perf_counter()
    
    
    facet_tags  = mark_boundaries(domain, grasp_facet_nodes[n,:],grasp_nodes2[n,:],grasp_node_normal)
    
    ds = ufl.Measure("ds", domain=domain, subdomain_data=facet_tags)

    L_p = ufl.dot(f_vec, v) * ufl.dx + ufl.dot(T, v) * ds(2)

    # Dirichlet BC on left jaw (tag 1)
    left_facets = facet_tags.find(1)
    left_dofs = fem.locate_dofs_topological(Vp, fdim, left_facets)

    bcs   = fem.dirichletbc(u_D, left_dofs, Vp)
    bcs_p = [bcs]
    
    
    
    """
    Solve linear system A u = b with PETSc LU.
    """
    a_form = fem.form(a_p)
    L_form = fem.form(L_p)

    A = assemble_matrix(a_form, bcs=bcs_p)
    A.assemble()

    b = assemble_vector(L_form)
    b.ghostUpdate(addv=PETSc.InsertMode.ADD, mode=PETSc.ScatterMode.REVERSE)
    set_bc(b, bcs_p)

    uh = fem.Function(Vp)
    solver = PETSc.KSP().create(domain.comm)
    solver.setType(PETSc.KSP.Type.PREONLY)
    solver.getPC().setType(PETSc.PC.Type.LU)
    solver.getPC().setFactorSolverType('mumps')
    
    solver.setOperators(A)
    solver.solve(b, uh.x.petsc_vec)
    uh.x.scatter_forward()
    
    
    
    # Get displacement at node where load is applied
    dof_coords = Vp.tabulate_dof_coordinates()
    distances  = np.linalg.norm(dof_coords - grasp_nodes2[n,:], axis=1)
    dof_index  = np.argmin(distances)
    u_node     = np.sqrt(uh.x.array[3*dof_index]**2 + uh.x.array[3*dof_index+1]**2 + uh.x.array[3*dof_index+2]**2)
    
    
    
    
    
    
   
    t_end   = time.perf_counter()
    print(f"FEA elapsed time: {t_end  - t_start:.6f} seconds")


    with io.XDMFFile(domain_f.comm, f"max_p_stress2_part_loading_{n+1}.xdmf", "w") as xdmf:
        xdmf.write_mesh(domain_f)
        uh.name   = "u"
        xdmf.write_function(uh)


    print(u_node)
2

Here’s a shorter answer without looking too much at your code.

The reason the LU-factorization is not reused in the LinearProblem-solves is that the system matrix is reassembled for each problem.solve()-call, since constants might have been changed etc. Since the matrix is reassembled each time, the LU-factorization is discarded.

You can get around this by subclassing the dolfinx LinearProblem and copy-pasting some code like below. The only change I’ve made is to move the matrix assembly to a different method and comment it out in the solve() method. You then call assemble_matrix() before the first solve and any time the system matrix should be changed.

The reason assemble_matrix() needs the boundary conditions is that the rows corresponding to the boundary conditions are changed to be 1 on the diagonal and 0 elsewhere and then the same is done to the columns corresponding to boundary conditions through lifting. This process is explained very well in this tutorial by @dokken. So, the actual values have no effect on the assemble_matrix()-step, only which dofs are included. Therefore you can do this step as well as defining your solver outside your loop, as long as you do the lifting at each step.

from petsc4py import PETSc

import dolfinx
import dolfinx.fem.petsc

from dolfinx.fem.function import Function as _Function
from typing import Sequence

from dolfinx.fem.bcs import bcs_by_block as _bcs_by_block
from dolfinx.fem.petsc import assemble_matrix, assemble_vector, apply_lifting
from dolfinx.fem.forms import extract_function_spaces as _extract_function_spaces

class MyLinearProblem(dolfinx.fem.petsc.LinearProblem):
    """
    A customized version of :class:`dolfinx.fem.petsc.LinearProblem` that
    exposes the `assemble_matrix` method for assembling the system matrix
    separately from the `solve` method. Saves time when solving multiple
    linear systems with the same system matrix but different right-hand sides.
    """
            
    def assemble_matrix(self) -> None:
        """Assemble the system matrix and preconditioner matrix (if any)."""
        # Assemble lhs
        self.A.zeroEntries()
        assemble_matrix(self.A, self.a, bcs=self.bcs)  # type: ignore[arg-type, misc]
        self.A.assemble()

        # Assemble preconditioner
        if self.P_mat is not None:
            self.P_mat.zeroEntries()
            assemble_matrix(self.P_mat, self.preconditioner, bcs=self.bcs)  # type: ignore[arg-type, misc]
            self.P_mat.assemble()


    def solve(self) -> _Function | Sequence[_Function]:
        """Solve the problem.

        This method updates the solution ``u`` function(s) stored in the
        problem instance.

        Note:
            The user is responsible for asserting convergence of the KSP
            solver e.g. ``problem.solver.getConvergedReason() > 0``.
            Alternatively, pass ``"ksp_error_if_not_converged" : True`` in
            ``petsc_options`` to raise a ``PETScError`` on failure.

        Returns:
            The solution function(s).
        """

        # # Assemble lhs
        # self.A.zeroEntries()
        # assemble_matrix(self.A, self.a, bcs=self.bcs)  # type: ignore[arg-type, misc]
        # self.A.assemble()

        # # Assemble preconditioner
        # if self.P_mat is not None:
        #     self.P_mat.zeroEntries()
        #     assemble_matrix(self.P_mat, self.preconditioner, bcs=self.bcs)  # type: ignore[arg-type, misc]
        #     self.P_mat.assemble()

        # Assemble rhs
        dolfinx.la.petsc._zero_vector(self.b)
        assemble_vector(self.b, self.L)  # type: ignore[arg-type]

        # Apply boundary conditions to the rhs
        if self.bcs is not None:
            if isinstance(self.u, Sequence):  # block or nest
                bcs1 = _bcs_by_block(_extract_function_spaces(self.a, 1), self.bcs)  # type: ignore[arg-type]
                apply_lifting(self.b, self.a, bcs=bcs1)  # type: ignore[arg-type]
                dolfinx.la.petsc._ghost_update(
                    self.b,
                    PETSc.InsertMode.ADD,  # type: ignore[attr-defined]
                    PETSc.ScatterMode.REVERSE,  # type: ignore[attr-defined]
                )
                bcs0 = _bcs_by_block(_extract_function_spaces(self.L), self.bcs)  # type: ignore[arg-type]
                dolfinx.fem.petsc.set_bc(self.b, bcs0)
            else:  # single form
                apply_lifting(self.b, [self.a], bcs=[self.bcs])  # type: ignore[arg-type]
                dolfinx.la.petsc._ghost_update(
                    self.b,
                    PETSc.InsertMode.ADD,  # type: ignore[attr-defined]
                    PETSc.ScatterMode.REVERSE,  # type: ignore[attr-defined]
                )
                for bc in self.bcs:
                    bc.set(self.b.array_w)
        else:
            dolfinx.la.petsc._ghost_update(self.b, PETSc.InsertMode.ADD, PETSc.ScatterMode.REVERSE)  # type: ignore[attr-defined]

        # Solve linear system and update ghost values in the solution
        self.solver.solve(self.b, self.x)
        dolfinx.la.petsc._ghost_update(self.x, PETSc.InsertMode.INSERT, PETSc.ScatterMode.FORWARD)  # type: ignore[attr-defined]
        dolfinx.fem.petsc.assign(self.x, self.u)
        return self.u