def compute_kernel(self):
_2L = 2 * self.L
kernel = np.zeros((_2L, _2L, _2L), dtype=self.as_type)
filters_f = self.src.filters.evaluate_grid(self.L)
sq_filters_f = np.array(filters_f**2, dtype=self.as_type)
for i in range(0, self.n, self.batch_size):
pts_rot = rotated_grids(self.L,
self.src.rots[:, :, i:i + self.batch_size])
weights = sq_filters_f[:, :,
self.src.filters.indices[i:i +
self.batch_size]]
weights *= self.src.amplitudes[i:i + self.batch_size]**2
if self.L % 2 == 0:
weights[0, :, :] = 0
weights[:, 0, :] = 0
pts_rot = m_reshape(pts_rot, (3, -1))
weights = m_flatten(weights)
kernel += 1 / (self.n * self.L**4) * anufft3(
weights, pts_rot, (_2L, _2L, _2L), real=True)
# Ensure symmetric kernel
kernel[0, :, :] = 0
kernel[:, 0, :] = 0
kernel[:, :, 0] = 0
logger.info('Computing non-centered Fourier Transform')
kernel = mdim_ifftshift(kernel, range(0, 3))
kernel_f = fft2(kernel, axes=(0, 1, 2))
kernel_f = np.real(kernel_f)
return FourierKernel(kernel_f, centered=False)
def compute_kernel(self):
_2L = 2 * self.L
kernel = np.zeros((_2L, _2L, _2L), dtype=self.dtype)
sq_filters_f = self.src.eval_filter_grid(self.L, power=2)
for i in range(0, self.n, self.batch_size):
_range = np.arange(i, min(self.n, i + self.batch_size), dtype=int)
pts_rot = rotated_grids(self.L, self.src.rots[_range, :, :])
weights = sq_filters_f[:, :, _range]
weights *= self.src.amplitudes[_range]**2
if self.L % 2 == 0:
weights[0, :, :] = 0
weights[:, 0, :] = 0
pts_rot = m_reshape(pts_rot, (3, -1))
weights = m_flatten(weights)
kernel += (1 / (self.n * self.L**4) *
anufft(weights, pts_rot, (_2L, _2L, _2L), real=True))
# Ensure symmetric kernel
kernel[0, :, :] = 0
kernel[:, 0, :] = 0
kernel[:, :, 0] = 0
logger.info("Computing non-centered Fourier Transform")
kernel = mdim_ifftshift(kernel, range(0, 3))
kernel_f = fft2(kernel, axes=(0, 1, 2))
kernel_f = np.real(kernel_f)
return FourierKernel(kernel_f, centered=False)
def circularize(self):
logger.info("Circularizing kernel")
kernel = np.real(ifftn(self.kernel))
kernel = mdim_fftshift(kernel)
for dim in range(self.ndim):
logger.info(f"Circularizing dimension {dim}")
kernel = self.circularize_1d(kernel, dim)
xx = fftn(mdim_ifftshift(kernel))
return xx
def compute_kernel(self):
# TODO: Most of this stuff is duplicated in MeanEstimator - move up the hierarchy?
n = self.n
L = self.L
_2L = 2 * self.L
kernel = np.zeros((_2L, _2L, _2L, _2L, _2L, _2L), dtype=self.as_type)
filters_f = self.src.filters.evaluate_grid(L)
sq_filters_f = np.array(filters_f**2, dtype=self.as_type)
for i in tqdm(range(0, n, self.batch_size)):
pts_rot = rotated_grids(L, self.src.rots[:, :,
i:i + self.batch_size])
weights = sq_filters_f[:, :,
self.src.filters.indices[i:i +
self.batch_size]]
weights *= self.src.amplitudes[i:i + self.batch_size]**2
if L % 2 == 0:
weights[0, :, :] = 0
weights[:, 0, :] = 0
# TODO: This is where this differs from MeanEstimator
pts_rot = m_reshape(pts_rot, (3, L**2, -1))
weights = m_reshape(weights, (L**2, -1))
batch_n = weights.shape[-1]
factors = np.zeros((_2L, _2L, _2L, batch_n), dtype=self.as_type)
# TODO: Numpy has got to have a functional shortcut to avoid looping like this!
for j in range(batch_n):
factors[:, :, :, j] = anufft3(weights[:, j],
pts_rot[:, :, j],
(_2L, _2L, _2L),
real=True)
factors = vol_to_vec(factors)
kernel += vecmat_to_volmat(factors @ factors.T) / (n * L**8)
# Ensure symmetric kernel
kernel[0, :, :, :, :, :] = 0
kernel[:, 0, :, :, :, :] = 0
kernel[:, :, 0, :, :, :] = 0
kernel[:, :, :, 0, :, :] = 0
kernel[:, :, :, :, 0, :] = 0
kernel[:, :, :, :, :, 0] = 0
logger.info('Computing non-centered Fourier Transform')
kernel = mdim_ifftshift(kernel, range(0, 6))
kernel_f = fftn(kernel)
# Kernel is always symmetric in spatial domain and therefore real in Fourier
kernel_f = np.real(kernel_f)
return FourierKernel(kernel_f, centered=False)
def compute_kernel(self):
# TODO: Most of this stuff is duplicated in MeanEstimator - move up the hierarchy?
n = self.n
L = self.L
_2L = 2 * self.L
kernel = np.zeros((_2L, _2L, _2L, _2L, _2L, _2L), dtype=self.dtype)
sq_filters_f = self.src.eval_filter_grid(self.L, power=2)
for i in tqdm(range(0, n, self.batch_size)):
_range = np.arange(i, min(n, i + self.batch_size))
pts_rot = rotated_grids(L, self.src.rots[_range, :, :])
weights = sq_filters_f[:, :, _range]
weights *= self.src.amplitudes[_range]**2
if L % 2 == 0:
weights[0, :, :] = 0
weights[:, 0, :] = 0
# TODO: This is where this differs from MeanEstimator
pts_rot = np.moveaxis(pts_rot, -1, 0).reshape(-1, 3, L**2)
weights = weights.T.reshape((-1, L**2))
batch_n = weights.shape[0]
factors = np.zeros((batch_n, _2L, _2L, _2L), dtype=self.dtype)
for j in range(batch_n):
factors[j] = anufft(weights[j],
pts_rot[j], (_2L, _2L, _2L),
real=True)
factors = Volume(factors).to_vec()
kernel += vecmat_to_volmat(factors.T @ factors) / (n * L**8)
# Ensure symmetric kernel
kernel[0, :, :, :, :, :] = 0
kernel[:, 0, :, :, :, :] = 0
kernel[:, :, 0, :, :, :] = 0
kernel[:, :, :, 0, :, :] = 0
kernel[:, :, :, :, 0, :] = 0
kernel[:, :, :, :, :, 0] = 0
logger.info("Computing non-centered Fourier Transform")
kernel = mdim_ifftshift(kernel, range(0, 6))
kernel_f = fftn(kernel)
# Kernel is always symmetric in spatial domain and therefore real in Fourier
kernel_f = np.real(kernel_f)
return FourierKernel(kernel_f, centered=False)