Simplification of the core struct and related methods

src/CounterfactualExplanations.jl

@@ -84,6 +84,7 @@ export Generator, @objective, @threshold
 ###
 include("counterfactuals/Counterfactuals.jl")
 export CounterfactualExplanation
+export generate_counterfactual
 export initialize!, update!
 export total_steps, converged, terminated, path, target_probs
 export animate_path
@@ -93,9 +94,6 @@ export animate_path
 include("data/Data.jl")
 using .Data
 
-include("generate_counterfactual/generate_counterfactual.jl")
-export generate_counterfactual
-
 include("evaluation/Evaluation.jl")
 using .Evaluation
 

src/counterfactuals/Counterfactuals.jl

@@ -5,16 +5,19 @@ using .Models
 using ChainRulesCore
 using Flux
 using MLUtils
+using MultivariateStats
 using Statistics
 using StatsBase
 
 include("core_struct.jl")
 include("convergence.jl")
 include("encodings.jl")
+include("generate_counterfactual.jl")
 include("info_extraction.jl")
 include("initialisation.jl")
 include("latent_space_mappings.jl")
 include("path_tracking.jl")
 include("printing.jl")
 include("search.jl")
 include("utils.jl")
+include("vectorised.jl")

src/counterfactuals/convergence.jl

@@ -4,47 +4,73 @@
 A convenience method to determine if the counterfactual search has terminated.
 """
 function terminated(ce::CounterfactualExplanation)
-    if ce.M isa Models.TreeModel
-        return in_target_class(ce)
-    end
     return converged(ce) || steps_exhausted(ce)
 end
 
 """
-    in_target_class(ce::CounterfactualExplanation)
+    converged(ce::CounterfactualExplanation)
 
-Check if the counterfactual is in the target class.
+A convenience method to determine if the counterfactual search has converged.
+The search is considered to have converged only if the counterfactual is valid.
 """
-function in_target_class(ce::CounterfactualExplanation)
-    return Models.predict_label(ce.M, ce.data, decode_state(ce))[1] == ce.target
+function converged(ce::CounterfactualExplanation)
+    return converged(ce, Val(ce.convergence[:converge_when]))
 end
 
 """
-    converged(ce::CounterfactualExplanation)
+    converged(ce::CounterfactualExplanation, ::Val{:decision_threshold})
 
-A convenience method to determine if the counterfactual search has converged. The search is considered to have converged only if the counterfactual is valid.
+Checks if the counterfactual search has converged when the convergence criterion is the decision threshold.
 """
-function converged(ce::CounterfactualExplanation)
-    if ce.generator isa GrowingSpheresGenerator || ce.generator isa FeatureTweakGenerator
-        conv = ce.search[:converged]
-    elseif ce.convergence[:converge_when] == :decision_threshold
-        conv = threshold_reached(ce)
-    elseif ce.convergence[:converge_when] == :generator_conditions
-        conv = threshold_reached(ce) && Generators.conditions_satisfied(ce.generator, ce)
-    elseif ce.convergence[:converge_when] == :max_iter
-        conv = false
-    elseif ce.convergence[:converge_when] == :invalidation_rate
-        ir = Generators.invalidation_rate(ce)
-        # gets the label from an array, not sure why it is an array though.
-        label = predict_label(ce.M, ce.data, ce.x′)[1]
-        conv = label == ce.target && ce.params[:invalidation_rate] > ir
-    elseif (ce.convergence[:converge_when] == :early_stopping)
-        conv = steps_exhausted(ce)
-    else
-        @error "Convergence criterion not recognized."
-    end
-
-    return conv
+function converged(ce::CounterfactualExplanation, ::Val{:decision_threshold})
+    return threshold_reached(ce)
+end
+
+"""
+    converged(ce::CounterfactualExplanation, ::Val{:generator_conditions})
+
+Checks if the counterfactual search has converged when the convergence criterion is generator_conditions.
+"""
+function converged(ce::CounterfactualExplanation, ::Val{:generator_conditions})
+    return threshold_reached(ce) && Generators.conditions_satisfied(ce.generator, ce)
+end
+
+"""
+    converged(ce::CounterfactualExplanation, ::Val{:max_iter})
+
+Checks if the counterfactual search has converged when the convergence criterion is maximum iterations.
+"""
+function converged(ce::CounterfactualExplanation, ::Val{:max_iter})
+    return false
+end
+
+"""
+    converged(ce::CounterfactualExplanation, ::Val{:invalidation_rate})
+
+Checks if the counterfactual search has converged when the convergence criterion is invalidation rate.
+"""
+function converged(ce::CounterfactualExplanation, ::Val{:invalidation_rate})
+    ir = Generators.invalidation_rate(ce)
+    label = predict_label(ce.M, ce.data, ce.x′)[1]
+    return label == ce.target && ce.generator.invalidation_rate > ir
+end
+
+"""
+    converged(ce::CounterfactualExplanation, ::Val{:early_stopping})
+
+Checks if the counterfactual search has converged when the convergence criterion is early stopping.
+"""
+function converged(ce::CounterfactualExplanation, ::Val{:early_stopping})
+    return steps_exhausted(ce)
+end
+
+"""
+    converged(ce::CounterfactualExplanation, ::Val{sym})
+
+Throws an error when the `converged()` method is called on an unrecognized convergence criterion.
+"""
+function converged(ce::CounterfactualExplanation, ::Val{sym}) where {sym}
+    @error "Convergence criterion not recognized: $sym"
 end
 
 """
