more general cumsum and a 'tile' (repmat) method

Rakefile

@@ -88,3 +88,7 @@ end
 file "src" do
   ln_s ".","src"
 end
+
+task :unit_tests do
+  ruby "-I.:lib test/unit_tests.rb"
+end

lib/narray_ext.rb

@@ -209,8 +209,119 @@ def randomn!
   end
 
   #SFloatOne = NArray.sfloat(1).fill!(1)
-end
 
+  #
+  # Cumulative sum along dimension +dim+; modifies this array in place.
+  #
+  # @param [Number] dim non-negative
+  #
+  # @return [NArray] self
+  #
+  def cumsum_general! dim=0
+    if self.dim > dim
+      if self.dim == 1
+        # use the built-in version for dimension 1
+        self.cumsum_1!
+      else
+        # for example, if this is a matrix and dim = 0, mask_0 selects the first
+        # column of the matrix and mask_1 selects the second column; then we
+        # just shuffle them along and accumulate.
+        mask_0 = (0...self.dim).map{|d| d == dim ? 0 : true}
+        mask_1 = (0...self.dim).map{|d| d == dim ? 1 : true}
+        while mask_1[dim] < self.shape[dim]
+          self[*mask_1] += self[*mask_0]
+          mask_0[dim] += 1
+          mask_1[dim] += 1
+        end
+      end
+    end
+    self
+  end
+
+  #
+  # Cumulative sum along dimension +dim+.
+  #
+  # @param [Number] dim non-negative
+  #
+  # @return [NArray]
+  #
+  def cumsum_general dim=0
+    self.dup.cumsum_general!(dim)
+  end
+
+  # The built-in cumsum only does vectors (dim 1).
+  alias cumsum_1 cumsum
+  alias cumsum cumsum_general
+  alias cumsum_1! cumsum!
+  alias cumsum! cumsum_general!
+
+
+  #
+  # Replicate this array to make a tiled array; this is the matlab function
+  # repmat.
+  #
+  # @param [Array<Number>] reps number of times to repeat in each dimension;
+  # note that reps.size is allowed to be different from self.dim, and dimensions
+  # of size 1 will be added to compensate
+  #
+  # @return [NArray] with same typecode as self 
+  #
+  def tile *reps
+    if self.dim == 0 || reps.member?(0)
+      # Degenerate case: 0 dimensions or dimension 0
+      res = NArray.new(self.typecode, 0)
+    else
+      if reps.size <= self.dim 
+        # Repeat any extra dims once.
+        reps = reps + [1]*(self.dim - reps.size) 
+        tile = self
+      else
+        # Have to add some more dimensions (with implicit shape[dim] = 1).
+        tile_shape = self.shape + [1]*(reps.size - self.dim) 
+        tile = self.reshape(*tile_shape)
+      end
+
+      # Allocate tiled matrix.
+      res_shape = (0...tile.dim).map{|i| tile.shape[i] * reps[i]}
+      res = NArray.new(self.typecode, *res_shape)
+
+      # Copy tiles.
+      # This probably isn't the most efficient way of doing this; just doing
+      # res[] = tile doesn't seem to work in general
+      nested_for_zero_to(reps) do |tile_pos|
+        tile_slice = (0...tile.dim).map{|i|
+          (tile.shape[i] * tile_pos[i])...(tile.shape[i] * (tile_pos[i]+1))}
+        res[*tile_slice] = tile
+      end
+    end
+    res
+  end
+
+  private
+
+  #
+  # This is effectively <tt>suprema.size</tt> nested 'for' loops, in which the
+  # outermost loop runs over <tt>0...suprema.first</tt>, and the innermost loop
+  # runs over <tt>0...suprema.last</tt>.
+  #
+  # For example, when +suprema+ is [3], it yields [0], [1] and [2], and when
+  # +suprema+ is [3,2] it yields [0,0], [0,1], [1,0], [1,1], [2,0] and [2,1].
+  #
+  # @param [Array<Integer>] suprema non-negative entries; does not yield if
+  #        empty 
+  #
+  # @return [nil]
+  # 
+  def nested_for_zero_to suprema
+    unless suprema.empty?
+      nums = suprema.map{|n| (0...n).to_a}
+      nums.first.product(*nums.drop(1)).each do |num|
+        yield num
+      end
+    end
+    nil
+  end
+end
 
 module NMath
   PI = Math::PI

