RF+DOC+TEST: tck format.

Author: MarcCote

nibabel/streamlines/tck.py

@@ -16,6 +16,7 @@
 from .test_tractogram import assert_tractogram_equal
 from ..tractogram import Tractogram, LazyTractogram
 from ..tractogram_file import TractogramFile, ExtensionWarning
+from .. import FORMATS
 from .. import trk
 
 DATA = {}
@@ -24,11 +25,13 @@
 def setup():
     global DATA
     DATA['empty_filenames'] = [pjoin(data_path, "empty" + ext)
-                               for ext in nib.streamlines.FORMATS.keys()]
+                               for ext in FORMATS.keys()]
     DATA['simple_filenames'] = [pjoin(data_path, "simple" + ext)
-                                for ext in nib.streamlines.FORMATS.keys()]
+                                for ext in FORMATS.keys()]
     DATA['complex_filenames'] = [pjoin(data_path, "complex" + ext)
-                                 for ext in nib.streamlines.FORMATS.keys()]
+                                 for ext, cls in FORMATS.items()
+                                 if (cls.SUPPORTS_DATA_PER_POINT or
+                                     cls.SUPPORTS_DATA_PER_STREAMLINE)]
 
     DATA['streamlines'] = [np.arange(1*3, dtype="f4").reshape((1, 3)),
                            np.arange(2*3, dtype="f4").reshape((2, 3)),
@@ -84,23 +87,23 @@ def test_is_supported_detect_format():
     assert_true(nib.streamlines.detect_format("") is None)
 
     # Valid file without extension
-    for tfile_cls in nib.streamlines.FORMATS.values():
+    for tfile_cls in FORMATS.values():
         f = BytesIO()
         f.write(tfile_cls.MAGIC_NUMBER)
         f.seek(0, os.SEEK_SET)
         assert_true(nib.streamlines.is_supported(f))
         assert_true(nib.streamlines.detect_format(f) is tfile_cls)
 
     # Wrong extension but right magic number
-    for tfile_cls in nib.streamlines.FORMATS.values():
+    for tfile_cls in FORMATS.values():
         with tempfile.TemporaryFile(mode="w+b", suffix=".txt") as f:
             f.write(tfile_cls.MAGIC_NUMBER)
             f.seek(0, os.SEEK_SET)
             assert_true(nib.streamlines.is_supported(f))
             assert_true(nib.streamlines.detect_format(f) is tfile_cls)
 
     # Good extension but wrong magic number
-    for ext, tfile_cls in nib.streamlines.FORMATS.items():
+    for ext, tfile_cls in FORMATS.items():
         with tempfile.TemporaryFile(mode="w+b", suffix=ext) as f:
             f.write(b"pass")
             f.seek(0, os.SEEK_SET)
@@ -113,13 +116,13 @@ def test_is_supported_detect_format():
     assert_true(nib.streamlines.detect_format(f) is None)
 
     # Good extension, string only
-    for ext, tfile_cls in nib.streamlines.FORMATS.items():
+    for ext, tfile_cls in FORMATS.items():
         f = "my_tractogram" + ext
         assert_true(nib.streamlines.is_supported(f))
         assert_equal(nib.streamlines.detect_format(f), tfile_cls)
 
     # Extension should not be case-sensitive.
-    for ext, tfile_cls in nib.streamlines.FORMATS.items():
+    for ext, tfile_cls in FORMATS.items():
         f = "my_tractogram" + ext.upper()
         assert_true(nib.streamlines.detect_format(f) is tfile_cls)
 
@@ -208,7 +211,7 @@ def test_save_tractogram_file(self):
 
     def test_save_empty_file(self):
         tractogram = Tractogram(affine_to_rasmm=np.eye(4))
-        for ext, cls in nib.streamlines.FORMATS.items():
+        for ext, cls in FORMATS.items():
             with InTemporaryDirectory():
                 filename = 'streamlines' + ext
                 nib.streamlines.save(tractogram, filename)
@@ -218,7 +221,7 @@ def test_save_empty_file(self):
     def test_save_simple_file(self):
         tractogram = Tractogram(DATA['streamlines'],
                                 affine_to_rasmm=np.eye(4))
-        for ext, cls in nib.streamlines.FORMATS.items():
+        for ext, cls in FORMATS.items():
             with InTemporaryDirectory():
                 filename = 'streamlines' + ext
                 nib.streamlines.save(tractogram, filename)
@@ -231,7 +234,7 @@ def test_save_complex_file(self):
                                         DATA['data_per_point'],
                                         affine_to_rasmm=np.eye(4))
 
