Sort out PROBE

src/CounterfactualExplanations.jl

@@ -14,6 +14,7 @@ export AbstractCounterfactualExplanation
 export AbstractFittedModel
 export AbstractGenerator
 export AbstractParallelizer
+export AbstractConvergence
 
 # Traits:
 include("traits/traits.jl")
@@ -79,11 +80,16 @@ export generator_catalogue
 export generate_perturbations, conditions_satisfied, mutability_constraints
 export Generator, @objective, @threshold
 
+### Convergence
+include("convergence/Convergence.jl")
+using .Convergence
+
 ### CounterfactualExplanation
 # argmin 
 ###
 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 +99,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/base_types.jl

@@ -15,3 +15,6 @@ Base.broadcastable(gen::AbstractGenerator) = Ref(gen)
 
 "An abstract type for parallelizers."
 abstract type AbstractParallelizer end
+
+"An abstract type that serves as the base type for convergence objects."
+abstract type AbstractConvergence end

src/convergence/Convergence.jl

@@ -0,0 +1,60 @@
+module Convergence
+
+using Distributions
+using Flux
+using LinearAlgebra
+using ..CounterfactualExplanations
+using ..Generators
+using ..Models
+using ..Objectives
+
+include("decision_threshold.jl")
+include("generator_conditions.jl")
+include("invalidation_rate.jl")
+include("max_iter.jl")
+
+"""
+    convergence_catalogue
+
+A dictionary containing all convergence criteria.
+"""
+const convergence_catalogue = Dict(
+    :decision_threshold => DecisionThresholdConvergence(),
+    :generator_conditions => GeneratorConditionsConvergence(),
+    :max_iter => MaxIterConvergence(),
+    :invalidation_rate => InvalidationRateConvergence(),
+)
+
+"""
+    get_convergence_type(convergence::AbstractConvergence)
+
+Returns the convergence object.
+"""
+function get_convergence_type(convergence::AbstractConvergence)
+    return convergence
+end
+
+"""
+    get_convergence_type(convergence::Symbol)
+
+Returns the convergence object from the dictionary of default convergence types.
+"""
+function get_convergence_type(convergence::Symbol)
+    return get(
+        convergence_catalogue,
+        convergence,
+        () -> error("Convergence criterion not recognized: $convergence."),
+    )
+end
+
+export convergence_catalogue
+export converged
+export get_convergence_type
+export hinge_loss, invalidation_rate
+export threshold_reached
+export DecisionThresholdConvergence
+export GeneratorConditionsConvergence
+export InvalidationRateConvergence
+export MaxIterConvergence
+
+end

src/convergence/decision_threshold.jl

@@ -0,0 +1,27 @@
+Base.@kwdef struct DecisionThresholdConvergence <: AbstractConvergence
+    decision_threshold::AbstractFloat = 0.5
+    max_iter::Int = 100
+    min_success_rate::AbstractFloat = 0.75
+end
+
+"""
+    converged(convergence::DecisionThresholdConvergence, ce::CounterfactualExplanation)
+
+Checks if the counterfactual search has converged when the convergence criterion is the decision threshold.
+"""
+function converged(
+    convergence::DecisionThresholdConvergence, ce::AbstractCounterfactualExplanation
+)
+    return threshold_reached(ce)
+end
+
+"""
+    threshold_reached(ce::CounterfactualExplanation)
+
+Determines if the predefined threshold for the target class probability has been reached.
+"""
+function threshold_reached(ce::AbstractCounterfactualExplanation)
+    γ = ce.convergence.decision_threshold
+    success_rate = sum(target_probs(ce) .>= γ) / ce.num_counterfactuals
+    return success_rate > ce.convergence.min_success_rate
+end

src/convergence/generator_conditions.jl

@@ -0,0 +1,17 @@
+Base.@kwdef struct GeneratorConditionsConvergence <: AbstractConvergence
+    decision_threshold::AbstractFloat = 0.5
+    gradient_tol::AbstractFloat = 1e-2
+    max_iter::Int = 100
+    min_success_rate::AbstractFloat = 0.75
+end
+
+"""
+    converged(convergence::GeneratorConditionsConvergence, ce::CounterfactualExplanation)
+
+Checks if the counterfactual search has converged when the convergence criterion is generator_conditions.
+"""
+function converged(
+    convergence::GeneratorConditionsConvergence, ce::AbstractCounterfactualExplanation
+)
+    return threshold_reached(ce) && Generators.conditions_satisfied(ce.generator, ce)
+end

src/convergence/invalidation_rate.jl

@@ -0,0 +1,18 @@
+Base.@kwdef struct InvalidationRateConvergence <: AbstractConvergence
+    invalidation_rate::AbstractFloat = 0.1
+    max_iter::Int = 100
+    variance::AbstractFloat = 0.01
+end
+
+"""
+    converged(convergence::InvalidationRateConvergence, ce::CounterfactualExplanation)
+
+Checks if the counterfactual search has converged when the convergence criterion is invalidation rate.
+"""
+function converged(
+    convergence::InvalidationRateConvergence, ce::AbstractCounterfactualExplanation
+)
+    ir = Objectives.invalidation_rate(ce)
+    label = Models.predict_label(ce.M, ce.data, ce.x′)[1]
+    return label == ce.target && convergence.invalidation_rate > ir
+end

src/convergence/max_iter.jl

@@ -0,0 +1,12 @@
+Base.@kwdef struct MaxIterConvergence <: AbstractConvergence
+    max_iter::Int = 100
+end
+
+"""
+    converged(convergence::MaxIterConvergence, ce::CounterfactualExplanation)
+
+Checks if the counterfactual search has converged when the convergence criterion is maximum iterations. This means the counterfactual search will not terminate until the maximum number of iterations has been reached independently of the other convergence criteria.
+"""
+function converged(convergence::MaxIterConvergence, ce::AbstractCounterfactualExplanation)
+    return ce.search[:iteration_count] == convergence.max_iter
+end

src/counterfactuals/Counterfactuals.jl

@@ -1,21 +1,24 @@
+using .Convergence
 using .DataPreprocessing
 using .GenerativeModels
 using .Generators
 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("growing_spheres.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")