From 8913bbadc8d02be5fa996102788f213378f227f1 Mon Sep 17 00:00:00 2001 From: Josh Carmichael Date: Thu, 13 Jun 2024 08:59:56 -0400 Subject: [PATCH 1/5] precompute radial_norm and gl_weighted_nodes in build. --- src/aspire/basis/ffb_2d.py | 38 +++++++++++++++++--------------------- 1 file changed, 17 insertions(+), 21 deletions(-) diff --git a/src/aspire/basis/ffb_2d.py b/src/aspire/basis/ffb_2d.py index 821800ee54..891ab2600d 100644 --- a/src/aspire/basis/ffb_2d.py +++ b/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.array(self._precomp["radial"]) / xp.array( + np.expand_dims(self.angular_norms, 1) + ) + + # precompute weighted nodes + self.gl_weighted_nodes = xp.array(self._precomp["gl_weights"]) * xp.array( + self._precomp["gl_nodes"] + ) + def _precomp(self): """ Precomute the basis functions on a polar Fourier grid @@ -123,11 +133,7 @@ def _evaluate(self, v): idx = ind + xp.arange(self.k_max[0], dtype=int) - # include the normalization factor of angular part into radial part - radial_norm = xp.array(self._precomp["radial"]) / xp.array( - np.expand_dims(self.angular_norms, 1) - ) - pf[:, 0, :] = v[:, xp.array(self._zero_angular_inds)] @ radial_norm[idx] + pf[:, 0, :] = v[:, xp.array(self._zero_angular_inds)] @ self.radial_norm[idx] ind = ind + idx.size ind_pos = ind @@ -142,7 +148,7 @@ 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: @@ -159,9 +165,7 @@ def _evaluate(self, v): # Only need "positive" frequencies. hsize = int(pf.shape[1] / 2) pf = pf[:, 0:hsize, :] - pf *= ( - xp.array(self._precomp["gl_weights"]) * xp.array(self._precomp["gl_nodes"]) - )[None, None, :] + 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 @@ -202,13 +206,7 @@ def _evaluate_t(self, x): pf = xp.concatenate((pf, pf.conjugate()), axis=2) # evaluate radial integral using the Gauss-Legendre quadrature rule - pf = ( - pf - * ( - xp.array(self._precomp["gl_weights"]) - * xp.array(self._precomp["gl_nodes"]) - )[None, :, None] - ) + pf *= self.gl_weighted_nodes[None, :, None] # 1D FFT on the angular dimension for each concentric circle pf = 2 * xp.pi / (2 * n_theta) * fft.fft(pf) @@ -221,10 +219,8 @@ def _evaluate_t(self, x): idx = ind + xp.arange(self.k_max[0]) # include the normalization factor of angular part into radial part - radial_norm = xp.array( - self._precomp["radial"] / np.expand_dims(self.angular_norms, 1) - ) - v[:, self._zero_angular_inds] = pf[:, :, 0].real @ radial_norm[idx].T + + v[:, self._zero_angular_inds] = pf[:, :, 0].real @ self.radial_norm[idx].T ind = ind + idx.size ind_pos = ind @@ -233,7 +229,7 @@ def _evaluate_t(self, x): 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 = v_ell.real From 665554d9bccf5f0eb689b8a9bcef3682dddc340f Mon Sep 17 00:00:00 2001 From: Josh Carmichael Date: Thu, 13 Jun 2024 09:07:46 -0400 Subject: [PATCH 2/5] remove comment --- src/aspire/basis/ffb_2d.py | 2 -- 1 file changed, 2 deletions(-) diff --git a/src/aspire/basis/ffb_2d.py b/src/aspire/basis/ffb_2d.py index 891ab2600d..1c7ed8cad1 100644 --- a/src/aspire/basis/ffb_2d.py +++ b/src/aspire/basis/ffb_2d.py @@ -218,8 +218,6 @@ def _evaluate_t(self, x): ind = 0 idx = ind + xp.arange(self.k_max[0]) - # include the normalization factor of angular part into radial part - v[:, self._zero_angular_inds] = pf[:, :, 0].real @ self.radial_norm[idx].T ind = ind + idx.size From 4443f8367969caf1170b6fc36ad8348ebcc13fce Mon Sep 17 00:00:00 2001 From: Josh Carmichael Date: Thu, 13 Jun 2024 10:42:35 -0400 Subject: [PATCH 3/5] use asarray --- src/aspire/basis/ffb_2d.py | 14 +++++++------- 1 file changed, 7 insertions(+), 7 deletions(-) diff --git a/src/aspire/basis/ffb_2d.py b/src/aspire/basis/ffb_2d.py index 1c7ed8cad1..f44e959149 100644 --- a/src/aspire/basis/ffb_2d.py +++ b/src/aspire/basis/ffb_2d.py @@ -59,12 +59,12 @@ def _build(self): self._precomp = self._precomp() # include the normalization factor of angular part into radial part - self.radial_norm = xp.array(self._precomp["radial"]) / xp.array( + 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.array(self._precomp["gl_weights"]) * xp.array( + self.gl_weighted_nodes = xp.asarray(self._precomp["gl_weights"]) * xp.asarray( self._precomp["gl_nodes"] ) @@ -115,7 +115,7 @@ 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.array(v) + v = xp.asarray(v) sz_roll = v.shape[:-1] v = v.reshape(-1, self.count) @@ -133,7 +133,7 @@ def _evaluate(self, v): idx = ind + xp.arange(self.k_max[0], dtype=int) - pf[:, 0, :] = v[:, xp.array(self._zero_angular_inds)] @ self.radial_norm[idx] + pf[:, 0, :] = v[:, xp.asarray(self._zero_angular_inds)] @ self.radial_norm[idx] ind = ind + idx.size ind_pos = ind @@ -141,7 +141,7 @@ def _evaluate(self, v): for ell in range(1, self.ell_max + 1): idx = ind + xp.arange(self.k_max[ell], dtype=int) idx_pos = ind_pos + xp.arange(self.k_max[ell], dtype=int) - idx_neg = idx_pos + xp.array(self.k_max[ell]) + idx_neg = idx_pos + xp.asarray(self.k_max[ell]) v_ell = (v[:, idx_pos] - 1j * v[:, idx_neg]) / 2.0 @@ -157,7 +157,7 @@ def _evaluate(self, v): pf[:, 2 * n_theta - ell, :] = -pf_ell.conjugate() ind = ind + idx.size - ind_pos = ind_pos + 2 * xp.array(self.k_max[ell]) + ind_pos = ind_pos + 2 * xp.asarray(self.k_max[ell]) # 1D inverse FFT in the degree of polar angle pf = 2 * xp.pi * fft.ifft(pf, axis=1) @@ -199,7 +199,7 @@ def _evaluate_t(self, x): n_images = x.shape[0] # resamping x in a polar Fourier gird using nonuniform discrete Fourier transform - pf = nufft(xp.array(x), 2 * pi * freqs) + pf = nufft(xp.asarray(x), 2 * pi * freqs) pf = pf.reshape(n_images, n_r, n_theta) # Recover "negative" frequencies from "positive" half plane. From 717ac96afb43e5af8ec27faf03004ec430ab8689 Mon Sep 17 00:00:00 2001 From: Josh Carmichael Date: Thu, 13 Jun 2024 13:36:16 -0400 Subject: [PATCH 4/5] Remove cupy.fill culprit. un-cupy indices. --- src/aspire/basis/ffb_2d.py | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/src/aspire/basis/ffb_2d.py b/src/aspire/basis/ffb_2d.py index f44e959149..116e009757 100644 --- a/src/aspire/basis/ffb_2d.py +++ b/src/aspire/basis/ffb_2d.py @@ -141,7 +141,7 @@ def _evaluate(self, v): for ell in range(1, self.ell_max + 1): idx = ind + xp.arange(self.k_max[ell], dtype=int) idx_pos = ind_pos + xp.arange(self.k_max[ell], dtype=int) - idx_neg = idx_pos + xp.asarray(self.k_max[ell]) + idx_neg = idx_pos + self.k_max[ell] v_ell = (v[:, idx_pos] - 1j * v[:, idx_neg]) / 2.0 @@ -157,7 +157,7 @@ def _evaluate(self, v): pf[:, 2 * n_theta - ell, :] = -pf_ell.conjugate() ind = ind + idx.size - ind_pos = ind_pos + 2 * xp.asarray(self.k_max[ell]) + ind_pos = ind_pos + 2 * self.k_max[ell] # 1D inverse FFT in the degree of polar angle pf = 2 * xp.pi * fft.ifft(pf, axis=1) From e49bf5dbb7d6d51a804669134549d0d0bce3bf4c Mon Sep 17 00:00:00 2001 From: Josh Carmichael Date: Thu, 13 Jun 2024 14:10:48 -0400 Subject: [PATCH 5/5] cupy.fill culprit in fle_2d. sparse indices. --- src/aspire/basis/fle_2d_utils.py | 12 ++++++++---- 1 file changed, 8 insertions(+), 4 deletions(-) diff --git a/src/aspire/basis/fle_2d_utils.py b/src/aspire/basis/fle_2d_utils.py index 23d1441a68..33f237165e 100644 --- a/src/aspire/basis/fle_2d_utils.py +++ b/src/aspire/basis/fle_2d_utils.py @@ -44,9 +44,9 @@ def precomp_transform_complex_to_real(ells): """ count = len(ells) num_nonzero = np.sum(ells == 0) + 2 * np.sum(ells != 0) - idx = xp.zeros(num_nonzero, dtype=int) - jdx = xp.zeros(num_nonzero, dtype=int) - vals = xp.zeros(num_nonzero, dtype=np.complex128) + idx = np.zeros(num_nonzero, dtype=int) + jdx = np.zeros(num_nonzero, dtype=int) + vals = np.zeros(num_nonzero, dtype=np.complex128) k = 0 for i in range(count): @@ -86,7 +86,11 @@ def precomp_transform_complex_to_real(ells): jdx[k] = i + 1 k = k + 1 - A = sparse.csr_matrix((vals, (idx, jdx)), shape=(count, count), dtype=np.complex128) + A = sparse.csr_matrix( + (xp.asarray(vals), (xp.asarray(idx), xp.asarray(jdx))), + shape=(count, count), + dtype=np.complex128, + ) return A.conjugate()