Merge pull request #61 from firedrakeproject/pbrubeck/cleanup-dualset

Tidy up DualSet

FIAT/brezzi_douglas_marini.py

@@ -9,81 +9,61 @@
                   polynomial_set, nedelec)
 from FIAT.check_format_variant import check_format_variant
 from FIAT.quadrature_schemes import create_quadrature
+from FIAT.quadrature import FacetQuadratureRule
 
 
 class BDMDualSet(dual_set.DualSet):
     def __init__(self, ref_el, degree, variant, interpolant_deg):
-
-        # Initialize containers for map: mesh_entity -> dof number and
-        # dual basis
-        entity_ids = {}
         nodes = []
-
         sd = ref_el.get_spatial_dimension()
-        t = ref_el.get_topology()
+        top = ref_el.get_topology()
+
+        entity_ids = {}
+        # set to empty
+        for dim in top:
+            entity_ids[dim] = {}
+            for entity in top[dim]:
+                entity_ids[dim][entity] = []
 
         if variant == "integral":
             facet = ref_el.get_facet_element()
-            # Facet nodes are \int_F v\cdot n p ds where p \in P_{q-1}
-            # degree is q - 1
-            Q = create_quadrature(facet, interpolant_deg + degree)
+            # Facet nodes are \int_F v\cdot n p ds where p \in P_{q}
+            # degree is q
+            Q_ref = create_quadrature(facet, interpolant_deg + degree)
             Pq = polynomial_set.ONPolynomialSet(facet, degree)
-            Pq_at_qpts = Pq.tabulate(Q.get_points())[(0,)*(sd - 1)]
-            nodes.extend(functional.IntegralMomentOfScaledNormalEvaluation(ref_el, Q, phi, f)
-                         for f in range(len(t[sd - 1]))
-                         for phi in Pq_at_qpts)
-
-            # internal nodes
-            if degree > 1:
-                Q = create_quadrature(ref_el, interpolant_deg + degree - 1)
-                qpts = Q.get_points()
-                Nedel = nedelec.Nedelec(ref_el, degree - 1, variant)
-                Nedfs = Nedel.get_nodal_basis()
-                Ned_at_qpts = Nedfs.tabulate(qpts)[(0,) * sd]
+            Pq_at_qpts = Pq.tabulate(Q_ref.get_points())[(0,)*(sd - 1)]
+            for f in top[sd - 1]:
+                cur = len(nodes)
+                Q = FacetQuadratureRule(ref_el, sd - 1, f, Q_ref)
+                Jdet = Q.jacobian_determinant()
+                n = ref_el.compute_scaled_normal(f) / Jdet
+                phis = n[None, :, None] * Pq_at_qpts[:, None, :]
                 nodes.extend(functional.FrobeniusIntegralMoment(ref_el, Q, phi)
-                             for phi in Ned_at_qpts)
+                             for phi in phis)
+                entity_ids[sd - 1][f] = list(range(cur, len(nodes)))
 
         elif variant == "point":
             # Define each functional for the dual set
             # codimension 1 facets
