Source code for seemps.tools

from __future__ import annotations
from math import cos, sin, sqrt
import numpy as np
import scipy.sparse as sp
from typing import Any
from .typing import DenseOperator, Operator


class InvalidOperation(TypeError):
    """Exception for operations with invalid or non-matching arguments."""

    def __init__(self, op: str, *args: Any):
        super().__init__(
            f"Invalid operation {op} between arguments of types {(type(x) for x in args)}"
        )


class Logger:
    active: bool = False

    def __call__(self, *args: Any, **kwdargs: Any):
        pass

    def __enter__(self) -> Logger:
        return self

    def __exit__(self, exc_type, exc_value, traceback):  # pyright: ignore[reportMissingParameterType]
        pass

    def __bool__(self) -> bool:
        return False

    def close(self):
        pass


DEBUG = 0
PREFIX = ""
NO_LOGGER = Logger()


class VerboseLogger(Logger):
    old_prefix: str
    level: int
    active: bool

    def __init__(self, level: int):
        global PREFIX
        self.old_prefix = PREFIX
        self.level = level
        if level <= DEBUG:
            self.active = True
            PREFIX = PREFIX + " "
        else:
            self.active = False

    def __bool__(self) -> bool:
        return self.active

    def __enter__(self) -> Logger:
        super().__enter__()
        return self

    def __call__(self, *args: Any, **kwdargs: Any):
        if self.active:
            txt = " ".join([str(a) for a in args])
            txt = " ".join([PREFIX + a for a in txt.split("\n")])
            print(txt, **kwdargs)

    def __exit__(self, exc_type, exc_value, traceback):  # pyright: ignore[reportMissingParameterType]
        self.close()
        super().__exit__(exc_type, exc_value, traceback)

    def close(self):
        global PREFIX
        PREFIX = self.old_prefix


def make_logger(level: int = 1) -> Logger:
    """Create an object that logs debug information. This object has a property
    `active` that determines whether logging is working. It also has a `__call__`
    method that allows invoking the object with the information to log, working
    as if it were a `print` statement."""
    if level > DEBUG:
        return NO_LOGGER
    return VerboseLogger(level)


# TODO: Find a faster way to do logs. Currently `log` always gets called
# We should find a way to replace calls to log in the code with an if-statement
# that checks `DEBUG`
def log(*args: Any, debug_level: int = 1) -> None:
    """Optionally log informative messages to the console.

    Logging is only active when :var:`~seemps.tools.DEBUG` is True or an
    integer above or equal to the given `debug_level`.

    Parameters
    ----------
    *args : str
        Strings to be output
    debug_level : int, default = 1
        Level of messages to log
    """
    if DEBUG and (DEBUG is True or DEBUG >= debug_level):
        print(*args)


def random_isometry(N: int, M: int | None = None) -> DenseOperator:
    """Returns a random isometry with size `(M, N)`.

    Parameters
    ----------
    N, M : int
        Size of the isometry, with `N` defaulting to `M`.

    Returns
    -------
    Operator
        A dense matrix for the isometry.
    """
    if M is None:
        M = N
    U = np.random.rand(N, M)
    U, _, V = np.linalg.svd(U, full_matrices=False)
    if M <= N:
        return U
    else:
        return V


σx = np.array([[0.0, 1.0], [1.0, 0.0]])
σz = np.array([[1.0, 0.0], [0.0, -1.0]])
σy = -1j * σz @ σx


def random_Pauli():
    """Random rotation generated by Pauli matrices."""
    r = np.random.rand(2)
    θ = (2 * r[0] - 1) * np.pi
    ϕ = r[1] * np.pi
    return cos(ϕ) * (cos(θ) * σx + sin(θ) * σy) + sin(ϕ) * σz




def creation(d: int) -> DenseOperator:
    """Bosonic creation operator for a Hilbert space with occupations 0 to `d-1`."""
    return np.diag(sqrt(np.arange(1, d)), -1).astype(complex)





def annihilation(d: int) -> DenseOperator:
    """Bosonic annihilation operator for a Hilbert space with occupations 0 to `d-1`."""
    return np.diag(sqrt(np.arange(1, d)), 1).astype(complex)



def mkron(A: Operator, *other_operators: Operator) -> Operator:
    for B in other_operators:
        A = sp.kron(A, B)
    return A