Allow providing an AbstractRNG for seed in spring_layout

Author: simonschoelly

src/layout.jl

@@ -2,6 +2,7 @@ using SparseArrays: SparseMatrixCSC, sparse
 using ArnoldiMethod: SR
 using Base: OneTo
 using LinearAlgebra: eigen
+using Random: AbstractRNG, default_rng
 
 """
 Position nodes uniformly at random in the unit square.
@@ -83,6 +84,14 @@ where C is a parameter we can adjust
 *g*
 a graph
 
+*locs_x_in*
+x coordinates of the initial locations. If not provided they are sampled
+from [-1, 1]. Can be modified.
+
+*locs_y_in*
+y coordinates of the initial locations. If not provided they are sampled
+from [-1, 1]. Can be modified.
+
 *C*
 Constant to fiddle with density of resulting layout
 
@@ -93,7 +102,8 @@ Number of iterations we apply the forces
 Initial "temperature", controls movement per iteration
 
 *seed*
-Integer seed for pseudorandom generation of locations (default = 0).
+Either an `Integer` seed or an `Random.AbstractRNG` for generation of initial locations.
+If neither is provided `Random.default_rng()` is used.
 
 **Examples**
 ```
@@ -102,13 +112,20 @@ julia> locs_x, locs_y = spring_layout(g)
 ```
 """
 function spring_layout(g::AbstractGraph,
-                       locs_x_in::AbstractVector{R1}=2*rand(nv(g)).-1.0,
-                       locs_y_in::AbstractVector{R2}=2*rand(nv(g)).-1.0;
+                       locs_x_in::AbstractVector{R1},
+                       locs_y_in::AbstractVector{R2};
                        C=2.0,
                        MAXITER=100,
                        INITTEMP=2.0) where {R1 <: Real, R2 <: Real}
-
     nvg = nv(g)
+
+    if length(locs_x_in) != nvg
+        throw(ArgumentError("The length of locs_x_in does not equal the number of vertices"))
+    end
+    if length(locs_y_in) != nvg
+        throw(ArgumentError("The length of locs_y_in does not equal the number of vertices"))
+    end
+
     adj_matrix = adjacency_matrix(g)
 
     # The optimal distance bewteen vertices
@@ -180,7 +197,14 @@ using Random: MersenneTwister
 
 function spring_layout(g::AbstractGraph, seed::Integer; kws...)
     rng = MersenneTwister(seed)
-    spring_layout(g, 2 .* rand(rng, nv(g)) .- 1.0, 2 .* rand(rng,nv(g)) .- 1.0; kws...)
+    spring_layout(g, rng; kws...)
+end
+
+function spring_layout(g::AbstractGraph, rng::AbstractRNG=default_rng(); kws...)
+    nvg = nv(g)
+    locs_x_in = 2.0 * rand(rng, nvg) .- 1.0
+    locs_y_in = 2.0 * rand(rng, nvg) .- 1.0
+    spring_layout(g, locs_x_in, locs_y_in; kws...)
 end
 
 """

test/Project.toml

@@ -2,6 +2,7 @@
 Cairo = "159f3aea-2a34-519c-b102-8c37f9878175"
 ImageMagick = "6218d12a-5da1-5696-b52f-db25d2ecc6d1"
 Random = "9a3f8284-a2c9-5f02-9a11-845980a1fd5c"
+StableRNGs = "860ef19b-820b-49d6-a774-d7a799459cd3"
 Test = "8dfed614-e22c-5e08-85e1-65c5234f0b40"
 VisualRegressionTests = "34922c18-7c2a-561c-bac1-01e79b2c4c92"
 

test/runtests.jl

@@ -8,6 +8,7 @@ using Cairo
 using GraphPlot.Colors
 using GraphPlot.Compose
 using Random
+using StableRNGs: StableRNG
 using Test
 using VisualRegressionTests
 using ImageMagick
@@ -125,13 +126,34 @@ end
 
 @testset "Spring Layout" begin
     g1 = path_digraph(3)
-    x1, y1 = spring_layout(g1, 0; C = 1)
-    # TODO spring_layout uses random values which have changed on higher Julia versions
-    #   we should therefore use StableRNGs.jl for these layouts
-    @static if VERSION < v"1.7"
-        @test all(isapprox.(x1, [1.0, -0.014799825222963192, -1.0]))
-        @test all(isapprox.(y1, [-1.0, 0.014799825222963303, 1.0]))
-    end
+    g2 = smallgraph(:house)
+
+    # Neither seed nor initial locations provided
+    x1, y1 = spring_layout(g1; MAXITER=10)
+    @test length(x1) == nv(g1)
+    @test length(y1) == nv(g1)
+
+    # Using a seed
+    x2, y2 = spring_layout(g1, 0; C = 1)
+    @test length(x2) == nv(g1)
+    @test length(y2) == nv(g1)
+
+    # Using a rng
+    rng = StableRNG(123)
+    x3, y3 = spring_layout(g2, rng; INITTEMP = 7.5)
+    @test x3 ≈ [0.6417685918857294, -1.0, 1.0, -0.5032029640625139, 0.585415479582793]
+    @test y3 ≈ [-1.0, -0.7760280912987298, 0.06519424728464562, 0.2702599482349506, 1.0]
+
+    # Using initial locations
+    locs_x_in = 1:5
+    locs_y_in = [-1.0, 2.0, 0.3, 0.4, -0.5]
+    x4, y4 = spring_layout(g2, locs_x_in, locs_y_in)
+    @test x4 ≈ [-1.0, -0.4030585026962391, -0.050263101475789274, 0.5149349966578818, 1.0]
+    @test y4 ≈ [-0.03307638042475203, 1.0, -0.8197758901868164, 0.15834883764718155, -1.0]
+
+    # Providing initial locations with the wrong lengths should throw an ArgumentError
+    @test_throws ArgumentError("The length of locs_x_in does not equal the number of vertices") spring_layout(g1, 1:5, [1,2,3])
+    @test_throws ArgumentError("The length of locs_y_in does not equal the number of vertices") spring_layout(g2, 1:5, [1,2,3])
 end
 
 @testset "Circular Layout" begin