From 49e10850d5fc92fdea09a2741bc62d66a7b599af Mon Sep 17 00:00:00 2001 From: Pablo Brubeck Date: Sun, 18 Feb 2024 16:48:37 +0000 Subject: [PATCH] docs --- firedrake/preconditioners/fdm.py | 157 +++++++++++++++++++------------ 1 file changed, 96 insertions(+), 61 deletions(-) diff --git a/firedrake/preconditioners/fdm.py b/firedrake/preconditioners/fdm.py index 89c30c3032..ef05ffc23b 100644 --- a/firedrake/preconditioners/fdm.py +++ b/firedrake/preconditioners/fdm.py @@ -380,21 +380,21 @@ def condense(self, A, J, bcs, fcp, pc_type="icc"): The matrix to statically condense. J : ufl.Form The bilinear form to statically condense. - bcs : - An iterable of boundary conditions. - fcp : - A dict with the form compiler parameters. - pc_type : + bcs : .BCBase[] + An iterable of boundary conditions to apply on ``A``. + fcp : dict + The form compiler parameters. + pc_type : PETSc.PC.Type The preconditioner type for the interior solver. Returns ------- - Smat : - A :class:`PETSc.Mat` with the original blocks of A, except that + Smat : PETSc.Mat + A matrix with the original blocks of ``A``, except that the matrix-free Schur complement replaces the interface-interface block. - Pmat : - A `dict` mapping pairs of function spaces to the preconditioner blocks - [[inv(A00), A01], [A10, inv(S)]]. + Pmat : dict + A dict mapping pairs of function spaces to the preconditioner blocks + ``[[inv(A00), A01], [A10, inv(S)]]``. """ Smats = {} V = J.arguments()[0].function_space() @@ -437,7 +437,7 @@ def condense(self, A, J, bcs, fcp, pc_type="icc"): *args_acc, x.dat(op2.READ, x.cell_node_map()), y.dat(op2.INC, y.cell_node_map())) - ctx = ParloopMatrixContext(parloop, x, y, bcs=bcs) + ctx = PythonMatrixContext(parloop, x, y, bcs=bcs) Smats[Vsub, Vsub] = PETSc.Mat().createPython(sizes, context=ctx, comm=comm) if Vsub == V0: Pmats[Vsub, Vsub] = Smats[Vsub, Vsub] @@ -668,15 +668,22 @@ def _element_mass_matrix(self): @PETSc.Log.EventDecorator("FDMSetValues") def set_values(self, A, Vrow, Vcol, addv, mat_type="aij"): - """ - Assemble the auxiliary operator in the FDM basis using sparse reference - tensors and diagonal mass matrices. - - :arg A: the :class:`PETSc.Mat` to assemble - :arg Vrow: the :class:`.FunctionSpace` test space - :arg Vcol: the :class:`.FunctionSpace` trial space - :arg addv: a `PETSc.Mat.InsertMode` - :arg mat_type: the `PETSc.Mat.Type` of the auxiliary operator + """Assemble the auxiliary operator in the FDM basis using sparse + reference tensors and diagonal mass matrices. + + Parameters + ---------- + A : PETSc.Mat + The (initialized) matrix to assemble. + Vrow : .FunctionSpace + The test space. + Vcol : .FunctionSpace + The trial space. + addv : PETSc.Mat.InsertMode + Flag indicating if we want to insert or add matrix values. + mat_type : PETSc.Mat.Type + The matrix type of auxiliary operator. This only used when ``A`` is a preallocator + to determine the nonzeros on the upper triangual part of an ``'sbaij'`` matrix. """ key = (Vrow.ufl_element(), Vcol.ufl_element()) on_diag = Vrow == Vcol @@ -795,7 +802,7 @@ def kernel(self, mat_type="aij", on_diag=False, addv=None): class TripleProductKernel(ElementKernel): - """Kernel builder that inserts a triple product of the form L * C * R, + """Kernel builder to assemble a triple product of the form L * C * R for each cell, where L, C, R are sparse matrices and the entries of C are updated on each cell.""" code = dedent(""" PetscErrorCode %(name)s(const Mat A, const Mat B, @@ -810,8 +817,8 @@ class TripleProductKernel(ElementKernel): PetscFunctionReturn(PETSC_SUCCESS); }""") - def __init__(self, A, B, C, name=None): - self.product = partial(A.matMatMult, B, C) + def __init__(self, L, C, R, name=None): + self.product = partial(L.matMatMult, C, R) super().__init__(self.product(), name=name) @@ -1169,15 +1176,15 @@ def matmult_kernel_code(a, prefix="form", fcp=None, matshell=False): Returns ------- - A 4-tuple of C code strings/string generators, - matmult_struct : + matmult_struct : str The C code to compute the matrix-vector product. - matmult_call : - A lambda to call the matrix-vector product on the first argument, - writing the output on the second argument. - ctx_struct : + matmult_call : callable + - ``x``: the pointer name of the input vector (`str`). + - ``y``: the pointer name of the output vector (`str`). + A lambda to generate the C code calling the matrix-vector product. + ctx_struct : str The signature of the kernel. - ctx_pack : + ctx_pack : str Code to update the coefficient array pointers to be called before applying the matshell. """ @@ -1193,14 +1200,18 @@ def matmult_kernel_code(a, prefix="form", fcp=None, matshell=False): kernel = kernels[-1].kinfo.kernel nargs = len(kernel.arguments) - len(a.arguments()) ncoef = nargs - len(extract_firedrake_constants(F)) + + matmult_struct = cache_generate_code(kernel, V._comm) + matmult_struct = matmult_struct.replace("void "+kernel.name, "static void "+kernel.name) + + ctx_coeff = "".join(f"appctx[{i}], " for i in range(ncoef)) + ctx_const = "".join(f", appctx[{i}]" for i in range(ncoef, nargs)) + matmult_call = lambda x, y: f"{kernel.name}({y}, {ctx_coeff}{x}{ctx_const});" + ctx_struct = "".join(f"const PetscScalar *restrict c{i}, " for i in range(nargs)) ctx_pointers = ", ".join(f"c{i}" for i in range(nargs)) ctx_pack = f"const PetscScalar *appctx[{nargs}] = {{ {ctx_pointers} }};" - ctx_coefficients = "".join("appctx[%d], " % i for i in range(ncoef)) - ctx_constants = "".join(", appctx[%d]" % i for i in range(ncoef, nargs)) - matmult_call = lambda x, y: f"{kernel.name}({y}, {ctx_coefficients}{x}{ctx_constants});" - matmult_struct = cache_generate_code(kernel, V._comm) - matmult_struct = matmult_struct.replace("void "+kernel.name, "static void "+kernel.name) + cache[key] = (matmult_struct, matmult_call, ctx_struct, ctx_pack) if matshell: @@ -1262,24 +1273,20 @@ def __init__(self, kernel, form, name=None, prefix="interior_", fcp=None, pc_typ A.setSizes(B.getSizes()) A.setUp() - use_cg = False - if use_cg: - ksp_type = PETSc.KSP.Type.CG - norm_type = PETSc.KSP.NormType.NATURAL - else: - ksp_type = PETSc.KSP.Type.MINRES - norm_type = PETSc.KSP.NormType.PRECONDITIONED - knoll = False - knoll = True + # Set up the local KSP for the cell interiors ksp = PETSc.KSP().create(comm=comm) ksp.setOptionsPrefix(prefix) ksp.setOperators(A, B) + + # Default solver options, these can be overriden via -interior_ksp_type, etc. + rtol = 1E-8 + atol = 1E-14 + ksp_type = PETSc.KSP.Type.MINRES + norm_type = PETSc.KSP.NormType.PRECONDITIONED ksp.pc.setType(pc_type) ksp.setType(ksp_type) ksp.setNormType(norm_type) - ksp.setTolerances(rtol=1E-8, atol=0) - ksp.setInitialGuessNonzero(knoll) - ksp.setInitialGuessKnoll(knoll) + ksp.setTolerances(rtol=rtol, atol=atol) ksp.setFromOptions() ksp.setUp() super().__init__(ksp, name=name) @@ -1352,12 +1359,18 @@ def __init__(self, kernel, name=None): fdofs = PETSc.IS().createBlock(V.value_size, restricted_dofs(V1.finat_element, V.finat_element), comm=comm) size = idofs.size + fdofs.size assert size == V.finat_element.space_dimension() * V.value_size + # Bilinear form on the space with interior and interface a = form if Q == V else form(*(t.reconstruct(function_space=V) for t in args)) - a00 = form if Q == V0 else form(*(t.reconstruct(function_space=V0) for t in args)) + # Generate code to apply the action of A within the Schur complement action A_struct, A_call, ctx_struct, ctx_pack = matmult_kernel_code(a, prefix="A", fcp=fcp) + + # Bilinear form on the interior + a00 = form if Q == V0 else form(*(t.reconstruct(function_space=V0) for t in args)) + # Generate code to apply A00 as a PETSc.Mat of type shell within the interior KSP A00_struct, *_ = matmult_kernel_code(a00, prefix="A_interior", fcp=fcp, matshell=True) A00_struct = A00_struct.replace("#include ", "") + # Replacement rules to use idofs, fdofs, A, and A00 on self.code rules = dict(A_struct=A_struct, A_call=A_call("x", "y"), ctx_struct=ctx_struct, ctx_pack=ctx_pack, A00_struct=A00_struct, size=size, isize=idofs.size, fsize=fdofs.size, idofs=", ".join(map(str, idofs.indices)), @@ -1367,10 +1380,23 @@ def __init__(self, kernel, name=None): fdofs.destroy() -class ParloopMatrixContext: +class PythonMatrixContext: + """Python matrix context that handles boundary conditions.""" - def __init__(self, parloop, x, y, bcs=None): - self._mult_parloop = parloop + def __init__(self, mult_callable, x, y, bcs=None): + """ + Parameters + ---------- + mult_callable : callable + The callable performing the matrix-vector product. + x : .Function + The tensor holding the input to the matrix-vector product. + y : .Function + The tensor holding the output to the matrix-vector product. + bcs : .BCBase[] or None + An iterable of boundary conditions to apply on ``x`` and ``y``. + """ + self._mult_callable = mult_callable self._x = x self._y = y Vrow = y.function_space() @@ -1408,11 +1434,11 @@ def _op(self, action, X, Y, W=None): @PETSc.Log.EventDecorator() def mult(self, mat, X, Y): - self._op(self._mult_parloop, X, Y) + self._op(self._mult_callable, X, Y) @PETSc.Log.EventDecorator() def multAdd(self, mat, X, Y, W): - self._op(self._mult_parloop, X, Y, W) + self._op(self._mult_callable, X, Y, W) def is_restricted(finat_element): @@ -1672,7 +1698,8 @@ def get_base_elements(e): class SparseAssembler: - + """Class to generate and cache python wrappers to insert sparse element + matrices directly with PETSc C code.""" _cache = {} @staticmethod @@ -1807,14 +1834,22 @@ def assemble_reference_tensor(self, V): @PETSc.Log.EventDecorator("FDMSetValues") def set_values(self, A, Vrow, Vcol, addv, mat_type="aij"): - """ - Assemble the stiffness matrix in the FDM basis using Kronecker products of interval matrices + """Assemble the stiffness matrix in the FDM basis using Kronecker + products of interval matrices. - :arg A: the :class:`PETSc.Mat` to assemble - :arg Vrow: the :class:`.FunctionSpace` test space - :arg Vcol: the :class:`.FunctionSpace` trial space - :arg addv: a `PETSc.Mat.InsertMode` - :arg mat_type: a `PETSc.Mat.Type` + Parameters + ---------- + A : PETSc.Mat + The (initialized) matrix to assemble. + Vrow : .FunctionSpace + The test space. + Vcol : .FunctionSpace + The trial space. + addv : PETSc.Mat.InsertMode + Flag indicating if we want to insert or add matrix values. + mat_type : PETSc.Mat.Type + The matrix type of auxiliary operator. This only used when ``A`` is a preallocator + to determine the nonzeros on the upper triangual part of an ``'sbaij'`` matrix. """ triu = A.getType() == "preallocator" and mat_type.endswith("sbaij") set_submat = SparseAssembler.setSubMatCSR(PETSc.COMM_SELF, triu=triu)