From d2992f5fe8ef34a0971237d50a3beba6400b1b7f Mon Sep 17 00:00:00 2001
From: Sebastian Pech <sebastian.pech@me.com>
Date: Sun, 10 Nov 2019 15:29:37 +0100
Subject: [PATCH] Add function to properly extract global solution

---
 src/problems.jl | 27 ++++++++++++++++++++-------
 1 file changed, 20 insertions(+), 7 deletions(-)

diff --git a/src/problems.jl b/src/problems.jl
index 2ae7c99..47c4ee7 100644
--- a/src/problems.jl
+++ b/src/problems.jl
@@ -337,25 +337,38 @@ function get_global_solution(problem::Problem, assembly::Assembly)
     end
 end
 
+"""
+    get_element_solution(problem, assembly, element)
+
+Return the global solution (u, la) at `element`s nodes.
+"""
+function get_element_solution(problem::Problem{P}, assembly::Assembly, element::Element) where P<:FieldProblem
+    u = assembly.u[get_gdofs(problem,element)]
+    la = assembly.la[get_gdofs(problem,element)]
+    field_dim = get_unknown_field_dimension(problem)
+    nnodes = div(length(u),field_dim)
+    un  = [(idx = (i-1)*field_dim;  u[idx+1:idx+field_dim]) for i in 1:nnodes]
+    lan = [(idx = (i-1)*field_dim; la[idx+1:idx+field_dim]) for i in 1:nnodes]
+    return un, lan
+end
+
 function update!(problem::Problem{P}, assembly::Assembly, elements::Vector{Element}, time::Float64) where P<:FieldProblem
-    u, la = get_global_solution(problem, assembly)
     field_name = get_unknown_field_name(problem)
     # update solution u for elements
     for element in elements
-        connectivity = get_connectivity(element)
-        update!(element, field_name, time => tuple(u[connectivity]...))
+        u, la = get_element_solution(problem,assembly,element)
+        update!(element, field_name, time => tuple(u...))
     end
 end
 
 function update!(problem::Problem{P}, assembly::Assembly, elements::Vector{Element}, time::Float64) where P<:BoundaryProblem
-    u, la = get_global_solution(problem, assembly)
     parent_field_name = get_parent_field_name(problem) # displacement
     field_name = get_unknown_field_name(problem) # lambda
     # update solution and lagrange multipliers for boundary elements
     for element in elements
-        connectivity = get_connectivity(element)
-        update!(element, parent_field_name, time => tuple(u[connectivity]...))
-        update!(element, field_name, time => tuple(la[connectivity]...))
+        u, la = get_element_solution(problem,assembly,element)
+        update!(element, parent_field_name, time => tuple(u...))
+        update!(element, field_name, time => tuple(la...))
     end
 end