Improve documentation of analyze and heatmap

Author: adrhill

docs/literate/example.jl

@@ -50,14 +50,14 @@ input = reshape(x, 28, 28, 1, :);
 
 # ## Calling the analyzer
 # We can now select an analyzer of our choice
-# and call [`analyze`](@ref) to get an `Explanation`:
+# and call [`analyze`](@ref) to get an [`Explanation`](@ref):
 analyzer = LRP(model)
 expl = analyze(input, analyzer);
 
-# This `Explanation` bundles the following data:
+# This [`Explanation`](@ref) bundles the following data:
 # * `expl.attribution`: the analyzer's attribution
 # * `expl.output`: the model output
-# * `expl.neuron_selection`: the neuron index of used for the attribution
+# * `expl.neuron_selection`: the neuron index used for the attribution
 # * `expl.analyzer`: a symbol corresponding the used analyzer, e.g. `:LRP`
 
 # Finally, we can visualize the `Explanation` through [`heatmap`](@ref):

docs/src/api.md

@@ -2,6 +2,7 @@
 All methods in ExplainableAI.jl work by calling `analyze` on an input and an analyzer:
 ```@docs
 analyze
+Explanation
 heatmap
 ```
 

src/ExplainableAI.jl

@@ -36,7 +36,7 @@ include("preprocessing.jl")
 export analyze
 
 # Analyzers
-export AbstractXAIMethod
+export AbstractXAIMethod, Explanation
 export Gradient, InputTimesGradient
 export NoiseAugmentation, SmoothGrad
 export InterpolationAugmentation, IntegratedGradients

src/analyze_api.jl

@@ -11,7 +11,7 @@ const BATCHDIM_MISSING = ArgumentError(
     analyze(input, method)
     analyze(input, method, neuron_selection)
 
-Return raw classifier output and explanation.
+Return an [`Explanation`](@ref) containing the attribution and the raw classifier output.
 If `neuron_selection` is specified, the explanation will be calculated for that neuron.
 Otherwise, the output neuron with the highest activation is automatically chosen.
 
@@ -57,10 +57,15 @@ function _analyze(
     return method(input, sel; kwargs...)
 end
 
-# for convenience, the anaylyzer can be called directly
+"""
+Return type of analyzers when calling `analyze`.
 
-# Explanations and outputs are returned in a wrapper.
-# Metadata such as the analyzer allows dispatching on functions like `heatmap`.
+Contains:
+* `attribution`: the analyzer's attribution
+* `output`: the model output
+* `neuron_selection`: the neuron index used for the attribution
+* `analyzer`: a symbol corresponding the used analyzer, e.g. `:LRP`
+"""
 struct Explanation{A,O,I,L}
     attribution::A
     output::O

src/heatmap.jl

@@ -8,11 +8,9 @@ const HEATMAPPING_PRESETS = Dict{Symbol,Tuple{ColorScheme,Symbol,Symbol}}(
 )
 
 """
-    heatmap(expl::Explanation; kwargs...)
-    heatmap(attr::AbstractArray; kwargs...)
-
-    heatmap(input, analyzer::AbstractXAIMethod)
-    heatmap(input, analyzer::AbstractXAIMethod, neuron_selection::Int)
+    heatmap(explanation)
+    heatmap(input, analyzer)
+    heatmap(input, analyzer, neuron_selection)
 
 Visualize explanation.
 Assumes Flux's WHCN convention (width, height, color channels, batch size).
@@ -37,7 +35,7 @@ Assumes Flux's WHCN convention (width, height, color channels, batch size).
 - `unpack_singleton::Bool`: When heatmapping a batch with a single sample, setting `unpack_singleton=true`
     will return an image instead of an Vector containing a single image.
 
-**Note:** these keyword arguments can't be used when calling `heatmap` with an analyzer.
+**Note:** keyword arguments can't be used when calling `heatmap` with an analyzer.
 """
 function heatmap(
     attr::AbstractArray{T,N};