Source code for discopy.matrix

The category of matrices with the direct sum as monoidal product.


.. autosummary::
    :template: class.rst


.. admonition:: Functions

    .. autosummary::
        :template: function.rst


See also

* :class:`discopy.tensor.Tensor` is a subclass of :class:`Matrix` with the
  Kronecker product as tensor.

from __future__ import annotations

from contextlib import contextmanager
from types import ModuleType
from typing import Union, Literal as L, Callable, TYPE_CHECKING

from discopy import monoidal, config, messages
from import (
from discopy.monoidal import Whiskerable
from discopy.utils import assert_isinstance, unbiased, NamedGeneric

    import sympy

[docs]@factory class Matrix(Composable[int], Whiskerable, NamedGeneric['dtype']): """ A matrix is an ``array`` with natural numbers as ``dom`` and ``cod``. .. admonition:: Summary .. autosummary:: id then tensor zero swap transpose conjugate dagger map round copy discard merge basis repeat trace lambdify subs grad Note ---- The class ``Matrix[dtype]`` has arrays with entries in any given ``dtype``. For example: >>> Matrix[complex].id(1) Matrix[complex]([1.+0.j], dom=1, cod=1) >>> assert Matrix[complex].id(1) != Matrix[float].id(1) The default data type is determined by underlying array datastructure of the backend used. An array is initialised with ``array`` as parameter and the dtype of the ``Matrix`` object is the data type of this array. >>> import numpy as np >>> assert Matrix([1, 0], dom=1, cod=2).dtype == np.int64 >>> assert Matrix([0.5, 0.5], dom=1, cod=2).dtype == np.float64 >>> assert Matrix([0.5j], dom=1, cod=1).dtype == np.complex128 The data type needs to have the structure of a rig (riNg with no negatives) i.e. with methods ``__add__`` and ``__mul__`` as well as an ``__init__`` that can accept both ``0`` and ``1`` as input. Examples -------- >>> m = Matrix([0, 1, 1, 0], 2, 2) >>> v = Matrix([0, 1], 1, 2) >>> v >> m >> v.dagger() Matrix[int64]([0], dom=1, cod=1) >>> m + m Matrix[int64]([0, 2, 2, 0], dom=2, cod=2) >>> assert m.then(m, m, m, m) == m >> m >> m >> m >> m The monoidal product for :py:class:`.Matrix` is the direct sum: >>> x = Matrix([2, 4], 2, 1) >>> x.array array([[2], [4]]) >>> x @ x Matrix[int64]([2, 0, 4, 0, 0, 2, 0, 4], dom=4, cod=2) >>> (x @ x).array array([[2, 0], [4, 0], [0, 2], [0, 4]]) """
[docs] def cast(self, dtype: type) -> Matrix: """ Cast a matrix to a given ``dtype``. Parameters: dtype : The target datatype. Example ------- >>> assert == Matrix[bool].id() """ return type(self)[dtype](self.array, self.dom, self.cod)
def __new__(cls, array, *args, **kwargs): with backend() as np: if cls.dtype is None: array = np.array(array) # The dtype of an np.arrays is a class that contains a type # attribute that is the actual type. However, other backends # have different structures, so this is the easiest option: dtype = getattr(array.dtype, "type", array.dtype) return cls.__new__(cls[dtype], array, *args, **kwargs) return object.__new__(cls) def __init__(self, array, dom: int, cod: int): assert_isinstance(dom, int) assert_isinstance(cod, int) self.dom, self.cod = dom, cod with backend() as np: self.array = np.array(array, dtype=self.dtype).reshape((dom, cod)) def __eq__(self, other): return isinstance(other, self.factory)\ and self.dtype == other.dtype\ and (self.dom, self.cod) == (other.dom, other.cod)\ and (self.array == other.array).all()
[docs] def is_close(self, other: Matrix, rtol: float = 1.e-8, atol: float = 1.e-8 ) -> bool: """ Whether a matrix is numerically close to an ``other``. Parameters: other : The other matrix with which to check closeness. rtol: The relative tolerance parameter (see Notes). Default value for results of order unity is 1.e-5 atol : The absolute tolerance parameter (see Notes). Default value for results of order unity is 1.e-8 Notes: (taken from np.isclose documentation) For finite values, isclose uses the following equation to test whether two floating point values are equivalent. absolute(`a` - `b`) <= (`atol` + `rtol` * absolute(`b`)) Unlike the built-in `math.isclose`, the above equation is not symmetric in `a` and `b` -- it assumes `b` is the reference value -- so that `isclose(a, b)` might be different from `isclose(b, a)`. Furthermore, the default value of atol is not zero, and is used to determine what small values should be considered close to zero. The default value is appropriate for expected values of order unity: if the expected values are significantly smaller than one, it can result in false positives. `atol` should be carefully selected for the use case at hand. A zero value for `atol` will result in `False` if either `a` or `b` is zero. `isclose` is not defined for non-numeric data types. `bool` is considered a numeric data-type for this purpose """ assert_isinstance(other, type(self)) assert_isinstance(self, type(other)) assert_isparallel(self, other) with backend() as np: return np.isclose(self.array, other.array, rtol, atol).all()
def __repr__(self): np_array = getattr(self.array, 'numpy', lambda: self.array)() return type(self).__name__ + f"({array2string(np_array.reshape(-1))},"\ f" dom={self.dom}, cod={self.cod})" def __iter__(self): for i in self.array: yield i def __bool__(self): return bool(self.array) def __int__(self): return int(self.array) def __float__(self): return float(self.array) def __complex__(self): return complex(self.array) @classmethod def id(cls, dom=0) -> Matrix: with backend('numpy') as np: return cls(np.identity(dom, dtype=cls.dtype or int), dom, dom) twist = id @unbiased def then(self, other: Matrix) -> Matrix: assert_isinstance(other, type(self)) assert_iscomposable(self, other) with backend() as np: array = np.matmul(self.array, other.array) return type(self)(array, self.dom, other.cod) def tensor(self, other: Matrix = None, *others: Matrix): if others or other is None: return monoidal.Diagram.tensor(self, other, *others) assert_isinstance(other, type(self)) dom, cod = self.dom + other.dom, self.cod + other.cod array =, cod).array array[:self.dom, :self.cod] = self.array array[self.dom:, self.cod:] = other.array return type(self)(array, dom, cod) def __add__(self, other): assert_isinstance(other, Matrix) assert_isparallel(self, other) return type(self)(self.array + other.array, self.dom, self.cod) def __radd__(self, other): return self if other == 0 else self.__add__(other)
[docs] @classmethod def zero(cls, dom: int, cod: int) -> Matrix: """ Returns the zero matrix of a given shape. Examples -------- >>> assert, 2) == Matrix([0, 0, 0, 0], 2, 2) """ with backend() as np: return cls(np.zeros((dom, cod), dtype=cls.dtype or int), dom, cod)
[docs] @classmethod def swap(cls, left: int, right: int) -> Matrix: """ The matrix that swaps left and right dimensions. Parameters: left : The left dimension. right : The right dimension. Example ------- >>> Matrix.swap(1, 1) Matrix[int64]([0, 1, 1, 0], dom=2, cod=2) >>> Matrix.swap(2,1) Matrix[int64]([0, 1, 0, 0, 0, 1, 1, 0, 0], dom=3, cod=3) """ dom = cod = left + right array =, cod).array array[:left, right:] = array[left:, :right] = return cls(array, dom, cod)
braid = swap def transpose(self) -> Matrix: return type(self)(self.array.transpose(), self.cod, self.dom) def conjugate(self) -> Matrix: return type(self)(self.array.conjugate(), self.dom, self.cod) def dagger(self) -> Matrix: return self.conjugate().transpose() def map(self, func: Callable, dtype: type | None = None) -> Matrix: array = list(map(func, self.array.reshape(-1))) return type(self)[dtype or self.dtype](array, self.dom, self.cod)
[docs] def round(self, decimals=0) -> Matrix: """ Rounds the entries of a matrix up to a number of decimals. """ with backend() as np: array = np.around(self.array, decimals=decimals) return type(self)(array, self.dom, self.cod)
@classmethod def copy(cls, x: int, n: int) -> Matrix: array = [[i + int(j % n * x) == j for j in range(n * x)] for i in range(x)] return cls(array, x, n * x) @classmethod def discard(cls, x: int) -> Matrix: return cls.copy(x, 0) @classmethod def merge(cls, x: int, n: int) -> Matrix: return cls.copy(x, n).dagger() @classmethod def ones(cls, x: int) -> Matrix: return cls.merge(x, 0)
[docs] @classmethod def basis(cls, x: int, i: int) -> Matrix: """ The ``i``-th basis vector of dimension ``x``. Parameters: x : The dimension of the basis vector. i : The index of the basis vector. Example ------- >>> Matrix.basis(4, 2) Matrix[int64]([0, 0, 1, 0], dom=1, cod=4) """ return cls([[int(i == j) for j in range(x)]], x ** 0, x)
[docs] def repeat(self) -> Matrix: """ The reflexive transitive closure of a boolean matrix. Example ------- >>> Matrix[bool]([0, 1, 1, 0], 2, 2).repeat() Matrix[bool]([True, True, True, True], dom=2, cod=2) """ if self.dtype != bool or self.dom != self.cod: raise TypeError(messages.MATRIX_REPEAT_ERROR) return sum(*n * [self]) for n in range(self.dom + 1))
[docs] def trace(self, n=1, left=False) -> Matrix: """ The trace of a Boolean matrix, computed with :meth:`Matrix.repeat`. Parameters: n : The number of dimensions to trace. Example ------- >>> assert Matrix[bool].swap(1, 1).trace() == Matrix[bool].id(1) """ A, B, C, D = (row >> self >> column for row in [ - n) @ self.ones(n), self.ones(self.dom - n) @] for column in [ - n) @ self.discard(n), self.discard(self.cod - n) @]) return A + (B >> D.repeat() >> C)
def lambdify( self, *symbols: "sympy.Symbol", dtype=None, **kwargs) -> Callable: from sympy import lambdify with backend() as np: array = lambdify(symbols, self.array, modules=np.module, **kwargs) dtype = dtype or self.dtype return lambda *xs: type(self)[dtype](array(*xs), self.dom, self.cod) def subs(self, *args) -> Matrix: return x: getattr(x, "subs", lambda y, *_: y)(*args))
[docs] def grad(self, var, **params) -> Matrix: """ Gradient with respect to variables. """ return x: getattr(x, "diff", lambda _: 0)(var, **params))
def array2string(array, **params): """ Numpy array pretty print. """ import numpy numpy.set_printoptions(threshold=config.NUMPY_THRESHOLD) return numpy.array2string(array, **dict(params, separator=', '))\ .replace('[ ', '[').replace(' ', ' ')
[docs]class Backend: """ A matrix backend. Parameters: module : The main module of the backend. array : The array class of the backend. """ def __init__(self, module: ModuleType, array: type = None): self.module, self.array = module, array or module.array def __getattr__(self, attr): return getattr(self.module, attr)
[docs]class NumPy(Backend): """ NumPy backend. """ def __init__(self): import numpy super().__init__(numpy)
[docs]class JAX(Backend): """ JAX backend. """ def __init__(self): import jax super().__init__(jax.numpy)
[docs]class PyTorch(Backend): """ PyTorch backend. """ def __init__(self): import torch super().__init__(torch, array=torch.as_tensor)
[docs]class TensorFlow(Backend): """ TensorFlow backend. """ def __init__(self): import tensorflow.experimental.numpy as tnp from tensorflow.python.ops.numpy_ops import np_config np_config.enable_numpy_behavior() super().__init__(tnp)
BACKENDS = { 'numpy': NumPy, 'jax': JAX, 'pytorch': PyTorch, 'tensorflow': TensorFlow, } BackendName = Union[tuple(L[x] for x in BACKENDS)]
[docs]@contextmanager def backend(name: BackendName = None, _stack=[config.DEFAULT_BACKEND], _cache=dict()): """ Context manager for matrix backend. Parameters: name : The name of the backend, default is ``"numpy"``. Example ------- >>> with backend('jax'): ... assert type(Matrix([0, 1, 1, 0], 2, 2).array).__module__\\ ... == 'jaxlib.xla_extension' """ name = name or _stack[-1] _stack.append(name) try: if name not in _cache: _cache[name] = BACKENDS[name]() yield _cache[name] finally: _stack.pop()
[docs]def set_backend(name: BackendName) -> None: """ Override the default backend. Parameters: name : The name of the backend. Example ------- >>> set_backend('jax') >>> assert type(Matrix([0, 1, 1, 0], 2, 2).array).__module__\\ ... == 'jaxlib.xla_extension' >>> set_backend('numpy') >>> assert type(Matrix([0, 1, 1, 0], 2, 2).array).__module__\\ ... == 'numpy' """ backend.__wrapped__.__defaults__[1][-1] = name
[docs]def get_backend() -> Backend: """ Get the current backend. Example ------- >>> set_backend('jax') >>> assert isinstance(get_backend(), JAX) >>> set_backend('numpy') >>> assert isinstance(get_backend(), NumPy) """ with backend() as result: return result