cuslines 2.2__tar.gz → 2.2.2__tar.gz

{cuslines-2.2 → cuslines-2.2.2}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: cuslines
-Version: 2.2
+Version: 2.2.2
 Summary: GPU-accelerated tractography package
 Project-URL: Homepage, https://github.com/dipy/GPUStreamlines
 Requires-Python: >=3.7

{cuslines-2.2 → cuslines-2.2.2}/cuslines/generic_tracker.py

@@ -3,6 +3,7 @@ import numpy as np
 from tqdm import tqdm
 from trx.trx_file_memmap import TrxFile
 from dipy.io.stateful_tractogram import Space, StatefulTractogram
+from dipy.tracking.streamlinespeed import compress_streamlines
 from nibabel.streamlines.array_sequence import ArraySequence
 from nibabel.streamlines.tractogram import Tractogram
 
@@ -16,11 +17,55 @@ class GenericTracker:
     def __exit__(self, exc_type, exc, tb):
         return False
 
+    def set_compression_parameters(self, pos_dtype=np.float32, linearize=False, tol_error=0.1, max_segment_length=10):
+        """
+        Set compression parameters to compress generated streamlines.
+        Only works with TRX.
+        
+        Parameters
+        ----------
+        pos_dtype : dtype, optional
+        Data type to use for the positions of the streamlines.
+        Default: np.float32
+
+        linearize : bool, optional
+        Whether to linearize the streamlines using [1].
+        Default: False
+
+        tol_error : float, optional
+        If linearize is true, tolerance error in mm.
+        Default: 0.1
+
+        max_segment_length : float, optional
+        If linearize is true, maximum length in mm of any given segment produced by the compression.
+        Default: 10
+    
+        References
+        ----------
+        [1] Caroline Presseau, Pierre-Marc Jodoin, Jean-Christophe Houde, and Maxime Descoteaux.
+            A new compression format for fiber tracking datasets.
+            NeuroImage, 109:73-83, 2015. URL: 10.1016/j.neuroimage.2014.12.058
+        """
+        self.pos_dtype = pos_dtype
+        self.linearize = linearize
+        self.tol_error = tol_error
+        self.max_segment_length = max_segment_length
+
+
     def _ngpus(self):
-        if hasattr(self, "ngpus"):
-            return self.ngpus
-        else:
-            return 1
+        return getattr(self, "ngpus", 1)
+
+    def _pos_dtype(self):
+        return getattr(self, "pos_dtype", np.float16)
+
+    def _linearize(self):
+        return getattr(self, "linearize", False)
+
+    def _tol_error(self):
+        return getattr(self, "tol_error", 0.1)
+
+    def _max_segment_length(self):
+        return getattr(self, "max_segment_length", np.inf)
 
     def _divide_chunks(self, seeds):
         global_chunk_sz = self.chunk_size * self._ngpus()
@@ -58,7 +103,7 @@ class GenericTracker:
         # trx files use memory mapping
         trx_reference = TrxFile(reference=ref_img)
         trx_reference.streamlines._data = trx_reference.streamlines._data.astype(
-            np.float32
+            self._pos_dtype()
         )
         trx_reference.streamlines._offsets = trx_reference.streamlines._offsets.astype(
             np.uint64
@@ -81,9 +126,20 @@ class GenericTracker:
                     self.seed_propagator.as_array_sequence(),
                     affine_to_rasmm=ref_img.affine,
                 )
+                if len(tractogram) == 0:
+                    continue
+
                 tractogram.to_world()
                 sls = tractogram.streamlines
 
+                if self._linearize():
+                    sls = ArraySequence(compress_streamlines(
+                        sls,
+                        tol_error=self._tol_error(),
+                        max_segment_length=self._max_segment_length(),
+                    ))
+                sls._data = sls._data.astype(self._pos_dtype())
+
                 new_offsets_idx = offsets_idx + len(sls._offsets)
                 new_sls_data_idx = sls_data_idx + len(sls._data)
 

{cuslines-2.2 → cuslines-2.2.2}/cuslines/metal/mt_tractography.py

