Source code for seemps.tools

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

SEED = 0x874665212
DEFAULT_RNG = np.random.default_rng(SEED)


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


# TODO: Document all environment variables
DEBUG = int(os.environ.get("SEEMPS_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, file: TextIO = sys.stderr, **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, file=file, **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, file=sys.stderr)


def random_isometry(
    N: int, M: int | None = None, rng: np.random.Generator = DEFAULT_RNG
) -> 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 = rng.normal(size=(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(rng: np.random.Generator = DEFAULT_RNG):
    """Random rotation generated by Pauli matrices."""
    r = rng.normal(size=2)
    θ = (2 * r[0] - 1) * np.pi
    ϕ = r[1] * np.pi
    return cos(ϕ) * (cos(θ) * σx + sin(θ) * σy) + sin(ϕ) * σz


[docs] 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)
[docs] 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