unwrap: Initialize image reliability in deterministic order (#617)

* unwrap: Initialize image reliability in deterministic order

* correct some typos

* Remove disused constructor

* add tests for reproducibility

for invalid input; valid input should be resistant to random noise

* unwrap test cleanup

* test with threads

---------

Co-authored-by: wheeheee <104880306+wheeheee@users.noreply.github.com>

Author: RomanHargrave

.github/workflows/CI.yml

@@ -14,6 +14,9 @@ jobs:
       actions: write
       contents: read
 
+    env:
+      JULIA_NUM_THREADS: 'auto'
+
     strategy:
       fail-fast: false
       matrix:

Project.toml

@@ -30,14 +30,16 @@ Polynomials = "4"
 Random = "1.6"
 Reexport = "1.0"
 SpecialFunctions = "2.0"
+StableRNGs = "1"
 Statistics = "1.6"
 Test = "1.6"
 julia = "1.10"
 
 [extras]
 DelimitedFiles = "8bb1440f-4735-579b-a4ab-409b98df4dab"
 OffsetArrays = "6fe1bfb0-de20-5000-8ca7-80f57d26f881"
+StableRNGs = "860ef19b-820b-49d6-a774-d7a799459cd3"
 Test = "8dfed614-e22c-5e08-85e1-65c5234f0b40"
 
 [targets]
-test = ["DelimitedFiles", "OffsetArrays", "Test"]
+test = ["DelimitedFiles", "OffsetArrays", "StableRNGs", "Test"]

src/unwrap.jl

@@ -46,11 +46,11 @@ wrapping discontinuities across all `ndims(m)` dimensions.
 
 A common usage for unwrapping across a singleton dimension is for a phase
 measurement over time, such as when
-comparing successive frames of a short-time-fourier-transform, as
+comparing successive frames of a short-time Fourier transform, as
 each frame is wrapped to stay within (-pi, pi].
 
 A common usage for unwrapping across multiple dimensions is for a phase
-measurement of a scene, such as when retrieving the phase information of
+measurement of a scene, such as when retrieving the phase information
 of an image, as each pixel is wrapped to stay within (-pi, pi].
 
 # Arguments
@@ -94,7 +94,7 @@ mutable struct Pixel{T}
         return pixel
     end
 end
-Pixel(v, rng) = Pixel{typeof(v)}(0, v, rand(rng), 1)
+Pixel(v, rel::Float64) = Pixel{typeof(v)}(0, v, rel, 1)
 @inline Base.length(p::Pixel) = p.head.groupsize
 
 struct Edge{T}
@@ -146,8 +146,13 @@ end
 # function to broadcast
 function init_pixels(wrapped_image::AbstractArray{T, N}, rng) where {T, N}
     pixel_image = similar(wrapped_image, Pixel{T})
-    Threads.@threads for i in eachindex(wrapped_image, pixel_image)
-        pixel_image[i] = Pixel(wrapped_image[i], rng)
+
+    # Initialize reliability values before going parallel. This ensures that
+    # reliability values are generated in a deterministic order.
+    rels = rand(rng, Float64, size(wrapped_image))
+
+    Threads.@threads for i in eachindex(wrapped_image, pixel_image, rels)
+        pixel_image[i] = Pixel(wrapped_image[i], rels[i])
     end
     return pixel_image
 end
@@ -248,8 +253,8 @@ function populate_edges!(edges::Vector{Edge{T}}, pixel_image::AbstractArray{Pixe
 end
 
 function calculate_reliability(pixel_image::AbstractArray{T, N}, circular_dims, range) where {T, N}
-    # get the shifted pixel indices in CartesinanIndex form
-    # This gets all the nearest neighbors (CartesionIndex{N}() = one(CartesianIndex{N}))
+    # get the shifted pixel indices in CartesianIndex form
+    # This gets all the nearest neighbors
     one_cart = oneunit(CartesianIndex{N})
     pixel_shifts = -one_cart:one_cart
     image_inds = CartesianIndices(pixel_image)

test/unwrap.jl

@@ -1,28 +1,32 @@
 using DSP, Test
-using Random: MersenneTwister
+using Random: MersenneTwister, default_rng, seed!
+using StableRNGs
+using Statistics: mean
 
 @testset "Unwrap 1D" begin
-    @test unwrap([0.1, 0.2, 0.3, 0.4]) ≈ [0.1, 0.2, 0.3, 0.4]
-    @test unwrap([0.1, 0.2 + 2pi, 0.3, 0.4]) ≈ [0.1, 0.2, 0.3, 0.4]
-    @test unwrap([0.1, 0.2 - 2pi, 0.3, 0.4]) ≈ [0.1, 0.2, 0.3, 0.4]
-    @test unwrap([0.1, 0.2 - 2pi, 0.3 - 2pi, 0.4]) ≈ [0.1, 0.2, 0.3, 0.4]
-    @test unwrap([0.1 + 2pi, 0.2, 0.3, 0.4]) ≈ [0.1 + 2pi, 0.2 + 2pi, 0.3 + 2pi, 0.4 + 2pi]
-    @test unwrap([0.1, 0.2 + 6pi, 0.3, 0.4]) ≈ [0.1, 0.2, 0.3, 0.4]
+    unwrapped = [0.1, 0.2, 0.3, 0.4]
+    @test unwrap([0.1, 0.2, 0.3, 0.4]) ≈ unwrapped
+    @test unwrap([0.1, 0.2 + 2pi, 0.3, 0.4]) ≈ unwrapped
+    @test unwrap([0.1, 0.2 - 2pi, 0.3, 0.4]) ≈ unwrapped
+    @test unwrap([0.1, 0.2 - 2pi, 0.3 - 2pi, 0.4]) ≈ unwrapped
+    @test unwrap([0.1 + 2pi, 0.2, 0.3, 0.4]) ≈ unwrapped .+ 2pi
+    @test unwrap([0.1, 0.2 + 6pi, 0.3, 0.4]) ≈ unwrapped
 
     test_v = [0.1, 0.2, 0.3 + 2pi, 0.4]
     res_v = unwrap(test_v)
-    @test test_v ≈ [0.1, 0.2, 0.3 + 2pi, 0.4]
+    @test res_v ≈ unwrapped
+    @test test_v == [0.1, 0.2, 0.3 + 2pi, 0.4]
+
     res_v .= 0
     unwrap!(res_v, test_v)
-    @test res_v ≈ [0.1, 0.2, 0.3, 0.4]
-    @test test_v ≈ [0.1, 0.2, 0.3 + 2pi, 0.4]
-    unwrap!(test_v)
-    @test test_v ≈ [0.1, 0.2, 0.3, 0.4]
+    @test res_v ≈ unwrapped
+    @test test_v == [0.1, 0.2, 0.3 + 2pi, 0.4]
+
+    @test unwrap!(test_v) === test_v
+    @test test_v ≈ unwrapped
 
     # test unwrapping within multi-dimensional array
-    wrapped = [0.1, 0.2 + 2pi, 0.3, 0.4]
-    unwrapped = [0.1, 0.2, 0.3, 0.4]
-    wrapped = hcat(wrapped, wrapped)
+    wrapped = repeat([0.1, 0.2 + 2pi, 0.3, 0.4], 1, 2)
     unwrapped = hcat(unwrapped, unwrapped)
     @test unwrap(wrapped, dims=2) ≈ wrapped
     @test unwrap(wrapped, dims=1) ≈ unwrapped
@@ -101,39 +105,46 @@ end
 
 @testset "Unwrap 3D" begin
     types = (Float32, Float64, BigFloat)
-    f(x, y, z) = 0.1x^2 - 2y + 2z
+
+    f_nowrap(x, y, z) = 0.1x^2 - 2y + 2z
     f_wraparound2(x, y, z) = 5*sin(x) + 2*cos(y) + z
     f_wraparound3(x, y, z) = 5*sin(x) + 2*cos(y) - 4*cos(z)
+
     for T in types
         grid = range(zero(T), stop=2convert(T, π), length=11)
-        f_uw = f.(grid, grid', reshape(grid, 1, 1, :))
-        f_wr = f_uw .% (2convert(T, π))
-        uw_test = unwrap(f_wr, dims=1:3)
-        offset = first(f_uw) - first(uw_test)
-        @test (uw_test.+offset) ≈ f_uw rtol=eps(T) #oop, nowrap
-        # test in-place version
-        unwrap!(f_wr, dims=1:3)
-        offset = first(f_uw) - first(f_wr)
-        @test (f_wr.+offset) ≈ f_uw rtol=eps(T) #ip, nowrap
-
-        f_uw = f_wraparound2.(grid, grid', reshape(grid, 1, 1, :))
-        f_wr = f_uw .% (2convert(T, π))
-        uw_test = unwrap(f_wr, dims=1:3)
-        offset = first(f_uw) - first(uw_test)
-        @test (uw_test.+offset) ≈ f_uw #oop, 2wrap
-        # test in-place version
-        unwrap!(f_wr, dims=1:3, circular_dims=(true, true, false))
-        offset = first(f_uw) - first(f_wr)
-        @test (f_wr.+offset) ≈ f_uw #ip, 2wrap
-
-        f_uw = f_wraparound3.(grid, grid', reshape(grid, 1, 1, :))
-        f_wr = f_uw .% (2convert(T, π))
-        uw_test = unwrap(f_wr, dims=1:3, circular_dims=(true, true, true))
-        offset = first(f_uw) - first(uw_test)
-        @test (uw_test.+offset) ≈ f_uw #oop, 3wrap
-        # test in-place version
-        unwrap!(f_wr, dims=1:3, circular_dims=(true, true, true))
-        offset = first(f_uw) - first(f_wr)
-        @test (f_wr.+offset) ≈ f_uw #oop, 3wrap
+
+        for (func, circular_dims) in (
+            (f_nowrap,      (false, false, false)),
+            (f_wraparound2, (true, true, false)),
+            (f_wraparound3, (true, true, true))
+        )
+            f_uw = func.(grid, grid', reshape(grid, 1, 1, :))
+            f_wr = f_uw .% (2convert(T, π))
+            uw_test = unwrap(f_wr; dims=1:3, circular_dims)
+            offset = first(f_uw) - first(uw_test)
+            @test (uw_test .+ offset) ≈ f_uw rtol = eps(T)  # oop
+
+            # test in-place version
+            unwrap!(f_wr; dims=1:3, circular_dims)
+            offset = first(f_uw) - first(f_wr)
+            @test (f_wr .+ offset) ≈ f_uw rtol = eps(T)     # ip
+        end
+    end
+end
+
+@testset "reproducible unwrap" begin
+    function measure(s, rng, seed=rand(UInt))
+        u1 = unwrap(s; dims=1:ndims(s), rng=seed!(rng, seed))
+        u2 = unwrap(s; dims=1:ndims(s), rng=seed!(rng, seed))
+        return sqrt(mean(abs2, u1 - u2))
+    end
+
+    x1 = 1234 * rand(10_000)
+    x2 = 1357 * rand(200, 200)
+    x3 = 1248 * rand(30, 30, 30)
+
+    for x in (x1, x2, x3)
+        @test measure(x, default_rng()) == measure(x, MersenneTwister()) ==
+              measure(x, StableRNG(10)) == 0.0
     end
 end