New Line Class with Back Project Method

.github/workflows/long_workflow.yml

@@ -1,14 +1,14 @@
 name: ASPIRE Python Long Running Test Suite
 
 on:
-  push:
-    branches:
-      - 'main'
-      - 'develop'
+  pull_request:
+    types: [opened, synchronize, reopened, ready_for_review]
 
 jobs:
   expensive_tests:
     runs-on: self-hosted
+    # Only run on review ready pull_requests
+    if: ${{ github.event_name == 'pull_request' && github.event.pull_request.draft == false }}
     timeout-minutes: 360
     steps:
     - uses: actions/checkout@v4
@@ -33,8 +33,9 @@ jobs:
         cat ${WORK_DIR}/config.yaml
     - name: Run
       run: |
+        export OMP_NUM_THREADS=1
         ASPIREDIR=${{ env.WORK_DIR }} python -c \
         "import aspire; print(aspire.config['ray']['temp_dir'])"
-        ASPIREDIR=${{ env.WORK_DIR }} python -m pytest -m "expensive" --durations=0
+        ASPIREDIR=${{ env.WORK_DIR }} python -m pytest -n8 -m "expensive" --durations=0
     - name: Cleanup
       run: rm -rf ${{ env.WORK_DIR }}

.github/workflows/workflow.yml

@@ -147,7 +147,10 @@ jobs:
         echo "Stash the WORK_DIR to GitHub env so we can clean it up later."
         echo "WORK_DIR=${WORK_DIR}" >> $GITHUB_ENV
         echo -e "ray:\n    temp_dir: ${WORK_DIR}\n" > ${WORK_DIR}/config.yaml
-        echo -e "common:\n    cache_dir: ${CI_CACHE_DIR}\n" >> ${WORK_DIR}/config.yaml
+        echo -e "common:" >> ${WORK_DIR}/config.yaml
+        echo -e "    cache_dir: ${CI_CACHE_DIR}" >> ${WORK_DIR}/config.yaml
+        echo -e "    numeric: cupy" >> ${WORK_DIR}/config.yaml
+        echo -e "    fft: cupy\n" >> ${WORK_DIR}/config.yaml
         echo "Log the config: ${WORK_DIR}/config.yaml"
         cat ${WORK_DIR}/config.yaml
     - name: Run

docs/source/installation.rst

@@ -129,25 +129,17 @@ an M1 laptop:
 Installing GPU Extensions
 *************************
 
-ASPIRE does support GPUs, depending on several external packages.  The
-collection of GPU extensions can be installed using ``pip``.
-Extensions are grouped based on CUDA versions.  To find the CUDA
-driver version, run ``nvidia-smi`` on the intended system.
+ASPIRE does support using a GPU, depending on several external
+packages.  The collection of GPU extensions can be installed using
+``pip``.  Extensions are grouped based on CUDA versions.  To find the
+CUDA driver version, run ``nvidia-smi`` on the intended system.
 
 .. list-table:: CUDA GPU Extension Versions
    :widths: 25 25
    :header-rows: 1
 
    * - CUDA Version
      - ASPIRE Extension
-   * - 10.2
-     - gpu-102
-   * - 11.0
-     - gpu-110
-   * - 11.1
-     - gpu-111
-   * - >=11.2
-     - gpu-11x
    * - >=12
      - gpu-12x
 
@@ -164,12 +156,15 @@ the command below would install GPU packages required for ASPIRE.
 
 
 By default if the required GPU extensions are correctly installed,
-ASPIRE should automatically begin using the GPU for select components
-(such as those using ``nufft``).
-
-Because GPU extensions depend on several third party packages and
-libraries, we can only offer limited support if one of the packages
-has a problem on your system.
+ASPIRE should automatically begin using the GPU calls to our ``nufft`` module.
+
+Using GPU in other areas of the code is still an experimental feature
+and requires a minor configuration setting to enable ``cupy``. See the
+:ref:`sphx_glr_auto_tutorials_configuration.py` for details. Because
+GPU extensions depend on several third party softwares and machines
+vary wildly, we can only offer limited support if one of the packages
+has a problem on your system.  We are currently expanding GPU code
+coverage.
 
 Generating Documentation
 ************************

