Add stiffness matrix demo (#94)

* Add stiffness matrix demo

* flake

Author: mscroggs

demo/stiffness_matrix.py

@@ -0,0 +1,20 @@
+import symfem
+from symfem.core.vectors import vdot
+from symfem.core.calculus import grad
+
+matrix = [[0 for i in range(4)] for j in range(4)]
+
+vertices = [(0, 0), (1, 0), (0, 1), (1, 1)]
+triangles = [[0, 1, 2], [1, 3, 2]]
+
+element = symfem.create_element("triangle", "Lagrange", 1)
+for triangle in triangles:
+    vs = [vertices[i] for i in triangle]
+    ref = symfem.create_reference("triangle", vertices=vs)
+    basis = element.map_to_cell(vs)
+    for test_i, test_f in zip(triangle, basis):
+        for trial_i, trial_f in zip(triangle, basis):
+            integrand = vdot(grad(test_f, 2), grad(trial_f, 2))
+            matrix[test_i][trial_i] += ref.integral(integrand)
+
+print(matrix)

docs/demos/custom_element.rst

@@ -2,5 +2,8 @@
 Defining a custom element
 #########################
 
+This demo shows how a custom element with custom functionals can be
+defined in Symfem.
+
 .. literalinclude:: ../../demo/custom_element.py
   :language: Python

docs/demos/index.rst

@@ -5,5 +5,6 @@ Symfem demos
 .. toctree::
    :titlesonly:
 
+   stiffness_matrix
    custom_element
 

docs/demos/stiffness_matrix.rst

@@ -0,0 +1,9 @@
+############################
+Computing a stiffness matrix
+############################
+
+This demo shows how a stiffness matrix over a simple mesh can be
+computed using Symfem.
+
+.. literalinclude:: ../../demo/stiffness_matrix.py
+  :language: Python

test/test_stiffness_matrix.py

@@ -0,0 +1,33 @@
+import symfem
+import sympy
+from symfem.core.vectors import vdot
+from symfem.core.calculus import grad
+
+
+def test_stiffness_matrix():
+    vertices = [(0, 0), (1, 0), (0, 1), (1, 1)]
+    triangles = [[0, 1, 2], [1, 3, 2]]
+
+    matrix = [[0 for i in vertices] for j in vertices]
+
+    element = symfem.create_element("triangle", "Lagrange", 1)
+    for triangle in triangles:
+        vs = [vertices[i] for i in triangle]
+        ref = symfem.create_reference("triangle", vertices=vs)
+        basis = element.map_to_cell(vs)
+        for test_i, test_f in zip(triangle, basis):
+            for trial_i, trial_f in zip(triangle, basis):
+                integrand = vdot(grad(test_f, 2), grad(trial_f, 2))
+                matrix[test_i][trial_i] += ref.integral(integrand)
+
+    half = sympy.Rational(1, 2)
+    actual_matrix = [
+        [1, -half, -half, 0],
+        [-half, 1, 0, -half],
+        [-half, 0, 1, -half],
+        [0, -half, -half, 1]
+    ]
+
+    for row1, row2 in zip(matrix, actual_matrix):
+        for i, j in zip(row1, row2):
+            assert i == j