furax

@register_dataclass_with_keys
@dataclass(frozen=True)
@dataclass_transform(frozen_default=True, field_specifiers=(field,))
class AbstractLinearOperator(ABC):
    """Base class for linear operators."""

    # Class-level tags (set by decorators)
    class_tags: ClassVar[OperatorTag] = OperatorTag.NONE

    in_structure: PyTree[jax.ShapeDtypeStruct] = field(
        kw_only=True, metadata={'static': True}, default=None
    )

    def __init_subclass__(cls, **kwargs: Any) -> None:
        dataclass(frozen=True)(cls)
        register_dataclass_with_keys(cls)

    def __post_init__(self) -> None:
        if self.in_structure is None:
            raise ValueError('The input structure of the operator is not defined.')

    @property
    def tags(self) -> OperatorTag:
        """Get the tags for this operator instance."""
        return self.class_tags

    # Operator properties with default False values
    @property
    def is_square(self) -> bool:
        return bool(self.tags & OperatorTag.SQUARE)

    @property
    def is_symmetric(self) -> bool:
        return bool(self.tags & OperatorTag.SYMMETRIC)

    @property
    def is_orthogonal(self) -> bool:
        return bool(self.tags & OperatorTag.ORTHOGONAL)

    @property
    def is_diagonal(self) -> bool:
        return bool(self.tags & OperatorTag.DIAGONAL)

    @property
    def is_tridiagonal(self) -> bool:
        return bool(self.tags & OperatorTag.TRIDIAGONAL)

    @property
    def is_lower_triangular(self) -> bool:
        return bool(self.tags & OperatorTag.LOWER_TRIANGULAR)

    @property
    def is_upper_triangular(self) -> bool:
        return bool(self.tags & OperatorTag.UPPER_TRIANGULAR)

    @property
    def is_positive_semidefinite(self) -> bool:
        return bool(self.tags & OperatorTag.POSITIVE_SEMIDEFINITE)

    @property
    def is_negative_semidefinite(self) -> bool:
        return bool(self.tags & OperatorTag.NEGATIVE_SEMIDEFINITE)

    @property
    def is_idempotent(self) -> bool:
        return bool(self.tags & OperatorTag.IDEMPOTENT)

    @overload
    def __call__(
        self, *, solver: lx.AbstractLinearSolver, **keywords: Any
    ) -> 'AbstractLinearOperator': ...

    @overload
    def __call__(self, x: PyTree[jax.ShapeDtypeStruct]) -> PyTree[jax.ShapeDtypeStruct]: ...

    def __call__(
        self, x: PyTree[jax.ShapeDtypeStruct] | None = None, **keywords: Any
    ) -> 'AbstractLinearOperator | PyTree[jax.ShapeDtypeStruct]':
        if keywords:
            raise TypeError('No keywords is allowed in AbstractLinearOperator __call__ method')
        if isinstance(x, AbstractLinearOperator):
            raise ValueError("Use '@' to compose operators")
        return self.mv(x)

    def __matmul__(self, other: Any) -> 'AbstractLinearOperator':
        if not isinstance(other, AbstractLinearOperator):
            return NotImplemented
        if not structure_equal(self.in_structure, other.out_structure):
            msg = (
                f'Incompatible linear operator structures: '
                f'self.in_structure={self.in_structure}, '
                f'other.out_structure={other.out_structure}'
            )
            raise ValueError(msg)
        if isinstance(other, CompositionOperator):
            return NotImplemented
        if isinstance(other, AbstractLazyInverseOperator):
            if other.operator is self:
                return IdentityOperator(in_structure=self.in_structure)

        return CompositionOperator([self, other])

    def __add__(self, other: Any) -> 'AbstractLinearOperator':
        if not isinstance(other, AbstractLinearOperator):
            return NotImplemented
        if not structure_equal(self.in_structure, other.in_structure):
            raise ValueError('Incompatible linear operator input structures')
        if not structure_equal(self.out_structure, other.out_structure):
            raise ValueError('Incompatible linear operator output structures')
        if isinstance(other, AdditionOperator):
            return NotImplemented

        return AdditionOperator([self, other])

    def __sub__(self, other: Any) -> 'AbstractLinearOperator':
        if not isinstance(other, AbstractLinearOperator):
            return NotImplemented
        if not structure_equal(self.in_structure, other.in_structure):
            raise ValueError('Incompatible linear operator input structures')
        if not structure_equal(self.out_structure, other.out_structure):
            raise ValueError('Incompatible linear operator output structures')

        result: AbstractLinearOperator = self + (-other)
        return result

    def __mul__(self, other: ScalarLike) -> 'AbstractLinearOperator':
        result = other * self
        assert isinstance(result, AbstractLinearOperator)  # mypy
        return result

    def __rmul__(self, other: ScalarLike) -> 'AbstractLinearOperator':
        other = jnp.asarray(other)
        if other.shape != ():
            raise ValueError('Can only multiply AbstractLinearOperators by scalars.')
        return HomothetyOperator(other, in_structure=self.out_structure) @ self

    def __truediv__(self, other: ScalarLike) -> 'AbstractLinearOperator':
        other = jnp.asarray(other)
        if other.shape != ():
            raise ValueError('Can only divide AbstractLinearOperators by scalars.')
        return HomothetyOperator(1 / other, in_structure=self.out_structure) @ self

    def __pos__(self) -> 'AbstractLinearOperator':
        return self

    def __neg__(self) -> 'AbstractLinearOperator':
        return (-1) * self

    @abstractmethod
    def mv(self, x: PyTree[Inexact[Array, ' _a']]) -> PyTree[Inexact[Array, ' _b']]: ...

    def reduce(self) -> 'AbstractLinearOperator':
        """Returns a linear operator with a reduced structure."""
        return self

    def as_matrix(self) -> Inexact[Array, 'a b']:
        """Returns the operator as a dense matrix.

        Input and output PyTrees are flattened and concatenated.
        """
        in_pytree = zeros_like(self.in_structure)
        in_leaves_ref, in_treedef = jax.tree.flatten(in_pytree)

        matrix = jnp.empty((self.out_size, self.in_size), dtype=self.out_promoted_dtype)
        jcounter = 0

        for ileaf, leaf in enumerate(in_leaves_ref):

            def body(index, carry):  # type: ignore[no-untyped-def]
                matrix, jcounter = carry
                zeros = in_leaves_ref.copy()
                zeros[ileaf] = leaf.ravel().at[index].set(1).reshape(leaf.shape)  # noqa: B023 (`body` consumed immediately, no deferred call)
                in_pytree = jax.tree.unflatten(in_treedef, zeros)
                out_pytree = self.mv(in_pytree)
                out_leaves = [leaf.ravel() for leaf in jax.tree.leaves(out_pytree)]
                matrix = matrix.at[:, jcounter].set(jnp.concatenate(out_leaves))
                jcounter += 1
                return matrix, jcounter

            matrix, jcounter = jax.lax.fori_loop(0, leaf.size, body, (matrix, jcounter))

        return matrix

    def transpose(self) -> 'AbstractLinearOperator':
        return TransposeOperator(self)

    @property
    def T(self) -> 'AbstractLinearOperator':
        return self.transpose()

    def inverse(self) -> 'AbstractLinearOperator':
        return InverseOperator(self)

    @property
    def I(self) -> 'AbstractLinearOperator':  # noqa: E743
        return self.inverse()

    @property
    def out_structure(self) -> PyTree[jax.ShapeDtypeStruct]:
        return jax.eval_shape(self.mv, self.in_structure)

    @property
    def in_size(self) -> int:
        """The number of elements in the input PyTree."""
        return sum(_.size for _ in jax.tree.leaves(self.in_structure))

    @property
    def out_size(self) -> int:
        """The number of elements in the output PyTree."""
        return sum(_.size for _ in jax.tree.leaves(self.out_structure))

    @property
    def in_promoted_dtype(self) -> DType[Any]:
        """Returns the promoted data type of the operator's input leaves."""
        leaves = jax.tree.leaves(self.in_structure)
        return jnp.result_type(*leaves)

    @property
    def out_promoted_dtype(self) -> DType[Any]:
        """Returns the promoted data type of the operator's output leaves."""
        leaves = jax.tree.leaves(self.out_structure)
        return jnp.result_type(*leaves)

