Source code for seemps.tools

from __future__ import annotations
from math import cos, sin, sqrt
import numpy as np


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

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


def take_from_list(O, i):
    if isinstance(O, list):
        return O[i]
    else:
        return O


class Logger:
    active: bool = False

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

    def __enter__(self) -> Logger:
        return self

    def __exit__(self, exc_type, exc_value, traceback):
        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

    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:
        return self

    def __call__(self, *args, **kwdargs):
        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):
        self.close()

    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, debug_level=1):
    """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, M=None):
    """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):
    """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):
    """Bosonic annihilation operator for a Hilbert space with occupations 0 to `d-1`."""
    return np.diag(sqrt(np.arange(1, d)), 1).astype(complex)