Fix bug in NoiseAugmentation constructor (#183)

Author: adrhill

Project.toml

@@ -1,7 +1,7 @@
 name = "ExplainableAI"
 uuid = "4f1bc3e1-d60d-4ed0-9367-9bdff9846d3b"
 authors = ["Adrian Hill <[email protected]>"]
-version = "0.10.0"
+version = "0.10.1-DEV"
 
 [deps]
 ADTypes = "47edcb42-4c32-4615-8424-f2b9edc5f35b"

README.md

@@ -69,10 +69,10 @@ heatmap(input, analyzer, 5)  # for heatmap
 
 | **Analyzer**                                  | **Heatmap for class "castle"** |**Heatmap for class "street sign"** |
 |:--------------------------------------------- |:------------------------------:|:----------------------------------:|
-| `InputTimesGradient`                          | ![][castle-ixg]                | ![][streetsign-ixg]                |
 | `Gradient`                                    | ![][castle-grad]               | ![][streetsign-grad]               |
 | `SmoothGrad`                                  | ![][castle-smoothgrad]         | ![][streetsign-smoothgrad]         |
 | `IntegratedGradients`                         | ![][castle-intgrad]            | ![][streetsign-intgrad]            |
+| `InputTimesGradient`                          | ![][castle-ixg]                | ![][streetsign-ixg]                |
 
 > [!TIP]
 > The heatmaps shown above were created using a VGG-16 vision model 

src/gradient.jl

@@ -76,12 +76,17 @@ function call_analyzer(
 end
 
 """
-    SmoothGrad(analyzer, [n=50, std=1.0f0, rng=GLOBAL_RNG])
-    SmoothGrad(analyzer, [n=50, distribution=Normal(0.0f0, 1.0f0), rng=GLOBAL_RNG])
+    SmoothGrad(analyzer)
+    SmoothGrad(analyzer, [n, std, rng]])
+    SmoothGrad(analyzer, [n, distribution, rng])
 
 Analyze model by calculating a smoothed sensitivity map.
 This is done by averaging sensitivity maps of a `Gradient` analyzer over random samples
-in a neighborhood of the input, typically by adding Gaussian noise with mean 0.
+in a neighborhood of the input.
+Defaults to 50 samples from the normal distribution with zero mean and `std=1.0f0`.
+
+For optimal results, $REF_SMILKOV_SMOOTHGRAD recommends setting `std` between 10% and 20% of the input range of each sample,
+e.g. `std = 0.1 * (maximum(input) - minimum(input))`.
 
 # References
 - $REF_SMILKOV_SMOOTHGRAD

src/input_augmentation.jl

@@ -9,14 +9,13 @@ end
 (s::AugmentationSelector)(out) = s.indices
 
 """
-    NoiseAugmentation(analyzer, n)
-    NoiseAugmentation(analyzer, n, std::Real)
-    NoiseAugmentation(analyzer, n, distribution::Sampleable)
+    NoiseAugmentation(analyzer, n, [std::Real, rng])
+    NoiseAugmentation(analyzer, n, [distribution::Sampleable, rng])
 
 A wrapper around analyzers that augments the input with `n` samples of additive noise sampled from a scalar `distribution`.
 This input augmentation is then averaged to return an `Explanation`.
+Defaults to the normal distribution with zero mean and `std=1.0f0`.
 
-Defaults to the normal distribution `Normal(0, std^2)` with `std=1.0f0`.
 For optimal results, $REF_SMILKOV_SMOOTHGRAD recommends setting `std` between 10% and 20% of the input range of each sample,
 e.g. `std = 0.1 * (maximum(input) - minimum(input))`.
 
@@ -32,17 +31,14 @@ struct NoiseAugmentation{A<:AbstractXAIMethod,D<:Sampleable,R<:AbstractRNG} <:
     rng::R
 
     function NoiseAugmentation(
-        analyzer::A, n::Int, distribution::D, rng::R
+        analyzer::A, n::Int, distribution::D, rng::R=GLOBAL_RNG
     ) where {A<:AbstractXAIMethod,D<:Sampleable,R<:AbstractRNG}
-        n < 2 &&
-            throw(ArgumentError("Number of noise samples `n` needs to be larger than one."))
+        n < 1 && throw(ArgumentError("Number of samples `n` needs to be larger than zero."))
         return new{A,D,R}(analyzer, n, distribution, rng)
     end
 end
-function NoiseAugmentation(analyzer, n, std::T=1.0f0, rng=GLOBAL_RNG) where {T<:Real}
-    return NoiseAugmentation(analyzer, n, Normal(zero(T), std^2), rng)
-end
-function NoiseAugmentation(analyzer, n, distribution::Sampleable, rng=GLOBAL_RNG)
+function NoiseAugmentation(analyzer, n::Int, std::T=1.0f0, rng=GLOBAL_RNG) where {T<:Real}
+    distribution = Normal(zero(T), std^2)
     return NoiseAugmentation(analyzer, n, distribution, rng)
 end
 

test/runtests.jl

@@ -21,6 +21,10 @@ using JET
         end
     end
 
+    @testset "Constructors" begin
+        @info "Testing constructors..."
+        include("test_constructors.jl")
+    end
     @testset "CNN" begin
         @info "Testing analyzers on CNN..."
         include("test_cnn.jl")

test/test_constructors.jl

@@ -0,0 +1,16 @@
+using ExplainableAI
+using Test
+using Distributions: Normal
+using StableRNGs: StableRNG
+
+@testset "Optional arguments" begin
+    distribution = Normal(0.0f0, 1.0f0)
+    rng = StableRNG(123)
+
+    @test_nowarn SmoothGrad(identity, 50, distribution)
+    @test_nowarn SmoothGrad(identity, 50, distribution, rng)
+
+    gradient_analyzer = Gradient(identity)
+    @test_nowarn NoiseAugmentation(gradient_analyzer, 50, distribution)
+    @test_nowarn NoiseAugmentation(gradient_analyzer, 50, distribution, rng)
+end