@@ -64,6 +90,7 @@ end
 A convenience method that determines if the predefined threshold for the target class probability has been reached for a specific sample `x`.
 """
 function threshold_reached(ce::CounterfactualExplanation, x::AbstractArray)
+    print(ce.convergence[:min_success_rate])
     γ = ce.convergence[:decision_threshold]
     success_rate = sum(target_probs(ce, x) .>= γ) / ce.num_counterfactuals
     return success_rate > ce.convergence[:min_success_rate]

src/counterfactuals/core_struct.jl

@@ -10,7 +10,6 @@ mutable struct CounterfactualExplanation <: AbstractCounterfactualExplanation
     data::DataPreprocessing.CounterfactualData
     M::Models.AbstractFittedModel
     generator::Generators.AbstractGenerator
-    generative_model_params::NamedTuple
     params::Dict
     search::Union{Dict,Nothing}
     convergence::Dict
@@ -28,12 +27,8 @@ end
 		max_iter::Int = 100,
 		num_counterfactuals::Int = 1,
 		initialization::Symbol = :add_perturbation,
-		generative_model_params::NamedTuple = (;),
 		min_success_rate::AbstractFloat=0.99,
         converge_when::Symbol=:decision_threshold,
-        invalidation_rate::AbstractFloat=0.5,
-        learning_rate::AbstractFloat=1.0,
-        variance::AbstractFloat=0.01,
 	)
 
 Outer method to construct a `CounterfactualExplanation` structure.
@@ -46,15 +41,11 @@ function CounterfactualExplanation(
     generator::Generators.AbstractGenerator;
     num_counterfactuals::Int=1,
     initialization::Symbol=:add_perturbation,
-    generative_model_params::NamedTuple=(;),
     max_iter::Int=100,
     decision_threshold::AbstractFloat=0.5,
     gradient_tol::AbstractFloat=parameters[:τ],
     min_success_rate::AbstractFloat=parameters[:min_success_rate],
     converge_when::Symbol=:decision_threshold,
-    invalidation_rate::AbstractFloat=0.5,
-    learning_rate::AbstractFloat=1.0,
-    variance::AbstractFloat=0.01,
 )
 
     # Assertions:
@@ -67,6 +58,8 @@ function CounterfactualExplanation(
         :invalidation_rate,
         :early_stopping,
     ]
+    @assert !(converge_when == :invalidation_rate && isnothing(generator.invalidation_rate)) "The convergence criterion is invalidation rate but no invalidation rate has been provided."
+    @assert !(converge_when == :invalidation_rate && isnothing(generator.variance)) "The convergence criterion is invalidation rate but no variance has been provided."
 
     # Factual:
     x = typeof(x) == Int ? select_factual(data, x) : x
@@ -79,9 +72,6 @@ function CounterfactualExplanation(
         :mutability => DataPreprocessing.mutability_constraints(data),
         :latent_space => generator.latent_space,
         :dim_reduction => generator.dim_reduction,
-        :invalidation_rate => invalidation_rate,
-        :learning_rate => learning_rate,
-        :variance => variance,
     )
     ids = findall(predict_label(M, data) .== target)
     n_candidates = minimum([size(data.y, 2), 1000])
@@ -107,7 +97,6 @@ function CounterfactualExplanation(
         data,
         M,
         deepcopy(generator),
-        generative_model_params,
         params,
         nothing,
         convergence,

src/counterfactuals/encodings.jl

@@ -1,6 +1,3 @@
-using MultivariateStats
-using StatsBase
-
 """
     encode_array(dt::MultivariateStats.AbstractDimensionalityReduction, x::AbstractArray)
 

