From 12bd763f4b3498f965a6ce1429bdfe2a7d94eac3 Mon Sep 17 00:00:00 2001
From: Christoph Pohl <christophpohl@gmx.com>
Date: Fri, 8 Feb 2019 12:20:39 +0100
Subject: [PATCH] Add DoF extraction function

---
 fenics_helpers/__init__.py             |  1 +
 fenics_helpers/dof_value_extraction.py | 27 ++++++++++++++++
 tests/test_dof_extraction.py           | 43 ++++++++++++++++++++++++++
 3 files changed, 71 insertions(+)
 create mode 100644 fenics_helpers/dof_value_extraction.py
 create mode 100644 tests/test_dof_extraction.py

diff --git a/fenics_helpers/__init__.py b/fenics_helpers/__init__.py
index 9398692..7d90b5b 100644
--- a/fenics_helpers/__init__.py
+++ b/fenics_helpers/__init__.py
@@ -1 +1,2 @@
 from . import timestepping
+from .dof_value_extraction import extract_dof_values
diff --git a/fenics_helpers/dof_value_extraction.py b/fenics_helpers/dof_value_extraction.py
new file mode 100644
index 0000000..2864ee9
--- /dev/null
+++ b/fenics_helpers/dof_value_extraction.py
@@ -0,0 +1,27 @@
+import dolfin
+from ufl.core.expr import Expr
+
+
+def _extract_dof_from_function(f):
+    # Functions have direct access to their DoF values
+    try:
+        values = f.vector().get_local()
+    # subfunctions need to be interpolated
+    except RuntimeError:
+        fspace = f.function_space()
+        try:
+            fspace = fspace.collapse()
+        except RuntimeError:
+            pass
+        values = dolfin.interpolate(f, fspace).vector().get_local()
+
+    return values
+
+
+def extract_dof_values(expression):
+    if isinstance(expression, dolfin.Function):
+        return _extract_dof_from_function(expression)
+    if isinstance(expression, Expr):
+        return dolfin.project(expression).vector().get_local()
+    else:
+        raise RuntimeError("Can't extract values from {}".format(type(expression)))
diff --git a/tests/test_dof_extraction.py b/tests/test_dof_extraction.py
new file mode 100644
index 0000000..6227641
--- /dev/null
+++ b/tests/test_dof_extraction.py
@@ -0,0 +1,43 @@
+from context import fenics_helpers as fh
+import unittest
+import dolfin
+import numpy as np
+
+
+class TestDofExtraction(unittest.TestCase):
+    def setUp(self):
+        mesh = dolfin.UnitIntervalMesh(100)
+        P1 = dolfin.FiniteElement("P", dolfin.interval, 1)
+        fs = dolfin.FunctionSpace(mesh, P1)
+        self.u = dolfin.Function(fs)
+
+        mixed_element = dolfin.MixedElement([P1, P1, P1])
+        mixed_fs = dolfin.FunctionSpace(mesh, mixed_element)
+        self.mixed_u = dolfin.Function(mixed_fs)
+
+    def test_function(self):
+        u_vec = fh.extract_dof_values(self.u)
+        self._is_numpy_array_of_size(u_vec, 101)
+
+    def test_subspace_function(self):
+        u1, u2, u3 = self.mixed_u.split()
+        u2_values = fh.extract_dof_values(u2)
+        self._is_numpy_array_of_size(u2_values, 101)
+
+    def test_expression(self):
+        my_expr = self.u ** 2 + self.u + dolfin.sin(self.u) + 73.0
+        my_expr_values = fh.extract_dof_values(my_expr)
+        self._is_numpy_array_of_size(my_expr_values, 101)
+
+    def test_vector_expression(self):
+        my_expr = dolfin.as_vector([self.u ** 2, self.u, dolfin.sin(self.u)])
+        my_expr_values = fh.extract_dof_values(my_expr)
+        self._is_numpy_array_of_size(my_expr_values, 3 * 101)
+
+    def _is_numpy_array_of_size(self, array, size):
+        self.assertIsInstance(array, np.ndarray)
+        self.assertEqual(array.size, size)
+
+
+if __name__ == "__main__":
+    unittest.main()