Merge pull request #2232 from CliMA/dy/matrix_mul_symbol

Use * to denote matrix multiplication

Author: dennisYatunin

examples/hybrid/staggered_nonhydrostatic_model.jl

@@ -331,11 +331,11 @@ function implicit_equation_jacobian!(j, Y, p, δtγ, t)
     #     norm_sqr(C123(ᶜuₕ) + C123(ᶜinterp(ᶠw))) / 2 =
     #     ACT12(ᶜuₕ) * ᶜuₕ / 2 + ACT3(ᶜinterp(ᶠw)) * ᶜinterp(ᶠw) / 2
     # ∂(ᶜK)/∂(ᶠw) = ACT3(ᶜinterp(ᶠw)) * ᶜinterp_matrix()
-    @. ∂ᶜK∂ᶠw = DiagonalMatrixRow(adjoint(CT3(ᶜinterp(ᶠw)))) ⋅ ᶜinterp_matrix()
+    @. ∂ᶜK∂ᶠw = DiagonalMatrixRow(adjoint(CT3(ᶜinterp(ᶠw)))) * ᶜinterp_matrix()
 
     # ᶜρₜ = -ᶜdivᵥ(ᶠinterp(ᶜρ) * ᶠw)
     # ∂(ᶜρₜ)/∂(ᶠw) = -ᶜdivᵥ_matrix() * ᶠinterp(ᶜρ) * ᶠg³³
-    @. ∂ᶜρₜ∂ᶠ𝕄 = -(ᶜdivᵥ_matrix()) ⋅ DiagonalMatrixRow(ᶠinterp(ᶜρ) * g³³(ᶠgⁱʲ))
+    @. ∂ᶜρₜ∂ᶠ𝕄 = -(ᶜdivᵥ_matrix()) * DiagonalMatrixRow(ᶠinterp(ᶜρ) * g³³(ᶠgⁱʲ))
 
     if :ρθ in propertynames(Y.c)
         ᶜρθ = Y.c.ρθ
@@ -349,14 +349,14 @@ function implicit_equation_jacobian!(j, Y, p, δtγ, t)
             # ᶜρθₜ = -ᶜdivᵥ(ᶠinterp(ᶜρθ) * ᶠw)
             # ∂(ᶜρθₜ)/∂(ᶠw) = -ᶜdivᵥ_matrix() * ᶠinterp(ᶜρθ) * ᶠg³³
             @. ∂ᶜ𝔼ₜ∂ᶠ𝕄 =
-                -(ᶜdivᵥ_matrix()) ⋅ DiagonalMatrixRow(ᶠinterp(ᶜρθ) * g³³(ᶠgⁱʲ))
+                -(ᶜdivᵥ_matrix()) * DiagonalMatrixRow(ᶠinterp(ᶜρθ) * g³³(ᶠgⁱʲ))
         else
             # ᶜρθₜ = -ᶜdivᵥ(ᶠinterp(ᶜρ) * ᶠupwind_product(ᶠw, ᶜρθ / ᶜρ))
             # ∂(ᶜρθₜ)/∂(ᶠw) =
             #     -ᶜdivᵥ_matrix() * ᶠinterp(ᶜρ) *
             #     ∂(ᶠupwind_product(ᶠw, ᶜρθ / ᶜρ))/∂(ᶠw)
             @. ∂ᶜ𝔼ₜ∂ᶠ𝕄 =
