Mutagenesis walk-through for ExplainMillX.jl
A complete, runnable example in three parts:
- Download the Mutagenesis dataset (cached locally).
- Train a Mill.jl model on it – a condensed version of the official JsonGrinder.jl tutorial:
- Explain one of the model's predictions using ExplainMillX's top-level
explain(...)convenience function. - Walk through what
explain(...)actually does internally, step by step: score importance withShapleyExplainer, prune withPruningStrategy/prune!, and print which atoms/bonds/descriptors the model actually needed to keep its prediction.
Run with: julia –project=examples examples/mutagenesis.jl
using HTTP
using JSON
using JsonGrinder
using Mill
using Flux
using MLUtils
using Statistics
using Printf
using Random
using ExplainMillX
Random.seed!(42)
# ---------------------------------------------------------------------
# 1. Download the data (cached locally so re-runs don't re-download)
# ---------------------------------------------------------------------
const DATA_DIR = joinpath(@__DIR__, "data")
const DATA_PATH = joinpath(DATA_DIR, "mutagenesis.json")
const DATA_URL = "https://raw.githubusercontent.com/CTUAvastLab/JsonGrinder.jl/master/docs/src/examples/mutagenesis/mutagenesis.json"
function ensure_data!(path=DATA_PATH, url=DATA_URL)
if !isfile(path)
mkpath(dirname(path))
@info "Downloading Mutagenesis dataset" url path
resp = HTTP.get(url)
write(path, resp.body)
end
path
end
ensure_data!()
# ---------------------------------------------------------------------
# 2. Train a Mill.jl model (condensed from the JsonGrinder.jl tutorial)
# ---------------------------------------------------------------------
dataset = JSON.parsefile(DATA_PATH)
jss_train, jss_test = dataset[1:100], dataset[101:end]
y_train = Flux.onehotbatch(getindex.(jss_train, "mutagenic"), 0:1)
y_test = Flux.onehotbatch(getindex.(jss_test, "mutagenic"), 0:1)2×88 OneHotMatrix(::Vector{UInt32}) with eltype Bool:
1 ⋅ ⋅ ⋅ ⋅ ⋅ ⋅ ⋅ ⋅ ⋅ ⋅ 1 ⋅ ⋅ ⋅ ⋅ 1 ⋅ 1 ⋅ 1 1 ⋅ 1 ⋅ 1 ⋅ ⋅ ⋅ ⋅ ⋅ ⋅ ⋅ 1 ⋅ ⋅ ⋅ ⋅ ⋅ 1 ⋅ ⋅ ⋅ ⋅ ⋅ ⋅ ⋅ 1 ⋅ ⋅ 1 1 ⋅ ⋅ ⋅ 1 1 ⋅ ⋅ 1 ⋅ ⋅ 1 1 1 ⋅ ⋅ ⋅ ⋅ ⋅ ⋅ 1 ⋅ ⋅ ⋅ 1 1 1 ⋅ ⋅ ⋅ ⋅ ⋅ ⋅ ⋅ ⋅ 1 ⋅
⋅ 1 1 1 1 1 1 1 1 1 1 ⋅ 1 1 1 1 ⋅ 1 ⋅ 1 ⋅ ⋅ 1 ⋅ 1 ⋅ 1 1 1 1 1 1 1 ⋅ 1 1 1 1 1 ⋅ 1 1 1 1 1 1 1 ⋅ 1 1 ⋅ ⋅ 1 1 1 ⋅ ⋅ 1 1 ⋅ 1 1 ⋅ ⋅ ⋅ 1 1 1 1 1 1 ⋅ 1 1 1 ⋅ ⋅ ⋅ 1 1 1 1 1 1 1 1 ⋅ 1Infer the JSON schema, then drop the label field (it isn't a feature).
sch = schema(jss_train)
delete!(sch, :mutagenic)Dict{Symbol, JsonGrinder.Schema} with 5 entries:
:lumo => LeafEntry
:inda => LeafEntry
:logp => LeafEntry
:ind1 => LeafEntry
:atoms => ArrayEntryall_stable=true (unlike the original tutorial's default extractor) makes every field – not just ones the schema happens to see as optional – extract into a MaybeHotMatrix/Union{Missing,...} representation rather than a plain, never-missing OneHotMatrix. ExplainMillX represents "this item is pruned away" as missing (masks.md §6), so every field needs to be able to hold missing for masking to have something to do to it.
e = suggestextractor(sch; all_stable=true)
x_train = extract(e, jss_train)
x_test = extract(e, jss_test)ProductNode 88 obs
├─── lumo: ArrayNode(99×88 MaybeHotMatrix with Union{Missing, Bool} elements) 88 obs
├─── inda: ArrayNode(2×88 MaybeHotMatrix with Union{Missing, Bool} elements) 88 obs
├─── logp: ArrayNode(63×88 MaybeHotMatrix with Union{Missing, Bool} elements) 88 obs
├─── ind1: ArrayNode(3×88 MaybeHotMatrix with Union{Missing, Bool} elements) 88 obs
╰── atoms: BagNode 88 obs
╰── ProductNode 2364 obs
├──── element: ArrayNode(7×2364 MaybeHotMatrix with Union{Missing, Bool} elements) 2364 obs
├────── bonds: BagNode 2364 obs
│ ╰── ProductNode 5084 obs
│ ┊
├───── charge: ArrayNode(1×2364 Array with Union{Missing, Float32} elements) 2364 obs
╰── atom_type: ArrayNode(29×2364 MaybeHotMatrix with Union{Missing, Bool} elements) 2364 obsall_imputing=true (again unlike the tutorial's default call) builds every Dense layer with an imputing weight matrix, so the model knows how to handle a missing input rather than only ever seeing complete molecules. This is required for ds[mask]'s pruned samples to be evaluable by the model at all.
model = reflectinmodel(sch, e; all_imputing=true, fsm = Dict("" => d -> Chain(Dense(d,10,relu), Dense(10,2))))
pred(m, x) = m(x)
loss(m, x, y) = Flux.Losses.logitcrossentropy(m(x), y)
opt_state = Flux.setup(Flux.Optimise.Descent(), model)
minibatch_iterator = Flux.DataLoader((x_train, y_train), batchsize=32, shuffle=true)
accuracy(p, y) = mean(Flux.onecold(p) .== Flux.onecold(y))
for i in 1:50
Flux.train!(loss, model, minibatch_iterator, opt_state)
@printf("Epoch %d train_accuracy: %.3f test_accuracy: %.3f\n",i, accuracy(pred(model, x_train), y_train), accuracy(pred(model, x_test), y_test))
end
# ---------------------------------------------------------------------
# 3. Explain one prediction
# ---------------------------------------------------------------------Epoch 1 train_accuracy: 0.610 test_accuracy: 0.727
Epoch 2 train_accuracy: 0.700 test_accuracy: 0.727
Epoch 3 train_accuracy: 0.640 test_accuracy: 0.727
Epoch 4 train_accuracy: 0.610 test_accuracy: 0.727
Epoch 5 train_accuracy: 0.610 test_accuracy: 0.727
Epoch 6 train_accuracy: 0.610 test_accuracy: 0.727
Epoch 7 train_accuracy: 0.610 test_accuracy: 0.727
Epoch 8 train_accuracy: 0.820 test_accuracy: 0.864
Epoch 9 train_accuracy: 0.820 test_accuracy: 0.864
Epoch 10 train_accuracy: 0.820 test_accuracy: 0.864
Epoch 11 train_accuracy: 0.820 test_accuracy: 0.864
Epoch 12 train_accuracy: 0.790 test_accuracy: 0.898
Epoch 13 train_accuracy: 0.820 test_accuracy: 0.864
Epoch 14 train_accuracy: 0.810 test_accuracy: 0.898
Epoch 15 train_accuracy: 0.820 test_accuracy: 0.852
Epoch 16 train_accuracy: 0.820 test_accuracy: 0.864
Epoch 17 train_accuracy: 0.820 test_accuracy: 0.864
Epoch 18 train_accuracy: 0.820 test_accuracy: 0.864
Epoch 19 train_accuracy: 0.820 test_accuracy: 0.864
Epoch 20 train_accuracy: 0.820 test_accuracy: 0.864
Epoch 21 train_accuracy: 0.690 test_accuracy: 0.739
Epoch 22 train_accuracy: 0.870 test_accuracy: 0.909
Epoch 23 train_accuracy: 0.820 test_accuracy: 0.864
Epoch 24 train_accuracy: 0.820 test_accuracy: 0.864
Epoch 25 train_accuracy: 0.820 test_accuracy: 0.864
Epoch 26 train_accuracy: 0.610 test_accuracy: 0.727
Epoch 27 train_accuracy: 0.830 test_accuracy: 0.875
Epoch 28 train_accuracy: 0.820 test_accuracy: 0.864
Epoch 29 train_accuracy: 0.840 test_accuracy: 0.875
Epoch 30 train_accuracy: 0.820 test_accuracy: 0.864
Epoch 31 train_accuracy: 0.820 test_accuracy: 0.864
Epoch 32 train_accuracy: 0.820 test_accuracy: 0.864
Epoch 33 train_accuracy: 0.880 test_accuracy: 0.886
Epoch 34 train_accuracy: 0.820 test_accuracy: 0.864
Epoch 35 train_accuracy: 0.820 test_accuracy: 0.864
Epoch 36 train_accuracy: 0.820 test_accuracy: 0.864
Epoch 37 train_accuracy: 0.880 test_accuracy: 0.818
Epoch 38 train_accuracy: 0.610 test_accuracy: 0.727
Epoch 39 train_accuracy: 0.820 test_accuracy: 0.864
Epoch 40 train_accuracy: 0.820 test_accuracy: 0.864
Epoch 41 train_accuracy: 0.820 test_accuracy: 0.864
Epoch 42 train_accuracy: 0.820 test_accuracy: 0.864
Epoch 43 train_accuracy: 0.820 test_accuracy: 0.864
Epoch 44 train_accuracy: 0.610 test_accuracy: 0.727
Epoch 45 train_accuracy: 0.820 test_accuracy: 0.864
Epoch 46 train_accuracy: 0.820 test_accuracy: 0.864
Epoch 47 train_accuracy: 0.820 test_accuracy: 0.864
Epoch 48 train_accuracy: 0.610 test_accuracy: 0.727
Epoch 49 train_accuracy: 0.820 test_accuracy: 0.864
Epoch 50 train_accuracy: 0.820 test_accuracy: 0.864
Test molecule #14, correctly classified by this model. Re-extracted (rather than reusing x_test) with store_input=Val(true) so every leaf's .metadata carries the original JSON values (element names, atom types, charges, ...) – applymask preserves .metadata through pruning, which is what lets the final explanation below be printed in human-readable form instead of as raw one-hot/matrix data.
sample_idx = 14
sample_json = jss_test[sample_idx]
ds = e(sample_json; store_input=Val(true))ProductNode 1 obs
├─── lumo: ArrayNode(99×1 MaybeHotMatrix with Union{Missing, Bool} elements) 1 obs
├─── inda: ArrayNode(2×1 MaybeHotMatrix with Union{Missing, Bool} elements) 1 obs
├─── logp: ArrayNode(63×1 MaybeHotMatrix with Union{Missing, Bool} elements) 1 obs
├─── ind1: ArrayNode(3×1 MaybeHotMatrix with Union{Missing, Bool} elements) 1 obs
╰── atoms: BagNode 1 obs
╰── ProductNode 30 obs
├──── element: ArrayNode(7×30 MaybeHotMatrix with Union{Missing, Bool} elements) 30 obs
├────── bonds: BagNode 30 obs
│ ╰── ProductNode 66 obs
│ ┊
├───── charge: ArrayNode(1×30 Array with Union{Missing, Float32} elements) 30 obs
╰── atom_type: ArrayNode(29×30 MaybeHotMatrix with Union{Missing, Bool} elements) 30 obsThe easiest way to explain the sample is to just call explain, which does everything below in one call: predicts the class, scores item importance, prunes, and reports how much of the sample was needed.
result = explain(ds, model)
display(result)
println()┌ Warning: explain: no tolerance specified, defaulting to rel_tol=0.9
└ @ ExplainMillX ~/work/ExplainMillX.jl/ExplainMillX.jl/src/explain.jl:25
The library allows to export the pruned sample to JSON, but to do that the sample needs to be extracted with argument store_input=Val(true) as used above. This is indeed pretty convenient.
JSON.print(explain_json(ds, result.mask, e), 4)
# ---------------------------------------------------------------------
# 4. Explaining one prediction -- step by step
# ---------------------------------------------------------------------{
"atoms": [
{
"charge": 0.812
},
{
"charge": 0.812
},
{
"charge": -0.388,
"element": "o"
},
{
"atom_type": 27,
"charge": 0.012,
"element": "c"
}
],
"ind1": 1,
"inda": 0,
"logp": 4.44
}Below, we show the explanation step by step, explaining the concrete machinery explain(...) above just did on our behalf – as literally as possible, so this section stays a faithful mirror of src/explain.jl rather than a simplified approximation of it. Intended for users who want to understand how the method actually works (or who need explainf's lower-level flexibility for something explain doesn't cover, e.g. a binary sigmoid head – see src/explain.jl's docstrings).
model's head is a genuine 2-class softmax (Dense(10, 2) above), so "confidence" is the gap between the predicted class's probability and the runner-up's – not the raw probability – exactly what explain computes internally. Defined locally, matching ExplainMillX._confgap's formula exactly, so the mechanism is visible here rather than hidden behind an internal function call.
confgap(p, c) = p[c] - maximum(p[1:end.!=c])
class = argmax(vec(model(ds)))
p0 = softmax(vec(model(ds)))
cg = confgap(p0, class)
@info "Explaining sample $sample_idx" confidence_gap = cg predicted = class true_label = argmax(y_test[:, sample_idx])┌ Info: Explaining sample 14
│ confidence_gap = 0.88930297f0
│ predicted = 2
└ true_label = 2
A real, honest finding from building this example: this quickly-trained model relies overwhelmingly on the four scalar molecular descriptors (lumo/logp/ind1/inda) for most predictions, so a looser rel_tol (e.g. 0.9) typically prunes the entire atoms/bonds structure away entirely – that's a genuine property of this model, not a limitation of the explainer. rel_tol is set tight (0.99) here so any residual atom-level signal has a chance to show up; because scoring is Monte Carlo (ShapleyExplainer), whether any specific atom's importance clears that bar can vary slightly run to run – the printed result below is whatever this run's explanation genuinely found, handled gracefully either way.
rel_tol = 0.99
threshold = rel_tol * cg
objective = o -> softmax(vec(o))[class]#5 (generic function with 1 method)explainf strips .metadata before scoring/pruning (Mill.dropmeta): stats/prune! re-evaluate model(ds[mask]) hundreds of times, the model never reads .metadata (it isn't part of .data), and applying a mask to a sample carrying metadata is measurably slower (roughly 2x per applymask call on a molecule this size) for zero benefit during the hot loop. The resulting mask is purely shape-derived, so it applies identically to the original, metadata-carrying ds afterward – which is exactly what lets Stage 3 below read real atom/bond values back out.
ds_nometa = Mill.dropmeta(ds)ProductNode 1 obs
├─── lumo: ArrayNode(99×1 MaybeHotMatrix with Union{Missing, Bool} elements) 1 obs
├─── inda: ArrayNode(2×1 MaybeHotMatrix with Union{Missing, Bool} elements) 1 obs
├─── logp: ArrayNode(63×1 MaybeHotMatrix with Union{Missing, Bool} elements) 1 obs
├─── ind1: ArrayNode(3×1 MaybeHotMatrix with Union{Missing, Bool} elements) 1 obs
╰── atoms: BagNode 1 obs
╰── ProductNode 30 obs
├──── element: ArrayNode(7×30 MaybeHotMatrix with Union{Missing, Bool} elements) 30 obs
├────── bonds: BagNode 30 obs
│ ╰── ProductNode 66 obs
│ ┊
├───── charge: ArrayNode(1×30 Array with Union{Missing, Float32} elements) 30 obs
╰── atom_type: ArrayNode(29×30 MaybeHotMatrix with Union{Missing, Bool} elements) 30 obsStage 1: score every maskable item's importance via Monte Carlo Shapley values (masks.md §8). stats builds the mask internally; it comes back scored and ready for pruning.
mk, acctree = stats(ShapleyExplainer(150), ds_nometa, model, objective; rng=MersenneTwister(11))
scores = nodescores(mk, acctree, score) # identity-keyed lookup, flatview.md §3IdDict{StructMask, Vector{Float64}} with 13 entries:
StructMask{Nothing, BitVector}(Bool[1], Bool[1], nothing) => [0.167197]
StructMask{Nothing, BitVector}(Bool[1, 0, 1, 0, 0, 0, 1, 1, 1, 1, 1, 1, 0, 1, 0, 0, 1, 0, 0, 0, 1, 0, 1, 0, 0, 0, 1, 1, 0, 1, 0, 1, 0, 0, 0, 0, 0, 1, 0, 0, 1, 1, 0, 1, 1, 1, 0, 1, 1, 0, 0, 1, 0, 1, 1, 0, 1, 1, 0, 1, 1, 1, 1, 0, 0, 1], Bool[0, 0, 0, 0, 1, 1, 0, 0, 1, 0, 1, 0, 0, 0, 0, 0, 0, 1, 0, 0, 0, 0, 1, 0, 1, 0, 0, 0, 0, 0, 1, 0, 0, 0, 1, 0, 1, 0, 0, 0, 0, 0, 0, 0, 0, 0, 1, 0, 1, 0, 0, 0, 0, 0, 1, 0, 1, 0, 0, 0, 0, 0, 1, 0, 1, 0], nothing) => [0.0616069, 0.0504963, -0.0578336, 0.0710557, -0.0153212, 0.0391668, 0.0112046, -0.0151794, 0.00812462, 0.0743775, 0.00336441, 0.0364797, 0.00954657, -0.0182135, 0.0318097, 0.0495887, 0.0499585, 0.0212526, 0.00477242, 0.0522761, -0.0324223, -0.0132901, 0.00051891, -0.0132985, -0.00751147, -0.0218695, 0.00184383, 0.0219735, -0.0420722, 0.0177575, -0.0372724, -0.075742, 0.0415495, 0.012949, 0.00630724, 0.000203869, -0.022224, 0.0381659, 0.0319637, 0.0529982, -0.0382313, 0.0210171, 0.00551792, -0.0133803, -0.00198466, 0.0290054, 0.00593201, -0.0430279, -0.0264403, 0.0618399, 0.0193084, 0.0551696, 0.0683105, 0.0153556, 0.0209837, 0.00544423, -0.0288089, -0.000876016, 0.00853339, -0.0136135, -0.00107243, 0.0357279, 0.0343952, 0.0267335, -0.00330008, -0.027074]
StructMask{Nothing, BitVector}(Bool[1, 0, 1, 0, 1, 1, 1, 1, 0, 1, 0, 1, 1, 0, 0, 1, 0, 0, 0, 0, 0, 1, 0, 1, 0, 0, 1, 1, 1, 1, 0, 0, 1, 0, 1, 0, 0, 0, 0, 0, 1, 1, 1, 0, 0, 0, 1, 0, 1, 1, 0, 0, 0, 0, 0, 1, 0, 1, 1, 1, 1, 1, 0, 0, 1, 0], Bool[0, 0, 0, 0, 1, 1, 0, 0, 1, 0, 1, 0, 0, 0, 0, 0, 0, 1, 0, 0, 0, 0, 1, 0, 1, 0, 0, 0, 0, 0, 1, 0, 0, 0, 1, 0, 1, 0, 0, 0, 0, 0, 0, 0, 0, 0, 1, 0, 1, 0, 0, 0, 0, 0, 1, 0, 1, 0, 0, 0, 0, 0, 1, 0, 1, 0], nothing) => [0.0308622, -0.0160049, 0.0146618, -0.0140694, -0.00897468, 0.0162264, -0.0739157, 0.0290449, 0.0316863, -0.0501108, -0.0451949, -0.00662563, 0.0130288, 0.000369277, 0.0774028, -0.0175609, -0.033493, 0.0069664, 0.026383, 0.00497459, -0.0294374, 0.0271778, -0.00722475, -0.0234697, 0.0113041, -0.0253408, 0.043474, -0.0165543, 0.0134104, -0.0583954, 0.00810927, -0.0644313, -0.0445018, 0.00597284, 0.0411286, -0.0541377, -0.0218666, -0.0172624, 0.0292327, -0.0554673, 0.0173906, 0.0711856, 0.0829243, 0.0048317, 0.0605551, 0.0578926, -0.0374935, -0.00387143, 0.00153581, 0.0299043, 0.0100215, 0.0165662, 0.0122528, -0.0156582, -0.0270592, 0.0237297, -0.0573537, 0.00492466, -0.0185208, -0.0108928, 0.0400743, -0.0335057, 0.00140623, 0.0474029, -0.0377058, 0.0178228]
StructMask{Nothing, BitVector}(Bool[0], Bool[1], nothing) => [0.0632141]
StructMask{StructMask{@NamedTuple{element::StructMask{Nothing, BitVector}, bonds::StructMask{StructMask{@NamedTuple{element::StructMask{Nothing, BitVector}, bond_type::StructMask{Nothing, BitVector}, charge::StructMask{Nothing, BitVector}, atom_type::StructMask{Nothing, BitVector}}, Vector{Bool}}, BitVector}, charge::StructMask{Nothing, BitVector}, atom_type::StructMask{Nothing, BitVector}}, Vector{Bool}}, BitVector}(Bool[0, 1, 1, 1, 0, 0, 1, 0, 1, 0, 1, 0, 1, 0, 0, 0, 0, 0, 1, 1, 1, 0, 0, 0, 1, 0, 1, 0, 1, 0], Bool[1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1], StructMask{@NamedTuple{element::StructMask{Nothing, BitVector}, bonds::StructMask{StructMask{@NamedTuple{element::StructMask{Nothing, BitVector}, bond_type::StructMask{Nothing, BitVector}, charge::StructMask{Nothing, BitVector}, atom_type::StructMask{Nothing, BitVector}}, Vector{Bool}}, BitVector}, charge::StructMask{Nothing, BitVector}, atom_type::StructMask{Nothing, BitVector}}, Vector{Bool}}(nothing, nothing, (element = StructMask{Nothing, BitVector}(Bool[0, 1, 1, 0, 0, 0, 1, 0, 0, 1, 1, 1, 1, 1, 0, 0, 0, 1, 1, 0, 0, 1, 0, 1, 0, 0, 0, 0, 0, 0], Bool[0, 1, 1, 1, 0, 0, 1, 0, 1, 0, 1, 0, 1, 0, 0, 0, 0, 0, 1, 1, 1, 0, 0, 0, 1, 0, 1, 0, 1, 0], nothing), bonds = StructMask{StructMask{@NamedTuple{element::StructMask{Nothing, BitVector}, bond_type::StructMask{Nothing, BitVector}, charge::StructMask{Nothing, BitVector}, atom_type::StructMask{Nothing, BitVector}}, Vector{Bool}}, BitVector}(Bool[0, 0, 1, 0, 1, 1, 0, 0, 1, 0, 1, 0, 0, 0, 1, 0, 0, 1, 0, 1, 0, 0, 1, 0, 1, 0, 0, 0, 0, 0, 1, 1, 1, 1, 1, 0, 1, 1, 1, 1, 0, 1, 0, 0, 0, 0, 1, 0, 1, 1, 1, 0, 0, 0, 1, 0, 1, 1, 0, 0, 0, 0, 1, 0, 1, 1], Bool[0, 0, 0, 1, 1, 1, 1, 1, 1, 1, 1, 1, 0, 0, 0, 0, 1, 1, 1, 0, 0, 0, 1, 1, 1, 0, 0, 0, 1, 1, 1, 0, 0, 0, 1, 1, 1, 0, 0, 0, 0, 0, 0, 0, 0, 0, 1, 1, 1, 0, 0, 0, 0, 0, 1, 1, 1, 0, 0, 0, 1, 1, 1, 0, 1, 0], StructMask{@NamedTuple{element::StructMask{Nothing, BitVector}, bond_type::StructMask{Nothing, BitVector}, charge::StructMask{Nothing, BitVector}, atom_type::StructMask{Nothing, BitVector}}, Vector{Bool}}(nothing, nothing, (element = StructMask{Nothing, BitVector}(Bool[1, 1, 1, 1, 0, 1, 0, 1, 1, 1, 1, 0, 1, 1, 0, 1, 1, 0, 0, 1, 0, 1, 0, 1, 1, 0, 0, 0, 0, 0, 0, 1, 1, 1, 1, 1, 1, 1, 1, 0, 1, 0, 0, 1, 1, 0, 1, 1, 0, 1, 0, 0, 1, 0, 0, 1, 0, 1, 0, 1, 0, 1, 1, 0, 0, 0], Bool[0, 0, 0, 0, 1, 1, 0, 0, 1, 0, 1, 0, 0, 0, 0, 0, 0, 1, 0, 0, 0, 0, 1, 0, 1, 0, 0, 0, 0, 0, 1, 0, 0, 0, 1, 0, 1, 0, 0, 0, 0, 0, 0, 0, 0, 0, 1, 0, 1, 0, 0, 0, 0, 0, 1, 0, 1, 0, 0, 0, 0, 0, 1, 0, 1, 0], nothing), bond_type = StructMask{Nothing, BitVector}(Bool[1, 0, 1, 0, 1, 1, 1, 1, 0, 1, 0, 1, 1, 0, 0, 1, 0, 0, 0, 0, 0, 1, 0, 1, 0, 0, 1, 1, 1, 1, 0, 0, 1, 0, 1, 0, 0, 0, 0, 0, 1, 1, 1, 0, 0, 0, 1, 0, 1, 1, 0, 0, 0, 0, 0, 1, 0, 1, 1, 1, 1, 1, 0, 0, 1, 0], Bool[0, 0, 0, 0, 1, 1, 0, 0, 1, 0, 1, 0, 0, 0, 0, 0, 0, 1, 0, 0, 0, 0, 1, 0, 1, 0, 0, 0, 0, 0, 1, 0, 0, 0, 1, 0, 1, 0, 0, 0, 0, 0, 0, 0, 0, 0, 1, 0, 1, 0, 0, 0, 0, 0, 1, 0, 1, 0, 0, 0, 0, 0, 1, 0, 1, 0], nothing), charge = StructMask{Nothing, BitVector}(Bool[0], Bool[1], nothing), atom_type = StructMask{Nothing, BitVector}(Bool[1, 0, 1, 0, 0, 0, 1, 1, 1, 1, 1, 1, 0, 1, 0, 0, 1, 0, 0, 0, 1, 0, 1, 0, 0, 0, 1, 1, 0, 1, 0, 1, 0, 0, 0, 0, 0, 1, 0, 0, 1, 1, 0, 1, 1, 1, 0, 1, 1, 0, 0, 1, 0, 1, 1, 0, 1, 1, 0, 1, 1, 1, 1, 0, 0, 1], Bool[0, 0, 0, 0, 1, 1, 0, 0, 1, 0, 1, 0, 0, 0, 0, 0, 0, 1, 0, 0, 0, 0, 1, 0, 1, 0, 0, 0, 0, 0, 1, 0, 0, 0, 1, 0, 1, 0, 0, 0, 0, 0, 0, 0, 0, 0, 1, 0, 1, 0, 0, 0, 0, 0, 1, 0, 1, 0, 0, 0, 0, 0, 1, 0, 1, 0], nothing)))), charge = StructMask{Nothing, BitVector}(Bool[1], Bool[1], nothing), atom_type = StructMask{Nothing, BitVector}(Bool[0, 1, 0, 0, 0, 1, 1, 0, 0, 0, 1, 0, 1, 1, 1, 0, 1, 0, 0, 1, 0, 1, 1, 1, 1, 0, 1, 1, 1, 0], Bool[0, 1, 1, 1, 0, 0, 1, 0, 1, 0, 1, 0, 1, 0, 0, 0, 0, 0, 1, 1, 1, 0, 0, 0, 1, 0, 1, 0, 1, 0], nothing)))) => [0.0135097, 0.00265741, 0.00154226, 0.020738, -0.00933494, -0.0178862, 0.0282703, 0.020601, -0.0195342, 2.20847e-5, -0.00585563, -0.00315248, -0.00267856, 0.012484, -0.00795806, 0.0203942, 0.0124837, 0.00699421, 0.0277331, 0.0415731, -0.000175487, 0.00836324, 0.00309195, -0.0129944, 0.0418446, -0.01529, 0.0266748, 0.00563637, 0.0241999, 0.0139911]
StructMask{Nothing, BitVector}(Bool[0], Bool[1], nothing) => [-0.0405711]
StructMask{Nothing, BitVector}(Bool[0, 1, 1, 0, 0, 0, 1, 0, 0, 1, 1, 1, 1, 1, 0, 0, 0, 1, 1, 0, 0, 1, 0, 1, 0, 0, 0, 0, 0, 0], Bool[0, 1, 1, 1, 0, 0, 1, 0, 1, 0, 1, 0, 1, 0, 0, 0, 0, 0, 1, 1, 1, 0, 0, 0, 1, 0, 1, 0, 1, 0], nothing) => [0.00536477, 0.0165955, -0.00141496, -0.000716351, 0.0288032, 0.01907, -0.0109439, 0.0204186, -0.0320588, -0.015088, -0.0329889, -0.0041571, 0.00594278, -0.0198576, 0.00652005, -0.038526, 0.0120942, -0.000379664, -0.0163922, 0.0125218, 0.00360763, 0.0290575, -0.012036, 0.0331534, -0.00419716, -0.0349528, -0.00440905, 0.0510103, -0.00636259, -0.00822159]
StructMask{Nothing, BitVector}(Bool[1, 1, 1, 1, 0, 1, 0, 1, 1, 1, 1, 0, 1, 1, 0, 1, 1, 0, 0, 1, 0, 1, 0, 1, 1, 0, 0, 0, 0, 0, 0, 1, 1, 1, 1, 1, 1, 1, 1, 0, 1, 0, 0, 1, 1, 0, 1, 1, 0, 1, 0, 0, 1, 0, 0, 1, 0, 1, 0, 1, 0, 1, 1, 0, 0, 0], Bool[0, 0, 0, 0, 1, 1, 0, 0, 1, 0, 1, 0, 0, 0, 0, 0, 0, 1, 0, 0, 0, 0, 1, 0, 1, 0, 0, 0, 0, 0, 1, 0, 0, 0, 1, 0, 1, 0, 0, 0, 0, 0, 0, 0, 0, 0, 1, 0, 1, 0, 0, 0, 0, 0, 1, 0, 1, 0, 0, 0, 0, 0, 1, 0, 1, 0], nothing) => [0.0649648, -0.0296105, -0.0189653, 0.00172106, 0.0293895, 0.0146818, 0.0249158, -0.0201987, 0.0342902, 0.00611028, 0.0710851, -0.0273961, 0.0934537, 0.0312484, 0.000430542, -0.00267963, -0.00250168, -0.00863872, -0.0123994, -0.0216886, -0.0253839, 0.00729609, -0.0231662, 0.0292081, 0.0498989, -0.00263753, 0.0165262, -0.032703, 0.0149456, -0.0271751, 0.031504, -0.0609509, -0.0168477, -0.0423555, -0.0390817, -0.0278362, 0.00314379, -0.0254397, -0.0256892, -0.0329311, -0.00630583, -0.0341553, -0.0307673, 0.0101285, 0.0272223, -0.00479982, 0.0218445, -0.0385135, 0.0343895, 0.0255816, -0.0106186, -0.0386143, 0.0242852, 0.0255976, 0.0634368, 0.0540459, 0.00607797, -0.0295372, -0.0387096, 0.0265264, 0.0063938, -0.0449601, -0.0301902, -0.0342071, -0.0109371, 0.0631993]
StructMask{Nothing, BitVector}(Bool[0, 1, 0, 0, 0, 1, 1, 0, 0, 0, 1, 0, 1, 1, 1, 0, 1, 0, 0, 1, 0, 1, 1, 1, 1, 0, 1, 1, 1, 0], Bool[0, 1, 1, 1, 0, 0, 1, 0, 1, 0, 1, 0, 1, 0, 0, 0, 0, 0, 1, 1, 1, 0, 0, 0, 1, 0, 1, 0, 1, 0], nothing) => [0.0174721, -0.0289968, 0.0120863, -0.00413541, -0.0112265, 0.0237004, 0.0200188, -0.0191852, 0.000778796, 0.00269878, -0.00985666, -0.0343644, -0.00281043, 0.0349941, -0.0122989, -0.00678818, -0.00972613, -0.00767608, 0.00088937, 0.0422146, 0.00434288, 0.00516921, 0.0164271, 0.000318921, 0.00289134, 0.010832, -0.0336588, -0.012759, 0.0243186, -0.000647951]
StructMask{Nothing, BitVector}(Bool[1], Bool[1], nothing) => [0.0279524]
StructMask{StructMask{@NamedTuple{element::StructMask{Nothing, BitVector}, bond_type::StructMask{Nothing, BitVector}, charge::StructMask{Nothing, BitVector}, atom_type::StructMask{Nothing, BitVector}}, Vector{Bool}}, BitVector}(Bool[0, 0, 1, 0, 1, 1, 0, 0, 1, 0, 1, 0, 0, 0, 1, 0, 0, 1, 0, 1, 0, 0, 1, 0, 1, 0, 0, 0, 0, 0, 1, 1, 1, 1, 1, 0, 1, 1, 1, 1, 0, 1, 0, 0, 0, 0, 1, 0, 1, 1, 1, 0, 0, 0, 1, 0, 1, 1, 0, 0, 0, 0, 1, 0, 1, 1], Bool[0, 0, 0, 1, 1, 1, 1, 1, 1, 1, 1, 1, 0, 0, 0, 0, 1, 1, 1, 0, 0, 0, 1, 1, 1, 0, 0, 0, 1, 1, 1, 0, 0, 0, 1, 1, 1, 0, 0, 0, 0, 0, 0, 0, 0, 0, 1, 1, 1, 0, 0, 0, 0, 0, 1, 1, 1, 0, 0, 0, 1, 1, 1, 0, 1, 0], StructMask{@NamedTuple{element::StructMask{Nothing, BitVector}, bond_type::StructMask{Nothing, BitVector}, charge::StructMask{Nothing, BitVector}, atom_type::StructMask{Nothing, BitVector}}, Vector{Bool}}(nothing, nothing, (element = StructMask{Nothing, BitVector}(Bool[1, 1, 1, 1, 0, 1, 0, 1, 1, 1, 1, 0, 1, 1, 0, 1, 1, 0, 0, 1, 0, 1, 0, 1, 1, 0, 0, 0, 0, 0, 0, 1, 1, 1, 1, 1, 1, 1, 1, 0, 1, 0, 0, 1, 1, 0, 1, 1, 0, 1, 0, 0, 1, 0, 0, 1, 0, 1, 0, 1, 0, 1, 1, 0, 0, 0], Bool[0, 0, 0, 0, 1, 1, 0, 0, 1, 0, 1, 0, 0, 0, 0, 0, 0, 1, 0, 0, 0, 0, 1, 0, 1, 0, 0, 0, 0, 0, 1, 0, 0, 0, 1, 0, 1, 0, 0, 0, 0, 0, 0, 0, 0, 0, 1, 0, 1, 0, 0, 0, 0, 0, 1, 0, 1, 0, 0, 0, 0, 0, 1, 0, 1, 0], nothing), bond_type = StructMask{Nothing, BitVector}(Bool[1, 0, 1, 0, 1, 1, 1, 1, 0, 1, 0, 1, 1, 0, 0, 1, 0, 0, 0, 0, 0, 1, 0, 1, 0, 0, 1, 1, 1, 1, 0, 0, 1, 0, 1, 0, 0, 0, 0, 0, 1, 1, 1, 0, 0, 0, 1, 0, 1, 1, 0, 0, 0, 0, 0, 1, 0, 1, 1, 1, 1, 1, 0, 0, 1, 0], Bool[0, 0, 0, 0, 1, 1, 0, 0, 1, 0, 1, 0, 0, 0, 0, 0, 0, 1, 0, 0, 0, 0, 1, 0, 1, 0, 0, 0, 0, 0, 1, 0, 0, 0, 1, 0, 1, 0, 0, 0, 0, 0, 0, 0, 0, 0, 1, 0, 1, 0, 0, 0, 0, 0, 1, 0, 1, 0, 0, 0, 0, 0, 1, 0, 1, 0], nothing), charge = StructMask{Nothing, BitVector}(Bool[0], Bool[1], nothing), atom_type = StructMask{Nothing, BitVector}(Bool[1, 0, 1, 0, 0, 0, 1, 1, 1, 1, 1, 1, 0, 1, 0, 0, 1, 0, 0, 0, 1, 0, 1, 0, 0, 0, 1, 1, 0, 1, 0, 1, 0, 0, 0, 0, 0, 1, 0, 0, 1, 1, 0, 1, 1, 1, 0, 1, 1, 0, 0, 1, 0, 1, 1, 0, 1, 1, 0, 1, 1, 1, 1, 0, 0, 1], Bool[0, 0, 0, 0, 1, 1, 0, 0, 1, 0, 1, 0, 0, 0, 0, 0, 0, 1, 0, 0, 0, 0, 1, 0, 1, 0, 0, 0, 0, 0, 1, 0, 0, 0, 1, 0, 1, 0, 0, 0, 0, 0, 0, 0, 0, 0, 1, 0, 1, 0, 0, 0, 0, 0, 1, 0, 1, 0, 0, 0, 0, 0, 1, 0, 1, 0], nothing)))) => [-0.0120504, 0.00205087, 0.0098847, 0.0372809, 0.0216493, 0.0121664, 0.0218495, 0.0300432, -0.00285337, -0.00813039, 0.0142509, 0.00325004, -0.0317011, 0.00499409, -0.000926673, 0.0099398, -0.033189, -0.0151161, -0.041076, -0.00504263, 0.0142892, -0.00920117, 0.0274311, -0.0143456, 0.00410302, -0.00444735, 0.0403075, -0.064047, 0.00412007, -0.0765662, -0.0395064, 0.0150802, 0.0245526, 0.0323317, 0.0219744, -0.0240421, -0.0172607, -0.0149857, -0.00217526, -0.0392414, 0.00103756, -0.0250746, -0.0249903, -0.00720338, -0.00182115, 0.0159324, 0.00555441, -0.00169447, -0.0207121, -0.01583, 0.0289533, -0.0868604, -0.0381776, -0.00113182, 0.0146682, 0.00594396, -0.000455783, -0.0698521, -0.0106328, 0.0376007, 0.00149931, -0.0167262, 0.0344597, -0.0122414, 0.0103925, 0.0154025]
StructMask{Nothing, BitVector}(Bool[0], Bool[1], nothing) => [-0.0170868]
StructMask{Nothing, BitVector}(Bool[1], Bool[1], nothing) => [0.0237967]Stage 2: prune down to a minimal subset that keeps the confidence gap above threshold – i.e. the model stays at least rel_tol as confident in the same prediction using only what survives pruning.
f = () -> confgap(softmax(vec(model(ds_nometa[mk]))), class) - threshold
strategy = PruningStrategy(HeuristicOrder(scores), true, true, true) # level-by-level: faster in practice (pruning.md §2.2)
prune!(mk, ds_nometa, model, f, strategy)
@info "Pruning result" remaining_confidence_gap = f() + threshold original_confidence_gap = cg┌ Info: Pruning result
│ remaining_confidence_gap = 0.8838201761245728
└ original_confidence_gap = 0.88930297f0
Stage 3: read the pruned mask back out in human-readable form. mk was built and searched entirely against ds_nometa (no metadata anywhere), but – per the note above – it applies just as well to the original, metadata-carrying ds, which is what makes .metadata available here to print real atom/bond values instead of raw one-hot/matrix data.
pruned = ds[mk]
println("\nWhat the model needed to keep its prediction:")
println(" lumo = ", only(pruned[:lumo].metadata))
println(" logp = ", only(pruned[:logp].metadata))
println(" ind1 = ", only(pruned[:ind1].metadata))
println(" inda = ", only(pruned[:inda].metadata))
atom_elements = pruned[:atoms].data[:element].metadata
atom_charges = pruned[:atoms].data[:charge].metadata
kept_atoms = findall(!ismissing, atom_elements)
if isempty(kept_atoms)
println(" atoms: none -- the scalar descriptors above were sufficient on their own")
else
println(" atoms kept ($(length(kept_atoms)) of $(length(atom_elements))):")
for i in kept_atoms
println(" atom $i: element=$(atom_elements[i]) charge=$(atom_charges[i])")
end
end
What the model needed to keep its prediction:
lumo = -2.055
logp = 4.44
ind1 = 1
inda = 0
atoms kept (9 of 9):
atom 1: element=h charge=0.142
atom 2: element=h charge=0.142
atom 3: element=n charge=0.812
atom 4: element=n charge=0.812
atom 5: element=h charge=0.142
atom 6: element=c charge=-0.088
atom 7: element=h charge=0.142
atom 8: element=h charge=0.142
atom 9: element=h charge=0.142
This page was generated using Literate.jl.