@@ -5,25 +5,16 @@ we wrap numpy arrays as Metal shared buffers with zero copies.
 """
 
 import numpy as np
-from tqdm import tqdm
 import logging
 from math import radians
 
 from cuslines.metal.mutils import (
-    REAL_SIZE,
     REAL_DTYPE,
-    aligned_array,
-    PAGE_SIZE,
-    checkMetalError,
 )
 
 from cuslines.metal.mt_direction_getters import MetalGPUDirectionGetter, MetalBootDirectionGetter
 from cuslines.metal.mt_propagate_seeds import MetalSeedBatchPropagator
-
-from trx.trx_file_memmap import TrxFile
-from nibabel.streamlines.tractogram import Tractogram
-from nibabel.streamlines.array_sequence import ArraySequence, MEGABYTE
-from dipy.io.stateful_tractogram import Space, StatefulTractogram
+from cuslines.generic_tracker import GenericTracker
 
 logger = logging.getLogger("GPUStreamlines")
 
@@ -64,7 +55,7 @@ def _buffer_as_array(buf, dtype, shape):
     return np.frombuffer(memview, dtype=dtype, count=count).reshape(shape)
 
 
-class MetalGPUTracker:
+class MetalGPUTracker(GenericTracker):
     def __init__(
         self,
         dg: MetalGPUDirectionGetter,
@@ -177,82 +168,3 @@ class MetalGPUTracker:
         self.dg.gen_pipeline = None
         self._allocated = False
         return False
-
-    def _divide_chunks(self, seeds):
-        global_chunk_sz = self.chunk_size  # single GPU
-        nchunks = (seeds.shape[0] + global_chunk_sz - 1) // global_chunk_sz
-        return global_chunk_sz, nchunks
-
-    def generate_sft(self, seeds, ref_img):
-        global_chunk_sz, nchunks = self._divide_chunks(seeds)
-        buffer_size = 0
-        generators = []
-
-        with tqdm(total=seeds.shape[0]) as pbar:
-            for idx in range(nchunks):
-                chunk = seeds[idx * global_chunk_sz : (idx + 1) * global_chunk_sz]
-                self.seed_propagator.propagate(chunk)
-                buffer_size += self.seed_propagator.get_buffer_size()
-                generators.append(self.seed_propagator.as_generator())
-                pbar.update(chunk.shape[0])
-
-        array_sequence = ArraySequence(
-            (item for gen in generators for item in gen), buffer_size
-        )
-        return StatefulTractogram(array_sequence, ref_img, Space.VOX)
-
-    def generate_trx(self, seeds, ref_img):
-        global_chunk_sz, nchunks = self._divide_chunks(seeds)
-
-        sl_len_guess = 100
-        sl_per_seed_guess = 2
-        n_sls_guess = sl_per_seed_guess * seeds.shape[0]
-
-        trx_reference = TrxFile(reference=ref_img)
-        trx_reference.streamlines._data = trx_reference.streamlines._data.astype(np.float32)
-        trx_reference.streamlines._offsets = trx_reference.streamlines._offsets.astype(np.uint64)
-
-        trx_file = TrxFile(
-            nb_streamlines=n_sls_guess,
-            nb_vertices=n_sls_guess * sl_len_guess,
-            init_as=trx_reference,
-        )
-        offsets_idx = 0
-        sls_data_idx = 0
-
-        with tqdm(total=seeds.shape[0]) as pbar:
-            for idx in range(int(nchunks)):
-                chunk = seeds[idx * global_chunk_sz : (idx + 1) * global_chunk_sz]
-                self.seed_propagator.propagate(chunk)
-                tractogram = Tractogram(
-                    self.seed_propagator.as_array_sequence(),
-                    affine_to_rasmm=ref_img.affine,
-                )
-                tractogram.to_world()
-                sls = tractogram.streamlines
-
-                new_offsets_idx = offsets_idx + len(sls._offsets)
-                new_sls_data_idx = sls_data_idx + len(sls._data)
-
-                if (
-                    new_offsets_idx > trx_file.header["NB_STREAMLINES"]
-                    or new_sls_data_idx > trx_file.header["NB_VERTICES"]
-                ):
-                    logger.info("TRX resizing...")
-                    trx_file.resize(
-                        nb_streamlines=new_offsets_idx * 2,
-                        nb_vertices=new_sls_data_idx * 2,
-                    )
-
-                trx_file.streamlines._data[sls_data_idx:new_sls_data_idx] = sls._data
-                trx_file.streamlines._offsets[offsets_idx:new_offsets_idx] = (
-                    sls_data_idx + sls._offsets
-                )
-                trx_file.streamlines._lengths[offsets_idx:new_offsets_idx] = sls._lengths
-
-                offsets_idx = new_offsets_idx
-                sls_data_idx = new_sls_data_idx
-                pbar.update(chunk.shape[0])
-
-        trx_file.resize()
-        return trx_file

{cuslines-2.2 → cuslines-2.2.2}/cuslines/numba/nu_tractography.py

@@ -2,9 +2,7 @@ import math
 from math import radians
 
 import numpy as np
-from dipy.io.stateful_tractogram import Space, StatefulTractogram
 from nibabel.streamlines.array_sequence import ArraySequence, MEGABYTE
-from tqdm import tqdm
 
 from cuslines.generic_tracker import GenericTracker
 from cuslines.numba_njit.num_streamlines_numba import getNumStreamlinesProb_generator

{cuslines-2.2 → cuslines-2.2.2}/cuslines/webgpu/wg_tractography.py

@@ -5,28 +5,22 @@ readbacks via device.queue.read_buffer() (similar to CUDA's cudaMemcpy).
 """
 
 import numpy as np
-from tqdm import tqdm
 import logging
 from math import radians
 
 from cuslines.webgpu.wgutils import (
-    REAL_SIZE,
     REAL_DTYPE,
     create_buffer_from_data,
 )
 
 from cuslines.webgpu.wg_direction_getters import WebGPUDirectionGetter, WebGPUBootDirectionGetter
 from cuslines.webgpu.wg_propagate_seeds import WebGPUSeedBatchPropagator
