def test_two_hs_symbol_sum():
"""Test sum_{ij} a_{ij} Psi_{ij}"""
i = IdxSym('i')
j = IdxSym('j')
a = IndexedBase('a')
hs1 = LocalSpace('1', dimension=3)
hs2 = LocalSpace('2', dimension=3)
hs = hs1 * hs2
Psi = IndexedBase('Psi')
a_ij = a[i, j]
Psi_ij = Psi[i, j]
KetPsi_ij = KetSymbol(StrLabel(Psi_ij), hs=hs)
term = a_ij * KetPsi_ij
expr1 = KetIndexedSum(
term,
ranges=(IndexOverFockSpace(i, hs=hs1), IndexOverFockSpace(j, hs=hs2)),
)
expr2 = KetIndexedSum(term,
ranges=(IndexOverRange(i, 0,
2), IndexOverRange(j, 0, 2)))
assert expr1.term.free_symbols == set(
[i, j, symbols('a'), symbols('Psi'), a_ij, Psi_ij])
assert expr1.free_symbols == set(
[symbols('a'), symbols('Psi'), a_ij, Psi_ij])
assert expr1.variables == [i, j]
assert (
ascii(expr1) == 'Sum_{i in H_1} Sum_{j in H_2} a_ij * |Psi_ij>^(1*2)')
assert unicode(expr1) == '∑_{i ∈ ℌ₁} ∑_{j ∈ ℌ₂} a_ij |Ψ_ij⟩^(1⊗2)'
assert (latex(expr1) ==
r'\sum_{i \in \mathcal{H}_{1}} \sum_{j \in \mathcal{H}_{2}} '
r'a_{i j} \left\lvert \Psi_{i j} \right\rangle^{(1 \otimes 2)}')
assert ascii(expr2) == 'Sum_{i,j=0}^{2} a_ij * |Psi_ij>^(1*2)'
assert unicode(expr2) == '∑_{i,j=0}^{2} a_ij |Ψ_ij⟩^(1⊗2)'
assert (latex(expr2) == r'\sum_{i,j=0}^{2} a_{i j} '
r'\left\lvert \Psi_{i j} \right\rangle^{(1 \otimes 2)}')
assert expr1.doit() == expr2.doit()
assert expr1.doit() == KetPlus(
a[0, 0] * KetSymbol('Psi_00', hs=hs),
a[0, 1] * KetSymbol('Psi_01', hs=hs),
a[0, 2] * KetSymbol('Psi_02', hs=hs),
a[1, 0] * KetSymbol('Psi_10', hs=hs),
a[1, 1] * KetSymbol('Psi_11', hs=hs),
a[1, 2] * KetSymbol('Psi_12', hs=hs),
a[2, 0] * KetSymbol('Psi_20', hs=hs),
a[2, 1] * KetSymbol('Psi_21', hs=hs),
a[2, 2] * KetSymbol('Psi_22', hs=hs),
)
def test_partial_expansion():
"""Test partially executing the sum (only for a subset of summation
indices)"""
i = IdxSym('i')
j = IdxSym('j')
k = IdxSym('k')
hs = LocalSpace('0', dimension=2)
Psi = IndexedBase('Psi')
def r(index_symbol):
return IndexOverFockSpace(index_symbol, hs=hs)
psi_ijk = KetSymbol(StrLabel(Psi[i, j, k]), hs=hs)
def psi(i_val, j_val, k_val):
return psi_ijk.substitute({i: i_val, j: j_val, k: k_val})
expr = KetIndexedSum(psi_ijk, ranges=(r(i), r(j), r(k)))
expr_expanded = expr.doit(indices=[i])
assert expr_expanded == KetIndexedSum(psi(0, j, k) + psi(1, j, k),
ranges=(r(j), r(k)))
expr_expanded = expr.doit(indices=[j])
assert expr_expanded == KetIndexedSum(psi(i, 0, k) + psi(i, 1, k),
ranges=(r(i), r(k)))
assert expr.doit(indices=[j]) == expr.doit(indices=['j'])
expr_expanded = expr.doit(indices=[i, j])
assert expr_expanded == KetIndexedSum(psi(0, 0, k) + psi(1, 0, k) +
psi(0, 1, k) + psi(1, 1, k),
ranges=r(k))
assert expr.doit(indices=[i, j]) == expr.doit(indices=[j, i])
assert expr.doit(indices=[i, j, k]) == expr.doit()
with pytest.raises(ValueError):
expr.doit(indices=[i], max_terms=10)