Make Captum interpretation optional

- Move Captum-backed method resolution behind lazy imports
- Keep interpretation method validation independent of Captum installation
- Add optional interpret dependency group for Captum support
- Update installation docs for core, interpret, and AnnData extras
- Preserve core package imports without requiring Captum

Author: Thomas-Rauter

docs/installation.md

@@ -2,16 +2,38 @@
 
 ## Install from PyPI
 
-Install `edge2torch` from PyPI with:
+Install the core `edge2torch` package from PyPI with:
 
 ```bash
 pip install edge2torch
 ```
 
-For optional `AnnData` support:
+The core installation supports compiling sparse PyTorch models from edgelists,
+aligning named input features, customizing compiled models, and training with
+ordinary PyTorch.
+
+## Optional interpretation support
+
+`interpret_model()` uses Captum and is installed as an optional dependency.
+
+To install `edge2torch` with interpretation support:
+
+```bash
+pip install "edge2torch[interpret]"
+```
+
+## Optional AnnData support
+
+For optional `AnnData` input support:
+
+```bash
+pip install "edge2torch[anndata]"
+```
+
+To install both interpretation and `AnnData` support:
 
 ```bash
-pip install "edge2torch[bio]"
+pip install "edge2torch[interpret,anndata]"
 ```
 
 ## Development installation
@@ -25,24 +47,36 @@ cd edge2torch
 pip install -e .
 ```
 
+For optional interpretation support during development:
+
+```bash
+pip install -e ".[interpret]"
+```
+
 For optional `AnnData` support during development:
 
 ```bash
-pip install -e .[bio]
+pip install -e ".[anndata]"
+```
+
+For both optional interpretation and `AnnData` support:
+
+```bash
+pip install -e ".[interpret,anndata]"
 ```
 
 ## Optional dependency groups
 
 Install development dependencies with:
 
 ```bash
-pip install -e .[dev]
+pip install -e ".[dev]"
 ```
 
 Install documentation dependencies with:
 
 ```bash