-
-from trx.trx_file_memmap import TrxFile
-from nibabel.streamlines.tractogram import Tractogram
-from nibabel.streamlines.array_sequence import ArraySequence, MEGABYTE
-from dipy.io.stateful_tractogram import Space, StatefulTractogram
+from cuslines.generic_tracker import GenericTracker
 
 logger = logging.getLogger("GPUStreamlines")
 
 
-class WebGPUTracker:
+class WebGPUTracker(GenericTracker):
     def __init__(
         self,
         dg: WebGPUDirectionGetter,
@@ -207,86 +201,3 @@ class WebGPUTracker:
         self.device = None
         self._allocated = False
         return False
-
-    def _divide_chunks(self, seeds):
-        global_chunk_sz = self.chunk_size  # single GPU
-        nchunks = (seeds.shape[0] + global_chunk_sz - 1) // global_chunk_sz
-        return global_chunk_sz, nchunks
-
-    def generate_sft(self, seeds, ref_img):
-        global_chunk_sz, nchunks = self._divide_chunks(seeds)
-        buffer_size = 0
-        generators = []
-
-        with tqdm(total=seeds.shape[0]) as pbar:
-            for idx in range(nchunks):
-                chunk = seeds[idx * global_chunk_sz : (idx + 1) * global_chunk_sz]
-                self.seed_propagator.propagate(chunk)
-                buffer_size += self.seed_propagator.get_buffer_size()
-                generators.append(self.seed_propagator.as_generator())
-                pbar.update(chunk.shape[0])
-
-        array_sequence = ArraySequence(
-            (item for gen in generators for item in gen), buffer_size
-        )
-        return StatefulTractogram(array_sequence, ref_img, Space.VOX)
-
-    def generate_trx(self, seeds, ref_img):
-        global_chunk_sz, nchunks = self._divide_chunks(seeds)
-
-        sl_len_guess = 100
-        sl_per_seed_guess = 2
-        n_sls_guess = sl_per_seed_guess * seeds.shape[0]
-
-        trx_reference = TrxFile(reference=ref_img)
-        trx_reference.streamlines._data = trx_reference.streamlines._data.astype(
-            np.float32
-        )
-        trx_reference.streamlines._offsets = trx_reference.streamlines._offsets.astype(
-            np.uint64
-        )
-
-        trx_file = TrxFile(
-            nb_streamlines=n_sls_guess,
-            nb_vertices=n_sls_guess * sl_len_guess,
-            init_as=trx_reference,
-        )
-        offsets_idx = 0
-        sls_data_idx = 0
-
-        with tqdm(total=seeds.shape[0]) as pbar:
-            for idx in range(int(nchunks)):
-                chunk = seeds[idx * global_chunk_sz : (idx + 1) * global_chunk_sz]
-                self.seed_propagator.propagate(chunk)
-                tractogram = Tractogram(
-                    self.seed_propagator.as_array_sequence(),
-                    affine_to_rasmm=ref_img.affine,
-                )
-                tractogram.to_world()
-                sls = tractogram.streamlines
-
-                new_offsets_idx = offsets_idx + len(sls._offsets)
-                new_sls_data_idx = sls_data_idx + len(sls._data)
-
-                if (
-                    new_offsets_idx > trx_file.header["NB_STREAMLINES"]
-                    or new_sls_data_idx > trx_file.header["NB_VERTICES"]
-                ):
-                    logger.info("TRX resizing...")
-                    trx_file.resize(
-                        nb_streamlines=new_offsets_idx * 2,
-                        nb_vertices=new_sls_data_idx * 2,
-                    )
-
-                trx_file.streamlines._data[sls_data_idx:new_sls_data_idx] = sls._data
-                trx_file.streamlines._offsets[offsets_idx:new_offsets_idx] = (
-                    sls_data_idx + sls._offsets
-                )
-                trx_file.streamlines._lengths[offsets_idx:new_offsets_idx] = sls._lengths
-
-                offsets_idx = new_offsets_idx
-                sls_data_idx = new_sls_data_idx
-                pbar.update(chunk.shape[0])
-
-        trx_file.resize()
-        return trx_file

{cuslines-2.2 → cuslines-2.2.2}/cuslines.egg-info/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: cuslines
-Version: 2.2
+Version: 2.2.2
 Summary: GPU-accelerated tractography package
 Project-URL: Homepage, https://github.com/dipy/GPUStreamlines
 Requires-Python: >=3.7

{cuslines-2.2 → cuslines-2.2.2}/run_gpu_streamlines.py

@@ -127,7 +127,7 @@ parser.add_argument(
 parser.add_argument(
     "--ngpus", type=int, default=1, help="number of GPUs to use if using gpu"
 )
-parser.add_argument("--write-method", type=str, default="trk", help="Can be trx or trk")
+parser.add_argument("--write-method", type=str, default="trx", help="Can be trx or trk")
 parser.add_argument(
     "--max-angle", type=float, default=60, help="max angle (in degrees)"
 )