[FIX] Add NUMBA backend; correctly handle symmetric/nonsymmetric spheres

cuslines/__init__.py

@@ -30,15 +30,14 @@ def _detect_backend():
         pass
     return None
 
-
 BACKEND = _detect_backend()
 
 if BACKEND == "metal":
     from cuslines.metal import (
         MetalBootDirectionGetter as BootDirectionGetter,
     )
     from cuslines.metal import (
-        MetalGPUTracker as GPUTracker,
+        MetalGPUTracker as Tracker,
     )
     from cuslines.metal import (
         MetalProbDirectionGetter as ProbDirectionGetter,
@@ -49,7 +48,7 @@ def _detect_backend():
 elif BACKEND == "cuda":
     from cuslines.cuda_python import (
         BootDirectionGetter,
-        GPUTracker,
+        GPUTracker as Tracker,
         ProbDirectionGetter,
         PttDirectionGetter,
     )
@@ -64,18 +63,24 @@ def _detect_backend():
         WebGPUPttDirectionGetter as PttDirectionGetter,
     )
     from cuslines.webgpu import (
-        WebGPUTracker as GPUTracker,
+        WebGPUTracker as Tracker,
     )
 else:
-    raise ImportError(
-        "No GPU backend available. Install either:\n"
-        "  - CUDA: pip install 'cuslines[cu13]' (NVIDIA GPU)\n"
-        "  - Metal: pip install 'cuslines[metal]' (Apple Silicon)\n"
-        "  - WebGPU: pip install 'cuslines[webgpu]' (cross-platform)"
+    from cuslines.numba import (
+        CPUBootDirectionGetter as BootDirectionGetter,
+    )
+    from cuslines.numba import (
+        CPUProbDirectionGetter as ProbDirectionGetter,
+    )
+    from cuslines.numba import (
+        CPUPttDirectionGetter as PttDirectionGetter,
+    )
+    from cuslines.numba import (
+        CPUTracker as Tracker,
     )
 
 __all__ = [
-    "GPUTracker",
+    "Tracker",
     "ProbDirectionGetter",
     "PttDirectionGetter",
     "BootDirectionGetter",

cuslines/cuda_c/generate_streamlines_cuda.cu

@@ -148,7 +148,7 @@ __device__ int get_direction_prob_d(curandStatePhilox4_32_10_t *st,
                                            dir.y*sphere_vertices[i].y+
                                            dir.z*sphere_vertices[i].z;
 
-                        if (FABS(dot) < cos_similarity) {
+                        if (APPLY_ABS_IF_SYM(dot) < cos_similarity) {
                                 __pmf_data_sh[i] = 0.0;
                         }
                 }

cuslines/cuda_c/globals.h

@@ -67,6 +67,12 @@
 
 #endif
 
+#if FULL_BASIS == 1
+  #define APPLY_ABS_IF_SYM(x) (x)
+#else
+  #define APPLY_ABS_IF_SYM(x) FABS(x)
+#endif
+
 #define MIN(x,y) (((x)<(y))?(x):(y))
 #define MAX(x,y) (((x)>(y))?(x):(y))
 #define POW2(n) (1 << (n))

cuslines/cuda_c/tracking_helpers.cu

@@ -206,9 +206,10 @@ __device__ int peak_directions_d(const REAL_T  *__restrict__ odf,
 
                                 int j = 0;
                                 for(; j < k; j++) {
-                                        const REAL_T cos = FABS(abc.x*dirs[j].x+
-                                                                abc.y*dirs[j].y+
-                                                                abc.z*dirs[j].z);
+                                        const REAL_T cos = APPLY_ABS_IF_SYM(
+                                                abc.x*dirs[j].x+
+                                                abc.y*dirs[j].y+
+                                                abc.z*dirs[j].z);
                                         if (cos > cos_similarity) {
                                                 break;
                                         }

cuslines/cuda_python/cu_direction_getters.py

@@ -82,6 +82,7 @@ def compile_program(self, gpu_tracker, debug: bool = False):
             "RNG_SEED": str(int(gpu_tracker.rng_seed)),
             "SAMPLM_NR": str(int(gpu_tracker.samplm_nr)),
             "NUM_EDGES": str(int(gpu_tracker.nedges)),
+            "FULL_BASIS": "1" if gpu_tracker.full_basis else "0",
             "EXCESS_ALLOC_FACT": str(int(EXCESS_ALLOC_FACT)),
             "MAX_SLINES_PER_SEED": str(int(MAX_SLINES_PER_SEED)),
             "MAX_SLINE_LEN": str(int(MAX_SLINE_LEN)),

cuslines/cuda_python/cu_tractography.py

@@ -10,6 +10,8 @@
 from tqdm import tqdm
 from trx.trx_file_memmap import TrxFile
 
+from cuslines.generic_tracker import GenericTracker
+
 from cuslines.cuda_python.cu_direction_getters import (
     BootDirectionGetter,
     GPUDirectionGetter,
@@ -31,7 +33,7 @@
 # Remove small/long streamlines on gpu
 
 
-class GPUTracker:
+class GPUTracker(GenericTracker):
     def __init__(
         self,
         dg: GPUDirectionGetter,
@@ -40,6 +42,7 @@ def __init__(
         stop_threshold: float,
         sphere_vertices: np.ndarray,
         sphere_edges: np.ndarray,
+        full_basis: bool = False,
         max_angle: float = radians(60),
         step_size: float = 0.5,
         min_pts=0,
@@ -70,6 +73,9 @@ def __init__(
             Vertices of the sphere used for direction sampling.
         sphere_edges : np.ndarray
             Edges of the sphere used for direction sampling.
+        full_basis : bool, optional
+            Whether to use full basis for spherical harmonics
+            default: False
         max_angle : float, optional
             Maximum angle (in radians) between steps
             default: radians(60)
@@ -143,6 +149,7 @@ def __init__(
         self.rng_seed = int(rng_seed)
         self.rng_offset = int(rng_offset)
         self.chunk_size = int(chunk_size)
+        self.full_basis = bool(full_basis)
 
         avail = checkCudaErrors(runtime.cudaGetDeviceCount())
         if self.ngpus > avail:
@@ -284,96 +291,3 @@ def __exit__(self, exc_type, exc, tb):
                 runtime.cudaStreamDestroy(self.streams[n]), hard_error=False
             )
         return False
-
-    def _divide_chunks(self, seeds):
-        global_chunk_sz = self.chunk_size * self.ngpus
-        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):
-                self.seed_propagator.propagate(
-                    seeds[idx * global_chunk_sz : (idx + 1) * global_chunk_sz]
-                )
-                buffer_size += self.seed_propagator.get_buffer_size()
-                generators.append(self.seed_propagator.as_generator())
-                pbar.update(
-                    seeds[idx * global_chunk_sz : (idx + 1) * global_chunk_sz].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)
-
-        # Will resize by a factor of 2 if these are exceeded
-        sl_len_guess = 100
-        sl_per_seed_guess = 2
-        n_sls_guess = sl_per_seed_guess * seeds.shape[0]
-
-        # trx files use memory mapping
-        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)):
-                self.seed_propagator.propagate(
-                    seeds[idx * global_chunk_sz : (idx + 1) * global_chunk_sz]
-                )
-                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 uses memmaps here
-                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(
-                    seeds[idx * global_chunk_sz : (idx + 1) * global_chunk_sz].shape[0]
-                )
-        trx_file.resize()
-
-        return trx_file

cuslines/generic_tracker.py

@@ -0,0 +1,116 @@
+import logging
+import numpy as np
+from tqdm import tqdm
+from trx.trx_file_memmap import TrxFile
+from dipy.io.stateful_tractogram import Space, StatefulTractogram
+from nibabel.streamlines.array_sequence import ArraySequence
+from nibabel.streamlines.tractogram import Tractogram
+
+logger = logging.getLogger("GPUStreamlines")
+
+
+class GenericTracker:
+    def __enter__(self):
+        return self
+
+    def __exit__(self, exc_type, exc, tb):
+        return False
+
+    def _ngpus(self):
+        if hasattr(self, "ngpus"):
+            return self.ngpus
+        else:
+            return 1
+
+    def _divide_chunks(self, seeds):
+        global_chunk_sz = self.chunk_size * self._ngpus()
+        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):
+                self.seed_propagator.propagate(
+                    seeds[idx * global_chunk_sz : (idx + 1) * global_chunk_sz]
+                )
+                buffer_size += self.seed_propagator.get_buffer_size()
+                generators.append(self.seed_propagator.as_generator())
+                pbar.update(
+                    seeds[idx * global_chunk_sz : (idx + 1) * global_chunk_sz].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)
+
+        # Will resize by a factor of 2 if these are exceeded
+        sl_len_guess = 100
+        sl_per_seed_guess = 2
+        n_sls_guess = sl_per_seed_guess * seeds.shape[0]
+
+        # trx files use memory mapping
+        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)):
+                self.seed_propagator.propagate(
+                    seeds[idx * global_chunk_sz : (idx + 1) * global_chunk_sz]
+                )
+                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 uses memmaps here
+                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(
+                    seeds[idx * global_chunk_sz : (idx + 1) * global_chunk_sz].shape[0]
+                )
+        trx_file.resize()
+
+        return trx_file