feat: allow naming scenarios and test sets (#712)

* feat: allow naming scenarios and test sets

* docstrings

Author: gdalle

DifferentiationInterfaceTest/Project.toml

@@ -1,7 +1,7 @@
 name = "DifferentiationInterfaceTest"
 uuid = "a82114a7-5aa3-49a8-9643-716bb13727a3"
 authors = ["Guillaume Dalle", "Adrian Hill"]
-version = "0.9.3"
+version = "0.9.4"
 
 [deps]
 ADTypes = "47edcb42-4c32-4615-8424-f2b9edc5f35b"

DifferentiationInterfaceTest/src/scenarios/modify.jl

@@ -15,6 +15,7 @@ function Base.zero(scen::Scenario{op,pl_op,pl_fun}) where {op,pl_op,pl_fun}
         res1=myzero(scen.res1),
         res2=myzero(scen.res2),
         smaller=isnothing(scen.smaller) ? nothing : zero(scen.smaller),
+        name=isnothing(scen.name) ? nothing : scen.name * " [zero]",
     )
 end
 
@@ -39,6 +40,7 @@ function change_function(
         else
             change_function(scen.smaller, new_f; keep_smaller=false)
         end,
+        name=isnothing(scen.name) ? nothing : scen.name * " [new function]",
     )
 end
 
@@ -63,6 +65,7 @@ function batchify(scen::Scenario{op,pl_op,pl_fun}) where {op,pl_op,pl_fun}
             res1=new_res1,
             res2,
             smaller=isnothing(smaller) ? nothing : batchify(smaller),
+            name=isnothing(scen.name) ? nothing : scen.name * " [batchified]",
         )
     elseif op == :hvp
         new_tang = (only(tang), -only(tang))
@@ -76,6 +79,7 @@ function batchify(scen::Scenario{op,pl_op,pl_fun}) where {op,pl_op,pl_fun}
             res1,
             res2=new_res2,
             smaller=isnothing(smaller) ? nothing : batchify(smaller),
+            name=isnothing(scen.name) ? nothing : scen.name * " [batchified]",
         )
     end
 end
@@ -160,6 +164,7 @@ function constantify(scen::Scenario{op,pl_op,pl_fun}) where {op,pl_op,pl_fun}
         res1=mymultiply(scen.res1, a),
         res2=mymultiply(scen.res2, a),
         smaller=isnothing(scen.smaller) ? nothing : constantify(scen.smaller),
+        name=isnothing(scen.name) ? nothing : scen.name * " [constantified]",
     )
 end
 
@@ -213,6 +218,7 @@ function cachify(scen::Scenario{op,pl_op,pl_fun}) where {op,pl_op,pl_fun}
         res1=scen.res1,
         res2=scen.res2,
         smaller=isnothing(scen.smaller) ? nothing : cachify(scen.smaller),
+        name=isnothing(scen.name) ? nothing : scen.name * " [cachified]",
     )
 end
 

DifferentiationInterfaceTest/src/scenarios/scenario.jl

@@ -13,8 +13,8 @@ This generic type should never be used directly: use the specific constructor co
 
 # Constructors
 
-    Scenario{op,pl_op}(f, x; tang, contexts, res1, res2)
-    Scenario{op,pl_op}(f!, y, x; tang, contexts, res1, res2)
+    Scenario{op,pl_op}(f, x; tang, contexts, res1, res2, name)
+    Scenario{op,pl_op}(f!, y, x; tang, contexts, res1, res2, name)
 
 # Fields
 
@@ -37,32 +37,57 @@ struct Scenario{op,pl_op,pl_fun,F,X,Y,T<:Union{Nothing,NTuple},C<:Tuple,R1,R2,S}
     res2::R2
     "private field (not part of the public API) containing a variant of the scenario to test preparation resizing"
     smaller::S
+    "name of the scenario for display in test sets and dataframes"
+    name::Union{String,Nothing}
 end
 
 function Scenario{op,pl_op,pl_fun}(
-    f::F; x::X, y::Y, tang::T, contexts::C, res1::R1, res2::R2, smaller::S=nothing
+    f::F;
+    x::X,
+    y::Y,
+    tang::T,
+    contexts::C,
+    res1::R1,
+    res2::R2,
+    smaller::S=nothing,
+    name=nothing,
 ) where {op,pl_op,pl_fun,F,X,Y,T,C,R1,R2,S<:Union{Nothing,Scenario}}
     @assert smaller isa Union{Nothing,Scenario{op,pl_op,pl_fun,F,X,Y,T,C,R1,R2}}
     return Scenario{op,pl_op,pl_fun,F,X,Y,T,C,R1,R2,S}(
-        f, x, y, tang, contexts, res1, res2, smaller
+        f, x, y, tang, contexts, res1, res2, smaller, name
     )
 end
 
 function Scenario{op,pl_op}(
-    f, x; tang=nothing, contexts=(), res1=nothing, res2=nothing, smaller=nothing
+    f,
+    x;
+    tang=nothing,
+    contexts=(),
+    res1=nothing,
+    res2=nothing,
+    smaller=nothing,
+    name=nothing,
 ) where {op,pl_op}
     @assert op in ALL_OPS
     @assert pl_op in (:in, :out)
     y = f(x, map(unwrap, contexts)...)
