Do some benchmarking

Project.toml

@@ -26,11 +26,12 @@ julia = "1.8"
 
 [extras]
 ChainRulesTestUtils = "cdddcdb0-9152-4a09-a978-84456f9df70a"
+FiniteDifferences = "26cc04aa-876d-5657-8c51-4c34ba976000"
 IndexFunArrays = "613c443e-d742-454e-bfc6-1d7f8dd76566"
 InteractiveUtils = "b77e0a4c-d291-57a0-90e8-8db25a27a240"
 Markdown = "d6f4376e-aef5-505a-96c1-9c027394607a"
 Test = "8dfed614-e22c-5e08-85e1-65c5234f0b40"
 Zygote = "e88e6eb3-aa80-5325-afca-941959d7151f"
 
 [targets]
-test = ["Test", "ChainRulesTestUtils", "Zygote", "IndexFunArrays", "Markdown", "InteractiveUtils"]
+test = ["Test", "ChainRulesTestUtils", "Zygote", "IndexFunArrays", "Markdown", "InteractiveUtils", "FiniteDifferences"]

examples/Manifest.toml

@@ -2,7 +2,7 @@
 
 julia_version = "1.9.3"
 manifest_format = "2.0"
-project_hash = "e1fb4fcb730bba1edcff897380dbaba98c9ecab6"
+project_hash = "5adba9be67adc706c964b884ad91ae500b32e8e9"
 
 [[deps.AbstractFFTs]]
 deps = ["LinearAlgebra"]
@@ -546,6 +546,12 @@ version = "2.21.1"
     BlockBandedMatrices = "ffab5731-97b5-5995-9138-79e8c1846df0"
     StaticArrays = "90137ffa-7385-5640-81b9-e52037218182"
 
+[[deps.FiniteDifferences]]
+deps = ["ChainRulesCore", "LinearAlgebra", "Printf", "Random", "Richardson", "SparseArrays", "StaticArrays"]
+git-tree-sha1 = "d77e4697046989f44dce3ed66269aaf1611a3406"
+uuid = "26cc04aa-876d-5657-8c51-4c34ba976000"
+version = "0.12.31"
+
 [[deps.FixedPointNumbers]]
 deps = ["Statistics"]
 git-tree-sha1 = "335bfdceacc84c5cdf16aadc768aa5ddfc5383cc"
@@ -1616,6 +1622,12 @@ git-tree-sha1 = "609c26951d80551620241c3d7090c71a73da75ab"
 uuid = "295af30f-e4ad-537b-8983-00126c2a3abe"
 version = "3.5.6"
 
+[[deps.Richardson]]
+deps = ["LinearAlgebra"]
+git-tree-sha1 = "48f038bfd83344065434089c2a79417f38715c41"
+uuid = "708f8203-808e-40c0-ba2d-98a6953ed40d"
+version = "1.4.2"
+
 [[deps.Rotations]]
 deps = ["LinearAlgebra", "Quaternions", "Random", "StaticArrays"]
 git-tree-sha1 = "0783924e4a332493f72490253ba4e668aeba1d73"

examples/Project.toml

@@ -4,6 +4,7 @@ Colors = "5ae59095-9a9b-59fe-a467-6f913c188581"
 ComponentArrays = "b0b7db55-cfe3-40fc-9ded-d10e2dbeff66"
 FFTW = "7a1cc6ca-52ef-59f5-83cd-3a7055c09341"
 FileIO = "5789e2e9-d7fb-5bc7-8068-2c6fae9b9549"
+FiniteDifferences = "26cc04aa-876d-5657-8c51-4c34ba976000"
 FourierTools = "b18b359b-aebc-45ac-a139-9c0ccbb2871e"
 ImageIO = "82e4d734-157c-48bb-816b-45c225c6df19"
 ImageShow = "4e3cecfd-b093-5904-9786-8bbb286a6a31"

examples/benchmark_CPU_CUDA.jl

