Упрощение матричного выражения с помощью MatrixSymbol в Python

Question

Я использую SymPy для вычисления символьных матриц, однако некоторые операторы очень большие. Кажется, есть способ еще больше упростить их. Я использовал simplify(), но мне не удалось получить то, что я хочу.

Например, изображение ниже представляет собой матрицу, которая была получена в результате выполнения длинного списка предыдущих матричных вычислений. Вывод серии матричных вычислений, что требует дальнейшего упрощения

В заключительном утверждении есть два сложения и одно умножение матриц. Интересно, есть ли способ также выполнить умножение матриц справа, чтобы мы могли просто получить 3 матричных суммирования?

Я понимаю, что это можно сделать, выполнив некоторые алгебраические манипуляции вручную, но меня больше интересует команда для выполнения, так что команда принимает весь оператор в качестве входных данных и выполняет все упрощения, включая любые умножения, дополнения и выходы Что мне нужно. Все это должно быть сделано с использованием sympy. Другими словами, если можно добавить сложение или умножение, я хочу, чтобы это было сделано, а не оставлялось отмененным.

Это MCVE, которая имитирует мою проблему

from sympy import *
init_printing()
J_22 = MatrixSymbol('J_22', 3, 3)
COV_b=Matrix([[2,1,1],[1,2,1],[1,1,2]])
(COV_b+J_22)*COV_b

Вывод этого кода

Тем не менее, я хотел бы иметь это в качестве вывода

Я понимаю, что в этом простом примере я мог бы просто решить проблему с помощью следующего кода

from sympy import *
init_printing()
J_22 = MatrixSymbol('J_22', 3, 3)
COV_b=Matrix([[2,1,1],[1,2,1],[1,1,2]])
(COV_b*COV_b+J_22*COV_b)

Однако это всего лишь простой пример, в реальной задаче это невозможно увидеть до того, как будет сгенерирован вывод. Следовательно, я хотел бы иметь возможность использовать команду, которая принимает выходные данные первого предоставленного кода в качестве входных данных и выводит требуемые выходные данные.

ОБНОВЛЕНИЕ: @WelcometoStackOverflow предоставляет функцию, которая значительно упрощает работу, но все еще оставляет отмену добавления матрицы.

from sympy import *
init_printing()
J_22 = MatrixSymbol('J_22', 3, 3)
COV_b=Matrix([[2,1,1],[1,2,1],[1,1,2]])
T=(COV_b+J_22)*COV_b+COV_b
def expand_matmul(expr):
    import itertools
    for a in preorder_traversal(expr):
        if isinstance(a, MatMul):
            terms = [f.args if isinstance(f, MatAdd) else [f] for f in a.args]
            expanded = Add(*[MatMul(*t) for t in itertools.product(*terms)])
            if a != expanded:
                expr = expr.xreplace({a: expanded})
                return expand_matmul(expr)
    return expr
expand_matmul(T)

Выход

с суммированием между первыми двумя матрицами, все еще не выполненным.

Answer 1

Это известная и старая проблема с выражениями SymPy: Не удается развернуть матричное выражение . Модуль Matrix Expressions полезен, но не наиболее активно поддерживается в SymPy. Я собрал функцию для расширения таких вещей.

def expand_matmul(expr):
    import itertools
    for a in preorder_traversal(expr):
        if isinstance(a, MatMul) and any(isinstance(f, MatAdd) for f in a.args):
            terms = [f.args if isinstance(f, MatAdd) else [f] for f in a.args]
            expanded = MatAdd(*[MatMul(*t) for t in itertools.product(*terms)])
            if a != expanded:
                expr = expr.xreplace({a: expanded})
                return expand_matmul(expr)
    return expr

Функция обходит дерево выражений с самых высоких уровней в поисках возможности расширения MatMul. Возвращенное выражение может выиграть от вызова метода doit для выполнения любого отмененного умножения из явных матриц, как в примере ниже.

J_22 = MatrixSymbol('J_22', 3, 3)
COV_b=Matrix([[2,1,1],[1,2,1],[1,1,2]])
T=(COV_b+J_22)*COV_b+COV_b  
pprint(expand_matmul(T).doit())

печать

⎡8  6  6⎤       ⎡2  1  1⎤
⎢       ⎥       ⎢       ⎥
⎢6  8  6⎥ + J₂₂⋅⎢1  2  1⎥
⎢       ⎥       ⎢       ⎥
⎣6  6  8⎦       ⎣1  1  2⎦