def _lower_contraction_to_addends(cls, expr, contraction_indices):
if isinstance(expr, ArrayAdd):
raise NotImplementedError()
if not isinstance(expr, ArrayTensorProduct):
return expr, contraction_indices
subranks = expr.subranks
cumranks = list(accumulate([0] + subranks))
contraction_indices_remaining = []
contraction_indices_args = [[] for i in expr.args]
backshift = set([])
for i, contraction_group in enumerate(contraction_indices):
for j in range(len(expr.args)):
if not isinstance(expr.args[j], ArrayAdd):
continue
if all(cumranks[j] <= k < cumranks[j+1] for k in contraction_group):
contraction_indices_args[j].append([k - cumranks[j] for k in contraction_group])
backshift.update(contraction_group)
break
else:
contraction_indices_remaining.append(contraction_group)
if len(contraction_indices_remaining) == len(contraction_indices):
return expr, contraction_indices
total_rank = get_rank(expr)
shifts = list(accumulate([1 if i in backshift else 0 for i in range(total_rank)]))
contraction_indices_remaining = [Tuple.fromiter(j - shifts[j] for j in i) for i in contraction_indices_remaining]
ret = ArrayTensorProduct(*[
ArrayContraction(arg, *contr) for arg, contr in zip(expr.args, contraction_indices_args)
])
return ret, contraction_indices_remaining
def _apply_recursively_over_nested_lists(func, arr):
if isinstance(arr, (tuple, list, Tuple)):
return tuple(
_apply_recursively_over_nested_lists(func, i) for i in arr)
elif isinstance(arr, Tuple):
return Tuple.fromiter(
_apply_recursively_over_nested_lists(func, i) for i in arr)
else:
return func(arr)
def test_Tuple():
t = (1, 2, 3, 4)
st = Tuple(*t)
assert set(sympify(t)) == set(st)
assert len(t) == len(st)
assert set(sympify(t[:2])) == set(st[:2])
assert isinstance(st[:], Tuple)
assert st == Tuple(1, 2, 3, 4)
assert st.func(*st.args) == st
p, q, r, s = symbols("p q r s")
t2 = (p, q, r, s)
st2 = Tuple(*t2)
assert st2.atoms() == set(t2)
assert st == st2.subs({p: 1, q: 2, r: 3, s: 4})
# issue 5505
assert all(isinstance(arg, Basic) for arg in st.args)
assert Tuple(p, 1).subs(p, 0) == Tuple(0, 1)
assert Tuple(p, Tuple(p, 1)).subs(p, 0) == Tuple(0, Tuple(0, 1))
assert Tuple(t2) == Tuple(Tuple(*t2))
assert Tuple.fromiter(t2) == Tuple(*t2)
assert Tuple.fromiter(x for x in range(4)) == Tuple(0, 1, 2, 3)
assert st2.fromiter(st2.args) == st2
def test_Tuple():
t = (1, 2, 3, 4)
st = Tuple(*t)
assert set(sympify(t)) == set(st)
assert len(t) == len(st)
assert set(sympify(t[:2])) == set(st[:2])
assert isinstance(st[:], Tuple)
assert st == Tuple(1, 2, 3, 4)
assert st.func(*st.args) == st
p, q, r, s = symbols('p q r s')
t2 = (p, q, r, s)
st2 = Tuple(*t2)
assert st2.atoms() == set(t2)
assert st == st2.subs({p: 1, q: 2, r: 3, s: 4})
# issue 5505
assert all(isinstance(arg, Basic) for arg in st.args)
assert Tuple(p, 1).subs(p, 0) == Tuple(0, 1)
assert Tuple(p, Tuple(p, 1)).subs(p, 0) == Tuple(0, Tuple(0, 1))
assert Tuple(t2) == Tuple(Tuple(*t2))
assert Tuple.fromiter(t2) == Tuple(*t2)
assert Tuple.fromiter(x for x in range(4)) == Tuple(0, 1, 2, 3)
assert st2.fromiter(st2.args) == st2