@@ -0,0 +1,112 @@
+### A Pluto.jl notebook ###
+# v0.19.30
+
+using Markdown
+using InteractiveUtils
+
+# ╔═╡ 2fd337ac-8ad0-11ee-3739-459b5825a8c5
+begin
+	using Pkg
+	Pkg.activate(".")
+	using Revise
+end
+
+# ╔═╡ 958b4e13-bb6c-4d4d-83f9-a922bbbfb842
+using WaveOpticsPropagation, FFTW, CUDA, Zygote
+
+# ╔═╡ 62c0385d-e712-43eb-adae-6731159f9f92
+begin
+	# use CUDA if functional
+	use_CUDA = Ref(true && CUDA.functional())
+	var"@mytime" = use_CUDA[] ? CUDA.var"@time" : var"@time"   
+	togoc(x) = use_CUDA[] ? CuArray(x) : x 
+end
+
+# ╔═╡ 8ee34fde-40a7-4cc1-a8ea-056a457901b0
+md"# Compared with FFT"
+
+# ╔═╡ a43a5d7e-5052-4ee1-a112-c8881e90b6a6
+sz = (2048, 2048)
+
+# ╔═╡ c33ed1de-feb3-4bab-80c5-cd422be96bb6
+array = randn(ComplexF32, sz);
+
+# ╔═╡ 5d7cf31a-7e08-45ea-80e4-a0e27a022258
+array_c = togoc(array);
+
+# ╔═╡ 4e651fc1-b6db-4c44-b830-f65c6f68f4a9
+p_cpu = plan_fft(array, flags=FFTW.ESTIMATE);
+
+# ╔═╡ 5f5113ed-d335-4038-a230-b739242aafbc
+p_cuda = plan_fft(array_c);
+
+# ╔═╡ f897cd86-ddc1-4b84-a11c-d99161c1df3d
+@mytime p_cpu * array;
+
+# ╔═╡ fcf4d2ab-512b-4a2c-a131-9d2519da1747
+@mytime p_cpu * array;
+
+# ╔═╡ f7633864-230a-4a06-a48e-8acf4adb58f1
+@mytime CUDA.@sync p_cuda * array_c;
+
+# ╔═╡ 0a7b6bc6-9a23-484b-ac26-6effb901ea3b
+@mytime CUDA.@sync p_cuda * array_c;
+
+# ╔═╡ 181eb167-f944-480d-907a-e54a2ba113df
+md"# Propagation"
+
+# ╔═╡ 864c37c5-6907-4b31-8a4e-8c3c9a684606
+AS = Angular_Spectrum(array, 1f0, 1f0, 1f0)[1];
+
+# ╔═╡ d9339c77-35cf-4575-b935-79d580badba5
+AS_c = Angular_Spectrum(array_c, 1f0, 1f0, 1f0)[1];
+
+# ╔═╡ ad07d942-4985-40dc-9371-d6579e5b92a3
+@mytime angular_spectrum(array, 1f0, 1f0, 1f0);
+
+# ╔═╡ 127e50d1-4538-4010-b7c3-716bffd804ed
+@mytime AS(array);
+
+# ╔═╡ b0d00f6f-2a6c-42c0-b878-1791b2073353
+@mytime CUDA.@sync angular_spectrum(array_c, 1f0, 1f0, 1f0);
+
+# ╔═╡ 20935643-b699-46c0-8a44-d948551d9c44
+@mytime CUDA.@sync AS_c(array_c);
+
+# ╔═╡ a0637882-140c-466f-ab7d-af9cdb979d1f
+md"# Gradient"
+
+# ╔═╡ ca2386b8-0882-4b31-90fa-7ab1df839457
+f(x) = sum(abs2, AS_c(x)[1])
+
+# ╔═╡ 17679b22-61ac-42c2-959c-5e28fe880d15
+f(array_c)
+
+# ╔═╡ 02c3e595-b978-4ad4-961d-5d7398d3d320
+@mytime CUDA.@sync Zygote.gradient(f, array_c)
+
+# ╔═╡ Cell order:
+# ╠═2fd337ac-8ad0-11ee-3739-459b5825a8c5
+# ╠═958b4e13-bb6c-4d4d-83f9-a922bbbfb842
+# ╠═62c0385d-e712-43eb-adae-6731159f9f92
+# ╟─8ee34fde-40a7-4cc1-a8ea-056a457901b0
+# ╠═a43a5d7e-5052-4ee1-a112-c8881e90b6a6
+# ╠═c33ed1de-feb3-4bab-80c5-cd422be96bb6
+# ╠═5d7cf31a-7e08-45ea-80e4-a0e27a022258
+# ╠═4e651fc1-b6db-4c44-b830-f65c6f68f4a9
+# ╠═5f5113ed-d335-4038-a230-b739242aafbc
+# ╠═f897cd86-ddc1-4b84-a11c-d99161c1df3d
+# ╠═fcf4d2ab-512b-4a2c-a131-9d2519da1747
+# ╠═f7633864-230a-4a06-a48e-8acf4adb58f1
+# ╠═0a7b6bc6-9a23-484b-ac26-6effb901ea3b
+# ╟─181eb167-f944-480d-907a-e54a2ba113df
+# ╠═864c37c5-6907-4b31-8a4e-8c3c9a684606
+# ╠═d9339c77-35cf-4575-b935-79d580badba5
+# ╠═ad07d942-4985-40dc-9371-d6579e5b92a3
+# ╠═127e50d1-4538-4010-b7c3-716bffd804ed
+# ╠═b0d00f6f-2a6c-42c0-b878-1791b2073353
+# ╠═20935643-b699-46c0-8a44-d948551d9c44
+# ╠═a0637882-140c-466f-ab7d-af9cdb979d1f
+# ╠═ca2386b8-0882-4b31-90fa-7ab1df839457
+# ╠═17679b22-61ac-42c2-959c-5e28fe880d15
+# ╠═02c3e595-b978-4ad4-961d-5d7398d3d320

test/angular_spectrum.jl

@@ -1,18 +1,18 @@
 @testset "Angular Spectrum" begin
 
     @testset "Test gradient with Finite Differences" begin
-        field = zeros(ComplexF32, (32, 32))
+        field = zeros(ComplexF64, (32, 32))
         field[14:16, 14:16] .= 1
 
-        gg(x) = sum(abs2.(x .- angular_spectrum(x, 100e-6, 633e-9, 100e-6)[1]))
+        gg(x) = sum(abs2.(x .- angular_spectrum(cis.(x), 100e-6, 633e-9, 100e-6)[1]))
 
         out2 = FiniteDifferences.grad(central_fdm(5, 1), gg, field)[1]
 
         out1 = gradient(gg, field)[1]
-        @test out1 ≈ out2
+        @test out1 .+ cis(1) ≈ out2  .+ cis(1)
         AS, _ = Angular_Spectrum(field, 100e-6, 633e-9, 100e-6)
 
-        f_AS(x) = sum(abs2.(x .- AS(x)[1]))
+        f_AS(x) = sum(abs2.(x .- AS(cis.(x))[1]))
 
         out3 = gradient(f_AS, field)[1]