Source code for qutip.dimensions

# This file is part of QuTiP: Quantum Toolbox in Python.
#
#    Copyright (c) 2011 and later, Paul D. Nation and Robert J. Johansson.
#    All rights reserved.
#
#    Redistribution and use in source and binary forms, with or without
#    modification, are permitted provided that the following conditions are
#    met:
#
#    1. Redistributions of source code must retain the above copyright notice,
#       this list of conditions and the following disclaimer.
#
#    2. Redistributions in binary form must reproduce the above copyright
#       notice, this list of conditions and the following disclaimer in the
#       documentation and/or other materials provided with the distribution.
#
#    3. Neither the name of the QuTiP: Quantum Toolbox in Python nor the names
#       of its contributors may be used to endorse or promote products derived
#       from this software without specific prior written permission.
#
#    THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS
#    "AS IS" AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT
#    LIMITED TO, THE IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A
#    PARTICULAR PURPOSE ARE DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT
#    HOLDER OR CONTRIBUTORS BE LIABLE FOR ANY DIRECT, INDIRECT, INCIDENTAL,
#    SPECIAL, EXEMPLARY, OR CONSEQUENTIAL DAMAGES (INCLUDING, BUT NOT
#    LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR SERVICES; LOSS OF USE,
#    DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER CAUSED AND ON ANY
#    THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, OR TORT
#    (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE
#    OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.
###############################################################################
"""
Internal use module for manipulating dims specifications.
"""

__all__ = [] # Everything should be explicitly imported, not made available
             # by default.

import numpy as np
from operator import getitem
from functools import partial

[docs]def is_scalar(dims): """ Returns True if a dims specification is effectively a scalar (has dimension 1). """ return np.prod(flatten(dims)) == 1
def is_vector(dims): return ( isinstance(dims, list) and isinstance(dims[0], (int, np.integer)) ) def is_vectorized_oper(dims): return ( isinstance(dims, list) and isinstance(dims[0], list) ) def type_from_dims(dims, enforce_square=False): bra_like, ket_like = map(is_scalar, dims) if bra_like: if is_vector(dims[1]): return 'bra' elif is_vectorized_oper(dims[1]): return 'operator-bra' if ket_like: if is_vector(dims[0]): return 'ket' elif is_vectorized_oper(dims[0]): return 'operator-ket' elif is_vector(dims[0]) and (dims[0] == dims[1] or not enforce_square): return 'oper' elif ( is_vectorized_oper(dims[0]) and ( ( dims[0] == dims[1] and dims[0][0] == dims[1][0] ) or not enforce_square ) ): return 'super' return 'other'
[docs]def flatten(l): """Flattens a list of lists to the first level. Given a list containing a mix of scalars and lists, flattens down to a list of the scalars within the original list. Examples -------- >>> flatten([[[0], 1], 2]) # doctest: +SKIP [0, 1, 2] """ if not isinstance(l, list): return [l] else: return sum(map(flatten, l), [])
[docs]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 = list(map(lambda elem: deep_remove(elem, to_remove), l)) return l
[docs]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) else: return [_collapse_dims_to_level(index, level=level - 1) for index in dims]
[docs]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)
[docs]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)
[docs]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)
[docs]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 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``. """ # We figure out the type of the dims specification, # relaxing the requirement that operators be square. # This means that dims_type need not coincide with # Qobj.type, but that works fine for our purposes here. dims_type = type_from_dims(dims, enforce_square=False) perm = enumerate_flat(dims) # If type is oper, ket or bra, we don't need to do anything. if dims_type in ('oper', 'ket', 'bra'): return flatten(perm) # If the type is other, we need to figure out if the # dims is superlike on its outputs and inputs # This is the case if the dims type for left or right # are, respectively, oper-like. if dims_type == 'other': raise NotImplementedError("Not yet implemented for type='other'.") # If we're still here, the story is more complicated. We'll # follow the strategy of creating a permutation by using # enumerate_flat then transforming the result to swap # input and output indices of vectorized matrices, then flattening # the result. We'll then rebuild indices using this permutation. if dims_type in ('operator-ket', 'super'): # Swap the input and output spaces of the right part of # perm. perm[1] = list(reversed(perm[1])) if dims_type in ('operator-bra', 'super'): # Ditto, but for the left indices. perm[0] = list(reversed(perm[0])) return flatten(perm)
[docs]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 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))
[docs]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 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)