Ok I ended up implementing something in line with the pseudo-code I was thinking:
from cdb.utils.indices import replace_index
def unify_variations_inplace(ex, symbol_name, target_indices):
# 1. Helper to find the sum node (integrand)
def find_sum_node(node):
if node.name == r'\sum': return node
for child in node.children():
res = find_sum_node(child)
if res: return res
return None
# Locate the node inside the original expression (ex)
sum_node = find_sum_node(ex.top())
if not sum_node:
return ex
# 2. Build the unified integrand as a new Ex
new_sum = Ex('0')
for term_node in sum_node.children():
t_ex = term_node.ex()
matches = list(t_ex[symbol_name])
if matches:
current_indices = [str(idx.name) for idx in matches[0].indices()]
for i in range(min(len(current_indices), len(target_indices))):
old_idx, new_idx = current_indices[i], target_indices[i]
if old_idx != new_idx:
t_ex = replace_index(t_ex, "", old_idx, new_idx)
# We skip canonicalise here to prevent index shuffling
new_sum += t_ex
# 3. REINSERT the result into the original expression tree
# This physically overwrites the sum inside the integral context
sum_node.replace(new_sum)
return ex
Then I use this like this:
dg{#}::LaTeXForm("\delta{g}").
action_test := \int{ W^{\mu}_{\nu \rho} g_{\mu \alpha} D^{\alpha} dg^{\rho \nu} + C_{\beta \mu} dg^{\beta \mu}}{x};
unify_variations_inplace(action_test, "dg", ['\psi', '\chi']);
factor_out(_, $dg^{\psi \chi}$);
The only shortcoming of this function is that as it is, it doesn't check if the target_indices are being already used inside the expression, but it shouldn't be too difficult to add