API Reference

All exported names, generated from docstrings. See the How-to Guides for task-oriented usage and the Explanation section for the concepts behind them.

Explaining a prediction

ExplainMillX.explainFunction
explain(ds, model, class; kwargs...) -> ExplanationResult
explain(ds, model; kwargs...) -> ExplanationResult

Explain why model classifies sample ds as class, by finding a small subset of ds's items that keeps the model's confidence in class within tolerance of its original value.

If class is omitted, the model's own predicted class (argmax of its softmax output on the full sample) is used.

This is the full pipeline in one call:

  1. Predict: evaluate model(ds) to get the baseline confidence gap for class (its softmax probability minus the runner-up's).
  2. Score: build a mask over ds and estimate each item's importance via scorer (e.g. ShapleyExplainer) – see masks.md §8.
  3. Prune: search for a small item subset keeping the confidence gap above a threshold derived from abs_tol/rel_tol – see docs/design/pruning.md.
  4. Report: package the result together with how much of the sample was pruned away, as an ExplanationResult.

model must produce a softmax-style output with two or more classes for a single observation (length(vec(model(ds))) >= 2); for a single scalar/sigmoid output, use explainf directly with a hand-written objective (examples/mutagenesis.jl shows exactly this for a binary classifier).

Only ds and model (the data being explained) are positional – every choice of how to explain is a keyword, including scorer, consistent with explainf's order/levelbylevel/random_removal/finetune already being keywords. The actual work happens in _explain, kept as a separate internal function specifically so a future scoring strategy that needs explain's own orchestration (not just its own stats method) to diverge can add an _explain method without touching this signature.

Arguments

  • ds::AbstractMillNode: the sample to explain.
  • model: the Mill model. model(ds) must return (something vec-able to) one softmax logit/probability per class.
  • class::Integer: which class to explain. Must be (one of) the model's actual predicted class(es) for ds – i.e. have a nonnegative confidence gap; explain raises an error otherwise, since there is nothing meaningful to preserve while pruning if the model didn't predict this class to begin with.

Keyword arguments

  • scorer = ShapleyExplainer(300): the scoring strategy. 300 Monte Carlo samples is a reasonable default in practice; pass e.g. scorer=ShapleyExplainer(1000) for a more precise (slower) estimate.
  • abs_tol, rel_tol: exactly one may be given (or neither, which warns once and defaults to rel_tol=0.9). rel_tol (in [0, 1]) keeps the confidence gap at that fraction of its original value; abs_tol keeps it within that absolute amount of the original value (and must not exceed it).
  • order, levelbylevel, random_removal, finetune, rng: forwarded to explainf – see its docstring for what each controls.

Example

model = reflectinmodel(ds, d -> Dense(d, nclasses); all_imputing=true)
result = explain(ds, model)  # explain the predicted class, default scorer
result.mask            # the pruned mask
ds[result.mask]         # the pruned sample
fraction_pruned(result) # how much of ds turned out to be unnecessary
ExplainMillX.explainfFunction
explainf(scorer, ds, model, fₛ, fₚ; kwargs...) -> (mask=..., n_total=..., n_kept=...)

Low-level entry point: explain ds under model using scoring strategy scorer and caller-supplied objectives, bypassing the classification convenience layer entirely (docs/design/design.doc §4.1). Use this directly for anything explain doesn't cover – e.g. a binary sigmoid head, a regression target, or a custom notion of "confidence" (see examples/mutagenesis.jl for a full worked example of exactly this).

fₛ(o) and fₚ(o) are both functions of the model's output o (matching stats's objective argument), not zero-argument closures – explainf builds the mask via stats first and only then has something for a zero-argument f() = fₚ(model(ds[mask])) to close over, which is what prune! actually requires.

