More efficient wicks implementation

adcgen/__init__.py

@@ -3,7 +3,7 @@
 from .eri_orbenergy import EriOrbenergy
 from .expression import ExprContainer, import_from_sympy_latex
 from .factor_intermediates import factor_intermediates
-from .func import evaluate_deltas, wicks
+from .func import evaluate_deltas
 from .generate_code import (generate_code, optimize_contractions, Contraction,
                             unoptimized_contraction)
 from .groundstate import GroundState
@@ -21,6 +21,7 @@
                             NonSymmetricTensor, KroneckerDelta, SymbolicTensor)
 from .tensor_names import tensor_names
 from .resolution_of_identity import apply_resolution_of_identity
+from .wicks import wicks
 from . import sort_expr as sort
 
 

adcgen/func.py

@@ -1,15 +1,10 @@
 from collections.abc import Sequence
 import itertools
 
-from sympy.physics.secondquant import (
-    F, Fd, FermionicOperator, NO
-)
-from sympy import S, Add, Expr, Mul
+from sympy import Add, Expr, Mul
 
-from .expression import ExprContainer
 from .misc import Inputerror
-from .rules import Rules
-from .indices import Index, Indices, get_symbols
+from .indices import Index, get_symbols
 from .sympy_objects import KroneckerDelta
 
 
