SHAPTreeExplainer additivity check fails using unconstrained RandomForestClassifier

[SamuelHomberg]

Describe the bug

When using SHAPTreeExplainer on a molpipeline.Pipeline with an unconstrained RandomForestClassifier (i.e., default max_depth=None), the explanation fails with an Additivity check failed in TreeExplainer! error for larger datasets.

Limiting the tree depth explicitly using max_depth (e.g., max_depth=15) resolves the issue.

Suggested Resolution

It would be highly beneficial to add a warning or a documentation note inside SHAPTreeExplainer advising users against using unconstrained RandomForestClassifier (or similar tree-based models) on large molecular fingerprint datasets, suggesting they explicitly set a reasonable max_depth.

MWE

Click here to expand

import numpy as np
from rdkit import Chem
from sklearn.ensemble import RandomForestClassifier
from molpipeline import Pipeline
from molpipeline.any2mol import AutoToMol
from molpipeline.mol2any import MolToMorganFP
from molpipeline.experimental.explainability import SHAPTreeExplainer

def create_synthetic_data(n_samples: int, n_bits: int = 2048, n_rules: int = 5, noise_level: float = 0.1, seed: int = 42):
    """
    Creates a reproducible, realistic synthetic dataset of SMILES and labels.

    The labels are determined by the presence of a few "magic" fingerprint bits,
    plus some noise. This creates a learnable signal that encourages deep tree growth.
    """
    rng = np.random.default_rng(seed)
    smiles = []
    # Generate diverse SMILES strings
    for _ in range(n_samples):
        n_c = rng.integers(5, 25)
        n_o = rng.integers(0, 5)
        n_n = rng.integers(0, 5)
        s = "C" * n_c + "O" * n_o + "N" * n_n
        s_list = list(s)
        rng.shuffle(s_list)
        smiles.append("".join(s_list))

    mols = [Chem.MolFromSmiles(s) for s in smiles]

    # Create fingerprints to base the labels on
    fp_gen = MolToMorganFP(n_bits=n_bits, radius=3, return_as="dense")
    fingerprints = fp_gen.fit_transform(mols)

    # Select some random fingerprint bits that will determine the label
    magic_indices = rng.choice(n_bits, size=n_rules, replace=False)

    # Create labels based on the presence of magic bits
    labels = np.zeros(n_samples, dtype=int)
    for i, fp in enumerate(fingerprints):
        # The label is 1 if at least one magic bit is present
        if np.sum(fp[magic_indices]) > 0:
            labels[i] = 1

    # Add some noise to make it more realistic
    n_flips = int(n_samples * noise_level)
    flip_indices = rng.choice(n_samples, size=n_flips, replace=False)
    labels[flip_indices] = 1 - labels[flip_indices]

    return mols, labels


def run_experiment(model, mols, labels, model_name=""):
    """Trains a pipeline and runs the SHAP explainer."""
    print(f"\n--- Running experiment with {model_name} ---")
    
    pipeline = Pipeline(
        [
            ("auto2mol", AutoToMol()),
            ("morgan", MolToMorganFP(n_bits=2048, radius=3, return_as="dense")),
            ("model", model),
        ],
    )
    
    print("Fitting pipeline...")
    pipeline.fit(mols, labels)
    
    print("Creating SHAP explainer...")
    explainer = SHAPTreeExplainer(pipeline)
    
    try:
        print("Try explaining...")
        # We only need to explain a single molecule to trigger the bug
        explanations = explainer.explain(mols[:1])
        print(f"SUCCESS: SHAP explanation completed for {model_name}.")
        return True
    except Exception as e:
        print(f"FAILURE: SHAP explanation failed for {model_name}.")
        print(f"Error: {e}")
        return False

def main():
    """
    Demonstrates the SHAPTreeExplainer failure with deep RandomForest trees
    and the success with a max_depth limit.
    """
    # A larger number of samples is required to grow deep trees and trigger the issue.
    n_samples = 5000
    print(f"Generating {n_samples} synthetic molecules...")
    mols, labels = create_synthetic_data(n_samples, n_bits=2048, seed=42)

    # 1. Demonstrate the failure with default RandomForestClassifier
    #    This creates very deep trees, leading to a recursion error in SHAP.
    rf_unconstrained = RandomForestClassifier(n_estimators=100, random_state=42, n_jobs=-1)
    run_experiment(rf_unconstrained, mols, labels, "Unconstrained RandomForest")

    # 2. Show the workaround by constraining the tree depth
    #    This prevents the recursion depth error.
    rf_constrained = RandomForestClassifier(n_estimators=100, max_depth=15, random_state=42, n_jobs=-1)
    run_experiment(rf_constrained, mols, labels, "Constrained RandomForest (max_depth=15)")

if __name__ == "__main__":
    main()

MWE output:

Generating 5000 synthetic molecules...

--- Running experiment with Unconstrained RandomForest ---
Fitting pipeline...
Creating SHAP explainer...
Try explaining...
FAILURE: SHAP explanation failed for Unconstrained RandomForest.
Error: Additivity check failed in TreeExplainer! Please ensure the data matrix you passed to the explainer is the 
same shape that the model was trained on. If your data shape is correct then please report this on GitHub. 
Consider retrying with the feature_perturbation='interventional' option. This check failed because for one of 
the samples the sum of the SHAP values was [sum_val=array([-6.90695819e+11]) ind=np.int64(0)], while 
the model output was [model_output=array([0.99])]. If this difference is acceptable you can set 
check_additivity=False to disable this check.

--- Running experiment with Constrained RandomForest (max_depth=15) ---
Fitting pipeline...
Creating SHAP explainer...
Try explaining...
SUCCESS: SHAP explanation completed for Constrained RandomForest (max_depth=15).

Note: At some point the indexing in the error message itself (shap/explainers/_tree.py:906f" was {sum_val[ind]:f}, while the model output was {model_output[ind]:f}. If this") lead to another error regarding indexing, however, I am currently unable to reproduce it.