From 380aac025ab3f1a4f7a4d153e302acb72567d1a9 Mon Sep 17 00:00:00 2001 From: pat-alt Date: Mon, 10 Aug 2026 10:14:36 +0200 Subject: [PATCH 1/4] first stab at docs --- docs/Project.toml | 5 + docs/make.jl | 7 +- docs/src/_intro.qmd | 4 +- docs/src/gpu.md | 126 +++++ docs/src/gpu.qmd | 138 ++++++ .../figure-commonmark/cell-8-output-1.svg | 436 ++++++++++++++++++ docs/src/index.md | 45 +- docs/src/introduction.md | 110 +++++ docs/src/introduction.qmd | 118 +++++ .../figure-commonmark/cell-6-output-1.svg | 383 +++++++++++++++ 10 files changed, 1367 insertions(+), 5 deletions(-) create mode 100644 docs/src/gpu.md create mode 100644 docs/src/gpu.qmd create mode 100644 docs/src/gpu_files/figure-commonmark/cell-8-output-1.svg create mode 100644 docs/src/introduction.md create mode 100644 docs/src/introduction.qmd create mode 100644 docs/src/introduction_files/figure-commonmark/cell-6-output-1.svg diff --git a/docs/Project.toml b/docs/Project.toml index 25b78593..ee7cc83c 100644 --- a/docs/Project.toml +++ b/docs/Project.toml @@ -1,3 +1,8 @@ [deps] +CounterfactualExplanations = "2f13d31b-18db-44c1-bc43-ebaf2cff0be0" CounterfactualTraining = "c5e5e904-ed0d-4ddb-967e-dbb7fd41a507" Documenter = "e30172f5-a6a5-5a46-863b-614d45cd2de4" +Flux = "587475ba-b771-5e3f-ad9e-33799f191a9c" +MLDatasets = "eb30cadb-4394-5ae3-aed4-317e484a6458" +Plots = "91a5bcdd-55d7-5caf-9e0b-520d859cae80" +Random = "9a3f8284-a2c9-5f02-9a11-845980a1fd5c" diff --git a/docs/make.jl b/docs/make.jl index 71dc1371..8051b2af 100755 --- a/docs/make.jl +++ b/docs/make.jl @@ -59,7 +59,12 @@ makedocs(; edit_link="main", assets=String[], ), - pages=["Home" => "index.md", "Reference" => "reference.md"], + pages=[ + "Home" => "index.md", + "Introduction" => "introduction.md", + "Training on GPU" => "gpu.md", + "Reference" => "reference.md", + ], ) deploydocs(; diff --git a/docs/src/_intro.qmd b/docs/src/_intro.qmd index 194159c9..37bc3e22 100644 --- a/docs/src/_intro.qmd +++ b/docs/src/_intro.qmd @@ -37,8 +37,8 @@ Pkg.add("CounterfactualTraining") ### Introduction - +A self-contained introductory example demonstrating counterfactual training with mutability protection is available on the [Introduction](introduction.md) page. ### Training on GPU - +The package provides two training approaches: a branch using [CounterfactualExplanations.jl](https://github.com/JuliaTrustworthyAI/CounterfactualExplanations.jl) as the backend for generating counterfactuals and a performance-focused `Native` submodule with GPU support. The former approach was developed during the initial research stage of the project. See [Training on GPU](gpu.md) for a full walkthrough with MNIST. diff --git a/docs/src/gpu.md b/docs/src/gpu.md new file mode 100644 index 00000000..fa4eb889 --- /dev/null +++ b/docs/src/gpu.md @@ -0,0 +1,126 @@ +# Training on GPU + + +This page demonstrates GPU training with the `Native` submodule of CounterfactualTraining.jl. + +## Two Training Approaches + +The package provides two training approaches that share the same `counterfactual_training` entry point (dispatch is based on the generator type): + +1. **Research branch** (non-native): Uses `CounterfactualExplanations.jl` structs (`GenericGenerator`, `CounterfactualExplanation`) and `TaijaParallel.parallelize` for per-sample counterfactual generation. Heavier and slower, but fully featured. Dispatched when `generator` is **not** a `NativeGenerator`. + +2. **Performance branch** (`Native` submodule): A lightweight, batched, GPU-friendly generator (`NativeGenerator`) that operates directly on `D×N` matrices without allocating `CounterfactualExplanation` objects per sample. Much faster and GPU-compatible. Dispatched when `generator isa NativeGenerator`. + +The `Native` branch was added as part of the performance optimization effort and is the recommended approach for most use cases. + +## Setup + +``` julia +using CounterfactualTraining +using CounterfactualTraining.Native +using CounterfactualExplanations +using Flux +using Plots +using Random +using MLDatasets + +Random.seed!(42) +gr() +``` + +## GPU Detection + +The code below detects available GPU backends (CUDA or AMDGPU) and falls back to CPU if none is available. This makes the example runnable in any environment, including CI. + +``` julia +has_gpu = false +device = identity +try + using CUDA + if CUDA.functional() + has_gpu = true + device = Flux.gpu + end +catch + try + using AMDGPU + if AMDGPU.functional() + has_gpu = true + device = Flux.gpu + end + catch + end +end + +if has_gpu + @info "GPU detected — training will run on GPU." +else + @info "No GPU available — training will run on CPU." +end +``` + +## Loading MNIST + +We load MNIST via `MLDatasets.jl`, flatten the 28×28 images to 784-dimensional vectors, and normalize pixel values to `[-1, 1]`. A subset of 5,000 samples is used to keep the example fast. + +``` julia +train_data = MLDatasets.MNIST(; split=:train) +X = Float32.(reshape(train_data.features, 784, :)) .* 2f0 .- 1f0 +y = train_data.targets .+ 1 # convert 0-9 labels to 1-10 + +idx = shuffle(1:size(X, 2))[1:5000] +X, y = X[:, idx], y[idx] + +domain = [(-1.0f0, 1.0f0) for _ in 1:784] +y_onehot = Flux.onehotbatch(y, 1:10) +train_set = Flux.DataLoader((X, y_onehot); batchsize=128, shuffle=true) +``` + +## Model + +We use a “chunky” MLP with two hidden layers, suitable for demonstrating GPU acceleration: + +``` julia +model = Chain( + Dense(784, 256, relu), + Dense(256, 128, relu), + Dense(128, 10) +) +``` + +## Training + +We train with `FullObjective` and `NativeGenerator`. The `device` keyword moves the model to the GPU when available; training data should already be on the device (the `DataLoader` handles batching). + +``` julia +gen = NativeGenerator() +obj = FullObjective() +opt_state = Flux.setup(Flux.Adam(1e-3), model) + +model, log = counterfactual_training( + obj, model, gen, train_set, opt_state; + device=device, nepochs=5, maxiter=10, burnin=0.0f0, + domain=domain, verbose=0 +) +``` + +## Counterfactual Visualization + +We generate counterfactuals for a few test samples, targeting a different class, and visualize the original vs. counterfactual images side by side. + +``` julia +X_test = X[:, 1:5] +targets = fill(2, 5) # target class 2 (digit 1) +data = CounterfactualData(X, y; domain=domain) +cfs, _, converged, _ = generate_counterfactuals!(model, X_test, targets, data, gen; maxiter=30) +``` + +``` julia +p1 = heatmap(reshape(X_test[:, 1], 28, 28), title="Original", color=:grays, aspect_ratio=:equal) +p2 = heatmap(reshape(cfs[:, 1], 28, 28), title="Counterfactual", color=:grays, aspect_ratio=:equal) +plot(p1, p2, layout=(1, 2), legend=false) +``` + +![](gpu_files/figure-commonmark/cell-8-output-1.svg) + +The counterfactual shows how the model would need the input to change to produce a different prediction. With counterfactual training, these changes should be plausible and respect the domain constraints of the pixel values. diff --git a/docs/src/gpu.qmd b/docs/src/gpu.qmd new file mode 100644 index 00000000..93aac131 --- /dev/null +++ b/docs/src/gpu.qmd @@ -0,0 +1,138 @@ +# Training on GPU + +This page demonstrates GPU training with the `Native` submodule of CounterfactualTraining.jl. + +## Two Training Approaches + +The package provides two training approaches that share the same `counterfactual_training` entry point (dispatch is based on the generator type): + +1. **Research branch** (non-native): Uses `CounterfactualExplanations.jl` structs (`GenericGenerator`, `CounterfactualExplanation`) and `TaijaParallel.parallelize` for per-sample counterfactual generation. Heavier and slower, but fully featured. Dispatched when `generator` is **not** a `NativeGenerator`. + +2. **Performance branch** (`Native` submodule): A lightweight, batched, GPU-friendly generator (`NativeGenerator`) that operates directly on `D×N` matrices without allocating `CounterfactualExplanation` objects per sample. Much faster and GPU-compatible. Dispatched when `generator isa NativeGenerator`. + +The `Native` branch was added as part of the performance optimization effort and is the recommended approach for most use cases. + +## Setup + +```{julia} +#| echo: true + +using CounterfactualTraining +using CounterfactualTraining.Native +using CounterfactualExplanations +using Flux +using Plots +using Random +using MLDatasets + +Random.seed!(42) +gr() +``` + +## GPU Detection + +The code below detects available GPU backends (CUDA or AMDGPU) and falls back to CPU if none is available. This makes the example runnable in any environment, including CI. + +```{julia} +#| echo: true + +has_gpu = false +device = identity +try + using CUDA + if CUDA.functional() + has_gpu = true + device = Flux.gpu + end +catch + try + using AMDGPU + if AMDGPU.functional() + has_gpu = true + device = Flux.gpu + end + catch + end +end + +if has_gpu + @info "GPU detected — training will run on GPU." +else + @info "No GPU available — training will run on CPU." +end +``` + +## Loading MNIST + +We load MNIST via `MLDatasets.jl`, flatten the 28×28 images to 784-dimensional vectors, and normalize pixel values to `[-1, 1]`. A subset of 5,000 samples is used to keep the example fast. + +```{julia} +#| echo: true + +train_data = MLDatasets.MNIST(; split=:train) +X = Float32.(reshape(train_data.features, 784, :)) .* 2f0 .- 1f0 +y = train_data.targets .+ 1 # convert 0-9 labels to 1-10 + +idx = shuffle(1:size(X, 2))[1:5000] +X, y = X[:, idx], y[idx] + +domain = [(-1.0f0, 1.0f0) for _ in 1:784] +y_onehot = Flux.onehotbatch(y, 1:10) +train_set = Flux.DataLoader((X, y_onehot); batchsize=128, shuffle=true) +``` + +## Model + +We use a "chunky" MLP with two hidden layers, suitable for demonstrating GPU acceleration: + +```{julia} +#| echo: true + +model = Chain( + Dense(784, 256, relu), + Dense(256, 128, relu), + Dense(128, 10) +) +``` + +## Training + +We train with `FullObjective` and `NativeGenerator`. The `device` keyword moves the model to the GPU when available; training data should already be on the device (the `DataLoader` handles batching). + +```{julia} +#| echo: true + +gen = NativeGenerator() +obj = FullObjective() +opt_state = Flux.setup(Flux.Adam(1e-3), model) + +model, log = counterfactual_training( + obj, model, gen, train_set, opt_state; + device=device, nepochs=5, maxiter=10, burnin=0.0f0, + domain=domain, verbose=0 +) +``` + +## Counterfactual Visualization + +We generate counterfactuals for a few test samples, targeting a different class, and visualize the original vs. counterfactual images side by side. + +```{julia} +#| echo: true + +X_test = X[:, 1:5] +targets = fill(2, 5) # target class 2 (digit 1) +data = CounterfactualData(X, y; domain=domain) +cfs, _, converged, _ = generate_counterfactuals!(model, X_test, targets, data, gen; maxiter=30) +``` + +```{julia} +#| echo: true +#| output: true + +p1 = heatmap(reshape(X_test[:, 1], 28, 28), title="Original", color=:grays, aspect_ratio=:equal) +p2 = heatmap(reshape(cfs[:, 1], 28, 28), title="Counterfactual", color=:grays, aspect_ratio=:equal) +plot(p1, p2, layout=(1, 2), legend=false) +``` + +The counterfactual shows how the model would need the input to change to produce a different prediction. With counterfactual training, these changes should be plausible and respect the domain constraints of the pixel values. diff --git a/docs/src/gpu_files/figure-commonmark/cell-8-output-1.svg b/docs/src/gpu_files/figure-commonmark/cell-8-output-1.svg new file mode 100644 index 00000000..011d2e77 --- /dev/null +++ b/docs/src/gpu_files/figure-commonmark/cell-8-output-1.svg @@ -0,0 +1,436 @@ + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + diff --git a/docs/src/index.md b/docs/src/index.md index 37a268e7..9df76a07 100644 --- a/docs/src/index.md +++ b/docs/src/index.md @@ -1,4 +1,6 @@ -```@meta + + +``` @meta CurrentModule = CounterfactualTraining ``` @@ -6,5 +8,44 @@ CurrentModule = CounterfactualTraining Documentation for [CounterfactualTraining](https://github.com/JuliaTrustworthyAI/CounterfactualTraining.jl). -```@index +# CounterfactualTraining.jl + +Teaching model plausible and actionable explanations. + +[![Stable](https://img.shields.io/badge/docs-stable-blue.svg)](https://JuliaTrustworthyAI.github.io/CounterfactualTraining.jl/stable/) +[![Dev](https://img.shields.io/badge/docs-dev-blue.svg)](https://JuliaTrustworthyAI.github.io/CounterfactualTraining.jl/dev/) +[![Build Status](https://github.com/JuliaTrustworthyAI/CounterfactualTraining.jl/actions/workflows/CI.yml/badge.svg?branch=main)](https://github.com/JuliaTrustworthyAI/CounterfactualTraining.jl/actions/workflows/CI.yml?query=branch%3Amain) +[![Coverage](https://codecov.io/gh/JuliaTrustworthyAI/CounterfactualTraining.jl/branch/main/graph/badge.svg)](https://codecov.io/gh/JuliaTrustworthyAI/CounterfactualTraining.jl) +[![Code Style: Blue](https://img.shields.io/badge/code%20style-blue-4495d1.svg)](https://github.com/invenia/BlueStyle) +[![Aqua QA](https://raw.githubusercontent.com/JuliaTesting/Aqua.jl/master/badge.svg)](https://github.com/JuliaTesting/Aqua.jl) + +## TL;DR + +CounterfactualTraining.jl can be used to train artificial neural networks that are inherently more explainable and robust than contenionally trained models. + +## Paper + +The package was initially developed as part of our [IEEE SaTML 2026](https://satml.org/2026/) paper *Counterfactual Training: Teaching Models Plausible and Actionable Explanations*: open the [preprint](https://arxiv.org/pdf/2601.16205). + +**Abstract**: We propose a novel training regime termed counterfactual training that leverages counterfactual explanations to increase the explanatory capacity of models. Counterfactual explanations have emerged as a popular post-hoc explanation method for opaque machine learning models: they inform how factual inputs would need to change in order for a model to produce some desired output. To be useful in real-world decision-making systems, counterfactuals should be plausible with respect to the underlying data and actionable with respect to the feature mutability constraints. Much existing research has therefore focused on developing post-hoc methods to generate counterfactuals that meet these desiderata. In this work, we instead hold models directly accountable for the desired end goal: counterfactual training employs counterfactuals during the training phase to minimize the divergence between learned representations and plausible, actionable explanations. We demonstrate empirically and theoretically that our proposed method facilitates training models that deliver inherently desirable counterfactual explanations and additionally exhibit improved adversarial robustness. + +## Package + +The package provides GPU-friendly training routines for models trained in Flux.jl. + +### Installation + +The package can be installed as follows: + +``` julia +using Pkg +Pkg.add("CounterfactualTraining") ``` + +### Introduction + +A self-contained introductory example demonstrating counterfactual training with mutability protection is available on the [Introduction](introduction.md) page. + +### Training on GPU + +The package provides two training approaches: a branch using [CounterfactualExplanations.jl](https://github.com/JuliaTrustworthyAI/CounterfactualExplanations.jl) as the backend for generating counterfactuals and a performance-focused `Native` submodule with GPU support. The former approach was developed during the initial research stage of the project. See [Training on GPU](gpu.md) for a full walkthrough with MNIST. diff --git a/docs/src/introduction.md b/docs/src/introduction.md new file mode 100644 index 00000000..0e0eae91 --- /dev/null +++ b/docs/src/introduction.md @@ -0,0 +1,110 @@ +# Introduction + + +This page demonstrates the core functionality of CounterfactualTraining.jl: training models that produce more plausible and actionable counterfactual explanations by holding them directly accountable during training. + +## Mutability Protection + +When generating counterfactual explanations, some features may be **immutable** (e.g., age, race) while others are **mutable** (e.g., debt, income). Without protection, the contrastive divergence penalty can inadvertently make models *more* sensitive to immutable features, which is undesirable. + +The key insight is that the contrastive divergence penalty gradient with respect to coefficient θ\_{y⁺,d} is: + +∂/∂θ\_{y⁺,d} div(x⁺, x′, y; θ) = x′\_d − x⁺\_d + +Since immutable features tend to differ between classes, the penalty would exacerbate the model’s existing sensitivity to those features. By **protecting** (zeroing gradients on) immutable features during counterfactual generation, the model learns to be relatively *more* sensitive to mutable features — producing counterfactuals that are actionable rather than relying on changes to features that cannot actually be changed. + +Domain constraints are also applied to keep counterfactuals within plausible ranges for each feature. + +## Setup + +``` julia +using CounterfactualTraining +using CounterfactualTraining.Native +using CounterfactualExplanations +using Flux +using Plots +using Random + +Random.seed!(42) +gr() +``` + +## Data + +We create a 2D synthetic dataset with two Gaussian clusters. Feature 1 (x-axis) is **mutable** (`:both`), while feature 2 (y-axis) is **immutable** (`:none`). + +``` julia +N = 150 +X1 = randn(Float32, 2, N) # Class 1: centered at origin +X2 = randn(Float32, 2, N) .+ 3.0f0 # Class 2: shifted by +3 +X = hcat(X1, X2) +y = vcat(fill(1, N), fill(2, N)) + +domain = [(-3.0f0, 6.0f0), (-3.0f0, 6.0f0)] +mutability = [:both, :none] +data = CounterfactualData(X, y; domain=domain, mutability=mutability) + +y_onehot = Flux.onehotbatch(y, 1:2) +train_set = Flux.DataLoader((X, y_onehot); batchsize=32, shuffle=true) +``` + +## Training + +We train two identical models with different objectives: + +- **Vanilla objective** (`VanillaObjective`): standard cross-entropy training, no counterfactual penalty. +- **Full objective** (`FullObjective`): counterfactual training with energy differential, regularization, and adversarial loss. + +Both use `NativeGenerator` for fast, batched counterfactual generation, and pass `mutability=[:both, :none]` to protect the immutable feature. + +``` julia +generator = NativeGenerator() + +# Model A: Vanilla (no CF penalty) +model_vanilla = Chain(Dense(2, 8, relu), Dense(8, 2)) +opt_vanilla = Flux.setup(Flux.Adam(1e-2), model_vanilla) +model_vanilla, log_vanilla = counterfactual_training( + VanillaObjective(needs_ce=true), model_vanilla, generator, train_set, opt_vanilla; + nepochs=20, maxiter=10, burnin=0.0f0, mutability=mutability, domain=domain, verbose=0 +) + +# Model B: Full (with CF penalty) +model_full = Chain(Dense(2, 8, relu), Dense(8, 2)) +opt_full = Flux.setup(Flux.Adam(1e-2), model_full) +model_full, log_full = counterfactual_training( + FullObjective(), model_full, generator, train_set, opt_full; + nepochs=20, maxiter=10, burnin=0.0f0, mutability=mutability, domain=domain, verbose=0 +) +``` + +## Counterfactual Generation + +We generate counterfactuals for 20 samples from class 1, targeting class 2, using each trained model. + +``` julia +X_test = X[:, 1:20] +targets = fill(2, 20) + +cfs_vanilla, _, _, _ = generate_counterfactuals!( + model_vanilla, X_test, targets, data, generator +) +cfs_full, _, _, _ = generate_counterfactuals!( + model_full, X_test, targets, data, generator +) +``` + +## Results + +``` julia +scatter(X[1, 1:N], X[2, 1:N], label="Class 1", color=:blue, ms=3) +scatter!(X[1, N+1:end], X[2, N+1:end], label="Class 2", color=:red, ms=3) +scatter!(cfs_vanilla[1, :], cfs_vanilla[2, :], label="CF (Vanilla)", + color=:green, markershape=:utriangle, ms=5) +scatter!(cfs_full[1, :], cfs_full[2, :], label="CF (Full)", + color=:purple, markershape=:star5, ms=5) +plot!(xlabel="Feature 1 (mutable)", ylabel="Feature 2 (immutable)") +``` + +![](introduction_files/figure-commonmark/cell-6-output-1.svg) + +With the `FullObjective`, counterfactuals should move primarily along the **mutable** feature (x-axis), leaving the **immutable** feature (y-axis) relatively unchanged. The model has learned to be less sensitive to the immutable feature, producing counterfactuals that are more actionable. diff --git a/docs/src/introduction.qmd b/docs/src/introduction.qmd new file mode 100644 index 00000000..e60ff9b8 --- /dev/null +++ b/docs/src/introduction.qmd @@ -0,0 +1,118 @@ +# Introduction + +This page demonstrates the core functionality of CounterfactualTraining.jl: training models that produce more plausible and actionable counterfactual explanations by holding them directly accountable during training. + +## Mutability Protection + +When generating counterfactual explanations, some features may be **immutable** (e.g., age, race) while others are **mutable** (e.g., debt, income). Without protection, the contrastive divergence penalty can inadvertently make models *more* sensitive to immutable features, which is undesirable. + +The key insight is that the contrastive divergence penalty gradient with respect to coefficient θ_{y⁺,d} is: + +∂/∂θ_{y⁺,d} div(x⁺, x′, y; θ) = x′_d − x⁺_d + +Since immutable features tend to differ between classes, the penalty would exacerbate the model's existing sensitivity to those features. By **protecting** (zeroing gradients on) immutable features during counterfactual generation, the model learns to be relatively *more* sensitive to mutable features — producing counterfactuals that are actionable rather than relying on changes to features that cannot actually be changed. + +Domain constraints are also applied to keep counterfactuals within plausible ranges for each feature. + +## Setup + +```{julia} +#| echo: true + +using CounterfactualTraining +using CounterfactualTraining.Native +using CounterfactualExplanations +using Flux +using Plots +using Random + +Random.seed!(42) +gr() +``` + +## Data + +We create a 2D synthetic dataset with two Gaussian clusters. Feature 1 (x-axis) is **mutable** (`:both`), while feature 2 (y-axis) is **immutable** (`:none`). + +```{julia} +#| echo: true + +N = 150 +X1 = randn(Float32, 2, N) # Class 1: centered at origin +X2 = randn(Float32, 2, N) .+ 3.0f0 # Class 2: shifted by +3 +X = hcat(X1, X2) +y = vcat(fill(1, N), fill(2, N)) + +domain = [(-3.0f0, 6.0f0), (-3.0f0, 6.0f0)] +mutability = [:both, :none] +data = CounterfactualData(X, y; domain=domain, mutability=mutability) + +y_onehot = Flux.onehotbatch(y, 1:2) +train_set = Flux.DataLoader((X, y_onehot); batchsize=32, shuffle=true) +``` + +## Training + +We train two identical models with different objectives: + +- **Vanilla objective** (`VanillaObjective`): standard cross-entropy training, no counterfactual penalty. +- **Full objective** (`FullObjective`): counterfactual training with energy differential, regularization, and adversarial loss. + +Both use `NativeGenerator` for fast, batched counterfactual generation, and pass `mutability=[:both, :none]` to protect the immutable feature. + +```{julia} +#| echo: true + +generator = NativeGenerator() + +# Model A: Vanilla (no CF penalty) +model_vanilla = Chain(Dense(2, 8, relu), Dense(8, 2)) +opt_vanilla = Flux.setup(Flux.Adam(1e-2), model_vanilla) +model_vanilla, log_vanilla = counterfactual_training( + VanillaObjective(needs_ce=true), model_vanilla, generator, train_set, opt_vanilla; + nepochs=20, maxiter=10, burnin=0.0f0, mutability=mutability, domain=domain, verbose=0 +) + +# Model B: Full (with CF penalty) +model_full = Chain(Dense(2, 8, relu), Dense(8, 2)) +opt_full = Flux.setup(Flux.Adam(1e-2), model_full) +model_full, log_full = counterfactual_training( + FullObjective(), model_full, generator, train_set, opt_full; + nepochs=20, maxiter=10, burnin=0.0f0, mutability=mutability, domain=domain, verbose=0 +) +``` + +## Counterfactual Generation + +We generate counterfactuals for 20 samples from class 1, targeting class 2, using each trained model. + +```{julia} +#| echo: true + +X_test = X[:, 1:20] +targets = fill(2, 20) + +cfs_vanilla, _, _, _ = generate_counterfactuals!( + model_vanilla, X_test, targets, data, generator +) +cfs_full, _, _, _ = generate_counterfactuals!( + model_full, X_test, targets, data, generator +) +``` + +## Results + +```{julia} +#| echo: true +#| output: true + +scatter(X[1, 1:N], X[2, 1:N], label="Class 1", color=:blue, ms=3) +scatter!(X[1, N+1:end], X[2, N+1:end], label="Class 2", color=:red, ms=3) +scatter!(cfs_vanilla[1, :], cfs_vanilla[2, :], label="CF (Vanilla)", + color=:green, markershape=:utriangle, ms=5) +scatter!(cfs_full[1, :], cfs_full[2, :], label="CF (Full)", + color=:purple, markershape=:star5, ms=5) +plot!(xlabel="Feature 1 (mutable)", ylabel="Feature 2 (immutable)") +``` + +With the `FullObjective`, counterfactuals should move primarily along the **mutable** feature (x-axis), leaving the **immutable** feature (y-axis) relatively unchanged. The model has learned to be less sensitive to the immutable feature, producing counterfactuals that are more actionable. diff --git a/docs/src/introduction_files/figure-commonmark/cell-6-output-1.svg b/docs/src/introduction_files/figure-commonmark/cell-6-output-1.svg new file mode 100644 index 00000000..b91d3579 --- /dev/null +++ b/docs/src/introduction_files/figure-commonmark/cell-6-output-1.svg @@ -0,0 +1,383 @@ + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + From ebbc7e26e3a13158e29e9eabc8dcefa2836961c7 Mon Sep 17 00:00:00 2001 From: pat-alt Date: Mon, 10 Aug 2026 10:18:09 +0200 Subject: [PATCH 2/4] close test gaps --- .github/workflows/CI.yml | 2 +- .gitignore | 2 + Project.toml | 4 ++ test/Project.toml | 2 + test/counterfactuals_tests.jl | 35 ++++++++++ test/loss_tests.jl | 58 +++++++++++++++++ test/native_edge_tests.jl | 118 ++++++++++++++++++++++++++++++++++ test/native_helpers_tests.jl | 85 ++++++++++++++++++++++++ test/objectives_tests.jl | 55 ++++++++++++++++ test/runtests.jl | 10 +++ test/training_tests.jl | 93 +++++++++++++++++++++++++++ test/utils_tests.jl | 37 +++++++++++ 12 files changed, 500 insertions(+), 1 deletion(-) create mode 100644 test/counterfactuals_tests.jl create mode 100644 test/loss_tests.jl create mode 100644 test/native_edge_tests.jl create mode 100644 test/native_helpers_tests.jl create mode 100644 test/objectives_tests.jl create mode 100644 test/training_tests.jl create mode 100644 test/utils_tests.jl diff --git a/.github/workflows/CI.yml b/.github/workflows/CI.yml index dd28615d..5c65251e 100644 --- a/.github/workflows/CI.yml +++ b/.github/workflows/CI.yml @@ -66,7 +66,7 @@ jobs: cache-name: cache-quarto with: path: _freeze - key: ${{ runner.os }}-${{ env.cache-name }}-${{ hashFiles('*.qmd') }} + key: ${{ runner.os }}-${{ env.cache-name }}-${{ hashFiles('**/*.qmd') }} restore-keys: | ${{ runner.os }}-${{ env.cache-name }}- - name: Configure doc environment diff --git a/.gitignore b/.gitignore index a937e136..a0b8650d 100644 --- a/.gitignore +++ b/.gitignore @@ -11,3 +11,5 @@ _freeze/ **.DS_Store + +**/*.quarto_ipynb diff --git a/Project.toml b/Project.toml index 349c606f..63cdbaf2 100644 --- a/Project.toml +++ b/Project.toml @@ -24,6 +24,8 @@ EnergySamplers = "1.0" Flux = "0.14, 0.15, 0.16" JLD2 = "0.4, 0.5" LinearAlgebra = "1" +MLDatasets = "0.7.21" +Plots = "1.41.6" ProgressMeter = "1.10.2" Random = "1" StatsBase = "0.34" @@ -34,6 +36,8 @@ julia = "1.10" [extras] Aqua = "4c88cf16-eb10-579e-8560-4a9242c79595" +MLDatasets = "eb30cadb-4394-5ae3-aed4-317e484a6458" +Plots = "91a5bcdd-55d7-5caf-9e0b-520d859cae80" Test = "8dfed614-e22c-5e08-85e1-65c5234f0b40" [targets] diff --git a/test/Project.toml b/test/Project.toml index 6313b18b..ffb4b187 100644 --- a/test/Project.toml +++ b/test/Project.toml @@ -3,6 +3,8 @@ Aqua = "4c88cf16-eb10-579e-8560-4a9242c79595" CounterfactualExplanations = "2f13d31b-18db-44c1-bc43-ebaf2cff0be0" Flux = "587475ba-b771-5e3f-ad9e-33799f191a9c" LinearAlgebra = "37e2e46d-f89d-539d-b4ee-838fcccc9c8e" +Random = "9a3f8284-a2c9-5f02-9a11-845980a1fd5c" +Statistics = "10745b16-79ce-11e8-11f9-7d13ad32a3b2" Test = "8dfed614-e22c-5e08-85e1-65c5234f0b40" [compat] diff --git a/test/counterfactuals_tests.jl b/test/counterfactuals_tests.jl new file mode 100644 index 00000000..7651d172 --- /dev/null +++ b/test/counterfactuals_tests.jl @@ -0,0 +1,35 @@ +@testset "Counterfactual generation (non-Native)" begin + # protect_immutable! (non-Native) - operates on Vector{Matrix} + # mutability is per-feature-row: [:both, :none, :increase, :decrease] + samples = [rand(Float32, 4, 1) for _ in 1:3] + cfs = [fill(2.0f0, 4, 1) for _ in 1:3] + mut = [:both, :none, :increase, :decrease] + CounterfactualTraining.protect_immutable!(samples, cfs, mut) + for i in 1:3 + @test all(samples[i][1, :] .>= 0.0f0) # :both -> keep sample (random) + @test all(samples[i][2, :] .== 2.0f0) # :none -> use cf + @test all(samples[i][3, :] .>= 2.0f0) # :increase -> max(cf=2, s) >= 2 + @test all(samples[i][4, :] .<= 2.0f0) # :decrease -> min(cf=2, s) <= 2 + end + + # mutability=nothing early return + samples2 = [rand(Float32, 4, 1) for _ in 1:3] + original = copy(samples2[1]) + CounterfactualTraining.protect_immutable!(samples2, cfs, nothing) + @test samples2[1] == original + + # isvalid + X = hcat(randn(Float32, 2, 30), randn(Float32, 2, 30) .+ 3.0f0) + y = vcat(fill(1, 30), fill(2, 30)) + train_set = Flux.DataLoader((X, Flux.onehotbatch(y, 1:2)); batchsize=32) + model = Chain(Dense(2, 4, relu), Dense(4, 2)) + gen_ce = GenericGenerator() + dl, pct, ces = CounterfactualTraining.generate!( + model, train_set, gen_ce; nsamples=5 + ) + @test length(dl) == length(train_set) + @test 0.0 <= pct <= 1.0 + data = CounterfactualData(X, y) + v = CounterfactualTraining.isvalid(ces[1], model, data) + @test v isa Bool +end diff --git a/test/loss_tests.jl b/test/loss_tests.jl new file mode 100644 index 00000000..c987dce4 --- /dev/null +++ b/test/loss_tests.jl @@ -0,0 +1,58 @@ +@testset "Loss functions" begin + m = Chain(Dense(3, 2)) + N = 5 + C = 2 + targets = zeros(Float32, C, N) + for i in 1:N + targets[rand(1:C), i] = 1.0f0 + end + samples = randn(Float32, 3, N) + counterfactual = randn(Float32, 3, N) + + # implausibility + impl = implausibility(m, counterfactual, samples, targets) + @test length(impl) == N + @test all(isfinite, impl) + + # reg_loss + reg = reg_loss(m, counterfactual, samples, targets) + @test length(reg) == N + @test all(>=(0.0f0), reg) + + # adv_loss - branch with valid adversarial examples + perturbations_small = zeros(Float32, 3, N) + adv1 = CounterfactualTraining.adv_loss(m, counterfactual, perturbations_small, targets) + @test isfinite(adv1) + + # adv_loss - branch with no valid adversarial examples + perturbations_large = fill(10.0f0, 3, N) + adv2 = CounterfactualTraining.adv_loss( + m, counterfactual, perturbations_large, targets; epsilon=0.5f0 + ) + @test adv2 == 0.0f0 + + # NormBound struct + nb = CounterfactualTraining.NormBound(epsilon=0.3, p=Inf) + @test nb.epsilon == 0.3 + @test nb.p == Inf + nb2 = CounterfactualTraining.NormBound(; epsilon=1.0, p=2) + @test nb2.epsilon == 1.0 + @test nb2.p == 2 + + # NormBound callable + @test nb(zeros(Float32, 3)) == true + @test nb(fill(10.0f0, 3)) == false + + # isadvexm + @test CounterfactualTraining.isadvexm(zeros(Float32, 3), 0.5, Inf) == true + @test CounterfactualTraining.isadvexm(ones(Float32, 3), 0.5, Inf) == false + @test CounterfactualTraining.isadvexm(zeros(Float32, 3), 0.5, 2) == true + @test CounterfactualTraining.isadvexm(fill(1.0f0, 3), 0.5, 2) == false + + # get/set global AE criterium + original = CounterfactualTraining.get_global_ae_criterium() + custom = CounterfactualTraining.NormBound(epsilon=1.0, p=2) + CounterfactualTraining.set_global_ae_criterium(custom) + @test CounterfactualTraining.get_global_ae_criterium() === custom + CounterfactualTraining.set_global_ae_criterium(original) +end diff --git a/test/native_edge_tests.jl b/test/native_edge_tests.jl new file mode 100644 index 00000000..0f0049c0 --- /dev/null +++ b/test/native_edge_tests.jl @@ -0,0 +1,118 @@ +@testset "generate_native! edge cases" begin + X = randn(Float32, 10, 100) + y = vcat(fill(1, 50), fill(2, 50)) + train_set = Flux.DataLoader( + (Float32.(X), Flux.onehotbatch(y, 1:2)); batchsize=32 + ) + model = Chain(Dense(10, 8, relu), Dense(8, 2)) + gen = NativeGenerator() + + # nsamples < N subsampling branch + dl, pct_valid, _ = generate_native!( + model, train_set, gen; nsamples=20, maxiter=5, verbose=0 + ) + @test length(dl) == length(train_set) + @test 0 ≤ pct_valid ≤ 1 + + # nsamples < length(train_set) warning+clamp branch + dl2, pct_valid2, _ = generate_native!( + model, train_set, gen; nsamples=2, maxiter=5, verbose=0 + ) + @test length(dl2) == length(train_set) + @test 0 ≤ pct_valid2 ≤ 1 +end + +@testset "Native training edge cases" begin + X = hcat(randn(Float32, 2, 100), randn(Float32, 2, 100) .+ 3.0f0) + y = vcat(fill(1, 100), fill(2, 100)) + train_set = Flux.DataLoader((X, Flux.onehotbatch(y, 1:2)); batchsize=32) + model = Chain(Dense(2, 8, relu), Dense(8, 2)) + opt_state = Flux.setup(Flux.Adam(1e-2), model) + gen = NativeGenerator() + obj = VanillaObjective(; needs_ce=true) + + # burnin > 0 + model_b, log_b = counterfactual_training( + obj, + model, + gen, + train_set, + opt_state; + nepochs=4, + verbose=0, + maxiter=10, + burnin=0.5f0, + ) + @test length(log_b) == 4 + @test log_b[1].implaus === nothing + @test isfinite(log_b[end].implaus) + + # val_set + val_set = Flux.DataLoader( + (X[:, 1:20], Flux.onehotbatch(y[1:20], 1:2)); batchsize=10 + ) + model_v = Chain(Dense(2, 8, relu), Dense(8, 2)) + opt_v = Flux.setup(Flux.Adam(1e-2), model_v) + model_v, log_v = counterfactual_training( + obj, + model_v, + gen, + train_set, + opt_v; + nepochs=2, + verbose=0, + maxiter=10, + burnin=0.0f0, + val_set=val_set, + ) + @test isfinite(log_v[end].acc_val) + + # checkpoint_dir - save then load + ckpt = mktempdir() + model_c = Chain(Dense(2, 8, relu), Dense(8, 2)) + opt_c = Flux.setup(Flux.Adam(1e-2), model_c) + model_c, log_c = counterfactual_training( + obj, + model_c, + gen, + train_set, + opt_c; + nepochs=2, + verbose=0, + maxiter=10, + burnin=0.0f0, + checkpoint_dir=ckpt, + ) + @test isfile(joinpath(ckpt, "checkpoint.jld2")) + + # re-run to exercise checkpoint load path (training already completed) + model_c2 = Chain(Dense(2, 8, relu), Dense(8, 2)) + opt_c2 = Flux.setup(Flux.Adam(1e-2), model_c2) + model_c2, log_c2 = counterfactual_training( + obj, + model_c2, + gen, + train_set, + opt_c2; + nepochs=2, + verbose=0, + maxiter=10, + burnin=0.0f0, + checkpoint_dir=ckpt, + ) + @test length(log_c2) == 0 + + # verbose=0 + model_v0 = Chain(Dense(2, 8, relu), Dense(8, 2)) + opt_v0 = Flux.setup(Flux.Adam(1e-2), model_v0) + counterfactual_training( + obj, model_v0, gen, train_set, opt_v0; nepochs=1, verbose=0, maxiter=5 + ) + + # verbose=3 + model_v3 = Chain(Dense(2, 8, relu), Dense(8, 2)) + opt_v3 = Flux.setup(Flux.Adam(1e-2), model_v3) + counterfactual_training( + obj, model_v3, gen, train_set, opt_v3; nepochs=1, verbose=3, maxiter=5 + ) +end diff --git a/test/native_helpers_tests.jl b/test/native_helpers_tests.jl new file mode 100644 index 00000000..002156e9 --- /dev/null +++ b/test/native_helpers_tests.jl @@ -0,0 +1,85 @@ +@testset "Native helpers (coverage)" begin + # check_batched_convergence - maxiter branch + probs = rand(Float32, 2, 5) + @test Native.check_batched_convergence(probs, [1, 2, 1, 2, 1], 10, 10, 0.75f0) == + trues(5) + + # check_batched_convergence - threshold branch + probs2 = Float32[ + 0.9 0.5 0.9 0.5 0.9 + 0.1 0.5 0.1 0.5 0.1 + ] + conv = Native.check_batched_convergence(probs2, fill(1, 5), 1, 10, 0.75f0) + @test conv == [true, false, true, false, true] + + # track_adversarial_examples! - p == Inf branch + X_ae = zeros(Float32, 2, 3) + Xp_ae = zeros(Float32, 2, 3) + Xp_ae[:, 2] .= 1.0f0 + last_valid = copy(X_ae) + Native.track_adversarial_examples!(last_valid, X_ae, Xp_ae, 0.5f0, Inf) + @test last_valid[:, 1] == Xp_ae[:, 1] + @test last_valid[:, 2] == X_ae[:, 2] + @test last_valid[:, 3] == Xp_ae[:, 3] + + # track_adversarial_examples! - finite p branch + # NOTE: native/training.jl:150 has a bug: `abs .^ p` tries to raise + # the function `abs` to a power instead of `abs.(perturbations) .^ p` + @test_broken begin + last_valid2 = copy(X_ae) + Native.track_adversarial_examples!(last_valid2, X_ae, Xp_ae, 0.5f0, 2) + last_valid2[:, 1] == Xp_ae[:, 1] + end + + # generator_loss + gen = NativeGenerator() + model = Chain(Dense(3, 5, relu), Dense(5, 2)) + X_gl = randn(Float32, 3, 8) + Xp_gl = copy(X_gl) + targets_oh = Flux.onehotbatch(fill(2, 8), 1:2) + loss_val = Native.generator_loss( + gen, model, Xp_gl, X_gl, targets_oh, fill(2, 8), 1, 1.0f-3, 0.9f0, 10 + ) + @test isfinite(loss_val) + gen_big = NativeGenerator(; λ=[1.0f0, 1.0f0]) + loss_big = Native.generator_loss( + gen_big, model, Xp_gl, X_gl, targets_oh, fill(2, 8), 1, 1.0f-3, 0.9f0, 10 + ) + @test isfinite(loss_big) + + # find_neighbours - normal case + X_nb = randn(Float32, 3, 6) + y_nb = [1, 1, 2, 2, 1, 2] + rng = MersenneTwister(42) + nb = Native.find_neighbours(X_nb, y_nb, [2, 1], [1, 2]; rng=rng) + @test size(nb) == (3, 2) + @test all(isfinite, nb) + + # find_neighbours - empty candidates branch + rng2 = MersenneTwister(42) + nb2 = Native.find_neighbours(X_nb, y_nb, [3], [1, 2]; rng=rng2) + @test size(nb2) == (3, 1) + @test all(isfinite, nb2) + + # protect_immutable! (Native) - all directions + nothing + neigh = ones(Float32, 4, 3) + cfs = fill(2.0f0, 4, 3) + mut = [:both, :none, :increase, :decrease] + Native.protect_immutable!(neigh, cfs, mut) + @test all(neigh[1, :] .== 1.0f0) # :both -> keep neighbour + @test all(neigh[2, :] .== 2.0f0) # :none -> use counterfactual + @test all(neigh[3, :] .== 2.0f0) # :increase -> max(cf=2,neigh=1)=2 + @test all(neigh[4, :] .== 1.0f0) # :decrease -> min(cf=2,neigh=1)=1 + # mutability=nothing early return + neigh2 = ones(Float32, 4, 3) + Native.protect_immutable!(neigh2, cfs, nothing) + @test all(neigh2 .== 1.0f0) + + # split_obs + parts = Native.split_obs(1:10, 3) + @test length(parts) >= 3 + @test sort(reduce(vcat, parts)) == collect(1:10) + parts2 = Native.split_obs(1:5, 10) + @test length(parts2) >= 5 + @test sort(reduce(vcat, parts2)) == collect(1:5) +end diff --git a/test/objectives_tests.jl b/test/objectives_tests.jl new file mode 100644 index 00000000..b5b42fda --- /dev/null +++ b/test/objectives_tests.jl @@ -0,0 +1,55 @@ +struct FooObj <: CounterfactualTraining.AbstractObjective end + +@testset "Objectives" begin + C = 2 + N = 8 + yhat = randn(Float32, C, N) + y = Flux.onehotbatch(rand(1:C, N), 1:C) + ed = rand(Float32, N) + reg = rand(Float32, N) + adv = rand(Float32, N) + ce = Flux.Losses.logitcrossentropy + + # VanillaObjective + @test_throws AssertionError VanillaObjective(ce, Float32[1.0, 2.0]) + obj_v = VanillaObjective(; needs_ce=true) + @test CounterfactualTraining.needs_counterfactuals(obj_v) == true + obj_v2 = VanillaObjective(; needs_ce=false) + @test CounterfactualTraining.needs_counterfactuals(obj_v2) == false + rv = obj_v(yhat, y) + @test isfinite(rv) + @test rv ≈ ce(yhat, y; agg=mean) * 1.0 + + # needs_counterfactuals fallback for custom subtype + @test CounterfactualTraining.needs_counterfactuals(FooObj()) == true + + # EnergyDifferentialObjective - inner constructor with short lambda + # NOTE: objectives.jl:99 has a bug (references undefined `obj`), so + # constructing with length<3 lambda throws UndefVarError + @test_throws UndefVarError EnergyDifferentialObjective(ce, Float32[1.0, 0.5]) + @test_throws UndefVarError EnergyDifferentialObjective(ce, Float32[1.0]) + obj_e = EnergyDifferentialObjective() + re = obj_e(yhat, y, ed, reg) + @test isfinite(re) + @test re ≈ ce(yhat, y; agg=mean) * obj_e.lambda[1] + + mean(Float32.(ed)) * obj_e.lambda[2] + + mean(Float32.(reg)) * obj_e.lambda[3] + + # AdversarialObjective + @test_throws AssertionError AdversarialObjective(ce, Float32[1.0]) + obj_a = AdversarialObjective() + ra = obj_a(yhat, y, Float32[], Float32[], adv) + @test isfinite(ra) + @test ra ≈ ce(yhat, y; agg=mean) * obj_a.lambda[1] + + mean(Float32.(adv)) * obj_a.lambda[2] + + # FullObjective + @test_throws AssertionError FullObjective(ce, Float32[1.0, 2.0]) + obj_f = FullObjective() + rf = obj_f(yhat, y, ed, reg, adv) + @test isfinite(rf) + @test rf ≈ ce(yhat, y; agg=mean) * obj_f.lambda[1] + + mean(Float32.(ed)) * obj_f.lambda[2] + + mean(Float32.(reg)) * obj_f.lambda[3] + + mean(Float32.(adv)) * obj_f.lambda[4] +end diff --git a/test/runtests.jl b/test/runtests.jl index ba036b07..a651dc81 100644 --- a/test/runtests.jl +++ b/test/runtests.jl @@ -5,6 +5,8 @@ using Aqua using CounterfactualExplanations using Flux using LinearAlgebra +using Random +using Statistics @testset "CounterfactualTraining.jl" begin @testset "Code quality (Aqua.jl)" begin @@ -218,4 +220,12 @@ using LinearAlgebra @test typeof(model.layers[1].weight) != Matrix{Float32} end end + + include("loss_tests.jl") + include("utils_tests.jl") + include("objectives_tests.jl") + include("native_helpers_tests.jl") + include("native_edge_tests.jl") + include("counterfactuals_tests.jl") + include("training_tests.jl") end diff --git a/test/training_tests.jl b/test/training_tests.jl new file mode 100644 index 00000000..45cc9b05 --- /dev/null +++ b/test/training_tests.jl @@ -0,0 +1,93 @@ +@testset "Counterfactual training (non-Native)" begin + X = hcat(randn(Float32, 2, 60), randn(Float32, 2, 60) .+ 3.0f0) + y = vcat(fill(1, 60), fill(2, 60)) + train_set = Flux.DataLoader((X, Flux.onehotbatch(y, 1:2)); batchsize=32) + model = Chain(Dense(2, 8, relu), Dense(8, 2)) + opt_state = Flux.setup(Flux.Adam(1e-2), model) + gen_ce = GenericGenerator() + obj = VanillaObjective(; needs_ce=true) + + # basic 2-epoch run + model, log = counterfactual_training( + obj, model, gen_ce, train_set, opt_state; nepochs=2, verbose=0 + ) + @test length(log) == 2 + @test isfinite(log[end].acc) + @test isfinite(log[end].train_loss) + @test 0.0 <= log[end].percent_valid <= 1.0 + + # burnin > 0 + model_b = Chain(Dense(2, 8, relu), Dense(8, 2)) + opt_b = Flux.setup(Flux.Adam(1e-2), model_b) + model_b, log_b = counterfactual_training( + obj, model_b, gen_ce, train_set, opt_b; nepochs=3, verbose=0, burnin=0.5 + ) + @test length(log_b) == 3 + @test log_b[1].implaus === nothing + @test isfinite(log_b[end].implaus) + + # val_set + val_set = Flux.DataLoader( + (X[:, 1:20], Flux.onehotbatch(y[1:20], 1:2)); batchsize=10 + ) + model_v = Chain(Dense(2, 8, relu), Dense(8, 2)) + opt_v = Flux.setup(Flux.Adam(1e-2), model_v) + model_v, log_v = counterfactual_training( + obj, + model_v, + gen_ce, + train_set, + opt_v; + nepochs=2, + verbose=0, + val_set=val_set, + ) + @test isfinite(log_v[end].acc_val) + + # callback (called post-burnin when ces !== nothing) + calls = Int[] + model_cb = Chain(Dense(2, 8, relu), Dense(8, 2)) + opt_cb = Flux.setup(Flux.Adam(1e-2), model_cb) + model_cb, log_cb = counterfactual_training( + obj, + model_cb, + gen_ce, + train_set, + opt_cb; + nepochs=2, + verbose=0, + callback=(m, ces) -> push!(calls, 1), + ) + @test length(calls) >= 1 + + # checkpoint_dir - save then load + ckpt = mktempdir() + model_c = Chain(Dense(2, 8, relu), Dense(8, 2)) + opt_c = Flux.setup(Flux.Adam(1e-2), model_c) + model_c, log_c = counterfactual_training( + obj, + model_c, + gen_ce, + train_set, + opt_c; + nepochs=2, + verbose=0, + checkpoint_dir=ckpt, + ) + @test isfile(joinpath(ckpt, "checkpoint.jld2")) + + # checkpoint load path (training already completed) + model_c2 = Chain(Dense(2, 8, relu), Dense(8, 2)) + opt_c2 = Flux.setup(Flux.Adam(1e-2), model_c2) + model_c2, log_c2 = counterfactual_training( + obj, + model_c2, + gen_ce, + train_set, + opt_c2; + nepochs=2, + verbose=0, + checkpoint_dir=ckpt, + ) + @test length(log_c2) >= 0 +end diff --git a/test/utils_tests.jl b/test/utils_tests.jl new file mode 100644 index 00000000..5c96a0c7 --- /dev/null +++ b/test/utils_tests.jl @@ -0,0 +1,37 @@ +@testset "Utils" begin + # unwrap + X = randn(Float32, 3, 20) + y = vcat(fill(1, 10), fill(2, 10)) + train_set = Flux.DataLoader((X, Flux.onehotbatch(y, 1:2)); batchsize=10) + X2, ycold = CounterfactualTraining.unwrap(train_set) + @test size(X2) == (3, 20) + @test sort(unique(ycold)) == [1, 2] + + # unwrap with labels (BUG: utils.jl:23 - ycold is Vector{Int}, can't hold Strings) + @test_broken begin + X3, ycold_labeled = CounterfactualTraining.unwrap(train_set; labels=["a", "b"]) + sort(unique(ycold_labeled)) == ["a", "b"] + end + + # accuracy + X_acc = hcat(randn(Float32, 2, 50), randn(Float32, 2, 50) .+ 3.0f0) + y_acc = vcat(fill(1, 50), fill(2, 50)) + train_set_acc = Flux.DataLoader( + (X_acc, Flux.onehotbatch(y_acc, 1:2)); batchsize=25 + ) + model_acc = Chain(Dense(2, 8, relu), Dense(8, 2)) + acc = CounterfactualTraining.accuracy(model_acc, train_set_acc) + @test 0.0 <= acc <= 1.0 + + # infer_domain_constraints + X_dom = randn(Float32, 4, 100) + bounds = CounterfactualTraining.infer_domain_constraints(X_dom) + @test length(bounds) == 4 + @test all(b -> b isa Tuple, bounds) + for i in 1:4 + lb, ub = bounds[i] + @test lb <= ub + @test lb <= minimum(X_dom[i, :]) + @test ub >= maximum(X_dom[i, :]) + end +end From 1c9152fca14ec0cd6411c5ee20bfb4e78fa39db0 Mon Sep 17 00:00:00 2001 From: pat-alt Date: Mon, 10 Aug 2026 10:22:22 +0200 Subject: [PATCH 3/4] changelog and bump version --- CHANGELOG.md | 7 +++++++ Project.toml | 2 +- 2 files changed, 8 insertions(+), 1 deletion(-) diff --git a/CHANGELOG.md b/CHANGELOG.md index 5fd7f37a..c845f1cc 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -4,6 +4,13 @@ All notable changes to this project will be documented in this file. The format is based on [Keep a Changelog](https://keepachangelog.com/en/1.1.0/), and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0.html). +## Version [0.2.1] - 2026-08-10 + +### Added + +- Adds two documentation sections. +- Adds more unit tests. + ## Version [0.2.0] - 2026-08-09 ### Changed diff --git a/Project.toml b/Project.toml index 63cdbaf2..646d3cc7 100644 --- a/Project.toml +++ b/Project.toml @@ -1,6 +1,6 @@ name = "CounterfactualTraining" uuid = "c5e5e904-ed0d-4ddb-967e-dbb7fd41a507" -version = "0.2.0" +version = "0.2.1" authors = ["Patrick Altmeyer"] [deps] From 9f677757103258283c9003dfd1778a15ec48c4db Mon Sep 17 00:00:00 2001 From: pat-alt Date: Mon, 10 Aug 2026 10:29:25 +0200 Subject: [PATCH 4/4] formatting --- test/counterfactuals_tests.jl | 4 +--- test/native_edge_tests.jl | 8 ++------ test/objectives_tests.jl | 10 ++++++---- test/training_tests.jl | 31 ++++--------------------------- test/utils_tests.jl | 4 +--- 5 files changed, 14 insertions(+), 43 deletions(-) diff --git a/test/counterfactuals_tests.jl b/test/counterfactuals_tests.jl index 7651d172..adfe1208 100644 --- a/test/counterfactuals_tests.jl +++ b/test/counterfactuals_tests.jl @@ -24,9 +24,7 @@ train_set = Flux.DataLoader((X, Flux.onehotbatch(y, 1:2)); batchsize=32) model = Chain(Dense(2, 4, relu), Dense(4, 2)) gen_ce = GenericGenerator() - dl, pct, ces = CounterfactualTraining.generate!( - model, train_set, gen_ce; nsamples=5 - ) + dl, pct, ces = CounterfactualTraining.generate!(model, train_set, gen_ce; nsamples=5) @test length(dl) == length(train_set) @test 0.0 <= pct <= 1.0 data = CounterfactualData(X, y) diff --git a/test/native_edge_tests.jl b/test/native_edge_tests.jl index 0f0049c0..d98fefb9 100644 --- a/test/native_edge_tests.jl +++ b/test/native_edge_tests.jl @@ -1,9 +1,7 @@ @testset "generate_native! edge cases" begin X = randn(Float32, 10, 100) y = vcat(fill(1, 50), fill(2, 50)) - train_set = Flux.DataLoader( - (Float32.(X), Flux.onehotbatch(y, 1:2)); batchsize=32 - ) + train_set = Flux.DataLoader((Float32.(X), Flux.onehotbatch(y, 1:2)); batchsize=32) model = Chain(Dense(10, 8, relu), Dense(8, 2)) gen = NativeGenerator() @@ -48,9 +46,7 @@ end @test isfinite(log_b[end].implaus) # val_set - val_set = Flux.DataLoader( - (X[:, 1:20], Flux.onehotbatch(y[1:20], 1:2)); batchsize=10 - ) + val_set = Flux.DataLoader((X[:, 1:20], Flux.onehotbatch(y[1:20], 1:2)); batchsize=10) model_v = Chain(Dense(2, 8, relu), Dense(8, 2)) opt_v = Flux.setup(Flux.Adam(1e-2), model_v) model_v, log_v = counterfactual_training( diff --git a/test/objectives_tests.jl b/test/objectives_tests.jl index b5b42fda..aeedb84d 100644 --- a/test/objectives_tests.jl +++ b/test/objectives_tests.jl @@ -31,7 +31,8 @@ struct FooObj <: CounterfactualTraining.AbstractObjective end obj_e = EnergyDifferentialObjective() re = obj_e(yhat, y, ed, reg) @test isfinite(re) - @test re ≈ ce(yhat, y; agg=mean) * obj_e.lambda[1] + + @test re ≈ + ce(yhat, y; agg=mean) * obj_e.lambda[1] + mean(Float32.(ed)) * obj_e.lambda[2] + mean(Float32.(reg)) * obj_e.lambda[3] @@ -40,15 +41,16 @@ struct FooObj <: CounterfactualTraining.AbstractObjective end obj_a = AdversarialObjective() ra = obj_a(yhat, y, Float32[], Float32[], adv) @test isfinite(ra) - @test ra ≈ ce(yhat, y; agg=mean) * obj_a.lambda[1] + - mean(Float32.(adv)) * obj_a.lambda[2] + @test ra ≈ + ce(yhat, y; agg=mean) * obj_a.lambda[1] + mean(Float32.(adv)) * obj_a.lambda[2] # FullObjective @test_throws AssertionError FullObjective(ce, Float32[1.0, 2.0]) obj_f = FullObjective() rf = obj_f(yhat, y, ed, reg, adv) @test isfinite(rf) - @test rf ≈ ce(yhat, y; agg=mean) * obj_f.lambda[1] + + @test rf ≈ + ce(yhat, y; agg=mean) * obj_f.lambda[1] + mean(Float32.(ed)) * obj_f.lambda[2] + mean(Float32.(reg)) * obj_f.lambda[3] + mean(Float32.(adv)) * obj_f.lambda[4] diff --git a/test/training_tests.jl b/test/training_tests.jl index 45cc9b05..bc0a1097 100644 --- a/test/training_tests.jl +++ b/test/training_tests.jl @@ -27,20 +27,11 @@ @test isfinite(log_b[end].implaus) # val_set - val_set = Flux.DataLoader( - (X[:, 1:20], Flux.onehotbatch(y[1:20], 1:2)); batchsize=10 - ) + val_set = Flux.DataLoader((X[:, 1:20], Flux.onehotbatch(y[1:20], 1:2)); batchsize=10) model_v = Chain(Dense(2, 8, relu), Dense(8, 2)) opt_v = Flux.setup(Flux.Adam(1e-2), model_v) model_v, log_v = counterfactual_training( - obj, - model_v, - gen_ce, - train_set, - opt_v; - nepochs=2, - verbose=0, - val_set=val_set, + obj, model_v, gen_ce, train_set, opt_v; nepochs=2, verbose=0, val_set=val_set ) @test isfinite(log_v[end].acc_val) @@ -65,14 +56,7 @@ model_c = Chain(Dense(2, 8, relu), Dense(8, 2)) opt_c = Flux.setup(Flux.Adam(1e-2), model_c) model_c, log_c = counterfactual_training( - obj, - model_c, - gen_ce, - train_set, - opt_c; - nepochs=2, - verbose=0, - checkpoint_dir=ckpt, + obj, model_c, gen_ce, train_set, opt_c; nepochs=2, verbose=0, checkpoint_dir=ckpt ) @test isfile(joinpath(ckpt, "checkpoint.jld2")) @@ -80,14 +64,7 @@ model_c2 = Chain(Dense(2, 8, relu), Dense(8, 2)) opt_c2 = Flux.setup(Flux.Adam(1e-2), model_c2) model_c2, log_c2 = counterfactual_training( - obj, - model_c2, - gen_ce, - train_set, - opt_c2; - nepochs=2, - verbose=0, - checkpoint_dir=ckpt, + obj, model_c2, gen_ce, train_set, opt_c2; nepochs=2, verbose=0, checkpoint_dir=ckpt ) @test length(log_c2) >= 0 end diff --git a/test/utils_tests.jl b/test/utils_tests.jl index 5c96a0c7..0ee1c870 100644 --- a/test/utils_tests.jl +++ b/test/utils_tests.jl @@ -16,9 +16,7 @@ # accuracy X_acc = hcat(randn(Float32, 2, 50), randn(Float32, 2, 50) .+ 3.0f0) y_acc = vcat(fill(1, 50), fill(2, 50)) - train_set_acc = Flux.DataLoader( - (X_acc, Flux.onehotbatch(y_acc, 1:2)); batchsize=25 - ) + train_set_acc = Flux.DataLoader((X_acc, Flux.onehotbatch(y_acc, 1:2)); batchsize=25) model_acc = Chain(Dense(2, 8, relu), Dense(8, 2)) acc = CounterfactualTraining.accuracy(model_acc, train_set_acc) @test 0.0 <= acc <= 1.0