CVTH3PE/app/solver/_wrap/__init__.py

"""
See also:
    https://github.com/cvxopt/cvxopt/blob/master/src/C/base.c
"""

from typing import (
    Any,
    BinaryIO,
    Generic,
    Literal,
    Optional,
    Protocol,
    Sequence,
    Tuple,
    TypeVar,
    Union,
    overload,
)

import numpy as np
from cvxopt import matrix as cvxopt_matrix
from cvxopt import sparse as cvxopt_sparse
from cvxopt import spmatrix as cvxopt_spmatrix
from numpy.typing import NDArray
from typing_extensions import Self, TypeAlias

Typecode: TypeAlias = Literal["i", "d", "z"]
# Integer sparse matrices are not implemented.
SparseTypecode: TypeAlias = Literal["d", "z"]
DenseT = TypeVar("DenseT", int, float, complex)
SparseT = TypeVar("SparseT", float, complex)
IndexType = Union[int, slice, Sequence[int], "Matrix[int]"]


class Matrix(Generic[DenseT], Protocol):
    """
    cvxopt.matrix interface
    """

    @property
    def size(self) -> Tuple[int, int]: ...

    @property
    def typecode(self) -> Typecode: ...

    def __mul__(self, other): ...
    def __add__(self, other): ...
    def __sub__(self, other): ...
    def __truediv__(self, other): ...
    def __mod__(self, other): ...
    def __len__(self) -> int: ...

    def transpose(self) -> Self: ...
    def ctrans(self) -> Self: ...
    def real(self) -> "Matrix[float]": ...
    def imag(self) -> "Matrix[float]": ...

    def tofile(self, f: BinaryIO) -> None: ...
    def fromfile(self, f: BinaryIO) -> None: ...

    def __getitem__(
        self, index: Union[IndexType, Tuple[IndexType, IndexType]]
    ) -> Union[DenseT, Self]: ...
    def __setitem__(
        self,
        index: Union[IndexType, Tuple[IndexType, IndexType]],
        value: Union[DenseT, "Matrix[Any]"],
    ) -> None: ...


@overload
def matrix(
    data: Any, size: Optional[Tuple[int, int]] = None, tc: Typecode = "d"
) -> Matrix[float]: ...


@overload
def matrix(
    data: Any, size: Optional[Tuple[int, int]] = None, tc: Typecode = "i"
) -> Matrix[int]: ...


@overload
def matrix(
    data: Any, size: Optional[Tuple[int, int]] = None, tc: Typecode = "z"
) -> Matrix[complex]: ...


def matrix(data: Any, size: Optional[Tuple[int, int]] = None, tc: Typecode = "d"):
    if size is None:
        return cvxopt_matrix(data, tc=tc)
    return cvxopt_matrix(data, size=size, tc=tc)


class SparseMatrix(Generic[SparseT], Protocol):
    """
    cvxopt.spmatrix interface
    """

    @property
    def size(self) -> Tuple[int, int]: ...

    @property
    def typecode(self) -> Typecode: ...

    @property
    def V(self) -> "Matrix[SparseT]": ...

    @property
    def I(self) -> "Matrix[int]": ...

    @property
    def J(self) -> "Matrix[int]": ...

    @property
    def CCS(self) -> "Matrix[int]": ...

    def __mul__(self, other): ...
    def __add__(self, other): ...
    def __sub__(self, other): ...
    def __truediv__(self, other): ...
    def __mod__(self, other): ...
    def __len__(self) -> int: ...

    def transpose(self) -> Self: ...
    def ctrans(self) -> Self: ...
    def real(self) -> "Matrix[float]": ...
    def imag(self) -> "Matrix[float]": ...

    def tofile(self, f: BinaryIO) -> None: ...
    def fromfile(self, f: BinaryIO) -> None: ...

    def __getitem__(
        self, index: Union[IndexType, Tuple[IndexType, IndexType]]
    ) -> Union[DenseT, Self]: ...
    def __setitem__(
        self,
        index: Union[IndexType, Tuple[IndexType, IndexType]],
        value: Union[DenseT, "Matrix[Any]"],
    ) -> None: ...


@overload
def spmatrix(
    x: Union[Sequence[float], float, Matrix[float], NDArray[np.floating[Any]]],
    I: Union[Sequence[int], NDArray[np.int_]],
    J: Union[Sequence[int], NDArray[np.int_]],
    size: Optional[Tuple[int, int]] = None,
    tc: SparseTypecode = "d",
) -> SparseMatrix[float]: ...


@overload
def spmatrix(
    x: Union[
        Sequence[complex], complex, Matrix[complex], NDArray[np.complexfloating[Any]]
    ],
    I: Union[Sequence[int], NDArray[np.int_]],
    J: Union[Sequence[int], NDArray[np.int_]],
    size: Optional[Tuple[int, int]] = None,
    tc: SparseTypecode = "z",
) -> SparseMatrix[complex]: ...


def spmatrix(
    x: Any,
    I: Any,
    J: Any,
    size: Optional[Tuple[int, int]] = None,
    tc: SparseTypecode = "d",
):
    if size is None:
        return cvxopt_spmatrix(x, I, J, tc=tc)
    return cvxopt_spmatrix(x, I, J, size=size, tc=tc)


@overload
def sparse(x: Any, tc: SparseTypecode = "d") -> SparseMatrix[float]: ...
@overload
def sparse(x: Any, tc: SparseTypecode = "z") -> SparseMatrix[complex]: ...


def sparse(x: Any, tc: SparseTypecode = "d"):
    return cvxopt_sparse(x, tc=tc)