101 get coverage up

.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

.gitignore

@@ -11,3 +11,5 @@ _freeze/
 **.DS_Store
 
 
+
+**/*.quarto_ipynb

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

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]
@@ -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]

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"

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(;

docs/src/_intro.qmd

@@ -37,8 +37,8 @@ Pkg.add("CounterfactualTraining")
 
 ### Introduction
 
-<!-- TODO: Add an introductory example demonstrating the package functionality; the example should use the data and visualizations produced here: https://github.com/JuliaTrustworthyAI/CounterfactualTraining.jl/blob/main/paper/sections/other/constraints.qmd -->
+A self-contained introductory example demonstrating counterfactual training with mutability protection is available on the [Introduction](introduction.md) page.
 
 ### Training on GPU
 
-<!-- TODO: Explain the two branches (initial research branch and performance-focused addition). Show how training can be done on the GPU. Ideally, add a somewhat more chunky model for MNIST, reproducing example here (only CE, not IG): https://github.com/JuliaTrustworthyAI/CounterfactualTraining.jl/blob/main/paper/experiments/mnist_chart.jl -->
+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.

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.

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.