-            for i in range(len(t[sd - 1])):
-                pts_cur = ref_el.make_points(sd - 1, i, sd + degree)
-                nodes.extend(functional.PointScaledNormalEvaluation(ref_el, i, pt)
+            for f in top[sd - 1]:
+                cur = len(nodes)
+                pts_cur = ref_el.make_points(sd - 1, f, sd + degree)
+                nodes.extend(functional.PointScaledNormalEvaluation(ref_el, f, pt)
                              for pt in pts_cur)
+                entity_ids[sd - 1][f] = list(range(cur, len(nodes)))
 
-            # internal nodes
-            if degree > 1:
-                Q = create_quadrature(ref_el, 2 * degree - 1)
-                qpts = Q.get_points()
-                Nedel = nedelec.Nedelec(ref_el, degree - 1, variant)
-                Nedfs = Nedel.get_nodal_basis()
-                Ned_at_qpts = Nedfs.tabulate(qpts)[(0,) * sd]
-                nodes.extend(functional.FrobeniusIntegralMoment(ref_el, Q, phi)
-                             for phi in Ned_at_qpts)
-
-        # sets vertices (and in 3d, edges) to have no nodes
-        for i in range(sd - 1):
-            entity_ids[i] = {}
-            for j in range(len(t[i])):
-                entity_ids[i][j] = []
-
-        cur = 0
-
-        # set codimension 1 (edges 2d, faces 3d) dof
-        pts_facet_0 = ref_el.make_points(sd - 1, 0, sd + degree)
-        pts_per_facet = len(pts_facet_0)
-
-        entity_ids[sd - 1] = {}
-        for i in range(len(t[sd - 1])):
-            entity_ids[sd - 1][i] = list(range(cur, cur + pts_per_facet))
-            cur += pts_per_facet
-
-        # internal nodes, if applicable
-        entity_ids[sd] = {0: []}
-
+        # internal nodes
         if degree > 1:
-            num_internal_nodes = len(Ned_at_qpts)
-            entity_ids[sd][0] = list(range(cur, cur + num_internal_nodes))
+            if interpolant_deg is None:
+                interpolant_deg = degree
+            cur = len(nodes)
+            Q = create_quadrature(ref_el, interpolant_deg + degree - 1)
+            Nedel = nedelec.Nedelec(ref_el, degree - 1, variant)
+            Nedfs = Nedel.get_nodal_basis()
+            Ned_at_qpts = Nedfs.tabulate(Q.get_points())[(0,) * sd]
+            nodes.extend(functional.FrobeniusIntegralMoment(ref_el, Q, phi)
+                         for phi in Ned_at_qpts)
+            entity_ids[sd][0] = list(range(cur, len(nodes)))
 
         super(BDMDualSet, self).__init__(nodes, ref_el, entity_ids)
 

FIAT/dual_set.py

@@ -7,9 +7,8 @@
 # SPDX-License-Identifier:    LGPL-3.0-or-later
 
 import numpy
-import collections
 
-from FIAT import polynomial_set
+from FIAT import polynomial_set, functional
 
 
 class DualSet(object):
@@ -105,63 +104,78 @@ def to_riesz(self, poly_set):
         ed = poly_set.get_embedded_degree()
         num_exp = es.get_num_members(poly_set.get_embedded_degree())
 
-        riesz_shape = tuple([num_nodes] + list(tshape) + [num_exp])
-
+        riesz_shape = (num_nodes, *tshape, num_exp)
         mat = numpy.zeros(riesz_shape, "d")
 
-        # Dictionaries mapping pts to which functionals they come from
-        pts_to_ells = collections.OrderedDict()
-        dpts_to_ells = collections.OrderedDict()
-
+        pts = set()
+        dpts = set()
+        Qs_to_ells = dict()
         for i, ell in enumerate(self.nodes):
-            for pt in ell.pt_dict:
-                if pt in pts_to_ells:
-                    pts_to_ells[pt].append(i)
-                else:
-                    pts_to_ells[pt] = [i]
-
-            for pt in ell.deriv_dict:
-                if pt in dpts_to_ells:
-                    dpts_to_ells[pt].append(i)
-                else:
-                    dpts_to_ells[pt] = [i]
+            if isinstance(ell, functional.IntegralMoment):
+                Q = ell.Q
+            else:
+                Q = None
+                pts.update(ell.pt_dict.keys())
+                dpts.update(ell.deriv_dict.keys())
+            if Q in Qs_to_ells:
+                Qs_to_ells[Q].append(i)
+            else:
+                Qs_to_ells[Q] = [i]
+
+        Qs_to_pts = {None: tuple(sorted(pts))}
+        for Q in Qs_to_ells:
+            if Q is not None:
+                cur_pts = tuple(map(tuple, Q.pts))
+                Qs_to_pts[Q] = cur_pts
+                pts.update(cur_pts)
 
         # Now tabulate the function values
-        pts = list(pts_to_ells.keys())
-        expansion_values = es.tabulate(ed, pts)
-
-        for j, pt in enumerate(pts):
-            which_ells = pts_to_ells[pt]
-
-            for k in which_ells:
-                pt_dict = self.nodes[k].pt_dict
-                wc_list = pt_dict[pt]
-
-                for i in range(num_exp):
-                    for (w, c) in wc_list:
-                        mat[k][c][i] += w*expansion_values[i, j]
+        pts = list(sorted(pts))
+        expansion_values = numpy.transpose(es.tabulate(ed, pts))
+
+        for Q in Qs_to_ells:
+            ells = Qs_to_ells[Q]
+            cur_pts = Qs_to_pts[Q]
+            indices = list(map(pts.index, cur_pts))
+            wshape = (len(ells), *tshape, len(cur_pts))
+            wts = numpy.zeros(wshape, "d")
+            if Q is None:
+                for i, k in enumerate(ells):
+                    ell = self.nodes[k]
+                    for pt, wc_list in ell.pt_dict.items():
+                        j = cur_pts.index(pt)
+                        for (w, c) in wc_list:
+                            wts[i][c][j] = w
+            else:
+                for i, k in enumerate(ells):
+                    ell = self.nodes[k]
+                    wts[i][ell.comp][:] = ell.f_at_qpts
+                qwts = Q.get_weights()
+                wts = numpy.multiply(wts, qwts, out=wts)
+            mat[ells] += numpy.dot(wts, expansion_values[indices])
 
         # Tabulate the derivative values that are needed
-        max_deriv_order = max([ell.max_deriv_order for ell in self.nodes])
+        max_deriv_order = max(ell.max_deriv_order for ell in self.nodes)
         if max_deriv_order > 0:
-            dpts = list(dpts_to_ells.keys())
+            dpts = list(sorted(dpts))
             # It's easiest/most efficient to get derivatives of the
             # expansion set through the polynomial set interface.
             # This is creating a short-lived set to do just this.
-            expansion = polynomial_set.ONPolynomialSet(self.ref_el, ed)
+            coeffs = numpy.eye(num_exp)
+            expansion = polynomial_set.PolynomialSet(self.ref_el, ed, ed, es, coeffs)
             dexpansion_values = expansion.tabulate(dpts, max_deriv_order)
 
-            for j, pt in enumerate(dpts):
-                which_ells = dpts_to_ells[pt]
-
-                for k in which_ells:
-                    dpt_dict = self.nodes[k].deriv_dict
-                    wac_list = dpt_dict[pt]
-
-                    for i in range(num_exp):
-                        for (w, alpha, c) in wac_list:
-                            mat[k][c][i] += w*dexpansion_values[alpha][i, j]
-
+            ells = [k for k, ell in enumerate(self.nodes) if len(ell.deriv_dict) > 0]
+            wshape = (len(ells), *tshape, len(dpts))
+            dwts = {alpha: numpy.zeros(wshape, "d") for alpha in dexpansion_values if sum(alpha) > 0}
+            for i, k in enumerate(ells):
+                ell = self.nodes[k]
+                for pt, wac_list in ell.deriv_dict.items():
+                    j = dpts.index(pt)
+                    for (w, alpha, c) in wac_list:
+                        dwts[alpha][i][c][j] = w
+            for alpha in dwts:
+                mat[ells] += numpy.dot(dwts[alpha], dexpansion_values[alpha].T)
         return mat
 
 

FIAT/functional.py

@@ -26,16 +26,7 @@ def index_iterator(shp):
     """Constructs a generator iterating over all indices in
     shp in generalized column-major order  So if shp = (2,2), then we
     construct the sequence (0,0),(0,1),(1,0),(1,1)"""
-    if len(shp) == 0:
-        return
-    elif len(shp) == 1:
-        for i in range(shp[0]):
-            yield [i]
-    else:
-        shp_foo = shp[1:]
-        for i in range(shp[0]):
-            for foo in index_iterator(shp_foo):
-                yield [i] + foo
+    return numpy.ndindex(shp)
 
 
 class Functional(object):
@@ -292,12 +283,11 @@ class IntegralMoment(Functional):
 
     def __init__(self, ref_el, Q, f_at_qpts, comp=tuple(), shp=tuple()):
         self.Q = Q
+        self.f_at_qpts = f_at_qpts
         qpts, qwts = Q.get_points(), Q.get_weights()
-        pt_dict = OrderedDict()
         self.comp = comp
-        for i in range(len(qpts)):
-            pt_cur = tuple(qpts[i])
-            pt_dict[pt_cur] = [(qwts[i] * f_at_qpts[i], comp)]
+        weights = numpy.multiply(f_at_qpts, qwts)
+        pt_dict = {tuple(pt): [(wt, comp)] for pt, wt in zip(qpts, weights)}
         Functional.__init__(self, ref_el, shp, pt_dict, {}, "IntegralMoment")
 
     def __call__(self, fn):
@@ -331,7 +321,7 @@ def __init__(self, ref_el, facet_no, Q, f_at_qpts):
 
         dpt_dict = OrderedDict()
 
-        alphas = [tuple([1 if j == i else 0 for j in range(sd)]) for i in range(sd)]
+        alphas = [tuple(1 if j == i else 0 for j in range(sd)) for i in range(sd)]
         for j, pt in enumerate(dpts):
             dpt_dict[tuple(pt)] = [(qwts[j]*n[i]*f_at_qpts[j], alphas[i], tuple()) for i in range(sd)]
 
@@ -484,24 +474,23 @@ def __init__(self, ref_el, Q, f_at_qpts):
                          "IntegralMomentOfDivergence")
 
 