ds is stripped of metadata internally (Mill.dropmeta) before scoring and pruning: 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 with metadata is measurably slower (applymask on a BagNode/ProductNode re-slices .metadata on every call via Mill's own getindex, roughly 2x more work per call, confirmed by benchmark) for zero benefit during the hot loop. The returned mask is purely shape-derived and applies identically to a metadata-carrying ds afterward – explain does exactly that for its final result.

Keyword arguments

  • order::Union{HeuristicOrder,GreedyForward,Nothing} = nothing: the pruning search order (docs/design/pruning.md §4). nothing (the default) builds a HeuristicOrder automatically from the scoring result via nodescores; pass GreedyForward() to use stepwise greedy selection instead, or a pre-built HeuristicOrder to reuse scores from elsewhere.
  • levelbylevel::Bool = true: search one hierarchy depth at a time (generally faster in practice, docs/design/pruning.md §2.2) rather than the whole tree at once.
  • random_removal::Bool = true, finetune::Bool = true: optional redundancy-removal and local-search post-passes (docs/design/pruning.md §4.5, §4.6).
  • rng::AbstractRNG = Random.default_rng(): random number generator threaded through the (stochastic) scoring strategy, for reproducibility.

Returns a NamedTuple (mask, n_total, n_kept) – see ExplanationResult for what these mean; explainf doesn't wrap them in that struct itself since it has no class/confidence-gap notion to report.

ExplainMillX.ExplanationResultType
ExplanationResult

The result of explain: the pruned mask, the sample it applies to, and a summary of how much of the sample was pruned away and how much prediction confidence was retained.

Fields

  • mask::StructMask: the final, pruned mask. Apply it with sample[mask] to get the pruned Mill sample, or read prunemask.(first.(collectmasks(mask))) for the raw per-item keep/drop decisions.
  • sample::AbstractMillNode: the (original, un-pruned) sample this explanation is for – i.e. exactly the ds passed to explain. Stored alongside mask so an ExplanationResult is self-contained: mask alone means nothing without knowing what it applies to (e.g. for explain_json, or if sample was extracted with store_input=Val(true) and you want its metadata later without having to keep ds around separately).
  • class::Int: the class index this explanation was computed for.
  • confidence_gap::Float64: the model's confidence gap for class on the original, unpruned sample (softmax probability of class minus the highest probability among all other classes).
  • remaining_confidence_gap::Float64: the same confidence gap evaluated on the pruned sample sample[mask]. Always >= threshold.
  • threshold::Float64: the minimum confidence gap pruning was required to preserve, derived from abs_tol/rel_tol (see explain).
  • n_total::Int: total number of maskable items in the sample (features, bag instances, categorical values, ... – see docs/design/masks.md §3).
  • n_kept::Int: how many of those items survived pruning.

Use fraction_kept/fraction_pruned/n_pruned for derived statistics.

ExplainMillX.n_prunedFunction
n_pruned(r::ExplanationResult) -> Int

Number of items pruned away, i.e. r.n_total - r.n_kept.

ExplainMillX.fraction_keptFunction
fraction_kept(r::ExplanationResult) -> Float64

Fraction (in [0, 1]) of the sample's maskable items that survived pruning. 1.0 for a sample with no maskable items at all.

ExplainMillX.fraction_prunedFunction
fraction_pruned(r::ExplanationResult) -> Float64

1 - fraction_kept(r): the fraction of the sample's maskable items that were pruned away.

JSON output

ExplainMillX.explain_jsonFunction
explain_json(ds::AbstractMillNode, mask::StructMask, extractor) -> Any

Reconstruct a pruned mask over ds as a JSON-shaped value (nested Dict/Vector/scalars, with pruned or absent items represented as nothing), using extractor to know how to interpret each part of the tree and ds's preserved .metadata to recover original values.

ds must have been extracted with store_input=Val(true) (directly, or via extract(extractor, samples; store_input=Val(true))) – otherwise there is nothing for this function to reconstruct, and it raises a clear error the first time it needs a leaf's metadata and finds nothing instead.

extractor must be the same extractor object used to produce ds (structurally): reconstruction dispatches on the extractor's type at every level, since the extractor – not the Mill node type alone – determines the JSON shape (e.g. whether a ProductNode was a plain object or, in general, something else; whether an ArrayNode came from a categorical, n-gram, or scalar field).

ds must represent a single sample (numobs(ds) == 1) – explain_json has no batching support (docs/design/jsonoutput.md §5), matching the rest of ExplainMillX.

PolymorphExtractor and ProductNode{<:Tuple} (JsonGrinder's union-typed/positional extraction) are not supported.

Example

ds = e(json_sample; store_input=Val(true))
result = explain(ds, model)   # default scorer, ShapleyExplainer(300)
explain_json(result, e)   # => a JSON-shaped Dict, pruned parts as `nothing`/absent
explain_json(ds, result.mask, e)   # equivalent, spelled out
explain_json(result::ExplanationResult, extractor) -> Any

Convenience form of explain_json for the common case: result already carries the sample it was computed for (result.sample), so this is exactly explain_json(result.sample, result.mask, extractor).

Scoring strategies

ExplainMillX.ShapleyExplainerType
ShapleyExplainer(n=1000)

Monte Carlo Shapley/Banzhaf-style scoring strategy: repeatedly randomizes the whole mask, evaluates objective(model(ds[mask])), and accumulates a MeanDiff per unit. Exists primarily to exercise the mask infrastructure end to end (random sampling, participation, hard pruning, paired bookkeeping traversal) – see docs/design/pruning.md for where this fits relative to a production scoring strategy.

ExplainMillX.statsFunction
stats(e::ShapleyExplainer, ds, model, objective; rng=Random.default_rng())

objective(output) -> Real is evaluated on the model's output for each randomly masked sample; no assumption about classification/class indices is made here (see design.doc §4.1 on objective injection).

Returns (mask, acctree); use score.(payload) on acctree's leaves (via foreach_paired) to read out per-unit importance.

ExplainMillX.MeanDiffType
MeanDiff

Running per-unit accumulator: separate means of the objective value observed when a unit was on vs. off across random subset samples. The difference (score) is an unbiased Monte Carlo (linear-approximation) estimate of the unit's Shapley/Banzhaf value – the same idea as ExplainMill.jl's Duff.Daf, reimplemented directly rather than depending on Duff.jl.

ExplainMillX.scoreFunction
score(a::MeanDiff) -> Float64

The estimated importance of the unit a accumulates statistics for: the mean objective value observed when the unit was on, minus the mean when it was off. 0.0 if the unit has never been observed both on and off.

ExplainMillX.leafscoresFunction
leafscores(m::StructMask, acc::AccTree)

Collect score.(payload) for every own-bearing node, in the same order collectmasks/foreach_mask would visit them.

Pruning strategies

ExplainMillX.PruningStrategyType
PruningStrategy{Order}(order, levelbylevel, random_removal, finetune)

Every axis is a required, explicit field – no hidden defaults – per docs/design/pruning.md §3.6 (a naming/behavior mismatch in the original was traced directly to a post-pass silently defaulting on inside an unrelated function).

  • order: HeuristicOrder(scores) or GreedyForward().
  • levelbylevel: false = search the whole tree at once. true = one pass per hierarchy depth, narrowing outward-in (see prune!).
  • random_removal: run randomremoval! as a post-pass.
  • finetune: run finetune! as a post-pass.
ExplainMillX.HeuristicOrderType
HeuristicOrder(scores::IdDict{StructMask,Vector{Float64}})

Order strategy: candidates are sorted by a precomputed importance score and added via bisection (addminimumbi!). scores is built once, by the caller, via nodescores(mask, acctree, score) right after scoring – prune!/PruningStrategy never reference AccTree or a scoring-strategy type at all (docs/design/pruning.md §3.3, §4, Part 2 of the driver discussion).

ExplainMillX.GreedyForwardType
GreedyForward

Order strategy: no precomputed importance signal. At each step, actually try every remaining candidate and keep whichever single addition improves the objective most (sfs!/addone!).

ExplainMillX.prune!Function
prune!(mask::StructMask, ds, model, f, strategy::PruningStrategy) -> mask

Search mask for a small item subset keeping f() ≥ 0, using strategy. Mutates mask in place (via FlatView aliasing) and returns it, per Julia's !-function convention.

ds/model are accepted uniformly for every strategy even though only a future gradient-based order would use them, to avoid a breaking signature change later (docs/design/pruning.md §3.1).

Design contract

prune! requires the full mask (every item on) to already satisfy f() ≥ 0 – if it doesn't, that's an unsatisfiable request (almost always a misconfigured tolerance upstream), and prune! raises immediately rather than attempting a search that cannot succeed. Given that precondition holds, prune! is guaranteed to return a mask with f() ≥ 0; if it doesn't, that indicates either a bug in the search primitives or a violated assumption about f (namely, that including strictly more of the sample never makes the objective worse – true for confidence-gap-style objectives, not guaranteed for an arbitrary injected f). Either way this is a loud, unconditional error, not a value for the caller to check.

ExplainMillX.search!Function
search!(f, fv::FlatView, order, candidates)

Decide fv's items using order's strategy, restricted to candidates.

Masks

ExplainMillX.StructMaskType
StructMask{C,V}

Mirrors the shape of a Mill sample. One node plays one of three roles, determined by which fields are populated (not by subtyping):

  • leaf: own set, children === nothing (e.g. an ArrayNode)
  • hybrid: own set, children a single StructMask (e.g. a BagNode)
  • router: own === nothing, children a Tuple/NamedTuple of StructMask (e.g. a ProductNode)

V<:AbstractVector signals at the type level whether the mask is binary (V<:AbstractVector{Bool}) or differentiable (V<:AbstractVector{<:Real}).

See docs/design/masks.md for the full design rationale.

ExplainMillX.create_structmaskFunction
create_structmask(ds::AbstractMillNode, mk)

Build a StructMask mirroring ds. mk is a leaf factory d -> own_vector of length d; the caller selects binary vs. differentiable masks entirely through what mk returns (see docs/design/masks.md §3.2).

What own's units mean is storage-format specific, matching the granularity at which it is natural to explain that format:

  • dense (Matrix): one unit per row (feature) – shared across all observations
  • sparse (SparseMatrixCSC): one unit per stored nonzero value
  • categorical (MaybeHotMatrix) / n-gram (NGramMatrix): one unit per observation
  • BagNode: one unit per instance (its child's observations)
  • ProductNode: no own unit; routes to named children
ExplainMillX.applymaskFunction
applymask(ds::AbstractMillNode, m::StructMask)

Hard-prune ds according to m, returning a new (smaller / missing-valued) sample. Never mutates ds. This is the one place per-storage-format knowledge is unavoidable (see docs/design/masks.md §6) – each Mill leaf-array type needs its own method because a dense Matrix, a SparseMatrixCSC, a MaybeHotMatrix, and an NGramMatrix each need genuinely different code to represent "this unit is absent."

Masked-out values become missing, matching Mill's native imputation support (sparse entries become 0, since sparsity already has no missing representation).

ExplainMillX.leafmaskFunction
leafmask(own::AbstractVector) -> StructMask

Build a leaf StructMask (no children) with maskable units own.

ExplainMillX.hybridmaskFunction
hybridmask(own::AbstractVector, children) -> StructMask

Build a hybrid StructMask that both has its own maskable units own (e.g. a BagNode's instances) and a nested children mask.

ExplainMillX.routermaskFunction
routermask(children) -> StructMask

Build a router StructMask (no own units, e.g. a ProductNode) that only routes to named/positional children masks.

ExplainMillX.isleafFunction
isleaf(m::StructMask) -> Bool

true if m has no nested children (m.children === nothing).

ExplainMillX.isrouterFunction
isrouter(m::StructMask) -> Bool

true if m has no own maskable units (m.own === nothing), i.e. it only routes to children (e.g. a ProductNode's mask).

ExplainMillX.prunemaskFunction
prunemask(fv::FlatView)

The aggregated boolean "is this item currently on" view across every item in fv, in flat order.

ExplainMillX.softvalueFunction
softvalue(m::StructMask) -> AbstractVector{<:Real}

The continuous [0,1]-valued view of m's mask, for gradient-based strategies. Errors if m is a binary (Vector{Bool}) mask – use prunemask for the boolean view instead.

ExplainMillX.participateFunction
participate(m::StructMask) -> Vector{Bool}

Per-item reachability: true for units of m currently reachable from the root given the current state of ancestor masks. See updateparticipation!.

participate(fv::FlatView)

The aggregated boolean reachability view across every item in fv, in flat order (see docs/design/flatview.md §5).

ExplainMillX.randomize!Function
randomize!([rng,] m::StructMask)

Independently sample each maskable unit in m (and its whole subtree) uniformly at random. Used by Monte Carlo scoring strategies.

ExplainMillX.updateparticipation!Function
updateparticipation!(ds::AbstractMillNode, m::StructMask)

Recompute, for every node in m, whether its units are currently reachable given the current state of ancestor masks (see docs/design/masks.md §5). Resets everything to participating, then propagates invalidity top-down via invalidate!, which is dispatched on ds's type since the index-space translation between a node and its children (e.g. bag membership, sparse column ownership) is storage-format specific.

Must be called after any mutation to a mask's own vector, before reading .participate anywhere in the tree.

ExplainMillX.foreach_maskFunction
foreach_mask(f, m::StructMask)

The sole sanctioned way to walk a StructMask tree for side effects. Calls f(node, level) on every node that has an own mask (node.own !== nothing). Memoized via an identity-keyed IdDict, so shared sub-structure (the same node reachable through more than one parent) is visited exactly once.

ExplainMillX.mapmaskFunction
mapmask(f, m::StructMask)

Structure-preserving transform: calls f(own_vector, level) on every node's own mask vector and rebuilds the tree with the returned vectors, preserving participate and shape. Memoized on the identity of the own vector, so masks sharing the same underlying vector remain shared after mapping.

ExplainMillX.collectmasksFunction
collectmasks(m::StructMask)

Collect all own-bearing nodes together with their depth, as node => level pairs. Mirrors ExplainMill.jl's collect_masks_with_levels.

ExplainMillX.AccTreeType
AccTree{P,C}

Per-unit scoring bookkeeping, kept entirely separate from StructMask (see docs/design/masks.md §8). Mirrors the shape of the StructMask it was built from, carrying whatever payload a scoring strategy needs – the model never sees this type.

ExplainMillX.create_acctreeFunction
create_acctree(m::StructMask, make_payload) -> AccTree

Build an AccTree mirroring the shape of m, calling make_payload(d) (d = number of maskable units) at every own-bearing node to construct that node's bookkeeping payload.

ExplainMillX.foreach_pairedFunction
foreach_paired(f, m::StructMask, acc::AccTree)

Walk m and acc in lockstep (same recursion shape as foreach_mask), calling f(node, payload, level) on every own-bearing node. Lets a scoring strategy update its bookkeeping without the model ever being aware bookkeeping exists.

Pruning internals

ExplainMillX.FlatViewType
FlatView

A flat, linear, mutable index over a collection of own-bearing StructMask nodes. Aliases the real own vectors in place – each flat index stores a reference to the actual StructMask node it belongs to plus a local index into that node's own – so reading/writing through a FlatView reads and writes the real mask directly, with no copying or reconciliation step.

Deliberately blind to tree shape (it only needs "which nodes, in what order") and to scoring/AccTree (see nodescores/heuristicscores for how those connect back in). See docs/design/flatview.md for the full rationale.

ExplainMillX.useditemsFunction
useditems(fv::FlatView)

Flat indices currently on – findall(prunemask(fv)). This is what redundancy-removal and fine-tuning search primitives iterate over.

ExplainMillX.nodescoresFunction
nodescores(mask::StructMask, acc::AccTree, scorefn) -> IdDict{StructMask,Vector{Float64}}

Build the identity-keyed lookup from docs/design/flatview.md §3 ("Option B"): walk mask and acc together (via foreach_paired, so they must have been built from one another and share structure) and record, per own-bearing node – keyed by object identity, not value – the plain Float64 scores obtained by applying scorefn to that node's accumulator payload.

The returned dictionary is only valid against FlatViews built over the exact same mask object (or a subset of its nodes) acc was paired with; see heuristicscores for how it's consumed and what happens on a mismatch.

ExplainMillX.heuristicscoresFunction
heuristicscores(fv::FlatView, scores::IdDict{StructMask,Vector{Float64}}) -> Vector{Float64}

For every flat index in fv, look up the score recorded in scores for the StructMask node (by identity, ===) and local index that flat index maps to. Works unchanged whether fv spans a whole tree or a subset of it (e.g. one level, for level-by-level pruning), since scores is keyed by node identity rather than position.

Raises KeyError if a node in fv isn't present in scores – this means fv was built over a mask tree that isn't (object-)identical to the one scores was built from; see docs/design/flatview.md §3.1.

ExplainMillX.addminimumbi!Function
addminimumbi!(f, fv::FlatView, order::AbstractVector{<:Integer})

Turn on items along order (a permutation of 1:length(fv), typically sorted by descending importance) via bisection, until f() ≥ 0. O(log n) evaluations of f, relying on order being roughly monotonic in effect – does not itself guarantee f() ≥ 0 is reachable at all (if turning on every item in order still leaves f() < 0, this returns having done its best; surfacing that as non-convergence is the driver's responsibility, see docs/design/pruning.md §3.7).

ExplainMillX.addone!Function
addone!(f, fv::FlatView, candidates=1:length(fv))

Single-step greedy forward selection: try turning on every currently-off item in candidates, keep whichever addition gives the largest f(). O(n) evaluations. Returns true if something was turned on.

candidates defaults to unrestricted (every index) rather than filtering by participate(fv) internally – restricting to participating items is the driver's responsibility (pass findall(participate(fv)) explicitly for level-by-level pruning), since for flat-granularity pruning participate may hold stale state from a prior scoring pass and shouldn't be consulted at all; see docs/design/pruning.md §4 (Part 4).

ExplainMillX.removeone!Function
removeone!(f, fv::FlatView, candidates=useditems(fv))

Single-step greedy backward selection: try turning off every item in candidates (currently-on items by default), keep whichever removal gives the largest f(). O(n) evaluations. Returns true if something was turned off.

ExplainMillX.sfs!Function
sfs!(f, fv::FlatView, candidates=1:length(fv))

Stepwise forward selection: starting from everything off, repeatedly addone! (restricted to candidates) until f() ≥ 0 or no candidate remains to add.

ExplainMillX.randomremoval!Function
randomremoval!(f, fv::FlatView; rng=Random.default_rng())

Redundancy-removal to a fixed point: repeatedly shuffle the currently-on items and try turning each off (in that random order), keeping a removal only if f() ≥ 0 still holds; stops once a full pass removes nothing more.

Fix vs. the original: ExplainMill.jl's removeexcess! took a shuffled candidate order but iterated useditems(flatmask) instead of the order it was given, so its randomization had no effect. Here the shuffled order is the order actually iterated.

ExplainMillX.greedyremoval!Function
greedyremoval!(f, fv::FlatView)

Redundancy-removal to a fixed point, using removeone!'s greedy (least-damaging) choice at each step instead of a random order: repeatedly remove the single least-damaging currently-on item, stopping either when no removal keeps f() ≥ 0 (reverting that last, over-aggressive removal) or when nothing remains that can be removed at all.

ExplainMillX.finetune!Function
finetune!(f, fv::FlatView, max_n=typemax(Int), candidates=1:length(fv))

A small local-search pass after a main search: alternates batched add/remove (finetuneadd!/finetuneremove!, add-side restricted to candidates), tracks visited item-sets to detect cycling (growing the batch size n when a state repeats), and reverts to the smallest, best-scoring feasible (f() ≥ 0) state visited by the end.

ExplainMillX.settobest!Function
settobest!(fv::FlatView, visited::Dict{Vector{Int},<:Real})

Set fv to the smallest, best-scoring feasible (f() ≥ 0) item set among visited's keys. Does nothing if no visited state is feasible.