src/generate_counterfactual/core.jl → src/counterfactuals/generate_counterfactual.jl

@@ -48,16 +48,12 @@ function generate_counterfactual(
     generator::AbstractGenerator;
     num_counterfactuals::Int=1,
     initialization::Symbol=:add_perturbation,
-    generative_model_params::NamedTuple=(;),
     max_iter::Int=100,
     decision_threshold::AbstractFloat=0.5,
     gradient_tol::AbstractFloat=parameters[:τ],
     min_success_rate::AbstractFloat=parameters[:min_success_rate],
     converge_when::Symbol=:decision_threshold,
     timeout::Union{Nothing,Int}=nothing,
-    invalidation_rate::AbstractFloat=0.1,
-    learning_rate::AbstractFloat=1.0,
-    variance::AbstractFloat=0.01,
 )
     # Initialize:
     ce = CounterfactualExplanation(
@@ -68,15 +64,11 @@ function generate_counterfactual(
         generator;
         num_counterfactuals=num_counterfactuals,
         initialization=initialization,
-        generative_model_params=generative_model_params,
         max_iter=max_iter,
         min_success_rate=min_success_rate,
         decision_threshold=decision_threshold,
         gradient_tol=gradient_tol,
         converge_when=converge_when,
-        invalidation_rate=invalidation_rate,
-        learning_rate=learning_rate,
-        variance=variance,
     )
 
     # Search:

src/counterfactuals/initialisation.jl

@@ -8,19 +8,7 @@ Initializes the starting point for the factual(s):
 """
 function initialize_state(ce::CounterfactualExplanation)
     @assert ce.initialization ∈ [:identity, :add_perturbation]
-
     s′ = ce.s′
-    data = ce.data
-
-    # No perturbation:
-    if ce.initialization == :identity
-        return s′
-    end
-
-    # If latent space, initial point is random anyway:
-    if ce.params[:latent_space]
-        return s′
-    end
 
     # Add random perturbation following Slack (2021): https://arxiv.org/abs/2106.02666
     if ce.initialization == :add_perturbation

src/counterfactuals/latent_space_mappings.jl

@@ -49,7 +49,7 @@ function map_to_latent(
 
     if ce.params[:latent_space]
         generative_model = DataPreprocessing.get_generative_model(
-            data; ce.generative_model_params...
+            data; generator.generative_model_params...
         )
         # map counterfactual to latent space: s′=z′∼p(z|x)
         s′, _, _ = GenerativeModels.rand(generative_model.encoder, s′)

src/counterfactuals/path_tracking.jl

@@ -17,7 +17,6 @@ end
 Returns the counterfactual probabilities for each step of the search.
 """
 function counterfactual_probability_path(ce::CounterfactualExplanation)
-    M = ce.M
     p = map(X -> counterfactual_probability(ce, X), path(ce))
     return p
 end

src/generators/gradient_based/base.jl

@@ -16,6 +16,9 @@ mutable struct GradientBasedGenerator <: AbstractGradientBasedGenerator
     latent_space::Bool
     dim_reduction::Bool
     opt::Flux.Optimise.AbstractOptimiser
+    invalidation_rate::Union{Nothing,AbstractFloat}
+    variance::Union{Nothing,AbstractFloat}
+    generative_model_params::NamedTuple
 end
 
 """
@@ -25,6 +28,9 @@ end
 		λ::Union{Nothing,AbstractFloat,Vector{AbstractFloat}}=nothing,
 		latent_space::Bool::false,
 		opt::Flux.Optimise.AbstractOptimiser=Flux.Descent(),
+        invalidation_rate::AbstractFloat=nothing,
+        variance::AbstractFloat=nothing,
+        generative_model_params::NamedTuple=(;),
 	)
 
 Default outer constructor for `GradientBasedGenerator`.
@@ -35,6 +41,9 @@ Default outer constructor for `GradientBasedGenerator`.
 - `λ::Union{Nothing,AbstractFloat,Vector{AbstractFloat}}=nothing`: The weight of the penalty function.
 - `latent_space::Bool=false`: Whether to use the latent space of a generative model to generate counterfactuals.
 - `opt::Flux.Optimise.AbstractOptimiser=Flux.Descent()`: The optimizer to use for the generator.
+- `invalidation_rate::AbstractFloat=nothing`: The invalidation rate of the counterfactual explanation.
+- `variance::AbstractFloat=nothing`: The variance term to be used when calculating the invalidation rate of the counterfactual explanation.
+- `generative_model_params::NamedTuple`: The parameters of the generative model associated with the generator.
 
 # Returns
 - `generator::GradientBasedGenerator`: A gradient-based counterfactual generator.
@@ -46,34 +55,16 @@ function GradientBasedGenerator(;
     latent_space::Bool=false,
     dim_reduction::Bool=false,
     opt::Flux.Optimise.AbstractOptimiser=Flux.Descent(),
+    invalidation_rate::Union{Nothing,AbstractFloat}=nothing,
+    variance::Union{Nothing,AbstractFloat}=nothing,
+    generative_model_params::NamedTuple=(;),
 )
     @assert !(isnothing(λ) && !isnothing(penalty)) "Penalty function(s) provided but no penalty weight(s) provided."
     @assert !(isnothing(λ) && !isnothing(penalty)) "Penalty weight(s) provided but no penalty function(s) provided."
+
     if typeof(penalty) <: Vector
         @assert length(λ) == length(penalty) || length(λ) == 1 "The number of penalty weights must match the number of penalty functions or be equal to one."
         length(λ) == 1 && (λ = fill(λ[1], length(penalty)))     # if only one penalty weight is provided, use it for all penalties
     end
-    return GradientBasedGenerator(loss, penalty, λ, latent_space, dim_reduction, opt)
-end
-
-"""
-	Generator(;
-		loss::Union{Nothing,Function}=nothing,
-		penalty::Penalty=nothing,
-		λ::Union{Nothing,AbstractFloat,Vector{AbstractFloat}}=nothing,
-		latent_space::Bool::false,
-		opt::Flux.Optimise.AbstractOptimiser=Flux.Descent(),
-	)
-
-An outer constructor that allows for more convenient creation of the `GradientBasedGenerator` type.
-"""
-function Generator(;
-    loss::Union{Nothing,Function}=nothing,
-    penalty::Penalty=nothing,
-    λ::Union{Nothing,AbstractFloat,Vector{<:AbstractFloat}}=nothing,
-    latent_space::Bool=false,
-    dim_reduction::Bool=false,
-    opt::Flux.Optimise.AbstractOptimiser=Flux.Descent(),
-)
-    return GradientBasedGenerator(loss, penalty, λ, latent_space, dim_reduction, opt)
+    return GradientBasedGenerator(loss, penalty, λ, latent_space, dim_reduction, opt, invalidation_rate, variance, generative_model_params)
 end