# Copyright (C) 2021-2022 Jørgen Schartum Dokken and Connor D. Pierce
#
# This file is part of DOLFINX_MPC
#
# SPDX-License-Identifier: MIT
from __future__ import annotations
__all__ = [
"gather_PETScVector",
"gather_PETScMatrix",
"compare_mpc_lhs",
"compare_mpc_rhs",
"gather_transformation_matrix",
"compare_CSR",
]
from typing import Any
from mpi4py import MPI
from petsc4py import PETSc
import dolfinx.common
import numpy as np
import scipy.sparse
import dolfinx_mpc
def _gather_slaves_global(constraint):
"""
Given a multi point constraint, return slaves for all processors with global dof numbering
"""
imap = constraint.function_space.dofmap.index_map
num_local_slaves = constraint.num_local_slaves
block_size = constraint.function_space.dofmap.index_map_bs
_slaves = constraint.slaves
if num_local_slaves > 0:
slave_blocks = _slaves[:num_local_slaves] // block_size
slave_rems = _slaves[:num_local_slaves] % block_size
glob_slaves = np.asarray(imap.local_to_global(slave_blocks), dtype=np.int64) * block_size + slave_rems
else:
glob_slaves = np.array([], dtype=np.int64)
slaves = np.hstack(MPI.COMM_WORLD.allgather(glob_slaves))
return slaves
[docs]
def gather_constants(constraint, root=0):
"""
Given a multi-point constraint, gather all constants
"""
imap = constraint.index_map()
constants = constraint._cpp_object.constants
l_range = imap.local_range
ranges = MPI.COMM_WORLD.gather(np.asarray(l_range, dtype=np.int64), root=root)
g_consts = MPI.COMM_WORLD.gather(constants[: l_range[1] - l_range[0]], root=root)
if MPI.COMM_WORLD.rank == root:
block_size = constraint.function_space().dofmap.index_map_bs
global_consts = np.zeros(imap.size_global * block_size, dtype=constraint.coefficients()[0].dtype)
for r, vals in zip(ranges, g_consts):
global_consts[r[0] : r[1]] = vals
return global_consts
else:
return
def petsc_to_local_CSR(A: PETSc.Mat, mpc: dolfinx_mpc.MultiPointConstraint): # type: ignore
"""
Convert a PETSc matrix to a local CSR matrix (scipy) including ghost entries
"""
global_indices = np.asarray(mpc.function_space.dofmap.index_map.global_indices(), dtype=PETSc.IntType) # type: ignore
sort_index = np.argsort(global_indices)
is_A = PETSc.IS().createGeneral(global_indices[sort_index]) # type: ignore
A_loc = A.createSubMatrices(is_A)[0]
ai, aj, av = A_loc.getValuesCSR()
A_csr = scipy.sparse.csr_matrix((av, aj, ai))
return A_csr[global_indices[:, None], global_indices]
[docs]
def gather_PETScMatrix(A: PETSc.Mat, root=0) -> scipy.sparse.csr_matrix: # type: ignore
"""
Given a distributed PETSc matrix, gather in on process 'root' in
a scipy CSR matrix
"""
ai, aj, av = A.getValuesCSR()
aj_all = MPI.COMM_WORLD.gather(aj, root=root) # type: ignore
av_all = MPI.COMM_WORLD.gather(av, root=root) # type: ignore
ai_all = MPI.COMM_WORLD.gather(ai, root=root) # type: ignore
if MPI.COMM_WORLD.rank == root:
ai_cum = [0]
for ai in ai_all: # type: ignore
offsets = ai[1:] + ai_cum[-1]
ai_cum.extend(offsets)
return scipy.sparse.csr_matrix((np.hstack(av_all), np.hstack(aj_all), ai_cum), shape=A.getSize()) # type: ignore
[docs]
def gather_PETScVector(vector: PETSc.Vec, root=0) -> np.ndarray: # type: ignore
"""
Gather a PETScVector from different processors on
process 'root' as an numpy array
"""
if vector.handle == 0:
raise RuntimeError("Vector has been destroyed prior to this call")
numpy_vec = np.zeros(vector.size, dtype=vector.array.dtype)
l_min = vector.owner_range[0]
l_max = vector.owner_range[1]
numpy_vec[l_min:l_max] += vector.array.astype(numpy_vec.dtype)
return np.asarray(sum(MPI.COMM_WORLD.allgather(numpy_vec)))
[docs]
def compare_CSR(A: scipy.sparse.csr_matrix, B: scipy.sparse.csr_matrix, atol=1e-10):
"""Compare CSR matrices A and B"""
diff = np.abs(A - B)
assert diff.max() < atol
[docs]
def compare_mpc_lhs(
A_org: PETSc.Mat, # type: ignore
A_mpc: PETSc.Mat, # type: ignore
mpc: dolfinx_mpc.MultiPointConstraint,
root: int = 0,
atol: np.floating[Any] = 5e3 * np.finfo(dolfinx.default_scalar_type).resolution,
):
"""
Compare an unmodified matrix for the problem with the one assembled with a
multi point constraint.
The unmodified matrix is multiplied with K^T A K, where K is the global transformation matrix.
"""
timer = dolfinx.common.Timer("~MPC: Compare matrices")
comm = mpc.V.mesh.comm
V = mpc.V
assert root < comm.size
is_complex = np.issubdtype(mpc.coefficients()[0].dtype, np.complexfloating) # type: ignore
scipy_dtype = np.complex128 if is_complex else np.float64
K = gather_transformation_matrix(mpc, root=root)
A_csr = gather_PETScMatrix(A_org, root=root)
# Get global slaves
glob_slaves = _gather_slaves_global(mpc)
A_mpc_csr = gather_PETScMatrix(A_mpc, root=root)
if MPI.COMM_WORLD.rank == root:
K = K.astype(scipy_dtype)
A_csr = A_csr.astype(scipy_dtype)
KTAK = np.conj(K.T) * A_csr * K
# Remove identity rows of MPC matrix
all_cols = np.arange(V.dofmap.index_map.size_global * V.dofmap.index_map_bs)
# Scipy >=1.17.0 requires special way of doing slicing to avoid memory overflow
# https://github.com/scipy/scipy/issues/24339
free_mask = np.flatnonzero(np.isin(all_cols, glob_slaves, invert=True).astype(np.int32))
mpc_without_slaves = A_mpc_csr[free_mask, :][:, free_mask]
# Compute difference
compare_CSR(KTAK, mpc_without_slaves, atol=atol)
timer.stop()
[docs]
def compare_mpc_rhs(
b_org: PETSc.Vec, # type: ignore
b: PETSc.Vec, # type: ignore
constraint: dolfinx_mpc.MultiPointConstraint,
root: int = 0,
):
"""
Compare an unconstrained RHS with an MPC rhs.
"""
glob_slaves = _gather_slaves_global(constraint)
b_org_np = gather_PETScVector(b_org, root=root)
b_np = gather_PETScVector(b, root=root)
K = gather_transformation_matrix(constraint, root=root)
# constants = gather_constants(constraint)
comm = constraint.V.mesh.comm
if comm.rank == root:
reduced_b = np.conj(K.T) @ b_org_np # - constants for RHS mpc
all_cols = np.arange(constraint.V.dofmap.index_map.size_global * constraint.V.dofmap.index_map_bs)
cols_except_slaves = np.flatnonzero(np.isin(all_cols, glob_slaves, invert=True).astype(np.int32))
assert np.allclose(b_np[glob_slaves], 0)
assert np.allclose(b_np[cols_except_slaves], reduced_b)