def _array_diag2contr_diagmatrix(expr: ArrayDiagonal):
if isinstance(expr.expr, ArrayTensorProduct):
args = list(expr.expr.args)
diag_indices = list(expr.diagonal_indices)
mapping = _get_mapping_from_subranks(
[_get_subrank(arg) for arg in args])
tuple_links = [[mapping[j] for j in i] for i in diag_indices]
contr_indices = []
total_rank = get_rank(expr)
replaced = [False for arg in args]
for i, (abs_pos, rel_pos) in enumerate(zip(diag_indices, tuple_links)):
if len(abs_pos) != 2:
continue
(pos1_outer, pos1_inner), (pos2_outer, pos2_inner) = rel_pos
arg1 = args[pos1_outer]
arg2 = args[pos2_outer]
if get_rank(arg1) != 2 or get_rank(arg2) != 2:
if replaced[pos1_outer]:
diag_indices[i] = None
if replaced[pos2_outer]:
diag_indices[i] = None
continue
pos1_in2 = 1 - pos1_inner
pos2_in2 = 1 - pos2_inner
if arg1.shape[pos1_in2] == 1:
if arg1.shape[pos1_inner] != 1:
darg1 = DiagMatrix(arg1)
else:
darg1 = arg1
args.append(darg1)
contr_indices.append(
((pos2_outer, pos2_inner), (len(args) - 1, pos1_inner)))
total_rank += 1
diag_indices[i] = None
args[pos1_outer] = OneArray(arg1.shape[pos1_in2])
replaced[pos1_outer] = True
elif arg2.shape[pos2_in2] == 1:
if arg2.shape[pos2_inner] != 1:
darg2 = DiagMatrix(arg2)
else:
darg2 = arg2
args.append(darg2)
contr_indices.append(
((pos1_outer, pos1_inner), (len(args) - 1, pos2_inner)))
total_rank += 1
diag_indices[i] = None
args[pos2_outer] = OneArray(arg2.shape[pos2_in2])
replaced[pos2_outer] = True
diag_indices_new = [i for i in diag_indices if i is not None]
cumul = list(accumulate([0] + [get_rank(arg) for arg in args]))
contr_indices2 = [
tuple(cumul[a] + b for a, b in i) for i in contr_indices
]
tc = _array_contraction(_array_tensor_product(*args), *contr_indices2)
td = _array_diagonal(tc, *diag_indices_new)
return td
return expr
def __new__(cls, expr, *contraction_indices, **kwargs):
contraction_indices = _sort_contraction_indices(contraction_indices)
expr = _sympify(expr)
if len(contraction_indices) == 0:
return expr
if isinstance(expr, ArrayContraction):
return cls._flatten(expr, *contraction_indices)
if isinstance(expr, (ZeroArray, ZeroMatrix)):
contraction_indices_flat = [j for i in contraction_indices for j in i]
shape = [e for i, e in enumerate(expr.shape) if i not in contraction_indices_flat]
return ZeroArray(*shape)
if isinstance(expr, PermuteDims):
return cls._handle_nested_permute_dims(expr, *contraction_indices)
if isinstance(expr, ArrayTensorProduct):
expr, contraction_indices = cls._sort_fully_contracted_args(expr, contraction_indices)
expr, contraction_indices = cls._lower_contraction_to_addends(expr, contraction_indices)
if len(contraction_indices) == 0:
return expr
if isinstance(expr, ArrayDiagonal):
return cls._handle_nested_diagonal(expr, *contraction_indices)
if isinstance(expr, ArrayAdd):
return ArrayAdd(*[ArrayContraction(i, *contraction_indices) for i in expr.args])
obj = Basic.__new__(cls, expr, *contraction_indices)
obj._subranks = _get_subranks(expr)
obj._mapping = _get_mapping_from_subranks(obj._subranks)
free_indices_to_position = {i: i for i in range(sum(obj._subranks)) if all([i not in cind for cind in contraction_indices])}
obj._free_indices_to_position = free_indices_to_position
shape = expr.shape
cls._validate(expr, *contraction_indices)
if shape:
shape = tuple(shp for i, shp in enumerate(shape) if not any(i in j for j in contraction_indices))
obj._shape = shape
return obj
def identify_hadamard_products(expr: Union[ArrayContraction, ArrayDiagonal]):
mapping = _get_mapping_from_subranks(expr.subranks)
editor: _EditArrayContraction
if isinstance(expr, ArrayContraction):
editor = _EditArrayContraction(expr)
elif isinstance(expr, ArrayDiagonal):
if isinstance(expr.expr, ArrayContraction):
editor = _EditArrayContraction(expr.expr)
diagonalized = ArrayContraction._push_indices_down(
expr.expr.contraction_indices, expr.diagonal_indices)
elif isinstance(expr.expr, ArrayTensorProduct):
editor = _EditArrayContraction(None)
editor.args_with_ind = [
_ArgE(arg) for i, arg in enumerate(expr.expr.args)
]
diagonalized = expr.diagonal_indices
else:
raise NotImplementedError("not implemented")
# Trick: add diagonalized indices as negative indices into the editor object:
for i, e in enumerate(diagonalized):
for j in e:
arg_pos, rel_pos = mapping[j]
editor.args_with_ind[arg_pos].indices[rel_pos] = -1 - i
map_contr_to_args: Dict[FrozenSet, List[_ArgE]] = defaultdict(list)
map_ind_to_inds = defaultdict(int)
for arg_with_ind in editor.args_with_ind:
for ind in arg_with_ind.indices:
map_ind_to_inds[ind] += 1
if None in arg_with_ind.indices:
continue
map_contr_to_args[frozenset(arg_with_ind.indices)].append(arg_with_ind)
k: FrozenSet[int]
v: List[_ArgE]
for k, v in map_contr_to_args.items():
if len(k) != 2:
# Hadamard product only defined for matrices:
continue
if len(v) == 1:
# Hadamard product with a single argument makes no sense:
continue
for ind in k:
if map_ind_to_inds[ind] <= 2:
# There is no other contraction, skip:
continue
# Check if expression is a trace:
if all([map_ind_to_inds[j] == len(v) and j >= 0 for j in k]):
# This is a trace
continue
# This is a Hadamard product:
def check_transpose(x):
x = [i if i >= 0 else -1 - i for i in x]
return x == sorted(x)
hp = hadamard_product(*[
i.element if check_transpose(i.indices) else Transpose(i.element)
for i in v
])
hp_indices = v[0].indices
if not check_transpose(v[0].indices):
hp_indices = list(reversed(hp_indices))
editor.insert_after(v[0], _ArgE(hp, hp_indices))
for i in v:
editor.args_with_ind.remove(i)
# Count the ranks of the arguments:
counter = 0
# Create a collector for the new diagonal indices:
diag_indices = defaultdict(list)
count_index_freq = Counter()
for arg_with_ind in editor.args_with_ind:
count_index_freq.update(Counter(arg_with_ind.indices))
free_index_count = count_index_freq[None]
# Construct the inverse permutation:
inv_perm1 = []
inv_perm2 = []
# Keep track of which diagonal indices have already been processed:
done = set([])
# Counter for the diagonal indices:
counter4 = 0
for arg_with_ind in editor.args_with_ind:
# If some diagonalization axes have been removed, they should be
# permuted in order to keep the permutation.
# Add permutation here
counter2 = 0 # counter for the indices
for i in arg_with_ind.indices:
if i is None:
inv_perm1.append(counter4)
counter2 += 1
counter4 += 1
continue
if i >= 0:
continue
# Reconstruct the diagonal indices:
diag_indices[-1 - i].append(counter + counter2)
if count_index_freq[i] == 1 and i not in done:
inv_perm1.append(free_index_count - 1 - i)
done.add(i)
elif i not in done:
inv_perm2.append(free_index_count - 1 - i)
done.add(i)
counter2 += 1
# Remove negative indices to restore a proper editor object:
arg_with_ind.indices = [
i if i is not None and i >= 0 else None
for i in arg_with_ind.indices
]
counter += len([i for i in arg_with_ind.indices if i is None or i < 0])
inverse_permutation = inv_perm1 + inv_perm2
permutation = _af_invert(inverse_permutation)
if isinstance(expr, ArrayContraction):
return editor.to_array_contraction()
else:
# Get the diagonal indices after the detection of HadamardProduct in the expression:
diag_indices_filtered = [
tuple(v) for v in diag_indices.values() if len(v) > 1
]
expr1 = editor.to_array_contraction()
expr2 = ArrayDiagonal(expr1, *diag_indices_filtered)
expr3 = PermuteDims(expr2, permutation)
return expr3
def split_multiple_contractions(self):
"""
Recognize multiple contractions and attempt at rewriting them as paired-contractions.
"""
from sympy import ask, Q
contraction_indices = self.contraction_indices
if isinstance(self.expr, ArrayTensorProduct):
args = list(self.expr.args)
else:
args = [self.expr]
# TODO: unify API, best location in ArrayTensorProduct
subranks = [get_rank(i) for i in args]
# TODO: unify API
mapping = _get_mapping_from_subranks(subranks)
reverse_mapping = {v:k for k, v in mapping.items()}
new_contraction_indices = []
for indl, links in enumerate(contraction_indices):
if len(links) <= 2:
new_contraction_indices.append(links)
continue
# Check multiple contractions:
#
# Examples:
#
# * `A_ij b_j0 C_jk` ===> `A*DiagMatrix(b)*C`
#
# Care for:
# - matrix being diagonalized (i.e. `A_ii`)
# - vectors being diagonalized (i.e. `a_i0`)
# Also consider the case of diagonal matrices being contracted:
current_dimension = self.expr.shape[links[0]]
tuple_links = [mapping[i] for i in links]
arg_indices, arg_positions = zip(*tuple_links)
args_updates = {}
if len(arg_indices) != len(set(arg_indices)):
# Maybe trace should be supported?
raise NotImplementedError
not_vectors = []
vectors = []
for arg_ind, arg_pos in tuple_links:
mat = args[arg_ind]
other_arg_pos = 1-arg_pos
other_arg_abs = reverse_mapping[arg_ind, other_arg_pos]
if (((1 not in mat.shape) and (not ask(Q.diagonal(mat)))) or
((current_dimension == 1) is True and mat.shape != (1, 1)) or
any([other_arg_abs in l for li, l in enumerate(contraction_indices) if li != indl])
):
not_vectors.append((arg_ind, arg_pos))
continue
args_updates[arg_ind] = diagonalize_vector(mat)
vectors.append((arg_ind, arg_pos))
vectors.append((arg_ind, 1-arg_pos))
if len(not_vectors) > 2:
new_contraction_indices.append(links)
continue
if len(not_vectors) == 0:
new_sequence = vectors[:1] + vectors[2:]
elif len(not_vectors) == 1:
new_sequence = not_vectors[:1] + vectors[:-1]
else:
new_sequence = not_vectors[:1] + vectors + not_vectors[1:]
for i in range(0, len(new_sequence) - 1, 2):
arg1, pos1 = new_sequence[i]
arg2, pos2 = new_sequence[i+1]
if arg1 == arg2:
raise NotImplementedError
continue
abspos1 = reverse_mapping[arg1, pos1]
abspos2 = reverse_mapping[arg2, pos2]
new_contraction_indices.append((abspos1, abspos2))
for ind, newarg in args_updates.items():
args[ind] = newarg
return ArrayContraction(
ArrayTensorProduct(*args),
*new_contraction_indices
)
def _support_function_tp1_recognize(contraction_indices, args):
subranks = [get_rank(i) for i in args]
coeff = reduce(lambda x, y: x * y,
[arg for arg, srank in zip(args, subranks) if srank == 0],
S.One)
mapping = _get_mapping_from_subranks(subranks)
new_contraction_indices = list(contraction_indices)
newargs = args[:] # make a copy of the list
removed = [None for i in newargs]
cumul = list(accumulate([0] + [get_rank(arg) for arg in args]))
new_perms = [
list(range(cumul[i], cumul[i + 1])) for i, arg in enumerate(args)
]
for pi, contraction_pair in enumerate(contraction_indices):
if len(contraction_pair) != 2:
continue
i1, i2 = contraction_pair
a1, e1 = mapping[i1]
a2, e2 = mapping[i2]
while removed[a1] is not None:
a1, e1 = removed[a1]
while removed[a2] is not None:
a2, e2 = removed[a2]
if a1 == a2:
trace_arg = newargs[a1]
newargs[a1] = Trace(trace_arg)._normalize()
new_contraction_indices[pi] = None
continue
if not isinstance(newargs[a1], MatrixExpr) or not isinstance(
newargs[a2], MatrixExpr):
continue
arg1 = newargs[a1]
arg2 = newargs[a2]
if (e1 == 1 and e2 == 1) or (e1 == 0 and e2 == 0):
arg2 = Transpose(arg2)
if e1 == 1:
argnew = arg1 * arg2
else:
argnew = arg2 * arg1
removed[a2] = a1, e1
new_perms[a1][e1] = new_perms[a2][1 - e2]
new_perms[a2] = None
newargs[a1] = argnew
newargs[a2] = None
new_contraction_indices[pi] = None
new_contraction_indices = [
i for i in new_contraction_indices if i is not None
]
newargs2 = [arg for arg in newargs if arg is not None]
if len(newargs2) == 0:
return coeff
tp = _a2m_tensor_product(*newargs2)
tc = ArrayContraction(tp, *new_contraction_indices)
new_perms2 = ArrayContraction._push_indices_up(
contraction_indices, [i for i in new_perms if i is not None])
permutation = _af_invert(
[j for i in new_perms2 for j in i if j is not None])
if permutation == [1, 0] and len(newargs2) == 1:
return Transpose(newargs2[0]).doit()
tperm = PermuteDims(tc, permutation)
return tperm