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)
class FourierKernelTestCase(TestCase):
def setUp(self):
self.kernel = FourierKernel(np.load(
os.path.join(DATA_DIR, 'fourier_kernel_16_16_16.npy')),
centered=False)
def tearDown(self):
pass
def test1(self):
toeplitz = self.kernel.toeplitz()
self.assertEqual(toeplitz.shape, (8, 8, 8, 8, 8, 8))
self.assertTrue(
np.allclose(
toeplitz[:, :, 0, 0, 0, 2],
np.array([[
-3.99237964e-04, -5.58560540e-04, -4.61112126e-04,
-1.33411959e-05, 1.80705421e-04, 5.17746266e-05,
-1.11463014e-04, -6.67081913e-05
],
[
-5.73768339e-04, -5.82092151e-04,
-3.57612298e-04, 6.28258349e-05, 2.10987753e-04,
-1.37316420e-05, -1.81071970e-04, -1.65530946e-05
],
[
-4.50893916e-04, -3.94886069e-04,
-1.31166336e-04, 1.83019380e-04, 1.86107689e-04,
-8.73301760e-05, -1.89300568e-04, 4.55726404e-05
],
[
-7.88352190e-05, -5.64255715e-05, 1.79488547e-04,
2.42688577e-04, 7.89244223e-05, -1.95445391e-04,
-1.64644473e-04, 8.46998155e-05
],
[
1.78478949e-04, 1.10234592e-04, 1.29420368e-04,
-9.66549123e-06, -1.59117772e-04,
-2.06353434e-04, 8.77526109e-05, 2.81196553e-04
],
[
-1.48982726e-05, -5.52274614e-05,
-9.01011008e-05, -1.95092929e-04,
-2.12131679e-04, -2.77411455e-05, 1.63609919e-04,
1.70907035e-04
],
[
-1.23197358e-04, -2.01485047e-04,
-2.22623392e-04, -1.35038150e-04, 2.55815121e-05,
1.63564575e-04, 7.31561740e-05, -2.58645450e-05
],
[
-1.00527395e-04, -5.97040562e-05, 6.75393894e-05,
2.12094223e-04, 2.47279182e-04, 1.78254995e-04,
-3.36604789e-05, -2.02765747e-04
]])))
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 __getattr__(self, name):
"""Lazy attributes instantiated on first-access"""
if name == 'kernel':
logger.info('Computing kernel')
kernel = self.kernel = self.compute_kernel()
return kernel
elif name == 'precond_kernel':
if self.preconditioner == 'circulant':
logger.info('Computing Preconditioner kernel')
precond_kernel = self.precond_kernel = FourierKernel(
1. / self.kernel.circularize(), centered=True)
else:
precond_kernel = self.precond_kernel = None
return precond_kernel
return super(Estimator, self).__getattr__(name)
def setUp(self):
self.kernel = FourierKernel(np.load(
os.path.join(DATA_DIR, 'fourier_kernel_16_16_16.npy')),
centered=False)