Source code for qutip.core.dimensions

"""
Internal use module for manipulating dims specifications.
"""

# Everything should be explicitly imported, not made available by default.

import numpy as np
import numbers
from operator import getitem
from functools import partial
from qutip.settings import settings


__all__ = ["to_tensor_rep", "from_tensor_rep", "Space", "Dimensions"]


def flatten(l):
    """Flattens a list of lists to the first level.

    Given a list containing a mix of scalars and lists or a dimension object,
    flattens it down to a list of the scalars within the original list.

    Parameters
    ----------
    l : scalar, list, Space, Dimension
        Object to flatten.

    Examples
    --------

    >>> flatten([[[0], 1], 2]) # doctest: +SKIP
    [0, 1, 2]

    Notes
    -----
    Any scalar will be returned wrapped in a list: ``flaten(1) == [1]``.
    A non-list iterable will not be treated as a list by flatten. For example, flatten would treat a tuple
    as a scalar.
    """
    if isinstance(l, (Space, Dimensions)):
        l = l.as_list()
    if not isinstance(l, list):
        return [l]
    else:
        return sum(map(flatten, l), [])


def deep_remove(l, *what):
    """Removes scalars from all levels of a nested list.

    Given a list containing a mix of scalars and lists,
    returns a list of the same structure, but where one or
    more scalars have been removed.

    Examples
    --------

    >>> deep_remove([[[[0, 1, 2]], [3, 4], [5], [6, 7]]], 0, 5) # doctest: +SKIP
    [[[[1, 2]], [3, 4], [], [6, 7]]]

    """
    if isinstance(l, list):
        # Make a shallow copy at this level.
        l = l[:]
        for to_remove in what:
            if to_remove in l:
                l.remove(to_remove)
            else:
                l = [deep_remove(elem, to_remove) for elem in l]
    return l


def unflatten(l, idxs):
    """Unflattens a list by a given structure.

    Given a list of scalars and a deep list of indices
    as produced by `flatten`, returns an "unflattened"
    form of the list. This perfectly inverts `flatten`.

    Examples
    --------

    >>> l = [[[10, 20, 30], [40, 50, 60]], [[70, 80, 90], [100, 110, 120]]] # doctest: +SKIP
    >>> idxs = enumerate_flat(l) # doctest: +SKIP
    >>> unflatten(flatten(l), idxs) == l # doctest: +SKIP
    True

    """
    acc = []
    for idx in idxs:
        if isinstance(idx, list):
            acc.append(unflatten(l, idx))
        else:
            acc.append(l[idx])
    return acc


def _enumerate_flat(l, idx=0):
    if not isinstance(l, list):
        # Found a scalar, so return and increment.
        return idx, idx + 1
    else:
        # Found a list, so append all the scalars
        # from it and recurse to keep the increment
        # correct.
        acc = []
        for elem in l:
            labels, idx = _enumerate_flat(elem, idx)
            acc.append(labels)
        return acc, idx


def _collapse_composite_index(dims):
    """
    Given the dimensions specification for a composite index
    (e.g.: [2, 3] for the right index of a ket with dims [[1], [2, 3]]),
    returns a dimensions specification for an index of the same shape,
    but collapsed to a single "leg." In the previous example, [2, 3]
    would collapse to [6].
    """
    return [np.prod(dims)]


def _collapse_dims_to_level(dims, level=1):
    """
    Recursively collapses all indices in a dimensions specification
    appearing at a given level, such that the returned dimensions
    specification does not represent any composite systems.
    """
    if level == 0:
        return _collapse_composite_index(dims)
    return [_collapse_dims_to_level(index, level=level - 1) for index in dims]


def collapse_dims_oper(dims):
    """
    Given the dimensions specifications for a ket-, bra- or oper-type
    Qobj, returns a dimensions specification describing the same shape
    by collapsing all composite systems. For instance, the bra-type
    dimensions specification ``[[2, 3], [1]]`` collapses to
    ``[[6], [1]]``.

    Parameters
    ----------

    dims : list of lists of ints
        Dimensions specifications to be collapsed.

    Returns
    -------

    collapsed_dims : list of lists of ints
        Collapsed dimensions specification describing the same shape
        such that ``len(collapsed_dims[0]) == len(collapsed_dims[1]) == 1``.
    """
    return _collapse_dims_to_level(dims, 1)