gallery/experiments/experimental_abinitio_pipeline_10081.py

@@ -59,7 +59,11 @@
 
 # Create a source object for the experimental images
 src = RelionSource(
-    starfile_in, pixel_size=pixel_size, max_rows=n_imgs, data_folder=data_folder
+    starfile_in,
+    pixel_size=pixel_size,
+    max_rows=n_imgs,
+    data_folder=data_folder,
+    symmetry_group="C4",
 )
 
 # Downsample the images
@@ -115,12 +119,13 @@
 # Volume Reconstruction
 # ----------------------
 #
-# Using the oriented source, attempt to reconstruct a volume.
-# Since this is a Cn symmetric molecule, as indicated by
-# ``symmetry="C4"`` above, the ``avgs`` images set will be repeated
-# for each of the 3 additional rotations during the back-projection
-# step.  This boosts the effective number of images used in the
-# reconstruction from ``n_classes`` to ``4*n_classes``.
+# Using the oriented source, attempt to reconstruct a volume.  Since
+# this is a Cn symmetric molecule, as specified by ``RelionSource(...,
+# symmetry_group="C4", ...)``, the ``symmetry_group`` source attribute
+# will flow through the pipeline to ``avgs``. Then each image will be
+# repeated for each of the 3 additional rotations during
+# back-projection.  This boosts the effective number of images used in
+# the reconstruction from ``n_classes`` to ``4*n_classes``.
 
 logger.info("Begin Volume reconstruction")
 

gallery/tutorials/configuration.py

@@ -102,6 +102,36 @@
     time.sleep(1)
 print("Done Loop 2\n")
 
+# %%
+# Enabling GPU Acceleration
+# -------------------------
+# Enabling GPU acceleration requires installing supporting software
+# packages and small config changes. Installing the supporting
+# software is most easily accomplished by installing ASPIRE with one
+# of the published GPU extensions, for example ``pip install
+# "aspire[dev,gpu_12x]"``.  Once the packages are installed users
+# should find that the NUFFT calls are automatically running on the
+# GPU.  Additional acceleration is achieved by enabling `cupy` for
+# `numeric` and `fft` components.
+#
+#    .. code-block:: yaml
+#
+#      common:
+#          # numeric module to use - one of numpy/cupy
+#          numeric: cupy
+#          # fft backend to use - one of pyfftw/scipy/cupy/mkl
+#          fft: cupy
+#
+# Alternatively, like other config options, this can be changed
+# dynamically with code.
+#
+#    .. code-block:: python
+#
+#      from aspire import config
+#
+#      config["common"]["numeric"] = "cupy"
+#      config["common"]["fft"] = "cupy"
+#
 
 # %%
 # Resolution

pyproject.toml

