Image wrapper indexing tests

tests/test_imagewrapper.py

@@ -1034,3 +1034,199 @@ def test_collapseExpansions(niters):
 
         for exp, col in zip(expected, collapsed):
             assert expEq(exp, col)
+
+
+def test_3D_indexing(shape=None, img=None):
+
+    # Test that a 3D image looks like a 3D image
+
+    if   shape is None:   shape = (21, 22, 23)
+    elif len(shape) == 2: shape = tuple(list(shape) + [1])
+        
+    if img is None:
+        data   = np.random.random(shape)
+        nibImg = nib.Nifti1Image(data, np.eye(4))
+        img    = imagewrap.ImageWrapper(nibImg, loadData=True)
+
+    assert tuple(img[:]      .shape) == tuple(shape)
+    assert tuple(img[:, :]   .shape) == tuple(shape)
+    assert tuple(img[:, :, :].shape) == tuple(shape)
+    assert tuple(img[:, 0, 0].shape) == (shape[0], )
+    assert tuple(img[0, :, 0].shape) == (shape[1], )
+    assert tuple(img[0, 0, :].shape) == (shape[2], )
+    assert tuple(img[0, :, :].shape) == (shape[1], shape[2])
+    assert tuple(img[:, 0, :].shape) == (shape[0], shape[2])
+    assert tuple(img[:, :, 0].shape) == (shape[0], shape[1])
+
+    assert type(img[0, 0, 0]) == np.float64
+
+    mask1 = np.zeros(shape, dtype=np.bool)
+
+    mask1[0, 0, 0] = True
+    mask1[1, 0, 0] = True
+
+    assert tuple(img[mask1].shape) == (2, )
+
+    img[0, 0, 0] =  999
+    img[:, 0, 0] = [999] * shape[0]
+    img[0, :, 0] = [999] * shape[1]
+    img[0, 0, :] = [999] * shape[2]
+    img[:, 0, 0] = np.array([999] * shape[0])
+    img[0, :, 0] = np.array([999] * shape[1])
+    img[0, 0, :] = np.array([999] * shape[2])
+        
+    img[0, :, :] = np.ones((shape[1], shape[2]))
+    img[:, 0, :] = np.ones((shape[0], shape[2]))
+    img[:, :, 0] = np.ones((shape[0], shape[1]))
+
+    img[0, :, :] = [[999] * shape[1]] * shape[2]
+    img[:, 0, :] = [[999] * shape[0]] * shape[2]
+    img[:, :, 0] = [[999] * shape[0]] * shape[1]
+
+def test_3D_4D_indexing():
+
+
+    # Testing ImageWrapper for an image with a
+    # trailing fourth dimension of length 1 -
+    # it should look like a 3D image, but
+    # should still accept (valid) 4D slicing.
+    
+    # __getitem__ and __setitem__ on
+    #   - 3D index
+    #   - 4D index
+    #   - 3D boolean array
+    #   - 4D boolean array
+    #
+    padShape = (21, 22, 23, 1)
+    shape    = padShape[:3]
+
+    data   = np.random.random(shape)
+    nibImg = nib.Nifti1Image(data, np.eye(4))
+    img    = imagewrap.ImageWrapper(nibImg, loadData=True)
+
+    test_3D_indexing(shape, img)
+
+    assert tuple(img[:, :, :, :].shape) == tuple(shape)
+    
+    assert tuple(img[:, 0, 0, 0].shape) == (shape[0], )
+    assert tuple(img[:, 0, 0, :].shape) == (shape[0], )
+    
+    assert tuple(img[:, :, 0, 0].shape) == (shape[0], shape[1])
+    assert tuple(img[:, :, 0, :].shape) == (shape[0], shape[1])
+
+    assert type(img[0, 0, 0, 0]) == np.float64
+    assert type(img[0, 0, 0, :]) == np.float64
+
+    mask = np.zeros(padShape, dtype=np.bool)
+    mask[0, 0, 0, 0] = True
+
+    assert type(img[mask])      == np.ndarray
+    assert      img[mask].shape == (1, ) 
+
+
+def test_3D_len_one_indexing(shape=None, img=None):
+
+    # Testing ImageWrapper for a 3D image with 
+    # a third dimension of length 1 - it should
+    # look like a 3D image, but should still 
+    # accept (valid) 2D slicing.
+
+    if   shape is None:  shape = (20, 20, 1)
+    elif len(shape) < 3: shape = tuple(list(shape) + [1])
+
+    if img is None:
+        data   = np.random.random(shape)
+        nibImg = nib.Nifti1Image(data, np.eye(4))
+        img    = imagewrap.ImageWrapper(nibImg, loadData=True)
+
+    test_3D_indexing(shape, img)
+
+    assert type(img[0, 0, :]) == np.ndarray
+    assert type(img[0, 0])    == np.ndarray
+    assert type(img[0, 0, 0]) == np.float64
+
+    mask = np.zeros(shape[:2], dtype=np.bool)
+    mask[0, 0] = True
+
+    assert type(img[mask])      == np.ndarray
+    assert      img[mask].shape == (1, )
+
+    mask = np.zeros(shape, dtype=np.bool)
+    mask[0, 0, 0] = True
+    
+    assert type(img[mask])      == np.ndarray
+    assert      img[mask].shape == (1, )
+ 
+
+def test_2D_indexing():
+
+    # Testing ImageWrapper for a 2D image - 
+    # it should look just like a 3D image 
+    # (the same as is tested above).
+
+    shape  = (20, 20)
+    data   = np.random.random(shape[:2])
+    nibImg = nib.Nifti1Image(data, np.eye(4))
+    img    = imagewrap.ImageWrapper(nibImg, loadData=True)
+
+    test_3D_len_one_indexing(shape, img)
+
+    
+def test_4D_indexing(shape=None, img=None):
+
+    if shape is None:
+        shape = (20, 21, 22, 23)
+        
+    if img is None:
+
+        data   = np.random.random(shape)
+        nibImg = nib.Nifti1Image(data, affine=np.eye(4))
+        img    = imagewrap.ImageWrapper(nibImg, loadData=True)
+
+    assert tuple(img[:]         .shape) == tuple(shape)
+    assert tuple(img[:, :]      .shape) == tuple(shape)
+    assert tuple(img[:, :, :]   .shape) == tuple(shape)
+    assert tuple(img[:, :, :, :].shape) == tuple(shape)
+    
+    assert tuple(img[:, 0, 0, 0].shape) == (shape[0], )
+    assert tuple(img[0, :, 0, 0].shape) == (shape[1], )
+    assert tuple(img[0, 0, :, 0].shape) == (shape[2], )
+    assert tuple(img[0, 0, 0, :].shape) == (shape[3], )
+
+    
+    assert tuple(img[0, :, :, :].shape) == (shape[1], shape[2], shape[3])
+    assert tuple(img[:, 0, :, :].shape) == (shape[0], shape[2], shape[3])
+    assert tuple(img[:, :, 0, :].shape) == (shape[0], shape[1], shape[3])
+    assert tuple(img[:, :, :, 0].shape) == (shape[0], shape[1], shape[2])
+
+    assert type(img[0, 0, 0, 0]) == np.float64 
+
+    mask1 = np.zeros(shape, dtype=np.bool)
+
+    mask1[0, 0, 0, 0] = True
+    mask1[1, 0, 0, 0] = True
+
+    assert tuple(img[mask1].shape) == (2, )
+
+    img[0, 0, 0, 0] =  999
+    
+    img[:, 0, 0, 0] = [999] * shape[0]
+    img[0, :, 0, 0] = [999] * shape[1]
+    img[0, 0, :, 0] = [999] * shape[2]
+    img[0, 0, 0, :] = [999] * shape[3]
+    
+    img[:, 0, 0, 0] = np.array([999] * shape[0])
+    img[0, :, 0, 0] = np.array([999] * shape[1])
+    img[0, 0, :, 0] = np.array([999] * shape[2])
+    img[0, 0, 0, :] = np.array([999] * shape[3])
+
+        
+    img[0, :, :, :] = np.ones((shape[1], shape[2], shape[3]))
+    img[:, 0, :, :] = np.ones((shape[0], shape[2], shape[3]))
+    img[:, :, 0, :] = np.ones((shape[0], shape[1], shape[3]))
+    img[:, :, :, 0] = np.ones((shape[0], shape[1], shape[2]))
+
+    img[0, :, :, :] = [[[999] * shape[1]] * shape[2]] * shape[3]
+    img[:, 0, :, :] = [[[999] * shape[0]] * shape[2]] * shape[3]
+    img[:, :, 0, :] = [[[999] * shape[0]] * shape[1]] * shape[3]
+    img[:, :, :, 0] = [[[999] * shape[0]] * shape[1]] * shape[2]