-    return Scenario{op,pl_op,:out}(f; x, y, tang, contexts, res1, res2, smaller)
+    return Scenario{op,pl_op,:out}(f; x, y, tang, contexts, res1, res2, smaller, name)
 end
 
 function Scenario{op,pl_op}(
-    f!, y, x; tang=nothing, contexts=(), res1=nothing, res2=nothing, smaller=nothing
+    f!,
+    y,
+    x;
+    tang=nothing,
+    contexts=(),
+    res1=nothing,
+    res2=nothing,
+    smaller=nothing,
+    name=nothing,
 ) where {op,pl_op}
     @assert op in ALL_OPS
     @assert pl_op in (:in, :out)
-    return Scenario{op,pl_op,:in}(f!; x, y, tang, contexts, res1, res2, smaller)
+    return Scenario{op,pl_op,:in}(f!; x, y, tang, contexts, res1, res2, smaller, name)
 end
 
 Base.:(==)(scen1::Scenario, scen2::Scenario) = false
@@ -85,7 +110,8 @@ function Base.:(==)(
     )
     eq_res1 = scen1.res1 == scen2.res1
     eq_res2 = scen1.res2 == scen2.res2
-    return (eq_x && eq_y && eq_tang && eq_contexts && eq_res1 && eq_res2)
+    eq_name = scen1.name == scen2.name
+    return (eq_x && eq_y && eq_tang && eq_contexts && eq_res1 && eq_res2 && eq_name)
 end
 
 operator(::Scenario{op}) where {op} = op
@@ -123,12 +149,16 @@ end
 function Base.show(
     io::IO, scen::Scenario{op,pl_op,pl_fun,F,X,Y,T}
 ) where {op,pl_op,pl_fun,F,X,Y,T}
-    print(io, "Scenario{$(repr(op)),$(repr(pl_op))} $(string(scen.f)) : $X -> $Y")
-    if op in (:pushforward, :pullback, :hvp)
-        print(io, " ($(length(scen.tang)) tangents)")
-    end
-    if length(scen.contexts) > 0
-        print(io, " ($(length(scen.contexts)) contexts)")
+    if isnothing(scen.name)
+        print(io, "Scenario{$(repr(op)),$(repr(pl_op))} $(string(scen.f)) : $X -> $Y")
+        if op in (:pushforward, :pullback, :hvp)
+            print(io, " ($(length(scen.tang)) tangents)")
+        end
+        if length(scen.contexts) > 0
+            print(io, " ($(length(scen.contexts)) contexts)")
+        end
+    else
+        print(io, scen.name)
     end
     return nothing
 end

DifferentiationInterfaceTest/src/test_differentiation.jl

@@ -17,6 +17,8 @@ This function always creates and runs a `@testset`, though its contents may vary
 
 # Keyword arguments
 
+- `testset_name=nothing`: how to display the test set
+
 **Test categories:**
 
 - `correctness=true`: whether to compare the differentiation results with the theoretical values specified in each scenario
@@ -66,6 +68,7 @@ Each setting tests/benchmarks a different subset of calls:
 function test_differentiation(
     backends::Vector{<:AbstractADType},
     scenarios::Vector{<:Scenario}=default_scenarios();
+    testset_name::Union{String,Nothing}=nothing,
     # test categories
     correctness::Bool=true,
     type_stability::Symbol=:none,
@@ -105,11 +108,15 @@ function test_differentiation(
     scenarios = filter(s -> !(operator(s) in excluded), scenarios)
     scenarios = sort(scenarios; by=s -> (operator(s), string(s.f)))
 
-    title_additions =
-        (correctness ? " + correctness" : "") *
-        ((type_stability != :none) ? " + type stability" : "") *
-        ((benchmark != :none) ? " + benchmarks" : "")
-    title = "Testing" * title_additions[3:end]
+    if isnothing(testset_name)
+        title_additions =
+            (correctness ? " + correctness" : "") *
+            ((type_stability != :none) ? " + type stability" : "") *
+            ((benchmark != :none) ? " + benchmarks" : "")
+        title = "Testing" * title_additions[3:end]
+    else
+        title = testset_name
+    end
 
     benchmark_data = DifferentiationBenchmarkDataRow[]
 
@@ -222,6 +229,7 @@ Specifying the set of scenarios is mandatory for this function.
 function benchmark_differentiation(
     backends,
     scenarios::Vector{<:Scenario};
+    testset_name::Union{String,Nothing}=nothing,
     benchmark::Symbol=:prepared,
     excluded::Vector{Symbol}=Symbol[],
     logging::Bool=false,
@@ -235,6 +243,7 @@ function benchmark_differentiation(
     return test_differentiation(
         backends,
         scenarios;
+        testset_name,
         correctness=false,
         type_stability=:none,
         allocations=:none,

DifferentiationInterfaceTest/test/runtests.jl

@@ -14,6 +14,9 @@ GROUP = get(ENV, "JULIA_DIT_TEST_GROUP", "All")
         @testset verbose = true "Formalities" begin
             include("formalities.jl")
         end
+        @testset verbose = true "Scenarios" begin
+            include("scenario.jl")
+        end
     end
 
     if GROUP == "Zero" || GROUP == "All"

DifferentiationInterfaceTest/test/scenario.jl

@@ -0,0 +1,17 @@
+using DifferentiationInterface
+using DifferentiationInterfaceTest
+using ForwardDiff: ForwardDiff
+using Test
+
+scen = Scenario{:gradient,:out}(
+    sum, zeros(10); res1=ones(10), name="My pretty little scenario"
+)
+@test string(scen) == "My pretty little scenario"
+
+testset = test_differentiation(
+    AutoForwardDiff(), [scen]; testset_name="My amazing test set"
+)
+
+data = benchmark_differentiation(
+    AutoForwardDiff(), [scen]; testset_name="My amazing test set"
+)