cuslines 2.2__tar.gz → 2.2.1__tar.gz

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

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

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

@@ -81,6 +81,9 @@ 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
 

{cuslines-2.2 → cuslines-2.2.1}/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.1}/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.1}/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.1}/cuslines.egg-info/PKG-INFO

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