@@ -127,173 +122,3 @@ def evaluate_deltas(
         return expr
     else:
         return expr
-
-
-def wicks(expr: Expr, rules: Rules | None = None,
-          simplify_kronecker_deltas: bool = False) -> Expr:
-    """
-    Evaluates Wicks theorem in the provided expression only returning fully
-    contracted contributions.
-    Adapted from the implementation in 'sympy.physics.secondquant'.
-
-    Parameters
-    ----------
-    expr: Expr
-        Expression containing the second quantized operator strings to
-        evaluate. This function expects plain sympy objects (Add/Mul/...)
-        and no container class.
-    rules : Rules, optional
-        Rules that are applied to the result before returning, e.g., in the
-        context of RE not all tensor blocks might be allowed in the result.
-    simplify_kronecker_deltas : bool, optional
-        If set, the KroneckerDeltas generated through the contractions
-        will be evaluated before returning.
-    """
-    assert isinstance(expr, Expr)
-    # normal ordered operator string has to evaluate to zero
-    # and a single second quantized operator can not be contracted
-    if isinstance(expr, (NO, FermionicOperator)):
-        return S.Zero
-
-    # break up any NO-objects, and evaluate commutators
-    expr = expr.doit(wicks=True).expand()
-    assert isinstance(expr, Expr)
-
-    if isinstance(expr, Add):
-        return Add(*(
-            wicks(term, rules=rules,
-                  simplify_kronecker_deltas=simplify_kronecker_deltas)
-            for term in expr.args
-        ))
-    elif not isinstance(expr, Mul):
-        # nether Add, Mul, NO, F, Fd -> maybe a number or tensor
-        return expr
-    # -> we have a Mul object
-    # we don't want to mess around with commuting part of Mul
-    # so we factorize it out before starting recursion
-    c_part: list[Expr] = []
-    op_string: list[FermionicOperator] = []
-    for factor in expr.args:
-        if factor.is_commutative:
-            c_part.append(factor)
-        else:
-            assert isinstance(factor, FermionicOperator)
-            op_string.append(factor)
-
-    if (n := len(op_string)) == 0:  # no operators
-        result = expr
-    elif n == 1:  # a single operator
-        return S.Zero
-    else:  # at least 2 operators
-        result = _contract_operator_string(op_string)
-        result = Mul(*c_part, result).expand()
-        assert isinstance(result, Expr)
-        if simplify_kronecker_deltas:
-            result = evaluate_deltas(result)
-
-    # apply rules to the result
-    if rules is None:
-        return result
-    assert isinstance(rules, Rules)
-    return rules.apply(ExprContainer(result)).inner
-
-
-def _contract_operator_string(op_string: list[FermionicOperator]) -> Expr:
-    """
-    Contracts the operator string only returning fully contracted
-    contritbutions.
-    Adapted from 'sympy.physics.secondquant'.
-    """
-    # check that we can get a fully contracted contribution
-    if not _has_fully_contracted_contribution(op_string):
-        return S.Zero
-
-    result = []
-    for i in range(1, len(op_string)):
-        c = _contraction(op_string[0], op_string[i])
-        if c is S.Zero:
-            continue
-        if not i % 2:  # introduce -1 for swapping operators
-            c *= S.NegativeOne
-
-        if len(op_string) - 2 > 0:  # at least one operator left
-            # remove the contracted operators from the string and recurse
-            remaining = op_string[1:i] + op_string[i+1:]
-            result.append(Mul(c, _contract_operator_string(remaining)))
-        else:  # no operators left
-            result.append(c)
-    return Add(*result)
-
-
-def _contraction(p: FermionicOperator, q: FermionicOperator) -> Expr:
-    """
-    Evaluates the contraction between two sqcond quantized fermionic
-    operators.
-    Adapted from 'sympy.physics.secondquant'.
-    """
-    assert isinstance(p, FermionicOperator)
-    assert isinstance(q, FermionicOperator)
-
-    p_idx, q_idx = p.args[0], q.args[0]
-    assert isinstance(p_idx, Index) and isinstance(q_idx, Index)
-    if p_idx.spin or q_idx.spin:
-        raise NotImplementedError("Contraction not implemented for indices "
-                                  "with spin.")
-    # get the space and ensure we have no unexpected space
-    space_p, space_q = p_idx.space[0], q_idx.space[0]
-    assert space_p in ["o", "v", "g"] and space_q in ["o", "v", "g"]
-
-    if isinstance(p, F) and isinstance(q, Fd):
-        if space_p == "o" or space_q == "o":
-            res = S.Zero
-        elif space_p == "v" or space_q == "v":
-            res = KroneckerDelta(p_idx, q_idx)
-        else:
-            res = Mul(
-                KroneckerDelta(p_idx, q_idx),
-                KroneckerDelta(q_idx, Index('a', above_fermi=True))
-            )
-    elif isinstance(p, Fd) and isinstance(q, F):
-        if space_p == "v" or space_q == "v":
-            res = S.Zero
-        elif space_p == "o" or space_q == "o":
-            res = KroneckerDelta(p_idx, q_idx)
-        else:
-            res = Mul(
-                KroneckerDelta(p_idx, q_idx),
-                KroneckerDelta(q_idx, Index('i', below_fermi=True))
-            )
-    else:  # vanish if 2xAnnihilator or 2xCreator
-        res = S.Zero
-    assert isinstance(res, Expr)
-    return res
-
-
-def _has_fully_contracted_contribution(op_string: list[FermionicOperator]
-                                       ) -> bool:
-    """
-    Takes a list of second quantized operators and checks whether a
-    non-vanishing fully contracted contribution can exist.
-    """
-    if len(op_string) % 2:  # odd number of operators
-        return False
-    # count the number of creation and annihilation operators per space
-    create = {space: 0 for space in Indices.base}
-    annihilate = {space: 0 for space in Indices.base}
-    for op in op_string:
-        if isinstance(op, Fd):
-            counter = create
-        else:
-            counter = annihilate
-        idx = op.args[0]
-        assert isinstance(idx, Index)
-        counter[idx.space] += 1
-    # check that we have a matching amount of creation and annihilation
-    # operators
-    for space, n_create in create.items():
-        if space == "general":
-            continue
-        n_annihilate = annihilate[space] + annihilate["general"]
-        if n_create - n_annihilate > 0:
-            return False
-    return True

adcgen/groundstate.py

@@ -4,14 +4,15 @@
 from sympy import Expr, Mul, Rational, S, latex
 
 from .expression import ExprContainer
-from .func import gen_term_orders, wicks
+from .func import gen_term_orders
 from .indices import Indices, n_ov_from_space
 from .logger import logger
 from .misc import cached_member, Inputerror, validate_input
 from .operators import Operators
 from .simplify import simplify
 from .sympy_objects import Amplitude
 from .tensor_names import tensor_names
+from .wicks import wicks
 
 
 class GroundState:

adcgen/intermediate_states.py

@@ -5,7 +5,7 @@
 from sympy import Expr, Mul, Rational, S, latex, nsimplify, diff, symbols
 
 from .expression import ExprContainer
-from .func import gen_term_orders, wicks, evaluate_deltas
+from .func import gen_term_orders, evaluate_deltas
 from .groundstate import GroundState
 from .indices import (
     n_ov_from_space, repeated_indices, Indices, generic_indices_from_space
@@ -15,6 +15,7 @@
 from .simplify import simplify
 from .sympy_objects import Amplitude
 from .tensor_names import tensor_names
+from .wicks import wicks
 
 
 class IntermediateStates:

adcgen/properties.py

@@ -4,13 +4,14 @@
 from sympy import Add, Expr, S, sqrt, sympify
 
 from .expression import ExprContainer
-from .func import gen_term_orders, wicks
+from .func import gen_term_orders
 from .indices import n_ov_from_space, generic_indices_from_space
 from .intermediate_states import IntermediateStates
 from .misc import Inputerror, cached_member, transform_to_tuple, validate_input
 from .rules import Rules
 from .secular_matrix import SecularMatrix
 from .simplify import simplify
+from .wicks import wicks
 
 
 class Properties:

adcgen/secular_matrix.py

@@ -4,7 +4,7 @@
 from sympy import Add, Expr, Mul, S, sqrt
 
 from .expression import ExprContainer
-from .func import gen_term_orders, wicks, evaluate_deltas
+from .func import gen_term_orders, evaluate_deltas
 from .groundstate import GroundState
 from .indices import (
     repeated_indices, Indices, generic_indices_from_space, n_ov_from_space
@@ -14,6 +14,7 @@
 from .operators import Operators
 from .rules import Rules
 from .simplify import simplify
+from .wicks import wicks
 
 
 class SecularMatrix: