Split tests to improve modularity and debloat runtests.jl (#46)

* Split tests to improve modularity and debloat runtests.jl

* Format

* Add imports

* Split Brusselator definition

* Fix definition

* Remove RecursiveSet

test/benchmark.jl

@@ -4,7 +4,7 @@ using SparseConnectivityTracer: SortedVector
 using Symbolics: Symbolics
 using NNlib: conv
 
-include("brusselator.jl")
+include("brusselator_definition.jl")
 
 function benchmark_brusselator(N::Integer, method=:tracer)
     dims = (N, N, 2)

test/brusselator.jl

@@ -1,20 +1,29 @@
-# Brusselator example taken from Gowda et al.
-# "Sparsity Programming: Automated Sparsity-Aware Optimizations in Differentiable Programming"
-# https://openreview.net/pdf?id=rJlPdcY38B
+using ReferenceTests
+using SparseConnectivityTracer
+using Test
 
-#! format: off
-brusselator_f(x, y, t) = ifelse((((x-0.3)^2 + (y-0.6)^2) <= 0.1^2) && (t >= 1.1), 5., 0.)
-limit(a, N) = a == N+1 ? 1 : a == 0 ? N : a
-function brusselator_2d_loop(du, u, p, t)
-    A, B, alpha, xyd, dx, N = p; alpha = alpha/dx^2
-    @inbounds for I in CartesianIndices((N, N))
-        i, j = Tuple(I)
-        x, y = xyd[I[1]], xyd[I[2]]
-        ip1, im1, jp1, jm1 = limit(i+1, N), limit(i-1, N), limit(j+1, N), limit(j-1, N)
-        du[i,j,1] = alpha*(u[im1,j,1] + u[ip1,j,1] + u[i,jp1,1] + u[i,jm1,1] - 4u[i,j,1]) +
-        B + u[i,j,1]^2*u[i,j,2] - (A + 1)*u[i,j,1] + brusselator_f(x, y, t)
-        du[i,j,2] = alpha*(u[im1,j,2] + u[ip1,j,2] + u[i,jp1,2] + u[i,jm1,2] - 4u[i,j,2]) +
-        A*u[i,j,1] - u[i,j,1]^2*u[i,j,2]
-    end
+include("brusselator_definition.jl")
+
+@testset "Set type $S" for S in (BitSet, Set{UInt64}, SortedVector{UInt64})
+    N = 6
+    dims = (N, N, 2)
+    A = 1.0
+    B = 1.0
+    alpha = 1.0
+    xyd = fill(1.0, N)
+    dx = 1.0
+    p = (A, B, alpha, xyd, dx, N)
+
+    u = rand(dims...)
+    du = similar(u)
+    f!(du, u) = brusselator_2d_loop(du, u, p, nothing)
+
+    C = connectivity_pattern(f!, du, u, S)
+    @test_reference "references/pattern/connectivity/Brusselator.txt" BitMatrix(C)
+    J = jacobian_pattern(f!, du, u, S)
+    @test_reference "references/pattern/jacobian/Brusselator.txt" BitMatrix(J)
+    @test C == J
+
+    C_ref = Symbolics.jacobian_sparsity(f!, du, u)
+    @test C == C_ref
 end
-#! format: on

test/brusselator_definition.jl

@@ -0,0 +1,20 @@
+# Brusselator example taken from Gowda et al.
+# "Sparsity Programming: Automated Sparsity-Aware Optimizations in Differentiable Programming"
+# https://openreview.net/pdf?id=rJlPdcY38B
+
+#! format: off
+brusselator_f(x, y, t) = ifelse((((x-0.3)^2 + (y-0.6)^2) <= 0.1^2) && (t >= 1.1), 5., 0.)
+limit(a, N) = a == N+1 ? 1 : a == 0 ? N : a
+function brusselator_2d_loop(du, u, p, t)
+    A, B, alpha, xyd, dx, N = p; alpha = alpha/dx^2
+    @inbounds for I in CartesianIndices((N, N))
+        i, j = Tuple(I)
+        x, y = xyd[I[1]], xyd[I[2]]
+        ip1, im1, jp1, jm1 = limit(i+1, N), limit(i-1, N), limit(j+1, N), limit(j-1, N)
+        du[i,j,1] = alpha*(u[im1,j,1] + u[ip1,j,1] + u[i,jp1,1] + u[i,jm1,1] - 4u[i,j,1]) +
+        B + u[i,j,1]^2*u[i,j,2] - (A + 1)*u[i,j,1] + brusselator_f(x, y, t)
+        du[i,j,2] = alpha*(u[im1,j,2] + u[ip1,j,2] + u[i,jp1,2] + u[i,jm1,2] - 4u[i,j,2]) +
+        A*u[i,j,1] - u[i,j,1]^2*u[i,j,2]
+    end
+end
+#! format: on

test/first_order.jl

@@ -0,0 +1,61 @@
+using ReferenceTests
+using SparseConnectivityTracer
+using SparseConnectivityTracer: tracer, trace_input, inputs, empty
+using SparseConnectivityTracer: SortedVector
+using Test
+
+@testset "Set type $S" for S in (BitSet, Set{UInt64}, SortedVector{UInt64})
+    CT = ConnectivityTracer{S}
+    JT = JacobianTracer{S}
+
+    x = rand(3)
+    xt = trace_input(CT, x)
+
+    # Matrix multiplication
+    A = rand(1, 3)
+    yt = only(A * xt)
+    @test connectivity_pattern(x -> only(A * x), x, S) ≈ [1 1 1]
+
+    # Custom functions
+    f(x) = [x[1]^2, 2 * x[1] * x[2]^2, sin(x[3])]
+    yt = f(xt)
+
+    @test connectivity_pattern(f, x, S) ≈ [1 0 0; 1 1 0; 0 0 1]
+    @test jacobian_pattern(f, x, S) ≈ [1 0 0; 1 1 0; 0 0 1]
+
+    @test connectivity_pattern(identity, rand(), S) ≈ [1;;]
+    @test jacobian_pattern(identity, rand(), S) ≈ [1;;]
+    @test connectivity_pattern(Returns(1), 1, S) ≈ [0;;]
+    @test jacobian_pattern(Returns(1), 1, S) ≈ [0;;]
+
+    # Test JacobianTracer on functions with zero derivatives
+    x = rand(2)
+    g(x) = [x[1] * x[2], ceil(x[1] * x[2]), x[1] * round(x[2])]
+    @test connectivity_pattern(g, x, S) ≈ [1 1; 1 1; 1 1]
+    @test jacobian_pattern(g, x, S) ≈ [1 1; 0 0; 1 0]
+
+    # Code coverage
+    @test connectivity_pattern(x -> [sincos(x)...], 1, S) ≈ [1; 1]
+    @test connectivity_pattern(typemax, 1, S) ≈ [0;;]
+    @test connectivity_pattern(x -> x^(2//3), 1, S) ≈ [1;;]
+    @test connectivity_pattern(x -> (2//3)^x, 1, S) ≈ [1;;]
+    @test connectivity_pattern(x -> x^ℯ, 1, S) ≈ [1;;]
+    @test connectivity_pattern(x -> ℯ^x, 1, S) ≈ [1;;]
+    @test connectivity_pattern(x -> round(x, RoundNearestTiesUp), 1, S) ≈ [1;;]
+
+    @test jacobian_pattern(x -> [sincos(x)...], 1, S) ≈ [1; 1]
+    @test jacobian_pattern(typemax, 1, S) ≈ [0;;]
+    @test jacobian_pattern(x -> x^(2//3), 1, S) ≈ [1;;]
+    @test jacobian_pattern(x -> (2//3)^x, 1, S) ≈ [1;;]
+    @test jacobian_pattern(x -> x^ℯ, 1, S) ≈ [1;;]
+    @test jacobian_pattern(x -> ℯ^x, 1, S) ≈ [1;;]
+    @test jacobian_pattern(x -> round(x, RoundNearestTiesUp), 1, S) ≈ [0;;]
+
+    # Base.show
+    @test_reference "references/show/ConnectivityTracer_$S.txt" repr(
+        "text/plain", tracer(CT, 2)
+    )
+    @test_reference "references/show/JacobianTracer_$S.txt" repr(
+        "text/plain", tracer(JT, 2)
+    )
+end

test/nnlib.jl

@@ -0,0 +1,14 @@
+using NNlib
+using ReferenceTests
+using SparseConnectivityTracer
+using Test
+
+@testset "Set type $S" for S in (BitSet, Set{UInt64}, SortedVector{UInt64})
+    x = rand(3, 3, 2, 1) # WHCN
+    w = rand(2, 2, 2, 1) # Conv((2, 2), 2 => 1)
+    C = jacobian_pattern(x -> NNlib.conv(x, w), x, S)
+    @test_reference "references/pattern/connectivity/NNlib/conv.txt" BitMatrix(C)
+    J = jacobian_pattern(x -> NNlib.conv(x, w), x, S)
+    @test_reference "references/pattern/jacobian/NNlib/conv.txt" BitMatrix(J)
+    @test C == J
+end