-        for ext, cls in nib.streamlines.FORMATS.items():
+        for ext, cls in FORMATS.items():
             with InTemporaryDirectory():
                 filename = 'streamlines' + ext
 
@@ -240,12 +243,17 @@ def test_save_complex_file(self):
                     nib.streamlines.save(complex_tractogram, filename)
 
                     # If streamlines format does not support saving data
-                    # per point or data per streamline, a warning message
+                    # per point or data per streamline, warning messages
                     # should be issued.
-                    if not (cls.SUPPORTS_DATA_PER_POINT and
-                            cls.SUPPORTS_DATA_PER_STREAMLINE):
-                        assert_equal(len(w), 1)
-                        assert_true(issubclass(w[0].category, Warning))
+                    nb_expected_warnings = 0
+                    if not cls.SUPPORTS_DATA_PER_POINT:
+                        nb_expected_warnings += 1
+                    if not cls.SUPPORTS_DATA_PER_STREAMLINE:
+                        nb_expected_warnings += 1
+
+                    assert_equal(len(w), nb_expected_warnings)
+                    for i in range(nb_expected_warnings):
+                        assert_true(issubclass(w[i].category, Warning))
 
                     tractogram = Tractogram(DATA['streamlines'],
                                             affine_to_rasmm=np.eye(4))

nibabel/streamlines/tests/test_streamlines.py

@@ -1,96 +1,142 @@
 import os
 import unittest
 import numpy as np
+from os.path import join as pjoin
 
 from nibabel.externals.six import BytesIO
 
-from nibabel.testing import suppress_warnings, clear_and_catch_warnings
-from nibabel.testing import assert_arrays_equal, check_iteration
-from nose.tools import assert_equal, assert_raises, assert_true
+from nose.tools import assert_equal
 
+from nibabel.testing import data_path
 from .test_tractogram import assert_tractogram_equal
-from ..tractogram import Tractogram, LazyTractogram
-from ..tractogram_file import DataError, HeaderError, HeaderWarning
+from ..tractogram import Tractogram
 
 from ..tck import TckFile
 
-DATA_PATH = os.path.join(os.path.dirname(__file__), 'data')
 
+DATA = {}
 
-class TestTCK(unittest.TestCase):
 
-    def setUp(self):
-        self.empty_tck_filename = os.path.join(DATA_PATH, "empty.tck")
-        # simple.tck contains only streamlines
-        self.simple_tck_filename = os.path.join(DATA_PATH, "simple.tck")
+def setup():
+    global DATA
+
+    DATA['empty_tck_fname'] = pjoin(data_path, "empty.tck")
+    # simple.trk contains only streamlines
+    DATA['simple_tck_fname'] = pjoin(data_path, "simple.tck")
+    DATA['simple_tck_big_endian_fname'] = pjoin(data_path,
+                                                "simple_big_endian.tck")
+    # standard.trk contains only streamlines
+    DATA['standard_tck_fname'] = pjoin(data_path, "standard.tck")
+    # complex.trk contains streamlines, scalars and properties
+    # DATA['complex_tck_fname'] = pjoin(data_path, "complex.tck")
+
+    DATA['streamlines'] = [np.arange(1*3, dtype="f4").reshape((1, 3)),
+                           np.arange(2*3, dtype="f4").reshape((2, 3)),
+                           np.arange(5*3, dtype="f4").reshape((5, 3))]
+
+    # DATA['fa'] = [np.array([[0.2]], dtype="f4"),
+    #               np.array([[0.3],
+    #                         [0.4]], dtype="f4"),
+    #               np.array([[0.5],
+    #                         [0.6],
+    #                         [0.6],
+    #                         [0.7],
+    #                         [0.8]], dtype="f4")]
+
+    # DATA['colors'] = [np.array([(1, 0, 0)]*1, dtype="f4"),
+    #                   np.array([(0, 1, 0)]*2, dtype="f4"),
+    #                   np.array([(0, 0, 1)]*5, dtype="f4")]
+
+    # DATA['mean_curvature'] = [np.array([1.11], dtype="f4"),
+    #                           np.array([2.11], dtype="f4"),
+    #                           np.array([3.11], dtype="f4")]
 
