Repository: kinnala/scikit-fem

Issue:
penalizing Dirichlet conditions
Following the introduction of `skfem.utils.enforce` in #591 and #596 and subsequent investigation which involved the penalty method ([2021-04-09](https://github.com/kinnala/scikit-fem/issues/591#issuecomment-816408235)) for a very special example, I realized that an implementation of the `skfem.utils.penalize` function would be quite similar.  It would add a large positive number to the diagonal elements of the stiffness matrix corresponding to the Dirichlet degrees of freedom and in the static case multiply the corresponding entries of the right-hand side vector by the same factor.  The mass matrix would be left untouched in the unsteady or eigenvalue cases.

Possibly applications of this besides pedagogical might follow from the feature that it would act very like `enforce` but not destroy symmetry; e.g. as Ern & Guermond (2004, &sect;8.4.4 &lsquo;Penalty method&rsquo;) put it:

> This technique, which is very simple to implement, is useful when one wishes to utilize a direct solution method based on the symmetry of the matrix such as the Choleski or the LDL<sup>T</sup> factorization

Though I don't think we have access to the latter and haven't used the former since `sksparse.cholmod` was dumped #447.

I'll take a stab at this, probably with overwriting turned off by default. #620 

Comments:

Patch:
diff --git a/README.md b/README.md
index 1c4ad3ae..4f0400a4 100644
--- a/README.md
+++ b/README.md
@@ -212,6 +212,8 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0
 
 ### Unreleased
 
+- Added: `utils.penalize`, an alternative to `condense` and `enforce` for
+  essential boundary conditions
 - Added: `InteriorBasis.point_source`, with `ex38`
 - Fixed: `MeshLine1.element_finder` 
 
diff --git a/docs/examples/ex19.py b/docs/examples/ex19.py
index 1440373c..1777fd74 100644
--- a/docs/examples/ex19.py
+++ b/docs/examples/ex19.py
@@ -63,33 +63,21 @@ M = asm(mass, basis)
 dt = .01
 print('dt =', dt)
 theta = 0.5                     # Crank–Nicolson
-A = M + theta * L * dt
-B = M - (1 - theta) * L * dt
+L0, M0 = penalize(L, M, D=basis.find_dofs())
+A = M0 + theta * L0 * dt
+B = M0 - (1 - theta) * L0 * dt
 
-boundary = basis.find_dofs()
-interior = basis.complement_dofs(boundary)
-
-# transpose as splu prefers CSC
-backsolve = splu(condense(A, D=boundary, expand=False).T).solve
+backsolve = splu(A.T).solve  # .T as splu prefers CSC
 
 u_init = (np.cos(np.pi * mesh.p[0, :] / 2 / halfwidth[0])
           * np.cos(np.pi * mesh.p[1, :] / 2 / halfwidth[1]))
 
 
-def step(t: float,
-         u: np.ndarray) -> Tuple[float, np.ndarray]:
-    u_new = np.zeros_like(u)              # zero Dirichlet conditions
-    _, b1 = condense(csr_matrix(A.shape),  # ignore condensed matrix
-                     B @ u, u_new, D=boundary, expand=False)
-    u_new[interior] = backsolve(b1)
-    return t + dt, u_new
-
-
 def evolve(t: float,
            u: np.ndarray) -> Iterator[Tuple[float, np.ndarray]]:
 
     while np.linalg.norm(u, np.inf) > 2**-3:
-        t, u = step(t, u)
+        t, u = t + dt, backsolve(B @ u)
         yield t, u
 
 
