# Copyright (C) 2020-2021 Jørgen S. Dokken
#
# This file is part of DOLFINX_MPC
#
# SPDX-License-Identifier: MIT
from __future__ import annotations
from collections.abc import Sequence
from typing import Optional, Union
from petsc4py import PETSc as _PETSc
import dolfinx.cpp as _cpp
import dolfinx.fem as _fem
from dolfinx_mpc import cpp
from .multipointconstraint import MultiPointConstraint
[docs]
def assemble_matrix(
form: _fem.Form,
constraint: Union[MultiPointConstraint, Sequence[MultiPointConstraint]],
bcs: Optional[Sequence[_fem.DirichletBC]] = None,
diagval: _PETSc.ScalarType = 1, # type: ignore
A: Optional[_PETSc.Mat] = None, # type: ignore
) -> _PETSc.Mat: # type: ignore
"""
Assemble a compiled DOLFINx bilinear form into a PETSc matrix with corresponding multi point constraints
and Dirichlet boundary conditions.
Args:
form: The compiled bilinear variational form
constraint: The multi point constraint
bcs: Sequence of Dirichlet boundary conditions
diagval: Value to set on the diagonal of the matrix
A: PETSc matrix to assemble into
Returns:
_PETSc.Mat: The matrix with the assembled bi-linear form #type: ignore
"""
bcs = [] if bcs is None else [bc._cpp_object for bc in bcs]
if not isinstance(constraint, Sequence):
assert form.function_spaces[0] == form.function_spaces[1]
constraint = (constraint, constraint)
# Generate matrix with MPC sparsity pattern
if A is None:
A = cpp.mpc.create_matrix(form._cpp_object, constraint[0]._cpp_object, constraint[1]._cpp_object)
A.zeroEntries()
# Assemble matrix in C++
cpp.mpc.assemble_matrix(A, form._cpp_object, constraint[0]._cpp_object, constraint[1]._cpp_object, bcs, diagval)
# Add one on diagonal for Dirichlet boundary conditions
if form.function_spaces[0] is form.function_spaces[1]:
A.assemblyBegin(_PETSc.Mat.AssemblyType.FLUSH) # type: ignore
A.assemblyEnd(_PETSc.Mat.AssemblyType.FLUSH) # type: ignore
_cpp.fem.petsc.insert_diagonal(A, form.function_spaces[0], bcs, diagval)
A.assemble()
return A
[docs]
def create_sparsity_pattern(form: _fem.Form, mpc: Union[MultiPointConstraint, Sequence[MultiPointConstraint]]):
"""
Create sparsity-pattern for MPC given a compiled DOLFINx form
Args:
form: The form
mpc: For square forms, the MPC. For rectangular forms a list of 2 MPCs on
axis 0 & 1, respectively
"""
if isinstance(mpc, Sequence):
assert len(mpc) == 2
for mpc_ in mpc:
mpc_._not_finalized() # type: ignore
return cpp.mpc.create_sparsity_pattern(form._cpp_object, mpc[0]._cpp_object, mpc[1]._cpp_object)
else:
mpc._not_finalized() # type: ignore
return cpp.mpc.create_sparsity_pattern(
form._cpp_object,
mpc._cpp_object, # type: ignore
mpc._cpp_object, # type: ignore
) # type: ignore
[docs]
def create_matrix_nest(a: Sequence[Sequence[_fem.Form]], constraints: Sequence[MultiPointConstraint]):
"""
Create a PETSc matrix of type "nest" with appropriate sparsity pattern
given the provided multi points constraints
Args:
a: The compiled bilinear variational form provided in a rank 2 list
constraints: An ordered list of multi point constraints
"""
assert len(constraints) == len(a)
A_ = [[None for _ in range(len(a[0]))] for _ in range(len(a))]
for i, a_row in enumerate(a):
for j, a_block in enumerate(a_row):
if a[i][j] is None:
continue
A_[i][j] = cpp.mpc.create_matrix(
a[i][j]._cpp_object, constraints[i]._cpp_object, constraints[j]._cpp_object
)
A = _PETSc.Mat().createNest( # type: ignore
A_, comm=constraints[0].function_space.mesh.comm
)
return A
[docs]
def assemble_matrix_nest(
A: _PETSc.Mat, # type: ignore
a: Sequence[Sequence[_fem.Form]],
constraints: Sequence[MultiPointConstraint],
bcs: Sequence[_fem.DirichletBC] = [],
diagval: _PETSc.ScalarType = 1, # type: ignore
):
"""
Assemble a compiled DOLFINx bilinear form into a PETSc matrix of type
"nest" with corresponding multi point constraints and Dirichlet boundary
conditions.
Args:
a: The compiled bilinear variational form provided in a rank 2 list
constraints: An ordered list of multi point constraints
bcs: Sequence of Dirichlet boundary conditions
diagval: Value to set on the diagonal of the matrix (Default 1)
A: PETSc matrix to assemble into
"""
for i, a_row in enumerate(a):
for j, a_block in enumerate(a_row):
if a_block is not None:
Asub = A.getNestSubMatrix(i, j)
assemble_matrix(a_block, (constraints[i], constraints[j]), bcs=bcs, diagval=diagval, A=Asub)