Merge pull request #761 from kinnala/multibasis-indexing

Make multibasis indexing more explicit; allow changing Dofs in assembly

README.md

@@ -225,6 +225,13 @@ with respect to documented and/or tested features.
 
 ### Unreleased
 
+- Added: `asm` will now include the extra parameter `idx` in `FormExtraParams`
+  that can be used to identify which basis is being assembled
+- Added: All basis class constructors now accept `dofs` keyword argument for
+  specifying custom `Dofs` object to be used in the assembly
+- Added: `Dofs` constructor accepts the `offset` keyword argument for
+  specifying a nonzero initial index for indexing the DOFs
+
 ### [4.0.1] - 2021-10-15
 
 - Fixed: `MappingIsoparametric` can now be pickled

docs/examples/ex04.py

@@ -77,96 +77,85 @@
 E1 = ElementQuad2()
 E = ElementVector(E1)
 
-ib = Basis(m, e, intorder=4)
-Ib = Basis(M, E, intorder=4)
+dofs1 = Dofs(m, e)
+dofs2 = Dofs(M, E, offset=dofs1.N)
+
+bases = [
+    Basis(m, e, intorder=4, dofs=dofs1),
+    Basis(M, E, intorder=4, dofs=dofs2),
+]
 
 mapping = MappingMortar.init_2D(m, M,
                                 m.boundaries['contact'],
                                 M.facets_satisfying(lambda x: x[0] == 1.0),
                                 np.array([0.0, 1.0]))
 
 mb = [
-    MortarFacetBasis(m, e, mapping=mapping, intorder=4, side=0),
-    MortarFacetBasis(M, E, mapping=mapping, intorder=4, side=1),
+    MortarFacetBasis(m, e, mapping=mapping, intorder=4, side=0, dofs=dofs1),
+    MortarFacetBasis(M, E, mapping=mapping, intorder=4, side=1, dofs=dofs2),
 ]
 
 # define bilinear forms
-E = 1000.0
-nu = 0.3
-Lambda, Mu = lame_parameters(E, nu)
+youngs_modulus = 1000.0
+poisson_ratio = 0.3
+Lambda, Mu = lame_parameters(youngs_modulus, poisson_ratio)
 
-weakform1 = linear_elasticity(Lambda, Mu)
-weakform2 = linear_elasticity(Lambda, Mu)
+weakform = linear_elasticity(Lambda, Mu)
 C = linear_stress(Lambda, Mu)
 
 alpha = 1000
 limit = 0.3
 
-# assemble the stiffness matrices
-K1 = asm(weakform1, ib)
-K2 = asm(weakform2, Ib)
-K_blocks = [[K1, 0.], [0., K2]]
-f = [None] * 2
-
+# mortar forms
+@BilinearForm
+def bilin_mortar(u, v, w):
+    # jumps
+    ju = (-1.) ** w.idx[0] * dot(u, w.n)
+    jv = (-1.) ** w.idx[1] * dot(v, w.n)
+    nxn = prod(w.n, w.n)
+    mu = .5 * ddot(nxn, C(sym_grad(u)))
+    mv = .5 * ddot(nxn, C(sym_grad(v)))
+    return ((1. / (alpha * w.h) * ju * jv - mu * jv - mv * ju)
+            * (np.abs(w.x[1]) <= limit))
 
 def gap(x):
     """Initial gap between the bodies."""
     return (1. - np.sqrt(1. - x[1] ** 2))
 
-# assemble the mortar matrices
-for i in range(2):
-    for j in range(2):
-
-        @BilinearForm
-        def bilin_mortar(u, v, w):
-            ju = (-1.) ** j * dot(u, w.n)
-            jv = (-1.) ** i * dot(v, w.n)
-            nxn = prod(w.n, w.n)
-            mu = .5 * ddot(nxn, C(sym_grad(u)))
-            mv = .5 * ddot(nxn, C(sym_grad(v)))
-            return ((1. / (alpha * w.h) * ju * jv - mu * jv - mv * ju)
-                    * (np.abs(w.x[1]) <= limit))
-
-        K_blocks[i][j] += asm(bilin_mortar, mb[j], mb[i])
 
