diff --git a/.gitignore b/.gitignore new file mode 100644 index 0000000..5fcfb04 --- /dev/null +++ b/.gitignore @@ -0,0 +1,29 @@ +# Files generated by invoking Julia with --code-coverage +*.jl.cov +*.jl.*.cov + +# Files generated by invoking Julia with --track-allocation +*.jl.mem + +# System-specific files and directories generated by the BinaryProvider and BinDeps packages +# They contain absolute paths specific to the host computer, and so should not be committed +deps/deps.jl +deps/build.log +deps/downloads/ +deps/usr/ +deps/src/ + +# Build artifacts for creating documentation generated by the Documenter package +docs/build/ +docs/site/ + +# File generated by Pkg, the package manager, based on a corresponding Project.toml +# It records a fixed state of all packages used by the project. As such, it should not be +# committed for packages, but should be committed for applications that require a static +# environment. +Manifest.toml + +*.swp + +*.vscode + diff --git a/LICENSE b/LICENSE new file mode 100644 index 0000000..69552e5 --- /dev/null +++ b/LICENSE @@ -0,0 +1,21 @@ +MIT License + +Copyright (c) 2023 roflmaostc and contributors + +Permission is hereby granted, free of charge, to any person obtaining a copy +of this software and associated documentation files (the "Software"), to deal +in the Software without restriction, including without limitation the rights +to use, copy, modify, merge, publish, distribute, sublicense, and/or sell +copies of the Software, and to permit persons to whom the Software is +furnished to do so, subject to the following conditions: + +The above copyright notice and this permission notice shall be included in all +copies or substantial portions of the Software. + +THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR +IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, +FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE +AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER +LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, +OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE +SOFTWARE. diff --git a/Project.toml b/Project.toml new file mode 100644 index 0000000..4a475f2 --- /dev/null +++ b/Project.toml @@ -0,0 +1,37 @@ +name = "WaveOpticsPropagation" +uuid = "c4c7a1f9-3adc-4a73-843a-d378b6c86436" +authors = ["Felix Wechsler (roflmaostc) "] +version = "0.1.0" + +[deps] +CUDA = "052768ef-5323-5732-b1bb-66c8b64840ba" +ChainRulesCore = "d360d2e6-b24c-11e9-a2a3-2a2ae2dbcce4" +EllipsisNotation = "da5c29d0-fa7d-589e-88eb-ea29b0a81949" +FFTW = "7a1cc6ca-52ef-59f5-83cd-3a7055c09341" +FourierTools = "b18b359b-aebc-45ac-a139-9c0ccbb2871e" +IndexFunArrays = "613c443e-d742-454e-bfc6-1d7f8dd76566" +NDTools = "98581153-e998-4eef-8d0d-5ec2c052313d" +Zygote = "e88e6eb3-aa80-5325-afca-941959d7151f" + +[compat] +CUDA = "4, 5" +ChainRulesCore = "1" +EllipsisNotation = "1" +FFTW = "1" +FourierTools = "0.4.1" +IndexFunArrays = "0.2" +NDTools = "0.5" +Zygote = "0.6" +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", "FiniteDifferences"] diff --git a/README.md b/README.md new file mode 100644 index 0000000..182b4ab --- /dev/null +++ b/README.md @@ -0,0 +1,48 @@ +# WaveOpticsPropagation.jl + +| **Build Status** | **Code Coverage** | **Docs stable version** | **Docs main branch** | +|:-----------------------------------------:|:-------------------------------:|:-----------------------:|:--------------------:| +| [![][CI-img]][CI-url] | [![][codecov-img]][codecov-url] | [![Documentation for stable version](https://img.shields.io/badge/docs-stable-blue.svg)](https://JuliaPhysics.github.io/WaveOpticsPropagation.jl/stable) | [![Documentation for development version](https://img.shields.io/badge/docs-main-blue.svg)](https://JuliaPhysics.github.io/WaveOpticsPropagation.jl/dev) + + + +Propagate waves efficiently, optically, physically, differentiably with [Julia Lang](https://julialang.org/). +Those functions are fast and memory efficiently implemented and hence are suited to be used in inverse problems. + +⚠️ Under heavy development. Expect things to break. + +## Installation +Not registered yet, hence install with: +```julia +julia> ]add https://github.com/JuliaPhysics/WaveOpticsPropagation.jl +``` + +## Features +### Implemented +* Propagate (electrical) fields based on wave propagation +* Propagations + * [x] Angular Spectrum Method of Plane Waves (AS) + * [x] Fraunhofer Diffraction + * [ ] [Scalable Angular Spectrum propagation](https://opg.optica.org/optica/viewmedia.cfm?uri=optica-10-11-1407&html=true) + * [ ] Fresnel Propagation with Scaling Behaviour (no priority yet, PR are welcome for that. In principle very similar to the other methods.) +* CUDA support +* Differentiable (mainly based on Zygote.jl and ChainRulesCore.jl) + +### Planned +In principle vectorial propagation in free space is just a propagation of each of the components. Right now, this is not a priority and is not implemented yet. +But of course, each vectorial component can be propagated separately. + +## Development +Contributions are very welcome! +File an [issue](https://github.com/JuliaPhysics/WaveOpticsPropagation.jl/issues) on [GitHub](https://github.com/JuliaPhysics/WaveOpticsPropagation.jl) if you encounter any problems. +Also file an issue if you want to discuss or propose features. + +## Related packages +There is the outdated [PhysicalOptics.jl](https://github.com/JuliaPhysics/PhysicalOptics.jl) which provided similar methods. +For geometrical ray tracing use [OpticSim.jl](https://github.com/brianguenter/OpticSim.jl). + +[CI-img]: https://github.com/JuliaPhysics/WaveOpticsPropagation.jl/actions/workflows/CI.yml/badge.svg +[CI-url]: https://github.com/JuliaPhysics/WaveOpticsPropagation.jl/actions/workflows/CI.yml + +[codecov-img]: https://codecov.io/gh/JuliaPhysics/WaveOpticsPropagation.jl/branch/main/graph/badge.svg?token=6XWI1M1MPB +[codecov-url]: https://codecov.io/gh/JuliaPhysics/WaveOpticsPropagation.jl diff --git a/docs/Project.toml b/docs/Project.toml new file mode 100644 index 0000000..0fca31b --- /dev/null +++ b/docs/Project.toml @@ -0,0 +1,3 @@ +[deps] +Documenter = "e30172f5-a6a5-5a46-863b-614d45cd2de4" +WaveOpticsPropagation = "c4c7a1f9-3adc-4a73-843a-d378b6c86436" diff --git a/docs/doc_server.sh b/docs/doc_server.sh new file mode 100755 index 0000000..7620444 --- /dev/null +++ b/docs/doc_server.sh @@ -0,0 +1,3 @@ +#!/bin/bash + +python3 -m http.server --bind localhost diff --git a/docs/make.jl b/docs/make.jl new file mode 100644 index 0000000..48b25c8 --- /dev/null +++ b/docs/make.jl @@ -0,0 +1,13 @@ +using WaveOpticsPropagation, Documenter + +DocMeta.setdocmeta!(WaveOpticsPropagation, :DocTestSetup, :(using WaveOpticsPropagation); recursive=true) +makedocs(modules = [WaveOpticsPropagation], + sitename = "WaveOpticsPropagation.jl", + pages = Any[ + "WaveOpticsPropagation.jl" => "index.md", + "Function Docstrings" => "functions.md" + ], + warnonly=true, + ) + +deploydocs(repo = "github.com/roflmaostc/WaveOpticsPropagation.jl.git", devbranch="main") diff --git a/docs/src/functions.md b/docs/src/functions.md new file mode 100644 index 0000000..210f501 --- /dev/null +++ b/docs/src/functions.md @@ -0,0 +1,9 @@ + +# Propagations +```@docs +AngularSpectrum +angular_spectrum +fraunhofer +Fraunhofer +ScalableAngularSpectrum +``` diff --git a/docs/src/index.md b/docs/src/index.md new file mode 100644 index 0000000..96e3ddf --- /dev/null +++ b/docs/src/index.md @@ -0,0 +1,46 @@ +# WaveOpticsPropagation.jl + +| **Build Status** | **Code Coverage** | +|:-----------------------------------------:|:-------------------------------:| +| [![][CI-img]][CI-url] | [![][codecov-img]][codecov-url] | + +Propagate waves efficiently, optically, physically, differentiably with [Julia Lang](https://julialang.org/). +Those functions are fast and memory efficiently implemented and hence are suited to be used in inverse problems. + +⚠️ Under heavy development. Expect things to break. + +## Installation +Not registered yet, hence install with: +```julia +julia> ]add https://github.com/JuliaPhysics/WaveOpticsPropagation.jl +``` + +## Features +### Implemented +* Propagate (electrical) fields based on wave propagation +* Propagations + * [x] Angular Spectrum Method of Plane Waves (AS) + * [x] Fraunhofer Diffraction + * [ ] [Scalable Angular Spectrum propagation](https://opg.optica.org/optica/viewmedia.cfm?uri=optica-10-11-1407&html=true) + * [ ] Fresnel Propagation with Scaling Behaviour (no priority yet, PR are welcome for that. In principle very similar to the other methods.) +* CUDA support +* Differentiable (mainly based on Zygote.jl and ChainRulesCore.jl) + +### Planned +In principle vectorial propagation in free space is just a propagation of each of the components. Right now, this is not a priority and is not implemented yet. +But of course, each vectorial component can be propagated separately. + +## Development +Contributions are very welcome! +File an [issue](https://github.com/roflmaostc/RadonKA.jl/issues) on [GitHub](https://github.com/roflmaostc/RadonKA.jl) if you encounter any problems. +Also file an issue if you want to discuss or propose features. + +## Related packages +There is the outdated [PhysicalOptics.jl](https://github.com/JuliaPhysics/PhysicalOptics.jl) which provided similar methods. +For geometrical ray tracing use [OpticSim.jl](https://github.com/brianguenter/OpticSim.jl). + +[CI-img]: https://github.com/JuliaPhysics/WaveOpticsPropagation.jl/actions/workflows/CI.yml/badge.svg +[CI-url]: https://github.com/JuliaPhysics/WaveOpticsPropagation.jl/actions/workflows/CI.yml + +[codecov-img]: https://codecov.io/gh/JuliaPhysics/WaveOpticsPropagation.jl/branch/main/graph/badge.svg?token=6XWI1M1MPB +[codecov-url]: https://codecov.io/gh/JuliaPhysics/WaveOpticsPropagation.jl diff --git a/examples/Project.toml b/examples/Project.toml new file mode 100644 index 0000000..12b1cf4 --- /dev/null +++ b/examples/Project.toml @@ -0,0 +1,22 @@ +[deps] +CUDA = "052768ef-5323-5732-b1bb-66c8b64840ba" +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" +IndexFunArrays = "613c443e-d742-454e-bfc6-1d7f8dd76566" +NDTools = "98581153-e998-4eef-8d0d-5ec2c052313d" +Napari = "ed46c112-542a-4e7e-ab0f-e39321fab6f0" +Optim = "429524aa-4258-5aef-a3af-852621145aeb" +Plots = "91a5bcdd-55d7-5caf-9e0b-520d859cae80" +Pluto = "c3e4b0f8-55cb-11ea-2926-15256bba5781" +PlutoUI = "7f904dfe-b85e-4ff6-b463-dae2292396a8" +PyPlot = "d330b81b-6aea-500a-939a-2ce795aea3ee" +Revise = "295af30f-e4ad-537b-8983-00126c2a3abe" +TestImages = "5e47fb64-e119-507b-a336-dd2b206d9990" +WaveOpticsPropagation = "c4c7a1f9-3adc-4a73-843a-d378b6c86436" +Zygote = "e88e6eb3-aa80-5325-afca-941959d7151f" diff --git a/examples/Scalable_Angular_Spectrum.html b/examples/Scalable_Angular_Spectrum.html new file mode 100644 index 0000000..f241440 --- /dev/null +++ b/examples/Scalable_Angular_Spectrum.html @@ -0,0 +1,16 @@ + + + + + + +
\ No newline at end of file diff --git a/examples/Scalable_Angular_Spectrum.jl b/examples/Scalable_Angular_Spectrum.jl new file mode 100644 index 0000000..e798838 --- /dev/null +++ b/examples/Scalable_Angular_Spectrum.jl @@ -0,0 +1,182 @@ +### A Pluto.jl notebook ### +# v0.19.30 + +using Markdown +using InteractiveUtils + +# This Pluto notebook uses @bind for interactivity. When running this notebook outside of Pluto, the following 'mock version' of @bind gives bound variables a default value (instead of an error). +macro bind(def, element) + quote + local iv = try Base.loaded_modules[Base.PkgId(Base.UUID("6e696c72-6542-2067-7265-42206c756150"), "AbstractPlutoDingetjes")].Bonds.initial_value catch; b -> missing; end + local el = $(esc(element)) + global $(esc(def)) = Core.applicable(Base.get, el) ? Base.get(el) : iv(el) + el + end +end + +# ╔═╡ 1ea0729c-ad94-11ee-3b52-bb21d59e9509 +begin + using Pkg + Pkg.activate(".") + Pkg.instantiate() + using Revise +end + +# ╔═╡ aafab549-c977-48f4-9dba-9d693e439b44 +using WaveOpticsPropagation, CUDA, IndexFunArrays, NDTools, ImageShow, FileIO, Zygote, Optim, Plots, PlutoUI, FourierTools, NDTools + +# ╔═╡ 4aabd650-792c-407c-9590-e4960ac776d6 +md"# Setup Environment and packages" + +# ╔═╡ 3ed16b3f-08a8-4fc3-97cf-a884c452dff0 +TableOfContents() + +# ╔═╡ 20734e1c-faf7-4117-8652-528eca92c6a2 +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 + +# ╔═╡ 7028f15f-7ec7-48b3-9b4d-66bd0c34bae6 +λ = 500f-9 + +# ╔═╡ f9869ff6-5b21-40a3-98cf-33e7192dfa8e +L = 128f-6 / 2 + +# ╔═╡ 1666eaa6-7857-415e-a43e-00849571db2a +N = 512 + +# ╔═╡ d418e87f-675f-4bc2-8eaf-70cf669aeb8d +y = fftpos(L, N, NDTools.CenterFT) + +# ╔═╡ 945ac45e-2773-4d34-878e-4f66a253afa1 +D_circ = N / 8 + +# ╔═╡ 4f7aca22-00d9-406c-90e1-f73394a718ee +U_circ = togoc(ComplexF32.(rr((N, N)) .< D_circ / 2) .* exp.(1im .* 2f0 * π ./ λ .* y .* sind(45f0)) .+ ComplexF32.(rr((N, N)) .< D_circ / 2) .* exp.(1im .* 2f0 * π ./ λ .* y' .* sind(-45f0))); + +# ╔═╡ 34c326ee-6663-432e-b884-a5419ce64827 +@bind M Slider(1:0.1:20, show_value=true) + +# ╔═╡ 4202033d-62e6-451f-a064-e61d420da5ff +z_circ = Float32(M / N / λ * L^2 * 2) + +# ╔═╡ 4dc0d972-88b7-4cb4-9ef3-925b154e45c0 +md"# First Example: Circular + +In the first example, one straight beam and one oblique beam are passing through a round aperture. + +The Fresnel number is $(round((D_circ / 2 * L / N)^2 / z_circ / λ, digits=3)) +" + +# ╔═╡ e5c50250-e3c0-4561-842a-d6c671999741 +simshow(Array(U_circ)) + +# ╔═╡ 484704ce-4497-4567-b387-aad06dee6a62 +@mytime SAS = ScalableAngularSpectrum(U_circ, z_circ, λ, L) + +# ╔═╡ 9e3f7e53-3fdb-4638-89bc-27bced937193 +# ╠═╡ disabled = true +#=╠═╡ +@time SAS_cpu = ScalableAngularSpectrum(Array(U_circ), z_circ, λ, L) + ╠═╡ =# + +# ╔═╡ 9ec9d456-a33b-4605-9025-c82676eca7e2 +U_circ_cpu = Array(U_circ); + +# ╔═╡ cf8d99c6-120d-45a6-9fcc-9d1cb6e40a9b +# ╠═╡ disabled = true +#=╠═╡ +@time SAS_cpu(U_circ_cpu); + ╠═╡ =# + +# ╔═╡ 342b000d-8865-4cd2-bb6b-2de0f6376e9c +@mytime CUDA.@sync U_p, t = SAS(U_circ) + +# ╔═╡ fc7dd2c9-7c7b-45a7-9fed-a319a62375cc +simshow(Array(abs2.(U_p)), γ=0.13, cmap=:inferno) + +# ╔═╡ 505bd6fa-2717-4297-b8b5-bf9452300655 +sum(abs2, U_p) + +# ╔═╡ bb1c90d3-62c0-4941-b8ac-00761856a3cd +sum(abs2, U_circ) + +# ╔═╡ 70afe821-c31c-4626-8b80-20695a46544f +L_box = 128f-6; + +# ╔═╡ cd2bdc4a-7abe-4184-a9c3-83b9820cf46d +N_box = 512; + +# ╔═╡ e227181f-017a-4601-914a-f63583755105 +y_box = fftpos(L_box, N_box, NDTools.CenterFT); + +# ╔═╡ f7926732-1e17-428b-9660-2bf3653a3f4a +D_box = L_box / 16f0 + +# ╔═╡ 4a30f342-cad6-46d5-a787-2144909c8765 +x_box = y_box'; + +# ╔═╡ f0cf4b3d-13fd-48d0-9890-1a8e21686a7e +U_box = togoc((x_box.^2 .<= (D_box / 2).^2) .* (y_box.^2 .<= (D_box / 2).^2) .* (exp.(1im .* 2f0 * π ./ λ .* y_box' .* sind(20f0)))); + +# ╔═╡ 11947eea-b2eb-4168-bc96-8e1fc7b74599 +@bind M_box Slider(1:0.5:20, show_value=true) + +# ╔═╡ aacc5f18-05b5-4a5e-9217-fa62966830e5 +z_box = M_box / N_box / λ * L_box^2 * 2 + +# ╔═╡ d325fa62-7895-44a9-87de-718e9d61f9bc +md"# Second Example: Quadratic + + +The Fresnel number is $(round((D_box)^2 / z_box / λ, digits=3)) +" + +# ╔═╡ 38bc91ff-3188-4129-aa6b-2af458cf59b1 +@mytime SAS2 = ScalableAngularSpectrum(U_box, z_box, λ, L_box, skip_final_phase=true) + +# ╔═╡ 2b11c87d-4746-4be6-81f6-f004d30beae4 +@mytime CUDA.@sync U_box_p, t_box = SAS2(U_box) + +# ╔═╡ 9e762dfc-4873-4fe3-b785-f6c4808417a6 +simshow(Array(abs2.(U_box_p)), γ=0.13, cmap=:inferno) + +# ╔═╡ Cell order: +# ╟─4aabd650-792c-407c-9590-e4960ac776d6 +# ╠═1ea0729c-ad94-11ee-3b52-bb21d59e9509 +# ╠═3ed16b3f-08a8-4fc3-97cf-a884c452dff0 +# ╠═aafab549-c977-48f4-9dba-9d693e439b44 +# ╠═20734e1c-faf7-4117-8652-528eca92c6a2 +# ╠═4dc0d972-88b7-4cb4-9ef3-925b154e45c0 +# ╠═7028f15f-7ec7-48b3-9b4d-66bd0c34bae6 +# ╠═f9869ff6-5b21-40a3-98cf-33e7192dfa8e +# ╠═1666eaa6-7857-415e-a43e-00849571db2a +# ╠═d418e87f-675f-4bc2-8eaf-70cf669aeb8d +# ╠═945ac45e-2773-4d34-878e-4f66a253afa1 +# ╠═4f7aca22-00d9-406c-90e1-f73394a718ee +# ╠═34c326ee-6663-432e-b884-a5419ce64827 +# ╠═4202033d-62e6-451f-a064-e61d420da5ff +# ╠═e5c50250-e3c0-4561-842a-d6c671999741 +# ╠═484704ce-4497-4567-b387-aad06dee6a62 +# ╠═9e3f7e53-3fdb-4638-89bc-27bced937193 +# ╠═9ec9d456-a33b-4605-9025-c82676eca7e2 +# ╠═cf8d99c6-120d-45a6-9fcc-9d1cb6e40a9b +# ╠═342b000d-8865-4cd2-bb6b-2de0f6376e9c +# ╠═fc7dd2c9-7c7b-45a7-9fed-a319a62375cc +# ╠═505bd6fa-2717-4297-b8b5-bf9452300655 +# ╠═bb1c90d3-62c0-4941-b8ac-00761856a3cd +# ╟─d325fa62-7895-44a9-87de-718e9d61f9bc +# ╠═70afe821-c31c-4626-8b80-20695a46544f +# ╠═cd2bdc4a-7abe-4184-a9c3-83b9820cf46d +# ╠═e227181f-017a-4601-914a-f63583755105 +# ╠═f7926732-1e17-428b-9660-2bf3653a3f4a +# ╠═4a30f342-cad6-46d5-a787-2144909c8765 +# ╠═f0cf4b3d-13fd-48d0-9890-1a8e21686a7e +# ╠═11947eea-b2eb-4168-bc96-8e1fc7b74599 +# ╠═aacc5f18-05b5-4a5e-9217-fa62966830e5 +# ╠═38bc91ff-3188-4129-aa6b-2af458cf59b1 +# ╠═2b11c87d-4746-4be6-81f6-f004d30beae4 +# ╠═9e762dfc-4873-4fe3-b785-f6c4808417a6 diff --git a/examples/Scalable_Angular_Spectrum.plutostate b/examples/Scalable_Angular_Spectrum.plutostate new file mode 100644 index 0000000..5dabfb0 Binary files /dev/null and b/examples/Scalable_Angular_Spectrum.plutostate differ diff --git a/examples/benchmark_CPU_CUDA.html b/examples/benchmark_CPU_CUDA.html new file mode 100644 index 0000000..828d2c6 --- /dev/null +++ b/examples/benchmark_CPU_CUDA.html @@ -0,0 +1,16 @@ + + + + + + +
\ No newline at end of file diff --git a/examples/benchmark_CPU_CUDA.jl b/examples/benchmark_CPU_CUDA.jl new file mode 100644 index 0000000..9117205 --- /dev/null +++ b/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 = AngularSpectrum(array, 1f0, 1f0, 1f0)[1]; + +# ╔═╡ d9339c77-35cf-4575-b935-79d580badba5 +AS_c = AngularSpectrum(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 diff --git a/examples/benchmark_CPU_CUDA.plutostate b/examples/benchmark_CPU_CUDA.plutostate new file mode 100644 index 0000000..d6fe287 Binary files /dev/null and b/examples/benchmark_CPU_CUDA.plutostate differ diff --git a/examples/double_slit.html b/examples/double_slit.html new file mode 100644 index 0000000..04f4707 --- /dev/null +++ b/examples/double_slit.html @@ -0,0 +1,16 @@ + + + + + + +
\ No newline at end of file diff --git a/examples/double_slit.jl b/examples/double_slit.jl new file mode 100644 index 0000000..93c5f4a --- /dev/null +++ b/examples/double_slit.jl @@ -0,0 +1,266 @@ +### A Pluto.jl notebook ### +# v0.19.30 + +using Markdown +using InteractiveUtils + +# This Pluto notebook uses @bind for interactivity. When running this notebook outside of Pluto, the following 'mock version' of @bind gives bound variables a default value (instead of an error). +macro bind(def, element) + quote + local iv = try Base.loaded_modules[Base.PkgId(Base.UUID("6e696c72-6542-2067-7265-42206c756150"), "AbstractPlutoDingetjes")].Bonds.initial_value catch; b -> missing; end + local el = $(esc(element)) + global $(esc(def)) = Core.applicable(Base.get, el) ? Base.get(el) : iv(el) + el + end +end + +# ╔═╡ 75796f02-ac0d-11ee-3913-11fa1f818dea +begin + using Pkg + Pkg.activate(".") + Pkg.instantiate() + using Revise +end + +# ╔═╡ 9896f57c-0540-4be7-91fb-43f08f9e084c +using WaveOpticsPropagation, CUDA, IndexFunArrays, NDTools, ImageShow, FileIO, Zygote, Optim, Plots, PlutoUI + +# ╔═╡ 5a88b44b-ecce-4002-a4ba-87e8b9d55c81 +md"# Load and initialize packages" + +# ╔═╡ 0c95865d-7998-4191-a637-457c8d04d6d4 +TableOfContents() + +# ╔═╡ 6dc592f5-7fd3-4782-a72d-23b5dbfa80e2 +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 + +# ╔═╡ aa9e7ca4-945a-4b8c-bc8b-29ab192350e3 +md"# Initialize physical parameters" + +# ╔═╡ a777cc7a-c79c-494c-9a70-83a90b2ee08d +L = 100f-3 + +# ╔═╡ 42ae161e-26b7-4d34-9e92-4c4af27fc447 +λ = 633f-9 + +# ╔═╡ 9facb2bf-f279-41de-8858-81dac7a032b5 +z = 1 + +# ╔═╡ e714ec1e-cceb-4e43-aa22-fa1f69d6c816 +L_new + +# ╔═╡ 3cde3fbe-9641-4ffd-9630-9c7bf167ec34 +N = 256 + +# ╔═╡ 771642e4-c03e-41a2-a4de-07fdfd4bd535 +x = WaveOpticsPropagation.fftpos(L, N) + +# ╔═╡ 7f79a274-8f9c-43b5-bb7e-96bb495fb044 +L / N + +# ╔═╡ ac7c756a-b182-4458-b5ea-017e23fc1009 +2 / N * L + +# ╔═╡ 6acb2991-cba5-4cf5-8d82-099f9ad2335a +@bind ΔS Slider(0:N÷4, show_value=true, default=12) + +# ╔═╡ e5fc1e39-9300-4c96-b747-d532245cfd97 +S = x[1 + ΔS * 2] - x[1] + +# ╔═╡ 4f9c514a-5f24-4d49-8a92-88f250d9ef92 +@bind ΔW Slider(0:N÷2, show_value=true, default=2) + +# ╔═╡ d1734ddf-ec04-4134-82a8-187e697c4810 +begin + slit = zeros(ComplexF32, (N, N)) + + mid = N ÷ 2+ 1 + slit[:, mid-ΔS-ΔW:mid-ΔS+ΔW] .= 1 + slit[:, mid+ΔS-ΔW:mid+ΔS+ΔW] .= 1 +end; + +# ╔═╡ 3ebcb7d8-0ed6-497a-95e3-6e40bfadd8d0 +simshow(slit) + +# ╔═╡ ac72f69a-935e-4ba4-b8bd-e16f27367acd +W = x[1 + ΔW * 2 + 1] - x[1] + +# ╔═╡ 3330fef8-e6ed-412a-9321-9dc705ce41fe +md"# Propagate with Fraunhofer Diffraction" + +# ╔═╡ 7e797b20-77ee-4888-9b23-74bb2713912f +@time output, t = fraunhofer(slit, z, λ, L) + +# ╔═╡ 0c74a49c-e4de-4172-ac75-4c2692c505fb +# creating this function is more efficient! +efficient_fraunhofer = Fraunhofer(slit, z, λ, L); + +# ╔═╡ 0e6f74f1-8723-4596-b754-973b253443a4 +@time output2, t2 = efficient_fraunhofer(slit) + +# ╔═╡ 5c0486d2-5fbd-4cd7-9fc5-583a8e188b2e +begin + intensity = abs2.(output) + intensity ./= maximum(intensity) +end; + +# ╔═╡ cd54d590-ea6f-4d3b-84ec-6f983a97c81c +simshow(intensity, γ=1) + +# ╔═╡ 5f3c7de8-a64e-42aa-b8a1-91f5869c9610 +md"# Compare to analytical solution" + +# ╔═╡ 4d836453-7e64-4f3b-9405-d3cdafc264dc +I_analytical(x) = sinc(W * x / λ / z)^2 * cos(π * S * x / λ / z)^2 + +# ╔═╡ 42e1a449-0e5c-4d67-a90c-10a54f479c8a +xpos_out = WaveOpticsPropagation.fftpos(t.L, N, NDTools.CenterFT) + +# ╔═╡ 5d36ba4d-e9fe-4b1b-b45d-c82b598dfe7a +begin + plot(xpos_out, intensity[(begin+end)÷2+1, :]) + plot!(xpos_out, I_analytical.(xpos_out)) +end + +# ╔═╡ ba6965bf-c7f1-4dfd-8381-e330582b3462 +all(≈(I_analytical.(xpos_out)[110:150] .+ 1, 1 .+ intensity[129, 110:150, :], rtol=1f-2)) + +# ╔═╡ 8193fd81-1580-4047-ac04-7c3837975c2d +md"# Optimization +We try to find the initial field from one diffraction pattern. + +We start with a initial guess of a double slit where distance and width are wrong. By using Fraunhofer as forward model, we are able to optimize for correct distance and width! +" + +# ╔═╡ 2c28b256-f60c-48b6-a34a-d2908c979e30 +intensity_cu = togoc(intensity); + +# ╔═╡ 3a19cbfe-6efc-45f7-9aa3-f319f34a534e +begin + rec0 = togoc(zeros(Float32, (N, N))); + + rec0[:, mid-40:mid-10] .= 1 + rec0[:, mid+10:mid+60] .= 1 +end; + +# ╔═╡ b19225e0-5536-4af1-bbbc-c6124e91536b +simshow(abs2.(Array(rec0))) + +# ╔═╡ e0aa0714-08e0-46f4-b6d4-8c5df044ddd9 +md"## Define Forward model and gradient +We use both methods: `fraunhofer` and `Fraunhofer`. +In principle `Fraunhofer` generates some buffers and stores the FFT plan and should be slightly faster. +" + +# ╔═╡ c8a3d8b0-f68d-4c38-ae1a-8081846bda23 +fwd = let z=z, L=L, λ=λ + x -> fraunhofer(x .+ 0im, z, λ, L) +end + +# ╔═╡ c9651b37-d911-4230-a2f5-f94e8f726fbc +fwd2 = let z=z, L=L, λ=λ, fr=Fraunhofer(rec0 .+ 0im, z, λ, L) + x -> fr(x .+ 0im) +end + +# ╔═╡ e5c3f3dc-5eb3-421e-90ab-a16ea73e5621 +f = let intensity_cu=intensity_cu, fwd=fwd + f(x) = sum(abs2, intensity_cu .- abs2.(fwd(abs2.(x) .+ 0im)[1])) +end + +# ╔═╡ cb0b276f-d479-4024-8ca4-116a210fe2ed +f2 = let intensity_cu=intensity_cu, fwd2=fwd2 + f(x) = sum(abs2, intensity_cu .- abs2.(fwd2(abs2.(x) .+ 0im)[1])) +end + +# ╔═╡ 30555d35-04a8-44b7-afd7-345730d97a61 +@mytime f(rec0); + +# ╔═╡ c6c30873-9419-4a54-a9f5-6bf883856102 +@mytime f2(rec0); + +# ╔═╡ 8ab5cfd9-15e0-4be9-9936-cd499555fbec +@mytime gradient(f, rec0) + +# ╔═╡ dd8a5b39-b28a-4b79-857f-8c0a8923f2b3 +@mytime gradient(f2, rec0); + +# ╔═╡ c9a37a37-eef0-498e-8105-dd29284bfa4e +g!(G, x) = G .= Zygote.gradient(f, x)[1] + +# ╔═╡ 6795095b-bbbe-4dec-a79b-12d604873474 +g2!(G, x) = G .= Zygote.gradient(f2, x)[1] + +# ╔═╡ b75deb02-31e2-4b3d-adad-2f380b107e1b +md"## Optimize" + +# ╔═╡ c539d936-acda-4970-9f5d-92538fda2bc5 +res = optimize(f, g!, rec0, LBFGS(), Optim.Options(iterations=400)) + +# ╔═╡ 9563f911-b861-4cb1-87bc-fbbb3a3b77d2 +md"## Results +Up to a shift we obtain the correct result! The shift is not possible to find out since we only measure the intensity +" + +# ╔═╡ 8c53353c-7339-4060-8d2e-caf8f3025b48 +simshow(abs2.(Array(res.minimizer))) + +# ╔═╡ ae232c00-b4b7-41ff-8c49-a55382743956 +simshow(abs2.(Array(slit))) + +# ╔═╡ Cell order: +# ╟─5a88b44b-ecce-4002-a4ba-87e8b9d55c81 +# ╠═75796f02-ac0d-11ee-3913-11fa1f818dea +# ╟─0c95865d-7998-4191-a637-457c8d04d6d4 +# ╠═9896f57c-0540-4be7-91fb-43f08f9e084c +# ╠═6dc592f5-7fd3-4782-a72d-23b5dbfa80e2 +# ╟─aa9e7ca4-945a-4b8c-bc8b-29ab192350e3 +# ╠═a777cc7a-c79c-494c-9a70-83a90b2ee08d +# ╠═42ae161e-26b7-4d34-9e92-4c4af27fc447 +# ╠═9facb2bf-f279-41de-8858-81dac7a032b5 +# ╠═e714ec1e-cceb-4e43-aa22-fa1f69d6c816 +# ╠═3cde3fbe-9641-4ffd-9630-9c7bf167ec34 +# ╠═d1734ddf-ec04-4134-82a8-187e697c4810 +# ╠═3ebcb7d8-0ed6-497a-95e3-6e40bfadd8d0 +# ╠═771642e4-c03e-41a2-a4de-07fdfd4bd535 +# ╠═ac72f69a-935e-4ba4-b8bd-e16f27367acd +# ╠═e5fc1e39-9300-4c96-b747-d532245cfd97 +# ╠═7f79a274-8f9c-43b5-bb7e-96bb495fb044 +# ╠═ac7c756a-b182-4458-b5ea-017e23fc1009 +# ╠═6acb2991-cba5-4cf5-8d82-099f9ad2335a +# ╠═4f9c514a-5f24-4d49-8a92-88f250d9ef92 +# ╟─3330fef8-e6ed-412a-9321-9dc705ce41fe +# ╠═7e797b20-77ee-4888-9b23-74bb2713912f +# ╠═0c74a49c-e4de-4172-ac75-4c2692c505fb +# ╠═0e6f74f1-8723-4596-b754-973b253443a4 +# ╠═5c0486d2-5fbd-4cd7-9fc5-583a8e188b2e +# ╠═cd54d590-ea6f-4d3b-84ec-6f983a97c81c +# ╟─5f3c7de8-a64e-42aa-b8a1-91f5869c9610 +# ╠═4d836453-7e64-4f3b-9405-d3cdafc264dc +# ╠═42e1a449-0e5c-4d67-a90c-10a54f479c8a +# ╠═5d36ba4d-e9fe-4b1b-b45d-c82b598dfe7a +# ╠═ba6965bf-c7f1-4dfd-8381-e330582b3462 +# ╟─8193fd81-1580-4047-ac04-7c3837975c2d +# ╠═2c28b256-f60c-48b6-a34a-d2908c979e30 +# ╠═3a19cbfe-6efc-45f7-9aa3-f319f34a534e +# ╠═b19225e0-5536-4af1-bbbc-c6124e91536b +# ╠═e0aa0714-08e0-46f4-b6d4-8c5df044ddd9 +# ╠═c8a3d8b0-f68d-4c38-ae1a-8081846bda23 +# ╠═c9651b37-d911-4230-a2f5-f94e8f726fbc +# ╠═e5c3f3dc-5eb3-421e-90ab-a16ea73e5621 +# ╠═cb0b276f-d479-4024-8ca4-116a210fe2ed +# ╠═30555d35-04a8-44b7-afd7-345730d97a61 +# ╠═c6c30873-9419-4a54-a9f5-6bf883856102 +# ╠═8ab5cfd9-15e0-4be9-9936-cd499555fbec +# ╠═dd8a5b39-b28a-4b79-857f-8c0a8923f2b3 +# ╠═c9a37a37-eef0-498e-8105-dd29284bfa4e +# ╠═6795095b-bbbe-4dec-a79b-12d604873474 +# ╟─b75deb02-31e2-4b3d-adad-2f380b107e1b +# ╠═c539d936-acda-4970-9f5d-92538fda2bc5 +# ╟─9563f911-b861-4cb1-87bc-fbbb3a3b77d2 +# ╠═8c53353c-7339-4060-8d2e-caf8f3025b48 +# ╠═ae232c00-b4b7-41ff-8c49-a55382743956 diff --git a/examples/double_slit.plutostate b/examples/double_slit.plutostate new file mode 100644 index 0000000..53efe3e Binary files /dev/null and b/examples/double_slit.plutostate differ diff --git a/examples/old/Project.toml b/examples/old/Project.toml new file mode 100644 index 0000000..81648c0 --- /dev/null +++ b/examples/old/Project.toml @@ -0,0 +1 @@ +[deps] diff --git a/examples/old/gaussian_beam_SLM_propagation.html b/examples/old/gaussian_beam_SLM_propagation.html new file mode 100644 index 0000000..217ca90 --- /dev/null +++ b/examples/old/gaussian_beam_SLM_propagation.html @@ -0,0 +1,16 @@ + + + + + + +
\ No newline at end of file diff --git a/examples/old/gaussian_beam_SLM_propagation.jl b/examples/old/gaussian_beam_SLM_propagation.jl new file mode 100644 index 0000000..724f6ca --- /dev/null +++ b/examples/old/gaussian_beam_SLM_propagation.jl @@ -0,0 +1,184 @@ +### A Pluto.jl notebook ### +# v0.19.30 + +using Markdown +using InteractiveUtils + +# This Pluto notebook uses @bind for interactivity. When running this notebook outside of Pluto, the following 'mock version' of @bind gives bound variables a default value (instead of an error). +macro bind(def, element) + quote + local iv = try Base.loaded_modules[Base.PkgId(Base.UUID("6e696c72-6542-2067-7265-42206c756150"), "AbstractPlutoDingetjes")].Bonds.initial_value catch; b -> missing; end + local el = $(esc(element)) + global $(esc(def)) = Core.applicable(Base.get, el) ? Base.get(el) : iv(el) + el + end +end + +# ╔═╡ 97891d2b-facd-4f51-9774-5afde6a874e9 +begin + using Pkg + Pkg.activate(".") + Pkg.instantiate() + using Revise +end + +# ╔═╡ 034e96b2-c88a-11ed-35f3-9fbfcefa62df +using WaveOpticsPropagation, ImageShow, FFTW, CUDA, FourierTools, NDTools, Plots, Colors + +# ╔═╡ 8082f7f3-3f47-4479-b3fa-20b2eb43ee78 +using LinearAlgebra + +# ╔═╡ b53bc181-42e6-410f-99fb-6fa28b9c4f09 +using IndexFunArrays + +# ╔═╡ f6d1a0dd-431f-4cbd-ae3c-3f031dc2beb5 +using PlutoUI + +# ╔═╡ febd1f55-b327-4e31-975b-afb635a82376 +using TestImages + +# ╔═╡ dfc85871-6697-44e1-a28a-86e7314a7d29 +FFTW.set_num_threads(4) + +# ╔═╡ 6e4cb613-6cec-4eb0-93a9-bc93b83bf755 +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 + toc(x) = x + toc(x::CuArray) = Array(x) + toc(x::LinearAlgebra.Adjoint{T, <:CuArray}) where T = Array(x) + +end + +# ╔═╡ 3dd7717c-0cbb-45f9-b03b-9316d37fd635 +ImageShow.simshow(x::CuArray; kwargs...) = simshow(toc(x); kwargs...) + +# ╔═╡ 15f89bfc-34c3-408e-9577-b79e88aea179 + + +# ╔═╡ 7e73e1e3-e881-4025-b88a-bb7b12f85e02 +N = 2000 + +# ╔═╡ e362673e-dbe9-4ad6-ba7a-d1f22da3f191 +field = togoc(zeros(ComplexF32, (N, N))) + +# ╔═╡ d6e11fcc-18ef-4c64-82b8-20abbb26d339 +function dmd_pattern() + dmd = togoc(zeros(ComplexF32, 768, 1024)); + + dmd[300:400, 300:400] .= 1 + dmd[390:410, 400:500] .= 1 + + dmd[450:451, 450:451] .= 1 + dmd[7 .+ (450:451), 450:451] .= 1 + + return dmd +end + +# ╔═╡ f0cd85d6-b8a4-4be5-8e99-0eebadc0d27b +simshow(dmd_pattern()) + +# ╔═╡ 5b586b1e-e785-4980-815f-0534f997ea41 +λ = 405f-9 / 1.5 + +# ╔═╡ 3ee91f4c-83a5-4e33-948f-88b36cf1828c +L = 14.0f-3 * 2000 / 1024 * 2 + +# ╔═╡ 25916008-4a44-4140-9558-6f055d01b101 +L / 2000 + +# ╔═╡ 49f045dd-b764-4e37-abc4-271a87db4677 +y = togoc(fftpos(L, N, CenterFT)) + +# ╔═╡ 2c301c7f-bc79-4141-8ae1-3022bfc54c66 +x = y'; + +# ╔═╡ b083214a-645f-499f-8aaa-e5511cbe64d8 +L ./ size(dmd) + +# ╔═╡ a11d219e-f6d2-4137-b0b7-be2ffa589893 +field_d = set_center!(copy(field), dmd_pattern()) .* cispi.(togoc(0.0000f0 .* rr2(field))); + +# ╔═╡ 5fc8ecab-6a7c-467f-a67d-d4ba7cd0b2e2 +simshow(field_d) + +# ╔═╡ 641e4caa-3d4a-46b1-af8f-a6ef8d246723 +@mytime field_p = angular_spectrum(field_d, 16f-3, λ, L)[1] + +# ╔═╡ b6db3c81-e960-4a29-b955-43721b3af769 +@bind iz PlutoUI.Slider(1:2, show_value=true) + +# ╔═╡ 3df00e4a-89f1-4ef3-a101-e02113fa471d +simshow(abs2.([toc(field_d);;; toc(field_p)][:, :, iz])) + +# ╔═╡ 6f31315a-f545-4989-91c9-b58db0e4ba18 +@view_image abs2.([toc(field_d);;; toc(field_p)]) + +# ╔═╡ 8e4c91eb-ed58-4d1f-9a43-48d2c262e4ea +begin + plot(abs.(resample(toc(field_d)[1000:1100, 938], (3000,)))) + plot!(abs2.(resample(toc(field_p)[1000:1100, 938], (3000,)))) +end; + +# ╔═╡ 661541a4-84f4-4f74-bbe7-3aab62fb52ce +begin + plot(abs2.(toc(field_d)[1000:1100, 938])) + plot!(abs2.(toc(field_p)[1000:1100, 938])) +end + +# ╔═╡ 3b90e7f5-daef-4181-81b7-db3d0a9a922b +z = range(100f-6, 10f-3, 5) + +# ╔═╡ 09ca6d0b-fbe0-4ff4-a900-2dd2179ab6bd +L / 1024 + +# ╔═╡ 8397198e-3b49-49db-af87-a932d4967a63 +heatmap(toc(x), toc(y), Float32.(simshow(abs2.(toc(field_p)), γ=0.4))) + +# ╔═╡ 3ba2a0a2-6535-464e-ae83-cddd68bff413 +toc(x) + +# ╔═╡ 0f232dcc-faf8-46d6-87bb-4fe318964968 +toc(x) + +# ╔═╡ 57cfdcf2-99c0-4db2-a740-ed74994116c9 +typeof(x) + +# ╔═╡ Cell order: +# ╠═97891d2b-facd-4f51-9774-5afde6a874e9 +# ╠═034e96b2-c88a-11ed-35f3-9fbfcefa62df +# ╠═8082f7f3-3f47-4479-b3fa-20b2eb43ee78 +# ╠═b53bc181-42e6-410f-99fb-6fa28b9c4f09 +# ╠═f6d1a0dd-431f-4cbd-ae3c-3f031dc2beb5 +# ╠═febd1f55-b327-4e31-975b-afb635a82376 +# ╠═3dd7717c-0cbb-45f9-b03b-9316d37fd635 +# ╠═dfc85871-6697-44e1-a28a-86e7314a7d29 +# ╠═6e4cb613-6cec-4eb0-93a9-bc93b83bf755 +# ╠═15f89bfc-34c3-408e-9577-b79e88aea179 +# ╠═7e73e1e3-e881-4025-b88a-bb7b12f85e02 +# ╠═e362673e-dbe9-4ad6-ba7a-d1f22da3f191 +# ╠═d6e11fcc-18ef-4c64-82b8-20abbb26d339 +# ╠═f0cd85d6-b8a4-4be5-8e99-0eebadc0d27b +# ╠═5b586b1e-e785-4980-815f-0534f997ea41 +# ╠═25916008-4a44-4140-9558-6f055d01b101 +# ╠═3ee91f4c-83a5-4e33-948f-88b36cf1828c +# ╠═49f045dd-b764-4e37-abc4-271a87db4677 +# ╠═2c301c7f-bc79-4141-8ae1-3022bfc54c66 +# ╠═b083214a-645f-499f-8aaa-e5511cbe64d8 +# ╠═a11d219e-f6d2-4137-b0b7-be2ffa589893 +# ╠═5fc8ecab-6a7c-467f-a67d-d4ba7cd0b2e2 +# ╠═641e4caa-3d4a-46b1-af8f-a6ef8d246723 +# ╠═b6db3c81-e960-4a29-b955-43721b3af769 +# ╠═3df00e4a-89f1-4ef3-a101-e02113fa471d +# ╠═6f31315a-f545-4989-91c9-b58db0e4ba18 +# ╠═8e4c91eb-ed58-4d1f-9a43-48d2c262e4ea +# ╠═661541a4-84f4-4f74-bbe7-3aab62fb52ce +# ╠═3b90e7f5-daef-4181-81b7-db3d0a9a922b +# ╠═09ca6d0b-fbe0-4ff4-a900-2dd2179ab6bd +# ╠═8397198e-3b49-49db-af87-a932d4967a63 +# ╠═3ba2a0a2-6535-464e-ae83-cddd68bff413 +# ╠═0f232dcc-faf8-46d6-87bb-4fe318964968 +# ╠═57cfdcf2-99c0-4db2-a740-ed74994116c9 diff --git a/examples/old/gaussian_beam_SLM_propagation.plutostate b/examples/old/gaussian_beam_SLM_propagation.plutostate new file mode 100644 index 0000000..9a9ffad Binary files /dev/null and b/examples/old/gaussian_beam_SLM_propagation.plutostate differ diff --git a/examples/old/propagation_over_distance.html b/examples/old/propagation_over_distance.html new file mode 100644 index 0000000..5e625a4 --- /dev/null +++ b/examples/old/propagation_over_distance.html @@ -0,0 +1,16 @@ + + + + + + +
\ No newline at end of file diff --git a/examples/old/propagation_over_distance.jl b/examples/old/propagation_over_distance.jl new file mode 100644 index 0000000..5124bd8 --- /dev/null +++ b/examples/old/propagation_over_distance.jl @@ -0,0 +1,196 @@ +### A Pluto.jl notebook ### +# v0.19.30 + +using Markdown +using InteractiveUtils + +# This Pluto notebook uses @bind for interactivity. When running this notebook outside of Pluto, the following 'mock version' of @bind gives bound variables a default value (instead of an error). +macro bind(def, element) + quote + local iv = try Base.loaded_modules[Base.PkgId(Base.UUID("6e696c72-6542-2067-7265-42206c756150"), "AbstractPlutoDingetjes")].Bonds.initial_value catch; b -> missing; end + local el = $(esc(element)) + global $(esc(def)) = Core.applicable(Base.get, el) ? Base.get(el) : iv(el) + el + end +end + +# ╔═╡ 6a9f04fe-ca20-11ed-1ed6-9967577ec81c +begin + using Pkg + Pkg.activate(".") + Pkg.instantiate() + using Revise +end + +# ╔═╡ cba50443-4d13-4c05-b413-6a6f3d3ff1c2 +using WaveOpticsPropagation, ImageShow, FFTW, CUDA, FourierTools, NDTools, Plots, Colors, PlutoUI + +# ╔═╡ f07896d0-808a-470d-8dd8-ed5770a333db +using LinearAlgebra, IndexFunArrays + +# ╔═╡ 5a6fa509-7535-4188-9cb8-2ab0097d5bb0 +using TestImages + +# ╔═╡ ec354723-61c0-4966-a4f7-60efbc53b917 +Plots.plotly() + +# ╔═╡ 4d1cb1e3-327c-4fbc-a501-18efc2c86096 +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 + toc(x) = Array(x) + toc(x::Array) = x + + ImageShow.simshow(x::CuArray; kwargs...) = simshow(toc(x); kwargs...) +end + +# ╔═╡ c3ba5b6b-33d2-4276-ab27-5009167d4d61 +N = 256 + +# ╔═╡ 7ef68ddb-474d-49ea-9510-91c0ab4d073d +field_img = togoc(ComplexF32.((rr2(Float32, (N, N), offset=(120, 200)) .< 10^2) .+ (rr2(Float32, (N, N), offset=(100, 100)) .< 40^2) .+ (rr2(Float32, (N, N), offset=(120, 170)) .< 3^2) )); + +# ╔═╡ e5fcbd81-c4f5-4bbc-8f22-dd248be8ec01 +L = 2f-3 + +# ╔═╡ 92d96248-06e5-446b-b856-163cbe8b4f45 +L / N + +# ╔═╡ da4d0740-46c4-460a-83ef-d78db575a3a4 +λ = 405f-9 / 1.5 + +# ╔═╡ 182fd5d8-857c-44d2-9451-8b6e8a9114e9 +zs = togoc(fftpos(16e-3, 256, CenterFT)); + +# ╔═╡ 38041742-1d46-4fa3-97b4-e5614f39ac33 +CUDA.@allowscalar zs[129] + +# ╔═╡ 851accf9-ea26-41f7-a2a5-4f0b9e53ad9c +@mytime field_img_p, tt = angular_spectrum(togoc(field_img), zs, λ, L); + +# ╔═╡ 62acd577-3402-4aee-9c62-4cc424b61758 +simshow(field_img) + +# ╔═╡ ddbe7f52-57dc-42d5-9edc-6377eb0bf31d +md"z=$(@bind iz PlutoUI.Slider(1:size(field_img_p, 3), show_value=true))" + +# ╔═╡ f9837f4a-e134-4d92-8984-51836bb89f74 +simshow((toc(abs2.(selectdim(field_img_p, 3, iz)))), γ=0.5) + +# ╔═╡ 4218847d-9a0c-4e4e-a8b6-6c76e0c3564e +begin + plot(fftpos(L, N, CenterFT), abs2.(toc(selectdim(field_img_p, 3, iz)))[120, :], title="z=$(1000 * toc(zs)[iz])mm", ticks=:native, marker=(:circle,1)) + plot!(fftpos(L, N, CenterFT), abs2.(toc(field_img)[120, :]), marker=(:circle,1)) +end + +# ╔═╡ 5fb47911-28b4-456c-b647-eed4af09f5e3 +sum(abs2.(field_img_p[:, :, 2])) + +# ╔═╡ a7325c95-d653-4e7e-ad9c-736a50ba9ea8 +sum(abs2.(field_img)) + +# ╔═╡ 03002431-9a33-4a6b-bb1a-d30f47c3bb45 + + +# ╔═╡ 5bb71ea0-cf7d-4be9-8df3-4d8bdab830d3 +simshow(togoc(selectdim(rs_kernel, 3, iz)), γ=0.1) + +# ╔═╡ 9cd1b7fd-d262-4ccc-8a2c-b82f0368ce35 +extrema(abs2.(rs_kernel)) + +# ╔═╡ 1b1144fe-b732-40bc-8382-727381f69b85 +size(field_img) + +# ╔═╡ 32ab279d-8c11-4753-9028-cb1983f1ed99 +function real_space_kernel(N, z, λ, L) + L = 2 * L + y = togoc(fftpos(L, N, CenterFT)) + x = y' + + r = sqrt.(x.^2 .+ y.^2 .+ z .^2 .+ 1f-7) + k = π / λ * 2 + kernel = z ./ r ./ 2 ./ π .* (1 ./ r .- 1im * k) .* exp.(1im .* k .* r) ./ r + return kernel +end + +# ╔═╡ dac63a43-73c2-44ab-bcaf-e3802b7e9714 +typeof(field_img) + +# ╔═╡ 5c7993ab-53bd-4991-900d-c8c64ebb76cb +field = CUDA.zeros(ComplexF32, 128, 128, 128); nothing + +# ╔═╡ 6ea1e637-112b-49f6-8526-b2a2ae20ef0c +sizeof(field) / 2^30 + +# ╔═╡ 31e0393f-9b1f-4dab-9a30-55b97d81dbac + + +# ╔═╡ 751c276f-6916-4db1-9a2f-5ec7379e45e3 +p = plan_fft!(field); nothing + +# ╔═╡ e61443e0-e3c6-4eff-98f7-e94d55f9f8fa +begin + GC.gc() + @mytime p * field; ; nothing +end + +# ╔═╡ 5ec4dbe5-9d88-45d4-9a47-3f2e888bb4c4 +begin + y = togoc(range(-10f0, 10f0, 100)) + x = togoc(range(-10f0, 10f0, 100)') + + z = togoc(reshape((range(0f0, 10f0, 100)), 1, 1, :)) +end + +# ╔═╡ 663ceff3-15be-4a61-8495-263af975de3c +@mytime exp.(1im .* z .* (x.^2 .+ y.^2)); + +# ╔═╡ d04ce01c-557f-441e-9968-c47df096349d +@mytime exp.(1im .* (x.^2 .+ y.^2)) .^ z; + +# ╔═╡ 81295271-bac4-44d7-8990-0741f0b8bb6e +NDTools.expand_dims + +# ╔═╡ 7feaec85-9604-4d3b-8f1d-c1da3df20a75 +@code_warntype WaveOpticsPropagation._prepare_angular_spectrum(field[:, :, 1], 1f0, 1f0, 1f0) + +# ╔═╡ Cell order: +# ╠═6a9f04fe-ca20-11ed-1ed6-9967577ec81c +# ╠═cba50443-4d13-4c05-b413-6a6f3d3ff1c2 +# ╠═f07896d0-808a-470d-8dd8-ed5770a333db +# ╠═ec354723-61c0-4966-a4f7-60efbc53b917 +# ╠═5a6fa509-7535-4188-9cb8-2ab0097d5bb0 +# ╠═4d1cb1e3-327c-4fbc-a501-18efc2c86096 +# ╠═c3ba5b6b-33d2-4276-ab27-5009167d4d61 +# ╠═7ef68ddb-474d-49ea-9510-91c0ab4d073d +# ╠═e5fcbd81-c4f5-4bbc-8f22-dd248be8ec01 +# ╠═92d96248-06e5-446b-b856-163cbe8b4f45 +# ╠═da4d0740-46c4-460a-83ef-d78db575a3a4 +# ╠═38041742-1d46-4fa3-97b4-e5614f39ac33 +# ╠═182fd5d8-857c-44d2-9451-8b6e8a9114e9 +# ╠═851accf9-ea26-41f7-a2a5-4f0b9e53ad9c +# ╠═62acd577-3402-4aee-9c62-4cc424b61758 +# ╠═f9837f4a-e134-4d92-8984-51836bb89f74 +# ╠═ddbe7f52-57dc-42d5-9edc-6377eb0bf31d +# ╠═4218847d-9a0c-4e4e-a8b6-6c76e0c3564e +# ╠═5fb47911-28b4-456c-b647-eed4af09f5e3 +# ╠═a7325c95-d653-4e7e-ad9c-736a50ba9ea8 +# ╠═03002431-9a33-4a6b-bb1a-d30f47c3bb45 +# ╠═5bb71ea0-cf7d-4be9-8df3-4d8bdab830d3 +# ╠═9cd1b7fd-d262-4ccc-8a2c-b82f0368ce35 +# ╠═1b1144fe-b732-40bc-8382-727381f69b85 +# ╠═32ab279d-8c11-4753-9028-cb1983f1ed99 +# ╠═dac63a43-73c2-44ab-bcaf-e3802b7e9714 +# ╠═6ea1e637-112b-49f6-8526-b2a2ae20ef0c +# ╠═5c7993ab-53bd-4991-900d-c8c64ebb76cb +# ╠═31e0393f-9b1f-4dab-9a30-55b97d81dbac +# ╠═751c276f-6916-4db1-9a2f-5ec7379e45e3 +# ╠═e61443e0-e3c6-4eff-98f7-e94d55f9f8fa +# ╠═5ec4dbe5-9d88-45d4-9a47-3f2e888bb4c4 +# ╠═663ceff3-15be-4a61-8495-263af975de3c +# ╠═d04ce01c-557f-441e-9968-c47df096349d +# ╠═81295271-bac4-44d7-8990-0741f0b8bb6e +# ╠═7feaec85-9604-4d3b-8f1d-c1da3df20a75 diff --git a/examples/old/propagation_over_distance.plutostate b/examples/old/propagation_over_distance.plutostate new file mode 100644 index 0000000..8924017 Binary files /dev/null and b/examples/old/propagation_over_distance.plutostate differ diff --git a/examples/optimize_phase.html b/examples/optimize_phase.html new file mode 100644 index 0000000..d6a20d8 --- /dev/null +++ b/examples/optimize_phase.html @@ -0,0 +1,16 @@ + + + + + + +
\ No newline at end of file diff --git a/examples/optimize_phase.jl b/examples/optimize_phase.jl new file mode 100644 index 0000000..a9cd67f --- /dev/null +++ b/examples/optimize_phase.jl @@ -0,0 +1,317 @@ +### A Pluto.jl notebook ### +# v0.19.30 + +using Markdown +using InteractiveUtils + +# This Pluto notebook uses @bind for interactivity. When running this notebook outside of Pluto, the following 'mock version' of @bind gives bound variables a default value (instead of an error). +macro bind(def, element) + quote + local iv = try Base.loaded_modules[Base.PkgId(Base.UUID("6e696c72-6542-2067-7265-42206c756150"), "AbstractPlutoDingetjes")].Bonds.initial_value catch; b -> missing; end + local el = $(esc(element)) + global $(esc(def)) = Core.applicable(Base.get, el) ? Base.get(el) : iv(el) + el + end +end + +# ╔═╡ 41c6fe0a-854c-11ee-0c0d-49cafb7df0d1 +begin + using Pkg + Pkg.activate(".") + Pkg.instantiate() + using Revise +end + +# ╔═╡ 22b6b431-8bc1-4630-9a5e-998874831b84 +using Zygote, WaveOpticsPropagation, FFTW, ImageShow, FiniteDifferences, CUDA, Optim, IndexFunArrays, FourierTools, TestImages + +# ╔═╡ a715d3ef-b3ef-496f-b50f-7ef688df5efc +using PlutoUI + +# ╔═╡ dfacac83-565f-431b-a42c-1072c419a4be +using Plots + +# ╔═╡ d18059d4-9ff2-4afa-b40f-5c37db59b57f +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 + +# ╔═╡ a1c76c16-1eff-4cde-adaa-b544f6be45a7 +sz = (512, 512) + +# ╔═╡ 74c17e61-15d7-48cd-9fce-689d1c5a3148 +#field = zeros(ComplexF32, sz); + +# ╔═╡ fdd08423-020d-4dbd-9198-b7b03264a169 +#gield .= cis.(rr2(sz, scale=0.05)) .* gaussian(sz, sigma=50); + +field = (rr2(sz) .< 150 .^2)# .* conv((normal(Float32, (sz), sigma=2)), rand(Float32, sz) .* cispi.(2 .* ( 2 .* rand(Float32, (sz...))))) + +# ╔═╡ b2f8e5bf-fc0e-4dbc-9816-3cb0072132c2 +begin + field2 = box(Float32, sz, (100, 100)) + 1im .* box(Float32, sz, (100, 100), offset=(200, 300)) + field2 .= (0.1f0 .+ Float32.(Gray.(testimage("cameraman")))) .* cispi.(Float32.(Gray.(testimage("barbara_gray_512")))) +end + +# ╔═╡ c34da860-e131-483f-a2d7-ce2dcb1ba1e6 +simshow(field2) + +# ╔═╡ e661ac7e-63f9-40d7-ba9f-1784bce61e16 +field_c = togoc(cat(field2, field, dims=3)); + +# ╔═╡ a1ba22f5-8015-4e19-ab9c-bc82a4d38e44 +sum(abs2, field_c) + +# ╔═╡ 680c0d66-97af-4d21-b90c-97cb6d8dfbd9 +simshow(field) + +# ╔═╡ e67607af-d100-42f9-a830-4c5889584b3d +#sum(abs2, AS_c2(field_c)[1][:, :, 1]) + +# ╔═╡ 82a210f8-910b-42b8-bdd2-34b11011d34a +field_c + +# ╔═╡ 1beb527b-19f0-4352-839a-9ac7e91a3f15 +λ = 633f-9 + +# ╔═╡ c11c6881-3d50-49f6-8799-8e16f172ac9a +L = 5f-3 + +# ╔═╡ 3a6eb1f4-0763-4aa6-92c1-0abe491dd558 +z = 30f-3 + +# ╔═╡ 7be3d272-cf76-4950-b80b-f743ea200d4a +AS_c, _ = AngularSpectrum(field_c[:, :, 2], z, λ, L) + +# ╔═╡ f796abe0-00cc-4178-a945-96ae5a245227 +simshow(Array(AS_c(field_c[:, :, 2])[1]), γ=0.3) + +# ╔═╡ 9b03154d-bf00-4e95-a5a1-5ae5e7cb48a4 +simshow(Array(abs2.(AS_c(field_c)[1]))) + +# ╔═╡ 0180ea0b-87a8-4233-86fe-09308a70738d +N = 10 + +# ╔═╡ 1418e664-6207-4df4-aa79-f18633b63b2e +begin + diffuser = conv((normal(Float32, (sz), sigma=1)), cispi.(2 .* rand(Float32, (sz..., N)))) + diffuser_c = togoc(diffuser)# .* box(sz, (120, 150)) + + diffuser_c[:, :, 1] .= cis.(rr2(sz, scale=0.05)) .* gaussian(sz, sigma=100); + diffuser_c[:, :, 2] .= cis.(.-1 .* rr2(sz, scale=0.05)) .* gaussian(sz, sigma=100); + diffuser_c[:, :, 3] .= 1 + + diffuser_c[:, :, 4] .*= cis.(0.3 .* xx((sz))) + diffuser_c[:, :, 5] .*= cis.(.- 0.3 .* xx((sz))) + diffuser_c[:, :, 6] .*= cis.(0.3 .* yy((sz))) + diffuser_c[:, :, 7] .*= cis.(.- 0.3 .* yy((sz))) + +end + +# ╔═╡ 52f7af37-1a44-4a7c-8a9e-31df9bbcda9a +simshow(Array(diffuser_c[:, :, 5])) + +# ╔═╡ b436fdc0-ae8d-4dbe-94e6-8b7c26339a70 +#AS_c, _ = AngularSpectrum(field_c[:, :, 1], z, λ, L) + +# ╔═╡ c97f6332-82d9-4441-80ad-8a49c7739119 +AS_c2, _ = AngularSpectrum(togoc(zeros(ComplexF32, (sz..., N))), togoc(repeat([z], N)), λ, L) + +# ╔═╡ 0bd359cc-0bd4-4748-8905-361969811248 +p = plan_fft(diffuser_c, (1,2)) + +# ╔═╡ 23e64d50-8a6c-418f-86eb-1b4f9917fb32 +fwd(x) = begin + #first = AS_c(x)[1] + first = x[:, :, 2] + second = first .* diffuser_c + third = x[:, :, 1] .* AS_c2(second)[1] + #fourth = abs2.(fftshift(p * ifftshift(third, (1,2)), (1,2))) ./ size(x, 1).^2 + fourth = abs2.(AS_c2(third)[1]) + return fourth#, second, third +end + +# ╔═╡ a5593e55-5b77-4c99-b9a9-5e060748d6af +simshow(Array(fwd(field_c)[2][:, :, 8])) + +# ╔═╡ d00e3c52-a25d-4f13-9f76-f9fb0e255737 +simshow(Array(AS_c(fwd(field_c)[1][:, :, 8])[1])) + +# ╔═╡ 86c47180-3afe-4545-912e-12911155a02a +simshow(Array(diffuser_c)[:, :, 8]) + +# ╔═╡ 2c6ec01e-df45-4482-9d52-ce04ed343c5d +simshow(Array(AS_c2(diffuser_c)[1][:, :, 8])) + +# ╔═╡ df218193-d0f0-4c85-9439-9b70e2e55cf8 +simshow(Array(diffuser_c)[:, :, 9]) + +# ╔═╡ 2531c673-ab41-4374-85fe-e4e256f05c98 +field_c |> size + +# ╔═╡ 8516d9be-440c-45c5-8f20-cce74a9783e9 +function make_fg!(fwd, measured, loss=:L2) + L2 = let measured=measured, fwd=fwd + function L2(x) + return sum(abs2, fwd(x) .- measured) + end + end + + + f = let + if loss == :L2 + L2 + end + end + + g! = let f=f + function g!(G, rec) + if !isnothing(G) + return G .= Zygote.gradient(f, rec)[1] + end + end + end + return f, g! +end + +# ╔═╡ bdeb034b-8c34-4134-943c-eed63314cf20 +measurement_c = fwd(field_c); + +# ╔═╡ 525a1be2-286b-4626-b122-377084c6d456 +simshow(Array(measurement_c)[:, :, 1], γ=1) + +# ╔═╡ 18eb702e-ecad-4648-9e3a-b85cef60deb4 +begin + rec0 = togoc(ones(ComplexF32, (sz..., 2))) .* (rr2(sz) .< 150 .^2) + #rec0[:, :, 1] .= field_c[:, :, 1] + #rec0[:, :, 2] .= conv(CuArray(normal(sz, sigma=20)), field_c[:, :, 2]) + #rec0 .= cis.(rr2(sz, scale=0.04, offset=(200, 200))) .* box(sz, (80, 120)); +end + +# ╔═╡ c9776bdb-04ab-44e7-9c8e-9a7b8eb080f0 +simshow(Array(rec0)) + +# ╔═╡ 504074c0-9ff8-4b5e-8c6a-d196c8111a81 +f, g! = make_fg!(fwd, measurement_c) + +# ╔═╡ ad585dda-69c5-4bb5-a1bf-1e6001447669 +sum(abs2, measurement_c .* diffuser_c) + +# ╔═╡ 2d6dadd6-57a8-4c98-bf12-56276cb4cffe +CUDA.@time res = Optim.optimize(f, g!, rec0, ConjugateGradient(), + Optim.Options(iterations = 500, + store_trace=true, + f_abstol=1e-10, + g_abstol=1e-10, )) + +# ╔═╡ 870e717a-577f-4097-a650-422e99fde221 +@time AS_c2(diffuser_c) + +# ╔═╡ a1bb8bb6-7b00-4478-84f3-1e350d530dbe +plot([t.iteration for t in res.trace], [t.value for t in res.trace], label="LBFGS", yaxis=:log) + +# ╔═╡ f252bf81-98d7-4383-9c29-fe4c2a034dc6 +simshow([Array(field_c[:, :, 1]) cispi(0.42) .* Array(res.minimizer)[:, :, 1]], γ=0.4) + +# ╔═╡ df93cc2e-4031-4011-a1d6-d76c62fd3000 +simshow(Array(field_c[:, :, 2])) + +# ╔═╡ 96c3bbfd-1b94-465e-b852-fadf7649ba57 +simshow(abs.([Array(field2) Array(res.minimizer)[:, :, 1]]), γ=1) + +# ╔═╡ 633dafe5-0f74-49a5-93bc-aaeec49d0254 +simshow(angle.([Array(field2) Array(res.minimizer)[:, :, 1]])) + +# ╔═╡ 13d98c90-296f-4fd4-bfc4-93ee81ae3484 +simshow(([Array(field2) Array(res.minimizer)[:, :, 2]])) + +# ╔═╡ afc890a3-25ae-4c4d-b886-a5433f317e73 +simshow(Array(res.minimizer)[:, :, 2]) + +# ╔═╡ 956c58a2-b182-48a9-b694-b8f6f9ba8077 +simshow(Array(field)) + +# ╔═╡ a9860e96-7888-4e68-9b9b-9adc0e260eef + + +# ╔═╡ 40c7f76c-2f40-403e-b160-d98ee9c54071 +@bind z_i Slider(1:size(measurement_c, 3)) + +# ╔═╡ 4823c477-38eb-47f0-89fa-6f02a0940d35 +#simshow(Array(fwd(res.minimizer)[:, :, z_i]) .- Array(measurement_c[:, :, z_i])) + +# ╔═╡ 00aaa61e-48a1-45dd-bdfe-5bd4879cd2d3 +simshow(Array(fwd(res.minimizer)[:, :, z_i]), γ=0.8) + +# ╔═╡ 6f57680b-7f7a-4d69-a62c-ccfce645151c +simshow(Array(measurement_c[:, :, z_i]), γ=0.8) + +# ╔═╡ 88239ba1-5528-46c3-9914-3ea62dfc0cfb + + +# ╔═╡ a3987c52-33d8-48f4-b373-4b6df1d6c17e +fwd(res.minimizer) + +# ╔═╡ Cell order: +# ╠═41c6fe0a-854c-11ee-0c0d-49cafb7df0d1 +# ╠═d18059d4-9ff2-4afa-b40f-5c37db59b57f +# ╠═22b6b431-8bc1-4630-9a5e-998874831b84 +# ╠═a715d3ef-b3ef-496f-b50f-7ef688df5efc +# ╠═dfacac83-565f-431b-a42c-1072c419a4be +# ╠═a1c76c16-1eff-4cde-adaa-b544f6be45a7 +# ╠═74c17e61-15d7-48cd-9fce-689d1c5a3148 +# ╠═fdd08423-020d-4dbd-9198-b7b03264a169 +# ╠═b2f8e5bf-fc0e-4dbc-9816-3cb0072132c2 +# ╠═c34da860-e131-483f-a2d7-ce2dcb1ba1e6 +# ╠═e661ac7e-63f9-40d7-ba9f-1784bce61e16 +# ╠═a1ba22f5-8015-4e19-ab9c-bc82a4d38e44 +# ╠═680c0d66-97af-4d21-b90c-97cb6d8dfbd9 +# ╠═7be3d272-cf76-4950-b80b-f743ea200d4a +# ╠═f796abe0-00cc-4178-a945-96ae5a245227 +# ╠═e67607af-d100-42f9-a830-4c5889584b3d +# ╠═82a210f8-910b-42b8-bdd2-34b11011d34a +# ╟─9b03154d-bf00-4e95-a5a1-5ae5e7cb48a4 +# ╠═1beb527b-19f0-4352-839a-9ac7e91a3f15 +# ╠═c11c6881-3d50-49f6-8799-8e16f172ac9a +# ╠═3a6eb1f4-0763-4aa6-92c1-0abe491dd558 +# ╠═0180ea0b-87a8-4233-86fe-09308a70738d +# ╠═1418e664-6207-4df4-aa79-f18633b63b2e +# ╠═52f7af37-1a44-4a7c-8a9e-31df9bbcda9a +# ╠═b436fdc0-ae8d-4dbe-94e6-8b7c26339a70 +# ╠═c97f6332-82d9-4441-80ad-8a49c7739119 +# ╠═0bd359cc-0bd4-4748-8905-361969811248 +# ╠═23e64d50-8a6c-418f-86eb-1b4f9917fb32 +# ╠═a5593e55-5b77-4c99-b9a9-5e060748d6af +# ╠═d00e3c52-a25d-4f13-9f76-f9fb0e255737 +# ╠═86c47180-3afe-4545-912e-12911155a02a +# ╠═2c6ec01e-df45-4482-9d52-ce04ed343c5d +# ╠═df218193-d0f0-4c85-9439-9b70e2e55cf8 +# ╠═2531c673-ab41-4374-85fe-e4e256f05c98 +# ╠═8516d9be-440c-45c5-8f20-cce74a9783e9 +# ╠═bdeb034b-8c34-4134-943c-eed63314cf20 +# ╠═525a1be2-286b-4626-b122-377084c6d456 +# ╠═18eb702e-ecad-4648-9e3a-b85cef60deb4 +# ╠═c9776bdb-04ab-44e7-9c8e-9a7b8eb080f0 +# ╠═504074c0-9ff8-4b5e-8c6a-d196c8111a81 +# ╠═ad585dda-69c5-4bb5-a1bf-1e6001447669 +# ╠═2d6dadd6-57a8-4c98-bf12-56276cb4cffe +# ╠═870e717a-577f-4097-a650-422e99fde221 +# ╠═a1bb8bb6-7b00-4478-84f3-1e350d530dbe +# ╠═f252bf81-98d7-4383-9c29-fe4c2a034dc6 +# ╠═df93cc2e-4031-4011-a1d6-d76c62fd3000 +# ╠═96c3bbfd-1b94-465e-b852-fadf7649ba57 +# ╠═633dafe5-0f74-49a5-93bc-aaeec49d0254 +# ╠═13d98c90-296f-4fd4-bfc4-93ee81ae3484 +# ╠═afc890a3-25ae-4c4d-b886-a5433f317e73 +# ╠═956c58a2-b182-48a9-b694-b8f6f9ba8077 +# ╠═a9860e96-7888-4e68-9b9b-9adc0e260eef +# ╠═40c7f76c-2f40-403e-b160-d98ee9c54071 +# ╠═4823c477-38eb-47f0-89fa-6f02a0940d35 +# ╠═00aaa61e-48a1-45dd-bdfe-5bd4879cd2d3 +# ╠═6f57680b-7f7a-4d69-a62c-ccfce645151c +# ╠═88239ba1-5528-46c3-9914-3ea62dfc0cfb +# ╠═a3987c52-33d8-48f4-b373-4b6df1d6c17e diff --git a/examples/optimize_phase.plutostate b/examples/optimize_phase.plutostate new file mode 100644 index 0000000..5a68ceb Binary files /dev/null and b/examples/optimize_phase.plutostate differ diff --git a/index.html b/index.html new file mode 100644 index 0000000..f66cc15 --- /dev/null +++ b/index.html @@ -0,0 +1,23 @@ + + + + + + + + + + + + +