-pip install -e .[docs]
+pip install -e ".[docs]"
 ```
 
 ## Notebook and documentation note
@@ -59,8 +93,14 @@ sudo apt install graphviz
 
 ## Verify the installation
 
-A minimal smoke test is:
+A minimal core-installation smoke test is:
 
 ```bash
 python -c "import edge2torch; print('edge2torch imported successfully')"
 ```
+
+To verify interpretation support, install `edge2torch[interpret]` and run:
+
+```bash
+python -c "from edge2torch import interpret_model; print(interpret_model)"
+```

pyproject.toml

@@ -15,7 +15,6 @@ authors = [
 dependencies = [
     "pandas",
     "torch",
-    "captum",
 ]
 classifiers = [
     "Development Status :: 3 - Alpha",
@@ -35,7 +34,14 @@ Repository = "https://github.com/Thomas-Rauter/edge2torch"
 Issues = "https://github.com/Thomas-Rauter/edge2torch/issues"
 
 [project.optional-dependencies]
-bio = [
+all = [
+    "captum",
+    "anndata",
+]
+interpret = [
+    "captum",
+]
+anndata = [
     "anndata",
 ]
 dev = [
@@ -44,6 +50,7 @@ dev = [
     "ruff",
     "mypy",
     "build",
+    "captum",
     "anndata",
 ]
 docs = [
@@ -59,6 +66,8 @@ docs = [
     "IPython",
     "seaborn",
     "black",
+    "captum",
+    "anndata",
 ]
 
 [tool.setuptools]

src/edge2torch/interpret/captum_classes.py

@@ -0,0 +1,92 @@
+"""
+Lazy Captum class resolution for optional interpretation support.
+
+Why this file exists
+--------------------
+Captum is an optional dependency of edge2torch. This module imports Captum
+classes only when interpretation is actually executed.
+
+Role in the package
+-------------------
+This is an internal dependency bridge. It maps validated Captum method names to
+Captum attribution classes. It should not validate public inputs, execute
+attribution, prepare input data, or map attribution results.
+"""
+
+from typing import Any
+
+from ..utils.errors import Edge2TorchError
+
+
+def get_captum_class(method: str) -> type[Any]:
+    """
+    Return the Captum attribution class for a supported method name.
+    """
+    try:
+        from captum.attr import (
+            LRP,
+            Deconvolution,
+            DeepLift,
+            DeepLiftShap,
+            FeatureAblation,
+            FeaturePermutation,
+            GradientShap,
+            GuidedBackprop,
+            InputXGradient,
+            IntegratedGradients,
+            InternalInfluence,
+            KernelShap,
+            LayerActivation,
+            LayerConductance,
+            LayerDeepLift,
+            LayerDeepLiftShap,
+            LayerFeatureAblation,
+            LayerFeaturePermutation,
+            LayerGradientShap,
+            LayerGradientXActivation,
+            LayerIntegratedGradients,
+            LayerLRP,
+            Lime,
+            Occlusion,
+            Saliency,
+            ShapleyValues,
+            ShapleyValueSampling,
+        )
+    except ImportError as exc:
+        raise Edge2TorchError(
+            "interpret_model() requires Captum, which is an optional "
+            "dependency. Install interpretation support with "
+            "'pip install \"edge2torch[interpret]\"'."
+        ) from exc
+
+    classes: dict[str, type[Any]] = {
+        "IntegratedGradients": IntegratedGradients,
+        "DeepLift": DeepLift,
+        "DeepLiftShap": DeepLiftShap,
+        "GradientShap": GradientShap,
+        "Saliency": Saliency,
+        "InputXGradient": InputXGradient,
+        "GuidedBackprop": GuidedBackprop,
+        "Deconvolution": Deconvolution,
+        "FeatureAblation": FeatureAblation,
+        "Occlusion": Occlusion,
+        "FeaturePermutation": FeaturePermutation,
+        "ShapleyValueSampling": ShapleyValueSampling,
+        "ShapleyValues": ShapleyValues,
+        "Lime": Lime,
+        "KernelShap": KernelShap,
+        "LRP": LRP,
+        "LayerActivation": LayerActivation,
+        "LayerConductance": LayerConductance,
+        "InternalInfluence": InternalInfluence,
+        "LayerGradientXActivation": LayerGradientXActivation,
+        "LayerDeepLift": LayerDeepLift,
+        "LayerDeepLiftShap": LayerDeepLiftShap,
+        "LayerGradientShap": LayerGradientShap,
+        "LayerIntegratedGradients": LayerIntegratedGradients,
+        "LayerFeatureAblation": LayerFeatureAblation,
+        "LayerFeaturePermutation": LayerFeaturePermutation,
+        "LayerLRP": LayerLRP,
+    }
+
+    return classes[method]

src/edge2torch/interpret/feature_attribution.py

@@ -23,9 +23,10 @@
 
 from ..compile.artifact import CompileArtifact
 from ..utils.errors import Edge2TorchError
+from .captum_classes import get_captum_class
 from .method_registry import (
-    FEATURE_INTERPRETERS_WITH_CONSTRUCTOR_KWARGS,
-    FEATURE_INTERPRETERS_WITHOUT_CONSTRUCTOR_KWARGS,
+    FEATURE_METHODS_WITH_CONSTRUCTOR_KWARGS,
+    FEATURE_METHODS_WITHOUT_CONSTRUCTOR_KWARGS,
 )
 
 # Level 2 functions (called by level 1 functions) ------------------------------
@@ -95,17 +96,15 @@ def _build_feature_interpreter(
     """
     Build a Captum interpreter for feature-level attribution.
     """
-    if method in FEATURE_INTERPRETERS_WITH_CONSTRUCTOR_KWARGS:
-        interpreter_class = FEATURE_INTERPRETERS_WITH_CONSTRUCTOR_KWARGS[method]
+    if method in FEATURE_METHODS_WITH_CONSTRUCTOR_KWARGS:
+        interpreter_class = get_captum_class(method)
         return interpreter_class(
             model,
             **constructor_kwargs,
         )
 
-    if method in FEATURE_INTERPRETERS_WITHOUT_CONSTRUCTOR_KWARGS:
-        interpreter_class = FEATURE_INTERPRETERS_WITHOUT_CONSTRUCTOR_KWARGS[
-            method
-        ]
+    if method in FEATURE_METHODS_WITHOUT_CONSTRUCTOR_KWARGS:
+        interpreter_class = get_captum_class(method)
         return interpreter_class(model)
 
     raise Edge2TorchError(

src/edge2torch/interpret/feedforward_node_attribution.py

@@ -25,9 +25,10 @@
 from ..compile.artifact import CompileArtifact
 from ..utils.constants import INTERNAL_NODE_PREFIX
 from ..utils.errors import Edge2TorchError
+from .captum_classes import get_captum_class
 from .method_registry import (
-    FEEDFORWARD_NODE_INTERPRETERS_WITH_CONSTRUCTOR_KWARGS,
-    FEEDFORWARD_NODE_INTERPRETERS_WITHOUT_CONSTRUCTOR_KWARGS,
+    FEEDFORWARD_NODE_METHODS_WITH_CONSTRUCTOR_KWARGS,
+    FEEDFORWARD_NODE_METHODS_WITHOUT_CONSTRUCTOR_KWARGS,
 )
 
 # Level 3 functions (called by level 2 functions) ------------------------------
@@ -113,20 +114,16 @@ def _build_feedforward_layer_interpreter(
     """
     Build a Captum interpreter for feedforward layer-level attribution.
     """
-    if method in FEEDFORWARD_NODE_INTERPRETERS_WITH_CONSTRUCTOR_KWARGS:
-        interpreter_class = (
-            FEEDFORWARD_NODE_INTERPRETERS_WITH_CONSTRUCTOR_KWARGS[method]
-        )
+    if method in FEEDFORWARD_NODE_METHODS_WITH_CONSTRUCTOR_KWARGS:
+        interpreter_class = get_captum_class(method)
         return interpreter_class(
             model,
             layer_block,
             **constructor_kwargs,
         )
 
-    if method in FEEDFORWARD_NODE_INTERPRETERS_WITHOUT_CONSTRUCTOR_KWARGS:
-        interpreter_class = (
-            FEEDFORWARD_NODE_INTERPRETERS_WITHOUT_CONSTRUCTOR_KWARGS[method]
-        )
+    if method in FEEDFORWARD_NODE_METHODS_WITHOUT_CONSTRUCTOR_KWARGS:
+        interpreter_class = get_captum_class(method)
         return interpreter_class(
             model,
             layer_block,

src/edge2torch/interpret/input_validation.py

@@ -29,10 +29,10 @@
 
 from ..utils.errors import Edge2TorchError
 from .method_registry import (
-    FEATURE_INTERPRETERS_WITHOUT_CONSTRUCTOR_KWARGS,
     FEATURE_METHODS,
-    FEEDFORWARD_NODE_INTERPRETERS_WITHOUT_CONSTRUCTOR_KWARGS,
+    FEATURE_METHODS_WITHOUT_CONSTRUCTOR_KWARGS,
     FEEDFORWARD_NODE_METHODS,
+    FEEDFORWARD_NODE_METHODS_WITHOUT_CONSTRUCTOR_KWARGS,
     SUPPORTED_METHODS,
 )
 
@@ -155,7 +155,7 @@ def _validate_interpret_options(
 
     if (
         target == "features"
-        and method in FEATURE_INTERPRETERS_WITHOUT_CONSTRUCTOR_KWARGS
+        and method in FEATURE_METHODS_WITHOUT_CONSTRUCTOR_KWARGS
         and constructor_kwargs
     ):
         raise Edge2TorchError(
@@ -165,7 +165,7 @@ def _validate_interpret_options(
 
     if (
         target == "nodes"
-        and method in FEEDFORWARD_NODE_INTERPRETERS_WITHOUT_CONSTRUCTOR_KWARGS
+        and method in FEEDFORWARD_NODE_METHODS_WITHOUT_CONSTRUCTOR_KWARGS
         and constructor_kwargs
     ):
         raise Edge2TorchError(