Fix exp_generic for general AbstractMatrix.

Project.toml

@@ -19,7 +19,8 @@ julia = "1"
 ForwardDiff = "f6369f11-7733-5829-9624-2563aa707210"
 Random = "9a3f8284-a2c9-5f02-9a11-845980a1fd5c"
 SafeTestsets = "1bc83da4-3b8d-516f-aca4-4fe02f6d838f"
+StaticArrays = "90137ffa-7385-5640-81b9-e52037218182"
 Test = "8dfed614-e22c-5e08-85e1-65c5234f0b40"
 
 [targets]
-test = ["ForwardDiff", "Test", "SafeTestsets", "Random"]
+test = ["ForwardDiff", "Test", "SafeTestsets", "StaticArrays", "Random"]

src/exp.jl

@@ -125,6 +125,8 @@ function naivemul!(C, A, B)
         mul!(C,A,B)
     end
 end
+_const(A) = A
+_const(A::Array) = Base.Experimental.Const(A)
 # Separated to make it easier to test.
 function naivemul!(C::AbstractMatrix{T}, A, B, Maxis, Naxis) where {T}
     N = last(Naxis)
@@ -138,15 +140,15 @@ function naivemul!(C::AbstractMatrix{T}, A, B, Maxis, Naxis) where {T}
                 while m < M - 3
                     Base.Cartesian.@nexprs 2 j -> Base.Cartesian.@nexprs 4 i -> Cmn_i_j = zero(T)
                     for k ∈ Kaxis
-                        Base.Cartesian.@nexprs 2 j -> Base.Cartesian.@nexprs 4 i -> Cmn_i_j = muladd(Base.Experimental.Const(A)[m+i,k],Base.Experimental.Const(B)[k,n+j],Cmn_i_j)
+                        Base.Cartesian.@nexprs 2 j -> Base.Cartesian.@nexprs 4 i -> Cmn_i_j = muladd(_const(A)[m+i,k],_const(B)[k,n+j],Cmn_i_j)
                     end
                     Base.Cartesian.@nexprs 2 j -> Base.Cartesian.@nexprs 4 i -> C[m+i,n+j] = Cmn_i_j
                     m += 4
                 end
                 for mm ∈ 1+m:M
                     Base.Cartesian.@nexprs 2 j -> Cmn_j = zero(T)
                     for k ∈ Kaxis
-                        Base.Cartesian.@nexprs 2 j -> Cmn_j = muladd(Base.Experimental.Const(A)[mm,k],Base.Experimental.Const(B)[k,n+j],Cmn_j)
+                        Base.Cartesian.@nexprs 2 j -> Cmn_j = muladd(_const(A)[mm,k],_const(B)[k,n+j],Cmn_j)
                     end
                     Base.Cartesian.@nexprs 2 j -> C[mm,n+j] = Cmn_j
                 end
@@ -156,15 +158,15 @@ function naivemul!(C::AbstractMatrix{T}, A, B, Maxis, Naxis) where {T}
             while m < M - 3
                 Base.Cartesian.@nexprs 4 i -> Cmn_i = zero(T)
                 for k ∈ Kaxis
-                    Base.Cartesian.@nexprs 4 i -> Cmn_i = muladd(Base.Experimental.Const(A)[m+i,k],Base.Experimental.Const(B)[k,N],Cmn_i)
+                    Base.Cartesian.@nexprs 4 i -> Cmn_i = muladd(_const(A)[m+i,k],_const(B)[k,N],Cmn_i)
                 end
                 Base.Cartesian.@nexprs 4 i -> C[m+i,N] = Cmn_i
                 m += 4
             end
             for mm ∈ 1+m:M
                 Cmn = zero(T)
                 for k ∈ Kaxis
-                    Cmn = muladd(Base.Experimental.Const(A)[mm,k], Base.Experimental.Const(B)[k,N], Cmn)
+                    Cmn = muladd(_const(A)[mm,k], _const(B)[k,N], Cmn)
                 end
                 C[mm,N] = Cmn
             end
@@ -225,7 +227,12 @@ function exp_generic_mutable(x::AbstractMatrix{T}, s, ::Val{13}) where {T}
 end
 function exp_generic!(y1, y2, y3, x, s, ::Val{13})
     if s > 0
-        exp_generic_core!(y1, y2, y3, x .* (1/(1 << s)), Val{13}())
+        _y3 = exp_generic_core!(y1, y2, y3, x .* (1/(1 << s)), Val{13}())
+        if typeof(_y3) === (y3)
+            y3 = _y3
+        else
+            y3 .= _y3
+        end
         for _ ∈ 1:s
             naivemul!(y1, y3, y3)
             y3, y1 = y1, y3
@@ -235,14 +242,17 @@ function exp_generic!(y1, y2, y3, x, s, ::Val{13})
         return exp_generic_core!(y1, y2, y3, x, Val{13}())
     end
 end
+# `lu!` is only defined for `StridedMatrix`, and `lu(::StaticArray)` (note `MArray<:StaticArray`) returns `::StaticArrays.LU !== LinearAlgebra.LU`
+_rdiv!(A, B::StridedMatrix) = rdiv!(A, lu!(B))
+_rdiv!(A, B) = A / B
+
 function exp_generic_core!(y1, y2, y3, x, ::Val{13})
     @inbounds for i ∈ eachindex(y3)
         y3[i] = -x[i]
     end
     exp_pade_p!(y1, y2, y3, Val{13}(), Val{13}())
     exp_pade_p!(y3, y2, x, Val{13}(), Val{13}())
-    rdiv!(y3, lu!(y1))
-    return y3
+    return _rdiv!(y3, y1)
 end
 
 function exp_pade_p!(y1::AbstractMatrix{T}, y2::AbstractMatrix{T}, x::AbstractMatrix, ::Val{13}, ::Val{13}) where {T}

test/runtests.jl

@@ -1,6 +1,6 @@
 using Test, LinearAlgebra, Random, SparseArrays, ExponentialUtilities
 using ExponentialUtilities: getH, getV, _exp!
-using ForwardDiff
+using ForwardDiff, StaticArrays
 
 @testset "Exp" begin
     n = 100
@@ -83,7 +83,17 @@ end
         B = rand(n,n);
         C = similar(A);
         @test ExponentialUtilities.naivemul!(C, A, B, axes(C)...) ≈ A*B
+        if n ≤ 16
+            Am = MMatrix{n,n}(A)
+            Bm = MMatrix{n,n}(B)
+            Cm = MMatrix{n,n}(A)
+            @test ExponentialUtilities.naivemul!(Cm, Am, Bm, axes(Cm)...) ≈ A*B
+        end
     end
+    A = @SMatrix rand(7,7);
+    Am = MMatrix(A);
+    Aa = Matrix(A);
+    @test exp(Aa) ≈ exp_generic(Aa) ≈ exp_generic(Am) ≈ exp_generic(A)
 end
 
 @testset "Phi" begin