# Required for Sphinx to follow autodoc_type_aliases
from __future__ import annotations
from . import data as _data
from .cy.qobjevo import QobjEvo
from .dimensions import Dimensions, Field, SumSpace
from .operators import qzero
from .qobj import Qobj
from .superoperator import operator_to_vector
from .. import settings
from ..typing import DimensionLike, LayerType
from functools import partial
from numbers import Number
from typing import Any, Union, overload
import numpy as np
__all__ = ['direct_sum', 'direct_sum_sparse',
'direct_component', 'set_direct_component']
ARRAY_ONE = np.full((1,), 1, dtype=_data.base.idxint_dtype)
def _is_like_ket(qobj: Any) -> bool:
return (
isinstance(qobj, Number)
or (
isinstance(qobj, (Qobj, QobjEvo))
and qobj.type in ['scalar', 'ket']
))
def _is_like_bra(qobj: Any) -> bool:
return (
isinstance(qobj, Number)
or (
isinstance(qobj, (Qobj, QobjEvo))
and qobj.type in ['scalar', 'bra']
))
def _is_like_oper(qobj: Any) -> bool:
return (
qobj is None
or isinstance(qobj, Number)
or (
isinstance(qobj, (Qobj, QobjEvo))
and qobj.type in ['scalar', 'oper', 'ket', 'bra']
))
def _is_like_operator_ket(qobj: Any) -> bool:
return (
isinstance(qobj, Number)
or (
isinstance(qobj, (Qobj, QobjEvo))
and qobj.type in ['scalar', 'oper', 'operator-ket']
))
def _is_like_operator_bra(qobj: Any) -> bool:
return (
isinstance(qobj, Number)
or (
isinstance(qobj, (Qobj, QobjEvo))
and qobj.type in ['scalar', 'oper', 'operator-bra']
))
def _is_like_super(qobj: Any) -> bool:
return (
qobj is None
or isinstance(qobj, Number)
or (
isinstance(qobj, (Qobj, QobjEvo))
and qobj.type in [
'scalar', 'super', 'operator-ket', 'operator-bra']
))
def _qobj_data(qobj: Number | Qobj, scalar_dtype: LayerType) -> _data.Data:
if qobj is None:
return None
if isinstance(qobj, Number):
return _data.identity[scalar_dtype](1, scale=qobj)
return qobj.data
def _qobj_dims(qobj: Number | Qobj | QobjEvo) -> Dimensions:
return (
qobj._dims if isinstance(qobj, (Qobj, QobjEvo))
else Dimensions(Field(), Field())
)
@overload
def direct_sum(
qobjs: list[Number | Qobj] | list[list[Number | Qobj]],
dtype: LayerType = None
) -> Qobj:
...
@overload
def direct_sum(
qobjs: list[Number | Qobj | QobjEvo] | list[list[Number | Qobj | QobjEvo]],
dtype: LayerType = None
) -> QobjEvo:
...
[docs]
def direct_sum(qobjs, dtype=None):
"""
Construct the direct sum of the provided component quantum objects.
The matrix representing the returned direct sum is a block matrix, where
the matrices of the components are all concatenated.
This function accepts either a 1D list of components or a 2D square matrix
(list of lists) of components. A 1D list is interpreted as a column or row
block vector, depending on the element types:
* If all elements are kets or scalars, the output will be a ket.
* If all elements are bras or scalars, the output will be a bra.
* If all elements are operator-kets, operators, or scalars, the operators
will be converted to operator-kets automatically, and the output will be
an operator-ket.
* If all elements are operator-bras, operators, or scalars, the operators
will be converted to operator-bras automatically, and the output will be
an operator-bra.
Python numbers are always accepted as scalars.
A 2D square matrix may contain any types of Qobj, but the horizontal and
vertical dimensions must match in each block row and block column. In a 2D
square matrix, entries may be ``None`` to indicate zero blocks. However,
each block row and each block column must contain at least one not-``None``
entry. For direct sums with all-zero rows or columns, consider using
:func:`.direct_sum_sparse`.
Parameters
----------
qobjs : list
1D list or 2D square matrix of components. The components may be Python
numbers, :class:`.Qobj`, or :class:`.QobjEvo`.
dtype : type or str, optional
Storage representation for the output ``Qobj``. Any data-layer known
to ``qutip.data.to`` is accepted.
Returns
-------
direct_sum: :class:`.Qobj` or :class:`.QobjEvo`
The assembled direct-sum object. Returns a time-dependent object
(:class:`.QobjEvo`) if any input component is time-dependent.
"""
if len(qobjs) == 0:
raise ValueError("No Qobjs provided for direct sum.")
# Step 1: check type of all provided qobj and make (row, col) -> qobj dict
linear = isinstance(qobjs[0], (Number, Qobj, QobjEvo))
if not linear and len(qobjs[0]) == 0:
raise ValueError("No Qobjs provided for direct sum.")
if not linear and not all(len(row) == len(qobjs[0]) for row in qobjs):
raise ValueError("Matrix of Qobjs in direct sum must be square.")
if linear and all(_is_like_ket(qobj) for qobj in qobjs):
num_rows = len(qobjs)
num_columns = 1
qobj_dict = {(row, 0): qobj for row, qobj in enumerate(qobjs)}
elif linear and all(_is_like_operator_ket(qobj) for qobj in qobjs):
# for convenience, we call operator_to_vector on operators provided
# in matrix form
def _ensure_vector(qobj):
if isinstance(qobj, (Qobj, QobjEvo)) and qobj.type == 'oper':
return operator_to_vector(qobj)
return qobj
num_rows = len(qobjs)
num_columns = 1
qobj_dict = {(row, 0): _ensure_vector(qobj)
for row, qobj in enumerate(qobjs)}
elif linear and all(_is_like_bra(qobj) for qobj in qobjs):
num_rows = 1
num_columns = len(qobjs)
qobj_dict = {(0, col): qobj for col, qobj in enumerate(qobjs)}
elif linear and all(_is_like_operator_bra(qobj) for qobj in qobjs):
# for convenience, we call operator_to_vector on operators provided
# in matrix form
def _ensure_vector(qobj):
if isinstance(qobj, (Qobj, QobjEvo)) and qobj.type == 'oper':
return operator_to_vector(qobj).dag()
return qobj
num_rows = 1
num_columns = len(qobjs)
qobj_dict = {(0, col): _ensure_vector(qobj)
for col, qobj in enumerate(qobjs)}
elif linear:
raise ValueError("Invalid combination of Qobj types for direct sum.")
else:
allsuper = all(_is_like_super(qobj) for row in qobjs for qobj in row)
alloper = all(_is_like_oper(qobj) for row in qobjs for qobj in row)
if not (allsuper or alloper):
raise ValueError(
"Invalid combination of Qobj types for direct sum.")
num_rows = len(qobjs)
num_columns = len(qobjs[0])
qobj_dict = {(row, col): qobjs[row][col]
for row in range(num_rows) for col in range(num_columns)
if qobjs[row][col] is not None}
# Step 2: determine dimensions of sum space
to_spaces = [None] * num_rows
from_spaces = [None] * num_columns
for (row, col), qobj in qobj_dict.items():
to_space = _qobj_dims(qobj)[0]
if to_spaces[row] is None:
to_spaces[row] = to_space
elif to_spaces[row] != to_space:
raise ValueError(
f"Direct sum: inconsistent dimensions in row {row}. Expected"
f" {to_spaces[row].as_list()}, got {to_space.as_list()}.")
from_space = _qobj_dims(qobj)[1]
if from_spaces[col] is None:
from_spaces[col] = from_space
elif from_spaces[col] != from_space:
raise ValueError(
f"Direct sum: inconsistent dimensions in column {col}."
f" Expected {from_spaces[col].as_list()}, got"
f" {from_space.as_list()}.")
for row, space in enumerate(to_spaces):
if space is None:
raise ValueError(f"Direct sum: empty row {row}.")
for col, space in enumerate(from_spaces):
if space is None:
raise ValueError(f"Direct sum: empty column {col}.")
sum_dimension = Dimensions(SumSpace(from_spaces), SumSpace(to_spaces))
# Step 3: create direct sum
return _do_direct_sum(qobj_dict, sum_dimension, dtype)
@overload
def direct_sum_sparse(
qobjs: dict[tuple[int, int], Number | Qobj],
sum_dimensions: DimensionLike,
dtype: LayerType = None
) -> Qobj:
...
@overload
def direct_sum_sparse(
qobjs: dict[tuple[int, int], Number | Qobj | QobjEvo],
sum_dimensions: DimensionLike,
dtype: LayerType = None
) -> QobjEvo:
...
[docs]
def direct_sum_sparse(qobjs, sum_dimensions, dtype=None):
r"""
Construct the direct sum of the provided component quantum objects.
This is a variant of :func:`.direct_sum` suitable for large, sparse block
matrices. The caller must provide the dimensions of the direct sum, for
example in list form as in :code:`[([2], [3]), ([2], [3])]` for an
operator on :math:`\mathbb C^2 \oplus \mathbb C^3`.
In more detail: the ``sum_dimensions`` parameter is a list of two entries.
The two entries of this list respectively correspond to the output and
input Hilbert spaces of the quantum object. Each entry is the "list
representation" of the dimensions of the corresponding Hilbert space, like
it is used throughout QuTiP. For use in this function, either one of the
two spaces or both of them must have the structure of a direct sum, which
is represented as a tuple in the "list representation".
List representation syntax overview:
* A normal, unmarked Hilbert space of dimension :math:`n` is represented
by the list :code:`[n]`.
* A product Hilbert space, where the constituents have the dimensions
:math:`n_1` up to :math:`n_k`, is represented by the list
:code:`[n1, n2, ..., nk]`.
* The direct sum of Hilbert spaces, which have the list representations
:code:`l1` up to :code:`[lk]`, is represented by the tuple
:code:`[l1, l2, ..., lk]`. Sum dimensions may be nested.
* For example, the direct sum of a 2D space and a 3D space has the list
representation :code:`([2], [3])`.
* As another example, the direct sum of a 2D space, and the product of two
3D spaces, is :code:`([2], [3, 3])`.
* Superoperators act on operator spaces; the list representation of an
operator space is the list :code:`[lOut, lIn]`, where :code:`lIn` and
:code:`lOut` are list representations of the input and output spaces
of the operators. Direct sum components may be either regular Hilbert
spaces or operator spaces, but not mixed within the same quantum object.
Parameters
----------
qobjs : dict
Mapping that assigns a component to the given (row, col) location in
the block matrix. Missing locations are filled with zero blocks. The
components may be Python numbers, :class:`.Qobj`, or :class:`.QobjEvo`.
sum_dimensions : list of tuples or lists
Dimensions of the resulting Qobj, like in the creation of a Qobj.
See detailed explanation above.
dtype : type or str, optional
Storage representation for the output ``Qobj``. Any data-layer known
to ``qutip.data.to`` is accepted.
Returns
-------
direct_sum: :class:`.Qobj` or :class:`.QobjEvo`
The assembled direct-sum object. Returns a time-dependent object
(:class:`.QobjEvo`) if any input component is time-dependent.
"""
sum_dimensions = Dimensions(sum_dimensions)
to_spaces, from_spaces = _spaces_from_dims(sum_dimensions)
for (row, col), qobj in qobjs.items():
_check_bounds(row, 0, len(to_spaces))
_check_bounds(col, 0, len(from_spaces))
if (
_qobj_dims(qobj)[0] != to_spaces[row] or
_qobj_dims(qobj)[1] != from_spaces[col]
):
raise ValueError("Direct sum: dimension mismatch for component at"
f" ({row}, {col}).")
return _do_direct_sum(dict(sorted(qobjs.items())), sum_dimensions, dtype)
def _do_direct_sum(
qobjs: dict[tuple[int, int], Number | Qobj | QobjEvo],
sum_dimensions: Dimensions,
dtype: LayerType = None
):
"""Assumes `qobjs` is sorted and performs no dimensions checks"""
if settings.core["default_dtype_scope"] == "full":
dtype = _data._parse_default_dtype(dtype, "sparse")
else:
dtype = _data._parse_default_dtype(dtype or "core", "sparse")
nonevo_count = sum(
not isinstance(qobj, QobjEvo) for qobj in qobjs.values()
)
block_rows = np.empty((nonevo_count,), dtype=_data.base.idxint_dtype)
block_cols = np.empty((nonevo_count,), dtype=_data.base.idxint_dtype)
blocks = np.empty((nonevo_count,), dtype=_data.Data)
i = 0
qobjevos = []
for (row, col), qobj in qobjs.items():
if not isinstance(qobj, QobjEvo):
block_rows[i] = row
block_cols[i] = col
blocks[i] = _qobj_data(qobj, dtype)
i += 1
else:
qobjevos.append((row, col, qobj))
if isinstance(sum_dimensions[0], SumSpace):
block_heights = np.array(
sum_dimensions[0]._space_dims, dtype=_data.base.idxint_dtype)
else:
block_heights = np.full(
(1,), sum_dimensions[0].size, dtype=_data.base.idxint_dtype)
if isinstance(sum_dimensions[1], SumSpace):
block_widths = np.array(
sum_dimensions[1]._space_dims, dtype=_data.base.idxint_dtype)
else:
block_widths = np.full(
(1,), sum_dimensions[1].size, dtype=_data.base.idxint_dtype)
result = Qobj(
_data.block_build(
block_rows, block_cols, blocks,
block_heights, block_widths, dtype=dtype
), dims=sum_dimensions, copy=False
)
for row, col, qobjevo in qobjevos:
result += qobjevo.linear_map(
partial(_blow_up_qobj, sum_dimensions=sum_dimensions,
row=row, col=col, dtype=dtype),
_skip_check=True
)
if qobjevos:
result.compress(_skip_coeff=True)
return result
@overload
def direct_component(sum_qobj: Qobj, *index: int | slice) -> Qobj:
...
@overload
def direct_component(sum_qobj: QobjEvo, *index: int | slice) -> Qobj | QobjEvo:
...
[docs]
def direct_component(sum_qobj, *index):
"""
Extract the component (block) at the given index from a direct sum
``Qobj``.
Parameters
----------
sum_qobj : :class:`.Qobj` or :class:`.QobjEvo`
A direct sum object from which to extract a component.
index : list of int or slice
One or two indices identifying the component to extract. If both the
row and column spaces are sums, two indices (``row``, ``col``) are
required. If only one side is a sum, a single index is accepted and
interpreted as the index into the summed side.
Returns
-------
component: :class:`.Qobj` or :class:`.QobjEvo`
The extracted component.
Notes
-----
Each index may be either an int indicating a single component or a slice
indicating multiple components. For convenience, slices can be created
using `np.s_[start:stop]`. Custom steps are not supported.
"""
if settings.core["default_dtype_scope"] == "full":
dtype = settings.core["default_dtype"] or sum_qobj.dtype
else:
dtype = sum_qobj.dtype
if isinstance(sum_qobj, QobjEvo):
result = sum_qobj.linear_map(lambda x: direct_component(x, *index),
_skip_check=True)
result.compress(_skip_coeff=True)
return result(0) if result.isconstant else result
to_index, from_index = _check_component_index(sum_qobj._dims, index)
component_dims, row_start, row_stop, col_start, col_stop =\
_component_info(sum_qobj._dims, to_index, from_index)
out_data = _data.block_extract(
sum_qobj.data, row_start, row_stop, col_start, col_stop, dtype=dtype
)
return Qobj(out_data, dims=component_dims, copy=False)
@overload
def set_direct_component(
sum_qobj: Qobj,
component: Qobj,
*index: int | slice,
dtype: LayerType = None
) -> Qobj:
...
@overload
def set_direct_component(
sum_qobj: Qobj | QobjEvo,
component: Qobj | QobjEvo,
*index: int | slice,
dtype: LayerType = None
) -> QobjEvo:
...
[docs]
def set_direct_component(sum_qobj, component, *index):
"""
Set (replace) a component in a direct sum ``Qobj``. The function returns a
new object where the component at the given index is replaced with
``component``.
Parameters
----------
sum_qobj : :class:`.Qobj` or :class:`.QobjEvo`
The direct sum object whose component will be replaced.
component : :class:`.Qobj` or :class:`.QobjEvo`
The new component to insert. ``None`` sets the component to zero.
index : list of int or slice
One or two indices identifying the component to extract. If both the
row and column spaces are sums, two indices (``row``, ``col``) are
required. If only one side is a sum, a single index is accepted and
interpreted as the index into the summed side.
Returns
-------
updated: :class:`.Qobj` or :class:`.QobjEvo`
The resulting direct sum object with the component set.
Notes
-----
Each index may be either an int indicating a single component or a slice
indicating multiple components. For convenience, slices can be created
using `np.s_[start:stop]`. Custom steps are not supported.
"""
if settings.core["default_dtype_scope"] == "full":
dtype = settings.core["default_dtype"] or sum_qobj.dtype
else:
dtype = sum_qobj.dtype
to_index, from_index = _check_component_index(sum_qobj._dims, index)
component_dims, row_start, _, col_start, _ =\
_component_info(sum_qobj._dims, to_index, from_index)
if component is not None and _qobj_dims(component) != component_dims:
raise ValueError("Canot set component of direct sum: dimension"
f" mismatch. Expected: {component_dims.as_list()},"
f" got: {_qobj_dims(component).as_list()}.")
if not (isinstance(sum_qobj, QobjEvo) or isinstance(component, QobjEvo)):
component_data = (
_qobj_data(component, dtype)
or _data.zeros[dtype](*component_dims.shape)
)
out_data = _data.block_overwrite(sum_qobj.data, component_data,
row_start, col_start, dtype=dtype)
return Qobj(out_data, dims=sum_qobj._dims, copy=False)
# If QobjEvo is involved, the result is the sum of two parts:
# - the sum_qobj with the component zeroed
# - the new component embedded into a large zero matrix
# Note that all "set_component" below act on simple Qobj
if isinstance(sum_qobj, QobjEvo):
zeroed = sum_qobj.linear_map(
lambda x: set_direct_component(x, None, *index),
_skip_check=True)
zeroed.compress(_skip_coeff=True)
zeroed = zeroed(0) if zeroed.isconstant else zeroed
else:
zeroed = set_direct_component(sum_qobj, None, *index)
blow_up_func = partial(_blow_up_qobj, sum_dimensions=sum_qobj._dims,
row=to_index, col=from_index, dtype=dtype)
if isinstance(component, QobjEvo):
blow_up = component.linear_map(blow_up_func, _skip_check=True)
else:
blow_up = blow_up_func(component)
return zeroed + blow_up
def _blow_up_qobj(x: Qobj, *, sum_dimensions, row, col, dtype):
# given small Qobj x, makes big Qobj with given dimensions
# resulting big Qobj is zero except for the block x starting at row, col
if x is None:
return qzero(sum_dimensions[0], sum_dimensions[1])
if type(row) is slice:
row = row.start
if type(col) is slice:
col = col.start
if isinstance(sum_dimensions[0], SumSpace):
data_row = sum_dimensions[0]._space_cumdim(row)
else:
data_row = 0
if isinstance(sum_dimensions[1], SumSpace):
data_col = sum_dimensions[1]._space_cumdim(col)
else:
data_col = 0
xh, xw = x.shape
sh, sw = sum_dimensions.shape
return Qobj(
_data.block_build(
ARRAY_ONE,
ARRAY_ONE,
np.full((1,), x.data, dtype=_data.Data),
np.array([data_row, xh, sh - xh - data_row],
dtype=_data.base.idxint_dtype),
np.array([data_col, xw, sw - xw - data_col],
dtype=_data.base.idxint_dtype),
dtype=dtype
),
dims=sum_dimensions, copy=False
)
def _check_bounds(given, min, max):
if type(given) is slice:
if given.step and given.step != 1:
raise IndexError("Slicing with custom step is not supported.")
if given.start >= given.stop:
raise IndexError("Empty slice is not supported.")
if given.start < min or given.stop > max:
raise IndexError(f"Index [{given.start}:{given.stop}] out of"
f" bounds ({min}, {max-1}) for component of"
" direct sum")
elif not (min <= given < max):
raise IndexError(f"Index ({given}) out of bounds ({min}, {max-1})"
" for component of direct sum.")
def _spaces_from_dims(sum_dimensions):
if isinstance(sum_dimensions[0], SumSpace):
to_spaces = sum_dimensions[0].spaces
else:
to_spaces = [sum_dimensions[0]]
if isinstance(sum_dimensions[1], SumSpace):
from_spaces = sum_dimensions[1].spaces
else:
from_spaces = [sum_dimensions[1]]
return to_spaces, from_spaces
def _check_component_index(sum_dimensions, index):
"""
Check that an appropriate number of indices is provided for the given
dimensions. 2 indices is always okay, 1 index is also okay if only one of
the spaces is a sum. Returns tuple (to_index, from_index) of 2 indices.
"""
to_spaces, from_spaces = _spaces_from_dims(sum_dimensions)
is_to_sum = len(to_spaces) > 1
is_from_sum = len(from_spaces) > 1
if is_to_sum and is_from_sum:
if len(index) != 2:
raise IndexError("Invalid number of indices provided for component"
" of direct sum (two indices required).")
else:
if len(index) == 1 and is_to_sum:
index = (index[0], 0)
elif len(index) == 1:
index = (0, index[0])
if len(index) != 2:
raise IndexError("Invalid number of indices provided for component"
" of direct sum (one or two indices required).")
return index
def _component_info(sum_dimensions, to_index, from_index):
"""
For the component (block) indicated by the indices, return
* the Dimensions of the block
* the index of the first row of the block in the full data array
* the index of the row below the block in the full data array
* the index of the first column of the block in the full data array
* the index of the column after the block in the full data array
"""
if isinstance(sum_dimensions[0], SumSpace):
_check_bounds(to_index, 0, len(sum_dimensions[0].spaces))
component_to = sum_dimensions[0].spaces[to_index]
if type(to_index) is slice:
data_row_start = sum_dimensions[0]._space_cumdim(to_index.start)
data_row_stop = sum_dimensions[0]._space_cumdim(to_index.stop)
else:
data_row_start = sum_dimensions[0]._space_cumdim(to_index)
data_row_stop = sum_dimensions[0]._space_cumdim(to_index + 1)
else:
_check_bounds(to_index, 0, 1)
component_to = sum_dimensions[0]
data_row_start = 0
data_row_stop = sum_dimensions[0].size
if type(to_index) is slice:
component_to = SumSpace(component_to)
if isinstance(sum_dimensions[1], SumSpace):
_check_bounds(from_index, 0, len(sum_dimensions[1].spaces))
component_from = sum_dimensions[1].spaces[from_index]
if type(from_index) is slice:
data_col_start = sum_dimensions[1]._space_cumdim(from_index.start)
data_col_stop = sum_dimensions[1]._space_cumdim(from_index.stop)
else:
data_col_start = sum_dimensions[1]._space_cumdim(from_index)
data_col_stop = sum_dimensions[1]._space_cumdim(from_index + 1)
else:
_check_bounds(from_index, 0, 1)
component_from = sum_dimensions[1]
data_col_start = 0
data_col_stop = sum_dimensions[1].size
if type(from_index) is slice:
component_from = SumSpace(component_from)
return (Dimensions(component_from, component_to),
data_row_start, data_row_stop,
data_col_start, data_col_stop)