1"""Functions for reordering operator expressions."""
2
3import warnings
4
5from sympy import Add, Mul, Pow, Integer
6from sympy.physics.quantum import Operator, Commutator, AntiCommutator
7from sympy.physics.quantum.boson import BosonOp
8from sympy.physics.quantum.fermion import FermionOp
9
10__all__ = [
11    'normal_order',
12    'normal_ordered_form'
13]
14
15
16def _expand_powers(factors):
17    """
18    Helper function for normal_ordered_form and normal_order: Expand a
19    power expression to a multiplication expression so that that the
20    expression can be handled by the normal ordering functions.
21    """
22
23    new_factors = []
24    for factor in factors.args:
25        if (isinstance(factor, Pow)
26                and isinstance(factor.args[1], Integer)
27                and factor.args[1] > 0):
28            for n in range(factor.args[1]):
29                new_factors.append(factor.args[0])
30        else:
31            new_factors.append(factor)
32
33    return new_factors
34
35
36def _normal_ordered_form_factor(product, independent=False, recursive_limit=10,
37                                _recursive_depth=0):
38    """
39    Helper function for normal_ordered_form_factor: Write multiplication
40    expression with bosonic or fermionic operators on normally ordered form,
41    using the bosonic and fermionic commutation relations. The resulting
42    operator expression is equivalent to the argument, but will in general be
43    a sum of operator products instead of a simple product.
44    """
45
46    factors = _expand_powers(product)
47
48    new_factors = []
49    n = 0
50    while n < len(factors) - 1:
51
52        if isinstance(factors[n], BosonOp):
53            # boson
54            if not isinstance(factors[n + 1], BosonOp):
55                new_factors.append(factors[n])
56
57            elif factors[n].is_annihilation == factors[n + 1].is_annihilation:
58                if (independent and
59                        str(factors[n].name) > str(factors[n + 1].name)):
60                    new_factors.append(factors[n + 1])
61                    new_factors.append(factors[n])
62                    n += 1
63                else:
64                    new_factors.append(factors[n])
65
66            elif not factors[n].is_annihilation:
67                new_factors.append(factors[n])
68
69            else:
70                if factors[n + 1].is_annihilation:
71                    new_factors.append(factors[n])
72                else:
73                    if factors[n].args[0] != factors[n + 1].args[0]:
74                        if independent:
75                            c = 0
76                        else:
77                            c = Commutator(factors[n], factors[n + 1])
78                        new_factors.append(factors[n + 1] * factors[n] + c)
79                    else:
80                        c = Commutator(factors[n], factors[n + 1])
81                        new_factors.append(
82                            factors[n + 1] * factors[n] + c.doit())
83                    n += 1
84
85        elif isinstance(factors[n], FermionOp):
86            # fermion
87            if not isinstance(factors[n + 1], FermionOp):
88                new_factors.append(factors[n])
89
90            elif factors[n].is_annihilation == factors[n + 1].is_annihilation:
91                if (independent and
92                        str(factors[n].name) > str(factors[n + 1].name)):
93                    new_factors.append(factors[n + 1])
94                    new_factors.append(factors[n])
95                    n += 1
96                else:
97                    new_factors.append(factors[n])
98
99            elif not factors[n].is_annihilation:
100                new_factors.append(factors[n])
101
102            else:
103                if factors[n + 1].is_annihilation:
104                    new_factors.append(factors[n])
105                else:
106                    if factors[n].args[0] != factors[n + 1].args[0]:
107                        if independent:
108                            c = 0
109                        else:
110                            c = AntiCommutator(factors[n], factors[n + 1])
111                        new_factors.append(-factors[n + 1] * factors[n] + c)
112                    else:
113                        c = AntiCommutator(factors[n], factors[n + 1])
114                        new_factors.append(
115                            -factors[n + 1] * factors[n] + c.doit())
116                    n += 1
117
118        elif isinstance(factors[n], Operator):
119
120            if isinstance(factors[n + 1], (BosonOp, FermionOp)):
121                new_factors.append(factors[n + 1])
122                new_factors.append(factors[n])
123                n += 1
124            else:
125                new_factors.append(factors[n])
126
127        else:
128            new_factors.append(factors[n])
129
130        n += 1
131
132    if n == len(factors) - 1:
133        new_factors.append(factors[-1])
134
135    if new_factors == factors:
136        return product
137    else:
138        expr = Mul(*new_factors).expand()
139        return normal_ordered_form(expr,
140                                   recursive_limit=recursive_limit,
141                                   _recursive_depth=_recursive_depth + 1,
142                                   independent=independent)
143
144
145def _normal_ordered_form_terms(expr, independent=False, recursive_limit=10,
146                               _recursive_depth=0):
147    """
148    Helper function for normal_ordered_form: loop through each term in an
149    addition expression and call _normal_ordered_form_factor to perform the
150    factor to an normally ordered expression.
151    """
152
153    new_terms = []
154    for term in expr.args:
155        if isinstance(term, Mul):
156            new_term = _normal_ordered_form_factor(
157                term, recursive_limit=recursive_limit,
158                _recursive_depth=_recursive_depth, independent=independent)
159            new_terms.append(new_term)
160        else:
161            new_terms.append(term)
162
163    return Add(*new_terms)
164
165
166def normal_ordered_form(expr, independent=False, recursive_limit=10,
167                        _recursive_depth=0):
168    """Write an expression with bosonic or fermionic operators on normal
169    ordered form, where each term is normally ordered. Note that this
170    normal ordered form is equivalent to the original expression.
171
172    Parameters
173    ==========
174
175    expr : expression
176        The expression write on normal ordered form.
177
178    recursive_limit : int (default 10)
179        The number of allowed recursive applications of the function.
180
181    Examples
182    ========
183
184    >>> from sympy.physics.quantum import Dagger
185    >>> from sympy.physics.quantum.boson import BosonOp
186    >>> from sympy.physics.quantum.operatorordering import normal_ordered_form
187    >>> a = BosonOp("a")
188    >>> normal_ordered_form(a * Dagger(a))
189    1 + Dagger(a)*a
190    """
191
192    if _recursive_depth > recursive_limit:
193        warnings.warn("Too many recursions, aborting")
194        return expr
195
196    if isinstance(expr, Add):
197        return _normal_ordered_form_terms(expr,
198                                          recursive_limit=recursive_limit,
199                                          _recursive_depth=_recursive_depth,
200                                          independent=independent)
201    elif isinstance(expr, Mul):
202        return _normal_ordered_form_factor(expr,
203                                           recursive_limit=recursive_limit,
204                                           _recursive_depth=_recursive_depth,
205                                           independent=independent)
206    else:
207        return expr
208
209
210def _normal_order_factor(product, recursive_limit=10, _recursive_depth=0):
211    """
212    Helper function for normal_order: Normal order a multiplication expression
213    with bosonic or fermionic operators. In general the resulting operator
214    expression will not be equivalent to original product.
215    """
216
217    factors = _expand_powers(product)
218
219    n = 0
220    new_factors = []
221    while n < len(factors) - 1:
222
223        if (isinstance(factors[n], BosonOp) and
224                factors[n].is_annihilation):
225            # boson
226            if not isinstance(factors[n + 1], BosonOp):
227                new_factors.append(factors[n])
228            else:
229                if factors[n + 1].is_annihilation:
230                    new_factors.append(factors[n])
231                else:
232                    if factors[n].args[0] != factors[n + 1].args[0]:
233                        new_factors.append(factors[n + 1] * factors[n])
234                    else:
235                        new_factors.append(factors[n + 1] * factors[n])
236                    n += 1
237
238        elif (isinstance(factors[n], FermionOp) and
239              factors[n].is_annihilation):
240            # fermion
241            if not isinstance(factors[n + 1], FermionOp):
242                new_factors.append(factors[n])
243            else:
244                if factors[n + 1].is_annihilation:
245                    new_factors.append(factors[n])
246                else:
247                    if factors[n].args[0] != factors[n + 1].args[0]:
248                        new_factors.append(-factors[n + 1] * factors[n])
249                    else:
250                        new_factors.append(-factors[n + 1] * factors[n])
251                    n += 1
252
253        else:
254            new_factors.append(factors[n])
255
256        n += 1
257
258    if n == len(factors) - 1:
259        new_factors.append(factors[-1])
260
261    if new_factors == factors:
262        return product
263    else:
264        expr = Mul(*new_factors).expand()
265        return normal_order(expr,
266                            recursive_limit=recursive_limit,
267                            _recursive_depth=_recursive_depth + 1)
268
269
270def _normal_order_terms(expr, recursive_limit=10, _recursive_depth=0):
271    """
272    Helper function for normal_order: look through each term in an addition
273    expression and call _normal_order_factor to perform the normal ordering
274    on the factors.
275    """
276
277    new_terms = []
278    for term in expr.args:
279        if isinstance(term, Mul):
280            new_term = _normal_order_factor(term,
281                                            recursive_limit=recursive_limit,
282                                            _recursive_depth=_recursive_depth)
283            new_terms.append(new_term)
284        else:
285            new_terms.append(term)
286
287    return Add(*new_terms)
288
289
290def normal_order(expr, recursive_limit=10, _recursive_depth=0):
291    """Normal order an expression with bosonic or fermionic operators. Note
292    that this normal order is not equivalent to the original expression, but
293    the creation and annihilation operators in each term in expr is reordered
294    so that the expression becomes normal ordered.
295
296    Parameters
297    ==========
298
299    expr : expression
300        The expression to normal order.
301
302    recursive_limit : int (default 10)
303        The number of allowed recursive applications of the function.
304
305    Examples
306    ========
307
308    >>> from sympy.physics.quantum import Dagger
309    >>> from sympy.physics.quantum.boson import BosonOp
310    >>> from sympy.physics.quantum.operatorordering import normal_order
311    >>> a = BosonOp("a")
312    >>> normal_order(a * Dagger(a))
313    Dagger(a)*a
314    """
315    if _recursive_depth > recursive_limit:
316        warnings.warn("Too many recursions, aborting")
317        return expr
318
319    if isinstance(expr, Add):
320        return _normal_order_terms(expr, recursive_limit=recursive_limit,
321                                   _recursive_depth=_recursive_depth)
322    elif isinstance(expr, Mul):
323        return _normal_order_factor(expr, recursive_limit=recursive_limit,
324                                    _recursive_depth=_recursive_depth)
325    else:
326        return expr
327