From 02637afc16999de070c0879aa96e486d0b685060 Mon Sep 17 00:00:00 2001
From: Jukka Aho <ahojukka5@gmail.com>
Date: Mon, 8 Apr 2019 18:26:34 +0300
Subject: [PATCH] Modifications to `add!` (#41)

With these changes it's possible to give `dofs1` and `dofs2` eg. in
sub-arrays or tuples, not only vectors. Also, `data` can be a long
vector in row-major order (`data_vec = vec(data)`).
---
 src/sparse.jl         | 9 ++++++---
 test/test_assembly.jl | 4 ++--
 2 files changed, 8 insertions(+), 5 deletions(-)

diff --git a/src/sparse.jl b/src/sparse.jl
index d473b6d..148fc8c 100644
--- a/src/sparse.jl
+++ b/src/sparse.jl
@@ -111,11 +111,14 @@ Matrix(A)
  0.0  0.0  0.0  0.0  0.0  8.0  9.0  10.0
 ```
 """
-function add!(A::SparseMatrixCOO, dofs1::Vector{Int}, dofs2::Vector{Int}, data)
-    n, m = size(data)
+function add!(A::SparseMatrixCOO, dofs1::AbstractVector{Int}, dofs2::AbstractVector{Int}, data)
+    n, m = length(dofs1), length(dofs2)
+    @assert length(data) == n*m
+    k = 1
     for j=1:m
         for i=1:n
-            add!(A, dofs1[i], dofs2[j], data[i,j])
+            add!(A, dofs1[i], dofs2[j], data[k])
+            k += 1
         end
     end
     return nothing
diff --git a/test/test_assembly.jl b/test/test_assembly.jl
index a6e8299..cd93634 100644
--- a/test/test_assembly.jl
+++ b/test/test_assembly.jl
@@ -78,8 +78,8 @@ end
     A = Assembly()
     B = Assembly()
     @test isempty(A)
-    add!(A.K, [1,2,3,4], [1,2,3,4], [1.0 2.0; 3.0 4.0])
-    add!(B.K, [1,2,3,4], [1,2,3,4], [1.0 2.0; 3.0 4.0])
+    add!(A.K, [1,2], [1,2], [1.0 2.0; 3.0 4.0])
+    add!(B.K, [1,2], [1,2], [1.0 2.0; 3.0 4.0])
     @test isapprox(A, B)
     @test !isempty(A)
     empty!(A)