Trefftz support for Firedrake

firedrake/__init__.py

@@ -119,6 +119,7 @@
 from firedrake.randomfunctiongen import *
 from firedrake.external_operators import *
 from firedrake.progress_bar import ProgressBar  # noqa: F401
+from firedrake.trefftz import *
 
 from firedrake.logging import *
 # Set default log level

firedrake/trefftz.py

@@ -0,0 +1,219 @@
+"""
+Provides a class to compute the Trefftz embedding of a given function space.
+It is also used to compute aggregation embedding of a given function space.
+"""
+from typing import Optional
+from firedrake.petsc import PETSc
+from firedrake.cython.dmcommon import FACE_SETS_LABEL, CELL_SETS_LABEL
+from firedrake.assemble import assemble
+from firedrake.mesh import Mesh
+from firedrake.functionspace import FunctionSpace
+from firedrake.function import Function
+from firedrake.ufl_expr import TestFunction, TrialFunction
+from firedrake.constant import Constant
+from ufl import dx, dS, inner, jump, grad, dot, CellDiameter, FacetNormal, Form
+import scipy.sparse as sp
+import numpy as np
+
+__all__ = ["TrefftzEmbedding", "TrefftzKSP", "AggregationEmbedding", "dumb_aggregation"]
+
+class TrefftzEmbedding:
+    """
+    Computes the Trefftz embedding of a given function space
+    Parameters
+    ----------
+    V : Ambient function space.
+    b : Bilinear form defining the Trefftz operator.
+    dim : Dimension of the embedding.
+        Default is the dimension of the function space.
+    tol : Tolerance for the singular values cutoff.
+        Default is 1e-12.
+    backend : Backend to use for the computation of the SVD.
+        Default is "scipy".
+    """
+    def __init__(self, V: FunctionSpace, b: Form, dim: Optional[int] = None,
+                 tol: Optional[float]=1e-12):
+        self.V = V
+        self.b = b
+        self.dim = V.dim() if not dim else dim + 1
+        self.tol = tol
+        self.svdsolver = "scipy"
+
+    def assemble(self) -> tuple[PETSc.Mat, np.array]:
+        """
+        Assemble the embedding, compute the SVD and return the embedding matrix
+        """
+        self.B = assemble(self.b).M.handle
+        if self.svdsolver == "scipy":
+            indptr, indices, data = self.B.getValuesCSR()
+            Bsp = sp.csr_matrix((data, indices, indptr), shape=self.B.getSize())
+            _, sig, VT = sp.linalg.svds(Bsp, k=self.dim-1, which="SM")
+            QT = sp.csr_matrix(VT[0:sum(sig < self.tol), :])
+            QTpsc = PETSc.Mat().createAIJ(size=QT.shape, csr=(QT.indptr, QT.indices, QT.data))
+            self.dimT = QT.shape[0]
+            self.sig = sig
+        else:
+            raise NotImplementedError("Only scipy backend is supported")
+        return QTpsc, sig
+
+
+class TrefftzKSP:
+    """
+    Wraps a PETSc KSP object to solve the reduced
+    system obtained by the Trefftz embedding.
+
+    There will bne no type hinting following petsc4py's style.
+    """
+    def __init__(self):
+        pass
+
+    @staticmethod
+    def get_appctx(ksp: PETSc.KSP):
+        """
+        Get the application context from the KSP
+        Parameters
+        ----------
+        ksp : The KSP object
+        """
+        from firedrake.dmhooks import get_appctx
+        return get_appctx(ksp.getDM()).appctx
+
+    def setUp(self, ksp: PETSc.KSP):
+        """
+        Set up the Trefftz KSP
+        Parameters
+        ----------
+        ksp : The KSP object
+        """
+        appctx = self.get_appctx(ksp)
+        self.QT, _ = appctx["trefftz_embedding"].assemble()
+
+    def solve(self, ksp: PETSc.KSP, b: PETSc.Vec, x:PETSc.Vec):
+        """
+        Solve the Trefftz KSP
+        Parameters
+        ----------
+        ksp : The KSP object
+        b : The right-hand side
+        x : The solution
+        """
+        A, P = ksp.getOperators()
+        self.Q = PETSc.Mat().createTranspose(self.QT)
+        ATF = self.QT @ A @ self.Q
+        PTF = self.QT @ P @ self.Q
+        bTF = self.QT.createVecLeft()
+        self.QT.mult(b, bTF)
+
+        tiny_ksp = PETSc.KSP().create()
+        tiny_ksp.setOperators(ATF, PTF)
+        tiny_ksp.setOptionsPrefix("trefftz_")
+        tiny_ksp.setFromOptions()
+        xTF = ATF.createVecRight()
+        tiny_ksp.solve(bTF, xTF)
+        self.QT.multTranspose(xTF, x)
+        ksp.setConvergedReason(tiny_ksp.getConvergedReason())
+
+
+class AggregationEmbedding(TrefftzEmbedding):
+    """
+    This class computes the aggregation embedding of a given function space.
+    Parameters
+    ----------
+    V : Ambient function space.
+    mesh : The mesh on which the aggregation is defined.
+    polyMesh : The function defining the aggregation.
+    dim : Dimension of the embedding.
+        Default is the dimension of the function space.
+    tol : Tolerance for the singular values cutoff.
+        Default is 1e-12.
+    """
+    def __init__(self, V: FunctionSpace, mesh: Mesh, polyMesh: Function,
+                 dim: Optional[int] = None, tol: Optional[float]=1e-12):
+        # Relabel facets that are inside an aggregated region
+        offset = 1 + mesh.topology_dm.getLabelSize(FACE_SETS_LABEL)
+        offset += mesh.topology_dm.getLabelSize(CELL_SETS_LABEL)
+        nPoly = int(max(polyMesh.dat.data[:]))  # Number of aggregates
+        getIdx = mesh._cell_numbering.getOffset
+        plex = mesh.topology_dm
+        pStart, pEnd = plex.getDepthStratum(2)
+        self.facet_index = []
+        for poly in range(nPoly+1):
+            facets = []
+            for i in range(pStart, pEnd):
+                if polyMesh.dat.data[getIdx(i)] == poly:
+                    for f in plex.getCone(i):
+                        if f in facets:
+                            plex.setLabelValue(FACE_SETS_LABEL, f, offset+poly)
+                            if offset+poly not in self.facet_index:
+                                self.facet_index = self.facet_index + [offset+poly]
+                    facets = facets + list(plex.getCone(i))
+        self.mesh = Mesh(plex)
+        h = CellDiameter(self.mesh)
+        n = FacetNormal(self.mesh)
+        W = FunctionSpace(self.mesh, V.ufl_element())
+        u = TrialFunction(W)
+        v = TestFunction(W)
+        self.b = Constant(0)*inner(u, v)*dx
+        for i in self.facet_index:
+            self.b += inner(jump(u), jump(v))*dS(i)
+        for k in range(1, V.ufl_element().degree()+1):
+            for i in self.facet_index:
+                self.b += ((0.5 * h("+") + 0.5 * h("-"))**(2*k)) *\
+                    inner(jump_normal(u, n("+"), k), jump_normal(v, n("+"), k))*dS(i)
+        super().__init__(W, self.b, dim, tol)
+
+
+def jump_normal(u: Function, n: FacetNormal, k: int):
+    """
+    Compute the jump of the normal derivative of a function u
+    Parameters
+    ----------
+    u : The function.
+    n : The normal vector.
+    k : The order of the normal derivative we aim to compute.
+    """
+    j = 0.5*dot(n, (grad(u)("+")-grad(u)("-")))
+    for _ in range(1, k):
+        j = 0.5*dot(n, (grad(j)-grad(j)))
+    return j
+
+
+def dumb_aggregation(mesh: Mesh) -> Function:
+    """
+    Compute a dumb aggregation of the mesh
+    Parameters
+    ----------
+    mesh : The mesh we aim to aggregate.
+    """
+    if mesh.comm.size > 1:
+        raise NotImplementedError("Parallel mesh aggregation not supported")
+    plex = mesh.topology_dm
+    pStart, pEnd = plex.getDepthStratum(2)
+    _, eEnd = plex.getDepthStratum(1)
+    adjacency = []
+    for i in range(pStart, pEnd):
+        ad = plex.getAdjacency(i)
+        local = []
+        for a in ad:
+            supp = plex.getSupport(a)
+            supp = supp[supp < eEnd]
+            for s in supp:
+                if s < pEnd and s != ad[0]:
+                    local = local + [s]
+        adjacency = adjacency + [(i, local)]
+    adjacency = sorted(adjacency, key=lambda x: len(x[1]))[::-1]
+    u = Function(FunctionSpace(mesh, "DG", 0))
+
+    getIdx = mesh._cell_numbering.getOffset
+    av = list(range(pStart, pEnd))
+    col = 0
+    for a in adjacency:
+        if a[0] in av:
+            for k in a[1]:
+                if k in av:
+                    av.remove(k)
+                    u.dat.data[getIdx(k)] = col
+            av.remove(a[0])
+            u.dat.data[getIdx(a[0])] = col
+            col = col + 1
+    return u

