ENH: speed up TRK loading by setting a suitable buffer size.

Author: MarcCote

nibabel/streamlines/array_sequence.py

@@ -547,7 +547,7 @@ def load(cls, filename):
         return seq
 
 
-def create_arraysequences_from_generator(gen, n):
+def create_arraysequences_from_generator(gen, n, buffer_sizes=None):
     """ Creates :class:`ArraySequence` objects from a generator yielding tuples
 
     Parameters
@@ -557,8 +557,13 @@ def create_arraysequences_from_generator(gen, n):
         array sequences.
     n : int
         Number of :class:`ArraySequences` object to create.
+    buffer_sizes : list of float, optional
+        Sizes (in Mb) for each ArraySequence's buffer.
     """
-    seqs = [ArraySequence() for _ in range(n)]
+    if buffer_sizes is None:
+        buffer_sizes = [4] * n
+
+    seqs = [ArraySequence(buffer_size=size) for size in buffer_sizes]
     for data in gen:
         for i, seq in enumerate(seqs):
             if data[i].nbytes > 0:

nibabel/streamlines/trk.py

@@ -372,8 +372,23 @@ def _read():
             tractogram = LazyTractogram.from_data_func(_read)
 
         else:
+
+            # Speed up loading by guessing a suitable buffer size.
+            with Opener(fileobj) as f:
+                old_file_position = f.tell()
+                f.seek(0, os.SEEK_END)
+                size = f.tell()
+                f.seek(old_file_position, os.SEEK_SET)
+
+            # Buffer size is in mega bytes.
+            mbytes = size // (1024 * 1024)
+            sizes = [mbytes, 4, 4]
+            if hdr["nb_scalars_per_point"] > 0:
+                sizes = [mbytes // 2, mbytes // 2, 4]
+
             trk_reader = cls._read(fileobj, hdr)
-            arr_seqs = create_arraysequences_from_generator(trk_reader, n=3)
+            arr_seqs = create_arraysequences_from_generator(trk_reader, n=3,
+                                                            buffer_sizes=sizes)
             streamlines, scalars, properties = arr_seqs
             properties = np.asarray(properties)  # Actually a 2d array.
             tractogram = Tractogram(streamlines)