expanding pows (#12872)

Co-authored-by: Daniel Diez <[email protected]>
KratosMultiphysics · Nov 26, 2024 · 8ef9268 · 8ef9268
1 parent 55ec898
commit 8ef9268
Show file tree

Hide file tree

Showing 2 changed files with 25 additions and 2 deletions.
diff --git a/kratos/python_scripts/sympy_fe_utilities.py b/kratos/python_scripts/sympy_fe_utilities.py
@@ -1,5 +1,9 @@
 import re
 import sympy
+from sympy.codegen.rewriting import create_expand_pow_optimization
+
+# We expand terms up to 8th power for optimization purposes
+expand_opt = create_expand_pow_optimization(8)
 
 def DefineMatrix(name, m, n):
     """
@@ -751,12 +755,14 @@ def _AuxiliaryOutputCollectionFactors(A, name, language, indentation_level, opti
     """
     symbol_name = "c" + name
     A_factors, A_collected = sympy.cse(A, sympy.numbered_symbols(symbol_name), optimizations)
-    A = A_collected[0] #overwrite lhs with the one with the collected components
+    A = A_collected[0] #overwrite A with the one with the collected components
+    A = expand_opt(A)
 
     Acoefficient_str = str("")
     for factor in A_factors:
         varname = str(factor[0])
         value = factor[1]
+        value = expand_opt(value)
         Acoefficient_str += OutputSymbolicVariableDeclaration(value, varname, language, indentation_level, replace_indices)
 
     A_out = Acoefficient_str + output_func(A, name, language, indentation_level, replace_indices, assignment_op)

diff --git a/kratos/tests/test_sympy_fe_utilities.py b/kratos/tests/test_sympy_fe_utilities.py
@@ -261,7 +261,7 @@ def testOutputScalar_CollectingFactors(self):
             "    const double cmyvariable2 = 2*y[2];",
             "    const double cmyvariable3 = 2*y[3];",
             "    const double cmyvariable4 = 2*y[4];",
-            "    myvariable=3*cmyvariable0*x[0] + 3*cmyvariable1*x[1] + 3*cmyvariable2*x[2] + 3*cmyvariable3*x[3] + 3*cmyvariable4*x[4] + 3*(pow(y[0], 2) + pow(y[1], 2) + pow(y[2], 2) + pow(y[3], 2) + pow(y[4], 2))*(x[0]*(cmyvariable0 + 3*x[0]) + x[1]*(cmyvariable1 + 3*x[1]) + x[2]*(cmyvariable2 + 3*x[2]) + x[3]*(cmyvariable3 + 3*x[3]) + x[4]*(cmyvariable4 + 3*x[4]));",
+            "    myvariable=3*cmyvariable0*x[0] + 3*cmyvariable1*x[1] + 3*cmyvariable2*x[2] + 3*cmyvariable3*x[3] + 3*cmyvariable4*x[4] + 3*(x[0]*(cmyvariable0 + 3*x[0]) + x[1]*(cmyvariable1 + 3*x[1]) + x[2]*(cmyvariable2 + 3*x[2]) + x[3]*(cmyvariable3 + 3*x[3]) + x[4]*(cmyvariable4 + 3*x[4]))*(y[0]*y[0] + y[1]*y[1] + y[2]*y[2] + y[3]*y[3] + y[4]*y[4]);",
             ""
         ])
         self.assertEqual(code, expected_code)
@@ -289,5 +289,22 @@ def testDefineSymmetricFourthOrderTensor(self):
                         aux_l = max(k,l)
                         self.assertEqual(A[i,j,k,l].name, "A_{}_{}_{}_{}".format(aux_i,aux_j,aux_k,aux_l))
 
+    def testAuxiliaryOutputCollectionFactorsWithPower(self):
+        # Note: create_expand_pow_optimization is not expanding expressions
+        # when there is a sum as shown in the last term of the test
+        x = KratosSympy.DefineVector('x', 5)
+        expression = (3*(x.T*(3*x))*(x.T*x) + 6*x.T*x)[0, 0]
+        code = KratosSympy._AuxiliaryOutputCollectionFactors(expression, "a", "c", 1, "basic", False, "=", KratosSympy.OutputScalar)
+        expected_code = "\n".join([
+            "    const double ca0 = x_0*x_0;",
+            "    const double ca1 = x_1*x_1;",
+            "    const double ca2 = x_2*x_2;",
+            "    const double ca3 = x_3*x_3;",
+            "    const double ca4 = x_4*x_4;",
+            "    a=6*ca0 + 6*ca1 + 6*ca2 + 6*ca3 + 6*ca4 + 9*pow(ca0 + ca1 + ca2 + ca3 + ca4, 2);",
+            ""
+        ])
+        self.assertEqual(code, expected_code)
+
 if __name__ == '__main__':
     KratosUnittest.main()