FDFI - Flow-Disentangled Feature Importance

A Python library for computing feature importance using disentangled methods, inspired by SHAP.

📖 Read the documentation

Current development version: 0.1.0

Overview

FDFI (Flow-Disentangled Feature Importance) is a Python module that provides interpretable machine learning explanations through disentangled feature importance methods. This package implements both OT-based DFI and flow-based FDFI methods.

Features

• 🎯 Three explainer variants, all model-agnostic and sharing one API: OTExplainer (Gaussian OT), EOTExplainer (a Gaussian approximation to entropic OT), and FlowExplainer (flow matching)
• 🔁 Cross-fitted inference: Crossfitting wraps any variant for valid standard errors at small sample sizes
• 📊 Rich Visualizations: Summary, waterfall, force, and dependence plots
• 🔧 Easy to Use: Simple API similar to SHAP
• 🧪 Statistical Inference: Confidence intervals and multiple testing correction (FDR/FWER)
• 🚀 Extensible: Built with modularity in mind for future enhancements

Installation

From PyPI

pip install fdfi
pip install "fdfi[flow]"

From Source

git clone https://github.com/jinhongdu-lab/FDFI.git
cd FDFI
pip install -e .

Dependencies

Use pyproject.toml extras:

pip install -e ".[dev]"
pip install -e ".[flow]"

Quick Start

import numpy as np
from fdfi.explainers import OTExplainer

# Define your model
def model(X):
    return X.sum(axis=1)

# Create background data
X_background = np.random.randn(100, 10)

# Create an explainer
explainer = OTExplainer(model, data=X_background, nsamples=50)

# Explain test instances
X_test = np.random.randn(10, 10)
results = explainer(X_test)

# Confidence intervals (post-hoc)
ci = explainer.conf_int(alpha=0.05, target="X", alternative="two-sided")

# With multiple testing correction (e.g., FDR control)
ci_fdr = explainer.conf_int(multitest_method="fdr_bh")
explainer.summary(multitest_method="fdr_bh")

Visualization

FDFI includes static Matplotlib plotting helpers for global scores, per-observation contributions, confidence intervals, diagnostics, and feature correlation.

from fdfi.plots import (
    confidence_interval_plot,
    correlation_heatmap,
    diagnostics_plot,
    summary_bar,
    summary_plot,
)

feature_names = [f"X{i}" for i in range(X_background.shape[1])]

# Background correlation structure
correlation_heatmap(X_background, feature_names, show=False)

# Global scores and standard errors from explainer output
summary_bar(results["phi_X"], results["se_X"], feature_names, show=False)

# Distribution of per-observation contributions
summary_plot(explainer.ueifs_X, features=X_test, feature_names=feature_names, show=False)

# Inference and quality checks
confidence_interval_plot(ci, feature_names=feature_names, show=False)
diagnostics_plot(explainer.diagnostics, feature_names=feature_names, show=False)

Confidence-interval defaults

By default, conf_int() uses:

• var_floor_method="mixture"
• margin_method="auto": the gap rule is used for fewer than 30 units; otherwise, the mixture rule is used.
• threshold_null=True: used for group-level inference.

This improves stability for weak effects and avoids ad hoc thresholding in many use cases. You can still override both methods explicitly if needed.

EOT Options (Entropic OT)

EOTExplainer supports a fixed or adaptive epsilon, three sampling methods, CPI or SCPI scores, and user-chosen losses:

from fdfi.explainers import EOTExplainer

explainer = EOTExplainer(
    model.predict,
    X_background,
    epsilon=0.1,                 # or auto_epsilon=True to choose epsilon from the data
    sampling_method="resample",  # 'resample', 'permutation', or 'normal'
    method="cpi",                # 'cpi' or 'scpi'
    loss="squared_error",
)
results = explainer(X_test)

Flow-DFI with FlowExplainer

FlowExplainer learns a nonlinear map by flow matching, supporting both CPI (Conditional Permutation Importance) and SCPI (Sobol-CPI):

• CPI: Half the average per-resample loss difference: $\frac12 E_b[L(Y,f(\tilde X_b))-L(Y,f(X))]$.
• SCPI: Average counterfactual predictions before applying the loss: $L(Y,E_b[f(\tilde X_b)])-L(Y,f(X))$.

This is the normalization and naming used in the FDFI paper. Conventional CPI without the factor $1/2$ is twice the package's CPI. Under squared-error loss, an exact disentangling map, and a Bayes predictor, CPI and the infinite-resample SCPI have the same population target. With the same finite resamples, $2,\widehat\phi^{CPI}=\widehat\phi^{SCPI}+\mathrm{Var}_b[f(\tilde X_b)]$.

from fdfi.explainers import FlowExplainer