-    @LinearForm
-    def lin_mortar(v, w):
-        jv = (-1.) ** i * dot(v, w.n)
-        mv = .5 * ddot(prod(w.n, w.n), C(sym_grad(v)))
-        return ((1. / (alpha * w.h) * gap(w.x) * jv - gap(w.x) * mv)
-                * (np.abs(w.x[1]) <= limit))
+@LinearForm
+def lin_mortar(v, w):
+    jv = (-1.) ** w.idx[0] * dot(v, w.n)
+    mv = .5 * ddot(prod(w.n, w.n), C(sym_grad(v)))
+    return ((1. / (alpha * w.h) * gap(w.x) * jv - gap(w.x) * mv)
+            * (np.abs(w.x[1]) <= limit))
 
-    f[i] = asm(lin_mortar, mb[i])
 
-import scipy.sparse
-K = (scipy.sparse.bmat(K_blocks)).tocsr()
+# assemble
+K = asm(weakform, bases) + asm(bilin_mortar, mb, mb)
+f = asm(lin_mortar, mb)
 
 # set boundary conditions and solve
-i1 = np.arange(K1.shape[0])
-i2 = np.arange(K2.shape[0]) + K1.shape[0]
+i1 = np.arange(bases[0].N)
+i2 = np.arange(bases[1].N) + bases[0].N
 
-D1 = ib.get_dofs(lambda x: x[0] == 0.0).all()
-D2 = Ib.get_dofs(lambda x: x[0] == 2.0).all()
+D1 = bases[0].get_dofs(lambda x: x[0] == 0.0).all()
+D2 = bases[1].get_dofs(lambda x: x[0] == 2.0).all()
 
 x = np.zeros(K.shape[0])
-
-f = np.hstack((f[0], f[1]))
-
-x = np.zeros(K.shape[0])
-D = np.concatenate((D1, D2 + ib.N))
+D = np.concatenate((D1, D2))
 I = np.setdiff1d(np.arange(K.shape[0]), D)
 
-x[ib.get_dofs(lambda x: x[0] == 0.0).nodal['u^1']] = 0.1
-x[ib.get_dofs(lambda x: x[0] == 0.0).facet['u^1']] = 0.1
+x[bases[0].get_dofs(lambda x: x[0] == 0.0).nodal['u^1']] = 0.1
+x[bases[0].get_dofs(lambda x: x[0] == 0.0).facet['u^1']] = 0.1
 
 x = solve(*condense(K, f, I=I, x=x))
 
-
 # create a displaced mesh
 sf = 1
 
-mdefo = m.translated(sf * x[i1][ib.nodal_dofs])
-Mdefo = M.translated(sf * x[i2][Ib.nodal_dofs])
+mdefo = m.translated(sf * x[dofs1.nodal_dofs])
+Mdefo = M.translated(sf * x[dofs2.nodal_dofs])
 
 
 # post processing
@@ -189,22 +178,22 @@ def mass(u, v, w):
         Ib_dg = Basis(Mdefo, E_dg, intorder=4)
 
         s[itr, jtr] = solve(asm(mass, ib_dg),
-                            asm(proj_cauchy, ib, ib_dg) @ x[i1])
+                            asm(proj_cauchy, bases[0], ib_dg) @ x[i1])
         S[itr, jtr] = solve(asm(mass, Ib_dg),
-                            asm(proj_cauchy, Ib, Ib_dg) @ x[i2])
+                            asm(proj_cauchy, bases[1], Ib_dg) @ x, I=i2)
 
-s[2, 2] = nu * (s[0, 0] + s[1, 1])
-S[2, 2] = nu * (S[0, 0] + S[1, 1])
+s[2, 2] = poisson_ratio * (s[0, 0] + s[1, 1])
+S[2, 2] = poisson_ratio * (S[0, 0] + S[1, 1])
 
 vonmises1 = np.sqrt(.5 * ((s[0, 0] - s[1, 1]) ** 2 +
                           (s[1, 1] - s[2, 2]) ** 2 +
                           (s[2, 2] - s[0, 0]) ** 2 +
-                          6. * s[0, 1]**2))
+                          6. * s[0, 1] ** 2))
 
 vonmises2 = np.sqrt(.5 * ((S[0, 0] - S[1, 1]) ** 2 +
                           (S[1, 1] - S[2, 2]) ** 2 +
                           (S[2, 2] - S[0, 0]) ** 2 +
-                          6. * S[0, 1]**2))
+                          6. * S[0, 1] ** 2))
 
 
 if __name__ == "__main__":

