<commit_before># This file is part of QNET.
#
# QNET is free software: you can redistribute it and/or modify
#    it under the terms of the GNU General Public License as published by
#    the Free Software Foundation, either version 3 of the License, or
#   (at your option) any later version.
#
#    QNET is distributed in the hope that it will be useful,
#    but WITHOUT ANY WARRANTY; without even the implied warranty of
#    MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
#    GNU General Public License for more details.
#
#    You should have received a copy of the GNU General Public License
#    along with QNET.  If not, see <http://www.gnu.org/licenses/>.
#
# Copyright (C) 2012-2017, QNET authors (see AUTHORS file)
#
###########################################################################
r"""
This module features classes and functions to define and manipulate symbolic
Operator expressions.  For more details see :ref:`operator_algebra`.

For a list of all properties and methods of an operator object, see the
documentation for the basic :py:class:`Operator` class.
"""
import re
from abc import ABCMeta, abstractmethod, abstractproperty
from collections import defaultdict, OrderedDict
from itertools import product as cartesian_product

from sympy import (exp, sqrt, I, sympify, Basic as SympyBasic,
        series as sympy_series, Add as SympyAdd)

from .scalar_types import SCALAR_TYPES
from .abstract_algebra import (
        Expression, Operation, assoc, orderby, filter_neutral,
        match_replace_binary, match_replace, set_union, substitute,
        CannotSimplify, cache_attr)
from .singleton import Singleton, singleton_object
from .hilbert_space_algebra import (
        TrivialSpace, HilbertSpace, LocalSpace, ProductSpace, BasisNotSetError,
        FullSpace)
from .pattern_matching import wc, pattern_head, pattern
from .ordering import (
        KeyTuple, expr_order_key, DisjunctCommutativeHSOrder,
        FullCommutativeHSOrder)
from ..printing import ascii, srepr

sympyOne = sympify(1)

# for hilbert space dimensions less than or equal to this,
# compute numerically PseudoInverse and NullSpaceProjector representations
DENSE_DIMENSION_LIMIT = 1000


__all__ = [
    'Adjoint', 'Create', 'Destroy', 'Displace', 'Jminus', 'Jplus',
    'Jz', 'LocalOperator', 'LocalSigma', 'NullSpaceProjector', 'Operator',
    'OperatorOperation', 'OperatorPlus', 'OperatorPlusMinusCC',
    'OperatorSymbol', 'OperatorTimes', 'OperatorTrace',
    'Phase', 'PseudoInverse', 'ScalarTimesOperator', 'SingleOperatorOperation',
    'Squeeze', 'Jmjmcoeff', 'Jpjmcoeff', 'Jzjmcoeff', 'LocalProjector',
    'X', 'Y', 'Z', 'adjoint', 'create_operator_pm_cc', 'decompose_space',
    'expand_operator_pm_cc', 'factor_coeff', 'factor_for_trace', 'get_coeffs',
    'scalar_free_symbols', 'simplify_scalar', 'space', 'II',
    'IdentityOperator', 'ZeroOperator']

__private__ = [  # anything not in __all__ must be in __private__
    'implied_local_space',  'delegate_to_method']



###############################################################################
# Algebraic properties
###############################################################################


def implied_local_space(*, arg_index=None, keys=None):
    """Return a simplification that converts the positional argument
    `arg_index` from (str, int) to LocalSpace, as well as any keyword argument
    with one of the given keys"""

    def args_to_local_space(cls, args, kwargs):
        """Convert (str, int) of selected args to LocalSpace"""
        if isinstance(args[arg_index], LocalSpace):
            new_args = args
        else:
            if isinstance(args[arg_index], (int, str)):
                hs = LocalSpace(args[arg_index])
            else:
                hs = args[arg_index]
                assert isinstance(hs, HilbertSpace)
            new_args = (tuple(args[:arg_index]) + (hs, ) +
                        tuple(args[arg_index+1:]))
        return new_args, kwargs

    def kwargs_to_local_space(cls, args, kwargs):
        """Convert (str, int) of selected kwargs to LocalSpace"""
        if all([isinstance(kwargs[key], LocalSpace) for key in keys]):
            new_kwargs = kwargs
        else:
            new_kwargs = {}
            for key, val in kwargs.items():
                if key in keys:
                    if isinstance(val, (int, str)):
                        val = LocalSpace(val)
                    assert isinstance(val, HilbertSpace)
                new_kwargs[key] = val
        return args, new_kwargs

    def to_local_space(cls, args, kwargs):
        """Convert (str, int) of selected args and kwargs to LocalSpace"""
        new_args, __ = args_to_local_space(args, kwargs, arg_index)
        __, new_kwargs = kwargs_to_local_space(args, kwargs, keys)
        return new_args, new_kwargs

    if (arg_index is not None) and (keys is None):
        return args_to_local_space
    elif (arg_index is None) and (keys is not None):
        return kwargs_to_local_space
    elif (arg_index is not None) and (keys is not None):
        return to_local_space
    else:
        raise ValueError("must give at least one of arg_index and keys")