-class FrobeniusIntegralMoment(Functional):
+class FrobeniusIntegralMoment(IntegralMoment):
 
     def __init__(self, ref_el, Q, f_at_qpts):
         # f_at_qpts is (some shape) x num_qpts
         shp = tuple(f_at_qpts.shape[:-1])
-        if len(Q.get_points()) != f_at_qpts.shape[-1]:
+        if len(Q.pts) != f_at_qpts.shape[-1]:
             raise Exception("Mismatch in number of quadrature points and values")
 
+        self.Q = Q
+        self.comp = slice(None, None)
+        self.f_at_qpts = f_at_qpts
         qpts, qwts = Q.get_points(), Q.get_weights()
-        pt_dict = {}
-
-        for i, (pt_cur, wt_cur) in enumerate(zip(map(tuple, qpts), qwts)):
-            pt_dict[pt_cur] = []
-            for alfa in index_iterator(shp):
-                qpidx = tuple(alfa + [i])
-                pt_dict[pt_cur].append((wt_cur * f_at_qpts[qpidx], tuple(alfa)))
+        weights = numpy.transpose(numpy.multiply(f_at_qpts, qwts), (-1,) + tuple(range(len(shp))))
+        alphas = list(index_iterator(shp))
 
-        super().__init__(ref_el, shp, pt_dict, {}, "FrobeniusIntegralMoment")
+        pt_dict = {tuple(pt): [(wt[alpha], alpha) for alpha in alphas] for pt, wt in zip(qpts, weights)}
+        Functional.__init__(self, ref_el, shp, pt_dict, {}, "FrobeniusIntegralMoment")
 
 
 class PointNormalEvaluation(Functional):