Fix NormalOp weighting behaviour

Author: nHackel

src/NormalOp.jl

@@ -1,5 +1,21 @@
 export normalOperator
 
+"""
+    NormalOp(T::Type; parent, weights)
+  
+Lazy normal operator of `parent` with an optional weighting operator `weights.`
+Computes `adjoint(parent) * weights * parent`.
+
+# Required Argument
+  * `T`                        - type of elements, .e.g. `Float64` for `ComplexF32`
+
+# Required Keyword argument
+  * `parent`                   - Base operator
+
+# Optional Keyword argument
+  * `weights`                  - Optional weights for normal operator. Must already be of form `weights = adjoint.(w) .* w`
+
+"""
 function LinearOperatorCollection.NormalOp(::Type{T}; parent, weights = opEye(eltype(parent), size(parent, 1), S = storage_type(parent))) where T <: Number
   return NormalOp(T, parent, weights)
 end
@@ -47,7 +63,6 @@ function NormalOpImpl(parent, weights, tmp)
   function produ!(y, parent, weights, tmp, x)
     mul!(tmp, parent, x)
     mul!(tmp, weights, tmp) # This can be dangerous. We might need to create two tmp vectors
-    mul!(tmp, weights, tmp)
     return mul!(y, adjoint(parent), tmp)
   end
 
@@ -63,6 +78,11 @@ function Base.copy(S::NormalOpImpl)
   return NormalOpImpl(copy(S.parent), S.weights, copy(S.tmp))
 end
 
+"""
+  normalOperator(parent (, weights); kwargs...)
+
+  Constructs a normal operator of the parent in an opinionated way, i.e. it tries to apply optimisations to the resulting operator.
+"""
 function normalOperator(parent, weights=opEye(eltype(parent), size(parent, 1), S= storage_type(parent)); kwargs...)
   return NormalOp(eltype(storage_type((parent))); parent = parent, weights = weights)
 end

src/ProdOp.jl

@@ -121,7 +121,12 @@ end
 # In this case we are converting the left argument into a 
 # weighting matrix, that is passed to normalOperator
 # TODO Port vom MRIOperators drops given weighting matrix, I just left it out for now
-normalOperator(S::ProdOp{T, <:WeightingOp, matT}; kwargs...) where {T, matT} = normalOperator(S.B, S.A; kwargs...)
+"""
+    normalOperator(prod::ProdOp{T, <:WeightingOp, matT}; kwargs...)
+
+Fuses weights of `ẀeightingOp` by computing `adjoint.(weights) .* weights`
+"""
+normalOperator(S::ProdOp{T, <:WeightingOp, matT}; kwargs...) where {T, matT} = normalOperator(S.B, WeightingOp(adjoint.(S.A.weights) .* S.A.weights); kwargs...)
 function normalOperator(S::ProdOp, W=opEye(eltype(S),size(S,1), S = storage_type(S)); kwargs...)
   arrayType = storage_type(S)
   tmp = arrayType(undef, size(S.A, 2))

test/testNormalOp.jl

@@ -1,28 +1,34 @@
 
 @testset "Normal Operator" begin
   for arrayType in arrayTypes
-    @testset "$arrayType" begin 
-      N = 512
+    for elType in [Float32, ComplexF32]
+      @testset "$arrayType" begin 
+        N = 512
 
-      Random.seed!(1234)
-      x = arrayType(rand(N))
-      A = arrayType(rand(N,N))
-      A_adj = arrayType(collect(adjoint(A))) # LinearOperators can't resolve storage_type otherwise
-      W = WeightingOp(arrayType(rand(N)))
-      WA = W*A
+        Random.seed!(1234)
+        x = arrayType(rand(elType, N))
+        A = arrayType(rand(elType, N,N))
+        A_adj = arrayType(collect(adjoint(A))) # LinearOperators can't resolve storage_type otherwise
+        W = WeightingOp(arrayType(rand(elType, N)))
+        WA = W*A
+        WHW = adjoint.(W.weights) .* W.weights
+        prod = ProdOp(W, A)
 
-      y1 = Array(A_adj*W*W*A*x)
-      y2 = Array(adjoint(WA) * WA * x)
-      y = Array(normalOperator(A,W)*x)
+        y1 = Array(A_adj*adjoint(W)*W*A*x)
+        y2 = Array(adjoint(WA) * WA * x)
+        y3 = Array(normalOperator(prod) * x)
+        y4 = Array(normalOperator(A, WHW)*x)
 
-      @test norm(y1 - y) / norm(y) ≈ 0 atol=0.01
-      @test norm(y2 - y) / norm(y) ≈ 0 atol=0.01
+        @test norm(y1 - y4) / norm(y4) ≈ 0 atol=0.01
+        @test norm(y2 - y4) / norm(y4) ≈ 0 atol=0.01
+        @test norm(y3 - y4) / norm(y4) ≈ 0 atol=0.01
 
 
-      y1 = Array(adjoint(A)*A*x)
-      y = Array(normalOperator(A)*x)
+        y1 = Array(adjoint(A)*A*x)
+        y = Array(normalOperator(A)*x)
 
-      @test norm(y1 - y) / norm(y) ≈ 0 atol=0.01
+        @test norm(y1 - y) / norm(y) ≈ 0 atol=0.01
+      end
     end
   end
 end