@orthogonal
@diagonal
@positive_semidefinite
@idempotent
class IdentityOperator(AbstractLinearOperator):
    """Operator that returns its input unchanged: I(x) = x.

    The identity operator is diagonal, orthogonal and positive semi-definite.
    Its transpose and inverse are itself.

    Examples:
        >>> I = IdentityOperator(in_structure=jax.ShapeDtypeStruct((3,), jnp.float32))
        >>> x = jnp.array([1.0, 2.0, 3.0])
        >>> I(x)
        Array([1., 2., 3.], dtype=float32)
    """

    def __matmul__(self, other: Any) -> AbstractLinearOperator:
        if not isinstance(other, AbstractLinearOperator):
            return NotImplemented
        return other

    def mv(self, x: PyTree[Inexact[Array, '...']]) -> PyTree[Inexact[Array, '...']]:
        return x

    def as_matrix(self) -> Inexact[Array, 'a b']:
        return jnp.identity(self.in_size, dtype=self.in_promoted_dtype)

@diagonal
class HomothetyOperator(AbstractLinearOperator):
    """Operator that multiplies its input by a scalar: H(x) = k * x.

    The homothety operator is diagonal, symmetric and positive semi-definite
    (for positive values). Two consecutive homotheties compose by multiplying
    their scalars: H(k1) @ H(k2) = H(k1 * k2).

    Attributes:
        value: The scalar multiplier.

    Examples:
        >>> H = HomothetyOperator(2.0, in_structure=jax.ShapeDtypeStruct((3,), jnp.float32))
        >>> x = jnp.array([1.0, 2.0, 3.0])
        >>> H(x)
        Array([2., 4., 6.], dtype=float32)
        >>> H.I(x)  # Inverse: divides by 2
        Array([0.5, 1. , 1.5], dtype=float32)
    """

    value: Scalar | int | float

    def __matmul__(self, other: Any) -> AbstractLinearOperator:
        if isinstance(other, HomothetyOperator):
            return HomothetyOperator(self.value * other.value, in_structure=self.in_structure)
        return super().__matmul__(other)

    def mv(self, x: PyTree[Inexact[Array, '...']]) -> PyTree[Inexact[Array, '...']]:
        return jax.tree.map(lambda leaf: self.value * leaf, x)

    def inverse(self) -> AbstractLinearOperator:
        return HomothetyOperator(1 / self.value, in_structure=self.in_structure)

    def as_matrix(self) -> Inexact[Array, 'a b']:
        return self.value * jnp.identity(self.in_size, dtype=self.out_promoted_dtype)

