Add shifted AS

Author: roflmaostc

README.md

@@ -30,6 +30,7 @@ Otherwise, just look into the [examples folder](examples/).
     * [x] Angular Spectrum Method of Plane Waves (AS)
     * [x] Fraunhofer Diffraction
     * [x] [Scalable Angular Spectrum propagation](https://opg.optica.org/optica/viewmedia.cfm?uri=optica-10-11-1407&html=true)
+    * [x] Shifted Angular Spectrum propagation
     * [ ] Fresnel Propagation with Scaling Behaviour (no priority yet, PR are welcome for that. In principle very similar to the other methods.)
 * [x] CUDA support
 * [x] Differentiable (mainly based on Zygote.jl and ChainRulesCore.jl)

docs/src/functions.md

@@ -4,6 +4,7 @@ They create efficient functions which avoid recalculating some parts.
 AngularSpectrum
 Fraunhofer
 ScalableAngularSpectrum
+ShiftedAngularSpectrum
 ```
 
 # Propagation

examples/shifted_angular_spectrum.jl

@@ -27,13 +27,13 @@ end
 md"# Define field with a ramp"
 
 # ╔═╡ 2f6871e8-7c11-49c0-ba9a-dc498e8eb39d
-N = 128
+N = 256
 
 # ╔═╡ 64b448ee-5ccc-4f87-8ee0-20d2d6a41a3b
 sz = (N, N)
 
 # ╔═╡ 90286b89-aedd-4ece-b9d6-e5c26c6ad635
-α = 0 * deg2rad(10f0)
+α = deg2rad(10f0)
 
 # ╔═╡ dc01bc87-ffd7-400f-bbf2-3b00a3a84b78
 L = 50f-6
@@ -60,7 +60,7 @@ md"# Compare AS and shifted AS"
 res_AS = AngularSpectrum(field, z, λ, L)[1](field)[1];
 
 # ╔═╡ 9cbafe25-3af6-4bcf-833c-8d3d7ca428a2
-res = shifted_angular_spectrum(field, z, λ, L, (α , 0), bandlimit=true)
+res = ShiftedAngularSpectrum(field, z, λ, L, (α , 0), bandlimit=true)[1](field)
 
 # ╔═╡ 4efdc02b-4f69-4893-a410-6c6bbb765bab
 shift = (z .* tan.(α) ./ L .* N)[1]
@@ -100,6 +100,12 @@ all(.≈(1 .+ FourierTools.shift(res[1], (shift, 0))[round(Int, shift)+1:end, :]
 # ╔═╡ d7eac41e-c5c0-46f8-9b2f-d2f116b65d95
 
 
+# ╔═╡ 4c3d7581-df26-4bec-9e0e-44279158d8b9
+
+
+# ╔═╡ ed3ec3b6-4825-4cd8-9d95-13d978d64584
+
+
 # ╔═╡ Cell order:
 # ╠═b11e7be2-b315-11ee-27e7-abecfdbe64b6
 # ╠═3a5d9f20-a01d-481b-9858-b8e523ba7a20
@@ -129,3 +135,5 @@ all(.≈(1 .+ FourierTools.shift(res[1], (shift, 0))[round(Int, shift)+1:end, :]
 # ╠═154b48f2-58b8-4e4f-8648-8ea771125be5
 # ╠═f3f7bc9e-a16f-49d3-920f-031326f5f1af
 # ╠═d7eac41e-c5c0-46f8-9b2f-d2f116b65d95
+# ╠═4c3d7581-df26-4bec-9e0e-44279158d8b9
+# ╟─ed3ec3b6-4825-4cd8-9d95-13d978d64584

src/shifted_angular_spectrum.jl

@@ -53,70 +53,164 @@ function _prepare_shifted_angular_spectrum(field::AbstractArray{CT}, z, λ, L, 
     shift = txy .* z
 
     ya = similar(field_new, real(eltype(field)), (size(field_new, 1), 1))
-    ya .= (fftpos(L_new[1], size(field_new, 1), CenterFT)) .+ shift[1]
+    Zygote.@ignore ya .= (fftpos(L_new[1], size(field_new, 1), CenterFT)) .+ shift[1]
     xa = similar(field_new, real(eltype(field)), (1, size(field_new, 2)))
-    xa .= (fftpos(L_new[2], size(field_new, 2), CenterFT))' .+ shift[2]
+    Zygote.@ignore xa .= (fftpos(L_new[2], size(field_new, 2), CenterFT))' .+ shift[2]
 
     ramp_before = ifftshift(exp.(1im .* 2 .* T(π) ./ λ .* (sxy[2] .* x .+ sxy[1] .* y)), (1,2))
     ramp_after = ifftshift(exp.(1im .* 2 .* T(π) ./ λ .* (sxy[2] .* xa .+ sxy[1] .* ya)), (1,2))
     return (;field_new, H, W, fftdims, ramp_before, ramp_after)
 end
 
 
+function shifted_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, ramp_before, ramp_after) = _prepare_shifted_angular_spectrum(field, z, λ, L, real(CT).(α); padding, 
+                                              pad_factor, bandlimit, bandlimit_border)
+
+	# propagate field
+    field_new_is = ifftshift(field_new, fftdims) ./ (ramp_before)
+    field_out = fftshift(ramp_after .* ifft(fft(field_new_is, fftdims) .* H .* W, fftdims), fftdims)
+    field_out_cropped = padding ? crop_center(field_out, size(field)) : field_out
+    shift = z .* tan.(α) ./ L
+	# return final field and some other variables
+    return field_out_cropped, (; L, shift)
+end
+
+
+ # highly optimized version with pre-planning
+struct ShiftedAngularSpectrum{A, T, T2, P, R}
+    HW::A
+    buffer::A
+    buffer2::A
+    L::T
+    shift::T2
+    p::P
+    padding::Bool
+    pad_factor::Int
+    ramp_before::R
+    ramp_after::R
+end
+
+
 """
-    shifted_angular_spectrum(field, z, λ, L, α; kwargs...)
+    ShiftedAngularSpectrum(field, z, λ, L, α; kwargs...)
 
-Returns the electrical field with physical length `L` and wavelength `λ` propagated with the shifted angular spectrum 
+Returns a method to propagate the electrical field with physical length `L` and wavelength `λ` with the shifted angular spectrum 
 method of plane waves (AS) by the propagation distance `z`.
-`α` is the shift angle with respect to the optical axis.
+`α` should be a tuple containing the offset angles with respect to the optical axis.
 
-This method is efficient but to avoid recalculating some arrays (such as the phase kernel), see [`ShiftedAngularSpectrum`](@ref). 
 
 # Arguments
 * `field`: Input field
 * `z`: propagation distance. Can be a single number or a vector of `z`s (Or `CuVector`). In this case the returning array has one dimension more.
 * `λ`: wavelength of field
 * `L`: field size (can be a scalar or a tuple) indicating field size
-* `α` is the shift angle with respect to the optical axis.
+* `α` is the tuple of shift angles with respect to the optical axis.
 
 
 # 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. 
-
+* `extract_ramp=true`: divides the field by phase ramp `exp.(1im * 2π / λ * (sin(α[2]) .* x .+ sin(α[1]) .* y))` and multiplies after
+                        propagation the ramp (with new real space coordinates) back to the field
 
 # Examples
 ```jldoctest
 julia> field = zeros(ComplexF32, (4,4)); field[3,3] = 1
-```
 
+julia> AS, _ = ShiftedAngularSpectrum(field, 100e-6, 633e-9, 100e-6, (deg2rad(10), 0));
+
+julia> AS(field)
+(ComplexF32[1.5269792f-5 + 1.7594219f-5im -3.996831f-5 - 7.624799f-5im -0.0047351345f0 + 0.002100923f0im -3.996831f-5 - 7.624799f-5im; -8.294997f-5 - 1.8230454f-5im 0.00028230582f0 + 8.1745195f-5im 0.0051693693f0 - 0.016958509f0im 0.00028230582f0 + 8.1745195f-5im; 0.0029884572f0 + 0.0040671355f0im -0.009686601f0 - 0.014245203f0im -0.82990384f0 + 0.5566719f0im -0.009686601f0 - 0.014245203f0im; -1.4191573f-7 + 9.41665f-5im 2.111472f-5 - 0.00031620878f0im -0.017670793f0 - 0.0014212304f0im 2.111472f-5 - 0.00031620878f0im], (L = 0.0001, shift = (0.17632698070846498, 0.0)))
+```
 
 # References
 * Matsushima, Kyoji. "Shifted angular spectrum method for off-axis numerical propagation." Optics Express 18.17 (2010): 18453-18463.
 """
-function shifted_angular_spectrum(field::AbstractArray{CT, 2}, z, λ, L, α; 
+function ShiftedAngularSpectrum(field::AbstractArray{CT, N}, z::Number, λ, L, α; 
+                          extract_ramp=true,
                           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, ramp_before, ramp_after) = 
+                          bandlimit_border=(0.8, 1)) where {CT, N}
+    
+        (; field_new, H, W, fftdims, ramp_before, ramp_after) = 
             _prepare_shifted_angular_spectrum(field, z, λ, L, real(CT).(α); padding, 
                                               pad_factor, bandlimit, bandlimit_border)
+    
+      
+        buffer2 = similar(field, complex(eltype(field_new)), (size(field_new, 1), size(field_new, 2)))
+        buffer = copy(buffer2)
+        
+        p = plan_fft!(buffer, (1, 2))
+        H .= H .* W
+        HW = H
+        shift = z .* tan.(α) ./ L
+        if !extract_ramp 
+            ramp_before = nothing
+            ramp_after = nothing
+        end
 
-	# propagate field
-    field_new_is = ifftshift(field_new, fftdims) ./ (ramp_before)
-#    ramp_after = 1
-    field_out = fftshift(ramp_after .* ifft(fft(field_new_is, fftdims) .* H .* W, fftdims), fftdims)
-    field_out_cropped = padding ? crop_center(field_out, size(field)) : field_out
-    shift = z .* tan.(α) ./ L
-	# return final field and some other variables
-    return field_out_cropped, (; L, shift)
+
+        return ShiftedAngularSpectrum{typeof(H), typeof(L), typeof(shift), typeof(p), typeof(ramp_before)}(HW, 
+                    buffer, buffer2, L, shift, p, padding, pad_factor, ramp_before, ramp_after), (;L, shift)
+    end
+
+
+
+"""
+    (shifted_as::ShiftedAngularSpectrum)(field)
+
+Uses the struct to efficiently store some pre-calculated objects.
+Propagate the field.
+"""
+function (as::ShiftedAngularSpectrum{A, T, T2, P, R})(field) where {A,T,T2,P,R}
+    fill!(as.buffer2, 0)
+    field_new = set_center!(as.buffer2, field, broadcast=true)
+    ifftshift!(as.buffer, field_new, (1, 2))
+    if !(R === Nothing) 
+        as.buffer ./= as.ramp_before
+    end
+    field_imd = as.p * as.buffer
+    field_imd .*= as.HW
+    field_imd = inv(as.p) * field_imd
+    if !(R === Nothing) 
+        field_imd .*= as.ramp_after
+    end
+    field_out = fftshift!(as.buffer2, 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, as.shift)
 end
 
 
-scalable_angular_spectrum(field::AbstractArray, z, λ, L; kwargs...) = throw(ArgumentError("Provided field needs to have a complex elementype"))
+function ChainRulesCore.rrule(as::ShiftedAngularSpectrum{A, T, T2, P, R}, field) where {A,T,T2,P,R}
+    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 
+        ifftshift!(as.buffer, field_new, (1, 2))
+        if !(R === Nothing) 
+            as.buffer .*= conj.(as.ramp_after)
+        end
+        field_imd = as.p * as.buffer 
+        field_imd .*= conj.(as.HW)
+        
+        field_imd = inv(as.p) * field_imd 
+        if !(R === Nothing) 
+            field_imd ./= conj.(as.ramp_before)
+        end
+        field_out = fftshift!(as.buffer2, 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

test/shifted_angular_spectrum.jl

@@ -8,23 +8,76 @@
     y = fftpos(L, N, CenterFT)
     field = box(Float32, sz, (20,20)) .* exp.(1im .* 2f0 * π ./ λ .* y .* sin(α));
     res_AS = AngularSpectrum(field, z, λ, L)[1](field)[1];
-    res = shifted_angular_spectrum(field, z, λ, L, (α , 0), bandlimit=true)
+    res = WaveOpticsPropagation.shifted_angular_spectrum(field, z, λ, L, (α , 0), bandlimit=true)
     shift = z * tan(α) / L * N
     shift2 = z * tan(α) / L
+    res2 = ShiftedAngularSpectrum(field, z, λ, L, (α, 0), bandlimit=true)[1](field)
     @test all(.≈(1 .+ FourierTools.shift(res[1], (shift, 0))[round(Int, shift)+1:end, :], 1 .+ res_AS[round(Int, shift)+1:end, :], rtol=5f-2))
+    @test all(.≈(1 .+ FourierTools.shift(res2[1], (shift, 0))[round(Int, shift)+1:end, :], 1 .+ res_AS[round(Int, shift)+1:end, :], rtol=5f-2))
 
 
     field = box(Float32, sz, (20,20)) .* exp.(1im .* 2f0 * π ./ λ .* y' .* sin(α));
     res_AS = AngularSpectrum(field, z, λ, L)[1](field)[1];
-    res = shifted_angular_spectrum(field, z, λ, L, (0, α), bandlimit=true)
+    res = WaveOpticsPropagation.shifted_angular_spectrum(field, z, λ, L, (0, α), bandlimit=true)
     shift = z * tan(α) / L * N
     shift2 = z * tan(α) / L
+    res2 = ShiftedAngularSpectrum(field, z, λ, L, (0, α), bandlimit=true)[1](field)
     @test all(.≈(1 .+ FourierTools.shift(res[1], (0, shift))[:, round(Int, shift)+1:end], 1 .+ res_AS[:, round(Int, shift)+1:end], rtol=5f-2))
+    @test all(.≈(1 .+ FourierTools.shift(res2[1], (0, shift))[:, round(Int, shift)+1:end], 1 .+ res_AS[:, round(Int, shift)+1:end], rtol=5f-2))
 
 
     res_AS = AngularSpectrum(field, z, λ, L)[1](field)[1];
-    res = shifted_angular_spectrum(field, z, λ, L, (0, 0), bandlimit=true)[1]
+    res = WaveOpticsPropagation.shifted_angular_spectrum(field, z, λ, L, (0, 0), bandlimit=true)[1]
+    res2 = ShiftedAngularSpectrum(field, z, λ, L, (0, 0), bandlimit=true)[1](field)[1]
     @test all(.≈(1 .+ res, 1 .+ res_AS, rtol=1f-1)) 
     @test ≈(1 .+ res, 1 .+ res_AS, rtol=1f-2) 
+    @test all(.≈(1 .+ res2, 1 .+ res_AS, rtol=1f-1)) 
+    @test ≈(1 .+ res2, 1 .+ res_AS, rtol=1f-2) 
+
+
+
+
+    @testset "Test gradient with Finite Differences" begin
+        field = zeros(ComplexF64, (24, 24))
+        field[14:16, 14:16] .= 1
+
+        α = deg2rad(10)
+        gg(x) = sum(abs2.(x .- WaveOpticsPropagation.shifted_angular_spectrum(cis.(x), 100e-6, 633e-9, 100e-6, (α, 0))[1]))
+
+        out2 = FiniteDifferences.grad(central_fdm(5, 1), gg, field)[1]
+
+        out1 = gradient(gg, field)[1]
+        @test out1 .+ cis(1) ≈ out2  .+ cis(1)
+        AS, _ = ShiftedAngularSpectrum(field, 100e-6, 633e-9, 100e-6, (α, 0))
+
+        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 .- WaveOpticsPropagation.shifted_angular_spectrum(cis.(x), 100e-6, 633e-9, 100e-6, (α, 0))[1]))
+        out2 = FiniteDifferences.grad(central_fdm(5, 1), gg, field)[1]
+        out1 = gradient(gg, field)[1]
+        @test out1 .+ cis(1) ≈ out2  .+ cis(1)
+        AS, _ = ShiftedAngularSpectrum(field, 100e-6, 633e-9, 100e-6, (α, 0))
+        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
+        AS, _ = ShiftedAngularSpectrum(field, 100e-6, 633e-9, 100e-6, (α, 0); extract_ramp=false)
+        f_AS(x) = sum(abs2.(x .- AS(cis.(x))[1]))
+        out2 = FiniteDifferences.grad(central_fdm(5, 1), f_AS, field)[1]
+        out3 = gradient(f_AS, field)[1]
+        @test out3 ≈ out2
+
+
+    end
 
 end