docs/examples/ex07.py

@@ -1,4 +1,4 @@
-"""Interior penalty method."""
+"""Discontinuous Galerkin method."""
 
 from skfem import *
 from skfem.helpers import grad, dot
@@ -14,15 +14,21 @@
 
 @BilinearForm
 def dgform(u, v, p):
-    ju = p.sign1 * u
-    jv = p.sign2 * v
+    ju = (-1.) ** p.idx[0] * u
+    jv = (-1.) ** p.idx[1] * v
     h = p.h
     n = p.n
     return ju * jv / (alpha * h) - dot(grad(u), n) * jv - dot(grad(v), n) * ju
 
+@BilinearForm
+def nitscheform(u, v, p):
+    h = p.h
+    n = p.n
+    return u * v / (alpha * h) - dot(grad(u), n) * v - dot(grad(v), n) * u
+
 A = asm(laplace, ib)
 B = asm(dgform, fb, fb)
-C = asm(dgform, bb)
+C = asm(nitscheform, bb)
 b = asm(unit_load, ib)
 
 x = solve(A + B + C, b)

docs/examples/ex40.py

@@ -33,7 +33,7 @@ def laplace(u, ut, v, vt, w):
 def hdg(u, ut, v, vt, w):
     from skfem.helpers import grad, dot
     # outwards normal
-    n = w.n * w.sign
+    n = w.n * (-1.) ** w.idx[0]
     return dot(n, grad(u)) * (vt - v) + dot(n, grad(v)) * (ut - u)\
         + 1e1 / w.h * (ut - u) * (vt - v)
 

docs/howto.rst

@@ -233,3 +233,63 @@ the same thing:
            4.59225774e-16, -4.41713127e-16,  4.63704316e-16,  1.25333771e-16,
            6.12372436e-01,  1.58113883e-01,  6.12372436e-01,  1.58113883e-01,
            0.00000000e+00,  0.00000000e+00,  0.00000000e+00,  0.00000000e+00])
+
+Assembling jump terms
+=====================
+
+The shorthand :func:`skfem.assembly.asm`
+supports special syntax for assembling the same form over a list or lists of
+bases and summing the result.  This can be helpful, e.g., when assembling the
+so-called "jump terms" used in penalty or discontinuous Galerkin methods,
+or when assembling the same form on mixed meshes with different element types.
+
+Consider the form
+
+.. math::
+
+   b(u,v) = \sum_{E \in \mathcal{E}_h} \int_{E} [u][v]\,\mathrm{d}s
+
+where :math:`\mathcal{E}_h` is the set of interior facets of a mesh
+and :math:`[u]` is the jump in the value of :math:`u` over the facet
+:math:`E`.
+We have
+:math:`[u] = u_1 - u_2` and :math:`[v] = v_1 - v_2`
+where the subscript denotes the value of the function restricted to one of the
+elements sharing a facet.  The form can be split as
+
+.. math::
+
+   b(u,v) = \sum_{E \in \mathcal{E}_h} \left(\int_{E} u_1 v_1\,\mathrm{d}s - \int_{E} u_1 v_2\,\mathrm{d}s - \int_{E} u_2 v_1\,\mathrm{d}s + \int_{E} u_2 v_2\,\mathrm{d}s\right)
+
+and normally we would assemble all four forms separately.
+
+Suppose a list of bases is provided during a call to :func:`skfem.assembly.asm`:
+
+.. doctest::
+
+   >>> import skfem as fem
+   >>> m = fem.MeshTri()
+   >>> e = fem.ElementTriP0()
+   >>> bases = [fem.InteriorFacetBasis(m, e, side=k) for k in [0, 1]]
+   >>> jumpform = fem.BilinearForm(lambda u, v, p: (-1) ** sum(p.idx) * u * v)
+   >>> fem.asm(jumpform, bases, bases).toarray()
+   array([[ 1.41421356, -1.41421356],
+          [-1.41421356,  1.41421356]])
+
+This is more or less equivalent to the following nested for-loop:
+
+.. doctest::
+
+   >>> import skfem as fem
+   >>> import numpy as np
+   >>> m = fem.MeshTri()
+   >>> e = fem.ElementTriP0()
+   >>> bases = [fem.InteriorFacetBasis(m, e, side=k) for k in [0, 1]]
+   >>> jumpform = fem.BilinearForm(lambda u, v, p: (-1) ** sum(p._idx) * u * v)
+   >>> A = np.zeros((2, 2))
+   >>> for idx1 in range(2):
+   ...     for idx2 in range(2):
+   ...         A += fem.asm(jumpform, bases[idx1], bases[idx2], _idx=(idx1, idx2)).toarray()
+   >>> A
+   array([[ 1.41421356, -1.41421356],
+          [-1.41421356,  1.41421356]])

