From 98fc052dc13f4f31a788f83ae2d01882b332b9cc Mon Sep 17 00:00:00 2001 From: 36000 Date: Tue, 24 Feb 2026 17:08:29 +0900 Subject: [PATCH 1/3] [ENH] Optimize TRX output --- cuslines/cuda_python/cu_tractography.py | 8 ++++++-- 1 file changed, 6 insertions(+), 2 deletions(-) diff --git a/cuslines/cuda_python/cu_tractography.py b/cuslines/cuda_python/cu_tractography.py index 7064cc9..551e60b 100644 --- a/cuslines/cuda_python/cu_tractography.py +++ b/cuslines/cuda_python/cu_tractography.py @@ -49,7 +49,7 @@ def __init__( ngpus: int = 1, rng_seed: int = 0, rng_offset: int = 0, - chunk_size: int = 100000, + chunk_size: int = 25000, ): """ Initialize GPUTracker with necessary data. @@ -91,6 +91,9 @@ def __init__( rng_offset : int, optional Offset for random number generator default: 0 + chunk_size : int, optional + Number of seeds to process in each chunk per GPU + default: 25000 """ self.dataf = np.ascontiguousarray(dataf, dtype=REAL_DTYPE) self.metric_map = np.ascontiguousarray(stop_map, dtype=REAL_DTYPE) @@ -258,7 +261,7 @@ def generate_trx(self, seeds, ref_img): # Will resize by a factor of 2 if these are exceeded sl_len_guess = 100 - sl_per_seed_guess = 3 + sl_per_seed_guess = 4 n_sls_guess = sl_per_seed_guess * seeds.shape[0] # trx files use memory mapping @@ -267,6 +270,7 @@ def generate_trx(self, seeds, ref_img): nb_streamlines=n_sls_guess, nb_vertices=n_sls_guess * sl_len_guess, ) + trx_file.streamlines._data = trx_file.streamlines._data.astype(np.float32) trx_file.streamlines._offsets = trx_file.streamlines._offsets.astype(np.uint64) offsets_idx = 0 sls_data_idx = 0 From 30d90e0b84d4fbd77999a6fee0d57f5ea4d11237 Mon Sep 17 00:00:00 2001 From: 36000 Date: Tue, 24 Feb 2026 17:18:44 +0900 Subject: [PATCH 2/3] use as init trx --- cuslines/cuda_python/cu_tractography.py | 11 +++++++---- 1 file changed, 7 insertions(+), 4 deletions(-) diff --git a/cuslines/cuda_python/cu_tractography.py b/cuslines/cuda_python/cu_tractography.py index 551e60b..5f5ac41 100644 --- a/cuslines/cuda_python/cu_tractography.py +++ b/cuslines/cuda_python/cu_tractography.py @@ -255,7 +255,6 @@ def generate_sft(self, seeds, ref_img): ) return StatefulTractogram(array_sequence, ref_img, Space.VOX) - # TODO: performance: consider a way to just output in VOX space directly def generate_trx(self, seeds, ref_img): global_chunk_sz, nchunks = self._divide_chunks(seeds) @@ -265,13 +264,17 @@ def generate_trx(self, seeds, ref_img): 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_file.streamlines._data.astype(np.float32) + trx_reference.streamlines._offsets = trx_file.streamlines._offsets.astype(np.uint64) + trx_file = TrxFile( - reference=ref_img, nb_streamlines=n_sls_guess, nb_vertices=n_sls_guess * sl_len_guess, + init_as=trx_reference ) - trx_file.streamlines._data = trx_file.streamlines._data.astype(np.float32) - trx_file.streamlines._offsets = trx_file.streamlines._offsets.astype(np.uint64) offsets_idx = 0 sls_data_idx = 0 From 1a4e752dfa13575bf81e87e48e7ef21aaefbc380 Mon Sep 17 00:00:00 2001 From: 36000 Date: Tue, 24 Feb 2026 17:20:38 +0900 Subject: [PATCH 3/3] bf --- cuslines/cuda_python/cu_tractography.py | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/cuslines/cuda_python/cu_tractography.py b/cuslines/cuda_python/cu_tractography.py index 5f5ac41..9b539be 100644 --- a/cuslines/cuda_python/cu_tractography.py +++ b/cuslines/cuda_python/cu_tractography.py @@ -267,8 +267,8 @@ def generate_trx(self, seeds, ref_img): trx_reference = TrxFile( reference=ref_img ) - trx_reference.streamlines._data = trx_file.streamlines._data.astype(np.float32) - trx_reference.streamlines._offsets = trx_file.streamlines._offsets.astype(np.uint64) + 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,