Change how geometric factors are applied

illinois-ceesd · Feb 12, 2024 · d9fea00 · d9fea00
1 parent cb44738
commit d9fea00
Showing 1 changed file with 8 additions and 17 deletions.
diff --git a/grudge/op.py b/grudge/op.py
@@ -347,14 +347,6 @@ def compute_tensor_product_grad(actx, grp, diff_mat, vec, ijm,
                         arg_names=("stiff_1d", f"vec_{xyz_axis}"))
                 )
 
-                # apply metric terms
-                grad[xyz_axis] = actx.einsum(
-                    'rej,ej->ej',
-                    ijm[xyz_axis],
-                    grad[xyz_axis],
-                    tagged=(FirstAxisIsElementsTag(),),
-                    arg_names=("inv_jac_t", f"vec_{xyz_axis}")
-                )
         else:
             diff_mat = get_diff_mat(actx, grp, grp)
 
@@ -371,15 +363,14 @@ def compute_tensor_product_grad(actx, grp, diff_mat, vec, ijm,
                     )
                 )
 
-                grad[xyz_axis] = actx.einsum(
-                    "rej,ej->ej",
-                    ijm[xyz_axis],
-                    grad[xyz_axis],
-                    tagged=(FirstAxisIsElementsTag(),),
-                    arg_names=("inv_jac_t", f"vec_{xyz_axis}")
-                )
-
-        return make_obj_array(grad)
+        grad = actx.np.stack(grad)
+        return actx.einsum(
+            "xrej,rej->xej",
+            ijm,
+            grad,
+            tagged=(FirstAxisIsElementsTag(),),
+            arg_names=("inv_jac_t", f"grad")
+        )
 
     # }}}