def delegate_to_method(mtd):
    """Create a simplification rule that delegates the instantiation to the
    method `mtd` of the operand (if defined)"""
    def _delegate_to_method(cls, ops, kwargs):
        assert len(ops) == 1
        op, = ops
        if hasattr(op, mtd):
            return getattr(op, mtd)()
        else:
            return ops, kwargs
    return _delegate_to_method


###############################################################################
# Abstract base classes
###############################################################################


class Operator(metaclass=ABCMeta):
    """The basic operator class, which fixes the abstract interface of operator
    objects and where possible also defines the default behavior under
    operations.  Any operator contains an associated HilbertSpace object,
    on which it is taken to act non-trivially.
    """

    @abstractproperty
    def space(self):
        """The :class:`HilbertSpace` on which the operator acts
        non-trivially"""
        raise NotImplementedError(self.__class__.__name__)

    def adjoint(self):
        """The Hermitian adjoint of the operator."""
        return Adjoint.create(self)

    conjugate = dag = adjoint

    def pseudo_inverse(self):
        """The pseudo-Inverse of the Operator, i.e., it inverts the operator on
        the orthogonal complement of its nullspace"""
        return PseudoInverse.create(self)

    def expand(self):
        """Expand out distributively all products of sums. Note that this does
        not expand out sums of scalar coefficients.

        :return: A fully expanded sum of operators.
        :rtype: Operator
        """
        return self._expand()

    @abstractmethod
    def _expand(self):
        raise NotImplementedError(self.__class__.__name__)

    def simplify_scalar(self):
        """Simplify all scalar coefficients within the Operator expression.

        :return: The simplified expression.
        :rtype: Operator
        """
        return self._simplify_scalar()

    def _simplify_scalar(self):
        return self

    def diff(self, sym, n=1, expand_simplify=True):
        """Differentiate by scalar parameter sym.

        :param ~sympy.cor.symbol.Symbol sym: What to differentiate by.
        :param int n: How often to differentiate
        :param bool expand_simplify: Whether to simplify the result.
        :return (Operator): The n-th derivative.<commit_msg>Add expand_in_basis method to Operators<commit_after># This file is part of QNET.
#
# QNET is free software: you can redistribute it and/or modify
#    it under the terms of the GNU General Public License as published by
#    the Free Software Foundation, either version 3 of the License, or
#   (at your option) any later version.
#
#    QNET is distributed in the hope that it will be useful,
#    but WITHOUT ANY WARRANTY; without even the implied warranty of
#    MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
#    GNU General Public License for more details.
#
#    You should have received a copy of the GNU General Public License
#    along with QNET.  If not, see <http://www.gnu.org/licenses/>.
#
# Copyright (C) 2012-2017, QNET authors (see AUTHORS file)
#
###########################################################################
r"""
This module features classes and functions to define and manipulate symbolic
Operator expressions.  For more details see :ref:`operator_algebra`.