class MoveAxisOperator(AbstractLinearOperator):
    """Operator that moves axes of pytree leaves: y = jnp.moveaxis(x, source, destination).

    This operator is orthogonal: its transpose and inverse are the reverse axis move.

    Attributes:
        source: The source axis or axes to move.
        destination: The destination axis or axes.

    Examples:
        >>> in_structure = jax.ShapeDtypeStruct((2, 3), jnp.float32)
        >>> op = MoveAxisOperator(0, 1, in_structure=in_structure)
        >>> op(jnp.array([[1., 1, 1], [2, 2, 2]]))
        Array([[1., 2.],
               [1., 2.],
               [1., 2.]], dtype=float32)
    """

    source: tuple[int, ...] = field(metadata={'static': True})
    destination: tuple[int, ...] = field(metadata={'static': True})

    def __init__(
        self,
        source: int | Sequence[int],
        destination: int | Sequence[int],
        *,
        in_structure: PyTree[jax.ShapeDtypeStruct],
    ) -> None:
        if isinstance(source, int):
            source = (source,)
        elif not isinstance(source, tuple):
            source = tuple(source)
        if isinstance(destination, int):
            destination = (destination,)
        elif not isinstance(destination, tuple):
            destination = tuple(destination)
        object.__setattr__(self, 'source', source)
        object.__setattr__(self, 'destination', destination)
        object.__setattr__(self, 'in_structure', in_structure)

    def mv(self, x: PyTree[Array, '...']) -> PyTree[Array, '...']:
        return jax.tree.map(lambda leaf: jnp.moveaxis(leaf, self.source, self.destination), x)

    def transpose(self) -> AbstractLinearOperator:
        return MoveAxisOperator(
            source=self.destination, destination=self.source, in_structure=self.out_structure
        )

    inverse = transpose

