Refactor heatmapping using symbols (#32)

Author: adrhill

src/heatmap.jl

@@ -1,59 +1,73 @@
-const SingleChannelImage = AbstractArray{<:Real,2}
-abstract type AbstractActivationNormalizer end
-abstract type AbstractColorReducer end
+# NOTE: Heatmapping assumes Flux's WHCN convention (width, height, color channels, batch size).
 
 """
     heatmap(expl; kwargs...)
 
 Visualize explanation.
+Assumes the Flux's WHCN convention (width, height, color channels, batch size).
+
+## Keyword arguments
+-`cs::ColorScheme`: ColorScheme that is applied. Defaults to `ColorSchemes.bwr`.
+-`reduce::Symbol`: How the color channels are reduced to a single number to apply a colorscheme.
+    Can be either `:sum` or `:maxabs`. `:sum` sums up all color channels for each pixel.
+    `:maxabs` selects the `maximum(abs, x)` over the color channel in each pixel.
+    Default is `:sum`.
+-`normalize::Symbol`: How the color channel reduced heatmap is normalized before the colorscheme is applied.
+    Can be either `:extrema` or `:centered`. Default for use with colorscheme `bwr` is `:centered`.
 """
 function heatmap(
     expl::AbstractArray;
     cs::ColorScheme=ColorSchemes.bwr,
-    normalizer::AbstractActivationNormalizer=MaxAbsNormalizer(),
-    reducer::AbstractColorReducer=SumReducer(),
-    nchannels::Int=3,
+    reduce::Symbol=:sum,
+    normalize::Symbol=:centered,
     permute::Bool=true,
 )
-    img = normalizer(reducer(drop_singleton_dims(expl), nchannels))
+    _size = size(expl)
+    length(_size) != 4 && throw(
+        DomainError(
+            _size,
+            """heatmap assumes Flux's WHCN convention (width, height, color channels, batch size) for the input.
+            Please reshape your attribution to match this format if your model doesn't adhere to this convention.""",
+        ),
+    )
+    _size[end] != 1 && throw(
+        DomainError(
+            _size[end],
+            """heatmap is only applicable to a single attribution, got a batch dimension of $(_size[end]).""",
+        ),
+    )
+    # drop batch dim -> reduce color channels -> normalize image -> apply color scheme
+    img = _normalize(_reduce(dropdims(expl; dims=4), reduce), normalize)
     permute && (img = permutedims(img))
-    return get(cs, img)
+    return ColorSchemes.get(cs, img)
 end
 
 # Normalize activations across pixels
-struct MaxAbsNormalizer <: AbstractActivationNormalizer end
-function (::MaxAbsNormalizer)(img::SingleChannelImage)
-    absmax = maximum(abs, img)
-    return img / (2 * absmax) .+ 0.5
-end
-
-struct RangeNormalizer <: AbstractActivationNormalizer end
-function (::RangeNormalizer)(img::SingleChannelImage)
-    min, max = extrema(img)
-    return (img .- min) / (max - min)
+function _normalize(attr, method::Symbol)
+    if method == :centered
+        min, max = (-1, 1) .* maximum(abs, attr)
+    elseif method == :extrema
+        min, max = extrema(attr)
+    else
+        throw(
+            ArgumentError(
+                "Color scheme normalizer :$method not supported, `normalize` should be :extrema or :centered",
+            ),
+        )
+    end
+    return (attr .- min) / (max - min)
 end
 
 # Reduces activation in a pixel with multiple color channels into a single activation
-struct MaxAbsReducer <: AbstractColorReducer end
-function (::MaxAbsReducer)(img, nchannels)
-    nchannels == 1 && return img
-    dim = find_color_channel(img, nchannels)
-    return dropdims(maximum(abs, img; dims=dim); dims=dim)
-end
-
-struct SumReducer <: AbstractColorReducer end
-function (::SumReducer)(img, nchannels)
-    nchannels == 1 && return img
-    dim = find_color_channel(img, nchannels)
-    return dropdims(sum(img; dims=dim); dims=dim)
-end
-
-function find_color_channel(img, nchannels)
-    colordims = findall(size(img) .== nchannels)
-    if length(colordims) == 0
-        throw(ArgumentError("No dimension with nchannels=$nchannels color channels found."))
-    elseif length(colordims) > 1
-        throw(ArgumentError("Several dimensions of length $nchannels found."))
+function _reduce(attr, method::Symbol)
+    if method == :maxabs
+        return dropdims(maximum(abs, attr; dims=3); dims=3)
+    elseif method == :sum
+        return dropdims(sum(attr; dims=3); dims=3)
     end
-    return first(colordims)
+    throw(
+        ArgumentError(
+            "Color channel reducer :$method not supported, `reduce` should be :maxabs or :sum",
+        ),
+    )
 end

test/references/heatmaps/reduce_maxabs_normalize_centered.txt

@@ -0,0 +1 @@
+▀▀

test/references/heatmaps/reduce_maxabs_normalize_extrema.txt

@@ -0,0 +1 @@
+▀▀

test/references/heatmaps/reduce_sum_normalize_centered.txt

@@ -0,0 +1 @@
+▀▀

test/references/heatmaps/reduce_sum_normalize_extrema.txt

@@ -0,0 +1 @@
+▀▀

test/test_heatmaps.jl

@@ -1,34 +1,24 @@
 using ExplainabilityMethods
-using ExplainabilityMethods: MaxAbsNormalizer, RangeNormalizer
-using ExplainabilityMethods: MaxAbsReducer, SumReducer
 
-# Defaults assume nchannels=3
-A = rand(Float32, 3, 4, 5)
-B = reshape(A, 3, 1, 4, 5, 1)
-@test heatmap(A; normalizer=MaxAbsNormalizer(), reducer=MaxAbsReducer()) ≈
-    heatmap(B; normalizer=MaxAbsNormalizer(), reducer=MaxAbsReducer())
-@test heatmap(A; normalizer=RangeNormalizer(), reducer=SumReducer()) ≈
-    heatmap(B; normalizer=RangeNormalizer(), reducer=SumReducer())
+# NOTE: Heatmapping assumes Flux's WHCN convention (width, height, color channels, batch size).
+shape = (2, 2, 3, 1)
+A = reshape(collect(Float32, 1:prod(shape)), shape)
 
-A = rand(Float32, 3, 4, 5)
-B = reshape(A, 3, 1, 4, 1, 5)
-@test heatmap(A; normalizer=MaxAbsNormalizer(), reducer=MaxAbsReducer()) ≈
-    heatmap(B; normalizer=MaxAbsNormalizer(), reducer=MaxAbsReducer())
-@test heatmap(A; normalizer=RangeNormalizer(), reducer=SumReducer()) ≈
-    heatmap(B; normalizer=RangeNormalizer(), reducer=SumReducer())
+reducers = [:sum, :maxabs]
+normalizers = [:extrema, :centered]
+for r in reducers
+    for n in normalizers
+        h = @inferred heatmap(A; reduce=r, normalize=n)
+        @test_reference "references/heatmaps/reduce_$(r)_normalize_$(n).txt" h
+    end
+end
 
-# Test with single channel
-A = rand(Float32, 4, 5)
-B = reshape(A, 4, 1, 1, 5, 1)
-@test heatmap(A; nchannels=1, normalizer=MaxAbsNormalizer(), reducer=MaxAbsReducer()) ≈
-    heatmap(B; nchannels=1, normalizer=MaxAbsNormalizer(), reducer=MaxAbsReducer())
-@test heatmap(A; nchannels=1, normalizer=RangeNormalizer(), reducer=SumReducer()) ≈
-    heatmap(B; nchannels=1, normalizer=RangeNormalizer(), reducer=SumReducer())
+@test_throws ArgumentError heatmap(A, reduce=:foo)
+@test_throws ArgumentError heatmap(A, normalize=:bar)
 
-# Test with 2 color channels
-A = rand(Float32, 2, 3, 4)
-B = reshape(A, 2, 1, 1, 3, 4)
-@test heatmap(A; nchannels=2, normalizer=MaxAbsNormalizer(), reducer=MaxAbsReducer()) ≈
-    heatmap(B; nchannels=2, normalizer=MaxAbsNormalizer(), reducer=MaxAbsReducer())
-@test heatmap(A; nchannels=2, normalizer=RangeNormalizer(), reducer=SumReducer()) ≈
-    heatmap(B; nchannels=2, normalizer=RangeNormalizer(), reducer=SumReducer())
+B = reshape(A, 2, 2, 3, 1, 1)
+@test_throws DomainError heatmap(B)
+B = reshape(A, 2, 2, 3)
+@test_throws DomainError heatmap(B)
+B = reshape(A, 2, 2, 1, 3)
+@test_throws DomainError heatmap(B)