From 83a341229900ab089267025fe89eb48edb7f0d6f Mon Sep 17 00:00:00 2001
From: RainerHeintzmann <heintzmann@gmail.com>
Date: Sun, 21 Jul 2024 08:44:03 +0200
Subject: [PATCH] bug fixes

---
 src/specific.jl  | 10 ++++-----
 test/runtests.jl | 56 ++++++++++++++++++++++++++++++++++++++++++------
 2 files changed, 54 insertions(+), 12 deletions(-)

diff --git a/src/specific.jl b/src/specific.jl
index 979ef99..52ef759 100644
--- a/src/specific.jl
+++ b/src/specific.jl
@@ -34,17 +34,17 @@ function generate_functions_expr()
             :((f, x, sz, sigma) -> .-x./sigma^2 .* f),
             :((f, x, sz, sigma) -> x.^2 ./sigma^3 .* f)
             ), 
-        (:(normal), (sigma=1.0,), :((x,sz, sigma) -> exp.(.- x.^2 ./(2*sigma^2)) ./ (sqrt(eltype(x)(2pi))*sigma)), Float32, *,
+        (:(normal), (sigma=1.0,), :((x,sz, sigma) -> exp.(.- x.^2 ./(2*sigma^2)) ./ (sqrt(eltype(x)(2pi))*abs(sigma))), Float32, *,
             :((f, x, sz, sigma) -> .-x./sigma^2 .* f),
-            :((f, x, sz, sigma) -> (x.^2 ./sigma^3 .+ 1/sigma) .* f)
+            :((f, x, sz, sigma) -> (x.^2 ./sigma^3 .- inv(sigma)) .* f)
             ),
         (:(sinc), NamedTuple(), :((x,sz) -> sinc.(x)), Float32, *,
             :((f, x, sz) -> ifelse.(x .== zero(eltype(x)), zeros(eltype(x), size(x)), (cospi.(x) .- f)./x))
             ),
         # the value "nothing" means that this default argument will not be handed over. But this works only for the last argument!
         (:(exp_ikx), (shift_by=nothing,), :((x,sz, shift_by=sz÷2) -> cis.(x.*(-eltype(x)(2pi)*shift_by/sz))), ComplexF32, *,
-            :((f, x, sz, shift_by) -> (-eltype(x)(2pi)*shift_by/sz) .* f),
-            :((f, x, sz, shift_by) -> (-eltype(x)(2pi)/sz) .*x .* f)
+            :((f, x, sz, shift_by) -> (-1im*eltype(x)(2pi)*shift_by/sz) .* f),
+            :((f, x, sz, shift_by) -> (-1im*eltype(x)(2pi)/sz) .*x .* f)
             ),
         (:(ramp), (slope=0,), :((x,sz, slope) -> slope.*x), Float32, +,
             :((f, x, sz, slope) ->  slope),
@@ -101,7 +101,7 @@ for F in generate_functions_expr()
                 # println("pb")
                 # @show dy
                 # @show $(F[6])(y, x, sz, args...; kwargs...)
-                mydx =  dy .* $(F[6])(y, x, sz, args...; kwargs...)
+                mydx =  conj.(dy) .* $(F[6])(y, x, sz, args...; kwargs...)
                 # targ = ntuple(d -> begin
                 #     mydarg = F[6+d]
                 #     dy .* $(mydarg)(y, x, sz, args...; kwargs...)
diff --git a/test/runtests.jl b/test/runtests.jl
index a1c6d51..f62b373 100644
--- a/test/runtests.jl
+++ b/test/runtests.jl
@@ -1,6 +1,9 @@
 using Test
 using IndexFunArrays
 using SeparableFunctions
+using FiniteDifferences
+using Zygote
+using Random
 
 function test_fct(T, fcts, sz, args...; kwargs...)
     ifa, fct = fcts
@@ -165,9 +168,26 @@ end
     @test maximum(abs.(res4 .- res5)) < 1e-6
 end
 
+function test_gradient(T, fct, sz, args...; kwargs...)
+    RT = real(T)
+    Random.seed!(1234)
+    dat = rand(T, sz...)
+    off0 = rand(RT, length(sz))
+    sca0 = rand(RT, length(sz))
+    args = ntuple((d)->RT.(args[d]), length(args))
+    loss = (off, sca, args...) -> sum(abs2.(fct(sz, off, sca, args..., kwargs...) .- dat))
+    # @show loss(off0, sca0, args...)
+    g = gradient(loss, off0, sca0, args...)
+    gn = grad(central_fdm(5, 1), loss, off0, sca0, args...) # 5th order method, 1st derivative
+    for r in 1:length(gn)
+        @test eltype(g[r]) == RT
+        # @show g[r]
+        # @show gn[r]
+        @test all(isapprox.(g[r], gn[r], rtol=1e-1))
+    end
+end
+
 @testset "gradient tests" begin
-    using FiniteDifferences
-    using Zygote
     rng = collect(1:0.1:2)
     sz = length(rng)
 
@@ -179,16 +199,38 @@ end
     @test g[1] ≈ gn[1]
     @test g[2] ≈ gn[2]
 
-    sz = (11, 22)
+    test_gradient(Float32, gaussian_nokw_sep, (11,22), (2.2, -0.8))
+    test_gradient(Float64, gaussian_nokw_sep, (6, 22, 7), 2.0)
+    test_gradient(Float32, gaussian_nokw_sep, (6,), 4.2f0)
+
+    test_gradient(Float64, normal_nokw_sep, (6, 22, 7), (2.0, -3.1, 1.2))
+
+    test_gradient(Float32, sinc_nokw_sep, (22, 11))
+    test_gradient(Float32, ramp_nokw_sep, (22, 11), (1.0, 2.0))
+    test_gradient(Float32, rr2_nokw_sep, (22, 11))
+    test_gradient(ComplexF32, exp_ikx_nokw_sep, (22, 11), (1.0, 2.0))
+
+    sz = (3,)
+    loss2 = (off, sca, shift_by) -> sum(imag.(exp_ikx_nokw_sep(sz, off, sca, shift_by)))
+    shift_by0 = 0.7
+    sca0 = 1.4
+    off0 = 0.3
+    loss2(off0, sca0, shift_by0)
+    g = gradient(loss2, off0, sca0, shift_by0)
+    gn = grad(central_fdm(5, 1), loss2, off0, sca0, shift_by0) # 5th order method, 1st derivative
+    @test all(isapprox.(g[1], gn[1], atol=5e-3))
+    @test all(isapprox.(g[2], gn[2], atol=1e-2))
+
+    sz = (11, 22, 7)
     loss2 = (off, sca, sigma) -> sum(gaussian_nokw_sep(sz, off, sca, sigma))
     sigma0 = 2.0
-    sca0 = (0.9, 1.2)
-    off0 = (1.1, 2.2)
+    sca0 = (0.9, 1.2, 0.4)
+    off0 = (0.9, 1.2, 0.4)
     loss2(off0, sca0, sigma0)
     g = gradient(loss2, off0, sca0, sigma0)
     gn = grad(central_fdm(5, 1), loss2, off0, sca0, sigma0) # 5th order method, 1st derivative
-    @test all(isapprox.(g[1], gn[1], atol=2e-3))
-    @test all(isapprox.(g[2], gn[2], atol=2e-3))
+    @test all(isapprox.(g[1], gn[1], atol=5e-3))
+    @test all(isapprox.(g[2], gn[2], atol=1e-2))
 
 end