tests/regression/test_trefftz.py

@@ -0,0 +1,85 @@
+from firedrake import *
+from firedrake.trefftz import TrefftzEmbedding, AggregationEmbedding, dumb_aggregation
+import pytest
+
+
+@pytest.mark.skipcomplex
+def test_trefftz_laplace():
+    order = 6
+    mesh = UnitSquareMesh(2, 2)
+    x, y = SpatialCoordinate(mesh)
+    h = CellDiameter(mesh)
+    n = FacetNormal(mesh)
+    V = FunctionSpace(mesh, "DG", order)
+    u = TrialFunction(V)
+    v = TestFunction(V)
+
+    def delta(u):
+        return div(grad(u))
+
+    a = inner(delta(u), delta(v)) * dx
+    alpha = 4
+    mean_dudn = 0.5 * dot(grad(u("+"))+grad(u("-")), n("+"))
+    mean_dvdn = 0.5 * dot(grad(v("+"))+grad(v("-")), n("+"))
+    aDG = inner(grad(u), grad(v)) * dx
+    aDG += inner((alpha*order**2/(h("+")+h("-")))*jump(u), jump(v))*dS
+    aDG += inner(-mean_dudn, jump(v))*dS-inner(mean_dvdn, jump(u))*dS
+    aDG += alpha*order**2/h*inner(u, v)*ds
+    aDG += -inner(dot(n, grad(u)), v)*ds - inner(dot(n, grad(v)), u)*ds
+    f = Function(V).interpolate(exp(x)*sin(y))
+    L = alpha*order**2/h*inner(f, v)*ds - inner(dot(n, grad(v)), f)*ds
+    # Solve the problem
+    uDG = Function(V)
+    uDG.rename("uDG")
+    embd = TrefftzEmbedding(V, a, tol=1e-8)
+    appctx = {"trefftz_embedding": embd}
+    uDG = Function(V)
+    solve(aDG == L, uDG, solver_parameters={"ksp_type": "python",
+          "ksp_python_type": "firedrake.TrefftzKSP"},
+          appctx=appctx)
+    assert (assemble(inner(uDG-f, uDG-f)*dx) < 1e-6)
+    assert (embd.dimT < V.dim()/2)
+
+
+@pytest.mark.skipcomplex
+def test_trefftz_aggregation():
+    try:
+        from netgen.occ import WorkPlane, OCCGeometry
+    except ImportError:
+        # Netgen is not installed
+        pytest.skip("Netgen/ngsPETSc not installed", allow_module_level=True)
+
+    Rectangle = WorkPlane().Rectangle(1, 1).Face()
+    geo = OCCGeometry(Rectangle, dim=2)
+    ngmesh = geo.GenerateMesh(maxh=0.3)
+    mesh = Mesh(ngmesh)
+
+    polymesh = dumb_aggregation(mesh)
+
+    order = 3
+    x, y = SpatialCoordinate(mesh)
+    h = CellDiameter(mesh)
+    n = FacetNormal(mesh)
+    V = FunctionSpace(mesh, "DG", order)
+    u = TrialFunction(V)
+    v = TestFunction(V)
+
+    alpha = 1e3
+    mean_dudn = 0.5 * dot(grad(u("+"))+grad(u("-")), n("+"))
+    mean_dvdn = 0.5 * dot(grad(v("+"))+grad(v("-")), n("+"))
+    aDG = inner(grad(u), grad(v)) * dx
+    aDG += inner((alpha*order**2/(h("+")+h("-")))*jump(u), jump(v))*dS
+    aDG += inner(-mean_dudn, jump(v))*dS-inner(mean_dvdn, jump(u))*dS
+    aDG += alpha*order**2/h*inner(u, v)*ds
+    aDG += -inner(dot(n, grad(u)), v)*ds - inner(dot(n, grad(v)), u)*ds
+    f = Function(V).interpolate(exp(x)*sin(y))
+    L = alpha*order**2/h*inner(f, v)*ds - inner(dot(n, grad(v)), f)*ds
+    agg_embd = AggregationEmbedding(V, mesh, polymesh)
+    appctx = {"trefftz_embedding": agg_embd}
+
+    uDG = Function(V)
+    solve(aDG == L, uDG, solver_parameters={"ksp_type": "python",
+          "ksp_python_type": "firedrake.TrefftzKSP"},
+          appctx=appctx)
+
+    assert (assemble(inner(uDG-f, uDG-f)*dx) < 1e-9)