@@ -61,11 +61,7 @@ dependencies = [
 "Source" = "https://github.com/ComputationalCryoEM/ASPIRE-Python"
 
 [project.optional-dependencies]
-gpu-102 = ["pycuda", "cupy-cuda102", "cufinufft==1.3"]
-gpu-110 = ["pycuda", "cupy-cuda110", "cufinufft==1.3"]
-gpu-111 = ["pycuda", "cupy-cuda111", "cufinufft==1.3"]
-gpu-11x = ["pycuda", "cupy-cuda11x", "cufinufft==1.3"]
-gpu-12x = ["pycuda", "cupy-cuda12x", "cufinufft==2.2.0"]
+gpu-12x = ["cupy-cuda12x", "cufinufft==2.2.0"]
 dev = [
     "black",
     "bumpversion",

src/aspire/abinitio/commonline_c3_c4.py

@@ -47,7 +47,7 @@ def __init__(
         n_theta=None,
         max_shift=0.15,
         shift_step=1,
-        epsilon=1e-3,
+        epsilon=1e-2,
         max_iters=1000,
         degree_res=1,
         seed=None,
@@ -691,7 +691,8 @@ def _J_sync_power_method(self, vijs):
         )
         while itr < max_iters and residual > epsilon:
             itr += 1
-            vec_new = self._signs_times_v(vijs, vec)
+            # Note, this appears to need double precision for accuracy in the following division.
+            vec_new = self._signs_times_v(vijs, vec).astype(np.float64, copy=False)
             vec_new = vec_new / norm(vec_new)
             residual = norm(vec_new - vec)
             vec = vec_new

src/aspire/basis/ffb_2d.py

@@ -58,6 +58,16 @@ def _build(self):
         # precompute the basis functions in 2D grids
         self._precomp = self._precomp()
 
+        # include the normalization factor of angular part into radial part
+        self.radial_norm = xp.asarray(self._precomp["radial"]) / xp.asarray(
+            np.expand_dims(self.angular_norms, 1)
+        )
+
+        # precompute weighted nodes
+        self.gl_weighted_nodes = xp.asarray(self._precomp["gl_weights"]) * xp.asarray(
+            self._precomp["gl_nodes"]
+        )
+
     def _precomp(self):
         """
         Precomute the basis functions on a polar Fourier grid
@@ -105,32 +115,31 @@ def _evaluate(self, v):
             coordinate basis. This is Image instance with resolution of `self.sz`
             and the first dimension correspond to remaining dimension of `v`.
         """
+        v = xp.asarray(v)
         sz_roll = v.shape[:-1]
         v = v.reshape(-1, self.count)
 
         # number of 2D image samples
         n_data = v.shape[0]
 
         # get information on polar grids from precomputed data
-        n_theta = np.size(self._precomp["freqs"], 2)
-        n_r = np.size(self._precomp["freqs"], 1)
+        n_theta = self._precomp["freqs"].shape[2]
+        n_r = self._precomp["freqs"].shape[1]
 
         # go through  each basis function and find corresponding coefficient
-        pf = np.zeros((n_data, 2 * n_theta, n_r), dtype=complex_type(self.dtype))
+        pf = xp.zeros((n_data, 2 * n_theta, n_r), dtype=complex_type(self.dtype))
 
         ind = 0
 
         idx = ind + np.arange(self.k_max[0], dtype=int)
 
-        # include the normalization factor of angular part into radial part
-        radial_norm = self._precomp["radial"] / np.expand_dims(self.angular_norms, 1)
-        pf[:, 0, :] = v[:, self._zero_angular_inds] @ radial_norm[idx]
-        ind = ind + np.size(idx)
+        pf[:, 0, :] = v[:, self._zero_angular_inds] @ self.radial_norm[idx]
+        ind = ind + idx.size
 
         ind_pos = ind
 
         for ell in range(1, self.ell_max + 1):
-            idx = ind + np.arange(self.k_max[ell], dtype=int)
+            idx = ind + xp.arange(self.k_max[ell], dtype=int)
             idx_pos = ind_pos + np.arange(self.k_max[ell], dtype=int)
             idx_neg = idx_pos + self.k_max[ell]
 
@@ -139,30 +148,25 @@ def _evaluate(self, v):
             if np.mod(ell, 2) == 1:
                 v_ell = 1j * v_ell
 
-            pf_ell = v_ell @ radial_norm[idx]
+            pf_ell = v_ell @ self.radial_norm[idx]
             pf[:, ell, :] = pf_ell
 
             if np.mod(ell, 2) == 0:
                 pf[:, 2 * n_theta - ell, :] = pf_ell.conjugate()
             else:
                 pf[:, 2 * n_theta - ell, :] = -pf_ell.conjugate()
 
-            ind = ind + np.size(idx)
+            ind = ind + idx.size
             ind_pos = ind_pos + 2 * self.k_max[ell]
 
         # 1D inverse FFT in the degree of polar angle
-        pf = 2 * pi * xp.asnumpy(fft.ifft(xp.asarray(pf), axis=1))
+        pf = 2 * xp.pi * fft.ifft(pf, axis=1)
 
         # Only need "positive" frequencies.
-        hsize = int(np.size(pf, 1) / 2)
+        hsize = int(pf.shape[1] / 2)
         pf = pf[:, 0:hsize, :]
-
-        for i_r in range(0, n_r):
-            pf[..., i_r] = pf[..., i_r] * (
-                self._precomp["gl_weights"][i_r] * self._precomp["gl_nodes"][i_r]
-            )
-
-        pf = np.reshape(pf, (n_data, n_r * n_theta))
+        pf *= self.gl_weighted_nodes[None, None, :]
+        pf = pf.reshape(n_data, n_r * n_theta)
 
         # perform inverse non-uniformly FFT transform back to 2D coordinate basis
         freqs = m_reshape(self._precomp["freqs"], (2, n_r * n_theta))
@@ -172,7 +176,7 @@ def _evaluate(self, v):
         # Return X as Image instance with the last two dimensions as *self.sz
         x = x.reshape((*sz_roll, *self.sz))
 
-        return x
+        return xp.asnumpy(x)
 
     def _evaluate_t(self, x):
         """
@@ -193,56 +197,51 @@ def _evaluate_t(self, x):
         n_images = x.shape[0]
 
         # resamping x in a polar Fourier gird using nonuniform discrete Fourier transform
-        pf = nufft(x, 2 * pi * freqs)
-        pf = np.reshape(pf, (n_images, n_r, n_theta))
+        pf = nufft(xp.asarray(x), 2 * pi * freqs)
+        pf = pf.reshape(n_images, n_r, n_theta)
 
         # Recover "negative" frequencies from "positive" half plane.
-        pf = np.concatenate((pf, pf.conjugate()), axis=2)
+        pf = xp.concatenate((pf, pf.conjugate()), axis=2)
 
         # evaluate radial integral using the Gauss-Legendre quadrature rule
-        for i_r in range(0, n_r):
-            pf[:, i_r, :] = pf[:, i_r, :] * (
-                self._precomp["gl_weights"][i_r] * self._precomp["gl_nodes"][i_r]
-            )
+        pf = pf * self.gl_weighted_nodes[None, :, None]
 
         #  1D FFT on the angular dimension for each concentric circle
-        pf = 2 * pi / (2 * n_theta) * xp.asnumpy(fft.fft(xp.asarray(pf)))
+        pf = 2 * xp.pi / (2 * n_theta) * fft.fft(pf)
 
         # This only makes it easier to slice the array later.
-        v = np.zeros((n_images, self.count), dtype=x.dtype)
+        v = xp.zeros((n_images, self.count), dtype=x.dtype)
 
         # go through each basis function and find the corresponding coefficient
         ind = 0
-        idx = ind + np.arange(self.k_max[0])
+        idx = ind + xp.arange(self.k_max[0])
 
-        # include the normalization factor of angular part into radial part
-        radial_norm = self._precomp["radial"] / np.expand_dims(self.angular_norms, 1)
-        v[:, self._zero_angular_inds] = pf[:, :, 0].real @ radial_norm[idx].T
-        ind = ind + np.size(idx)
+        v[:, self._zero_angular_inds] = pf[:, :, 0].real @ self.radial_norm[idx].T
+        ind = ind + idx.size
 
         ind_pos = ind
         for ell in range(1, self.ell_max + 1):
-            idx = ind + np.arange(self.k_max[ell])
-            idx_pos = ind_pos + np.arange(self.k_max[ell])
+            idx = ind + xp.arange(self.k_max[ell])
+            idx_pos = ind_pos + xp.arange(self.k_max[ell])
             idx_neg = idx_pos + self.k_max[ell]
 
-            v_ell = pf[:, :, ell] @ radial_norm[idx].T
+            v_ell = pf[:, :, ell] @ self.radial_norm[idx].T
 
             if np.mod(ell, 2) == 0:
-                v_pos = np.real(v_ell)
-                v_neg = -np.imag(v_ell)
+                v_pos = v_ell.real
+                v_neg = -v_ell.imag
             else:
-                v_pos = np.imag(v_ell)
-                v_neg = np.real(v_ell)
+                v_pos = v_ell.imag
+                v_neg = v_ell.real
 
             v[:, idx_pos] = v_pos
             v[:, idx_neg] = v_neg
 
-            ind = ind + np.size(idx)
+            ind = ind + idx.size
 
             ind_pos = ind_pos + 2 * self.k_max[ell]
 
-        return v
+        return xp.asnumpy(v)
 
     def filter_to_basis_mat(self, f, **kwargs):
         """