skfem/assembly/__init__.py

@@ -70,8 +70,9 @@ def asm(form: Form,
 
     """
     nargs = [[arg] if not isinstance(arg, list) else arg for arg in args]
-    out = sum(map(lambda bases: form.coo_data(*bases, **kwargs),
-                  product(*nargs)))
+    out = sum(map(lambda a: form.coo_data(*a[1], idx=a[0], **kwargs),
+                  zip(product(*(range(len(x)) for x in nargs)),
+                      product(*nargs))))
     assert not isinstance(out, int)
     return out.todefault()
 

skfem/assembly/basis/abstract_basis.py

@@ -30,23 +30,23 @@ class AbstractBasis:
     basis: List[Tuple[DiscreteField, ...]] = []
     X: ndarray
     W: ndarray
-    _sign: float = 1.
     dofs: Dofs
+    _side: int = 1
 
     def __init__(self,
                  mesh: Mesh,
                  elem: Element,
                  mapping: Optional[Mapping] = None,
                  intorder: Optional[int] = None,
                  quadrature: Optional[Tuple[ndarray, ndarray]] = None,
-                 refdom: Type[Refdom] = Refdom):
+                 refdom: Type[Refdom] = Refdom,
+                 dofs: Optional[Dofs] = None):
 
         if mesh.refdom != elem.refdom:
             raise ValueError("Incompatible Mesh and Element.")
 
         self.mapping = mesh._mapping() if mapping is None else mapping
-
-        self.dofs = Dofs(mesh, elem)
+        self.dofs = Dofs(mesh, elem) if dofs is None else dofs
 
         # global degree-of-freedom location
         try:

skfem/assembly/basis/boundary_facet_basis.py

@@ -10,6 +10,7 @@
 
 from .abstract_basis import AbstractBasis
 from .cell_basis import CellBasis
+from ..dofs import Dofs
 
 
 class BoundaryFacetBasis(AbstractBasis):
@@ -22,7 +23,8 @@ def __init__(self,
                  intorder: Optional[int] = None,
                  quadrature: Optional[Tuple[ndarray, ndarray]] = None,
                  facets: Optional[ndarray] = None,
-                 _side: int = 0):
+                 _side: int = 0,
+                 dofs: Optional[Dofs] = None):
         """Precomputed global basis on boundary facets.
 
         Parameters
@@ -42,22 +44,25 @@ def __init__(self,
             Optional tuple of quadrature points and weights.
         facets
             Optional subset of facet indices.
+        dofs
+            Optional :class:`~skfem.assembly.Dofs` object.
 
         """
         super(BoundaryFacetBasis, self).__init__(mesh,
                                                  elem,
                                                  mapping,
                                                  intorder,
                                                  quadrature,
-                                                 mesh.brefdom)
+                                                 mesh.brefdom,
+                                                 dofs)
 
         # facets where the basis is evaluated
         if facets is None:
             self.find = np.nonzero(self.mesh.f2t[1] == -1)[0]
         else:
             self.find = facets
         self.tind = self.mesh.f2t[_side, self.find]
-        self._sign = 1 - 2. * _side
+        self._side = _side  # for debugging
 
         # boundary refdom to global facet
         x = self.mapping.G(self.X, find=self.find)