Have a trajectory learner example and reform the documentation build flow

.github/workflows/master-cd.yml

@@ -29,42 +29,19 @@ jobs:
     if: success()
 
     steps:
-      - name: Checkout
-        uses: actions/checkout@v3
-      - name: Get history and tags for SCM versioning to work
-        run: |
-          git fetch --prune --unshallow
-          git fetch --depth=1 origin +refs/tags/*:refs/tags/*
-      - name: Set up Python
-        uses: actions/setup-python@v4
-        with:
-          python-version: "3.10"
-      - name: Install dependencies
-        shell: bash -l {0}
-        run: |
-          python -m pip install --upgrade pip
-          python -m pip install .[doc]
-          python -m pip install finufft
-
-      - name: Build API documentation
-        run: |
-          python -m sphinx docs docs_build
-
-      - name: Get the badge from CI
+      - name: Get the docs_build artifact
         uses: actions/download-artifact@v4
         with:
-          name: coverage_badge
-          path: docs_build/_static
-          github-token: ${{ secrets.GITHUB_TOKEN }}
+          name: docs_final
+          path: docs_build
           run-id: ${{ github.event.workflow_run.id }}
+
+      - name: Display structure of docs
+        run: ls -R docs_build/
 
-      - name: Display data
-        run: ls -R
-        working-directory: docs_build/_static
       - name: Upload artifact
         uses: actions/upload-pages-artifact@v1
         with:
-          # Upload entire repository
           path: 'docs_build'
 
       - name: Deploy to GitHub Pages

.github/workflows/style.yml

@@ -1,4 +1,4 @@
-name: Style checking
+name: Style
 
 on:
   push:

.github/workflows/test-ci.yml

@@ -170,7 +170,7 @@ jobs:
           rm -rf venv
 
   test-examples:
-    runs-on: cpu
+    runs-on: gpu
     if: ${{ !contains(github.event.head_commit.message, 'style')}}
 
     steps:
@@ -194,6 +194,15 @@ jobs:
           python -m pip install --upgrade pip
           python -m pip install -e .[test,dev]
           python -m pip install finufft pooch brainweb-dl torch
+
+      - name: Install GPU related interfaces
+        run: |
+          export CUDA_BIN_PATH=/usr/local/cuda-11.8/
+          export PATH=/usr/local/cuda-11.8/bin/:${PATH}
+          export LD_LIBRARY_PATH=/usr/local/cuda-11.8/lib64/:${LD_LIBRARY_PATH}
+          pip install cupy-cuda11x
+          pip install torch --index-url https://download.pytorch.org/whl/cu118
+          python -m pip install gpuNUFFT cufinufft
 
       - name: Run examples
         shell: bash
@@ -268,3 +277,80 @@ jobs:
         with:
           name: coverage_badge
           path: coverage_badge.svg
+
+  BuildDocs:
+    name: Build API Documentation
+    runs-on: gpu
+    if: ${{ contains(github.event.head_commit.message, 'docs_build')}} or ${{github.ref == 'refs/heads/master'}}
+    steps:
+      - name: Checkout
+        uses: actions/checkout@v3
+
+      - name: Get history and tags for SCM versioning to work
+        run: |
+          git fetch --prune --unshallow
+          git fetch --depth=1 origin +refs/tags/*:refs/tags/*
+
+      - name: Set up Python
+        uses: actions/setup-python@v4
+        with:
+          python-version: "3.10"
+
+      - name: Install dependencies
+        shell: bash -l {0}
+        run: |
+          python -m pip install --upgrade pip
+          python -m pip install .[doc]
+          python -m pip install finufft
+
+      - name: Install GPU related interfaces
+        run: |
+          export CUDA_BIN_PATH=/usr/local/cuda-11.8/
+          export PATH=/usr/local/cuda-11.8/bin/:${PATH}
+          export LD_LIBRARY_PATH=/usr/local/cuda-11.8/lib64/:${LD_LIBRARY_PATH}
+          pip install cupy-cuda11x
+          pip install torch --index-url https://download.pytorch.org/whl/cu118
+          python -m pip install gpuNUFFT cufinufft
+
+      - name: Build API documentation
+        run: |
+          python -m sphinx docs docs_build
+
+      - name: Display data
+        run: ls -R
+        working-directory: docs_build/_static
+
+      - name: Upload artifact
+        id: artifact-upload-step
+        uses: actions/upload-artifact@v4
+        with:
+          # Upload the docs
+          name: docs
+          path: 'docs_build'
+          retention-days: 5
+
+  CompileDocs:
+    name: Compile the coverage badge in docs
+    runs-on: ubuntu-latest
+    if: ${{ github.ref == 'refs/heads/master' }}
+    needs: [BuildDocs, coverage]
+    steps:
+      - name: Get the docs_build artifact
+        uses: actions/download-artifact@v4
+        with:
+          name: docs
+          path: docs_build
+
+      - name: Get the badge from CI
+        uses: actions/download-artifact@v4
+        with:
+          name: coverage_badge
+          path: docs_build/_static
+          github-token: ${{ secrets.GITHUB_TOKEN }}
+
+      - name: ReUpload artifacts
+        uses: actions/upload-artifact@v4
+        with:
+          name: docs_final
+          retention-days: 20
+          path: docs_build

docs/conf.py

@@ -77,8 +77,9 @@
     "examples_dirs": ["../examples/"],
     "gallery_dirs": ["generated/autoexamples"],
     "filename_pattern": "/example_",
-    "ignore_pattern": r"/(__init__|conftest|utils)\.py",
+    "ignore_pattern": r"(__init__|conftest|utils).py",
     "nested_sections": True,
+    "parallel": 3,
 }
 
 intersphinx_mapping = {

examples/GPU/README.rst

@@ -0,0 +1,6 @@
+.. _gpu_examples:
+
+GPU Examples
+------------
+
+This is a collection of examples showing features of mri-nufft, particularly those that are GPU-accelerated.

examples/example_density.py → examples/GPU/example_density.py

@@ -23,7 +23,7 @@
 # Create sample data
 # ------------------
 
-mri_2D = bwdl.get_mri(4, "T1")[80, ...].astype(np.float32)
+mri_2D = np.flipud(bwdl.get_mri(4, "T1")[80, ...]).astype(np.float32)
 
 print(mri_2D.shape)
 
@@ -136,18 +136,16 @@
 #    The Pipe method is currently only implemented for gpuNUFFT.
 
 # %%
-if check_backend("gpunufft"):
-    flat_traj = traj.reshape(-1, 2)
-    nufft = get_operator("gpunufft")(
-        traj, shape=mri_2D.shape, density={"name": "pipe", "osf": 2}
-    )
-    adjoint_manual = nufft.adj_op(kspace)
-    fig, axs = plt.subplots(1, 3, figsize=(15, 5))
-    axs[0].imshow(abs(mri_2D))
-    axs[0].set_title("Ground Truth")
-    axs[1].imshow(abs(adjoint))
-    axs[1].set_title("no density compensation")
-    axs[2].imshow(abs(adjoint_manual))
-    axs[2].set_title("Pipe density compensation")
-
-    print(nufft.density)
+flat_traj = traj.reshape(-1, 2)
+nufft = get_operator("gpunufft")(
+    traj, shape=mri_2D.shape, density={"name": "pipe", "osf": 2}
+)
+adjoint_manual = nufft.adj_op(kspace)
+fig, axs = plt.subplots(1, 3, figsize=(15, 5))
+axs[0].imshow(abs(mri_2D))
+axs[0].set_title("Ground Truth")
+axs[1].imshow(abs(adjoint))
+axs[1].set_title("no density compensation")
+axs[2].imshow(abs(adjoint_manual))
+axs[2].set_title("Pipe density compensation")
+print(nufft.density)