narray.gemspec

@@ -0,0 +1,19 @@
+# -*- encoding: utf-8 -*-
+
+Gem::Specification.new do |s|
+  s.name              = 'narray'
+  s.version           = '0.5.9.9'
+  s.platform          = Gem::Platform::RUBY
+  s.authors           = ['Masahiro Tanaka']
+  s.email             = ['[email protected]']
+  s.homepage          = 'http://narray.rubyforge.org/'
+  s.summary           = %q{N-dimensional Numerical Array class for Ruby}
+  s.description       = %q{Numerical N-dimensional Array class}
+
+  s.rubyforge_project = 'narray'
+
+  s.files      = Dir.glob('lib/*.rb') + Dir.glob('./*{.h,.c}')
+  s.extensions = "./extconf.rb"
+  s.require_paths << '.'
+end
+

test/testcumsum.rb

@@ -0,0 +1,44 @@
+require 'test/unit_tests'
+
+class TestCumSum < Test::Unit::TestCase
+  def test_narray_cumsum
+    # Degenerate case: dimension 0.
+    assert_equal NArray.float(0), NArray.float(0).cumsum
+    assert_equal NArray.int(0),   NArray.int(0).cumsum
+
+    # Single-element vector.
+    v = NArray.float(1).fill!(42)
+    assert_equal v, v.cumsum
+    assert_equal v.typecode, v.cumsum.typecode
+    v = NArray.int(1).fill!(42)
+    assert_equal v, v.cumsum
+    assert_equal v.typecode, v.cumsum.typecode
+
+    # Vector.
+    v = NArray.float(2).indgen!
+    assert_narray_close NArray[   0.0,  1.0], v.cumsum
+    v = NArray.float(3).indgen! + 1
+    assert_narray_close NArray[   1.0,  3.0,  6.0], v.cumsum
+    assert_equal v, v.cumsum(1) # dim 1 doesn't exist; cumsum has no effect
+
+    # Matrix.
+    m = NArray.float(3,2).indgen! + 1
+    assert_narray_close NArray[[  1.0,  3.0,  6.0],
+                               [  4.0,  9.0, 15.0]], m.cumsum
+    assert_narray_close NArray[[  1.0,  2.0,  3.0],
+                               [  5.0,  7.0,  9.0]], m.cumsum(1)
+    assert_equal m, m.cumsum(2) # dim 2 doesn't exist; cumsum has no effect
+
+    # Array with dim 3 with one extra dim.
+    d3 = NArray.int(3,2,1).indgen! - 1
+    assert_equal NArray[[[  -1,  -1,   0],
+                         [   2,   5,   9]]], d3.cumsum(0)
+    assert_equal NArray[[[  -1,   0,   1],
+                         [   1,   3,   5]]], d3.cumsum(1)
+    assert_equal NArray[[[  -1,   0,   1],
+                         [   2,   3,   4]]], d3.cumsum(2)
+    assert_equal NArray[[[  -1,   0,   1],
+                         [   2,   3,   4]]], d3.cumsum(3) # dim 3 doesn't exist
+  end
+end
+

test/testtile.rb

