`FLEBasis2D` for images of odd resolution

[chris-langfield]

The FLEBasis2D code (#693), in its original port from the fle_2d package, used a different grid convention for odd resolutions than ASPIRE's FBBasis2D code. This became apparent when calling evaluate() on one-hot stacks of coefficients resulted in slightly different basis functions between the two codes for odd resolutions. See #738.

@j-c-c originally pointed out that we could substitute ASPIRE's grid_2d to generate the correct grid to match FB, which is here:

ASPIRE-Python/src/aspire/basis/fle_2d.py

Lines 206 to 212 in fb5e445

	if self.match_fb:
	# creates correct odd-resolution grid
	# matching other FB classes
	grid = grid_2d(self.nres, dtype=self.dtype)
	self.xs = grid["x"]
	self.ys = grid["y"]
	self.rs = grid["r"]

This grid fixes the slow FLE matrix multiplication method of evaluating coefficients, as the grid generated here is used in the construction of the dense matrix:

ASPIRE-Python/src/aspire/basis/fle_2d.py

Lines 638 to 657 in fb5e445

	def create_dense_matrix(self):
	"""
	Directly computes the transformation matrix from Cartesian coordinates to
	FLE coordinates without any shortcuts.
	:return: A NumPy array of size `(self.nres**2, self.count)` containing the matrix
	entries.
	"""
	# See Eqns. 3 and 4, Section 1.2
	ts = np.arctan2(self.ys, self.xs)
	B = np.zeros((self.nres, self.nres, self.count), dtype=np.complex128)
	for i in range(self.count):
	B[:, :, i] = self.basis_functions[i](self.rs, ts) * self.h
	B = B.reshape(self.nres**2, self.count)
	B = transform_complex_to_real(np.conj(B), self.ells)
	B = B.reshape(self.nres**2, self.count)
	if self.match_fb:
	B[:, self.flip_sign_indices] *= -1.0
	B = B[:, self.fb_compat_indices]
	return B

Unfortunately, this change does not fix the fast evaluate and evaluate_t results. These use a different grid, generated in FLEBasis2D._compute_nufft_gridpoints():

ASPIRE-Python/src/aspire/basis/fle_2d.py

Lines 261 to 285 in fb5e445

	# create gridpoints
	nodes = 1 - (2 * np.arange(self.num_radial_nodes, dtype=self.dtype) + 1) / (
	2 * self.num_radial_nodes
	)
	nodes = (np.cos(np.pi * nodes) + 1) / 2
	nodes = (
	self.greatest_lambda - self.smallest_lambda
	) * nodes + self.smallest_lambda
	nodes = nodes.reshape(self.num_radial_nodes, 1)
	radius = self.nres // 2
	h = 1 / radius
	phi = (
	2
	* np.pi
	* np.arange(self.num_angular_nodes // 2, dtype=self.dtype)
	/ self.num_angular_nodes
	)
	x = np.cos(phi).reshape(1, self.num_angular_nodes // 2)
	y = np.sin(phi).reshape(1, self.num_angular_nodes // 2)
	x = x * nodes * h
	y = y * nodes * h
	self.grid_x = x.flatten()
	self.grid_y = y.flatten()

These are used in step1_t of evaluate_t and step1 of evaluate:

step1_t()

ASPIRE-Python/src/aspire/basis/fle_2d.py

Lines 516 to 523 in fb5e445

	z = np.zeros(
	(num_img, self.num_radial_nodes, self.num_angular_nodes),
	dtype=complex_type(self.dtype),
	)
	_z = (
	nufft(im, np.stack((self.grid_x, self.grid_y)), epsilon=self.epsilon)
	* self.h**2
	)

step1()

ASPIRE-Python/src/aspire/basis/fle_2d.py

Lines 624 to 630 in fb5e445

	z = z[:, :, : self.num_angular_nodes // 2].reshape(num_img, -1)
	im = anufft(
	z.astype(complex_type(self.dtype)),
	np.stack((self.grid_x, self.grid_y)),
	(self.nres, self.nres),
	epsilon=self.epsilon,
	)

---

Now that the FLE code orders the functions the same way by default, using the notebook in #738 can provide a way to visually compare the output of FBBasis2D and the fast FLEBasis2D methods. The knobs and levers that seem to adjust the exact distribution of energy over the grid seem to be the variables in the FLEBasis2D._compute_nufft_gridpoints() method. Picking different numbers of radial and angular points, or adjusting the normalization of the gridpoints, has the effect of dilating or contracting the evaluated functions, or shifting their centers.

I think there is some combination of adjustments to the nufft grid that will achieve the same result as the grid fix for the slow matrix multiplication method.