BF+TEST - fix and extend trackvis scalar, property checks on write

Author: matthew-brett

nibabel/tests/test_trackvis.py

@@ -42,18 +42,62 @@ def test_write():
            tv.write, out_f, [],{'hdr_size': 0})
 
 
+def test_write_scalars_props():
+    # Test writing of scalar array with streamlines
+    N = 6
+    M = 2
+    P = 4
+    points = np.arange(N*3).reshape((N,3))
+    scalars = np.arange(N*M).reshape((N,M)) + 100
+    props = np.arange(P) + 1000
+    # If scalars not same size for each point, error
+    out_f = StringIO()
+    streams = [(points, None, None),
+               (points, scalars, None)]
+    assert_raises(tv.DataError, tv.write, out_f, streams)
+    out_f.seek(0)
+    streams = [(points, np.zeros((N,M+1)), None),
+               (points, scalars, None)]
+    assert_raises(tv.DataError, tv.write, out_f, streams)
+    # Or if scalars different N compared to points
+    bad_scalars = np.zeros((N+1,M))
+    out_f.seek(0)
+    streams = [(points, bad_scalars, None),
+               (points, bad_scalars, None)]
+    assert_raises(tv.DataError, tv.write, out_f, streams)
+    # Similarly properties must have the same length for each streamline
+    out_f.seek(0)
+    streams = [(points, scalars, None),
+               (points, scalars, props)]
+    assert_raises(tv.DataError, tv.write, out_f, streams)
+    out_f.seek(0)
+    streams = [(points, scalars, np.zeros((P+1,))),
+               (points, scalars, props)]
+    assert_raises(tv.DataError, tv.write, out_f, streams)
+    # If all is OK, then we get back what we put in
+    out_f.seek(0)
+    streams = [(points, scalars, props),
+               (points, scalars, props)]
+    tv.write(out_f, streams)
+    out_f.seek(0)
+    back_streams, hdr = tv.read(out_f)
+    for actual, expected in zip(streams, back_streams):
+        for a_el, e_el in zip(actual, expected):
+            assert_array_equal(a_el, e_el)
+
+
 def streams_equal(stream1, stream2):
-    if not np.all(stream1[0] == stream1[0]):
+    if not np.all(stream1[0] == stream2[0]):
         return False
     if stream1[1] is None:
         if not stream2[1] is None:
             return False
     if stream1[2] is None:
         if not stream2[2] is None:
             return False
-    if not np.all(stream1[1] == stream1[1]):
+    if not np.all(stream1[1] == stream2[1]):
         return False
-    if not np.all(stream1[2] == stream1[2]):
+    if not np.all(stream1[2] == stream2[2]):
         return False
     return True
 

nibabel/trackvis.py

@@ -108,6 +108,7 @@ def read(fileobj, as_generator=False):
           where M is the number of scalars per point
        #. properties : None or ndarray shape (P,)
           where P is the number of properties
+
     hdr : structured array
        structured array with trackvis header fields
 
@@ -218,7 +219,9 @@ def write(fileobj, streamlines,  hdr_mapping=None, endianness=None):
        If `streamlines` has a ``len`` (for example, it is a list or a tuple),
        then we can write the number of streamlines into the header.  Otherwise
        we write 0 for the number of streamlines (a valid trackvis header) and
-       write streamlines into the file until the iterable is exhausted
+       write streamlines into the file until the iterable is exhausted.
+       M - the number of scalars - has to be the same for each streamline in
+       `streamlines`.  Similarly for P.
     hdr_mapping : None, ndarray or mapping, optional
        Information for filling header fields.  Can be something
        dict-like (implementing ``items``) or a structured numpy array
@@ -255,13 +258,13 @@ def write(fileobj, streamlines,  hdr_mapping=None, endianness=None):
     if not streams0 is None:
         pts, scalars, props = streams0
         # calculate number of scalars
-        if scalars:
+        if not scalars is None:
             n_s = scalars.shape[1]
         else:
             n_s = 0
         hdr['n_scalars'] = n_s
         # calculate number of properties
-        if props:
+        if not props is None:
             n_p = props.size
             hdr['n_properties'] = n_p
         else:
@@ -283,17 +286,23 @@ def write(fileobj, streamlines,  hdr_mapping=None, endianness=None):
         # the endianness is OK.
         if pts.dtype != f4dt:
             pts = pts.astype(f4dt)
-        if n_s:
+        if n_s == 0:
+            if not (scalars is None or len(scalars) == 0):
+                raise DataError('Expecting 0 scalars per point')
+        else:
             if scalars.shape != (n_pts, n_s):
-                raise ValueError('Scalars should be shape (%s, %s)'
+                raise DataError('Scalars should be shape (%s, %s)'
                                  % (n_pts, n_s))
             if scalars.dtype != f4dt:
                 scalars = scalars.astype(f4dt)
                 pts = np.c_[pts, scalars]
         fileobj.write(pts.tostring())
-        if n_p:
+        if n_p == 0:
+            if not (props is None or len(props) == 0):
+                raise DataError('Expecting 0 properties per point')
+        else:
             if props.size != n_p:
-                raise ValueError('Properties should be size %s' % n_p)
+                raise DataError('Properties should be size %s' % n_p)
             if props.dtype != f4dt:
                 props = props.astype(f4dt)
             fileobj.write(props.tostring())