Assembling the Mass Matrix in Parallel

dolfinx
1

I’m looking to assemble the mass matrix over a mesh that has been distributed over multiple processes.

In my example problem, I create a 2D rectangular Mesh over 2 processes.

Here is my Minimum Working Example (MWE) code.

# Imports
import dolfinx
import ufl

from mpi4py import MPI

import numpy as np

# Parameters
x0, x1 = (0.0, 1.0)
y0, y1 = (0.0, 1.0)
nx, ny = (2, 2)

# MPI Attributes
comm = MPI.COMM_WORLD
size = comm.Get_size()
rank = comm.Get_rank()

# Mesh

# The 2D rectangular domain is split into a FEM mesh
domain = dolfinx.mesh.create_rectangle(
    comm = comm,
    points = [[x0, y0], [x1, y1]],
    n = [nx, ny],
    cell_type = dolfinx.mesh.CellType.triangle,
    # Describes how the rectangular partitions are divided into triangles
    diagonal = dolfinx.mesh.DiagonalType.right,
    # Determines whether processes know information about adjacent cells on other processes
    ghost_mode = dolfinx.cpp.mesh.GhostMode.none,
    dtype = np.float64
)

# Dimensionality of the space should be 2 dimensional
ddim = domain.topology.dim
    
# Dimensionality of a facet, should be 1 dimensional
fdim = ddim - 1

# This step is important for detecting boundary facets and for proper partitioning of the mesh across multiple processes
domain.topology.create_connectivity(fdim, ddim)

# Mesh Boundaries

# FEniCS does not correctly perform oriented integrals over a closed boundary
# Sub-Meshes on the perimeter of the 2D rectangle are needed in order to perform oriented boundary integrals

# This is the absolute tolerance in determining which facets live on which boundaries
tolerance = np.finfo(np.float64).eps

# These are indicator functions used to mark the boundary facets of the mesh
# FEniCS does not have an elegant way of performing a directional integral on disjoint boundaries like a rectangle
boundary_indicators = [
    (1, lambda x: np.isclose(x[0], x0, 0.0, tolerance)), # Left
    (2, lambda x: np.isclose(x[1], y0, 0.0, tolerance)), # Bottom
    (3, lambda x: np.isclose(x[0], x1, 0.0, tolerance)), # Right
    (4, lambda x: np.isclose(x[1], y1, 0.0, tolerance)) # Top
]

facet_indices = list()

facet_markers = list()

for marker, indicator in boundary_indicators:
    facets = dolfinx.mesh.locate_entities(domain, fdim, indicator)
    
    facet_indices.append(facets)
    facet_markers.append(np.full(facets.shape, marker, dtype = np.int64))
    
facet_indices = np.hstack(facet_indices)
facet_markers = np.hstack(facet_markers)
sorted_indices = np.argsort(facet_indices)

facet_tags = dolfinx.mesh.meshtags(
    domain,
    fdim,
    facet_indices[sorted_indices],
    facet_markers[sorted_indices]
)

# Setup UFL Expressions

# These are specifically UFL Expression objects located as ufl.core.expr.Expr
# They are both immutable, and hashable, objects

# UFL Represenation of the spacial coordinates as a vector
x_coord = ufl.SpatialCoordinate(domain)

# UFL Representation of the time coordinate as a scalar constant over the the entire mesh
t_var = ufl.variable(ufl.Constant(domain, shape = ()))

# UFL Expression of the Manufactured Solution
solution = (1.0 - ufl.exp(-tau * t_var)) * ufl.sin(x_coord[0] * x_coord[1])

# UFL Expression of the Manufactured Source Function
# This is what forces the solution of the 2D Heat Equation
source = tau * ufl.exp(-tau * t_var) * ufl.sin(x_coord[0] * x_coord[1]) + alpha * solution * (x_coord[0] ** 2 + x_coord[1] ** 2)
    
# Define UFL Measure Objects

# We define UFL Measure objects explicitly, instead of using ufl.dx and ufl.ds because we may want to control the quadrature degree in the numerical integration
# This has a direct impact on the computational complexity of integration

# UFL Measure for integrating over the area of the 2D Cells of the Mesh
dx = ufl.Measure(
    integral_type = "cell",
    domain = domain,
    subdomain_id = "everywhere",
)

# UFL Measure for integrating over the 1D Exterior Facet boundary of the Mesh
# This is needed because we are explicitly setting which boundary sub-domains to integrate over
ds = ufl.Measure(
    integral_type = "exterior_facet",
    domain = domain,
    subdomain_data = facet_tags,
)

# Function Space

# The FEniCS FunctionSpace is a generalized wrapper over the UFL FunctionSpace object
# For the 2D Heat equation, the weak form of the PDE has the first order as the highest derivative
# This lets us use first order polynomials as our basis functions
function_space = dolfinx.fem.functionspace(
    domain,
    dolfinx.fem.ElementMetaData(family = "CG", degree = 1)
)
    
# Degrees of Freedom Index Map
dof_imap = function_space.dofmap.index_map

# Number of Degrees of Freedom
n_dofs_local = dof_imap.size_local # Number of local degrees of freedom
n_dofs_global = dof_imap.size_global # Number of global degrees of freedom

test_function = ufl.TestFunction(function_space)

trial_function = ufl.TrialFunction(function_space)

# Construct Mass Matrix
mass_form = dolfinx.fem.form(
    ufl.inner(test_function, trial_function) * dx,
    dtype = np.float64
)

mass_matrix = dolfinx.fem.assemble_matrix(mass_form)
mass_matrix.scatter_reverse()

# We need to cast the dolfinx.la.MatrixCSR matrix into a SciPy Sparse CSR matrix
# SciPy CSR matrices are serializable (can be pickled)
A = mass_matrix.to_scipy()

# Ranges of the Degrees of Freedom on this process
local_range = dof_imap.local_range

ranges = comm.gather(local_range, root = 0)
data = comm.gather(A, root = 0)

If I run the example on one process, I get a 9 by 9 matrix with 41 non-zero elements.

When run with two nodes, I get the following.
Process 0 Matrix – 4 by 8 – 17 non-zero elements
Process 1 Matrix – 5 by 8 – 24 non-zero elements

I can successfully get all the entries of the Mass Matrix to the root process, however I don’t know how to recover the structure of the mass matrix as if the program was run on one process.

Please let me know what I’m missing. Thank You.

2

See for instance: dolfinx_mpc/python/dolfinx_mpc/utils/test.py at main · jorgensd/dolfinx_mpc · GitHub

May I ask why you would want this?

1 Like
3
  • I am attempting to solve the 2D heat equation with a custom ODE solver.
  • The ODE solving algorithm needs to know the complete mass matrix, as well as the complete derivative matrix.
  • On each time marching step of the ODE algorithm, it calls Dolfinx subroutines to solve a series of integrals over the mesh.
  • I would like these integrals to be performed in parallel
  • My ODE algorithm is written in a purely numerical context using JAX
  • While there are discrete time steps in the ODE solver, the output solution has support over the full time domain.
  • UFL does not have a good way of implementing time dependence, so I use ufl.replace() and generate different UFL forms that are then solved with calls to dolfinx.fem functions
4

Can I assume that I can also use the gather_PETScVector() function of Dolfinx MPC to reconstruct vectors on the master process?

5

Yes, you can do that.

1 Like