In sympy, define known actions of operator on kets and use it to simplify

Question

I have a ladder operator â which acts on kets in the following way:

â |n〉= √n |n−1〉
​
â⁺ |n〉= √(n+1) |n+1〉

I want to use these relations to simplify more complicated expressions, for example â⁺â²|n〉. This is my Python code:

import sympy
from sympy import *
from sympy.physics.quantum import *

n = symbols('n')
a = Operator('a')
ad = Dagger(a)

down = Eq( a *Ket(n), sqrt(n  )*Ket(n-1) )
up   = Eq( ad*Ket(n), sqrt(n+1)*Ket(n+1) )

expr = ad * a**2 * Ket(n)

The standard and most basic way to use known equalities in sympy is subs.

But in this case it only works on |n〉, but not on |n+1〉 or other kets containing the symbol n. As can be seen in this example:

In[1]: expr.subs( down.lhs, down.rhs ).subs( down.lhs, down.rhs ).subs( up.lhs, up.rhs )

Out[1]: √n â⁺ â |n−1〉

Answer 1

I achieved the best result by writing a simple algorithm for this myself. The algorithm is very similar to the one in this answer and you can use it like this:

expr.apply_operator(down,up).factor()

Result

√n (n−1) |n−1〉

The algorithm only supports “pure” operators (without any operation) and adjoint operators on the left-hand side. With some minor changes to is_operator and is_expandable_pow_of the algorithm could be expanded to allow inverse operators on the left-hand side, for example.

from sympy.core.operations import AssocOp

def apply_operator(expr, eqns):
    if not isinstance(expr, Basic):
        raise TypeError("The expression to simplify is not a sympy expression.")

    if not isinstance(eqns, list) and not isinstance(eqns, tuple):
        eqns = (eqns,)


    rules = []


    class Rule(object):
        operator = None
        ketSymbol = None
        result = None
        generic = False


    def is_operator(op):
        return isinstance(op, Operator) \
        or isinstance(op, Dagger) \
        and isinstance(op.args[0], Operator)


    for eqn in eqns:
        if not isinstance(eqn, Eq):
            raise TypeError("One of the equations is not a valid sympy equation.")

        lhs = eqn.lhs
        rhs = eqn.rhs

        if not isinstance(lhs, Mul) \
        or len(lhs.args) != 2 \
        or not is_operator(lhs.args[0]) \
        or not isinstance(lhs.args[1], KetBase):
            raise ValueError("The left-hand side has to be an operator applied to a ket.")

        rule = Rule()
        rule.operator = lhs.args[0]
        rule.ketSymbol = lhs.args[1].args[0]
        rule.result = rhs

        if not isinstance(rule.ketSymbol, Symbol):
            raise ValueError("The left-hand ket has to contain a simple symbol.")

        for ket in preorder_traversal(rhs):
            if isinstance(ket, KetBase):
                for symb in preorder_traversal(ket):
                    if symb == rule.ketSymbol:
                        rule.generic = True
                        break

        rules.append(rule)


    def is_expandable_pow_of(base, expr):
        return isinstance(expr, Pow) \
            and base == expr.args[0] \
            and isinstance(expr.args[1], Number) \
            and expr.args[1] >= 1


    def is_ket_of_rule(ket, rule):
        if not isinstance(ket, KetBase):
            return False

        if rule.generic:
            for sym in preorder_traversal(ket):
                if sym == rule.ketSymbol:
                    return True
            return False

        else:
            return ket.args[0] == rule.ketSymbol


    def walk_tree(expr):
        if isinstance(expr, Number):
            return expr

        if not isinstance(expr, AssocOp) and not isinstance(expr, Function):
            return expr.copy()

        elif not isinstance(expr, Mul):
            return expr.func(*(walk_tree(node) for node in expr.args))

        else:
            args = [arg for arg in expr.args]

            for rule in rules:
                A = rule.operator
                ketSym = rule.ketSymbol

                for i in range(len(args)-1):
                    x = args[i]
                    y = args[i+1]

                    if A == x and is_ket_of_rule(y, rule):
                        ev = rule.result

                        if rule.generic:
                            ev = ev.subs(rule.ketSymbol, y.args[0])

                        args = args[0:i] + [ev] + args[i+2:]
                        return walk_tree( Mul(*args).expand() )

                    if is_expandable_pow_of(A, x) and is_ket_of_rule(y, rule):
                        xpow = Pow(A, x.args[1] - 1)
                        ev = rule.result

                        if rule.generic:
                            ev = ev.subs(rule.ketSymbol, y.args[0])

                        args = args[0:i] + [xpow, ev] + args[i+2:]
                        return walk_tree( Mul(*args).expand() )


            return expr.copy()


    return walk_tree(expr)


Basic.apply_operator = lambda self, *eqns: apply_operator(self, eqns)

(No rights reserved.)