class RavelOperator(AbstractRavelOrReshapeOperator):
    """Operator that flattens pytree leaves: y = x.ravel().

    By default, all dimensions are flattened. Use ``first_axis`` and ``last_axis``
    to flatten only a subset of contiguous axes.

    This operator is orthogonal: its transpose restores the original shape.

    Attributes:
        first_axis: The first axis to flatten (default: 0).
        last_axis: The last axis to flatten (default: -1).

    Examples:
        To flatten the leaves of a pytree:

        >>> in_structure = jax.ShapeDtypeStruct((2, 3), jnp.float32)
        >>> op = RavelOperator(in_structure=in_structure)
        >>> op.out_structure
        ShapeDtypeStruct(shape=(6,), dtype=float32)

        To flatten the first two axes of the leaves of a pytree:

        >>> import furax as fx
        >>> x = [jnp.ones((2, 2)), jnp.ones((2, 2, 8))]
        >>> op = RavelOperator(0, 1, in_structure=fx.tree.as_structure(x))
        >>> op.out_structure
        [ShapeDtypeStruct(shape=(4,), dtype=float32),
        ShapeDtypeStruct(shape=(4, 8), dtype=float32)]


        To flatten the last two axes of the leaves of a pytree:

        >>> import furax as fx
        >>> x = [jnp.ones((2, 2, 3)), jnp.ones((2, 8))]
        >>> op = RavelOperator(-2, -1, in_structure=fx.tree.as_structure(x))
        >>> op.out_structure
        [ShapeDtypeStruct(shape=(2, 6), dtype=float32),
        ShapeDtypeStruct(shape=(16,), dtype=float32)]
    """

    first_axis: int = field(default=0, metadata={'static': True})
    last_axis: int = field(default=-1, metadata={'static': True})

    def __post_init__(self) -> None:
        first_axis = self.first_axis
        last_axis = self.last_axis
        in_structure = self.in_structure

        if 0 <= last_axis < first_axis or last_axis < first_axis < 0:
            raise ValueError(
                f'the first axis ({first_axis}) to be flattened should be before the last one '
                f'({last_axis}).'
            )

        if first_axis < 0 <= last_axis or last_axis < 0 <= first_axis:
            for leaf in jax.tree.leaves(in_structure):
                first = leaf.ndim + first_axis if first_axis < 0 else first_axis
                last = leaf.ndim + last_axis if last_axis < 0 else last_axis
                if first > last:
                    raise ValueError(
                        f'there are no dimensions between {first_axis} and {last_axis} '
                        f'to be flattened in leaf of shape {leaf.shape}.'
                    )

    def mv(self, x: PyTree[Inexact[Array, ' _a']]) -> PyTree[Inexact[Array, ' _b']]:
        def func(leaf: Inexact[Array, ' _a']) -> Inexact[Array, ' _b']:
            first_axis = leaf.ndim + self.first_axis if self.first_axis < 0 else self.first_axis
            last_axis = leaf.ndim + self.last_axis if self.last_axis < 0 else self.last_axis
            if first_axis > last_axis:
                assert False, 'unreachable'  # noqa: B011
            if first_axis == last_axis:
                return leaf
            new_shape = leaf.shape[:first_axis] + (-1,) + leaf.shape[last_axis + 1 :]
            return leaf.reshape(new_shape)

        return jax.tree.map(func, x)