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)