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