-        self.streamlines = [np.arange(1*3, dtype="f4").reshape((1, 3)),
-                            np.arange(2*3, dtype="f4").reshape((2, 3)),
-                            np.arange(5*3, dtype="f4").reshape((5, 3))]
+    # DATA['mean_torsion'] = [np.array([1.22], dtype="f4"),
+    #                         np.array([2.22], dtype="f4"),
+    #                         np.array([3.22], dtype="f4")]
 
-        self.empty_tractogram = Tractogram()
-        self.simple_tractogram = Tractogram(self.streamlines)
+    # DATA['mean_colors'] = [np.array([1, 0, 0], dtype="f4"),
+    #                        np.array([0, 1, 0], dtype="f4"),
+    #                        np.array([0, 0, 1], dtype="f4")]
+
+    # DATA['data_per_point'] = {'colors': DATA['colors'],
+    #                           'fa': DATA['fa']}
+    # DATA['data_per_streamline'] = {'mean_curvature': DATA['mean_curvature'],
+    #                                'mean_torsion': DATA['mean_torsion'],
+    #                                'mean_colors': DATA['mean_colors']}
+
+    DATA['empty_tractogram'] = Tractogram(affine_to_rasmm=np.eye(4))
+    DATA['simple_tractogram'] = Tractogram(DATA['streamlines'],
+                                           affine_to_rasmm=np.eye(4))
+    # DATA['complex_tractogram'] = Tractogram(DATA['streamlines'],
+    #                                         DATA['data_per_streamline'],
+    #                                         DATA['data_per_point'],
+    #                                         affine_to_rasmm=np.eye(4))
+
+
+class TestTCK(unittest.TestCase):
 
     def test_load_empty_file(self):
         for lazy_load in [False, True]:
-            tck = TckFile.load(self.empty_tck_filename, lazy_load=lazy_load)
-            assert_tractogram_equal(tck.tractogram, self.empty_tractogram)
+            trk = TckFile.load(DATA['empty_tck_fname'], lazy_load=lazy_load)
+            assert_tractogram_equal(trk.tractogram, DATA['empty_tractogram'])
 
     def test_load_simple_file(self):
         for lazy_load in [False, True]:
-            tck = TckFile.load(self.simple_tck_filename, lazy_load=lazy_load)
-            assert_tractogram_equal(tck.tractogram, self.simple_tractogram)
+            trk = TckFile.load(DATA['simple_tck_fname'], lazy_load=lazy_load)
+            assert_tractogram_equal(trk.tractogram, DATA['simple_tractogram'])
 
     def test_write_empty_file(self):
-        tractogram = Tractogram()
+        tractogram = Tractogram(affine_to_rasmm=np.eye(4))
 
         tck_file = BytesIO()
         tck = TckFile(tractogram)
         tck.save(tck_file)
         tck_file.seek(0, os.SEEK_SET)
 
-        loaded_tck = TckFile.load(tck_file)
-        assert_tractogram_equal(loaded_tck.tractogram, tractogram)
+        new_tck = TckFile.load(tck_file)
+        assert_tractogram_equal(new_tck.tractogram, tractogram)
 
-        loaded_tck_orig = TckFile.load(self.empty_tck_filename)
-        assert_tractogram_equal(loaded_tck.tractogram, loaded_tck_orig.tractogram)
+        new_tck_orig = TckFile.load(DATA['empty_tck_fname'])
+        assert_tractogram_equal(new_tck.tractogram, new_tck_orig.tractogram)
 
         tck_file.seek(0, os.SEEK_SET)
-        assert_equal(tck_file.read(), open(self.empty_tck_filename, 'rb').read())
+        assert_equal(tck_file.read(),
+                     open(DATA['empty_tck_fname'], 'rb').read())
 
     def test_write_simple_file(self):