# Create explainer with CPI (default)
explainer = FlowExplainer(
    model.predict,
    X_background,
    fit_flow=True,
    method='cpi',     # 'cpi', 'scpi', or 'both'
    num_steps=200,    # flow training steps
    nsamples=50,      # counterfactual samples
    sampling_method='resample',  # 'resample', 'permutation', 'normal', 'condperm'
)

results = explainer(X_test)
# results['phi_Z']: Z-space importance
# results['phi_X']: same as phi_Z (Z-space methods)

# Confidence intervals
ci = explainer.conf_int(alpha=0.05, target="Z", alternative="two-sided")

Explainer diagnostics

Disentangled explainers (OTExplainer, EOTExplainer, and FlowExplainer) report two diagnostics with qualitative labels (GOOD / MODERATE / POOR) using consistent [FDFI][DIAG] logging:

• Latent independence (median dCor) — lower is better (thresholds: <0.10 good, <0.25 moderate).
• Distribution fidelity (MMD) — lower is better (thresholds: <0.05 good, <0.15 moderate).

Example log:

[FDFI][DIAG] Flow Model Diagnostics
[FDFI][DIAG] Latent independence (median dCor): 0.0421 [GOOD]  → lower is better
[FDFI][DIAG] Distribution fidelity (MMD):       0.0187 [GOOD]  → lower is better

Access diagnostics directly:

diag = explainer.diagnostics
print(diag["latent_independence_median"], diag["latent_independence_label"])
print(diag["distribution_fidelity_mmd"], diag["distribution_fidelity_label"])

For advanced users, flow models can be trained separately:

from fdfi.models import FlowMatchingModel

# Train flow model externally
flow_model = FlowMatchingModel(X_background, dim=X_background.shape[1])
flow_model.fit(num_steps=500, verbose='final')

# Set pre-trained flow
explainer = FlowExplainer(model.predict, X_background, fit_flow=False)
explainer.set_flow(flow_model)

Project Structure

FDFI/
├── fdfi/                  # Main package directory
│   ├── __init__.py       # Package initialization
│   ├── explainers.py     # Explainer classes
│   ├── losses.py         # Loss registry (regression + classification)
│   ├── models.py         # FlowMatchingModel for FlowExplainer
│   ├── plots.py          # Visualization functions
│   └── utils.py          # Utility functions
├── tests/                 # Test suite
├── docs/                  # Documentation & tutorials
│   ├── tutorials/        # Jupyter notebook tutorials
│   └── case_studies/     # Applied analyses: HIV-1 VRC01 neutralization, UCI CTG
├── pyproject.toml        # Package configuration
└── README.md            # This file

Development Status

FDFI is under active research development. The package includes implemented OT/EOT/Flow explainers, statistical inference helpers, diagnostics, plotting utilities, tests, and documentation. Some advanced modeling components continue to evolve as the methodology develops.

Testing

Run the test suite:

# Install development dependencies
pip install -e ".[dev]"

# Run tests
pytest

# Run tests with coverage
pytest --cov=fdfi --cov-report=html

Documentation

Full documentation is hosted at fdfi.readthedocs.io, or build it locally:

pip install -e ".[docs]"
cd docs && python -m sphinx -b html . _build/html
open _build/html/index.html

Tutorial notebooks live in docs/tutorials/:

• Quickstart Tutorial
• OT Explainer Tutorial
• EOT Explainer Tutorial
• Flow Explainer Tutorial
• Confidence Intervals
• Visualization Tutorial

Contributing

Contributions are welcome! Please feel free to submit a Pull Request.

License

This project is licensed under the MIT License - see the LICENSE file for details.

References

FDFI is based on:

• Du, J.-H., Roeder, K., & Wasserman, L. (2025). Disentangled Feature Importance. arXiv preprint arXiv:2507.00260.
• Chen, X., Guo, Y., & Du, J.-H. (2026). Flow-Disentangled Feature Importance. In The Thirteenth International Conference on Learning Representations (ICLR).

Related work:

• SHAP: A game theoretic approach to explain machine learning models

Citation

If you use FDFI in your research, please cite:

@software{dfi2026,
  title={FDFI: Python Library for Disentangled Feature Importance},
  author={FDFI Team},
  year={2026},
  url={https://github.com/jinhongdu-lab/FDFI}
}

@article{du2025disentangled,
  title={Disentangled Feature Importance},
  author={Du, Jin-Hong and Roeder, Kathryn and Wasserman, Larry},
  journal={arXiv preprint arXiv:2507.00260},
  year={2025}
}

@inproceedings{chen2026flow,
  title={Flow-Disentangled Feature Importance},
  author={Chen, Xingshu and Guo, Yifeng and Du, Jin-Hong},
  booktitle={International Conference on Learning Representations},
  volume={2026},
  pages={57346--57388},
  year={2026}
}

Contact

For questions and issues, please use the GitHub issue tracker.