def adjoint(self, inputs, outputs):
"""The adjoint operator.
Reads from inputs and writes to outputs.
"""
self.init_kernel()
if self.implementation == Impl['halide'] and \
(len(self.shape) == 2 or (len(self.shape) == 3 and self.dims == 2)):
# Halide implementation
tmpin = np.asfortranarray(inputs[0].astype(np.float32))
Halide('At_conv.cpp').At_conv(tmpin, self.kernel, self.tmpout) # Call
np.copyto(outputs[0], self.tmpout)
else:
# Default numpy using FFT
U = fftd(inputs[0], self.dims)
U *= self.adjoint_kernel
np.copyto(outputs[0], ifftd(U, self.dims).real)
def adjoint(self, inputs, outputs):
"""The adjoint operator.
Reads from inputs and writes to outputs.
"""
self.init_kernel()
if self.implementation == Impl['halide'] and \
(len(self.shape) == 2 or (len(self.shape) == 3 and self.dims == 2)):
# Halide implementation
tmpin = np.asfortranarray(inputs[0].astype(np.float32))
Halide('At_conv.cpp').At_conv(tmpin, self.kernel, self.tmpout) # Call
np.copyto(outputs[0], self.tmpout)
else:
# Default numpy using FFT
U = fftd(inputs[0], self.dims)
U *= self.adjoint_kernel
np.copyto(outputs[0], ifftd(U, self.dims).real)
def forward(self, inputs, outputs):
"""The forward operator.
Reads from inputs and writes to outputs.
"""
self.init_kernel()
if self.implementation == Impl['halide'] and \
(len(self.shape) == 2 or (len(self.shape) == 3 and self.dims == 2)):
# Halide implementation
Halide('A_conv').A_conv(inputs[0], self.kernel,
self.tmpout) # Call
np.copyto(outputs[0], self.tmpout)
else:
# Default numpy using FFT
X = fftd(inputs[0], self.dims)
X *= self.forward_kernel
np.copyto(outputs[0], ifftd(X, self.dims).real)
def solve(self, b, rho=None, v=None, lin_solver="lsqr", *args, **kwargs):
# KtK Operator is diagonal
if self.diag is not None:
Ktb = np.zeros(self.K.input_size)
self.K.adjoint(b, Ktb)
if rho is None:
Ktb /= self.diag
else:
Ktb += (rho / 2.) * v
Ktb /= (self.diag + rho / 2.)
return Ktb
# KtK operator is diagonal in frequency domain.
elif self.freq_diag is not None:
Ktb = np.zeros(self.K.input_size)
self.K.adjoint(b, Ktb)
# Frequency inversion
if self.implementation == Impl['halide'] and \
(len(self.freq_shape) == 2 or
(len(self.freq_shape) == 2 and self.freq_dims == 2)):
Halide('fft2_r2c.cpp').fft2_r2c(np.asfortranarray(np.reshape(
Ktb.astype(np.float32), self.freq_shape)), 0, 0, self.ftmp_halide)
Ktb = 1j * self.ftmp_halide[..., 1]
Ktb += self.ftmp_halide[..., 0]
if rho is None:
Ktb /= self.freq_diag
else:
Halide('fft2_r2c.cpp').fft2_r2c(np.asfortranarray(np.reshape(
v.astype(np.float32), self.freq_shape)), 0, 0, self.ftmp_halide)
vhat = self.ftmp_halide[..., 0] + 1j * self.ftmp_halide[..., 1]
Ktb *= 1.0 / rho
Ktb += vhat
Ktb /= (1.0 / rho * self.freq_diag + 1.0)
# Do inverse tranform
Ktb = np.asfortranarray(np.stack((Ktb.real, Ktb.imag), axis=-1))
Halide('ifft2_c2r.cpp').ifft2_c2r(Ktb, self.ftmp_halide_out)
return self.ftmp_halide_out.ravel()
else:
# General frequency inversion
Ktb = fftd(np.reshape(Ktb, self.freq_shape), self.freq_dims)
if rho is None:
Ktb /= self.freq_diag
else:
Ktb *= 2.0 / rho
Ktb += fftd(np.reshape(v, self.freq_shape), self.freq_dims)
Ktb /= (2.0 / rho * self.freq_diag + 1.0)
return (ifftd(Ktb, self.freq_dims).real).ravel()
elif lin_solver == "lsqr":
return self.solve_lsqr(b, rho, v, *args, **kwargs)
elif lin_solver == "cg":
return self.solve_cg(b, rho, v, *args, **kwargs)
else:
raise Exception("Unknown least squares solver.")
def solve(self, b, rho=None, v=None, lin_solver="lsqr", *args, **kwargs):
# KtK Operator is diagonal
if self.diag is not None:
Ktb = np.zeros(self.K.input_size)
self.K.adjoint(b, Ktb)
if rho is None:
Ktb /= self.diag
else:
Ktb += (rho / 2.) * v
Ktb /= (self.diag + rho / 2.)
return Ktb
# KtK operator is diagonal in frequency domain.
elif self.freq_diag is not None:
Ktb = np.zeros(self.K.input_size)
self.K.adjoint(b, Ktb)
# Frequency inversion
if self.implementation == Impl['halide'] and \
(len(self.freq_shape) == 2 or
(len(self.freq_shape) == 2 and self.freq_dims == 2)):
Halide('fft2_r2c.cpp').fft2_r2c(
np.asfortranarray(
np.reshape(Ktb.astype(np.float32), self.freq_shape)),
0, 0, self.ftmp_halide)
Ktb = 1j * self.ftmp_halide[..., 1]
Ktb += self.ftmp_halide[..., 0]
if rho is None:
Ktb /= self.freq_diag
else:
Halide('fft2_r2c.cpp').fft2_r2c(
np.asfortranarray(
np.reshape(v.astype(np.float32), self.freq_shape)),
0, 0, self.ftmp_halide)
vhat = self.ftmp_halide[...,
0] + 1j * self.ftmp_halide[..., 1]
Ktb *= 1.0 / rho
Ktb += vhat
Ktb /= (1.0 / rho * self.freq_diag + 1.0)
# Do inverse tranform
Ktb = np.asfortranarray(np.stack((Ktb.real, Ktb.imag),
axis=-1))
Halide('ifft2_c2r.cpp').ifft2_c2r(Ktb, self.ftmp_halide_out)
return self.ftmp_halide_out.ravel()
else:
# General frequency inversion
Ktb = fftd(np.reshape(Ktb, self.freq_shape), self.freq_dims)
if rho is None:
Ktb /= self.freq_diag
else:
Ktb *= 2.0 / rho
Ktb += fftd(np.reshape(v, self.freq_shape), self.freq_dims)
Ktb /= (2.0 / rho * self.freq_diag + 1.0)
return (ifftd(Ktb, self.freq_dims).real).ravel()
elif lin_solver == "lsqr":
return self.solve_lsqr(b, rho, v, *args, **kwargs)
elif lin_solver == "cg":
return self.solve_cg(b, rho, v, *args, **kwargs)
else:
raise Exception("Unknown least squares solver.")
def solve(self, b, rho=None, v=None, lin_solver="lsqr", *args, **kwargs):
# KtK Operator is diagonal
if self.diag is not None:
Ktb = np.empty(self.K.input_size, dtype=b.dtype)
self.K.adjoint(b, Ktb)
if rho is None:
Ktb /= self.diag
else:
ne.evaluate(
'(Ktb + v * half_rho) / (d + half_rho)',
{
'Ktb': Ktb,
'half_rho': rho * 0.5,
'v': np.zeros(Ktb.shape) if v is None else v,
'd': self.diag,
},
out=Ktb,
casting='unsafe',
)
return Ktb
# KtK operator is diagonal in frequency domain.
elif self.freq_diag is not None:
Ktb = np.empty(self.K.input_size, dtype=np.float32, order='F')
self.K.adjoint(b, Ktb)
# Frequency inversion
if self.implementation == Impl['halide'] and \
(len(self.freq_shape) == 2 or
(len(self.freq_shape) == 2 and self.freq_dims == 2)):
Halide('fft2_r2c').fft2_r2c(Ktb.reshape(self.freq_shape), 0, 0,
self.ftmp_halide)
if rho is None:
ne.evaluate('F_Ktb / d', {
'F_Ktb': self.ftmp_halide,
'd': self.freq_diag,
},
out=self.ftmp_halide,
casting='unsafe')
else:
F_Ktb = self.ftmp_halide.copy()
Halide('fft2_r2c').fft2_r2c(np.reshape(v, self.freq_shape),
0, 0, self.ftmp_halide)
ne.evaluate('(F_Ktb / rho + x) / (d / rho + 1.0)', {
'F_Ktb': F_Ktb,
'x': self.ftmp_halide,
'rho': rho,
'd': self.freq_diag,
},
out=self.ftmp_halide,
casting='unsafe')
# Do inverse tranform
Halide('ifft2_c2r').ifft2_c2r(self.ftmp_halide,
self.ftmp_halide_out)
return self.ftmp_halide_out.ravel()
else:
# General frequency inversion
Ktb = fftd(np.reshape(Ktb, self.freq_shape), self.freq_dims)
if rho is None:
Ktb /= self.freq_diag
else:
Ktb *= 2.0 / rho
Ktb += fftd(np.reshape(v, self.freq_shape), self.freq_dims)
Ktb /= (2.0 / rho * self.freq_diag + 1.0)
return (ifftd(Ktb, self.freq_dims).real).ravel()
elif lin_solver == "lsqr":
return self.solve_lsqr(b, rho, v, *args, **kwargs)
elif lin_solver == "cg":
return self.solve_cg(b, rho, v, *args, **kwargs)
else:
raise Exception("Unknown least squares solver.")