From 32ee9f1e33798e6a7238e193557a75665550f4c4 Mon Sep 17 00:00:00 2001
From: "Ziyi (Francis) Yin" <ziyi.yin@gatech.edu>
Date: Tue, 3 May 2022 16:29:39 -0400
Subject: [PATCH] fix mul! add associated test

---
 src/TimeModeling/LinearOperators/operators.jl |  2 +-
 src/TimeModeling/Types/abstract.jl            |  5 ++---
 test/test_linear_algebra.jl                   | 13 +++++++++++++
 3 files changed, 16 insertions(+), 4 deletions(-)

diff --git a/src/TimeModeling/LinearOperators/operators.jl b/src/TimeModeling/LinearOperators/operators.jl
index 44b9a17e9..4a69a855a 100644
--- a/src/TimeModeling/LinearOperators/operators.jl
+++ b/src/TimeModeling/LinearOperators/operators.jl
@@ -174,7 +174,7 @@ adjoint(L::LazyScal) = LazyScal(L.s, adjoint(L.P))
 
 mul!(out::SourceType{T}, F::judiPropagator{T, O}, q::SourceType{T}) where {T<:Number, O} = begin y = F*q; copyto!(out, y) end
 mul!(out::SourceType{T}, F::judiPropagator{T, O}, q::Vector{T}) where {T<:Number, O} = begin y = F*q; copyto!(out, y) end
-mul!(out::SourceType{T}, F::joLinearFunction{T, T}, q::SourceType{T}) where {T<:Number} = begin y = F*q; copyto!(out, y) end
+mul!(out::Union{SourceType{T1},SourceType{T2}}, F::joAbstractLinearOperator{T2, T1}, q::SourceType{T2}) where {T1<:Number, T2<:Number} = begin y = F*q; copyto!(out, y) end
 mul!(out::SourceType{T}, F::judiAbstractJacobian{T, :born, FT}, q::Vector{T}) where {T<:Number, FT} = begin y = F*q[:]; copyto!(out, y) end
 mul!(out::SourceType{T}, F::judiAbstractJacobian{T, :born, FT}, q::Array{T, N}) where {T<:Number, FT, N} = begin y = F*q[:]; copyto!(out, y) end
 mul!(out::Array{T, 2}, F::judiAbstractJacobian{T, :adjoint_born, FT}, q::SourceType{T}) where {T<:Number, FT} =  begin y = F*q; copyto!(out, y) end
diff --git a/src/TimeModeling/Types/abstract.jl b/src/TimeModeling/Types/abstract.jl
index 1ce4128b3..d529657b3 100644
--- a/src/TimeModeling/Types/abstract.jl
+++ b/src/TimeModeling/Types/abstract.jl
@@ -74,9 +74,8 @@ vec(x::judiMultiSourceVector) = vcat(vec.(x.data)...)
 time_sampling(ms::judiMultiSourceVector) = [1 for i=1:ms.nsrc]
 ############################################################################################################################
 # Linear algebra `*`
-(msv::judiMultiSourceVector{T})(x::Vector{T}) where {T<:AbstractFloat} = x
-(msv::judiMultiSourceVector{T})(x::judiMultiSourceVector{T}) where {T} = x
-(msv::judiMultiSourceVector{mT})(x::Vector{T}) where {mT, T<:Array} = begin y = deepcopy(msv); y.data .= x; return y end
+(msv::judiMultiSourceVector{mT})(x::AbstractVector{T}) where {mT<:Number, T<:Number} = x
+(msv::judiMultiSourceVector{mT})(x::AbstractVector{T}) where {mT, T<:Array} = begin y = deepcopy(msv); y.data .= x; return y end
 
 function *(J::Union{Matrix{vDT}, joAbstractLinearOperator}, x::judiMultiSourceVector{vDT}) where vDT
     outvec = try J.fop(x) catch; J*vec(x) end
diff --git a/test/test_linear_algebra.jl b/test/test_linear_algebra.jl
index 81993b966..af7bf2639 100644
--- a/test/test_linear_algebra.jl
+++ b/test/test_linear_algebra.jl
@@ -89,6 +89,19 @@
         @test isapprox(dobs_out, Md'*dobs; rtol=ftol)
 
         # Check JUDI-JOLI compat
+        
+        # check JOLI operator w/ judiVector
+        mul!(dobs_out, J*D, dm)
+        mul!(dobs, J*D, dm.data)
+        dlin = J*D*dm
+        @test isapprox(dobs_out, dlin; rtol=ftol)
+        @test isapprox(dobs_out, dobs; rtol=ftol)
+        # check JOLI operator w/ judiVector
+        DFT = joDFT(n; DDT=Float32, RDT=ComplexF32)
+        mul!(dm, adjoint(J*DFT), dlin)
+        dm1 = adjoint(J*DFT) * dlin
+        @test isapprox(dm1.data, dm; rtol=ftol)
+
         dm1 = M * J' * Md * dobs
         @test length(dm1) == prod(model0.n)
         @test dm1[1] == 0