Skip to content

[FIX] allow full basis in accelerated tractography #40

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Closed
36000 wants to merge 1 commit into from
Closed

[FIX] allow full basis in accelerated tractography #40

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
1 commit
Select commit
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
2 changes: 1 addition & 1 deletion cuslines/cuda_c/generate_streamlines_cuda.cu
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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;
}
}
Expand Down
6 changes: 6 additions & 0 deletions cuslines/cuda_c/globals.h
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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))
Expand Down
7 changes: 4 additions & 3 deletions cuslines/cuda_c/tracking_helpers.cu
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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;
}
Expand Down
1 change: 1 addition & 0 deletions cuslines/cuda_python/cu_direction_getters.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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)),
Expand Down
5 changes: 5 additions & 0 deletions cuslines/cuda_python/cu_tractography.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -40,6 +40,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,
Comment on lines 41 to 45
Copy link

Copilot AI Apr 28, 2026

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Adding full_basis in the middle of the GPUTracker.__init__ positional parameters is a backward-incompatible API change: any downstream code calling GPUTracker(..., sphere_edges, max_angle, step_size, ...) positionally will silently shift arguments (e.g., max_angle becomes full_basis=True), leading to incorrect tracking without an obvious error. To preserve compatibility, move full_basis to the end of the parameter list or make it keyword-only (e.g., introduce a * before full_basis).

Copilot uses AI. Check for mistakes.
Comment on lines 41 to 45
Copy link

Copilot AI Apr 28, 2026

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

This adds full_basis only to the CUDA backend GPUTracker constructor, which makes the top-level cuslines.GPUTracker API backend-dependent (Metal/WebGPU trackers do not accept this argument). The repo docs state the three backends "share the same API surface", so this introduces an inconsistency for backend-agnostic callers. Consider adding the same full_basis parameter (even if currently ignored) to MetalGPUTracker and WebGPUTracker, and ideally propagate it into their shader code paths as well, or otherwise document/encapsulate the divergence.

Copilot uses AI. Check for mistakes.
min_pts=0,
Expand Down Expand Up @@ -70,6 +71,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)
Expand Down Expand Up @@ -143,6 +147,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:
Expand Down
Loading