Notebooks

+ + + + diff --git a/pluto_export.json b/pluto_export.json new file mode 100644 index 0000000..e01a0e7 --- /dev/null +++ b/pluto_export.json @@ -0,0 +1 @@ +{"notebooks":{"examples/old/propagation_over_distance.jl":{"id":"examples/old/propagation_over_distance.jl","hash":"-ka3zcuzxh-YOaK8EUIr0SmXuIVxHctKYKenKPon4YQ","html_path":"examples/old/propagation_over_distance.html","statefile_path":"examples/old/propagation_over_distance.plutostate","notebookfile_path":"examples/old/propagation_over_distance.jl","current_hash":"-ka3zcuzxh-YOaK8EUIr0SmXuIVxHctKYKenKPon4YQ","desired_hash":"-ka3zcuzxh-YOaK8EUIr0SmXuIVxHctKYKenKPon4YQ","frontmatter":{"title":"propagation_over_distance"}},"examples/double_slit.jl":{"id":"examples/double_slit.jl","hash":"65MVOYMWkMajfnqj0P95FhfpmNiBJ_F55yUptntUmJ0","html_path":"examples/double_slit.html","statefile_path":"examples/double_slit.plutostate","notebookfile_path":"examples/double_slit.jl","current_hash":"65MVOYMWkMajfnqj0P95FhfpmNiBJ_F55yUptntUmJ0","desired_hash":"65MVOYMWkMajfnqj0P95FhfpmNiBJ_F55yUptntUmJ0","frontmatter":{"title":"double_slit"}},"examples/Scalable_Angular_Spectrum.jl":{"id":"examples/Scalable_Angular_Spectrum.jl","hash":"yPENux3rEWBsNTitRHZeiNJwjnMIVeKKc2NlM4wFMjQ","html_path":"examples/Scalable_Angular_Spectrum.html","statefile_path":"examples/Scalable_Angular_Spectrum.plutostate","notebookfile_path":"examples/Scalable_Angular_Spectrum.jl","current_hash":"yPENux3rEWBsNTitRHZeiNJwjnMIVeKKc2NlM4wFMjQ","desired_hash":"yPENux3rEWBsNTitRHZeiNJwjnMIVeKKc2NlM4wFMjQ","frontmatter":{"title":"Scalable_Angular_Spectrum"}},"examples/optimize_phase.jl":{"id":"examples/optimize_phase.jl","hash":"JOFkhlmgMJi7yRRNupWarIM0XStLM4idJeb_rOIT6gc","html_path":"examples/optimize_phase.html","statefile_path":"examples/optimize_phase.plutostate","notebookfile_path":"examples/optimize_phase.jl","current_hash":"JOFkhlmgMJi7yRRNupWarIM0XStLM4idJeb_rOIT6gc","desired_hash":"JOFkhlmgMJi7yRRNupWarIM0XStLM4idJeb_rOIT6gc","frontmatter":{"title":"optimize_phase"}},"test/gaussian_beam.jl":{"id":"test/gaussian_beam.jl","hash":"XxVlr_jDDoP_tENKAHP_AUBLZUnYnWGZH1FEzRhoMiw","html_path":"test/gaussian_beam.html","statefile_path":"test/gaussian_beam.plutostate","notebookfile_path":"test/gaussian_beam.jl","current_hash":"XxVlr_jDDoP_tENKAHP_AUBLZUnYnWGZH1FEzRhoMiw","desired_hash":"XxVlr_jDDoP_tENKAHP_AUBLZUnYnWGZH1FEzRhoMiw","frontmatter":{"title":"gaussian_beam"}},"examples/benchmark_CPU_CUDA.jl":{"id":"examples/benchmark_CPU_CUDA.jl","hash":"IjCQ4s_kPUCOalYeWByrVXfh4sb8Gl-DKgyB8xK-jB4","html_path":"examples/benchmark_CPU_CUDA.html","statefile_path":"examples/benchmark_CPU_CUDA.plutostate","notebookfile_path":"examples/benchmark_CPU_CUDA.jl","current_hash":"IjCQ4s_kPUCOalYeWByrVXfh4sb8Gl-DKgyB8xK-jB4","desired_hash":"IjCQ4s_kPUCOalYeWByrVXfh4sb8Gl-DKgyB8xK-jB4","frontmatter":{"title":"benchmark_CPU_CUDA"}},"examples/old/gaussian_beam_SLM_propagation.jl":{"id":"examples/old/gaussian_beam_SLM_propagation.jl","hash":"Fd2W88vhOghN87HsQH7OvTQn7KbdzTwkiyeB26S2R18","html_path":"examples/old/gaussian_beam_SLM_propagation.html","statefile_path":"examples/old/gaussian_beam_SLM_propagation.plutostate","notebookfile_path":"examples/old/gaussian_beam_SLM_propagation.jl","current_hash":"Fd2W88vhOghN87HsQH7OvTQn7KbdzTwkiyeB26S2R18","desired_hash":"Fd2W88vhOghN87HsQH7OvTQn7KbdzTwkiyeB26S2R18","frontmatter":{"title":"gaussian_beam_SLM_propagation"}},"test/double_slit.jl":{"id":"test/double_slit.jl","hash":"pXsuKzsLET9WEGPWCRlcQP6JiJRkdndh8u925JYP4vc","html_path":"test/double_slit.html","statefile_path":"test/double_slit.plutostate","notebookfile_path":"test/double_slit.jl","current_hash":"pXsuKzsLET9WEGPWCRlcQP6JiJRkdndh8u925JYP4vc","desired_hash":"pXsuKzsLET9WEGPWCRlcQP6JiJRkdndh8u925JYP4vc","frontmatter":{"title":"double_slit"}}},"pluto_version":"0.19.36","julia_version":"1.10.0","format_version":"1","title":null,"description":null,"collections":null,"binder_url":"https://mybinder.org/v2/gh/fonsp/pluto-on-binder/v0.19.36","slider_server_url":null} \ No newline at end of file diff --git a/pluto_state_cache/0_19_36-ka3zcuzxh-YOaK8EUIr0SmXuIVxHctKYKenKPon4YQ.plutostate b/pluto_state_cache/0_19_36-ka3zcuzxh-YOaK8EUIr0SmXuIVxHctKYKenKPon4YQ.plutostate new file mode 100644 index 0000000..8924017 Binary files /dev/null and b/pluto_state_cache/0_19_36-ka3zcuzxh-YOaK8EUIr0SmXuIVxHctKYKenKPon4YQ.plutostate differ diff --git a/pluto_state_cache/0_19_3665MVOYMWkMajfnqj0P95FhfpmNiBJ_F55yUptntUmJ0.plutostate b/pluto_state_cache/0_19_3665MVOYMWkMajfnqj0P95FhfpmNiBJ_F55yUptntUmJ0.plutostate new file mode 100644 index 0000000..53efe3e Binary files /dev/null and b/pluto_state_cache/0_19_3665MVOYMWkMajfnqj0P95FhfpmNiBJ_F55yUptntUmJ0.plutostate differ diff --git a/pluto_state_cache/0_19_36Fd2W88vhOghN87HsQH7OvTQn7KbdzTwkiyeB26S2R18.plutostate b/pluto_state_cache/0_19_36Fd2W88vhOghN87HsQH7OvTQn7KbdzTwkiyeB26S2R18.plutostate new file mode 100644 index 0000000..9a9ffad Binary files /dev/null and b/pluto_state_cache/0_19_36Fd2W88vhOghN87HsQH7OvTQn7KbdzTwkiyeB26S2R18.plutostate differ diff --git a/pluto_state_cache/0_19_36IjCQ4s_kPUCOalYeWByrVXfh4sb8Gl-DKgyB8xK-jB4.plutostate b/pluto_state_cache/0_19_36IjCQ4s_kPUCOalYeWByrVXfh4sb8Gl-DKgyB8xK-jB4.plutostate new file mode 100644 index 0000000..d6fe287 Binary files /dev/null and b/pluto_state_cache/0_19_36IjCQ4s_kPUCOalYeWByrVXfh4sb8Gl-DKgyB8xK-jB4.plutostate differ diff --git a/pluto_state_cache/0_19_36JOFkhlmgMJi7yRRNupWarIM0XStLM4idJeb_rOIT6gc.plutostate b/pluto_state_cache/0_19_36JOFkhlmgMJi7yRRNupWarIM0XStLM4idJeb_rOIT6gc.plutostate new file mode 100644 index 0000000..5a68ceb Binary files /dev/null and b/pluto_state_cache/0_19_36JOFkhlmgMJi7yRRNupWarIM0XStLM4idJeb_rOIT6gc.plutostate differ diff --git a/pluto_state_cache/0_19_36XxVlr_jDDoP_tENKAHP_AUBLZUnYnWGZH1FEzRhoMiw.plutostate b/pluto_state_cache/0_19_36XxVlr_jDDoP_tENKAHP_AUBLZUnYnWGZH1FEzRhoMiw.plutostate new file mode 100644 index 0000000..19d8f20 Binary files /dev/null and b/pluto_state_cache/0_19_36XxVlr_jDDoP_tENKAHP_AUBLZUnYnWGZH1FEzRhoMiw.plutostate differ diff --git a/pluto_state_cache/0_19_36pXsuKzsLET9WEGPWCRlcQP6JiJRkdndh8u925JYP4vc.plutostate b/pluto_state_cache/0_19_36pXsuKzsLET9WEGPWCRlcQP6JiJRkdndh8u925JYP4vc.plutostate new file mode 100644 index 0000000..ba7d4c6 Binary files /dev/null and b/pluto_state_cache/0_19_36pXsuKzsLET9WEGPWCRlcQP6JiJRkdndh8u925JYP4vc.plutostate differ diff --git a/pluto_state_cache/0_19_36yPENux3rEWBsNTitRHZeiNJwjnMIVeKKc2NlM4wFMjQ.plutostate b/pluto_state_cache/0_19_36yPENux3rEWBsNTitRHZeiNJwjnMIVeKKc2NlM4wFMjQ.plutostate new file mode 100644 index 0000000..5dabfb0 Binary files /dev/null and b/pluto_state_cache/0_19_36yPENux3rEWBsNTitRHZeiNJwjnMIVeKKc2NlM4wFMjQ.plutostate differ diff --git a/src/WaveOpticsPropagation.jl b/src/WaveOpticsPropagation.jl new file mode 100644 index 0000000..8539f1f --- /dev/null +++ b/src/WaveOpticsPropagation.jl @@ -0,0 +1,20 @@ +module WaveOpticsPropagation + +using EllipsisNotation +using FFTW +using ChainRulesCore +using Zygote +using NDTools +using FourierTools +using IndexFunArrays +using CUDA + +include("utils.jl") +include("propagation.jl") +include("angular_spectrum.jl") +include("scalable_angular_spectrum.jl") +include("fraunhofer.jl") +include("beams.jl") +include("conv.jl") + +end diff --git a/src/angular_spectrum.jl b/src/angular_spectrum.jl new file mode 100644 index 0000000..18b5979 --- /dev/null +++ b/src/angular_spectrum.jl @@ -0,0 +1,207 @@ +export angular_spectrum +export AngularSpectrum + + +_transform_z(::Type{T}, z::Number) where {T<:Number} = T(z) +_transform_z(::Type{T}, z::AbstractArray{T}) where T = reshape(z, 1, 1, :) +_transform_z(::Type{T}, z::AbstractArray{T2}) where {T, T2} = reshape(T.(z), 1, 1, :) + +function _prepare_angular_spectrum(field::AbstractArray{CT}, z, λ, L; + padding=true, pad_factor=2, + bandlimit=true, + bandlimit_border=(0.8, 1.0)) where {CT<:Complex} + + fftdims = (1,2) + T = real(CT) + λ = T(λ) + z = _transform_z(T, z) + L = T.(L isa Number ? (L, L) : L) + + bandlimit_border = real(CT).(bandlimit_border) + L_new = padding ? pad_factor .* L : L + + # applies zero padding + if ndims(field) == 2 + pad_factor2 = pad_factor + elseif ndims(field) == 3 + pad_factor2 = (pad_factor * size(field, 1), pad_factor * size(field, 2), size(field, 3)) + end + field_new = padding ? pad(field, pad_factor2) : field + + # helpful propagation variables + (; k, f_x, f_y) = Zygote.@ignore _propagation_variables(field_new, λ, L_new) + + # transfer function kernel of angular spectrum + H = exp.(1im .* k .* z .* sqrt.(CT(1) .- abs2.(f_x .* λ) .- abs2.(f_y .* λ))) + + # bandlimit according to Matsushima + # as addition we introduce a smooth bandlimit with a Hann window + # and fuzzy logic + Δu = 1 ./ L_new + u_limit = 1 ./ (sqrt.((2 .* Δu .* z).^2 .+ 1) .* λ) + + W = let + if bandlimit + # bandlimit filter + u_limit1 = u_limit isa AbstractArray ? u_limit[1:1, ..] : u_limit[1] + u_limit2 = u_limit isa AbstractArray ? u_limit[2:2, ..] : u_limit[2] + + W = .*(hann.(scale.(abs2.(f_y) ./ u_limit1 .^2 .+ abs2.(f_x) * λ^2, + bandlimit_border[1], bandlimit_border[2])), + hann.(scale.(abs2.(f_x) ./ u_limit2 .^2 .+ abs2.(f_y) * λ^2, + bandlimit_border[1], bandlimit_border[2]))) + else + # use an array here too, to avoid type instabilities + W = similar(field, real(eltype(field)), size(f_y, 1), size(f_x, 1)) + W .= 1 + end + end + + return (;field_new, H, W, fftdims) +end + + +""" + angular_spectrum(field, z, λ, L; kwargs...) + +Returns the electrical field with physical length `L` and wavelength `λ` propagated with the angular spectrum +method of plane waves (AS) by the propagation distance `z`. + +This method is efficient but to avoid recalculating some arrays (such as the phase kernel), see [`AngularSpectrum`](@ref). + +# Arguments +* `field`: Input field +* `z`: propagation distance +* `λ`: wavelength of field +* `L`: field size (can be a scalar or a tuple) indicating field size + + +# Keyword Arguments +* `padding=true`: applies padding to avoid convolution wraparound +* `pad_factor=2`: padding of 2. Larger numbers are not recommended since they don't provide better results. +* `bandlimit=true`: applies the bandlimit to avoid circular wraparound due to undersampling + of the complex propagation kernel [1] +* `bandlimit_border=(0.8, 1)`: applies a smooth bandlimit cut-off instead of hard-edge. + + +# References +* Matsushima, Kyoji, and Tomoyoshi Shimobaba. "Band-limited angular spectrum method for numerical simulation of free-space propagation in far and near fields." Optics express 17.22 (2009): 19662-19673. +""" +function angular_spectrum(field::AbstractArray{CT, 2}, z, λ, L; + padding=true, pad_factor=2, + bandlimit=true, + bandlimit_border=(0.8, 1.0)) where {CT<:Complex} + + @assert size(field, 1) == size(field, 2) "input field needs to be quadradically shaped and not $(size(field, 1)), $(size(field, 2))" + + (; field_new, H, W, fftdims) = _prepare_angular_spectrum(field, z, λ, L; padding, + pad_factor, bandlimit, bandlimit_border) + + # propagate field + field_out = fftshift(ifft(fft(ifftshift(field_new, fftdims), fftdims) .* H .* W, fftdims), fftdims) + + field_out_cropped = padding ? crop_center(field_out, size(field)) : field_out + + # return final field and some other variables + return field_out_cropped, (; H, L) +end + + +angular_spectrum(field::AbstractArray, z, λ, L; kwargs...) = throw(ArgumentError("Provided field needs to have a complex elementype")) + + + # highly optimized version with pre-planning +struct AngularSpectrum3{A, T, P} + HW::A + buffer::A + buffer2::A + L::T + p::P + padding::Bool + pad_factor::Int +end + +""" + AngularSpectrum(field, z, λ, L; kwargs...) + +Returns a function for efficient reuse of pre-calculated kernels. + +See [`angular_spectrum`](@ref) for the full documentation. + + +# Example +```jldoctest +julia> field = zeros(ComplexF32, (4,4)); field[3,3] = 1 +1 + +julia> as = AngularSpectrum(field, 100e-9, 632e-9, 10e-6) +WaveOpticsPropagation.AngularSpectrum3{Matrix{ComplexF64}, Float64, FFTW.cFFTWPlan{ComplexF32, -1, true, 2, UnitRange{Int64}}}(ComplexF64[0.5451947489704718 + 0.8383094212133275im 0.5456108987195186 + 0.8380386310895693im … 0.5468597886165149 + 0.8372242062878382im 0.5456108987195186 + 0.8380386310895693im; 0.5456108987195186 + 0.8380386310895693im 0.5460271219052333 + 0.8377674988586556im … 0.5472762321570075 + 0.8369520450515843im 0.5460271219052333 + 0.8377674988586556im; … ; 0.5468597886165149 + 0.8372242062878382im 0.5472762321570075 + 0.8369520450515843im … 0.5485260036074586 + 0.8361334961394803im 0.5472762321570075 + 0.8369520450515843im; 0.5456108987195186 + 0.8380386310895693im 0.5460271219052333 + 0.8377674988586556im … 0.5472762321570075 + 0.8369520450515843im 0.5460271219052333 + 0.8377674988586556im], ComplexF64[0.0 + 0.0im 0.0 + 0.0im … 2047.5009765625 + 4.571175720473986e-41im 2047.5009765625 + 4.571175720473986e-41im; 2058.2890625 + 4.571175720473986e-41im 0.0 + 0.0im … 2056.2578125 + 4.571175720473986e-41im 2058.1640625 + 4.571175720473986e-41im; … ; 0.0 + 0.0im 0.0 + 0.0im … 2047.5009765625 + 4.571175720473986e-41im 2047.5009765625 + 4.571175720473986e-41im; 0.0 + 0.0im 0.0 + 0.0im … 2058.0234375 + 4.571175720473986e-41im 0.0 + 0.0im], ComplexF64[0.0 + 0.0im 0.0 + 0.0im … 0.0 + 0.0im 0.0 + 0.0im; 0.0 + 0.0im 0.0 + 0.0im … 0.0 + 0.0im 0.0 + 0.0im; … ; 0.0 + 0.0im 0.0 + 0.0im … 0.0 + 0.0im 0.0 + 0.0im; 0.0 + 0.0im 0.0 + 0.0im … 0.0 + 0.0im 0.0 + 0.0im], 1.0e-5, FFTW in-place forward plan for 8×8 array of ComplexF32 +(dft-rank>=2/1 + (dft-direct-8-x8 "n1fv_8_avx2_128") + (dft-direct-8-x8 "n1fv_8_avx2_128")), true, 2) + +julia> as(field) +4×4 Matrix{ComplexF64}: + 7.07805e-8-3.53903e-7im -2.54379e-7+1.20194e-6im 0.000417542-0.000277363im -2.54379e-7+1.20194e-6im + -2.52505e-7+1.20063e-6im 8.57634e-7-4.05761e-6im -0.00142377+0.000938398im 8.57634e-7-4.05761e-6im + 0.000417545-0.00027737im -0.00142378+0.000938403im 0.549778+0.835303im -0.00142378+0.000938403im + -2.52505e-7+1.20063e-6im 8.57634e-7-4.05761e-6im -0.00142377+0.000938398im 8.57634e-7-4.05761e-6im +``` +""" +function AngularSpectrum(field::AbstractArray{CT, N}, z, λ, L; + padding=true, pad_factor=2, + bandlimit=true, + bandlimit_border=(0.8, 1)) where {CT, N} + + (; field_new, H, W, fftdims) = _prepare_angular_spectrum(field, z, λ, L; padding, + pad_factor, bandlimit, bandlimit_border) + + + if z isa AbstractVector + buffer2 = similar(field, complex(eltype(field_new)), (size(field_new, 1), size(field_new, 2), length(z))) + buffer = copy(buffer2) + else + buffer2 = similar(field, complex(eltype(field_new)), (size(field_new, 1), size(field_new, 2))) + buffer = copy(buffer2) + end + + p = plan_fft!(buffer, (1, 2)) + H .= H .* W + HW = H + + return AngularSpectrum3{typeof(H), typeof(L), typeof(p)}(HW, buffer, buffer2, L, p, padding, pad_factor), L + end + +""" + (as:AngularSpectrum3)(field) + +Uses the struct to efficiently store some pre-calculated objects. +Propagate the field. +""" +function (as::AngularSpectrum3)(field) + fill!(as.buffer2, 0) + field_new = set_center!(as.buffer2, field, broadcast=true) + field_imd = as.p * ifftshift!(as.buffer, field_new, (1, 2)) + field_imd .*= as.HW + field_out = fftshift!(as.buffer2, inv(as.p) * field_imd, (1, 2)) + field_out_cropped = as.padding ? crop_center(field_out, size(field), return_view=true) : field_out + return field_out_cropped, (; as.L) +end + + +function ChainRulesCore.rrule(as::AngularSpectrum3, field) + field_and_tuple = as(field) + function as_pullback(ȳ) + f̄ = NoTangent() + y2 = ȳ.backing[1] + + fill!(as.buffer2, 0) + field_new = as.padding ? set_center!(as.buffer2, y2, broadcast=true) : y2 + field_imd = as.p * ifftshift!(as.buffer, field_new, (1, 2)) + field_imd .*= conj.(as.HW) + field_out = fftshift!(as.buffer2, inv(as.p) * field_imd, (1, 2)) + field_out_cropped = as.padding ? crop_center(field_out, size(field), return_view=true) : field_out + return f̄, field_out_cropped + end + return field_and_tuple, as_pullback +end diff --git a/src/beams.jl b/src/beams.jl new file mode 100644 index 0000000..21f2cb7 --- /dev/null +++ b/src/beams.jl @@ -0,0 +1,18 @@ +export gauss_beam + + +gauss_beam(y, x, z, λ, w_0) = throw(ArgumentError("All arguments need to have the same datatype (Float32, Float64, ...).")) + +function gauss_beam(y::T, x::T, z::T, λ::T, w_0::T) where T + k = π / λ * 2 + z_R = π * w_0^2 / λ + r² = x ^ 2 + y ^ 2 + # don't put exp(i * k * z) into the same exp, it causes some strange wraps + return w_0 / gauss_w(z, z_R, w_0) * exp(-r² / gauss_w(z, z_R, w_0)^2) * + exp.(1im * k * z) * + exp(1im * (k * r² / 2 / gauss_R(z, z_R) - gauss_ψ(z, z_R))) +end + +@inline gauss_R(z, z_R) = z * (1 + (z_R /z)^2) +@inline gauss_ψ(z, z_R) = atan(z, z_R) +@inline gauss_w(z, z_R, w_0) = w_0 * sqrt(1 + (z / z_R)^2) diff --git a/src/conv.jl b/src/conv.jl new file mode 100644 index 0000000..fdbf9b4 --- /dev/null +++ b/src/conv.jl @@ -0,0 +1,7 @@ + + +function conv(u_in, v, dims=(1,2)) + u = pad(u_in, size(v)[1:2]) + + return crop_center(ifft(fft(u, dims) .* fft(v, dims), dims), size(u_in)) +end diff --git a/src/fraunhofer.jl b/src/fraunhofer.jl new file mode 100644 index 0000000..8c739d1 --- /dev/null +++ b/src/fraunhofer.jl @@ -0,0 +1,115 @@ +export fraunhofer +export Fraunhofer + + +""" + fraunhofer(field, z, λ, L) + +Returns the electrical field with physical length `L` and wavelength `λ` +propagated with the Fraunhofer propagation (a single FFT) by the propagation distance `z`. +This is based on a far field approximation `z >> λ` + +This method is efficient but to save memory and avoiding recalculating some arrays (such as the phase kernel), see [`Fraunhofer`](@ref). + +# Arguments +* `field`: Input field +* `z`: propagation distance +* `λ`: wavelength of field +* `L`: field size indicating field size + +# Keyword Arguments +* `skip_final_phase` skip the final phase which is multiplied to the propagated field at the end + +""" +function fraunhofer(U, z, λ, L; skip_final_phase=true) + @assert size(U, 1) == size(U, 2) + L_new = λ * z / L * size(U, 1) + Ns = size(U)[1:2] + + p = Zygote.@ignore plan_fft(U, (1,2)) + + if skip_final_phase + out = fftshift(p * ifftshift(U)) ./ √(size(U, 1) * size(U, 2)) + else + k = eltype(U)(2π) / λ + # output coordinates + y = similar(U, real(eltype(U)), (Ns[1], 1)) + Zygote.@ignore y .= (fftpos(L, Ns[1], CenterFT)) + x = similar(U, real(eltype(U)), (1, Ns[2])) + Zygote.@ignore x .= (fftpos(L, Ns[2], CenterFT))' + phasefactor = (-1im) .* exp.(1im * k / (2 * z) .* (x.^2 .+ y.^2)) + out = phasefactor .* fftshift(p * ifftshift(U)) ./ √(size(U, 1) * size(U, 2)) + end + + return out, (;L=L_new) +end + +""" + Fraunhofer(U, z, λ, L; skip_final_phase=true) + + +This returns a function for efficient reuse of pre-calculated kernels. +See [`fraunhofer`](@ref) for the full documentation. + +""" +function Fraunhofer(U, z, λ, L; skip_final_phase=true) + @assert size(U, 1) == size(U, 2) + L_new = λ * z / L * size(U, 1) + Ns = size(U)[1:2] + + if skip_final_phase + phasefactor = nothing + else + k = eltype(U)(2π) / λ + # output coordinates + y = similar(U, real(eltype(U)), (Ns[1], 1)) + y .= (fftpos(L, Ns[1], CenterFT)) + x = similar(U, real(eltype(U)), (1, Ns[2])) + x .= (fftpos(L, Ns[2], CenterFT))' + phasefactor = (-1im) .* exp.(1im * k / (2 * z) .* (x.^2 .+ y.^2)) + end + + buffer = zero.(U) + + FFTplan = plan_fft!(buffer, (1,2)) + + return FraunhoferOp{typeof(L), typeof(buffer), typeof(phasefactor), + typeof(FFTplan)}(buffer, phasefactor, L_new, FFTplan) +end + +struct FraunhoferOp{T, B, PF, P} + buffer::B + phasefactor::PF + L::T + FFTplan::P +end + +function (fraunhofer::FraunhoferOp)(field) + buffer = fraunhofer.buffer + ifftshift!(buffer, field) + fraunhofer.FFTplan * buffer + buffer ./= √(size(field, 1) * size(field, 2)) + out = fftshift(buffer) + if !isnothing(fraunhofer.phasefactor) + out .*= fraunhofer.phasefactor + end + return out, (; fraunhofer.L) +end + + +function ChainRulesCore.rrule(f::FraunhoferOp, U) + field_and_tuple = f(U) + + function f_pullback(ȳ) + buffer = f.buffer + y2 = ȳ.backing[1] + if !isnothing(f.phasefactor) + y2 = y2 .* conj.((f.phasefactor)) + end + ifftshift!(buffer, y2) + buffer .*= √(size(U, 1) * size(U, 2)) + res = fftshift(inv(f.FFTplan) * buffer) + return NoTangent(), res + end + return field_and_tuple, f_pullback +end diff --git a/src/propagation.jl b/src/propagation.jl new file mode 100644 index 0000000..8ad5916 --- /dev/null +++ b/src/propagation.jl @@ -0,0 +1,64 @@ +""" + _propagation_variables(field, λ, L) + +Internal method to create variables we need for propagation such as frequencies in Fourier space, etc.. +""" +function _propagation_variables(field::AbstractArray{T, M}, λ::Number, L::NTuple{2, T2}) where {T, T2, M} + # wave number + k = real(T)(2π) / λ + # number of samples + Ns = size(field)[1:2] + # sample spacing + dx = L ./ Ns + + # frequency spacing + df = 1 ./ L + + # total size in frequency space + Lf = Ns .* df + + # frequencies centered around first entry + # 1D vectors each + f_y = similar(field, real(eltype(field)), (Ns[1], 1)) + f_y .= fftfreq(Ns[1], Lf[1]) + f_x = similar(field, real(eltype(field)), (1, Ns[2])) + f_x .= fftfreq(Ns[2], Lf[2])' + + # y and x positions in real space, use correct spacing -> fftpos + y = similar(field, real(eltype(field)), (Ns[1], 1)) + y .= (fftpos(L[1], Ns[1], CenterFT)) + x = similar(field, real(eltype(field)), (1, Ns[2])) + x .= (fftpos(L[2], Ns[2], CenterFT))' + + return (; k, dx, df, f_x, f_y, x, y) +end + + +function _propagation_variables(field::AbstractArray{T, M}, λ, L::Number) where {T, M} + _propagation_variables(field, λ, (L, L)) +end + + + +function _real_space_kernel(field::AbstractArray{T, N}, z, λ, L) where {T, N} + L = L isa Number ? (L, L) : L + # wave number + k = T(2π) / λ + # number of samples + Ns = size(field)[1:2] + + # y and x positions in real space, use correct spacing -> fftpos + y = similar(field, real(eltype(field)), (Ns[1], 1)) + y .= (fftpos(L[1], Ns[1], CenterFT)) + x = similar(field, real(eltype(field)), (1, Ns[2])) + x .= (fftpos(L[2], Ns[2], CenterFT))' + + + #y = ifftshift(y, (1,)) + #x = ifftshift(x, (2,)) + + r = sqrt.(x.^2 .+ y.^2 .+ z .^2) + kernel = T(2 ./ π) .* z ./ r .* (1 ./ r .- 1im * k) .* exp.(1im .* k .* r) ./ r + kernel .*= CuArray(rr2(size(field)) .< 60 .^2) + return kernel +end diff --git a/src/scalable_angular_spectrum.jl b/src/scalable_angular_spectrum.jl new file mode 100644 index 0000000..2c1f9cc --- /dev/null +++ b/src/scalable_angular_spectrum.jl @@ -0,0 +1,216 @@ +export ScalableAngularSpectrum + + # highly optimized version with pre-planning +struct ScalableAngularSpectrum{AP, AP2, AB, T, P} + ΔH::AP + H₁::AP + H₂::AP2 # could be nothing! + buffer::AB + buffer2::AB + L::T + FFTplan::P + pad_factor::Int +end + +""" + _propagation_variables(field, z, λ, L) + +Internal method to create variables we need for propagation such as frequencies in Fourier space, etc.. +""" +function _propagation_variables(field::AbstractArray{T, M}, z, λ, L) where {T, M} + @assert size(field, 1) == size(field, 2) "Quadratic fields only working currently" + + # wave number + k = T(2π) / λ + # number of samples + N = size(field, 1) + # sample spacing + dx = L / N + # frequency spacing + df = 1 / L + # total size in frequency space + Lf = N * df + + # frequencies centered around first entry + # 1D vectors each + f_y = similar(field, real(eltype(field)), (N,)) + f_y .= fftfreq(N, Lf) + f_x = f_y' + + # y and x positions in real space + #y = ifftshift(range(-L/2, L/2, length=N)) + y = similar(field, real(eltype(field)), (N,)) + y .= fftpos(L, N, CenterFT) + y = ifftshift(y) + x = y' + + return (; k, dx, df, f_x, f_y, x, y) +end + + +# helper function for smooth window +function find_width_window(ineq_1D::AbstractVector, bandlimit_border) + ineq_1D = Array(ineq_1D) + bs = ineq_1D .≤ 1 + ind_x_first = findfirst(bs) + ind_x_last = findlast(bs) + + if isnothing(ind_x_first) || isnothing(ind_x_last) + return (1,1) + end + + diff_b = round(Int, 0.5 * (1 - bandlimit_border[1]) * (Tuple(ind_x_last)[1] - Tuple(ind_x_first)[1])) + diff_e = round(Int, 0.5 * (1 - bandlimit_border[2]) * (Tuple(ind_x_last)[1] - + Tuple(ind_x_first)[1])) + + ineq_v_b = ineq_1D[ind_x_first + diff_b] + ineq_v_e = ineq_1D[ind_x_first + diff_e] + + return ineq_v_b, ineq_v_e +end + +""" + ScalableAngularSpectrum(field, z, λ, L; kwargs...) + +Returns the electrical field with physical length `L` and wavelength `λ` propagated with the scalable angular spectrum method of plane waves (SAS) by the propagation distance `z`. + + +# Arguments +* `field`: Input field +* `z`: propagation distance +* `λ`: wavelength of field +* `L`: field size (can be a scalar or a tuple) indicating field size + + +# References +* [Rainer Heintzmann, Lars Loetgering, and Felix Wechsler, "Scalable angular spectrum propagation," Optica 10, 1407-1416 (2023)](https://opg.optica.org/optica/viewmedia.cfm?uri=optica-10-11-1407&html=true) +""" +function ScalableAngularSpectrum(ψ₀::AbstractArray{T}, z, λ, L ; + skip_final_phase=true, bandlimit_soft_px=20, + bandlimit_border=(0.8, 1)) where {T} + @assert bandlimit_soft_px ≥ 0 "bandlimit_soft_px must be ≥ 0" + @assert size(ψ₀, 1) == size(ψ₀, 2) "Restricted to auadratic fields." + + pad_factor = 2 + set_pad_zero = false + + N = size(ψ₀, 1) + z_limit = (- 4 * L * sqrt(8*L^2 / N^2 + λ^2) * sqrt(L^2 * inv(8 * L^2 + N^2 * λ^2)) / (λ * (-1+2 * sqrt(2) * sqrt(L^2 * inv(8 * L^2 + N^2 * λ^2))))) + + # vignetting limit + z > z_limit && @warn "Propagated field might be affected by vignetting" + L_new = pad_factor * L + + # applies zero padding + ψ_p = select_region(ψ₀, new_size=size(ψ₀) .* pad_factor) + k, dx, df, f_x, f_y, x, y = _propagation_variables(ψ_p, z, λ, L_new) + M = λ * z * N / L^2 / 2 + + # calculate anti_aliasing_filter for precompensation + cx = λ .* f_x + cy = λ .* f_y + tx = L_new / 2 / z .+ abs.(λ .* f_x) + ty = L_new / 2 / z .+ abs.(λ .* f_y) + + # smooth window function + smooth_f(x, α, β) = hann(scale(x, α, β)) + # find boundary for soft hann + ineq_x = fftshift(cx[1, :].^2 .* (1 .+ tx[1, :].^2) ./ tx[1, :].^2 .+ cy[1, :].^2) + limits = find_width_window(ineq_x, bandlimit_border) + + # bandlimit filter for precompensation + W = .*(smooth_f.(cx.^2 .* (1 .+ tx.^2) ./ tx.^2 .+ cy.^2, limits...), + smooth_f.(cy.^2 .* (1 .+ ty.^2) ./ ty.^2 .+ cx.^2, limits...)) + + # ΔH is the core part of Fresnel and AS + H_AS = sqrt.(0im .+ 1 .- abs2.(f_x .* λ) .- abs2.(f_y .* λ)) + H_Fr = 1 .- abs2.(f_x .* λ) / 2 .- abs2.(f_y .* λ) / 2 + # take the difference here, key part of the ScaledAS + ΔH = W .* exp.(1im .* k .* z .* (H_AS .- H_Fr)) + + + # new sample coordinates + dq = λ * z / L_new + Q = dq * N * pad_factor + q_y = similar(ψ_p, pad_factor * N) + # fftpos generates coordinates from -L/2 to L/2 but excluding the last + # final bit + q_y .= fftpos(dq * pad_factor * N, pad_factor * N, CenterFT) + q_y = ifftshift(q_y) + q_x = q_y' + + # calculate phases of Fresnel + H₁ = exp.(1im .* k ./ (2 .* z) .* (x .^ 2 .+ y .^ 2)) + + # skip final phase because often undersampled + if skip_final_phase + H₂ = nothing + else + H₂ = (exp.(1im .* k .* z) .* + exp.(1im .* k ./ (2 .* z) .* (q_x .^ 2 .+ q_y .^2))) + end + + FFTplan = plan_fft!(ψ_p, (1,2)) + + buffer = similar(ψ_p) + buffer2 = similar(buffer) + return ScalableAngularSpectrum{typeof(ΔH), typeof(H₂), typeof(buffer), + real(T), typeof(FFTplan)}( + ΔH, H₁, H₂, buffer, buffer2, Q, FFTplan, pad_factor) +end + + + +function (sas::ScalableAngularSpectrum)(ψ::AbstractArray{T}; return_view=false) where T + p = sas.FFTplan + fill!(sas.buffer2, 0) + ψ_p = set_center!(sas.buffer2, ψ) + ψ_p = ifftshift!(sas.buffer, ψ_p) + ψ_p_f = p * ψ_p + ψ_p_f .*= sas.ΔH + ψ_precomp = inv(p) * ψ_p_f + ψ_precomp .*= sas.H₁ + ψ_p_final = p * ψ_precomp + ψ_p_final .*= 1 / (1im * sqrt(T(size(ψ_precomp, 1) * size(ψ_precomp, 2)))) + + if !isnothing(sas.H₂) + ψ_p_final .*= sas.H₂ + end + fftshift!(sas.buffer2, ψ_p_final) + + ψ_final = crop_center(sas.buffer2, size(ψ); return_view) + return ψ_final, (; sas.L) +end + + + +function ChainRulesCore.rrule(sas::ScalableAngularSpectrum, ψ::AbstractArray{T}) where T + field_and_tuple = sas(ψ) + function sas_pullback(ȳ) + p = sas.FFTplan + y2 = ȳ.backing[1] + + fill!(sas.buffer2, 0) + + set_center!(sas.buffer2, y2) + ifftshift!(sas.buffer, sas.buffer2) + + if !isnothing(sas.H₂) + sas.buffer .*= conj.(sas.H₂) + end + sas.buffer .*= conj.(1 / (1im * sqrt(T(size(sas.buffer2, 1) * size(sas.buffer2, 2))))) + sas.buffer .*= (T(size(sas.buffer2, 1) * size(sas.buffer2, 2))) + inv(p) * sas.buffer + sas.buffer .*= conj.(sas.H₁) + (p) * sas.buffer + sas.buffer .*= conj.(sas.ΔH) + inv(p) * sas.buffer + fftshift!(sas.buffer2, sas.buffer) + + final = crop_center(sas.buffer2, size(ψ)) + + + return NoTangent(), final + end + return field_and_tuple, sas_pullback +end diff --git a/src/utils.jl b/src/utils.jl new file mode 100644 index 0000000..48a92b2 --- /dev/null +++ b/src/utils.jl @@ -0,0 +1,237 @@ +export crop_center +export set_center! +export pad + + +hann(x) = sin(π * x / 2)^2 +scale(x, α, β) = clamp(1 - (x - α) / (β - α), 0, 1) +smooth_f(x, α, β) = hann(scale(x, α, β)) + +""" + get_indices_around_center(i_in, i_out) + +A function which provides two output indices `i1` and `i2` +where `i2 - i1 = i_out` + +The indices are chosen in a way that the range `i1:i2` +cuts the interval `1:i_in` in a way that the center frequency +stays at the center position. + +Works for both odd and even indices. +""" +function get_indices_around_center(i_in, i_out) + if (mod(i_in, 2) == 0 && mod(i_out, 2) == 0 + || mod(i_in, 2) == 1 && mod(i_out, 2) == 1) + x = (i_in - i_out) ÷ 2 + return 1 + x, i_in - x + elseif mod(i_in, 2) == 1 && mod(i_out, 2) == 0 + x = (i_in - 1 - i_out) ÷ 2 + return 1 + x, i_in - x - 1 + elseif mod(i_in, 2) == 0 && mod(i_out, 2) == 1 + x = (i_in - (i_out - 1)) ÷ 2 + return 1 + x, i_in - (x - 1) + end +end + +""" + pad(arr::AbstractArray{T, N}, new_size; value=zero(T)) + +Pads `arr` with `values` around such that the resulting array size is `new_size`. + +See also [`crop_center`](@ref), [`set_center!`](@ref). + +```jldoctest +julia> pad(ones((2,2)), 2) +4×4 Matrix{Float64}: + 0.0 0.0 0.0 0.0 + 0.0 1.0 1.0 0.0 + 0.0 1.0 1.0 0.0 + 0.0 0.0 0.0 0.0 +``` +""" +function pad(arr::AbstractArray{T, N}, new_size::NTuple; value=zero(T)) where {T, N} + n_arr = similar(arr, new_size) + fill!(n_arr, value) + # don't do broadcast. Does not fit with the meaning of pad + set_center!(n_arr, arr, broadcast=false) +end + + +""" + pad(arr::AbstractArray{T, N}, M; value=zero(T)) + +Pads `arr` with `values` around such that the resulting array size is `round(Int, M .* size(arr))`. +If `M isa Integer`, no rounding is needed. + + +!!! warn "Automatic Differentation" + If `M isa AbstractFloat`, automatic differentation might fail because of size failures + + +See also [`crop_center`](@ref), [`set_center!`](@ref). + +```jldoctest +julia> pad(ones((2,2)), 2) +4×4 Matrix{Float64}: + 0.0 0.0 0.0 0.0 + 0.0 1.0 1.0 0.0 + 0.0 1.0 1.0 0.0 + 0.0 0.0 0.0 0.0 +``` +""" +function pad(arr::AbstractArray{T, N}, M::Number; value=zero(T)) where {T, N} + pad(arr, round.(Int, size(arr) .* M), value=value) +end + + +""" + crop_center(arr, new_size) + +Extracts a center of an array. +`new_size_array` must be list of sizes indicating the output +size of each dimension. Centered means that a center frequency +stays at the center position. Works for even and uneven. +If `length(new_size_array) < length(ndims(arr))` the remaining dimensions +are untouched and copied. + + +See also [`pad`](@ref), [`set_center!`](@ref). + +# Examples +```jldoctest +julia> crop_center([1 2; 3 4], (1,)) +1×2 Matrix{Int64}: + 3 4 + +julia> crop_center([1 2; 3 4], (1,1)) +1×1 Matrix{Int64}: + 4 + +julia> crop_center([1 2; 3 4], (2,2)) +2×2 Matrix{Int64}: + 1 2 + 3 4 +``` +""" +function crop_center(arr, new_size::NTuple{N}; return_view=true) where {N} + M = ndims(arr) + @assert N ≤ M "Can't specify more dimensions than the array has." + @assert all(new_size .≤ size(arr)[1:N]) "You can't extract a larger array than the input array." + @assert Base.require_one_based_indexing(arr) "Require one based indexing arrays" + + out_indices = ntuple(i -> let + if i ≤ N + inds = get_indices_around_center(size(arr, i), + new_size[i]) + inds[1]:inds[2] + else + 1:size(arr, i) + end + end, + Val(M)) + + + # if return_view + # return @inbounds view(arr, out_indices...) + # else + return @inbounds arr[out_indices...] + # end +end + + +""" + set_center!(arr_large, arr_small; broadcast=false) + +Puts the `arr_small` central into `arr_large`. + +The convention, where the center is, is the same as the definition +as for FFT based centered. + +Function works both for even and uneven arrays. +See also [`crop_center`](@ref), [`pad`](@ref), [`set_center!`](@ref). + +# Keyword +* If `broadcast==false` then a lower dimensional `arr_small` will not be broadcasted +along the higher dimensions. +* If `broadcast==true` it will be broadcasted along higher dims. + + +See also [`crop_center`](@ref), [`pad`](@ref). + + +# Examples +```jldoctest +julia> set_center!([1, 1, 1, 1, 1, 1], [5, 5, 5]) +6-element Vector{Int64}: + 1 + 1 + 5 + 5 + 5 + 1 + +julia> set_center!([1, 1, 1, 1, 1, 1], [5, 5, 5, 5]) +6-element Vector{Int64}: + 1 + 5 + 5 + 5 + 5 + 1 + +julia> set_center!(ones((3,3)), [5]) +3×3 Matrix{Float64}: + 1.0 1.0 1.0 + 1.0 5.0 1.0 + 1.0 1.0 1.0 + +julia> set_center!(ones((3,3)), [5], broadcast=true) +3×3 Matrix{Float64}: + 1.0 1.0 1.0 + 5.0 5.0 5.0 + 1.0 1.0 1.0 +``` +""" +function set_center!(arr_large::AbstractArray{T, N}, arr_small::AbstractArray{T1, M}; + broadcast=false) where {T, T1, M, N} + @assert N ≥ M "Can't put a higher dimensional array in a lower dimensional one." + + if broadcast == false + inds = ntuple(i -> begin + a, b = get_indices_around_center(size(arr_large, i), size(arr_small, i)) + a:b + end, + Val(N)) + arr_large[inds..., ..] .= arr_small + else + inds = ntuple(i -> begin + a, b = get_indices_around_center(size(arr_large, i), size(arr_small, i)) + a:b + end, + Val(M)) + arr_large[inds..., ..] .= arr_small + end + + + return arr_large +end + +function ChainRulesCore.rrule(::typeof(crop_center), arr, new_size::NTuple{N, <:Int}) where N + y = crop_center(arr, new_size) + function crop_center_pullback(ȳ) + c̄ = pad(ȳ, size(arr), value=zero(eltype(arr))) + return NoTangent(), c̄, NoTangent() + end + return y, crop_center_pullback +end + + +function ChainRulesCore.rrule(::typeof(pad), arr, M::Union{NTuple, Number}; value=zero(eltype(arr))) + y = pad(arr, M; value=value) + function pad_pullback(ȳ) + @assert size(y) == size(ȳ) + p̄ = crop_center(ȳ, size(arr)) + return NoTangent(), p̄, NoTangent() + end + return y, pad_pullback +end diff --git a/test/angular_spectrum.jl b/test/angular_spectrum.jl new file mode 100644 index 0000000..298fe36 --- /dev/null +++ b/test/angular_spectrum.jl @@ -0,0 +1,48 @@ +@testset "Angular Spectrum" begin + + @testset "Test gradient with Finite Differences" begin + field = zeros(ComplexF64, (24, 24)) + field[14:16, 14:16] .= 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 .+ cis(1) ≈ out2 .+ cis(1) + AS, _ = AngularSpectrum(field, 100e-6, 633e-9, 100e-6) + + f_AS(x) = sum(abs2.(x .- AS(cis.(x))[1])) + + out3 = gradient(f_AS, field)[1] + + @test out3 ≈ out1 + + + field = zeros(ComplexF64, (15, 15)) + field[5:6, 3:8] .= 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 .+ cis(1) ≈ out2 .+ cis(1) + AS, _ = AngularSpectrum(field, 100e-6, 633e-9, 100e-6) + + f_AS(x) = sum(abs2.(x .- AS(cis.(x))[1])) + + out3 = gradient(f_AS, field)[1] + + @test out3 ≈ out1 + + end + + @testset "double slit" begin + include("double_slit.jl") + end + + @testset "Gaussian beam" begin + include("gaussian_beam.jl") + end +end diff --git a/test/double_slit.html b/test/double_slit.html new file mode 100644 index 0000000..04f4707 --- /dev/null +++ b/test/double_slit.html @@ -0,0 +1,16 @@ + + + + + + +
\ No newline at end of file diff --git a/test/double_slit.jl b/test/double_slit.jl new file mode 100644 index 0000000..98e6a8c --- /dev/null +++ b/test/double_slit.jl @@ -0,0 +1,89 @@ +### A Pluto.jl notebook ### +# v0.19.30 + +using Markdown +using InteractiveUtils + +# ╔═╡ 9745c33c-c8d4-11ed-216d-cb75e0aaec28 +# ╠═╡ skip_as_script = true +#=╠═╡ +begin + using Pkg, Revise + Pkg.activate("../examples/.") +end + ╠═╡ =# + +# ╔═╡ 4acbbe98-f033-42c5-b6d0-60a5ddf2570e +using Test, WaveOpticsPropagation, IndexFunArrays, NDTools, FourierTools + +# ╔═╡ 2b5650e9-db71-459d-ae13-7b52ca8d2a85 +# ╠═╡ skip_as_script = true +#=╠═╡ +using Plots, ImageShow + ╠═╡ =# + +# ╔═╡ 83144a55-ac66-4b92-b9c7-36daacbab791 +""" + double_slit + +`d` is distance between slits in pixel +`b` is the width of the slits in pixel +""" +function double_slit(N, d, b, offset, z, λ, L) + + slit_init = (box(ComplexF32, (N, N), (N, b), offset=(11, offset + N÷2 + 1 - d ÷2 )) .+ + box(ComplexF32, (N, N), (N, b), offset=(11, offset + N÷2 + 1 + d ÷2 ))) + + xpos = fftpos(L, N, NDTools.CenterFT) .- offset .* L ./ N + sinθ = sin.(atan.(xpos, z)) + d_m = d * L / N + b_m = b * L / N + slit_ana = cos.(π .* d_m .* sinθ ./ λ).^2 .* sinc.(b_m .* sinθ ./ λ).^2 + + slit_prop = (angular_spectrum(slit_init, z, λ, L)[1]) + slit_prop2 = (AngularSpectrum(slit_init, z, λ, L)[1](slit_init))[1] + @test slit_prop2 ≈ slit_prop + + slit_prop = abs2.(slit_prop) + slit_prop2 = abs2.(slit_prop2) + + slit_prop ./= maximum(slit_prop[11, :]) + slit_prop2 ./= maximum(slit_prop2[11, :]) + + @test all(.≈(0.01f0 .+ slit_ana, 0.01f0 .+ slit_prop[11, :], rtol=0.3)) + @test findmax(slit_ana) == findmax(slit_prop[11, :]) + + return slit_init, slit_prop, slit_prop2, slit_ana +end + +# ╔═╡ c12b2aac-5e48-47b2-a0f8-e626a0020b20 +slit_init, slit_prop, slit_prop2, slit_ana = double_slit(100, 10, 3, 0, 25e-3, 1000e-9, 1e-3) + +# ╔═╡ 33a6b4e4-ee04-45fe-96b4-6fa8614b5c6c +#=╠═╡ +begin + #plot(abs2.(slit_init[11, :])) + plot(slit_prop[11, :]) + plot!(slit_ana) +end + ╠═╡ =# + +# ╔═╡ 6b2554ac-613a-4b06-b9a5-fb96a81816ee +double_slit(512, 10, 3, -30, 5e-3, 632e-9, 1e-3) + +# ╔═╡ 05587986-fe48-4ac8-91eb-5fedfbf6634e +double_slit(512, 10, 3, 30, 5e-3, 1000e-9, 1e-3) + +# ╔═╡ 8c9133cd-ae62-4a3d-81b8-4e30c795a04f + double_slit(100, 10, 3, 0, 25e-3, 1000e-9, 1e-3) + +# ╔═╡ Cell order: +# ╠═9745c33c-c8d4-11ed-216d-cb75e0aaec28 +# ╠═4acbbe98-f033-42c5-b6d0-60a5ddf2570e +# ╠═2b5650e9-db71-459d-ae13-7b52ca8d2a85 +# ╠═83144a55-ac66-4b92-b9c7-36daacbab791 +# ╠═c12b2aac-5e48-47b2-a0f8-e626a0020b20 +# ╠═33a6b4e4-ee04-45fe-96b4-6fa8614b5c6c +# ╠═6b2554ac-613a-4b06-b9a5-fb96a81816ee +# ╠═05587986-fe48-4ac8-91eb-5fedfbf6634e +# ╠═8c9133cd-ae62-4a3d-81b8-4e30c795a04f diff --git a/test/double_slit.plutostate b/test/double_slit.plutostate new file mode 100644 index 0000000..ba7d4c6 Binary files /dev/null and b/test/double_slit.plutostate differ diff --git a/test/fraunhofer.jl b/test/fraunhofer.jl new file mode 100644 index 0000000..0072e74 --- /dev/null +++ b/test/fraunhofer.jl @@ -0,0 +1,49 @@ +@testset "Angular Spectrum" begin + L = 100f-3 + λ = 633f-9 + z = 1 + N = 256 + ΔS = 18 + ΔW = 2 + slit = zeros(ComplexF32, (N, N)) + + x = WaveOpticsPropagation.fftpos(L, N) + W = x[1 + ΔW * 2 + 1] - x[1] + S = x[1 + ΔS * 2] - x[1] + + L_new_ref = λ * z / L * N + mid = N ÷ 2+ 1 + slit[:, mid-ΔS-ΔW:mid-ΔS+ΔW] .= 1 + slit[:, mid+ΔS-ΔW:mid+ΔS+ΔW] .= 1 + output, t = fraunhofer(slit, z, λ, L) + L_new = t.L + + efficient_fraunhofer = Fraunhofer(slit, z, λ, L); + output2, L_new2 = efficient_fraunhofer(slit) + + + @test L_new ≈ L_new_ref + @test L_new2.L ≈ L_new_ref + + @test output2 ≈ output + + I_analytical(x) = sinc(W * x / λ / z)^2 * cos(π * S * x / λ / z)^2 + intensity = abs2.(output) + intensity ./= maximum(intensity) + xpos_out = WaveOpticsPropagation.fftpos(L_new, N, NDTools.CenterFT) + + @test all(≈(I_analytical.(xpos_out)[110:150] .+ 1, 1 .+ intensity[129, 110:150, :], rtol=1f-2)) + + arr = randn(ComplexF32, (N, N)) + fr = Fraunhofer(arr, z, λ, L) + f(x) = sum(abs2, arr .- fr(x)[1]) + f2(x) = sum(abs2, arr .- fraunhofer(x, z, λ, L)[1]) + @test Zygote.gradient(f, arr)[1] ≈ Zygote.gradient(f2, arr)[1] + + arr = randn(ComplexF32, (15, 15)) + fr = Fraunhofer(arr, z, λ, L, skip_final_phase=false) + f(x) = sum(abs2, arr .- fr(x)[1]) + f2(x) = sum(abs2, arr .- fraunhofer(x, z, λ, L, skip_final_phase=false)[1]) + @test Zygote.gradient(f, arr)[1] ≈ Zygote.gradient(f2, arr)[1] + +end diff --git a/test/gaussian_beam.html b/test/gaussian_beam.html new file mode 100644 index 0000000..e1d565a --- /dev/null +++ b/test/gaussian_beam.html @@ -0,0 +1,16 @@ + + + + + + +
\ No newline at end of file diff --git a/test/gaussian_beam.jl b/test/gaussian_beam.jl new file mode 100644 index 0000000..9b40e70 --- /dev/null +++ b/test/gaussian_beam.jl @@ -0,0 +1,75 @@ +### A Pluto.jl notebook ### +# v0.19.30 + +using Markdown +using InteractiveUtils + +# ╔═╡ 9902bc0e-c8dc-11ed-05fe-a9f16c85ba7b +# ╠═╡ skip_as_script = true +#=╠═╡ +begin + using Pkg, Revise + Pkg.activate("../examples/.") +end + ╠═╡ =# + +# ╔═╡ 748c98db-f175-4b8d-87d4-eae2ec6f26ab +# ╠═╡ skip_as_script = true +#=╠═╡ +using Plots, ImageShow, PlutoUI + ╠═╡ =# + +# ╔═╡ 38f7c686-48c2-4051-9e2e-8f2aad6e95ee +using Test, WaveOpticsPropagation, IndexFunArrays, NDTools, FourierTools + +# ╔═╡ 0299632f-e805-4e27-82e8-d67a756409f2 + + +# ╔═╡ 96986249-3edb-4fae-b4a5-ec7ae18fdb59 +function test_gauss_consistency(λ, L, N, z, z_init, w_0; do_test=true) + y = reshape(Float32.(fftpos(L[1], N[1], CenterFT)) ,:, 1) + x = reshape(Float32.(fftpos(L[2], N[2], CenterFT)), 1, :) + field = gauss_beam.(y, x, z_init, λ, w_0); + field_as = angular_spectrum(field, z, λ, L)[1]; + field_z = gauss_beam.(y, x, z_init + z, λ, w_0); + + + if do_test + @test ≈(field_as .+ maximum(abs.(field_as)) , maximum(abs.(field_as)) .+ field_z, + rtol=0.004) + @test all(≈(field_as .+ maximum(abs.(field_as)) , maximum(abs.(field_as)) .+ field_z, + rtol=0.1)) + @test sum(abs2.(field_z)) ≈ sum(abs2.(field_as)) ≈ sum(abs2.(field)) + end + + return field, field_as, field_z +end + +# ╔═╡ c732c1d7-6ae8-41ab-ae2e-d635fa9d66a3 +res1 = test_gauss_consistency(405f-9, (10f-4, 20.0f-4), (512, 512), 1f-2, 1f-10, 0.01f-3) + +# ╔═╡ a5e7878f-4ea1-4f75-87f8-fafd09f5ad82 +res2 = test_gauss_consistency(10f-9, (1f-4, 0.5f-4), (400, 400), 0.02f-2, 0.05f-2, 0.003f-3, do_test = true) + +# ╔═╡ 22e070fd-b27f-4cb2-8ad7-71f409e881bf +# ╠═╡ skip_as_script = true +#=╠═╡ +simshow([res1[2] res1[3]], γ=0.1) + ╠═╡ =# + +# ╔═╡ 1f026add-c513-42fa-a26c-5da40badf8a8 +# ╠═╡ skip_as_script = true +#=╠═╡ +simshow([res2[1] res2[2] res2[3]], γ=0.2) + ╠═╡ =# + +# ╔═╡ Cell order: +# ╠═9902bc0e-c8dc-11ed-05fe-a9f16c85ba7b +# ╠═748c98db-f175-4b8d-87d4-eae2ec6f26ab +# ╠═38f7c686-48c2-4051-9e2e-8f2aad6e95ee +# ╠═0299632f-e805-4e27-82e8-d67a756409f2 +# ╠═96986249-3edb-4fae-b4a5-ec7ae18fdb59 +# ╠═c732c1d7-6ae8-41ab-ae2e-d635fa9d66a3 +# ╠═a5e7878f-4ea1-4f75-87f8-fafd09f5ad82 +# ╠═22e070fd-b27f-4cb2-8ad7-71f409e881bf +# ╠═1f026add-c513-42fa-a26c-5da40badf8a8 diff --git a/test/gaussian_beam.plutostate b/test/gaussian_beam.plutostate new file mode 100644 index 0000000..19d8f20 Binary files /dev/null and b/test/gaussian_beam.plutostate differ diff --git a/test/runtests.jl b/test/runtests.jl new file mode 100644 index 0000000..95a65fd --- /dev/null +++ b/test/runtests.jl @@ -0,0 +1,16 @@ +using WaveOpticsPropagation +using Test +using Zygote +using ChainRulesTestUtils +using IndexFunArrays, NDTools, FourierTools +using FiniteDifferences + +@testset "WaveOpticsPropagation.jl" begin + include("utils.jl") + include("angular_spectrum.jl") + include("scalable_angular_spectrum.jl") + include("fraunhofer.jl") +end + +return nothing + diff --git a/test/scalable_angular_spectrum.jl b/test/scalable_angular_spectrum.jl new file mode 100644 index 0000000..361de92 --- /dev/null +++ b/test/scalable_angular_spectrum.jl @@ -0,0 +1,23 @@ +@testset "Scalable Angular Spectrum" begin + + @testset "Test gradient with Finite Differences" begin + field = zeros(ComplexF64, (24, 24)) + field[14:16, 14:16] .= 1 + ss = ScalableAngularSpectrum(cis.(field), 100e-6, 633e-9, 100e-6) + gg(x) = sum(abs2.(x .- ss(cis.(x))[1])) + out2 = FiniteDifferences.grad(central_fdm(5, 1), gg, field)[1] + out1 = gradient(gg, field)[1] + @test out1 .+ cis(1) ≈ out2 .+ cis(1) + + field = zeros(ComplexF64, (19, 19)) + field[14:16, 14:16] .= 1 + ss = ScalableAngularSpectrum(cis.(field), 100e-6, 633e-9, 100e-6, skip_final_phase=false) + gg(x) = sum(abs2.(x .- ss(cis.(x))[1])) + out2 = FiniteDifferences.grad(central_fdm(5, 1), gg, field)[1] + out1 = gradient(gg, field)[1] + @test out1 .+ cis(1) ≈ out2 .+ cis(1) + + end + + +end diff --git a/test/utils.jl b/test/utils.jl new file mode 100644 index 0000000..836cece --- /dev/null +++ b/test/utils.jl @@ -0,0 +1,43 @@ + +@testset "Pad" begin + @test pad(ones((2, 2)), 2) == [0.0 0.0 0.0 0.0; 0.0 1.0 1.0 0.0; 0.0 1.0 1.0 0.0; 0.0 0.0 0.0 0.0] + @test pad(ones((1,)), 2) == [0.0, 1.0] + @test pad(ones((1,)), 3) == [0.0, 1.0, 0.0] + @test pad(ones((1,)), 4) == [0.0, 0.0, 1.0, 0.0] + @test pad(ones((1,)), 5) == [0.0, 0.0, 1.0, 0.0, 0.0] + @test pad(ones((2, 2)), 3) == [0.0 0.0 0.0 0.0 0.0 0.0; 0.0 0.0 0.0 0.0 0.0 0.0; 0.0 0.0 1.0 1.0 0.0 0.0; 0.0 0.0 1.0 1.0 0.0 0.0; 0.0 0.0 0.0 0.0 0.0 0.0; 0.0 0.0 0.0 0.0 0.0 0.0] + @test pad(ones((3, 3)), 3) == [0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0; 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0; 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0; 0.0 0.0 0.0 1.0 1.0 1.0 0.0 0.0 0.0; 0.0 0.0 0.0 1.0 1.0 1.0 0.0 0.0 0.0; 0.0 0.0 0.0 1.0 1.0 1.0 0.0 0.0 0.0; 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0; 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0; 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0] + + @test pad(ones((2, 2)), (2, 2, 2)) == [0.0 0.0; 0.0 0.0;;; 1.0 1.0; 1.0 1.0] + @test pad(ones((1, 1)), (1, 3, 4)) == [0.0 0.0 0.0;;; 0.0 0.0 0.0;;; 0.0 1.0 0.0;;; 0.0 0.0 0.0] + @test pad(ones((1, 1)), (1, 3, 4), value = 10) == [10.0 10.0 10.0;;; 10.0 10.0 10.0;;; 10.0 1.0 10.0;;; 10.0 10.0 10.0] + @test pad(ones((2, 2)), 2.5) == [0.0 0.0 0.0 0.0 0.0; 0.0 1.0 1.0 0.0 0.0; 0.0 1.0 1.0 0.0 0.0; 0.0 0.0 0.0 0.0 0.0; 0.0 0.0 0.0 0.0 0.0] +end + +@testset "set_center!" begin + @test set_center!([-1, 1, 1, 13, 1, 4], [5, 5, 5]) == [-1, 1, 5, 5, 5, 4] + @test set_center!(ones((3, 3)), [5]) == [1.0 1.0 1.0; 1.0 5.0 1.0; 1.0 1.0 1.0] + @test set_center!(ones((3, 3)), [5], broadcast = true) == [1.0 1.0 1.0; 5.0 5.0 5.0; 1.0 1.0 1.0] +end + +@testset "crop center" begin + @test crop_center([1 2; 3 4], (1,)) == [3 4] + @test crop_center([1 2 3; 3 4 5], (1,)) == [3 4 5] + @test crop_center([1 2 3; 3 4 5], (1, 2)) == [3 4] + @test crop_center([1 2 3; 3.3 4 5], (1, 2)) == [3.3 4.0] + @test crop_center([1 2 3; 3.3 4 5], (2, 2)) == [1.0 2.0; 3.3 4.0] +end + + +@testset "test rrule for pad and crop_center" begin + test_rrule(pad, randn((4,2)), 2, check_thunked_output_tangent=false) + test_rrule(pad, randn((2,8)), 2.5, check_thunked_output_tangent=false) + test_rrule(pad, randn((1,9)), 1, check_thunked_output_tangent=false) + test_rrule(pad, randn((1,9)), (2,20), check_thunked_output_tangent=false) + test_rrule(pad, randn((1,9)), (1,9), check_thunked_output_tangent=false) + + test_rrule(crop_center, randn((6,6)), (1,1), check_thunked_output_tangent=false) + test_rrule(crop_center, randn((6,4)), (2,3), check_thunked_output_tangent=false) + test_rrule(crop_center, randn((5,5)), (3,3), check_thunked_output_tangent=false) + test_rrule(crop_center, randn((5,5)), (5,5), check_thunked_output_tangent=false) +end