Cuda Python

.gitignore

@@ -0,0 +1,6 @@
+# Python bytecode
+**/*.pyc
+**/__pycache__/
+*.pyo
+*.pyd
+

Dockerfile

@@ -16,17 +16,20 @@ RUN wget https://github.com/Kitware/CMake/releases/download/v3.24.0/cmake-3.24.0
     && mkdir /opt/cmake \
     && /tmp/cmake-install.sh --skip-license --prefix=/opt/cmake \
     && rm /tmp/cmake-install.sh
-ENV PATH /opt/cmake/bin:${PATH}
+ENV PATH=/opt/cmake/bin:${PATH}
 
 RUN curl -L "https://repo.anaconda.com/miniconda/Miniconda3-latest-Linux-x86_64.sh" \
     -o "/tmp/Miniconda3.sh"
 RUN bash /tmp/Miniconda3.sh -b -p /opt/anaconda
 RUN rm -rf /tmp/Miniconda3.sh
 RUN cd /opt && eval "$(/opt/anaconda/bin/conda shell.bash hook)"
-ENV PATH /opt/anaconda/bin:${PATH}
-ENV LD_LIBRARY_PATH /opt/anaconda/lib:${LD_LIBRARY_PATH}
+ENV PATH=/opt/anaconda/bin:${PATH}
+ENV LD_LIBRARY_PATH=/opt/anaconda/lib:${LD_LIBRARY_PATH}
 
 # python prereqs
+RUN conda tos accept --override-channels --channel conda-forge
+RUN conda tos accept --override-channels --channel https://repo.anaconda.com/pkgs/main
+RUN conda tos accept --override-channels --channel https://repo.anaconda.com/pkgs/r
 RUN conda install -c conda-forge git
 RUN pip install numpy>=2.0.0
 RUN pip install scipy>=1.13.0 cython nibabel dipy tqdm

README.md

@@ -48,63 +48,7 @@ Destroy GPUTracker...
 
 Note that if you experience memory errors, you can adjust the `--chunk-size` flag.
 
-To run on more seeds, we suggest enabling the `--use-fast-write` flag in the GPU script to not get bottlenecked by writing files. Here is a comparison running on 500K seeds on 1 GPU with and without this flag:
-
-Without `--use-fast-write`:
-```
-$ python run_gpu_streamlines.py --output-prefix small --nseeds 500000 --ngpus 1
-parsing arguments
-Fitting Tensor
-Computing anisotropy measures (FA,MD,RGB)
-slowadcodf
-Bootstrap direction getter
-streamline gen
-Creating GPUTracker with 1 GPUs...
-Generated 143891 streamlines from 100000 seeds, time: 7.978902339935303 s
-Saved streamlines to small.1_5.trk, time 11.439777851104736 s
-Generated 151932 streamlines from 100000 seeds, time: 10.155118703842163 s
-Saved streamlines to small.2_5.trk, time 12.438884019851685 s
-Generated 146971 streamlines from 100000 seeds, time: 9.822870016098022 s
-Saved streamlines to small.3_5.trk, time 12.377111673355103 s
-Generated 153824 streamlines from 100000 seeds, time: 11.133368968963623 s
-Saved streamlines to small.4_5.trk, time 13.317519187927246 s
-Generated 162004 streamlines from 100000 seeds, time: 13.19784665107727 s
-Saved streamlines to small.5_5.trk, time 14.21276593208313 s
-Completed processing 500000 seeds.
-Initialization time: 14.789637088775635 sec
-Streamline generation total time: 116.0746865272522 sec
-        Streamline processing: 52.28810667991638 sec
-        File writing: 63.7860586643219 sec
-Destroy GPUTracker...
-```
-
-With `--use-fast-write`:
-```
-$ python run_gpu_streamlines.py --output-prefix small --nseeds 500000 --ngpus 1 --use-fast-write
-parsing arguments
-Fitting Tensor
-Computing anisotropy measures (FA,MD,RGB)
-slowadcodf
-Bootstrap direction getter
-streamline gen
-Creating GPUTracker with 1 GPUs...
-Generated 143891 streamlines from 100000 seeds, time: 7.962322473526001 s
-Saved streamlines to small.1_5_*.trk, time 0.1053612232208252 s
-Generated 151932 streamlines from 100000 seeds, time: 10.148677825927734 s
-Saved streamlines to small.2_5_*.trk, time 0.1606450080871582 s
-Generated 146971 streamlines from 100000 seeds, time: 9.811130285263062 s
-Saved streamlines to small.3_5_*.trk, time 0.571892499923706 s
-Generated 153824 streamlines from 100000 seeds, time: 11.186563968658447 s
-Saved streamlines to small.4_5_*.trk, time 0.3091111183166504 s
-Generated 162004 streamlines from 100000 seeds, time: 13.282683610916138 s
-Saved streamlines to small.5_5_*.trk, time 0.7107999324798584 s
-Completed processing 500000 seeds.
-Initialization time: 14.705361366271973 sec
-Streamline generation total time: 54.24975609779358 sec
-        Streamline processing: 52.39137816429138 sec
-        File writing: 1.8578097820281982 sec
-Destroy GPUTracker...
-```
+To run on more seeds, we suggest setting the `--write-method trx` flag in the GPU script to not get bottlenecked by writing files.
 
 ## Running on AWS with Docker
 First, set up an AWS instance with GPU and ssh into it (we recommend a P3 instance with at least 1 V100 16 GB GPU and a Deep Learning AMI Ubuntu 18.04 v 33.0.). Then do the following:

cuslines/__init__.py

@@ -0,0 +1,13 @@
+from .cuda_python import (
+    GPUTracker,
+    ProbDirectionGetter,
+    PttDirectionGetter,
+    BootDirectionGetter
+)
+
+__all__ = [
+    "GPUTracker",
+    "ProbDirectionGetter",
+    "PttDirectionGetter",
+    "BootDirectionGetter"
+]