From 6ae28fc6894966f087968e12f2f20ce6b1c712ba Mon Sep 17 00:00:00 2001
From: Jukka Aho <jukka.aho@kapsi.fi>
Date: Thu, 12 Nov 2015 21:54:46 +0200
Subject: [PATCH] rewrite assembly, see #69. a lot of tests probably fail but
 the most important ones pass

---
 src/JuliaFEM.jl               |   1 +
 src/assembly.jl               |  60 +-------
 src/dirichlet.jl              |  12 +-
 src/elasticity.jl             |   4 +
 src/equations.jl              | 256 +++++++++++++++-------------------
 src/heat.jl                   |  57 ++++----
 src/solvers.jl                |  41 +++---
 src/sparse.jl                 |  11 ++
 test/test_global_assembly.jl  |  31 ++--
 test/test_heat.jl             |  37 +++--
 test/test_potential_energy.jl | 151 +++++++++-----------
 test/test_virtual_work.jl     |  26 ++--
 12 files changed, 317 insertions(+), 370 deletions(-)

diff --git a/src/JuliaFEM.jl b/src/JuliaFEM.jl
index 474a8e5..c9b8719 100644
--- a/src/JuliaFEM.jl
+++ b/src/JuliaFEM.jl
@@ -63,6 +63,7 @@ include("lagrange.jl") # Lagrange elements
 
 ### EQUATIONS ###
 include("integrate.jl")  # default integration points for elements
+include("sparse.jl")
 include("equations.jl")
 include("problems.jl")
 
diff --git a/src/assembly.jl b/src/assembly.jl
index 5f82b46..0bbbe3b 100644
--- a/src/assembly.jl
+++ b/src/assembly.jl
@@ -3,61 +3,9 @@
 
 # Functions to handle global assembly of problem
 
-""" Global assembly. """
-type GlobalAssembly <: Assembly
-    ndofs :: Int
-    mass_matrix :: SparseMatrixCSC
-    stiffness_matrix :: SparseMatrixCSC
-    force_vector :: SparseMatrixCSC
-end
-
-""" Initialize global assembly of size ndofs. """
-function initialize_global_assembly(ndofs::Int=1)
-    mass_matrix = spzeros(ndofs, ndofs)
-    stiffness_matrix = spzeros(ndofs, ndofs)
-    force_vector = spzeros(ndofs, 1)
-    return GlobalAssembly(ndofs, mass_matrix, stiffness_matrix, force_vector)
-end
-
-""" Initialize global assembly, get dimension from problem. """
-function initialize_global_assembly(problem::Problem)
-    dim, ndofs = size(problem)
-    return initialize_global_assembly(ndofs)
-end
-
-""" Initialize or empty workspace for global assembly. """
-function initialize_global_assembly!(assembly::GlobalAssembly, problem::Problem)
-    ndofs = prod(size(problem))
-    if ndofs != assembly.ndofs
-        # if problem size changes, automatically initialize new work space
-        assembly.ndofs = ndofs
-        assembly.mass_matrix = spzeros(ndofs, ndofs)
-        assembly.stiffness_matrix = spzeros(ndofs, ndofs)
-        assembly.force_vector = spzeros(ndofs, 1)
-        return
-    end
-    # otherwise, empty workspace ready for next iteration
-    fill!(assembly.mass_matrix, 0.0)
-    fill!(assembly.stiffness_matrix, 0.0)
-    fill!(assembly.force_vector, 0.0)
-    return
-end
-
-""" Calculate global assembly for a problem. """
-function calculate_global_assembly!(assembly::GlobalAssembly, problem::Problem, time::Number=Inf)
-
-    unknown_field_name = get_unknown_field_name(problem)
-    initialize_global_assembly!(assembly, problem) # zero all
-    dim, ndofs = size(problem)
-    info("assembling problem for $unknown_field_name")
-    info("dimension of unknown field: $dim, problem dofs: $ndofs")
-    local_assembly = initialize_local_assembly()
-    for (i, equation) in enumerate(get_equations(problem))
-        calculate_local_assembly!(local_assembly, equation, unknown_field_name, time)
-        conn = get_connectivity(get_element(equation))
-        gdofs = get_gdofs(problem, equation)
-        assembly.mass_matrix[gdofs, gdofs] += local_assembly.mass_matrix
-        assembly.stiffness_matrix[gdofs, gdofs] += local_assembly.stiffness_matrix
-        assembly.force_vector[gdofs] += local_assembly.force_vector
+function assemble!(assembly::Assembly, problem::Problem, time::Number=0.0)
+    empty!(assembly)
+    for equation in get_equations(problem)
+        assemble!(assembly, equation, time, problem)
     end
 end
diff --git a/src/dirichlet.jl b/src/dirichlet.jl
index 68f50df..f80c8b8 100644
--- a/src/dirichlet.jl
+++ b/src/dirichlet.jl
@@ -5,6 +5,10 @@
 
 abstract DirichletEquation <: Equation
 
+function get_unknown_field_name(equation::DirichletEquation)
+    return "reaction force"
+end
+
 ### Dirichlet problem + equations
 
 type DirichletProblem <: BoundaryProblem
@@ -56,19 +60,19 @@ function DBC2D2(element::Seg2)
 end
 Base.size(equation::DBC2D2) = (1, 2)
 
-function calculate_local_assembly!(assembly::LocalAssembly, equation::DirichletEquation, unknown_field_name::ASCIIString, time::Number=0.0, problem=nothing)
-    initialize_local_assembly!(assembly, equation)
+function assemble!(assembly::Assembly, equation::DirichletEquation, time::Number=0.0, problem=nothing)
+    gdofs = get_gdofs(equation)
     element = get_element(equation)
     basis = get_basis(element)
     detJ = det(basis)
     for ip in get_integration_points(equation)
         w = ip.weight * detJ(ip)
         N = basis(ip, time)
