From fa18e0f123945ae8d35d3cf41f67c55d499d1f13 Mon Sep 17 00:00:00 2001
From: agisga <alexej.go@googlemail.com>
Date: Tue, 26 May 2015 16:00:09 -0500
Subject: [PATCH] NMatrix#block_diagonal added

---
 lib/nmatrix/shortcuts.rb | 55 ++++++++++++++++++++++++++++++++++++++++
 spec/shortcuts_spec.rb   | 22 ++++++++++++++++
 2 files changed, 77 insertions(+)

diff --git a/lib/nmatrix/shortcuts.rb b/lib/nmatrix/shortcuts.rb
index 6ce64cb6..9c384df7 100644
--- a/lib/nmatrix/shortcuts.rb
+++ b/lib/nmatrix/shortcuts.rb
@@ -275,6 +275,61 @@ def diagonal(entries, opts={})
     alias :diag :diagonal
     alias :diagonals :diagonal
 
+    # Generate a block-diagonal NMatrix from the supplied 2D square matrices.
+    #
+    # * *Arguments*
+    #   - +*params+ -> An array that collects all arguments passed to the method. The method
+    #                  can receive any number of arguments. Optionally, the last entry of +params+ is 
+    #                  a hash of options from NMatrix#initialize. All other entries of +params+ are 
+    #                  the blocks of the desired block-diagonal matrix, which are supplied
+    #                  as square 2D NMatrix objects.
+    # * *Returns*
+    #   - NMatrix of block-diagonal form filled with specified matrices
+    #     as the blocks along the diagonal.
+    #
+    # * *Example* 
+    #
+    #  a = NMatrix.new([2,2],[1,2,3,4])
+    #  b = NMatrix.new([1,1],[123],dtype: :int32)
+    #  c = NMatrix.new([3,3],[1,2,3,1,2,3,1,2,3], dtype: :float64)
+    #  m = NMatrix.block_diagonal(a,b,c,dtype: :int64, stype: :yale)
+    #      => 
+    #      [
+    #        [1, 2,   0, 0, 0, 0]
+    #        [3, 4,   0, 0, 0, 0]
+    #        [0, 0, 123, 0, 0, 0]
+    #        [0, 0,   0, 1, 2, 3]
+    #        [0, 0,   0, 1, 2, 3]
+    #        [0, 0,   0, 1, 2, 3]
+    #      ]
+    #
+    def block_diagonal(*params)
+      options = params.last.is_a?(Hash) ? params.pop : {}
+
+      block_sizes = [] #holds the size of each matrix block
+      params.each do |b|
+        raise ArgumentError, "Only NMatrix objects allowed" unless b.is_a?(NMatrix)
+        raise ArgumentError, "Only 2D matrices allowed" unless b.shape.size == 2
+        raise ArgumentError, "Only square matrices allowed" unless b.shape[0] == b.shape[1]
+        block_sizes << b.shape[0]
+      end
+
+      bdiag = NMatrix.zeros(block_sizes.sum, options)
+      (0...params.length).each do |n|
+        # First determine the size and position of the n'th block in the block-diagonal matrix
+        k = block_sizes[n]
+        block_pos = block_sizes[0...n].sum
+        # populate the n'th block in the block-diagonal matrix
+        (0...k).each do |i|
+          (0...k).each do |j|
+            bdiag[block_pos+i,block_pos+j] = params[n][i,j]
+          end
+        end
+      end
+
+      return bdiag
+    end
+    alias :block_diag :block_diagonal
 
     #
     # call-seq:
diff --git a/spec/shortcuts_spec.rb b/spec/shortcuts_spec.rb
index bb397174..677864fc 100644
--- a/spec/shortcuts_spec.rb
+++ b/spec/shortcuts_spec.rb
@@ -65,6 +65,28 @@
     expect(m[3,3]).to eq(arr[3])
   end
 
+  ALL_DTYPES.each do |dtype|
+    [:dense, :yale, :list].each do |stype|
+      context "#block_diagonal #{dtype} #{stype}" do
+        it "block_diagonal() creates a block-diagonal NMatrix" do
+          a = NMatrix.new([2,2],[1,2,
+                                 3,4])
+          b = NMatrix.new([1,1],[123.0])
+          c = NMatrix.new([3,3],[1,2,3,
+                                 1,2,3,
+                                 1,2,3])
+          m = NMatrix.block_diagonal(a,b,c, dtype: dtype, stype: stype)
+          expect(m).to eq(NMatrix.new([6,6], [1, 2,   0, 0, 0, 0,
+                                              3, 4,   0, 0, 0, 0,
+                                              0, 0, 123, 0, 0, 0,
+                                              0, 0,   0, 1, 2, 3,
+                                              0, 0,   0, 1, 2, 3,
+                                              0, 0,   0, 1, 2, 3], dtype: dtype, stype: stype))
+        end
+      end
+    end
+  end
+
   context "::random" do
     it "creates a matrix of random numbers" do
       m = NMatrix.random(2)