-                -(ᶜdivᵥ_matrix()) ⋅ DiagonalMatrixRow(
+                -(ᶜdivᵥ_matrix()) * DiagonalMatrixRow(
                     ᶠinterp(ᶜρ) *
                     vec_data(ᶠno_flux(ᶠupwind_product(ᶠw + εw, ᶜρθ / ᶜρ))) /
                     vec_data(CT3(ᶠw + εw)) * g³³(ᶠgⁱʲ),
@@ -381,10 +381,12 @@ function implicit_equation_jacobian!(j, Y, p, δtγ, t)
                 # ∂(ᶜp)/∂(ᶠw) = ∂(ᶜp)/∂(ᶜK) * ∂(ᶜK)/∂(ᶠw)
                 # ∂(ᶜp)/∂(ᶜK) = -ᶜρ * R_d / cv_d
                 @. ∂ᶜ𝔼ₜ∂ᶠ𝕄 =
-                    -(ᶜdivᵥ_matrix()) ⋅ (
+                    -(ᶜdivᵥ_matrix()) * (
                         DiagonalMatrixRow(ᶠinterp(ᶜρe + ᶜp) * g³³(ᶠgⁱʲ)) +
-                        DiagonalMatrixRow(CT3(ᶠw)) ⋅ ᶠinterp_matrix() ⋅
-                        DiagonalMatrixRow(-(ᶜρ * R_d / cv_d)) ⋅ ∂ᶜK∂ᶠw
+                        DiagonalMatrixRow(CT3(ᶠw)) *
+                        ᶠinterp_matrix() *
+                        DiagonalMatrixRow(-(ᶜρ * R_d / cv_d)) *
+                        ∂ᶜK∂ᶠw
                     )
             else
                 # ᶜρeₜ =
@@ -397,15 +399,17 @@ function implicit_equation_jacobian!(j, Y, p, δtγ, t)
                 # ∂(ᶜp)/∂(ᶠw) = ∂(ᶜp)/∂(ᶜK) * ∂(ᶜK)/∂(ᶠw)
                 # ∂(ᶜp)/∂(ᶜK) = -ᶜρ * R_d / cv_d
                 @. ∂ᶜ𝔼ₜ∂ᶠ𝕄 =
-                    -(ᶜdivᵥ_matrix()) ⋅ DiagonalMatrixRow(ᶠinterp(ᶜρ)) ⋅ (
+                    -(ᶜdivᵥ_matrix()) *
+                    DiagonalMatrixRow(ᶠinterp(ᶜρ)) *
+                    (
                         DiagonalMatrixRow(
                             vec_data(
                                 ᶠno_flux(
                                     ᶠupwind_product(ᶠw + εw, (ᶜρe + ᶜp) / ᶜρ),
                                 ),
                             ) / vec_data(CT3(ᶠw + εw)) * g³³(ᶠgⁱʲ),
                         ) +
-                        ᶠno_flux_row(ᶠupwind_product_matrix(ᶠw)) ⋅
+                        ᶠno_flux_row(ᶠupwind_product_matrix(ᶠw)) *
                         (-R_d / cv_d * ∂ᶜK∂ᶠw)
                     )
             end
@@ -414,11 +418,11 @@ function implicit_equation_jacobian!(j, Y, p, δtγ, t)
             # ∂ᶜ𝔼ₜ∂ᶠ𝕄 has 3 diagonals instead of 5
             if isnothing(ᶠupwind_product)
                 @. ∂ᶜ𝔼ₜ∂ᶠ𝕄 =
-                    -(ᶜdivᵥ_matrix()) ⋅
+                    -(ᶜdivᵥ_matrix()) *
                     DiagonalMatrixRow(ᶠinterp(ᶜρe + ᶜp) * g³³(ᶠgⁱʲ))
             else
                 @. ∂ᶜ𝔼ₜ∂ᶠ𝕄 =
-                    -(ᶜdivᵥ_matrix()) ⋅ DiagonalMatrixRow(
+                    -(ᶜdivᵥ_matrix()) * DiagonalMatrixRow(
                         ᶠinterp(ᶜρ) * vec_data(
                             ᶠno_flux(ᶠupwind_product(ᶠw + εw, (ᶜρe + ᶜp) / ᶜρ)),
                         ) / vec_data(CT3(ᶠw + εw)) * g³³(ᶠgⁱʲ),
@@ -443,9 +447,9 @@ function implicit_equation_jacobian!(j, Y, p, δtγ, t)
             #     -ᶜdivᵥ_matrix() * ᶠinterp(ᶜρe_int + ᶜp) * ᶠg³³ +
             #     ᶜinterp_matrix() * adjoint(ᶠgradᵥ(ᶜp)) * ᶠg³³
             @. ∂ᶜ𝔼ₜ∂ᶠ𝕄 =
-                -(ᶜdivᵥ_matrix()) ⋅
+                -(ᶜdivᵥ_matrix()) *
                 DiagonalMatrixRow(ᶠinterp(ᶜρe_int + ᶜp) * g³³(ᶠgⁱʲ)) +
-                ᶜinterp_matrix() ⋅
+                ᶜinterp_matrix() *
                 DiagonalMatrixRow(adjoint(ᶠgradᵥ(ᶜp)) * g³³(ᶠgⁱʲ))
         else
             # ᶜρe_intₜ =
@@ -456,12 +460,12 @@ function implicit_equation_jacobian!(j, Y, p, δtγ, t)
             #     ∂(ᶠupwind_product(ᶠw, (ᶜρe_int + ᶜp) / ᶜρ))/∂(ᶠw) +
             #     ᶜinterp_matrix() * adjoint(ᶠgradᵥ(ᶜp)) * ᶠg³³
             @. ∂ᶜ𝔼ₜ∂ᶠ𝕄 =
-                -(ᶜdivᵥ_matrix()) ⋅ DiagonalMatrixRow(
+                -(ᶜdivᵥ_matrix()) * DiagonalMatrixRow(
                     ᶠinterp(ᶜρ) * vec_data(
                         ᶠno_flux(ᶠupwind_product(ᶠw + εw, (ᶜρe_int + ᶜp) / ᶜρ)),
                     ) / vec_data(CT3(ᶠw + εw)) * g³³(ᶠgⁱʲ),
                 ) +
-                ᶜinterp_matrix() ⋅
+                ᶜinterp_matrix() *
                 DiagonalMatrixRow(adjoint(ᶠgradᵥ(ᶜp)) * g³³(ᶠgⁱʲ))
         end
     end
@@ -480,7 +484,8 @@ function implicit_equation_jacobian!(j, Y, p, δtγ, t)
         # ∂(ᶠgradᵥ(ᶜp))/∂(ᶜρθ) =
         #     ᶠgradᵥ_matrix() * γ * R_d * (ᶜρθ * R_d / p_0)^(γ - 1)
         @. ∂ᶠ𝕄ₜ∂ᶜ𝔼 =
-            -DiagonalMatrixRow(1 / ᶠinterp(ᶜρ)) ⋅ ᶠgradᵥ_matrix() ⋅
+            -DiagonalMatrixRow(1 / ᶠinterp(ᶜρ)) *
+            ᶠgradᵥ_matrix() *
             DiagonalMatrixRow(γ * R_d * (ᶜρθ * R_d / p_0)^(γ - 1))
 
         if flags.∂ᶠ𝕄ₜ∂ᶜρ_mode == :exact
@@ -489,20 +494,20 @@ function implicit_equation_jacobian!(j, Y, p, δtγ, t)
             # ∂(ᶠwₜ)/∂(ᶠinterp(ᶜρ)) = ᶠgradᵥ(ᶜp) / ᶠinterp(ᶜρ)^2
             # ∂(ᶠinterp(ᶜρ))/∂(ᶜρ) = ᶠinterp_matrix()
             @. ∂ᶠ𝕄ₜ∂ᶜρ =
-                DiagonalMatrixRow(ᶠgradᵥ(ᶜp) / ᶠinterp(ᶜρ)^2) ⋅ ᶠinterp_matrix()
+                DiagonalMatrixRow(ᶠgradᵥ(ᶜp) / ᶠinterp(ᶜρ)^2) * ᶠinterp_matrix()
         elseif flags.∂ᶠ𝕄ₜ∂ᶜρ_mode == :hydrostatic_balance
             # same as above, but we assume that ᶠgradᵥ(ᶜp) / ᶠinterp(ᶜρ) =
             # -ᶠgradᵥ(ᶜΦ)
             @. ∂ᶠ𝕄ₜ∂ᶜρ =
-                -DiagonalMatrixRow(ᶠgradᵥ(ᶜΦ) / ᶠinterp(ᶜρ)) ⋅ ᶠinterp_matrix()
+                -DiagonalMatrixRow(ᶠgradᵥ(ᶜΦ) / ᶠinterp(ᶜρ)) * ᶠinterp_matrix()
         end
     elseif :ρe in propertynames(Y.c)
         # ᶠwₜ = -ᶠgradᵥ(ᶜp) / ᶠinterp(ᶜρ) - ᶠgradᵥ(ᶜK + ᶜΦ)
         # ∂(ᶠwₜ)/∂(ᶜρe) = ∂(ᶠwₜ)/∂(ᶠgradᵥ(ᶜp)) * ∂(ᶠgradᵥ(ᶜp))/∂(ᶜρe)
         # ∂(ᶠwₜ)/∂(ᶠgradᵥ(ᶜp)) = -1 / ᶠinterp(ᶜρ)
         # ∂(ᶠgradᵥ(ᶜp))/∂(ᶜρe) = ᶠgradᵥ_matrix() * R_d / cv_d
         @. ∂ᶠ𝕄ₜ∂ᶜ𝔼 =
-            -DiagonalMatrixRow(1 / ᶠinterp(ᶜρ)) ⋅ (ᶠgradᵥ_matrix() * R_d / cv_d)
+            -DiagonalMatrixRow(1 / ᶠinterp(ᶜρ)) * (ᶠgradᵥ_matrix() * R_d / cv_d)
 
         if flags.∂ᶠ𝕄ₜ∂ᶜρ_mode == :exact
             # ᶠwₜ = -ᶠgradᵥ(ᶜp) / ᶠinterp(ᶜρ) - ᶠgradᵥ(ᶜK + ᶜΦ)
@@ -515,24 +520,26 @@ function implicit_equation_jacobian!(j, Y, p, δtγ, t)
             # ∂(ᶠwₜ)/∂(ᶠinterp(ᶜρ)) = ᶠgradᵥ(ᶜp) / ᶠinterp(ᶜρ)^2
             # ∂(ᶠinterp(ᶜρ))/∂(ᶜρ) = ᶠinterp_matrix()
             @. ∂ᶠ𝕄ₜ∂ᶜρ =
-                -DiagonalMatrixRow(1 / ᶠinterp(ᶜρ)) ⋅ ᶠgradᵥ_matrix() ⋅
+                -DiagonalMatrixRow(1 / ᶠinterp(ᶜρ)) *
+                ᶠgradᵥ_matrix() *
                 DiagonalMatrixRow(R_d * (-(ᶜK + ᶜΦ) / cv_d + T_tri)) +
-                DiagonalMatrixRow(ᶠgradᵥ(ᶜp) / ᶠinterp(ᶜρ)^2) ⋅ ᶠinterp_matrix()
+                DiagonalMatrixRow(ᶠgradᵥ(ᶜp) / ᶠinterp(ᶜρ)^2) * ᶠinterp_matrix()
         elseif flags.∂ᶠ𝕄ₜ∂ᶜρ_mode == :hydrostatic_balance
             # same as above, but we assume that ᶠgradᵥ(ᶜp) / ᶠinterp(ᶜρ) =
             # -ᶠgradᵥ(ᶜΦ) and that ᶜK is negligible compared ot ᶜΦ
             @. ∂ᶠ𝕄ₜ∂ᶜρ =
-                -DiagonalMatrixRow(1 / ᶠinterp(ᶜρ)) ⋅ ᶠgradᵥ_matrix() ⋅
+                -DiagonalMatrixRow(1 / ᶠinterp(ᶜρ)) *
+                ᶠgradᵥ_matrix() *
                 DiagonalMatrixRow(R_d * (-(ᶜΦ) / cv_d + T_tri)) -
-                DiagonalMatrixRow(ᶠgradᵥ(ᶜΦ) / ᶠinterp(ᶜρ)) ⋅ ᶠinterp_matrix()
+                DiagonalMatrixRow(ᶠgradᵥ(ᶜΦ) / ᶠinterp(ᶜρ)) * ᶠinterp_matrix()
         end
     elseif :ρe_int in propertynames(Y.c)
         # ᶠwₜ = -ᶠgradᵥ(ᶜp) / ᶠinterp(ᶜρ) - ᶠgradᵥ(ᶜK + ᶜΦ)
         # ∂(ᶠwₜ)/∂(ᶜρe_int) = ∂(ᶠwₜ)/∂(ᶠgradᵥ(ᶜp)) * ∂(ᶠgradᵥ(ᶜp))/∂(ᶜρe_int)
         # ∂(ᶠwₜ)/∂(ᶠgradᵥ(ᶜp)) = -1 / ᶠinterp(ᶜρ)
         # ∂(ᶠgradᵥ(ᶜp))/∂(ᶜρe_int) = ᶠgradᵥ_matrix() * R_d / cv_d
         @. ∂ᶠ𝕄ₜ∂ᶜ𝔼 =
-            DiagonalMatrixRow(-1 / ᶠinterp(ᶜρ)) ⋅ (ᶠgradᵥ_matrix() * R_d / cv_d)
+            DiagonalMatrixRow(-1 / ᶠinterp(ᶜρ)) * (ᶠgradᵥ_matrix() * R_d / cv_d)
 
         if flags.∂ᶠ𝕄ₜ∂ᶜρ_mode == :exact
             # ᶠwₜ = -ᶠgradᵥ(ᶜp) / ᶠinterp(ᶜρ) - ᶠgradᵥ(ᶜK + ᶜΦ)
@@ -544,16 +551,16 @@ function implicit_equation_jacobian!(j, Y, p, δtγ, t)
             # ∂(ᶠwₜ)/∂(ᶠinterp(ᶜρ)) = ᶠgradᵥ(ᶜp) / ᶠinterp(ᶜρ)^2
             # ∂(ᶠinterp(ᶜρ))/∂(ᶜρ) = ᶠinterp_matrix()
             @. ∂ᶠ𝕄ₜ∂ᶜρ =
-                -DiagonalMatrixRow(1 / ᶠinterp(ᶜρ)) ⋅
+                -DiagonalMatrixRow(1 / ᶠinterp(ᶜρ)) *
                 (ᶠgradᵥ_matrix() * R_d * T_tri) +
-                DiagonalMatrixRow(ᶠgradᵥ(ᶜp) / ᶠinterp(ᶜρ)^2) ⋅ ᶠinterp_matrix()
+                DiagonalMatrixRow(ᶠgradᵥ(ᶜp) / ᶠinterp(ᶜρ)^2) * ᶠinterp_matrix()
         elseif flags.∂ᶠ𝕄ₜ∂ᶜρ_mode == :hydrostatic_balance
             # same as above, but we assume that ᶠgradᵥ(ᶜp) / ᶠinterp(ᶜρ) =
             # -ᶠgradᵥ(ᶜΦ)
             @. ∂ᶠ𝕄ₜ∂ᶜρ =
-                DiagonalMatrixRow(-1 / ᶠinterp(ᶜρ)) ⋅
+                DiagonalMatrixRow(-1 / ᶠinterp(ᶜρ)) *
                 (ᶠgradᵥ_matrix() * R_d * T_tri) -
-                DiagonalMatrixRow(ᶠgradᵥ(ᶜΦ) / ᶠinterp(ᶜρ)) ⋅ ᶠinterp_matrix()
+                DiagonalMatrixRow(ᶠgradᵥ(ᶜΦ) / ᶠinterp(ᶜρ)) * ᶠinterp_matrix()
         end
     end
 
@@ -571,13 +578,14 @@ function implicit_equation_jacobian!(j, Y, p, δtγ, t)
     # ∂(ᶠwₜ)/∂(ᶠgradᵥ(ᶜK + ᶜΦ)) = -1
     # ∂(ᶠgradᵥ(ᶜK + ᶜΦ))/∂(ᶜK) = ᶠgradᵥ_matrix()
     if :ρθ in propertynames(Y.c) || :ρe_int in propertynames(Y.c)
-        @. ∂ᶠ𝕄ₜ∂ᶠ𝕄 = -(ᶠgradᵥ_matrix()) ⋅ ∂ᶜK∂ᶠw
+        @. ∂ᶠ𝕄ₜ∂ᶠ𝕄 = -(ᶠgradᵥ_matrix()) * ∂ᶜK∂ᶠw
     elseif :ρe in propertynames(Y.c)
         @. ∂ᶠ𝕄ₜ∂ᶠ𝕄 =
             -(
-                DiagonalMatrixRow(1 / ᶠinterp(ᶜρ)) ⋅ ᶠgradᵥ_matrix() ⋅
+                DiagonalMatrixRow(1 / ᶠinterp(ᶜρ)) *
+                ᶠgradᵥ_matrix() *
                 DiagonalMatrixRow(-(ᶜρ * R_d / cv_d)) + ᶠgradᵥ_matrix()
-            ) ⋅ ∂ᶜK∂ᶠw
+            ) * ∂ᶜK∂ᶠw
     end
 
     I = one(∂R∂Y)

src/MatrixFields/MatrixFields.jl

@@ -12,8 +12,8 @@ matrices, one for each column of the `Field`. Such `Field`s are called
 for them:
 - Constructors, e.g., `matrix_field = @. BidiagonalMatrixRow(field1, field2)`
 - Linear combinations, e.g., `@. 3 * matrix_field1 + matrix_field2 / 3`
-- Matrix-vector multiplication, e.g., `@. matrix_field ⋅ field`
-- Matrix-matrix multiplication, e.g., `@. matrix_field1 ⋅ matrix_field2`
+- Matrix-vector multiplication, e.g., `@. matrix_field * field`
+- Matrix-matrix multiplication, e.g., `@. matrix_field1 * matrix_field2`
 - Compatibility with `LinearAlgebra.I`, e.g., `@. matrix_field = (4I,)` or
     `@. matrix_field - (4I,)`
 - Integration with `RecursiveApply`, e.g., the entries of `matrix_field` can be
@@ -107,6 +107,37 @@ include("field_matrix_solver.jl")
 include("field_matrix_iterative_solver.jl")
 include("field_matrix_with_solver.jl")
 
+const FieldOrStencilStyleType = Union{
+    Fields.Field,
+    Base.Broadcast.Broadcasted{<:Fields.AbstractFieldStyle},
+    Operators.StencilBroadcasted,
+    LazyOperatorBroadcasted,
+}
+
+Base.Broadcast.broadcasted(
+    ::typeof(*),
+    field_or_broadcasted::FieldOrStencilStyleType,
+    args...,
+) =
+    unrolled_reduce(args; init = field_or_broadcasted) do arg1, arg2
+        arg1_isa_matrix =
+            eltype(arg1) <: BandMatrixRow || arg1 isa LazyOperatorBroadcasted
+        use_matrix_mul_op = arg1_isa_matrix && arg2 isa FieldOrStencilStyleType
+        op = use_matrix_mul_op ? MultiplyColumnwiseBandMatrixField() : ⊠
+        Base.Broadcast.broadcasted(op, arg1, arg2)
+    end
+Base.Broadcast.broadcasted(
+    ::typeof(*),
+    single_value_or_broadcasted::SingleValueStyleType,
+    field_or_broadcasted::FieldOrStencilStyleType,
+    args...,
+) = Base.Broadcast.broadcasted(
+    ⊠,
+    single_value_or_broadcasted,
+    Base.Broadcast.broadcasted(*, field_or_broadcasted, args...),
+)
+# TODO: Generalize this to handle, e.g., @. scalar * scalar * matrix * matrix.
+
 function Base.show(io::IO, field::ColumnwiseBandMatrixField)
     print(io, eltype(field), "-valued Field")
     if eltype(eltype(field)) <: Number

src/MatrixFields/field_name_dict.jl

@@ -422,7 +422,7 @@ Base.Broadcast.broadcasted(
     arg3,
     args...,
 ) =
-    foldl((arg1, arg2, arg3, args...)) do arg1′, arg2′
+    unrolled_reduce((arg1, arg2, arg3, args...)) do arg1′, arg2′
         Base.Broadcast.broadcasted(f, arg1′, arg2′)
     end
 
@@ -553,7 +553,7 @@ function Base.Broadcast.broadcasted(
             elseif entry2 isa ScalingFieldMatrixEntry
                 Base.Broadcast.broadcasted(*, entry1, (scaling_value(entry2),))
             else
-                Base.Broadcast.broadcasted(⋅, entry1, entry2)
+                Base.Broadcast.broadcasted(*, entry1, entry2)
             end
         end
         length(summand_bcs) == 1 ? summand_bcs[1] :