diff --git a/doc/changes/2545.feature b/doc/changes/2545.feature new file mode 100644 index 0000000000..3514cff823 --- /dev/null +++ b/doc/changes/2545.feature @@ -0,0 +1 @@ +Implements a numpy.einsum version for Qobj dimensions (Evaluates the Einstein summation convention on the operands.) \ No newline at end of file diff --git a/doc/changes/2557.misc b/doc/changes/2557.misc new file mode 100644 index 0000000000..4f470efc77 --- /dev/null +++ b/doc/changes/2557.misc @@ -0,0 +1 @@ +Improve performance of qt.Qobj by using static numpy version check diff --git a/qutip/core/data/dense.pyx b/qutip/core/data/dense.pyx index 22c4360aa1..120bd032a0 100644 --- a/qutip/core/data/dense.pyx +++ b/qutip/core/data/dense.pyx @@ -37,10 +37,11 @@ __all__ = [ class OrderEfficiencyWarning(EfficiencyWarning): pass +is_numpy1 = np.lib.NumpyVersion(np.__version__) < '2.0.0b1' cdef class Dense(base.Data): def __init__(self, data, shape=None, copy=True): - if np.lib.NumpyVersion(np.__version__) < '2.0.0b1': + if is_numpy1: # np2 accept None which act as np1's False copy = builtins.bool(copy) base = np.array(data, dtype=np.complex128, order='K', copy=copy) diff --git a/qutip/core/dimensions.py b/qutip/core/dimensions.py index f00afc1ff8..56fe353964 100644 --- a/qutip/core/dimensions.py +++ b/qutip/core/dimensions.py @@ -352,6 +352,34 @@ def _frozen(*args, **kwargs): raise RuntimeError("Dimension cannot be modified.") +def einsum(subscripts, *operands): + """ + Implementation of numpy.einsum for Qobj. + Evaluates the Einstein summation convention on the operands. + Parameters + ---------- + subscripts: str + Specifies the subscripts for summation as comma + separated list of subscript labels. + operands: list of array_like + These are the arrays for the operation. + + Returns + ------- + Qobj (numpy.complex128) + Result of einsum as Qobj (numpy.complex128 if result is scalar) + """ + operands_array = [to_tensor_rep(op) for op in operands] + result = np.einsum(subscripts, *operands_array) + if result.shape == (): + return result + dims = [ + [d for d in result.shape[:result.ndim // 2]], + [d for d in result.shape[result.ndim // 2:]] + ] + return from_tensor_rep(result, dims) + + class MetaSpace(type): def __call__(cls, *args: SpaceLike, rep: str = None) -> "Space": """ diff --git a/qutip/tests/core/test_dimensions.py b/qutip/tests/core/test_dimensions.py index 7be9d27600..96b2ed4cde 100644 --- a/qutip/tests/core/test_dimensions.py +++ b/qutip/tests/core/test_dimensions.py @@ -6,7 +6,7 @@ from qutip.core.dimensions import ( flatten, unflatten, enumerate_flat, deep_remove, deep_map, dims_idxs_to_tensor_idxs, dims_to_tensor_shape, dims_to_tensor_perm, - collapse_dims_super, collapse_dims_oper, Dimensions + collapse_dims_super, collapse_dims_oper, einsum, Dimensions ) @@ -189,3 +189,17 @@ def test_dims_comparison(): assert Dimensions([[1], [2]])[0] != Dimensions([[1], [2]])[1] assert not Dimensions([[1], [2]])[1] != Dimensions([[1], [2]])[1] assert not Dimensions([[1], [2]])[0] != Dimensions([[1], [2]])[0] + + +@pytest.mark.parametrize(["subscripts", "operands", "expected"], [ + pytest.param("ii", [qutip.sigmaz()], 0), + pytest.param("ij", [qutip.sigmax()], qutip.sigmax()), + pytest.param("ij->ji", [qutip.sigmay()], qutip.sigmay().trans()), + pytest.param("ij,ji", [qutip.sigmaz(), qutip.sigmaz()], 2), + pytest.param("ijij", [qutip.tensor(qutip.thermal_dm(2,1), qutip.thermal_dm(2,1))], 1), + pytest.param("ikjl,jm->ikml", [qutip.tensor(qutip.sigmaz(), qutip.sigmaz()), + qutip.sigmaz()], qutip.tensor(qutip.qeye(2), qutip.sigmaz())), + pytest.param("ijkl->kjil", [qutip.tensor(qutip.sigmam(), qutip.sigmaz())], qutip.tensor(qutip.sigmap(), qutip.sigmaz())) +]) +def test_einsum(subscripts, operands, expected): + assert einsum(subscripts, *operands) == expected \ No newline at end of file