# Source code for pylops.basicoperators.diagonal

__all__ = ["Diagonal"]

from typing import Optional, Union

import numpy as np

from pylops import LinearOperator
from pylops.utils._internal import _value_or_sized_to_tuple
from pylops.utils.backend import get_array_module, to_cupy_conditional
from pylops.utils.decorators import reshaped
from pylops.utils.typing import DTypeLike, InputDimsLike, NDArray

[docs]class Diagonal(LinearOperator):
r"""Diagonal operator.

Applies element-wise multiplication of the input vector with the vector
diag in forward and with its complex conjugate in adjoint mode.

This operator can also broadcast; in this case the input vector is
reshaped into its dimensions dims and the element-wise multiplication
with diag is perfomed along axis. Note that the
vector diag will need to have size equal to dims[axis].

Parameters
----------
diag : :obj:numpy.ndarray
Vector to be used for element-wise multiplication.
dims : :obj:list, optional
Number of samples for each dimension
(None if only one dimension is available)
axis : :obj:int, optional

Axis along which multiplication is applied.
dtype : :obj:str, optional
Type of elements in input array.
name : :obj:str, optional

Name of operator (to be used by :func:pylops.utils.describe.describe)

Attributes
----------
shape : :obj:tuple
Operator shape
explicit : :obj:bool
Operator contains a matrix that can be solved explicitly (True) or
not (False)

Notes
-----
Element-wise multiplication between the model :math:\mathbf{x} and/or
data :math:\mathbf{y} vectors and the array :math:\mathbf{d}
can be expressed as

.. math::

y_i = d_i x_i  \quad \forall i=1,2,\ldots,N

This is equivalent to a matrix-vector multiplication with a matrix
containing the vector :math:\mathbf{d} along its main diagonal.

For real-valued diag, the Diagonal operator is self-adjoint as the
adjoint of a diagonal matrix is the diagonal matrix itself. For
complex-valued diag, the adjoint is equivalent to the element-wise
multiplication with the complex conjugate elements of diag.

"""

def __init__(
self,
diag: NDArray,
dims: Optional[Union[int, InputDimsLike]] = None,
axis: int = -1,
dtype: DTypeLike = "float64",
name: str = "D",
) -> None:
self.diag = diag
self.axis = axis
origdims = dims
dims = self.diag.shape if dims is None else _value_or_sized_to_tuple(dims)
super().__init__(dtype=np.dtype(dtype), dims=dims, dimsd=dims, name=name)

ncp = get_array_module(diag)
self.complex = True if ncp.iscomplexobj(self.diag) else False
if origdims is not None:
diagdims = np.ones_like(self.dims)
diagdims[axis] = self.dims[axis]
self.diag = self.diag.reshape(diagdims)

@reshaped
def _matvec(self, x: NDArray) -> NDArray:
if type(self.diag) is not type(x):
self.diag = to_cupy_conditional(x, self.diag)
y = self.diag * x
return y

@reshaped
def _rmatvec(self, x: NDArray) -> NDArray:
if type(self.diag) is not type(x):
self.diag = to_cupy_conditional(x, self.diag)
if self.complex:
else:
return y

def matrix(self) -> NDArray:
"""Return diagonal matrix as dense :obj:numpy.ndarray

Returns
-------
densemat : :obj:numpy.ndarray
Dense matrix.

"""
ncp = get_array_module(self.diag)
densemat = ncp.diag(self.diag.squeeze())
return densemat

def todense(self) -> NDArray:
"""Fast implementation of todense based on known structure of the
operator

Returns
-------
densemat : :obj:numpy.ndarray
Dense matrix.
"""
if self.diag.ndim == 1:
return self.matrix()
else:
dims = list(self.dims)
dims[self.axis] = 1
matrix = np.diag(np.tile(self.diag, dims).ravel())
return matrix