For a list of all properties and methods of an operator object, see the
documentation for the basic :py:class:`Operator` class.
"""
import re
from abc import ABCMeta, abstractmethod, abstractproperty
from collections import defaultdict, OrderedDict
from itertools import product as cartesian_product

from sympy import (exp, sqrt, I, sympify, Basic as SympyBasic,
        series as sympy_series, Add as SympyAdd)

from .scalar_types import SCALAR_TYPES
from .abstract_algebra import (
        Expression, Operation, assoc, orderby, filter_neutral,
        match_replace_binary, match_replace, set_union, substitute,
        CannotSimplify, cache_attr)
from .singleton import Singleton, singleton_object
from .hilbert_space_algebra import (
        TrivialSpace, HilbertSpace, LocalSpace, ProductSpace, BasisNotSetError,
        FullSpace)
from .pattern_matching import wc, pattern_head, pattern
from .ordering import (
        KeyTuple, expr_order_key, DisjunctCommutativeHSOrder,
        FullCommutativeHSOrder)
from ..printing import ascii, srepr

sympyOne = sympify(1)

# for hilbert space dimensions less than or equal to this,
# compute numerically PseudoInverse and NullSpaceProjector representations
DENSE_DIMENSION_LIMIT = 1000


__all__ = [
    'Adjoint', 'Create', 'Destroy', 'Displace', 'Jminus', 'Jplus',
    'Jz', 'LocalOperator', 'LocalSigma', 'NullSpaceProjector', 'Operator',
    'OperatorOperation', 'OperatorPlus', 'OperatorPlusMinusCC',
    'OperatorSymbol', 'OperatorTimes', 'OperatorTrace',
    'Phase', 'PseudoInverse', 'ScalarTimesOperator', 'SingleOperatorOperation',
    'Squeeze', 'Jmjmcoeff', 'Jpjmcoeff', 'Jzjmcoeff', 'LocalProjector',
    'X', 'Y', 'Z', 'adjoint', 'create_operator_pm_cc', 'decompose_space',
    'expand_operator_pm_cc', 'factor_coeff', 'factor_for_trace', 'get_coeffs',
    'scalar_free_symbols', 'simplify_scalar', 'space', 'II',
    'IdentityOperator', 'ZeroOperator']

__private__ = [  # anything not in __all__ must be in __private__
    'implied_local_space',  'delegate_to_method']



###############################################################################
# Algebraic properties
###############################################################################


def implied_local_space(*, arg_index=None, keys=None):
    """Return a simplification that converts the positional argument
    `arg_index` from (str, int) to LocalSpace, as well as any keyword argument
    with one of the given keys"""

    def args_to_local_space(cls, args, kwargs):
        """Convert (str, int) of selected args to LocalSpace"""
        if isinstance(args[arg_index], LocalSpace):
            new_args = args
        else:
            if isinstance(args[arg_index], (int, str)):
                hs = LocalSpace(args[arg_index])
            else:
                hs = args[arg_index]
                assert isinstance(hs, HilbertSpace)
            new_args = (tuple(args[:arg_index]) + (hs, ) +
                        tuple(args[arg_index+1:]))
        return new_args, kwargs

    def kwargs_to_local_space(cls, args, kwargs):
        """Convert (str, int) of selected kwargs to LocalSpace"""
        if all([isinstance(kwargs[key], LocalSpace) for key in keys]):
            new_kwargs = kwargs
        else:
            new_kwargs = {}
            for key, val in kwargs.items():
                if key in keys:
                    if isinstance(val, (int, str)):
                        val = LocalSpace(val)
                    assert isinstance(val, HilbertSpace)
                new_kwargs[key] = val
        return args, new_kwargs

    def to_local_space(cls, args, kwargs):
        """Convert (str, int) of selected args and kwargs to LocalSpace"""
        new_args, __ = args_to_local_space(args, kwargs, arg_index)
        __, new_kwargs = kwargs_to_local_space(args, kwargs, keys)
        return new_args, new_kwargs

    if (arg_index is not None) and (keys is None):
        return args_to_local_space
    elif (arg_index is None) and (keys is not None):
        return kwargs_to_local_space
    elif (arg_index is not None) and (keys is not None):
        return to_local_space
    else:
        raise ValueError("must give at least one of arg_index and keys")


def delegate_to_method(mtd):
    """Create a simplification rule that delegates the instantiation to the
    method `mtd` of the operand (if defined)"""
    def _delegate_to_method(cls, ops, kwargs):
        assert len(ops) == 1
        op, = ops
        if hasattr(op, mtd):
            return getattr(op, mtd)()
        else:
            return ops, kwargs
    return _delegate_to_method


###############################################################################
# Abstract base classes
###############################################################################


class Operator(metaclass=ABCMeta):
    """The basic operator class, which fixes the abstract interface of operator
    objects and where possible also defines the default behavior under
    operations.  Any operator contains an associated HilbertSpace object,
    on which it is taken to act non-trivially.
    """

    @abstractproperty
    def space(self):
        """The :class:`HilbertSpace` on which the operator acts
        non-trivially"""
        raise NotImplementedError(self.__class__.__name__)

    def adjoint(self):
        """The Hermitian adjoint of the operator."""
        return Adjoint.create(self)

    conjugate = dag = adjoint

    def pseudo_inverse(self):
        """The pseudo-Inverse of the Operator, i.e., it inverts the operator on
        the orthogonal complement of its nullspace"""
        return PseudoInverse.create(self)

    def expand(self):
        """Expand out distributively all products of sums. Note that this does
        not expand out sums of scalar coefficients.

        :return: A fully expanded sum of operators.
        :rtype: Operator
        """
        return self._expand()

    def expand_in_basis(self, basis_states=None, hermitian=False):
        """Write the operator as an expansion into all :class:`KetBras
        <qnet.algebra.state_algebra.KetBra>` spanned by `basis_states`.

        Args:
            basis_states (list or None): List of basis states
                (:class:`~qnet.algebra.state_algebra.Ket` instances) into which
                to expand the operator. If None, use the operator's
                `space.basis_states`
            hermitian (bool): If True, assume that the operator is Hermitian
                and represent all elements in the lower triangle of the
                expansion via :class:`OperatorPlusMinusCC`. This is meant to
                enhance readability

        Raises:
            BasisNotSetError: If `basis_states` is None and the operator's
                Hilbert space has no well-defined basis

        Example:

            >>> hs = LocalSpace(1, basis=('g', 'e'))
            >>> op = LocalSigma('g', 'e', hs=hs) + LocalSigma('e', 'g', hs=hs)
            >>> print(ascii(op))
            sigma_e,g^(1) + sigma_g,e^(1)
            >>> print(ascii(op.expand_in_basis()))
            |e><g|_(1) + |g><e|_(1)
            >>> print(ascii(op.expand_in_basis(hermitian=True)))
            |g><e|_(1) + c.c.
        """
        from qnet.algebra.state_algebra import KetBra  # avoid circ. import
        if basis_states is None:
            basis_states = list(self.space.basis_states)
        else:
            basis_states = list(basis_states)
        diag_terms = []
        terms = []
        for i, ket_i in enumerate(basis_states):
            for j, ket_j in enumerate(basis_states):
                if i > j and hermitian:
                    continue
                op_ij = (ket_i.dag * self * ket_j).expand()
                ketbra = KetBra(ket_i, ket_j)
                term = op_ij * ketbra
                if term is not ZeroOperator:
                    if i == j:
                        diag_terms.append(op_ij * ketbra)
                    else:
                        terms.append(op_ij * ketbra)
        if hermitian:
            res = OperatorPlus(*diag_terms)
            if len(terms) > 0:
                res = res + OperatorPlusMinusCC(OperatorPlus(*terms))
            return res
        else:
            return OperatorPlus(*diag_terms) + OperatorPlus(*terms)

    @abstractmethod
    def _expand(self):
        raise NotImplementedError(self.__class__.__name__)

    def simplify_scalar(self):
        """Simplify all scalar coefficients within the Operator expression.

        :return: The simplified expression.
        :rtype: Operator
        """
        return self._simplify_scalar()

    def _simplify_scalar(self):
        return self

    def diff(self, sym, n=1, expand_simplify=True):
        """Differentiate by scalar parameter sym.

        :param ~sympy.cor.symbol.Symbol sym: What to differentiate by.
        :param int n: How often to differentiate
        :param bool expand_simplify: Whether to simplify the result.
        :return (Operator): The n-th derivative.
{'fraction_non_alphanumeric': 0.12823702363451317, 'fraction_numerical': 0.0023182177993893477, 'mean_word_length': 2.9880496054114993, 'pattern_counts': {'":': 0, '<': 9, '<?xml version=': 0, '>': 26, 'https://': 0, 'lorem ipsum': 0, 'www.': 2, 'xml': 0}, 'pii_count': 0, 'substrings_counts': 0, 'word_list_counts': {'cursed_substrings.json': 0, 'profanity_word_list.json': 0, 'sexual_word_list.json': 0, 'zh_pornsignals.json': 0}}
{'dir': 'git-commits-cleaned', 'id': '646933', 'n_tokens_mistral': 4981, 'n_tokens_neox': 4698, 'n_words': 1617}