RF+TST: refactor design stacking for generality

Allow stacking empty designs, designs that are just arrays.  This makes
fits what the docstring describes, and is a little easier to use.

nipy/modalities/fmri/design.py

@@ -330,11 +330,17 @@ def stack2designs(old_X, new_X, old_contrasts={}, new_contrasts={}):
        The new contrast matrices reflecting changes to the columns.
     """
     contrasts = {}
+    old_X = np.asarray(old_X)
+    new_X = np.asarray(new_X)
+    if old_X.size == 0:
+        return new_X, new_contrasts
+    if new_X.size == 0:
+        return old_X, old_contrasts
 
     if old_X.ndim == 1:
-        old_X = old_X.reshape((old_X.shape[0], 1))
+        old_X = old_X[:, None]
     if new_X.ndim == 1:
-        new_X = new_X.reshape((new_X.shape[0], 1))
+        new_X = new_X[:, None]
 
     X = np.hstack([old_X, new_X])
 
@@ -388,13 +394,14 @@ def stack_contrasts(contrasts, name, keys):
 
 
 def stack_designs(*pairs):
-    """
-    Stack a sequence of design / contrast dictionary
-    pairs. Uses multiple calls to stack2designs
+    """ Stack a sequence of design / contrast dictionary pairs
+
+    Uses multiple calls to :func:`stack2designs`
 
     Parameters
     ----------
-    pairs : sequence filled with (np.ndarray, dict) or np.ndarray
+    \*pairs : sequence
+        Elements of either (np.ndarray, dict) or (np.ndarray,) or np.ndarray
 
     Returns
     -------
@@ -403,16 +410,14 @@ def stack_designs(*pairs):
     contrasts : dict
        The new contrast matrices reflecting changes to the columns.
     """
-    p = pairs[0]
-    if len(p) == 1:
-        X = p[0]; contrasts={}
-    else:
-        X, contrasts = p
-
-    for q in pairs[1:]:
-        if len(q) == 1:
-            new_X = q[0]; new_con = {}
-        else:
-            new_X, new_con = q
+    X = []
+    contrasts = {}
+    for p in pairs:
+        if isinstance(p, np.ndarray):
+            new_X = p; new_con = {}
+        elif len(p) == 1:  # Length one sequence
+            new_X = p[0]; new_con = {}
+        else:  # Length 2 sequence
+            new_X, new_con = p
         X, contrasts = stack2designs(X, new_X, contrasts, new_con)
     return X, contrasts

nipy/modalities/fmri/tests/test_design.py

@@ -3,7 +3,7 @@
 
 import numpy as np
 
-from ..design import event_design, block_design
+from ..design import event_design, block_design, stack2designs, stack_designs
 from ..utils import (events, lambdify_t, T, convolve_functions,
                      blocks)
 from ..hrf import glover
@@ -158,3 +158,59 @@ def mk_blk_tc(ev_inds):
     spec_1d = make_recarray(zip(onsets, offsets, fac_1),
                             ('mister', 'end', 'smt'))
     assert_raises(ValueError, block_design, spec_1d, t)
+
+
+def assert_des_con_equal(one, two):
+    des1, con1 = one
+    des2, con2 = two
+    assert_array_equal(des1, des2)
+    assert_equal(set(con1), set(con2))
+    for key in con1:
+        assert_array_equal(con1[key], con2[key])
+
+
+def test_stack_designs():
+    # Test stack_designs function
+    N = 10
+    X1 = np.ones((N, 1))
+    con1 = dict(con1 = np.array([1]))
+    X2 = np.eye(N)
+    con2 = dict(con2 = np.array([1] + [0] * (N -1)))
+    sX, sc = stack_designs((X1, con1), (X2, con2))
+    X1_X2 = np.c_[X1, X2]
+    exp = (X1_X2,
+           dict(con1=[1] + [0] * N, con2=[0, 1] + [0] * (N - 1)))
+    assert_des_con_equal((sX, sc), exp)
+    # Result same when stacking just two designs
+    sX, sc = stack2designs(X1, X2, {}, con2)
+    # Stacking a design with empty design is OK
+    assert_des_con_equal(stack2designs([], X2, con1, con2),
+                         (X2, con2))
+    assert_des_con_equal(stack_designs(([], con1), (X2, con2)),
+                         (X2, con2))
+    assert_des_con_equal(stack2designs(X1, [], con1, con2),
+                         (X1, con1))
+    assert_des_con_equal(stack_designs((X1, con1), ([], con2)),
+                         (X1, con1))
+    # Stacking one design returns output unmodified
+    assert_des_con_equal(stack_designs((X1, con1)), (X1, con1))
+    # Can stack design without contrasts
+    assert_des_con_equal(stack_designs(X1, X2), (X1_X2, {}))
+    assert_des_con_equal(stack_designs(X1, (X2, con2)),
+                         (X1_X2, {'con2': [0, 1] + [0] * (N - 1)}))
+    assert_des_con_equal(stack_designs((X1, con1), X2),
+                         (X1_X2, {'con1': [1] + [0] * N}))
+    # No-contrasts can also be 1-length tuple
+    assert_des_con_equal(stack_designs((X1,), (X2, con2)),
+                         (X1_X2, {'con2': [0, 1] + [0] * (N - 1)}))
+    assert_des_con_equal(stack_designs((X1, con1), (X2,)),
+                         (X1_X2, {'con1': [1] + [0] * N}))
+    # Stack three
+    X3 = np.arange(N)[:, None]
+    con3 = dict(con3=np.array([1]))
+    assert_des_con_equal(
+        stack_designs((X1, con1), (X2, con2), (X3, con3)),
+        (np.c_[X1, X2, X3],
+         dict(con1=[1, 0] + [0] * N,
+              con2=[0, 1] + [0] * N,
+              con3=[0] * N + [0, 1])))