def collapse_dims_super(dims):
    """
    Given the dimensions specifications for an operator-ket-, operator-bra- or
    super-type Qobj, returns a dimensions specification describing the same
    shape by collapsing all composite systems. For instance, the super-type
    dimensions specification ``[[[2, 3], [2, 3]], [[2, 3], [2, 3]]]`` collapses
    to ``[[[6], [6]], [[6], [6]]]``.

    Parameters
    ----------

    dims : list of lists of ints
        Dimensions specifications to be collapsed.

    Returns
    -------

    collapsed_dims : list of lists of ints
        Collapsed dimensions specification describing the same shape
        such that ``len(collapsed_dims[i][j]) == 1`` for ``i`` and ``j``
        in ``range(2)``.
    """
    return _collapse_dims_to_level(dims, 2)


def enumerate_flat(l):
    """Labels the indices at which scalars occur in a flattened list.

    Given a list containing a mix of scalars and lists,
    returns a list of the same structure, where each scalar
    has been replaced by an index into the flattened list.

    Examples
    --------

    >>> print(enumerate_flat([[[10], [20, 30]], 40])) # doctest: +SKIP
    [[[0], [1, 2]], 3]

    """
    return _enumerate_flat(l)[0]


def deep_map(fn, collection, over=(tuple, list)):
    if isinstance(collection, over):
        return type(collection)(deep_map(fn, el, over) for el in collection)
    else:
        return fn(collection)


def dims_to_tensor_perm(dims):
    """
    Given the dims of a Qobj instance, returns a list representing
    a permutation from the flattening of that dims specification to
    the corresponding tensor indices.

    Parameters
    ----------

    dims : list, Dimensions
        Dimensions specification for a Qobj.

    Returns
    -------

    perm : list
        A list such that ``data[flatten(dims)[idx]]`` gives the
        index of the tensor ``data`` corresponding to the ``idx``th
        dimension of ``dims``.
    """
    if isinstance(dims, list):
        dims = Dimensions(dims)
    return dims._get_tensor_perm()


def dims_to_tensor_shape(dims):
    """
    Given the dims of a Qobj instance, returns the shape of the
    corresponding tensor. This helps, for instance, resolve the
    column-stacking convention for superoperators.

    Parameters
    ----------

    dims : list, Dimensions
        Dimensions specification for a Qobj.

    Returns
    -------

    tensor_shape : tuple
        NumPy shape of the corresponding tensor.
    """
    perm = dims_to_tensor_perm(dims)
    dims = flatten(dims)
    return tuple(map(partial(getitem, dims), perm))


def dims_idxs_to_tensor_idxs(dims, indices):
    """
    Given the dims of a Qobj instance, and some indices into
    dims, returns the corresponding tensor indices. This helps
    resolve, for instance, that column-stacking for superoperators,
    oper-ket and oper-bra implies that the input and output tensor
    indices are reversed from their order in dims.

    Parameters
    ----------

    dims : list, Dimensions
        Dimensions specification for a Qobj.

    indices : int, list or tuple
        Indices to convert to tensor indices. Can be specified
        as a single index, or as a collection of indices.
        In the latter case, this can be nested arbitrarily
        deep. For instance, [0, [0, (2, 3)]].

    Returns
    -------

    tens_indices : int, list or tuple
        Container of the same structure as indices containing
        the tensor indices for each element of indices.
    """
    perm = dims_to_tensor_perm(dims)
    return deep_map(partial(getitem, perm), indices)


