generalized cumsum and some unit tests for it

Author: jdleesmiller

lib/narray_ext.rb

@@ -209,8 +209,52 @@ 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] self
+  #
+  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!
+end
 
 module NMath
   PI = Math::PI

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
+