Repository: kinnala/scikit-fem

Issue:
overwrite_a: bool = False
In investigating the behaviour of the new #596 `skfem.utils.enforce` function [at the end of last week](https://github.com/kinnala/scikit-fem/issues/591#issuecomment-816408235), I wasted a couple of hours tracking down discrepancies with condensing and penalizing Dirichlet conditions because of not realizing that in a call like 
```python
L0, M0 = enforce(L, M, D=basis.find_dofs())
```
`L` and `M` would also be mutated, though I do now read that

https://github.com/kinnala/scikit-fem/blob/13f5630bd0077a1484a6ee392c219b4570dcb591/skfem/utils.py#L290-L293

I had been assuming that as discussed in #591 on [2021-03-25](https://github.com/kinnala/scikit-fem/issues/591#issuecomment-806409380)

> it should have the same signature as `condense`

I understand that ([ibid](https://github.com/kinnala/scikit-fem/issues/591#issue-840495893))

> Advantages of this approach include the fact that neither `A` nor `b` need to be rewritten in memory.

but I don't think any of the NumPy and SciPy routines that scikit-fem is built on mutate their arguments, and other functions from those libraries that do, for example, the dense [`scipy.linalg.solve`](https://docs.scipy.org/doc/scipy/reference/generated/scipy.linalg.solve.html#scipy.linalg.solve) only does so if `overwrite_a` or `overwrite_b` is `True` and both are `False` by default.

I think it does happen fairly often that an assembled finite element matrix is reused for the different things.  Of course mutation can be avoided in such cases by copying first, but only if one realizes that there is a threat of mutation.  Would it be safer to introduce overwrite-flags that are `False` by default?

I had a quick look at the implementation of `skfem.utils.enforce` and it doesn't look as though this would be too intrusive.

https://github.com/kinnala/scikit-fem/blob/13f5630bd0077a1484a6ee392c219b4570dcb591/skfem/utils.py#L326-L331

I think the assembly of `idx` and `d` would be the same but one only modifies `A` if `overwrite_A` and otherwise returns a fresh matrix with the condition enforced.

Comments:

Patch:
diff --git a/skfem/utils.py b/skfem/utils.py
index 315b5a80..e31e68af 100644
--- a/skfem/utils.py
+++ b/skfem/utils.py
@@ -22,12 +22,10 @@ from skfem.element import ElementVector
 Solution = Union[ndarray, Tuple[ndarray, ndarray]]
 LinearSolver = Callable[..., ndarray]
 EigenSolver = Callable[..., Tuple[ndarray, ndarray]]
-EnforcedSystem = Union[spmatrix,
-                       Tuple[spmatrix, ndarray],
-                       Tuple[spmatrix, spmatrix]]
-CondensedSystem = Union[spmatrix,
-                        Tuple[spmatrix, ndarray],
-                        Tuple[spmatrix, spmatrix],
+LinearSystem = Union[spmatrix,
+                     Tuple[spmatrix, ndarray],
+                     Tuple[spmatrix, spmatrix]]
+CondensedSystem = Union[LinearSystem,
                         Tuple[spmatrix, ndarray, ndarray],
                         Tuple[spmatrix, ndarray, ndarray, ndarray],
                         Tuple[spmatrix, spmatrix, ndarray, ndarray]]
@@ -248,7 +246,7 @@ def _flatten_dofs(S: Optional[DofsCollection]) -> Optional[ndarray]:
 
 
 def _init_bc(A: spmatrix,
-             b: Optional[ndarray] = None,
+             b: Optional[Union[ndarray, spmatrix]] = None,
              x: Optional[ndarray] = None,
              I: Optional[DofsCollection] = None,
              D: Optional[DofsCollection] = None) -> Tuple[Optional[ndarray],
@@ -284,7 +282,8 @@ def enforce(A: spmatrix,
             x: Optional[ndarray] = None,
             I: Optional[DofsCollection] = None,
             D: Optional[DofsCollection] = None,
-            diag: float = 1.) -> EnforcedSystem:
+            diag: float = 1.,
+            overwrite: bool = False) -> LinearSystem:
     r"""Enforce degrees-of-freedom of a linear system.
 
     .. note::
@@ -307,39 +306,46 @@ def enforce(A: spmatrix,
     D
         Specify either this or ``I``: The set of degree-of-freedom indices to
         enforce (rows/diagonal set to zero/one).
+    overwrite
+        Optionally, the original system is both modified (for performance) and
+        returned (for compatibility with :func:`skfem.utils.solve`).  By
+        default, ``False``.
 
     Returns
     -------
-    EnforcedSystem
+    LinearSystem
         A linear system with the enforced rows/diagonals set to zero/one.
 
     """
     b, x, I, D = _init_bc(A, b, x, I, D)
 
+    Aout = A if overwrite else A.copy()
+
     # set rows on lhs to zero
-    start = A.indptr[D]
-    stop = A.indptr[D + 1]
+    start = Aout.indptr[D]
+    stop = Aout.indptr[D + 1]
     count = stop - start
     idx = np.ones(count.sum(), dtype=np.int64)
     idx[np.cumsum(count)[:-1]] -= count[:-1]
     idx = np.repeat(start, count) + np.cumsum(idx) - 1
-    A.data[idx] = 0.
+    Aout.data[idx] = 0.
 
     # set diagonal value
-    d = A.diagonal()
+    d = Aout.diagonal()
     d[D] = diag
-    A.setdiag(d)
+    Aout.setdiag(d)
 
     if b is not None:
         if isinstance(b, spmatrix):
-            # eigenvalue problem
-            b = enforce(b, D=D, diag=0.)
+            # mass matrix (eigen- or initial value problem)
+            bout = enforce(b, D=D, diag=0., overwrite=overwrite)
         else:
             # set rhs to the given value
-            b[D] = x[D]
-        return A, b
+            bout = b if overwrite else b.copy()
+            bout[D] = x[D]
+        return Aout, bout
 
-    return A
+    return Aout
 
 
 def condense(A: spmatrix,


Test patch:
diff --git a/tests/test_utils.py b/tests/test_utils.py
index 48dcb86a..f5c11c27 100644
--- a/tests/test_utils.py
+++ b/tests/test_utils.py
@@ -44,10 +44,13 @@ class TestEnforce(TestCase):
         M = mass.assemble(basis)
         D = m.boundary_nodes()
 
-        assert_almost_equal(enforce(A, D=D).todense(), np.eye(A.shape[0]))
-        assert_almost_equal(enforce(M, D=D, diag=0.).todense(),
+        assert_almost_equal(enforce(A, D=D).toarray(), np.eye(A.shape[0]))
+        assert_almost_equal(enforce(M, D=D, diag=0.).toarray(),
                             np.zeros(M.shape))
 
+        enforce(A, D=D, overwrite=True)
+        assert_almost_equal(A.toarray(), np.eye(A.shape[0]))
+
 
 if __name__ == '__main__':
     unittest.main()