[docs]def to_tensor_rep(q_oper): """ Transform a ``Qobj`` to a numpy array whose shape is the flattened dimensions. Parameters ---------- q_oper: Qobj Object to reshape Returns ------- ndarray: Numpy array with one dimension for each index in dims. Examples -------- >>> ket.dims [[2, 3], [1]] >>> to_tensor_rep(ket).shape (2, 3, 1) >>> oper.dims [[2, 3], [2, 3]] >>> to_tensor_rep(oper).shape (2, 3, 2, 3) >>> super_oper.dims [[[2, 3], [2, 3]], [[2, 3], [2, 3]]] >>> to_tensor_rep(super_oper).shape (2, 3, 2, 3, 2, 3, 2, 3) """ dims = q_oper._dims data = q_oper.full().reshape(dims._get_tensor_shape()) return data.transpose(dims._get_tensor_perm())
[docs]def from_tensor_rep(tensorrep, dims): """ Reverse operator of :func:`to_tensor_rep`. Create a Qobj From a N-dimensions numpy array and dimensions with N indices. Parameters ---------- tensorrep: ndarray Numpy array with one dimension for each index in dims. dims: list of list, Dimensions Dimensions of the Qobj. Returns ------- Qobj Re constructed Qobj """ from . import Qobj dims = Dimensions(dims) data = tensorrep.transpose(np.argsort(dims._get_tensor_perm())) return Qobj(data.reshape(dims.shape), dims=dims)
def _frozen(*args, **kwargs): raise RuntimeError("Dimension cannot be modified.") class MetaSpace(type): def __call__(cls, *args, rep=None): """ Select which subclass is instantiated. """ if cls is Space and len(args) == 1 and isinstance(args[0], list): # From a list of int. return cls.from_list(*args, rep=rep) elif len(args) == 1 and isinstance(args[0], Space): # Already a Space return args[0] if cls is Space: if len(args) == 0: # Empty space: a Field. cls = Field elif len(args) == 1 and args[0] == 1: # Space(1): a Field. cls = Field elif len(args) == 1 and isinstance(args[0], Dimensions): # Making a Space out of a Dimensions object: Super Operator. cls = SuperSpace elif len(args) > 1 and all(isinstance(arg, Space) for arg in args): # list of space: tensor product space. cls = Compound if settings.core['auto_tidyup_dims']: if cls is Compound and all(isinstance(arg, Field) for arg in args): cls = Field if cls is SuperSpace and args[0].type == "scalar": cls = Field args = tuple([ tuple(arg) if isinstance(arg, list) else arg for arg in args ]) if cls is Field: return cls.field_instance if cls is SuperSpace: args = (*args, rep or 'super') if args not in cls._stored_dims: instance = cls.__new__(cls) instance.__init__(*args) cls._stored_dims[args] = instance return cls._stored_dims[args] def from_list(cls, list_dims, rep=None): if len(list_dims) == 0: raise ValueError("Empty list can't be used as dims.") elif ( sum(isinstance(entry, list) for entry in list_dims) not in [0, len(list_dims)] ): raise ValueError(f"Format dims not understood {list_dims}.") elif not isinstance(list_dims[0], list): # Tensor spaces = [Space(size) for size in list_dims] elif len(list_dims) == 1: # [[2, 3]]: tensor with an extra layer of list. spaces = [Space(size) for size in list_dims[0]] elif len(list_dims) % 2 == 0: # Superoperators or tensor of Superoperators spaces = [ Space(Dimensions( Space(list_dims[i+1]), Space(list_dims[i]) ), rep=rep) for i in range(0, len(list_dims), 2) ] else: raise ValueError(f'Format not understood {list_dims}') if len(spaces) == 1: return spaces[0] elif len(spaces) >= 2: return Space(*spaces) else: raise ValueError(f'Format not understood {list_dims}') class Space(metaclass=MetaSpace): _stored_dims = {} def __init__(self, dims): idims = int(dims) if idims <= 0 or idims != dims: raise ValueError("Dimensions must be integers > 0") # Size of the hilbert space self.size = dims self.issuper = False # Super representation, should be an empty string except for SuperSpace self.superrep = None # Does the size and dims match directly: size == prod(dims) self._pure_dims = True self.__setitem__ = _frozen def __eq__(self, other): return self is other or ( type(other) is type(self) and other.size == self.size ) def __hash__(self): return hash(self.size) def __repr__(self): return f"Space({self.size})" def as_list(self): return [self.size] def __str__(self): return str(self.as_list()) def dims2idx(self, dims): """ Transform dimensions indices to full array indices. """ if not isinstance(dims, list) or len(dims) != 1: raise ValueError("Dimensions must be a list of one element") if not (0 <= dims[0] < self.size): raise IndexError("Dimensions out of range") if not isinstance(dims[0], numbers.Integral): raise TypeError("Dimensions must be integers") return dims[0] def idx2dims(self, idx): """ Transform full array indices to dimensions indices. """ if not (0 <= idx < self.size): raise IndexError("Index out of range") return [idx] def step(self): """ Get the step in the array between for each dimensions index. If element ``[i, j, k]`` is ``ket.full()[m, 0]`` then element ``[i, j+1, k]`` is ``ket.full()[m + ket._dims.step()[1], 0]``. """ return [1] def flat(self): """ Dimensions as a flat list. """ return [self.size] def remove(self, idx): """ Remove a Space from a Dimensons or complex Space. ``Space([2, 3, 4]).remove(1) == Space([2, 4])`` """ raise RuntimeError("Cannot delete a flat space.") def replace(self, idx, new): """ Reshape a Space from a Dimensons or complex Space. ``Space([2, 3, 4]).replace(1, 5) == Space([2, 5, 4])`` """ if idx != 0: raise ValueError( "Cannot replace a non-zero index in a flat space." ) return Space(new) def replace_superrep(self, super_rep): return self def scalar_like(self): return Field() class Field(Space): field_instance = None def __init__(self): self.size = 1 self.issuper = False self.superrep = None self._pure_dims = True self.__setitem__ = _frozen def __eq__(self, other): return type(other) is Field def __hash__(self): return hash(0) def __repr__(self): return "Field()" def as_list(self): return [1] def step(self): return [1] def flat(self): return [1] def remove(self, idx): return self def replace(self, idx, new): return Space(new) Field.field_instance = Field.__new__(Field) Field.field_instance.__init__() class Compound(Space): _stored_dims = {} def __init__(self, *spaces): self.spaces = [] if len(spaces) <= 1: raise ValueError("Compound need multiple space to join.") for space in spaces: if isinstance(space, Compound): self.spaces += space.spaces else: self.spaces += [space] self.spaces = tuple(self.spaces) self.size = np.prod([space.size for space in self.spaces]) self.issuper = all(space.issuper for space in self.spaces) if not self.issuper and any(space.issuper for space in self.spaces): raise TypeError( "Cannot create compound space of super and non super." ) self._pure_dims = all(space._pure_dims for space in self.spaces) superrep = [space.superrep for space in self.spaces] if all(superrep[0] == rep for rep in superrep): self.superrep = superrep[0] else: raise TypeError( "Cannot create compound space of of super operators " "with different representation." ) self.__setitem__ = _frozen def __eq__(self, other): return self is other or ( type(other) is type(self) and self.spaces == other.spaces ) def __hash__(self): return hash(self.spaces) def __repr__(self): parts_rep = ", ".join(repr(space) for space in self.spaces) return f"Compound({parts_rep})" def as_list(self): return sum([space.as_list() for space in self.spaces], []) def dims2idx(self, dims): if len(dims) != len(self.spaces): raise ValueError("Length of supplied dims does not match the number of subspaces.") pos = 0 step = 1 for space, dim in zip(self.spaces[::-1], dims[::-1]): pos += space.dims2idx([dim]) * step step *= space.size return pos def idx2dims(self, idx): dims = [] for space in self.spaces[::-1]: idx, dim = divmod(idx, space.size) dims = space.idx2dims(dim) + dims return dims def step(self): steps = [] step = 1 for space in self.spaces[::-1]: steps = [step * N for N in space.step()] + steps step *= space.size return steps def flat(self): return sum([space.flat() for space in self.spaces], []) def remove(self, idx): new_spaces = [] for space in self.spaces: n_indices = len(space.flat()) idx_space = [i for i in idx if i < n_indices] idx = [i-n_indices for i in idx if i >= n_indices] new_space = space.remove(idx_space) if new_spaces: return Compound(*new_spaces) return Field() def replace(self, idx, new): new_spaces = [] for space in self.spaces: n_indices = len(space.flat()) if 0 <= idx < n_indices: new_spaces.append(space.replace(idx, new)) else: new_spaces.append(space) idx -= n_indices return Compound(*new_spaces) def replace_superrep(self, super_rep): return Compound( *[space.replace_superrep(super_rep) for space in self.spaces] ) def scalar_like(self): return [space.scalar_like() for space in self.spaces] class SuperSpace(Space): _stored_dims = {} def __init__(self, oper, rep='super'): self.oper = oper self.superrep = rep self.size = oper.shape[0] * oper.shape[1] self.issuper = True self._pure_dims = oper._pure_dims self.__setitem__ = _frozen def __eq__(self, other): return ( self is other or self.oper == other or ( type(other) is type(self) and self.oper == other.oper and self.superrep == other.superrep ) ) def __hash__(self): return hash((self.oper, self.superrep)) def __repr__(self): return f"Super({repr(self.oper)}, rep={self.superrep})" def as_list(self): return self.oper.as_list() def dims2idx(self, dims): posl, posr = self.oper.dims2idx(dims) return posl + posr * self.oper.shape[0] def idx2dims(self, idx): posl = idx % self.oper.shape[0] posr = idx // self.oper.shape[0] return self.oper.idx2dims(posl, posr) def step(self): stepl, stepr = self.oper.step() step = self.oper.shape[0] return stepl + [step * N for N in stepr] def flat(self): return sum(self.oper.flat(), []) def remove(self, idx): new_dims = self.oper.remove(idx) if new_dims.type == 'scalar': return Field() return SuperSpace(new_dims, rep=self.superrep) def replace(self, idx, new): return SuperSpace(self.oper.replace(idx, new), rep=self.superrep) def replace_superrep(self, super_rep): return SuperSpace(self.oper, rep=super_rep) def scalar_like(self): return self.oper.scalar_like() class MetaDims(type): def __call__(cls, *args, rep=None): if len(args) == 1 and isinstance(args[0], Dimensions): return args[0] elif len(args) == 1 and len(args[0]) == 2: args = ( Space(args[0][1], rep=rep), Space(args[0][0], rep=rep) ) elif len(args) != 2: raise NotImplementedError('No Dual, Ket, Bra...', args) elif ( settings.core["auto_tidyup_dims"] and args[0] == args[1] == Field() ): return Field() if args not in cls._stored_dims: instance = cls.__new__(cls) instance.__init__(*args) cls._stored_dims[args] = instance return cls._stored_dims[args] class Dimensions(metaclass=MetaDims): _stored_dims = {} _type = None def __init__(self, from_, to_): self.from_ = from_ self.to_ = to_ self.shape = to_.size, from_.size self.issuper = from_.issuper self._pure_dims = from_._pure_dims and to_._pure_dims self.issquare = False if self.from_.size == 1 and self.to_.size == 1: self.type = 'scalar' self.issquare = True self.superrep = None elif self.from_.size == 1: self.issuper = self.to_.issuper self.type = 'operator-ket' if self.issuper else 'ket' self.superrep = self.to_.superrep elif self.to_.size == 1: self.issuper = self.from_.issuper self.type = 'operator-bra' if self.issuper else 'bra' self.superrep = self.from_.superrep elif self.from_ == self.to_: self.issuper = self.from_.issuper self.type = 'super' if self.issuper else 'oper' self.superrep = self.from_.superrep self.issquare = True else: if from_.issuper != to_.issuper: raise NotImplementedError( "Operator with both space and superspace dimensions are " "not supported. Please open an issue if you have an use " f"case for these: {from_}, {to_}]" ) self.type = 'super' if self.from_.issuper else 'oper' if self.from_.superrep == self.to_.superrep: self.superrep = self.from_.superrep else: self.superrep = 'mixed' self.__setitem__ = _frozen def __eq__(self, other): if isinstance(other, Dimensions): return ( self is other or ( self.to_ == other.to_ and self.from_ == other.from_ ) ) return NotImplemented def __ne__(self, other): if isinstance(other, Dimensions): return not ( self is other or ( self.to_ == other.to_ and self.from_ == other.from_ ) ) return NotImplemented def __matmul__(self, other): if self.from_ != other.to_: raise TypeError(f"incompatible dimensions {self} and {other}") args = other.from_, self.to_ if args in Dimensions._stored_dims: return Dimensions._stored_dims[args] return Dimensions(*args) def __hash__(self): return hash((self.to_, self.from_)) def __repr__(self): return f"Dimensions({repr(self.from_)}, {repr(self.to_)})" def __str__(self): return str(self.as_list()) def as_list(self): """ Return the list representation of the Dimensions object. """ return [self.to_.as_list(), self.from_.as_list()] def __getitem__(self, key): if key == 0: return self.to_ elif key == 1: return self.from_ raise IndexError("Dimensions index out of range") def dims2idx(self, dims): """ Transform dimensions indices to full array indices. """ return self.to_.dims2idx(dims[0]), self.from_.dims2idx(dims[1]) def idx2dims(self, idxl, idxr): """ Transform full array indices to dimensions indices. """ return [self.to_.idx2dims(idxl), self.from_.idx2dims(idxr)] def step(self): """ Get the step in the array between for each dimensions index. If element ``[i, j, k]`` is ``ket.full()[m, 0]`` then element ``[i, j+1, k]`` is ``ket.full()[m + ket._dims.step()[1], 0]``. """ return [self.to_.step(), self.from_.step()] def flat(self): """ Dimensions as a flat list. """ return [self.to_.flat(), self.from_.flat()] def _get_tensor_shape(self): """ Get the shape to of the Nd tensor with one dimensions for each Dimension index. The order of the space values are not in the order of the Dimension index. """ # dims_to_tensor_shape stepl = self.to_.step() flatl = self.to_.flat() stepr = self.from_.step() flatr = self.from_.flat() return tuple(np.concatenate([ np.array(flatl)[np.argsort(stepl)[::-1]], np.array(flatr)[np.argsort(stepr)[::-1]], ])) def _get_tensor_perm(self): """ Get the permutation of a tensor created using ``_get_tensor_shape`` to reorder the tensor dimensions with those of the Dimensions object. """ # dims_to_tensor_perm stepl = self.to_.step() stepr = self.from_.step() return list(np.concatenate([ np.argsort(stepl)[::-1], np.argsort(stepr)[::-1] + len(stepl) ])) def remove(self, idx): """ Remove a Space from a Dimensons or complex Space. ``Space([2, 3, 4]).remove(1) == Space([2, 4])`` """ if not isinstance(idx, list): idx = [idx] if not idx: return self idx = sorted(idx) n_indices = len(self.to_.flat()) idx_to = [i for i in idx if i < n_indices] idx_from = [i-n_indices for i in idx if i >= n_indices] return Dimensions( self.from_.remove(idx_from), self.to_.remove(idx_to), ) def replace(self, idx, new): """ Reshape a Space from a Dimensons or complex Space. ``Space([2, 3, 4]).replace(1, 5) == Space([2, 5, 4])`` """ n_indices = len(self.to_.flat()) if idx < n_indices: new_to = self.to_.replace(idx, new) new_from = self.from_ else: new_to = self.to_ new_from = self.from_.replace(idx-n_indices, new) return Dimensions(new_from, new_to) def replace_superrep(self, super_rep): if not self.issuper and super_rep is not None: raise TypeError("Can't set a superrep of a non super object.") return Dimensions( self.from_.replace_superrep(super_rep), self.to_.replace_superrep(super_rep) ) def scalar_like(self): return [self.to_.scalar_like(), self.from_.scalar_like()]