diff --git a/skfem/__init__.py b/skfem/__init__.py
index 60af8049..5cad2124 100644
--- a/skfem/__init__.py
+++ b/skfem/__init__.py
@@ -20,6 +20,7 @@ __all__ = all_mesh + all_assembly + all_element + [  # noqa
     'build_pc_diag',
     'condense',
     'enforce',
+    'penalize',
     'project',
     'projection',
     'solve',
diff --git a/skfem/utils.py b/skfem/utils.py
index e31e68af..8c5343a1 100644
--- a/skfem/utils.py
+++ b/skfem/utils.py
@@ -288,8 +288,9 @@ def enforce(A: spmatrix,
 
     .. note::
 
-        The original system is both modified (for performance) and returned
-        (for compatibility with :func:`skfem.utils.solve`).
+        The original system is both returned
+        (for compatibility with :func:`skfem.utils.solve`) and optionally (if
+        `overwrite`) modified (for performance).
 
     Parameters
     ----------
@@ -348,6 +349,65 @@ def enforce(A: spmatrix,
     return Aout
 
 
+def penalize(A: spmatrix,
+             b: Optional[Union[ndarray, spmatrix]] = None,
+             x: Optional[ndarray] = None,
+             I: Optional[DofsCollection] = None,
+             D: Optional[DofsCollection] = None,
+             epsilon: Optional[float] = None,
+             overwrite: bool = False) -> LinearSystem:
+    r"""Penalize degrees-of-freedom of a linear system.
+
+    Parameters
+    ----------
+    A
+        The system matrix
+    b
+        Optionally, the right hand side vector.
+    x
+        The values of the penalized degrees-of-freedom. If not given, assumed
+        to be zero.
+    I
+        Specify either this or ``D``: The set of degree-of-freedom indices to
+        solve for.
+    D
+        Specify either this or ``I``: The set of degree-of-freedom indices to
+        enforce (rows/diagonal set to zero/one).
+    epsilon
+        Very small value, the reciprocal of which penalizes deviations from
+        the Dirichlet condition
+    overwrite
+        Optionally, the original system is both modified (for performance) and
+        returned (for compatibility with :func:`skfem.utils.solve`).  By
+        default, ``False``.
+
+    Returns
+    -------
+    LinearSystem
+        A linear system with the penalized diagonal and RHS entries set to
+        very large values, 1/epsilon and x/epsilon, respectively.
+
+    """
+    b, x, I, D = _init_bc(A, b, x, I, D)
+
+    Aout = A if overwrite else A.copy()
+
+    d = Aout.diagonal()
+    if epsilon is None:
+        epsilon = 1e-10 / np.linalg.norm(d[D], np.inf)
+    d[D] = 1/epsilon
+    Aout.setdiag(d)
+
+    if b is None:
+        return Aout
+
+    bout = b if overwrite else b.copy()
+    # Nothing needs doing for mass matrix, but RHS vector needs penalty factor
+    if not isinstance(b, spmatrix):
+        bout[D] = x[D] / epsilon
+    return Aout, bout
+
+
 def condense(A: spmatrix,
              b: Optional[Union[ndarray, spmatrix]] = None,
              x: Optional[ndarray] = None,


Test patch:
diff --git a/tests/test_manufactured.py b/tests/test_manufactured.py
index 0eeb4c4a..1e0a90d2 100644
--- a/tests/test_manufactured.py
+++ b/tests/test_manufactured.py
@@ -1,8 +1,11 @@
 """Solve problems that have manufactured solutions."""
 
+from skfem.utils import penalize
 from unittest import TestCase
 from pathlib import Path
 
+import pytest
+
 import numpy as np
 from skfem import (LinearForm, Functional, asm, condense, solve, projection,
                    enforce)
@@ -283,12 +286,6 @@ class SolveCirclePoissonTet(SolveCirclePoisson):
         return self.mesh_type.init_ball().scaled(0.5)
 
 
-class SolveCirclePoissonTet2(SolveCirclePoissonTet):
-
-    mesh_type = MeshTet2
-    element_type = ElementTetP2
-
-
 class SolveCirclePoissonTet2(SolveCirclePoisson):
 
     mesh_type = MeshTet2
@@ -297,49 +294,47 @@ class SolveCirclePoissonTet2(SolveCirclePoisson):
     maxval = 0.0405901240018571
 
 
-class SolveInhomogeneousLaplace(TestCase):
+@pytest.mark.parametrize(
+    "mesh_elem", [(MeshTri, ElementTriP2()), (MeshQuad, ElementQuad2())]
+)
+@pytest.mark.parametrize("impose", [enforce, penalize])
+def test_solving_inhomogeneous_laplace(mesh_elem, impose):
     """Adapted from example 14."""
 
-    mesh = MeshTri
-    elem = ElementTriP2()
-
-    def runTest(self):
-        m = self.mesh().refined(4)
-        basis = InteriorBasis(m, self.elem)
-        boundary_basis = FacetBasis(m, self.elem)
-        boundary_dofs = boundary_basis.get_dofs().flatten()
-
+    mesh, elem = mesh_elem
 
-        def dirichlet(x):
-            """return a harmonic function"""
-            return ((x[0] + 1.j * x[1]) ** 2).real
+    m = mesh().refined(4)
+    basis = InteriorBasis(m, elem)
+    boundary_basis = FacetBasis(m, elem)
+    boundary_dofs = boundary_basis.get_dofs().flatten()
 
+    def dirichlet(x):
+        """return a harmonic function"""
+        return ((x[0] + 1.j * x[1]) ** 2).real
 
-        u = basis.zeros()
-        A = laplace.assemble(basis)
-        u[boundary_dofs] = projection(dirichlet,
-                                      boundary_basis,
-                                      I=boundary_dofs)
-        u = solve(*enforce(A, x=u, D=boundary_dofs))
+    u = basis.zeros()
+    A = laplace.assemble(basis)
+    u[boundary_dofs] = projection(dirichlet,
+                                  boundary_basis,
+                                  I=boundary_dofs)
+    u = solve(*impose(A, x=u, D=boundary_dofs))
 
 
-        @Functional
-        def gradu(w):
-            gradu = w['sol'].grad
-            return dot(gradu, gradu)
+    @Functional
+    def gradu(w):
+        gradu = w['sol'].grad
+        return dot(gradu, gradu)
 
-        self.assertAlmostEqual(
-            gradu.assemble(basis, sol=basis.interpolate(u)),
-            8 / 3,
-            delta=1e-10,
-        )
+    np.testing.assert_almost_equal(
+        gradu.assemble(basis, sol=basis.interpolate(u)),
+        8 / 3,
+        decimal=9
+    )
 
 
-class SolveInhomogeneousLaplaceQuad(SolveInhomogeneousLaplace):
-
-    mesh = MeshQuad
-    elem = ElementQuad2()
-
+if __name__ == "__main__":
+    import pytest
+    import unittest
 
-if __name__ == '__main__':
-    main()
+    unittest.main()
+    pytest.main()