-        tractogram = Tractogram(self.streamlines)
+        tractogram = Tractogram(DATA['streamlines'],
+                                affine_to_rasmm=np.eye(4))
 
         tck_file = BytesIO()
         tck = TckFile(tractogram)
         tck.save(tck_file)
         tck_file.seek(0, os.SEEK_SET)
 
-        loaded_tck = TckFile.load(tck_file)
-        assert_tractogram_equal(loaded_tck.tractogram, tractogram)
+        new_tck = TckFile.load(tck_file)
+        assert_tractogram_equal(new_tck.tractogram, tractogram)
 
-        loaded_tck_orig = TckFile.load(self.simple_tck_filename)
-        assert_tractogram_equal(loaded_tck.tractogram, loaded_tck_orig.tractogram)
+        new_tck_orig = TckFile.load(DATA['simple_tck_fname'])
+        assert_tractogram_equal(new_tck.tractogram, new_tck_orig.tractogram)
 
         tck_file.seek(0, os.SEEK_SET)
-        assert_equal(tck_file.read(), open(self.simple_tck_filename, 'rb').read())
+        assert_equal(tck_file.read(),
+                     open(DATA['simple_tck_fname'], 'rb').read())
 
     def test_load_write_file(self):
-        for filename in [self.empty_tck_filename, self.simple_tck_filename]:
+        for fname in [DATA['empty_tck_fname'],
+                      DATA['simple_tck_fname']]:
             for lazy_load in [False, True]:
-                tck = TckFile.load(filename, lazy_load=lazy_load)
+                tck = TckFile.load(fname, lazy_load=lazy_load)
                 tck_file = BytesIO()
                 tck.save(tck_file)
 
-                loaded_tck = TckFile.load(filename, lazy_load=False)
+                loaded_tck = TckFile.load(fname, lazy_load=False)
                 assert_tractogram_equal(loaded_tck.tractogram, tck.tractogram)
 
-                # tck_file.seek(0, os.SEEK_SET)
-                # assert_equal(tck_file.read(), open(filename, 'rb').read())
-
         # Save tractogram that has an affine_to_rasmm.
         for lazy_load in [False, True]:
-            tck = TckFile.load(self.simple_tck_filename, lazy_load=lazy_load)
+            tck = TckFile.load(DATA['simple_tck_fname'], lazy_load=lazy_load)
             affine = np.eye(4)
             affine[0, 0] *= -1  # Flip in X
             tractogram = Tractogram(tck.streamlines, affine_to_rasmm=affine)

nibabel/streamlines/tests/test_tck.py

@@ -296,8 +296,6 @@ def test_load_write_file(self):
                 new_trk = TrkFile.load(fname, lazy_load=False)
                 assert_tractogram_equal(new_trk.tractogram, trk.tractogram)
 
-                trk_file.seek(0, os.SEEK_SET)
-
     def test_load_write_LPS_file(self):
         # Load the RAS and LPS version of the standard.
         trk_RAS = TrkFile.load(DATA['standard_trk_fname'], lazy_load=False)

nibabel/streamlines/tests/test_trk.py

@@ -14,6 +14,10 @@ class HeaderWarning(Warning):
     """ Base class for warnings about tractogram file header. """
 
 
+class DataWarning(Warning):
+    """ Base class for warnings about tractogram file data. """
+
+
 class HeaderError(Exception):
     """ Raised when a tractogram file header contains invalid information. """
 

nibabel/streamlines/tractogram_file.py

@@ -220,7 +220,6 @@ class TrkFile(TractogramFile):
 
     Notes
     -----
-
     TrackVis (so its file format: TRK) considers the streamline coordinate
     (0,0,0) to be in the corner of the voxel whereas NiBabel's streamlines
     internal representation (Voxel space) assumes (0,0,0) to be in the
@@ -597,7 +596,7 @@ def _read_header(fileobj):
             # Keep the file position where the data begin.
             header['_offset_data'] = f.tell()
 
-        # Set the file position where it was, if it was previously open
+        # Set the file position where it was, if it was previously open.
         if start_position is not None:
             fileobj.seek(start_position, os.SEEK_SET)