Move core interface into XAIBase.jl package (#154)

* Add XAIBase dependency  

* Add heatmapping preset field to Explanations  

* Update `heatmap` kwarg `cs` to `colorscheme`  

* Remove tests that were moved to XAIBase   

* Update compat entries of dependencies and test dependencies  

* Update documentation

Author: adrhill

Project.toml

@@ -1,7 +1,7 @@
 name = "ExplainableAI"
 uuid = "4f1bc3e1-d60d-4ed0-9367-9bdff9846d3b"
 authors = ["Adrian Hill <[email protected]>"]
-version = "0.6.3"
+version = "1.0.0-DEV"
 
 [deps]
 ColorSchemes = "35d6a980-a343-548e-a6ea-1d62b119f2f4"
@@ -12,8 +12,10 @@ ImageTransformations = "02fcd773-0e25-5acc-982a-7f6622650795"
 MacroTools = "1914dd2f-81c6-5fcd-8719-6d5c9610ff09"
 Markdown = "d6f4376e-aef5-505a-96c1-9c027394607a"
 Random = "9a3f8284-a2c9-5f02-9a11-845980a1fd5c"
+Reexport = "189a3867-3050-52da-a836-e630ba90ab69"
 Statistics = "10745b16-79ce-11e8-11f9-7d13ad32a3b2"
 Tullio = "bc48ee85-29a4-5162-ae0b-a64e1601d4bc"
+XAIBase = "9b48221d-a747-4c1b-9860-46a1d8ba24a7"
 Zygote = "e88e6eb3-aa80-5325-afca-941959d7151f"
 
 [weakdeps]
@@ -29,6 +31,11 @@ Flux = "0.13, 0.14"
 ImageCore = "0.9, 0.10"
 ImageTransformations = "0.9, 0.10"
 MacroTools = "0.5"
+Markdown = "1"
+Random = "1"
+Reexport = "1"
+Statistics = "1"
 Tullio = "0.3"
+XAIBase = "1.2"
 Zygote = "0.6"
 julia = "1.6"

docs/make.jl

@@ -1,6 +1,6 @@
-using ExplainableAI
 using Documenter
 using Literate
+using ExplainableAI
 
 LITERATE_DIR = joinpath(@__DIR__, "src/literate")
 OUT_DIR = joinpath(@__DIR__, "src/generated")
@@ -22,7 +22,7 @@ end
 convert_literate(LITERATE_DIR, OUT_DIR)
 
 makedocs(;
-    modules=[ExplainableAI],
+    modules=[XAIBase, ExplainableAI],
     authors="Adrian Hill",
     sitename="ExplainableAI.jl",
     format=Documenter.HTML(; prettyurls=get(ENV, "CI", "false") == "true", assets=String[]),

docs/src/literate/example.jl

@@ -48,20 +48,22 @@ analyzer = LRP(model)
 expl = analyze(input, analyzer);
 
 # The return value `expl` is of type [`Explanation`](@ref) and bundles the following data:
-# * `expl.val`: the numerical output of the analyzer, e.g. an attribution or gradient
-# * `expl.output`: the model output for the given analyzer input
-# * `expl.neuron_selection`: the neuron index used for the explanation
-# * `expl.analyzer`: a symbol corresponding the used analyzer, e.g. `:LRP`
-# * `expl.extras`: an optional named tuple that can be used by analyzers
-#   to return additional information.
+# * `expl.val`: numerical output of the analyzer, e.g. an attribution or gradient
+# * `expl.output`: model output for the given analyzer input
+# * `expl.output_selection`: index of the output used for the explanation
+# * `expl.analyzer`: symbol corresponding the used analyzer, e.g. `:Gradient` or `:LRP`
+# * `expl.heatmap`: symbol indicating a preset heatmapping style,
+#     e.g. `:attibution`, `:sensitivity` or `:cam`
+# * `expl.extras`: optional named tuple that can be used by analyzers
+#     to return additional information.
 #
 # We used an LRP analyzer, so `expl.analyzer` is `:LRP`.
 expl.analyzer
 
 # By default, the explanation is computed for the maximally activated output neuron.
 # Since our digit is a 9 and Julia's indexing is 1-based,
 # the output neuron at index `10` of our trained model is maximally activated.
-expl.neuron_selection
+expl.output_selection
 
 # Finally, we obtain the result of the analyzer in form of an array.
 expl.val

docs/src/literate/heatmapping.jl

@@ -39,9 +39,9 @@ heatmap(input, analyzer)
 using ColorSchemes
 
 expl = analyze(input, analyzer)
-heatmap(expl; cs=ColorSchemes.jet)
+heatmap(expl; colorscheme=:jet)
 #-
-heatmap(expl; cs=ColorSchemes.inferno)
+heatmap(expl; colorscheme=:inferno)
 
 # Refer to the [ColorSchemes.jl catalogue](https://juliagraphics.github.io/ColorSchemes.jl/stable/basics/)
 # for a gallery of available color schemes.
@@ -84,9 +84,9 @@ heatmap(expl; rangescale=:centered)
 heatmap(expl; rangescale=:extrema)
 # However, for the `inferno` color scheme, which is not centered around zero,
 # `:extrema` leads to a heatmap with higher contrast.
-heatmap(expl; rangescale=:centered, cs=ColorSchemes.inferno)
+heatmap(expl; rangescale=:centered, colorscheme=:inferno)
 #-
-heatmap(expl; rangescale=:extrema, cs=ColorSchemes.inferno)
+heatmap(expl; rangescale=:extrema, colorscheme=:inferno)
 
 # For the full list of `heatmap` keyword arguments, refer to the [`heatmap`](@ref) documentation.
 
@@ -110,11 +110,13 @@ mosaic(heatmaps; nrow=10)
 #
 # If this bevahior is not desired,
 # `heatmap` can be called with the keyword-argument `process_batch=true`:
-heatmaps = heatmap(batch, analyzer; process_batch=true)
+expl = analyze(batch, analyzer)
+heatmaps = heatmap(expl; process_batch=true)
 mosaic(heatmaps; nrow=10)
 
 # This can be useful when comparing heatmaps for fixed output neurons:
-heatmaps = heatmap(batch, analyzer, 7; process_batch=true) # heatmaps for digit "6"
+expl = analyze(batch, analyzer, 7) # explain digit "6"
+heatmaps = heatmap(expl; process_batch=true)
 mosaic(heatmaps; nrow=10)
 
 #md # !!! note "Output type consistency"

src/ExplainableAI.jl

@@ -1,5 +1,8 @@
 module ExplainableAI
 
+using Reexport
+@reexport using XAIBase
+
 using Base.Iterators
 using MacroTools: @forward
 using Distributions: Distribution, Sampleable, Normal
@@ -15,8 +18,6 @@ using ColorSchemes
 
 include("compat.jl")
 include("bibliography.jl")
-include("neuron_selection.jl")
-include("analyze_api.jl")
 include("flux_types.jl")
 include("flux_layer_utils.jl")
 include("flux_chain_utils.jl")
@@ -31,7 +32,6 @@ include("lrp/lrp.jl")
 include("lrp/show.jl")
 include("lrp/composite_presets.jl") # uses lrp/show.jl
 include("lrp/crp.jl")
-include("heatmap.jl")
 include("preprocessing.jl")
 export analyze
 

src/analyze_api.jl

@@ -22,7 +22,7 @@ end
 
 function (analyzer::Gradient)(input, ns::AbstractNeuronSelector)
     grad, output, output_indices = gradient_wrt_input(analyzer.model, input, ns)
-    return Explanation(grad, output, output_indices, :Gradient, nothing)
+    return Explanation(grad, output, output_indices, :Gradient, :sensitivity, nothing)
 end
 
 """
@@ -42,7 +42,7 @@ end
 function (analyzer::InputTimesGradient)(input, ns::AbstractNeuronSelector)
     grad, output, output_indices = gradient_wrt_input(analyzer.model, input, ns)
     attr = input .* grad
-    return Explanation(attr, output, output_indices, :InputTimesGradient, nothing)
+    return Explanation(attr, output, output_indices, :InputTimesGradient, :attribution, nothing)
 end
 
 """

src/gradient.jl

@@ -1,3 +1,13 @@
+"""
+    AugmentationSelector(index)
+
+Neuron selector that passes through an augmented neuron selection.
+"""
+struct AugmentationSelector{I} <: AbstractNeuronSelector
+    indices::I
+end
+(s::AugmentationSelector)(out) = s.indices
+
 """
     augment_batch_dim(input, n)
 
@@ -109,6 +119,7 @@ function (aug::NoiseAugmentation)(input, ns::AbstractNeuronSelector)
         output,
         output_indices,
         augmented_expl.analyzer,
+        augmented_expl.heatmap,
         nothing,
     )
 end
@@ -148,7 +159,7 @@ function (aug::InterpolationAugmentation)(
     # Average gradients and compute explanation
     expl = (input - input_ref) .* reduce_augmentation(augmented_expl.val, aug.n)
 
-    return Explanation(expl, output, output_indices, augmented_expl.analyzer, nothing)
+    return Explanation(expl, output, output_indices, augmented_expl.analyzer, augmented_expl.heatmap, nothing)
 end
 
 """

src/heatmap.jl

@@ -94,7 +94,7 @@ function (crp::CRP)(input::AbstractArray{T,N}, ns::AbstractNeuronSelector) where
             end
         end
     end
-    return Explanation(R_return, last(as), ns(last(as)), :CRP, nothing)
+    return Explanation(R_return, last(as), ns(last(as)), :CRP, :attribution, nothing)
 end
 
 #===================#