-        assembly.stiffness_matrix += w * N'*N
+        add!(assembly.stiffness_matrix, gdofs, gdofs, w*N'*N)
         if !isa(problem, Void)
             X = basis("geometry", ip, time)
             u = problem.field_value(X)
-            assembly.force_vector += w * N'*u
+            add!(assembly.force_vector, gdofs, w*N'*u)
         end
     end
 end
diff --git a/src/elasticity.jl b/src/elasticity.jl
index 0b77984..0d2f5b0 100644
--- a/src/elasticity.jl
+++ b/src/elasticity.jl
@@ -6,6 +6,10 @@
 abstract ElasticityProblem <: Problem
 abstract ElasticityEquation <: Equation
 
+function get_unknown_field_name(equation::ElasticityEquation)
+    return "displacement"
+end
+
 ### Formulation ###
 
 """ Calculate internal energy for elasticity equation.
diff --git a/src/equations.jl b/src/equations.jl
index c62bf02..9c89b80 100644
--- a/src/equations.jl
+++ b/src/equations.jl
@@ -5,203 +5,177 @@
 
 abstract Equation
 
-abstract Assembly
-
-""" Local element assembly. """
-type LocalAssembly <: Assembly
-    ndofs :: Int
-    mass_matrix :: Matrix
-    stiffness_matrix :: Matrix
-    force_vector :: Matrix
-    potential_energy
-    residual_vector :: Vector
+type Assembly
+    mass_matrix :: SparseMatrixIJV
+    stiffness_matrix :: SparseMatrixIJV
+    force_vector :: SparseMatrixIJV
+    lhs :: SparseMatrixIJV
+    rhs :: SparseMatrixIJV
 end
 
-""" Initialize workspace for local matrices for dimension ndofs. """
-function initialize_local_assembly(ndofs::Int=1)
-    mass_matrix = zeros(ndofs, ndofs)
-    stiffness_matrix = zeros(ndofs, ndofs)
-    force_vector = zeros(ndofs, 1)
-    potential_energy = 0.0
-    residual_vector = zeros(ndofs)
-    return LocalAssembly(ndofs, mass_matrix, stiffness_matrix, force_vector,
-                         potential_energy, residual_vector)
+function Assembly()
+    return Assembly(
+        SparseMatrixIJV(),
+        SparseMatrixIJV(),
+        SparseMatrixIJV(),
+        SparseMatrixIJV(),
+        SparseMatrixIJV())
 end
 
-""" Initialize workspace for local matrices, get dimension from equation. """
-function initialize_local_assembly(equation::Equation)
-    ndofs = prod(size(equation))
-    return initialize_local_assembly(ndofs)
+function Base.empty!(assembly::Assembly)
+    empty!(assembly.mass_matrix)
+    empty!(assembly.stiffness_matrix)
+    empty!(assembly.force_vector)
+    empty!(assembly.lhs)
+    empty!(assembly.rhs)
 end
 
-""" Initialize or zero workspace. """
-function initialize_local_assembly!(assembly::LocalAssembly, equation::Equation)
-    ndofs = prod(size(equation))
-    if ndofs != assembly.ndofs
-        # if problem size changes, automatically initialize new work space
-        assembly.ndofs = ndofs
-        assembly.mass_matrix = zeros(ndofs, ndofs)
-        assembly.stiffness_matrix = zeros(ndofs, ndofs)
-        assembly.force_vector = zeros(ndofs, 1)
-        assembly.potential_energy = 0.0
-        assembly.residual_vector = zeros(ndofs)
-        return
-    end
-    # otherwise, empty workspace ready for next iteration
-    fill!(assembly.mass_matrix, 0.0)
-    fill!(assembly.stiffness_matrix, 0.0)
-    fill!(assembly.force_vector, 0.0)
-    assembly.potential_energy = 0.0
-    fill!(assembly.residual_vector, 0.0)
-    return
+function get_mass_matrix
 end
 
-has_mass_matrix(equation::Equation) = false
-function get_mass_matrix(equation::Equation, ip, time=0.0, problem=nothing)
-    get_mass_matrix(equation, ip, time)
-end
-function get_mass_matrix(equation::Equation, ip, time=0.0)
-    get_mass_matrix(equation, ip)
-end
-function get_mass_matrix(equation::Equation, ip)
-    nothing
+function get_stiffness_matrix
 end
 
-has_stiffness_matrix(equation::Equation) = false
-function get_stiffness_matrix(equation::Equation, ip, time=0.0, problem=nothing)
-    get_stiffness_matrix(equation, ip, time)
-end
-function get_stiffness_matrix(equation::Equation, ip, time=0.0)
-    get_stiffness_matrix(equation, ip)
-end
-function get_stiffness_matrix(equation::Equation, ip)
-    nothing
+function get_force_vector
 end
 
-has_force_vector(equation::Equation) = false
-function get_force_vector(equation::Equation, ip, time=0.0, problem=nothing)
-    get_force_vector(equation, ip, time)
-end
-function get_force_vector(equation::Equation, ip, time=0.0)
-    get_force_vector(equation, ip)
-end
-function get_force_vector(equation::Equation, ip)
-    nothing
+function get_potential_energy
 end
 
-has_residual_vector(equation::Equation) = false
-function get_residual_vector(equation::Equation, ip, time=0.0, problem=nothing)
-    get_residual_vector(equation, ip, time)
-end
-function get_residual_vector(equation::Equation, ip, time=0.0)
-    get_residual_vector(equation, ip)
-end
-function get_residual_vector(equation::Equation, ip)
-    nothing
+function get_residual_vector
 end
 
-has_potential_energy(equation::Equation) = false
-function get_potential_energy(equation::Equation, ip, time=0.0, problem=nothing)
-    get_potential_energy(equation, ip, time)
-end
-function get_potential_energy(equation::Equation, ip, time=0.0)
-    get_potential_energy(equation, ip)
-end
-function get_potential_energy(equation::Equation, ip)
-    nothing
+function has_mass_matrix(equation::Equation)
+    default_args = Tuple{typeof(equation), IntegrationPoint, Float64}
+    return method_exists(get_mass_matrix, default_args)
 end
 
-get_element(equation::Equation) = equation.element
-get_integration_points(equation::Equation) = equation.integration_points
+function has_stiffness_matrix(equation::Equation)
+    default_args = Tuple{typeof(equation), IntegrationPoint, Float64}
+    return method_exists(get_stiffness_matrix, default_args)
+end
 
+function has_force_vector(equation::Equation)
+    default_args = Tuple{typeof(equation), IntegrationPoint, Float64}
+    return method_exists(get_force_vector, default_args)
+end
 
-""" Return a local assembly for element. """
-function calculate_local_assembly!(assembly::LocalAssembly, equation::Equation,
-                                   unknown_field_name::ASCIIString, time::Number=0.0,
-                                   problem=nothing)
+function has_potential_energy(equation::Equation)
+    default_args = Tuple{typeof(equation), IntegrationPoint, Float64}
+    return method_exists(get_potential_energy, default_args)
+end
 
-    initialize_local_assembly!(assembly, equation) # zero all
+function has_residual_vector(equation::Equation)
+    default_args = Tuple{typeof(equation), IntegrationPoint, Float64}
+    return method_exists(get_residual_vector, default_args)
+end
+
+function get_element(equation::Equation)
+    return equation.element
+end
+
+function get_integration_points(equation::Equation)
+    return equation.integration_points
+end
+
+function Base.size(equation::Equation, i::Int)
+    return size(equation)[i]
+end
+
+""" Return global degrees of freedom of element in matrix level.
+
+Notes
+-----
+This is calculated from connectivity and equation dimension.
+"""
+
+function get_gdofs(equation::Equation)
+    element = get_element(equation)
+    conn = get_connectivity(element)
+    dim = size(equation, 1)
+    gdofs = vec(vcat([dim*conn'-i for i=dim-1:-1:0]...))
+    return gdofs
+end
+
+""" Assemble element. """
+function assemble!(assembly::Assembly, equation::Equation, time::Number=0.0, problem=nothing)
 
     element = get_element(equation)
+    gdofs = get_gdofs(equation)
     basis = get_basis(element)
     detJ = det(basis)
+    unknown_field_name = get_unknown_field_name(equation)
 
-    # 1. if equations are defined we just integrate them
+    # 1. if equations are defined we just integrate them, without caring how they are done
     if has_mass_matrix(equation) || has_stiffness_matrix(equation) || has_force_vector(equation)
         for ip in get_integration_points(equation)
             s = ip.weight*detJ(ip)
             if has_mass_matrix(equation)
-                assembly.mass_matrix += s*get_mass_matrix(equation, ip, time, problem)
+                add!(assembly.mass_matrix, gdofs, gdofs, s*get_mass_matrix(equation, ip, time))
             end
             if has_stiffness_matrix(equation)
-                assembly.stiffness_matrix += s*get_stiffness_matrix(equation, ip, time, problem)
+                add!(assembly.stiffness_matrix, gdofs, gdofs, s*get_stiffness_matrix(equation, ip, time))
             end
             if has_force_vector(equation)
-                assembly.force_vector += s*get_force_vector(equation, ip, time, problem)
+                add!(assembly.force_vector, gdofs, s*get_force_vector(equation, ip, time))
             end
         end
         # external loads -- if any nodal loads is defined add to force vector
         if haskey(element, "$unknown_field_name nodal load")
-            assembly.force_vector += vec(element["$unknown_field_name nodal load"](time))
+            add!(assembly.force_vector, gdofs, vec(element["$unknown_field_name nodal load"](time)))
         end
     end
 
-    # 2. variational / energy form - user has defined some potential energy / variational form
+    # 2. energy form -- user has defined potential energy W -> min!
     if has_potential_energy(equation)
-        element = get_element(equation)
         field = element[unknown_field_name](time)
-        function potential_energy(data::Vector)
-            # calculate potential energy for some setting. this is needed by forwarddiff
-            assembly.potential_energy = 0.0
+
+        """ Wrapper for potential energy for ForwardDiff. """
+        function calc_W(data::Vector)
+            W = 0.0
             df = similar(field, data)
             # integrate potential energy
             for ip in get_integration_points(equation)
+                s = ip.weight*detJ(ip)
                 dw = get_potential_energy(equation, ip, time; variation=df)
-                assembly.potential_energy += ip.weight * dw * detJ(ip)
+                W += s*dw
             end
             # external energy -- if any nodal loads is defined, decrease from potential energy
             if haskey(element, "$unknown_field_name nodal load")
                 P = element["$unknown_field_name nodal load"](time)
-                assembly.potential_energy -= dot(vec(P), vec(df))
+                W -= dot(vec(P), vec(df))
             end
-            if isa(assembly.potential_energy, Array)
-                return assembly.potential_energy[1]
-            end
-            return assembly.potential_energy
+            return isa(W, Array) ? W[1] : W
         end
-        hessian, allresults = ForwardDiff.hessian(potential_energy, vec(field),
-                                                  AllResults, cache=autodiffcache)
-        assembly.stiffness_matrix += hessian
-        assembly.force_vector -= ForwardDiff.gradient(allresults)  # <--- minus explained in tutorial
-        assembly.potential_energy = ForwardDiff.value(allresults)
-        #info("potential energy of system: $(assembly.potential_energy)")
+
+        hessian, allresults = ForwardDiff.hessian(calc_W, vec(field), AllResults, cache=autodiffcache)
+        add!(assembly.stiffness_matrix, gdofs, gdofs, hessian)
+        add!(assembly.force_vector, gdofs, -ForwardDiff.gradient(allresults))
     end
 
-    # 3. virtual work form - user has defined residual vector δW_int(u,δu) + δW_ext(u,δu) = 0 ∀ v
+    # 3. virtual work -- user has defined some residual r = p - f = 0
     if has_residual_vector(equation)
-        element = get_element(equation)
         field = element[unknown_field_name](time)
-        function residual_vector(data::Vector)
-            fill!(assembly.residual_vector, 0.0)
-            #@debug("field: $field, length = $(size(field))")
-            #@debug("data: $data, size = $(size(data))")
-            #df = similar(field, data)
-            df = Increment(reshape(data, size(equation)...))
-            # integrate W
-            for ip in get_integration_points(equation)
-                dr = get_residual_vector(equation, ip, time; variation=df)
-                assembly.residual_vector += ip.weight*dr*detJ(ip)
-            end
-            # external loads -- if any nodal loads is defined, remove from residual
-            if haskey(element, "$unknown_field_name nodal load")
-                assembly.residual_vector -= vec(element["$unknown_field_name nodal load"](time))
-            end
-            return assembly.residual_vector
-        end
-        jacobian, allresults = ForwardDiff.jacobian(residual_vector, vec(field),
-                                                    AllResults, cache=autodiffcache)
-        assembly.stiffness_matrix += jacobian
-        assembly.force_vector -= ForwardDiff.value(allresults)  # <-- minus explained in tutorial
-    end
 
+        """ Wrapper for virtual work for ForwardDiff. """
+        function calc_R(data::Vector)
+            R = zeros(length(data))
+            df = similar(field, data)
+            # integrate residual vector
+            for ip in get_integration_points(equation)
+                s = ip.weight*detJ(ip)
+                dr = get_residual_vector(equation, ip, time; variation=df)
+                R += s*dr
+            end
+            # external loads -- if any nodal loads is defined, decrease from residual
+            if haskey(element, "$unknown_field_name nodal load")
+                R -= vec(element["$unknown_field_name nodal load"](time))
+            end
+            return R
+        end
+       
+        jacobian, allresults = ForwardDiff.jacobian(calc_R, vec(field), AllResults, cache=autodiffcache)
+        add!(assembly.stiffness_matrix, gdofs, gdofs, jacobian)
+        add!(assembly.force_vector, gdofs, -ForwardDiff.value(allresults))
+    end
 end
diff --git a/src/heat.jl b/src/heat.jl
index daaf3ca..6c56410 100644
--- a/src/heat.jl
+++ b/src/heat.jl
@@ -6,6 +6,10 @@
 abstract HeatProblem <: Problem
 abstract HeatEquation <: Equation
 
+function get_unknown_field_name(equation::HeatEquation)
+    return "temperature"
+end
+
 ### Formulation ###
 
 """ Heat equations.
@@ -32,55 +36,36 @@ References
 https://en.wikipedia.org/wiki/Heat_equation
 
 """
-function calculate_local_assembly!(assembly::LocalAssembly, equation::HeatEquation,
-                                            unknown_field_name::ASCIIString, time::Number=Inf,
-                                            problem=nothing)
-
-    initialize_local_assembly!(assembly, equation)
+function assemble!(assembly::Assembly, equation::HeatEquation, time::Number=0.0, problem=nothing)
 
     element = get_element(equation)
+    gdofs = get_gdofs(equation)
     basis = get_basis(element)
     dbasis = grad(basis)
     detJ = det(basis)
     for ip in get_integration_points(equation)
-        w = ip.weight * detJ(ip)
+        w = ip.weight*detJ(ip)
         N = basis(ip, time)
         if haskey(element, "density")
             rho = basis("density", ip, time)
-            assembly.mass_matrix += w * rho*N'*N
+            add!(assembly.mass_matrix, gdofs, gdofs, w*rho*N'*N)
         end
         if haskey(element, "temperature thermal conductivity")
             dN = dbasis(ip, time)
             k = basis("temperature thermal conductivity", ip, time)
-            assembly.stiffness_matrix += w * k*dN'*dN
+            add!(assembly.stiffness_matrix, gdofs, gdofs, w*k*dN'*dN)
         end
         if haskey(element, "temperature load")
             f = basis("temperature load", ip, time)
-            assembly.force_vector += w * N'*f
+            add!(assembly.force_vector, gdofs, w*N'*f)
         end
         if haskey(element, "temperature flux")
             g = basis("temperature flux", ip, time)
-            assembly.force_vector += w * N'*g
+            add!(assembly.force_vector, gdofs, w*N'*g)
         end
     end
 end
 
-### Problems ###
-
-type PlaneHeatProblem <: HeatProblem
-    unknown_field_name :: ASCIIString
-    unknown_field_dimension :: Int
-    equations :: Array{HeatEquation, 1}
-    element_mapping :: Dict{DataType, DataType}
-end
-
-""" Default constructor for problem takes no arguments. """
-function PlaneHeatProblem()
-    element_mapping = Dict(
-        Quad4 => DC2D4,
-        Seg2 => DC2D2)
-    return PlaneHeatProblem("temperature", 1, [], element_mapping)
-end
 
 ### Equations ###
 
@@ -101,7 +86,7 @@ Base.size(equation::DC2D4) = (1, 4)
 """ Diffusive heat transfer for 2-node linear segment. """
 type DC2D2 <: HeatEquation
     element :: Seg2
-    integration_points :: Array{IntegrationPoint, 1}
+    integration_points :: Vector{IntegrationPoint}
 end
 function DC2D2(element::Seg2)
     integration_points = get_default_integration_points(element)
@@ -112,3 +97,21 @@ function DC2D2(element::Seg2)
 end
 Base.size(equation::DC2D2) = (1, 2)
 
+### Problems ###
+
+type PlaneHeatProblem <: HeatProblem
+    unknown_field_name :: ASCIIString
+    unknown_field_dimension :: Int
+    equations :: Vector{Equation}
+    #element_mapping :: Dict{Element, Equation}
+    # FIXME: Why is not working ^
+    element_mapping :: Dict{Any, Any}
+end
+
+""" Default constructor for problem takes no arguments. """
+function PlaneHeatProblem()
+    element_mapping = Dict(
+        Quad4 => DC2D4,
+        Seg2 => DC2D2)
+    return PlaneHeatProblem("temperature", 1, [], element_mapping)
+end
diff --git a/src/solvers.jl b/src/solvers.jl
index 580e90b..d376360 100644
--- a/src/solvers.jl
+++ b/src/solvers.jl
@@ -9,18 +9,19 @@ abstract Solver
 Solve field equations for single element with some dofs fixed. This can be used
 to test nonlinear element formulations.
 """
-function solve!(equation::Equation, unknown_field_name::ASCIIString,
-                 free_dofs::Array{Int, 1}, time::Number=0.0; 
+function solve!(equation::Equation, free_dofs::Vector{Int}, time::Number=0.0; 
                  max_iterations::Int=10, tolerance::Float64=1.0e-12, dump_matrices::Bool=false)
+    unknown_field_name = get_unknown_field_name(equation)
     element = get_element(equation)
     x0 = element[unknown_field_name](0.0)
     x = zeros(prod(size(equation)))
     dx = fill!(similar(x), 0.0)
-    la = initialize_local_assembly()
+    ass = Assembly()
     for i=1:max_iterations
-        calculate_local_assembly!(la, equation, unknown_field_name)
-        A = la.stiffness_matrix[free_dofs, free_dofs]
-        b = la.force_vector[free_dofs]
+        empty!(ass)
+        assemble!(ass, equation)
+        A = full(ass.stiffness_matrix)[free_dofs, free_dofs]
+        b = full(ass.force_vector)[free_dofs]
         if dump_matrices
             dump(full(A))
             dump(full(b)')
@@ -39,30 +40,34 @@ to test nonlinear element formulations. Dirichlet boundary is assumed to be homo
 and degrees of freedom are eliminated. So if boundary condition is known in nodal
 points and everything is zero this should be quite good.
 """
-function solve!(problem::Problem, free_dofs::Array{Int, 1}, time::Number=1.0;
-                max_iterations::Int=10, tolerance::Float64=1.0e-12, dump_matrices::Bool=false)
+function solve!(problem::Problem, free_dofs::Vector{Int}, time::Number=1.0; max_iterations::Int=10, tolerance::Float64=1.0e-12, dump_matrices::Bool=false)
     info("start solver")
-    ga = initialize_global_assembly(problem)
-    x = zeros(ga.ndofs)
-    dx = fill!(similar(x), 0.0)
+    assembly = Assembly()
+#   x = zeros(ga.ndofs)
+#   dx = fill!(similar(x), 0.0)
+#   FIXME: better.
+    x = nothing
+    dx = nothing
     field_name = get_unknown_field_name(problem)
     dim = get_unknown_field_dimension(problem)
     for i=1:max_iterations
-        info("calculate global assembly")
-        calculate_global_assembly!(ga, problem)
-        info("done")
-        A = ga.stiffness_matrix[free_dofs, free_dofs]
-        b = ga.force_vector[free_dofs]
+        assemble!(assembly, problem, time)
+        A = sparse(assembly.stiffness_matrix)
+        b = sparse(assembly.force_vector)
         if dump_matrices
             dump(full(A))
             dump(full(b)')
         end
-        dx[free_dofs] = lufact(A) \ full(b)
+        if isa(dx, Void)
+            x = zeros(length(b))
+            dx = zeros(length(b))
+        end
+        dx[free_dofs] = lufact(A[free_dofs,free_dofs]) \ full(b)[free_dofs]
         info("Difference in solution norm: $(norm(dx))")
         x += dx
         for equation in get_equations(problem)
             element = get_element(equation)
-            gdofs = get_gdofs(problem, equation)
+            gdofs = get_gdofs(equation)
             data = reshape(full(x[gdofs]), size(equation))
             push!(element[field_name], data)
         end
diff --git a/src/sparse.jl b/src/sparse.jl
index f812ea5..e834800 100644
--- a/src/sparse.jl
+++ b/src/sparse.jl
@@ -14,6 +14,10 @@ function SparseMatrixIJV()
     SparseMatrixIJV([], [], [])
 end
 
+function Base.sparse(A::SparseMatrixIJV)
+    return sparse(A.I, A.J, A.V)
+end
+
 function Base.push!(A::SparseMatrixIJV, I::Int, J::Int, V::Float64)
     push!(A.I, I)
     push!(A.J, J)
@@ -63,3 +67,10 @@ function add!(A::SparseMatrixIJV, dofs1::Vector{Int}, dofs2::Vector{Int}, data::
     append!(A.V, vec(data))
 end
 
+""" Sparse vector version. """
+function add!(A::SparseMatrixIJV, dofs::Vector{Int}, data::Array{Float64})
+    append!(A.I, dofs)
+    append!(A.J, ones(Int, length(dofs)))
+    append!(A.V, vec(data))
+end
+
diff --git a/test/test_global_assembly.jl b/test/test_global_assembly.jl
index c48aced..7318bd3 100644
--- a/test/test_global_assembly.jl
+++ b/test/test_global_assembly.jl
@@ -1,37 +1,42 @@
 # This file is a part of JuliaFEM.
 # License is MIT: see https://github.com/JuliaFEM/JuliaFEM.jl/blob/master/LICENSE.md
 
-module GlobalAssemblyTests
+module AssemblyTests
 
 using JuliaFEM.Test
 using JuliaFEM: Quad4, Seg2, FieldSet, Field, PlaneHeatProblem
-using JuliaFEM: initialize_global_assembly, calculate_global_assembly!
+using JuliaFEM: Assembly, assemble!
 
 """assemble a simple two element problem and solve"""
-function test_asssembly()
+function test_assembly()
+    info("create elements")
     el1 = Quad4([1, 2, 3, 4])
     el1["geometry"] = Vector[[0.0, 0.0], [1.0, 0.0], [1.0, 1.0], [0.0, 1.0]]
     el1["temperature thermal conductivity"] = 6.0
-    el1["temperature load"] = [12.0, 12.0, 12.0, 12.0]
-    el1["density"] = 10
-
+    el1["temperature load"] = 12.0
+    el1["density"] = 36.0
     el2 = Seg2([1, 2])
     el2["geometry"] = Vector[[0.0, 0.0], [1.0, 0.0]]
-
     # Boundary load, linear ramp 0 -> 600 at time 0 -> 1
-    el2["temperature flux"] = FieldSet(Field[Field(0.0, 0.0), Field(1.0, 600.0)])
+    el2["temperature flux"] = ((0.0 => 0.0), (1.0 => 600.0))
+    info("element created")
 
     problem = PlaneHeatProblem()
+    info("problem created. pushing elements")
     push!(problem, el1)
     push!(problem, el2)
 
-    global_assembly = initialize_global_assembly(problem)
-    calculate_global_assembly!(global_assembly, problem)
+    info("creating assembly from equations")
+    assembly = Assembly()
+    assemble!(assembly, problem, 1.0)
+    info("solving")
+
     free_dofs = [1, 2]
-    A = lufact(global_assembly.stiffness_matrix[free_dofs, free_dofs])
-    b = full(global_assembly.force_vector)[free_dofs]
+    A = full(assembly.stiffness_matrix)[free_dofs, free_dofs]
+    b = full(assembly.force_vector)[free_dofs]
     u = A \ b
-    @test isapprox(u, roughly([101.0, 101.0]))
+    info("solution u=$u")
+    @test isapprox(u, [101.0, 101.0])
 end
 
 end
diff --git a/test/test_heat.jl b/test/test_heat.jl
index 56f5e1a..e654cb2 100644
--- a/test/test_heat.jl
+++ b/test/test_heat.jl
@@ -6,13 +6,8 @@
 module HeatTests  # always wrap tests to module ending with "Tests"
 
 using JuliaFEM.Test  # always use JuliaFEM.Test, not Base.Test
+using JuliaFEM: Seg2, Quad4, DC2D4, DC2D2, Assembly, assemble!
 
-using JuliaFEM: Seg2, Quad4, Field, FieldSet, DC2D4,
-                initialize_local_assembly, calculate_local_assembly!,
-                DC2D2
-
-
-"tests on [0x1]x[0x1] domain"
 function test_one_element()  # always start test function with name test_
 
     # volume element
@@ -31,20 +26,34 @@ function test_one_element()  # always start test function with name test_
     # Set constant source f=12 with k=6. Accurate solution is
     # T=1 on free boundary, u(x,y) = -1/6*(1/2*f*x^2 - f*x)
     equation = DC2D4(element)
-    la = initialize_local_assembly()
-    calculate_local_assembly!(la, equation, "temperature")
+    #la = initialize_local_assembly()
+    #calculate_local_assembly!(la, equation, "temperature")
+    assembly = Assembly()
+    assemble!(assembly, equation)
     fdofs = [1, 2]
-    A = la.stiffness_matrix
-    b = la.force_vector
+    A = full(assembly.stiffness_matrix)
+    b = full(assembly.force_vector)
     @test isapprox(A[fdofs, fdofs] \ b[fdofs], [1.0, 1.0])
 
     # Set constant flux g=6 on boundary. Accurate solution is
     # u(x,y) = x which equals T=1 on boundary.
-    boundary_equation = DC2D2(boundary_element);
+    boundary_equation = DC2D2(boundary_element)
+    empty!(assembly)
 
-    calculate_local_assembly!(la, boundary_equation, "temperature")
-    b = la.force_vector
-    @test isapprox(A[fdofs, fdofs] \ b[fdofs], [1.0, 1.0])  # always use @test to test things.
+    time = 1.0
+    assemble!(assembly, equation, time)
+    info("after first element: $(length(assembly.force_vector.V))")
+    info(full(assembly.force_vector)')
+    assemble!(assembly, boundary_equation, time)
+    info("after second element: $(length(assembly.force_vector.V))")
+    info(full(assembly.force_vector)')
+    #calculate_local_assembly!(la, boundary_equation, "temperature")
+    #b = la.force_vector
+    A = full(assembly.stiffness_matrix)
+    b = full(assembly.force_vector)
+    T = A[fdofs, fdofs] \ b[fdofs]
+    info("T = $T")
+    @test isapprox(T, [2.0, 2.0])  # always use @test to test things.
 
 end
 
diff --git a/test/test_potential_energy.jl b/test/test_potential_energy.jl
index 68f45ae..f461376 100644
--- a/test/test_potential_energy.jl
+++ b/test/test_potential_energy.jl
@@ -6,25 +6,27 @@ module ElementTests
 using JuliaFEM.Test
 
 using JuliaFEM
-using JuliaFEM: Equation, Quad4, IntegrationPoint, initialize_local_assembly,
-                get_element, get_basis, grad, calculate_local_assembly!,
-                PlaneHeatProblem, Seg2, HeatEquation, Problem, solve!
+using JuliaFEM: Equation, Quad4, IntegrationPoint, Assembly, assemble!,
+                get_element, get_basis, grad, get_unknown_field_name,
+                PlaneHeatProblem, Seg2, Problem, solve!,
+                get_default_integration_points, Equation
 
+abstract MyEquation <: Equation
 
-""" Diffusive heat transfer for 4-node bilinear element, with a nonlinear source term. """
-type DC2D4NL <: Equation
-    element :: Quad4
-    integration_points :: Array{IntegrationPoint, 1}
+function JuliaFEM.get_unknown_field_name(equation::MyEquation)
+    return "temperature"
 end
 
-function DC2D4NL(element::Quad4, initial_temperature=zeros(4))
-    integration_points = [
-        IntegrationPoint(1.0/sqrt(3.0)*[-1, -1], 1.0),
-        IntegrationPoint(1.0/sqrt(3.0)*[ 1, -1], 1.0),
-        IntegrationPoint(1.0/sqrt(3.0)*[ 1,  1], 1.0),
-        IntegrationPoint(1.0/sqrt(3.0)*[-1,  1], 1.0)]
+""" Diffusive heat transfer for 4-node bilinear element, with a nonlinear source term. """
+type DC2D4NL <: MyEquation
+    element :: Quad4
+    integration_points :: Vector{IntegrationPoint}
+end
+
+function DC2D4NL(element::Quad4)
+    integration_points = get_default_integration_points(element)
     if !haskey(element, "temperature")
-        element["temperature"] = initial_temperature
+        element["temperature"] = zeros(4)
     end
     DC2D4NL(element, integration_points)
 end
@@ -34,17 +36,15 @@ function Base.size(equation::DC2D4NL)
 end
 
 """ Nonlinear flux term. """
-type DC2D2NL <: Equation
+type DC2D2NL <: MyEquation
     element :: Seg2
-    integration_points :: Array{IntegrationPoint, 1}
+    integration_points :: Vector{IntegrationPoint}
 end
 
-function DC2D2NL(element::Seg2, initial_temperature=zeros(2))
-    #integration_points = [
-    #    IntegrationPoint([0.0], 2.0)]
+function DC2D2NL(element::Seg2)
     integration_points = JuliaFEM.line5()
     if !haskey(element, "temperature")
-        element["temperature"] = initial_temperature
+        element["temperature"] = zeros(2)
     end
     DC2D2NL(element, integration_points)
 end
@@ -63,34 +63,23 @@ function JuliaFEM.get_potential_energy(equation::DC2D4NL, ip, time; variation=no
     c = basis("temperature nonlinearity coefficient", ip, time)
     gradT = grad(basis)("temperature", ip, time, variation)
     Wint = (k + c*T) * 1/2*vecdot(gradT, gradT)
-    #Wint = k*1/2*vecdot(gradT, gradT)
     Wext = f*T
-    #Wext = 0.0
     return Wint - Wext
 end
 
-function JuliaFEM.has_potential_energy(eq::DC2D4NL)
-    return true
-end
-
 function JuliaFEM.get_potential_energy(equation::DC2D2NL, ip, time; variation=nothing)
     element = get_element(equation)
     basis = get_basis(element)
     T = basis("temperature", ip, time, variation)[1]
-    Wint = 0.0
-    sig = 5.7e-8
-    eps = basis("emissivity", ip, time)[1]
     T_ext = basis("temperature external", ip, time)[1]
-    q0 = eps*sig*((T_ext+273.15)^4 - (T+273.15)^4)
+    coeff = basis("temperature coefficient", ip, time)[1]
+    q0 = coeff*(T_ext^4 - T^4)
+    Wint = 0.0
     Wext = q0*T
     W = Wint - Wext
     return W
 end
 
-function JuliaFEM.has_potential_energy(eq::DC2D2NL)
-    return true
-end
-
 function test_potential_energy_method()
 
     # create model -- start
@@ -99,36 +88,37 @@ function test_potential_energy_method()
     element["temperature thermal conductivity"] = 6.0
     element["temperature load"] = [0.0, 0.0, 0.0, 0.0]
     element["temperature nodal load"] = [3.0, 3.0, 0.0, 0.0]
-    element["temperature nonlinearity coefficient"] = [6.0, 6.0, 6.0, 6.0]
+    element["temperature nonlinearity coefficient"] = 6.0
     equation = DC2D4NL(element)
     # create model -- end
 
-    la = initialize_local_assembly() # create workspace for local matrices
+    ass = Assembly()
+    info("unknown field name: $(get_unknown_field_name(equation))")
+
     T = zeros(4) # create workspace for solution vector
     dT = zeros(4) # 
     fd = [1, 2] # free dofs
-    tic()
     # start loops, in principle solve ∂r(u)/∂uΔu = -r(u) and update.
     for i=1:10
-        calculate_local_assembly!(la, equation, "temperature")  # calculate local matrices
-        dT[fd] = la.stiffness_matrix[fd,fd] \ la.force_vector[fd]
+        empty!(ass)
+        assemble!(ass, equation)  # calculate local matrices
+        dT[fd] = full(ass.stiffness_matrix)[fd,fd] \ full(ass.force_vector)[fd]
         T += dT
         push!(element["temperature"], T) # add new increment to model
-        info("T = $T")
         @printf("increment %2d, |du| = %8.5f\n", i, norm(dT))
         err = last(element["temperature"])[1] - 2/3
         isapprox(err, 0.0) && break
     end
-    toc()
     err = last(element["temperature"])[1] - 2/3
     info("error: $err")
     @test isapprox(err, 0.0)
 end
 
+
 type TestProblem <: Problem
     unknown_field_name :: ASCIIString
     unknown_field_dimension :: Int
-    equations :: Array{Equation, 1}
+    equations :: Vector{Equation}
     element_mapping :: Dict{DataType, DataType}
 end
 
@@ -143,55 +133,48 @@ function test_potential_energy_method_2()
 
     # create model -- start
     N = Vector[[0.0, 0.0], [1.0, 0.0], [1.0, 1.0], [0.0, 1.0]]
-    element = Quad4([1, 2, 3, 4])
-    element["geometry"] = Vector[N[1], N[2], N[3], N[4]]
-    element["temperature thermal conductivity"] = 6.0
-    element["temperature load"] = [0.0, 0.0, 0.0, 0.0]
-    element["temperature nonlinearity coefficient"] = [0.0, 0.0, 0.0, 0.0]
-    #equation1 = DC2D4NL(element, initial_temperature=ones(4))
-    equation1 = DC2D4NL(element)
+    element1 = Quad4([1, 2, 3, 4])
+    element1["geometry"] = Vector[N[1], N[2], N[3], N[4]]
+    element1["temperature thermal conductivity"] = 6.0
+    element1["temperature load"] = [0.0, 0.0, 0.0, 0.0]
+    element1["temperature nonlinearity coefficient"] = [0.0, 0.0, 0.0, 0.0]
+    element1["temperature"] = ones(4)
 
-    boundary_element = Seg2([1, 2])
-    boundary_element["geometry"] = Vector[N[1], N[2]]
-    boundary_element["emissivity"] = 0.5
-    boundary_element["temperature external"] = 10.0
-    #equation2 = DC2D2NL(boundary_element, initial_temperature=ones(4))
-    equation2 = DC2D2NL(boundary_element)
+    element2 = Seg2([1, 2])
+    element2["geometry"] = Vector[N[1], N[2]]
+    element2["temperature coefficient"] = 3.0e-8 # ~ 5.7e-8 * 0.5
+    element2["temperature external"] = 100.0
+    element2["temperature"] = ones(2)
     # create model -- end
+    
+    equation1 = DC2D4NL(element1)
+    equation2 = DC2D2NL(element2)
+ 
+    ass = Assembly()
+    info("unknown field name: $(get_unknown_field_name(equation1))")
 
-    element["temperature"] = ones(4)
-    boundary_element["temperature"] = ones(2)
-
-    equations = [equation1, equation2]
-    la = initialize_local_assembly() # create workspace for local matrices
     T = zeros(4) # create workspace for solution vector
     dT = zeros(4) # 
     fd = [1, 2] # free dofs
-    info("equation 1")
-    calculate_local_assembly!(la, equation1, "temperature")
-    info("stiffness matrix: $(la.stiffness_matrix)")
-#   info("force vector: $(la.force_vector)")
-    info("equation 2")
-    calculate_local_assembly!(la, equation2, "temperature")
-#   info("stiffness matrix: $(la.stiffness_matrix)")
-    info("force vector: $(la.force_vector)")
+    # start loops, in principle solve ∂r(u)/∂uΔu = -r(u) and update.
+    for i=1:10
+        empty!(ass)
+        assemble!(ass, equation1)
+        assemble!(ass, equation2)
+        dT[fd] = full(ass.stiffness_matrix)[fd,fd] \ full(ass.force_vector)[fd]
+        T += dT
+        push!(element1["temperature"], T)
+        push!(element2["temperature"], T[fd])
+        @printf("increment %2d, |du| = %8.5f\n", i, norm(dT))
+        err = last(element1["temperature"])[1] - 0.5
+        isapprox(err, 0.0) && break
+    end
 
-    info("Creating problem")
-    #problem = PlaneHeatProblem("temperature", 1, equations, Dict())
-    problem = TestProblem(equations)
-
-    free_dofs = [1, 2]
-    tic()
-    solve!(problem, free_dofs; max_iterations=10)
-    toc()
-    temp = get_basis(boundary_element)("temperature", [0.0])[1]
-    info("temperature = $temp")
-    #err = last(element["temperature"])[1] - 2/3
-    #info("error: $err")
-    # 0.3888756709834147 tulee jostakin syysta...
-    # tai -0.39411350336960116
-    info(boundary_element["temperature"])
-    @test isapprox(temp, 2.93509690572300E+00)  # tested using Code Aster
+    err = last(element1["temperature"])[1] - 0.5
+    info("error: $err")
+    @test isapprox(err, 0.0)
+    
+    # @test isapprox(temp, 2.93509690572300E+00)  # tested using Code Aster
 end
 
 end
diff --git a/test/test_virtual_work.jl b/test/test_virtual_work.jl
index 5d1acef..3957687 100644
--- a/test/test_virtual_work.jl
+++ b/test/test_virtual_work.jl
@@ -5,30 +5,32 @@ module TestAutoDiffWeakForm
 
 using JuliaFEM.Test
 using JuliaFEM
-using JuliaFEM: Quad4, Equation, IntegrationPoint,
+using JuliaFEM: Quad4, Equation, IntegrationPoint, assemble!,
+                Assembly,
                 solve!, get_field, get_element, get_basis,
-                grad
+                grad, get_default_integration_points
 
 """ Plane stress formulation for 4-node bilinear element. """
 type CPS4 <: Equation
     element :: Quad4
-    integration_points :: Array{IntegrationPoint, 1}
+    integration_points :: Vector{IntegrationPoint}
+end
+
+function JuliaFEM.get_unknown_field_name(equation::CPS4)
+    return "displacement"
 end
 
 function CPS4(element::Quad4)
-    integration_points = [
-        IntegrationPoint(1.0/sqrt(3.0)*[-1, -1], 1.0),
-        IntegrationPoint(1.0/sqrt(3.0)*[ 1, -1], 1.0),
-        IntegrationPoint(1.0/sqrt(3.0)*[ 1,  1], 1.0),
-        IntegrationPoint(1.0/sqrt(3.0)*[-1,  1], 1.0)]
+    integration_points = get_default_integration_points(element)
     if !haskey(element, "displacement")
-        # initial field must be defined if using autodiff
         element["displacement"] = zeros(2, 4)
     end
     CPS4(element, integration_points)
 end
 
-JuliaFEM.size(eq::CPS4) = (2, 4)
+function Base.size(eq::CPS4)
+    return (2, 4)
+end
 
 function JuliaFEM.get_residual_vector(equation::CPS4, ip, time; variation=nothing)
     element = get_element(equation)
@@ -58,8 +60,6 @@ function JuliaFEM.get_residual_vector(equation::CPS4, ip, time; variation=nothin
     return vec(r)
 end
 
-JuliaFEM.has_residual_vector(equation::CPS4) = true
-
 function test_residual_form()
     # create model -- start
     element = Quad4([1, 2, 3, 4])
@@ -71,7 +71,7 @@ function test_residual_form()
     # create model -- end
 
     free_dofs = [3, 4, 5, 6]
-    solve!(equation, "displacement", free_dofs)  # launch a newton solver for single element
+    solve!(equation, free_dofs)  # launch a newton solver for single element
     disp = get_basis(element)("displacement", [1.0, 1.0])[2]
     println("displacement at tip: $disp")
     # verified using Code Aster.