@@ -0,0 +1,87 @@
+require 'test/unit_tests'
+
+class TestTile < Test::Unit::TestCase
+  def test_narray_tile
+    # Degenerate case: tile on a dimension 0 array.
+    assert_equal NArray.float(0), NArray.float(0).tile
+    assert_equal NArray.float(0), NArray.float(0).tile(0)
+    assert_equal NArray.float(0), NArray.float(0).tile(0,0)
+    assert_equal NArray.float(0), NArray.float(0).tile(1,1)
+
+    # Degenerate case: tile 0 times on some dimension.
+    assert_equal NArray.float(0), NArray.float(1).tile(0)
+    assert_equal NArray.float(0), NArray.float(1).tile(0,0)
+    assert_equal NArray.float(0), NArray.float(2,3).tile(0)
+    assert_equal NArray.float(0), NArray.float(3,4,2).tile(1,2,0)
+    assert_equal NArray.float(0), NArray.float(3,4,2).tile(1,0,2)
+
+    # Degenerate case: tile with no args returns copy of original.
+    assert_equal NArray.float(1).fill!(1), NArray.float(1).fill!(1).tile
+    assert_equal NArray.float(1,2).indgen!, NArray.float(1,2).indgen!.tile
+
+    # Tile a scalar.
+    assert_equal NArray[1.0, 1.0],
+                 NArray.float(1).fill!(1).tile(2) # row vector
+    assert_equal NArray[[1.0],
+                        [1.0]],
+                 NArray.float(1).fill!(1).tile(1,2) # column vector
+    assert_equal NArray[[[1.0]],
+                        [[1.0]]],
+                 NArray.float(1).fill!(1).tile(1,1,2) # add a dimension
+    assert_equal NArray[[1.0, 1.0, 1.0],
+                        [1.0, 1.0, 1.0]],
+                 NArray.float(1).fill!(1).tile(3,2) # matrix
+
+    # Tile a vector.
+    v = NArray.float(2).indgen!
+    assert_equal NArray[0.0, 1.0, 0.0, 1.0], v.tile(2)
+    assert_equal NArray[0.0, 1.0, 0.0, 1.0, 0.0, 1.0], v.tile(3)
+    assert_equal NArray[[0.0, 1.0],
+                        [0.0, 1.0]], v.tile(1,2)
+    assert_equal NArray[[0.0, 1.0],
+                        [0.0, 1.0],
+                        [0.0, 1.0]], v.tile(1,3)
+
+    # Tile a matrix.
+    m = NArray.float(2,3).indgen!
+    assert_equal NArray[[0.0, 1.0],
+                        [2.0, 3.0],
+                        [4.0, 5.0]], m.tile
+    assert_equal NArray[[0.0, 1.0, 0.0, 1.0],
+                        [2.0, 3.0, 2.0, 3.0],
+                        [4.0, 5.0, 4.0, 5.0]], m.tile(2)
+    assert_equal m.tile(2), m.tile(2,1)
+    assert_equal NArray[[0.0, 1.0],
+                        [2.0, 3.0],
+                        [4.0, 5.0],
+                        [0.0, 1.0],
+                        [2.0, 3.0],
+                        [4.0, 5.0]], m.tile(1,2)
+    assert_equal NArray[[[0.0, 1.0],
+                         [2.0, 3.0],
+                         [4.0, 5.0]],
+                        [[0.0, 1.0],
+                         [2.0, 3.0],
+                         [4.0, 5.0]]], m.tile(1,1,2)
+
+    # Tile another matrix.
+    m = NArray.float(3,2).indgen!
+    assert_equal NArray[[0.0, 1.0, 2.0],
+                        [3.0, 4.0, 5.0]], m.tile
+    assert_equal NArray[[0.0, 1.0, 2.0, 0.0, 1.0, 2.0],
+                        [3.0, 4.0, 5.0, 3.0, 4.0, 5.0]], m.tile(2)
+    assert_equal NArray[[0.0, 1.0, 2.0],
+                        [3.0, 4.0, 5.0],
+                        [0.0, 1.0, 2.0],
+                        [3.0, 4.0, 5.0],
+                        [0.0, 1.0, 2.0],
+                        [3.0, 4.0, 5.0]], m.tile(1,3)
+    assert_equal NArray[[0.0, 1.0, 2.0, 0.0, 1.0, 2.0],
+                        [3.0, 4.0, 5.0, 3.0, 4.0, 5.0],
+                        [0.0, 1.0, 2.0, 0.0, 1.0, 2.0],
+                        [3.0, 4.0, 5.0, 3.0, 4.0, 5.0],
+                        [0.0, 1.0, 2.0, 0.0, 1.0, 2.0],
+                        [3.0, 4.0, 5.0, 3.0, 4.0, 5.0]], m.tile(2,3)
+  end
+end
+

test/unit_tests.rb

@@ -0,0 +1,13 @@
+require 'test/unit'
+require 'narray'
+
+$delta = 1e-6 # numerical tolerance
+
+def assert_narray_close exp, obs
+  assert exp.shape == obs.shape && ((exp - obs).abs < $delta).all?,
+    "#{exp.inspect} expected; got\n#{obs.inspect}"
+end
+
+load 'test/testcumsum.rb'
+load 'test/testtile.rb'
+