From 2936ff1d1d3b2d813b518381d3cd07a50c67a0f8 Mon Sep 17 00:00:00 2001
From: Pablo Brubeck <brubeck@protonmail.com>
Date: Mon, 12 Aug 2024 11:55:53 +0100
Subject: [PATCH] Reduced Demkowicz elements

---
 FIAT/demkowicz.py  | 31 +++++++++++++++++++++++++++++--
 FIAT/restricted.py |  4 ++++
 2 files changed, 33 insertions(+), 2 deletions(-)

diff --git a/FIAT/demkowicz.py b/FIAT/demkowicz.py
index 7020fb78..5a29cdd8 100644
--- a/FIAT/demkowicz.py
+++ b/FIAT/demkowicz.py
@@ -8,6 +8,7 @@
 import scipy
 
 from FIAT.dual_set import DualSet
+from FIAT.restricted import RestrictedDualSet
 from FIAT.functional import PointEvaluation, FrobeniusIntegralMoment
 from FIAT.polynomial_set import make_bubbles, ONPolynomialSet, PolynomialSet
 from FIAT.quadrature import FacetQuadratureRule
@@ -68,6 +69,7 @@ class DemkowiczDual(DualSet):
     def __init__(self, ref_el, degree, sobolev_space, kind=2):
         nodes = []
         entity_ids = {}
+        reduced_dofs = {}
         top = ref_el.get_topology()
         sd = ref_el.get_spatial_dimension()
         formdegree = {"H1": 0, "HCurl": 1, "HDiv": sd-1, "L2": sd}[sobolev_space]
@@ -79,14 +81,17 @@ def __init__(self, ref_el, degree, sobolev_space, kind=2):
             if dim < formdegree or degree <= dim - formdegree:
                 for entity in top[dim]:
                     entity_ids[dim][entity] = []
+                reduced_dofs[dim] = 0
             elif dim == 0 and formdegree == 0:
                 for entity in sorted(top[dim]):
                     cur = len(nodes)
                     pts = ref_el.make_points(dim, entity, degree)
                     nodes.extend(PointEvaluation(ref_el, pt) for pt in pts)
                     entity_ids[dim][entity] = list(range(cur, len(nodes)))
+                reduced_dofs[dim] = len(nodes)
             else:
-                Q_ref, Phis = self._reference_duals(dim, degree, formdegree, sobolev_space, kind)
+                Q_ref, Phis, rdofs = self._reference_duals(dim, degree, formdegree, sobolev_space, kind)
+                reduced_dofs[dim] = rdofs
                 mapping = dual_mapping if dim == sd else trace
                 for entity in sorted(top[dim]):
                     cur = len(nodes)
@@ -94,6 +99,7 @@ def __init__(self, ref_el, degree, sobolev_space, kind=2):
                     nodes.extend(FrobeniusIntegralMoment(ref_el, Q, phi) for phi in phis)
                     entity_ids[dim][entity] = list(range(cur, len(nodes)))
 
+        self._reduced_dofs = reduced_dofs
         super(DemkowiczDual, self).__init__(nodes, ref_el, entity_ids)
 
     def _reference_duals(self, dim, degree, formdegree, sobolev_space, kind):
@@ -117,6 +123,7 @@ def _reference_duals(self, dim, degree, formdegree, sobolev_space, kind):
                 dtrial = dtrial[:, None, :]
             K = self._bubble_derivative_moments(facet, degree, formdegree, kind, Qpts, Qwts, dtrial)
 
+        reduced_dofs = K.shape[0]
         if formdegree > 0:
             if dim == 2 and formdegree == 1 and sobolev_space == "HDiv":
                 trial = perp(trial)
@@ -125,7 +132,7 @@ def _reference_duals(self, dim, degree, formdegree, sobolev_space, kind):
             K = numpy.vstack((K, M))
 
         duals = numpy.tensordot(K, P_at_qpts[(0,) * dim], axes=(1, 0))
-        return Q, duals
+        return Q, duals, reduced_dofs
 
     def _bubble_derivative_moments(self, facet, degree, formdegree, kind, Qpts, Qwts, trial):
         """Integrate trial expressions against an orthonormal basis for
@@ -164,6 +171,26 @@ def _bubble_derivative_moments(self, facet, degree, formdegree, kind, Qpts, Qwts
             dtest = numpy.tensordot(S.T, dtest, axes=(1, 0))
         return inner(dtest, trial, Qwts)
 
+    def get_indices(self, restriction_domain, take_closure=True):
+        """Return the list of dofs with support on the given restriction domain.
+        Allows for reduced Demkowicz elements, excluding the exterior
+        derivative of the previous space in the de Rham complex.
+
+        :arg restriction_domain: can be 'reduced', 'interior', 'vertex',
+                                 'edge', 'face' or 'facet'
+        :kwarg take_closure: Are we taking the closure of the restriction domain?
+        """
+        if restriction_domain == "reduced":
+            indices = []
+            entity_ids = self.get_entity_ids()
+            for dim in entity_ids:
+                reduced_dofs = self._reduced_dofs[dim]
+                for entity, ids in entity_ids[dim].items():
+                    indices.extend(ids[:reduced_dofs])
+            return indices
+        else:
+            return super(DemkowiczDual, self).get_indices(restriction_domain, take_closure=take_closure)
+
 
 class FDMDual(DualSet):
 
diff --git a/FIAT/restricted.py b/FIAT/restricted.py
index 26c82d91..6f964b39 100644
--- a/FIAT/restricted.py
+++ b/FIAT/restricted.py
@@ -40,6 +40,10 @@ def get_indices(self, restriction_domain, take_closure=True):
         # Call get_indices on the parent class to support multiple restriction domains
         return type(self._dual).get_indices(self, restriction_domain, take_closure=take_closure)
 
+    def __getattr__(self, name):
+        val = getattr(self._dual, name)
+        return val
+
 
 class RestrictedElement(CiarletElement):
     """Restrict the given element to the specified list of dofs."""