def _make_multiplier_matrix(self, conjugate=False):
# Shape [B1,...Bb, 1, 1]
multiplier_matrix = array_ops.expand_dims(
array_ops.expand_dims(self.multiplier, -1), -1)
if conjugate:
multiplier_matrix = math_ops.conj(multiplier_matrix)
return multiplier_matrix
def _diag_part(self):
# [U D V^T]_{ii} = sum_{jk} U_{ij} D_{jk} V_{ik}
# = sum_{j} U_{ij} D_{jj} V_{ij}
product = self.u * math_ops.conj(self.v)
if self.diag_update is not None:
product = product * array_ops.expand_dims(self.diag_update,
axis=-2)
return (math_ops.reduce_sum(product, axis=-1) +
self.base_operator.diag_part())
def _adjoint_diag(diag_operator):
diag = diag_operator.diag
if np.issubdtype(diag.dtype, np.complexfloating):
diag = math_ops.conj(diag)
return linear_operator_diag.LinearOperatorDiag(
diag=diag,
is_non_singular=diag_operator.is_non_singular,
is_self_adjoint=diag_operator.is_self_adjoint,
is_positive_definite=diag_operator.is_positive_definite,
is_square=True)
def _adjoint_circulant(circulant_operator):
spectrum = circulant_operator.spectrum
if np.issubdtype(spectrum.dtype, np.complexfloating):
spectrum = math_ops.conj(spectrum)
# Conjugating the spectrum is sufficient to get the adjoint.
return linear_operator_circulant.LinearOperatorCirculant(
spectrum=spectrum,
is_non_singular=circulant_operator.is_non_singular,
is_self_adjoint=circulant_operator.is_self_adjoint,
is_positive_definite=circulant_operator.is_positive_definite,
is_square=True)
def _adjoint_scaled_identity(identity_operator):
multiplier = identity_operator.multiplier
if np.issubdtype(multiplier.dtype, np.complexfloating):
multiplier = math_ops.conj(multiplier)
return linear_operator_identity.LinearOperatorScaledIdentity(
num_rows=identity_operator._num_rows, # pylint: disable=protected-access
multiplier=multiplier,
is_non_singular=identity_operator.is_non_singular,
is_self_adjoint=identity_operator.is_self_adjoint,
is_positive_definite=identity_operator.is_positive_definite,
is_square=True)
def _solve(self, rhs, adjoint=False, adjoint_arg=False):
rhs = linalg.adjoint(rhs) if adjoint_arg else rhs
spectrum = _to_complex(self.spectrum)
if adjoint:
spectrum = math_ops.conj(spectrum)
rhs, spectrum = self._broadcast_batch_dims(rhs, spectrum)
rhs_vb = self._vectorize_then_blockify(rhs)
fft_rhs_vb = self._fft(rhs_vb)
solution_vb = self._ifft(fft_rhs_vb / spectrum)
x = self._unblockify_then_matricize(solution_vb)
return _ops.cast(x, self.dtype)
def _matmul(self, x, adjoint=False, adjoint_arg=False):
x = linalg.adjoint(x) if adjoint_arg else x
# With F the matrix of a DFT, and F^{-1}, F^H the inverse and Hermitian
# transpose, one can show that F^{-1} = F^{H} is the IDFT matrix. Therefore
# matmul(x) = F^{-1} diag(spectrum) F x,
# = F^{H} diag(spectrum) F x,
# so that
# matmul(x, adjoint=True) = F^{H} diag(conj(spectrum)) F x.
spectrum = _to_complex(self.spectrum)
if adjoint:
spectrum = math_ops.conj(spectrum)
x = _ops.cast(x, spectrum.dtype)
x, spectrum = self._broadcast_batch_dims(x, spectrum)
x_vb = self._vectorize_then_blockify(x)
fft_x_vb = self._fft(x_vb)
block_vector_result = self._ifft(spectrum * fft_x_vb)
y = self._unblockify_then_matricize(block_vector_result)
return _ops.cast(y, self.dtype)
def _trace(self):
if self.is_self_adjoint:
return self.operator.trace()
return math_ops.conj(self.operator.trace())
def _determinant(self):
if self.is_self_adjoint:
return self.operator.determinant()
return math_ops.conj(self.operator.determinant())
def _solve(self, rhs, adjoint=False, adjoint_arg=False):
diag_term = math_ops.conj(self._diag) if adjoint else self._diag
rhs = linalg.adjoint(rhs) if adjoint_arg else rhs
inv_diag_mat = array_ops.expand_dims(1. / diag_term, -1)
return rhs * inv_diag_mat
def _matmul(self, x, adjoint=False, adjoint_arg=False):
diag_term = math_ops.conj(self._diag) if adjoint else self._diag
x = linalg.adjoint(x) if adjoint_arg else x
diag_mat = array_ops.expand_dims(diag_term, -1)
return diag_mat * x
def _diag_part(self):
reflection_axis = ops.convert_to_tensor(self.reflection_axis)
normalized_axis = reflection_axis / linalg.norm(
reflection_axis, axis=-1, keepdims=True)
return 1. - 2 * normalized_axis * math_ops.conj(normalized_axis)
def _solve_matmul_internal(self,
x,
solve_matmul_fn,
adjoint=False,
adjoint_arg=False):
# We heavily rely on Roth's column Lemma [1]:
# (A x B) * vec X = vec BXA^T
# where vec stacks all the columns of the matrix under each other.
# In our case, we use a variant of the lemma that is row-major
# friendly: (A x B) * vec' X = vec' AXB^T
# Where vec' reshapes a matrix into a vector. We can repeatedly apply this
# for a collection of kronecker products.
# Given that (A x B)^-1 = A^-1 x B^-1 and (A x B)^T = A^T x B^T, we can
# use the above to compute multiplications, solves with any composition of
# transposes.
output = x
if adjoint_arg:
if np.issubdtype(self.dtype, np.complexfloating):
output = math_ops.conj(output)
else:
output = linalg.transpose(output)
for o in reversed(self.operators):
# Statically compute the reshape.
if adjoint:
operator_dimension = o.range_dimension_tensor()
else:
operator_dimension = o.domain_dimension_tensor()
output_shape = _prefer_static_shape(output)
if ops.get_static_value(operator_dimension) is not None:
operator_dimension = ops.get_static_value(operator_dimension)
if tensor_shape.TensorShape(
output.shape
)[-2] is not None and tensor_shape.TensorShape(
output.shape)[-1] is not None:
dim = int(
tensor_shape.TensorShape(output.shape)[-2] *
output_shape[-1] // operator_dimension)
else:
dim = _ops.cast(output_shape[-2] * output_shape[-1] //
operator_dimension,
dtype=dtypes.int32)
output_shape = _prefer_static_concat_shape(
output_shape[:-2], [dim, operator_dimension])
output = array_ops.reshape(output, shape=output_shape)
# Conjugate because we are trying to compute A @ B^T, but
# `LinearOperator` only supports `adjoint_arg`.
if np.issubdtype(self.dtype, np.complexfloating):
output = math_ops.conj(output)
output = solve_matmul_fn(o,
output,
adjoint=adjoint,
adjoint_arg=True)
if adjoint_arg:
col_dim = _prefer_static_shape(x)[-2]
else:
col_dim = _prefer_static_shape(x)[-1]
if adjoint:
row_dim = self.domain_dimension_tensor()
else:
row_dim = self.range_dimension_tensor()
matrix_shape = [row_dim, col_dim]
output = array_ops.reshape(
output,
_prefer_static_concat_shape(
_prefer_static_shape(output)[:-2], matrix_shape))
if tensor_shape.TensorShape(x.shape).is_fully_defined():
if adjoint_arg:
column_dim = tensor_shape.TensorShape(x.shape)[-2]
else:
column_dim = tensor_shape.TensorShape(x.shape)[-1]
broadcast_batch_shape = common_shapes.broadcast_shape(
tensor_shape.TensorShape(x.shape)[:-2], self.batch_shape)
if adjoint:
matrix_dimensions = [self.domain_dimension, column_dim]
else:
matrix_dimensions = [self.range_dimension, column_dim]
tensorshape_util.set_shape(
output, broadcast_batch_shape.concatenate(matrix_dimensions))
return output
def _eigvals(self):
eigvals = self.operator.eigvals()
if not self.operator.is_self_adjoint:
eigvals = math_ops.conj(eigvals)
return eigvals
def _diag_part(self):
normalized_axis = self.reflection_axis / linalg.norm(
self.reflection_axis, axis=-1, keepdims=True)
return 1. - 2 * normalized_axis * math_ops.conj(normalized_axis)
def _matvec(self, x, adjoint=False):
diag_term = math_ops.conj(self._diag) if adjoint else self._diag
return diag_term * x
def _diag_part(self):
reflection_axis = ops.convert_to_tensor(
self.reflection_axis)
normalized_axis = nn.l2_normalize(reflection_axis, axis=-1)
return 1. - 2 * normalized_axis * math_ops.conj(normalized_axis)
def __init__(self,
spectrum,
block_depth,
input_output_dtype=dtypes.complex64,
is_non_singular=None,
is_self_adjoint=None,
is_positive_definite=None,
is_square=True,
name="LinearOperatorCirculant"):
r"""Initialize an `_BaseLinearOperatorCirculant`.
Args:
spectrum: Shape `[B1,...,Bb, N]` `Tensor`. Allowed dtypes: `float16`,
`float32`, `float64`, `complex64`, `complex128`. Type can be different
than `input_output_dtype`
block_depth: Python integer, either 1, 2, or 3. Will be 1 for circulant,
2 for block circulant, and 3 for nested block circulant.
input_output_dtype: `dtype` for input/output.
is_non_singular: Expect that this operator is non-singular.
is_self_adjoint: Expect that this operator is equal to its hermitian
transpose. If `spectrum` is real, this will always be true.
is_positive_definite: Expect that this operator is positive definite,
meaning the quadratic form `x^H A x` has positive real part for all
nonzero `x`. Note that we do not require the operator to be
self-adjoint to be positive-definite. See:
https://en.wikipedia.org/wiki/Positive-definite_matrix\
#Extension_for_non_symmetric_matrices
is_square: Expect that this operator acts like square [batch] matrices.
name: A name to prepend to all ops created by this class.
Raises:
ValueError: If `block_depth` is not an allowed value.
TypeError: If `spectrum` is not an allowed type.
"""
allowed_block_depths = [1, 2, 3]
self._name = name
if block_depth not in allowed_block_depths:
raise ValueError("Expected block_depth to be in %s. Found: %s." %
(allowed_block_depths, block_depth))
self._block_depth = block_depth
with ops.name_scope(name, values=[spectrum]):
self._spectrum = self._check_spectrum_and_return_tensor(spectrum)
# Check and auto-set hints.
if not np.issubdtype(self.spectrum.dtype, np.complexfloating):
if is_self_adjoint is False:
raise ValueError(
"A real spectrum always corresponds to a self-adjoint operator.")
is_self_adjoint = True
if is_square is False:
raise ValueError(
"A [[nested] block] circulant operator is always square.")
is_square = True
# If _ops.TensorShape(spectrum.shape) = [s0, s1, s2], and block_depth = 2,
# block_shape = [s1, s2]
s_shape = array_ops.shape(self.spectrum)
self._block_shape_tensor = s_shape[-self.block_depth:]
# Add common variants of spectrum to the graph.
self._spectrum_complex = _to_complex(self.spectrum)
self._abs_spectrum = math_ops.abs(self.spectrum)
self._conj_spectrum = math_ops.conj(self._spectrum_complex)
super(_BaseLinearOperatorCirculant, self).__init__(
dtype=dtypes.as_dtype(input_output_dtype),
graph_parents=[self.spectrum],
is_non_singular=is_non_singular,
is_self_adjoint=is_self_adjoint,
is_positive_definite=is_positive_definite,
is_square=is_square,
name=name)
