From 0a455b044f4e0edda8c63bb4922933f747d8379c Mon Sep 17 00:00:00 2001
From: Jukka Aho <jukka.aho@kapsi.fi>
Date: Thu, 13 Oct 2016 00:59:38 +0300
Subject: [PATCH] with_mass_matrix keyword for assemble

---
 src/solvers.jl | 28 ++++++++++++++++++----------
 1 file changed, 18 insertions(+), 10 deletions(-)

diff --git a/src/solvers.jl b/src/solvers.jl
index 5e8b66c..35a0a56 100644
--- a/src/solvers.jl
+++ b/src/solvers.jl
@@ -348,23 +348,31 @@ function solve!(solver::Solver; empty_assemblies_before_solution=true, show_info
 end
 
 """ Default assembler for solver. """
-function assemble!(solver::Solver; show_info=true, timing=true)
+function assemble!(solver::Solver; show_info=true, timing=true, with_mass_matrix=false)
     show_info && info("Assembling problems ...")
-    t0 = Base.time()
-    assembly_times = Dict()
-    nproblems = 0
-    ndofs = 0
-    for problem in solver.problems
+
+    function do_assemble(problem)
         t00 = Base.time()
         empty!(problem.assembly)
         assemble!(problem, solver.time)
-        nproblems += 1
+        if with_mass_matrix && is_field_problem(problem)
+            assemble!(problem, solver.time, Val{:mass_matrix})
+        end
+        t11 = Base.time()
+        return t11-t00
+    end
+
+    t0 = Base.time()
+    assembly_times = map(do_assemble, solver.problems)
+    nproblems = length(assembly_times)
+
+    ndofs = 0
+    for problem in solver.problems
         Ks = size(problem.assembly.K, 2)
         Cs = size(problem.assembly.C1, 2)
         ndofs = max(ndofs, Ks, Cs)
-        t11 = Base.time()
-        assembly_times[problem.name] = t11-t00
     end
+
     solver.ndofs = ndofs
     t1 = round(Base.time()-t0, 2)
     show_info && info("Assembled $nproblems problems in $t1 seconds. ndofs = $ndofs.")
@@ -372,7 +380,7 @@ function assemble!(solver::Solver; show_info=true, timing=true)
         info("Assembly times:")
         for (i, problem) in enumerate(solver.problems)
             pn = problem.name
-            pt = round(assembly_times[pn], 2)
+            pt = round(assembly_times[i], 2)
             info("$i $pn $pt")
         end
     end