From 6eab785c1ccae7bbb9e2cd98c3a55933758bdbf8 Mon Sep 17 00:00:00 2001
From: lijas <elias.isak.borjesson@gmail.com>
Date: Tue, 13 Sep 2022 13:39:24 +0200
Subject: [PATCH] Add assemble function for non-square matrices (#471)

---
 src/assembler.jl      | 24 +++++++++++++++++++-----
 test/test_assemble.jl | 10 ++++++++++
 2 files changed, 29 insertions(+), 5 deletions(-)

diff --git a/src/assembler.jl b/src/assembler.jl
index 45193b5e83..6958b78e3f 100644
--- a/src/assembler.jl
+++ b/src/assembler.jl
@@ -40,12 +40,26 @@ end
 Assembles the element matrix `Ke` into `a`.
 """
 function assemble!(a::Assembler{T}, dofs::AbstractVector{Int}, Ke::AbstractMatrix{T}) where {T}
-    n_dofs = length(dofs)
+    assemble!(a, dofs, dofs, Ke)
+end
+
+"""
+    assemble!(a::Assembler, rowdofs, coldofs, Ke)
+
+Assembles the matrix `Ke` into `a` according to the dofs specified by `rowdofs` and `coldofs`.
+"""
+function assemble!(a::Ferrite.Assembler{T}, rowdofs::AbstractVector{Int}, coldofs::AbstractVector{Int}, Ke::AbstractMatrix{T}) where {T}
+    nrows = length(rowdofs)
+    ncols = length(coldofs)
+
+    @assert(size(Ke,1) == nrows)
+    @assert(size(Ke,2) == ncols)
+
     append!(a.V, Ke)
-    @inbounds for j in 1:n_dofs
-        append!(a.I, dofs)
-        for i in 1:n_dofs
-            push!(a.J, dofs[j])
+    @inbounds for i in 1:ncols
+        append!(a.I, rowdofs)
+        for _ in 1:nrows
+            push!(a.J, coldofs[i])
         end
     end
 end
diff --git a/test/test_assemble.jl b/test/test_assemble.jl
index bba16961b9..a789c47df0 100644
--- a/test/test_assemble.jl
+++ b/test/test_assemble.jl
@@ -18,4 +18,14 @@
     @test K[1,1] == Ke[1,1]
     @test K[1,5] == Ke[1,3]
     @test K[5,1] == Ke[3,1]
+
+    # assemble with different row and col dofs
+    rdofs = [1,4,6]
+    cdofs = [1,7]
+    a = start_assemble()
+    Ke = rand(length(rdofs), length(cdofs))
+    assemble!(a, rdofs, cdofs, Ke)
+    K = end_assemble(a)
+    @test (K[rdofs,cdofs] .== Ke) |> all
+
 end