Source code for modeva.models.directrs.icl_moe_api

"""
MoDeVa model wrappers for ICL-MoE on DirectRS.

MoDirectRSICLRegressor — ICL-MoE post-processor for regression.
MoDirectRSICLClassifier — ICL-MoE post-processor for binary classification.

All interpretability methods return ``ValidationResult`` objects following
the MoDeVa standard.  Use ``result.plot()`` to visualize and ``result.value``
to access the raw numerical data.
"""

import numpy as np
import pandas as pd
from scipy.special import expit as _sigmoid
from sklearn.base import RegressorMixin, ClassifierMixin

from ..base import ModelBaseRegressor, ModelBaseClassifier
from ...utils.results import ValidationResult
from .icl_moe import ICLMoECore
from . import icl_moe_visualize as iviz


class _ICLMoEInterpretMixin:
    """Shared interpretability methods for ICL-MoE models.

    All methods call the underlying ``ICLMoECore`` for computation,
    build mocharts visualizations via ``icl_moe_visualize.py``, and wrap
    everything in ``ValidationResult``.
    """

    def explain_local(self, X, feature_names=None, sample_index=0):
        """Per-sample additive decomposition in original feature space.

        Parameters
        ----------
        X : array-like of shape (n_samples, n_features)
        feature_names : list of str, optional
        sample_index : int, default=0
            Sample to show in the default plot.

        Returns
        -------
        ValidationResult
            ``result.value`` contains all samples (intercept, contributions
            arrays).  ``result.plot()`` shows the selected sample.
        """
        raw = self._icl_core.explain_local(X, feature_names)
        names = raw['feature_names']
        c = raw['contributions'][sample_index]

        table = pd.DataFrame({
            "Name": names,
            "Effect": c,
        })

        options = iviz.viz_local_explanation(raw, sample_index)

        return ValidationResult(
            key="icl_moe_local",
            model=self.name,
            inputs={"sample_index": sample_index,
                    "feature_names": feature_names},
            value=raw,
            table=table,
            options=options,
        )

    def importance_global(self, X, feature_names=None, mode="contrib_abs"):
        """Global feature importance via explain_local aggregation.

        Parameters
        ----------
        X : array-like of shape (n_samples, n_features)
        feature_names : list of str, optional
        mode : str, default="contrib_abs"
            "contrib_abs" or "contrib_rms".

        Returns
        -------
        ValidationResult
            Horizontal bar plot of feature importance.
        """
        raw = self._icl_core.importance_global(X, feature_names, mode)
        names = raw['feature_names']

        table = pd.DataFrame({
            "Name": names,
            "Importance": raw['importance'],
        })

        options = iviz.viz_importance_global(raw)

        return ValidationResult(
            key="icl_moe_importance",
            model=self.name,
            inputs={"mode": mode, "feature_names": feature_names},
            value=raw,
            table=table,
            options=options,
        )

    def get_neighbor_analysis(self, X, sample_index=0, feature_names=None):
        """Return neighbor analysis for a query point.

        Parameters
        ----------
        X : array-like of shape (n_samples, n_features)
        sample_index : int, default=0
        feature_names : list of str, optional

        Returns
        -------
        ValidationResult
            Accessible plots via ``result.plot('weight_bar')`` and
            ``result.plot('neighbor_scatter')``, or ``result.plot()`` for all.
        """
        raw = self._icl_core.get_neighbor_analysis(
            X, sample_index, feature_names)

        k = len(raw['neighbor_weights'])
        table = pd.DataFrame({
            "Neighbor": [f"n{i}" for i in raw['neighbor_indices']],
            "Weight": raw['neighbor_weights'],
            "Target": raw['neighbor_y'],
        })

        options = iviz.viz_neighbor_analysis(raw)

        return ValidationResult(
            key="icl_moe_neighbors",
            model=self.name,
            inputs={"sample_index": sample_index,
                    "feature_names": feature_names},
            value=raw,
            table=table,
            options=options,
        )

    def get_leaf_gating_analysis(self, X, sample_index=0, tree_index=0,
                                 feature_names=None):
        """Leaf gating weights for a query point.

        Parameters
        ----------
        X : array-like of shape (n_samples, n_features)
        sample_index : int, default=0
        tree_index : int, default=0
        feature_names : list of str, optional

        Returns
        -------
        ValidationResult
            Bar plot of per-leaf gating weights.
        """
        raw = self._icl_core.get_leaf_gating_analysis(
            X, sample_index, tree_index, feature_names)

        n_leaves = len(raw['leaf_weights'])
        table = pd.DataFrame({
            "Leaf": [f"leaf_{i}" for i in range(n_leaves)],
            "Weight": raw['leaf_weights'],
            "Prediction": raw['leaf_predictions'],
            "Cover": raw['leaf_covers'],
        })

        options = iviz.viz_leaf_gating(raw)

        return ValidationResult(
            key="icl_moe_leaf_gating",
            model=self.name,
            inputs={"sample_index": sample_index,
                    "tree_index": tree_index,
                    "feature_names": feature_names},
            value=raw,
            table=table,
            options=options,
        )


[docs] class MoDirectRSICLRegressor(_ICLMoEInterpretMixin, RegressorMixin, ModelBaseRegressor): """ICL-MoE post-processor for DirectRS regressors. Takes a fitted ``MoDirectRSRegressor`` and builds a soft-gated mixture of experts using kNN attention in the DirectRS stretch embedding space. Parameters ---------- directrs_model : MoDirectRSRegressor A fitted DirectRS regressor (must have ``._core`` set). name : str, optional Model identifier. Default: "ICL-MoE". variant : str, default="hierarchical" One of "point_expert", "local_linear", "leaf_expert", "hierarchical". k : int, default=50 Number of nearest neighbors for kNN gating. tau : float, default=1.0 Temperature for softmax gating weights. ridge_lambda : float, default=1.0 Regularisation for local linear / residual ridge. top_m : int, default=5 Number of top leaf experts per tree (variant C / hierarchical). """ def __init__(self, directrs_model, name=None, variant="hierarchical", k=50, tau=1.0, ridge_lambda=1.0, top_m=5): self.directrs_model = directrs_model self.name = name or "ICL-MoE" self.variant = variant self.k = k self.tau = tau self.ridge_lambda = ridge_lambda self.top_m = top_m self._icl_core = None
[docs] def fit(self, X, y, sample_weight=None, verbose=False): """Fit ICL-MoE on training data. Parameters ---------- X : np.ndarray of shape (n_samples, n_features) y : np.ndarray of shape (n_samples,) sample_weight : ignored (kept for API compatibility) verbose : bool, default=False Returns ------- self """ X, y, _ = self._validate_fit_inputs(X, y, sample_weight) if self.directrs_model._core is None: raise ValueError("directrs_model must be fitted before " "creating ICL-MoE. Call directrs_model.fit() first.") self._icl_core = ICLMoECore( self.directrs_model._core, variant=self.variant, k=self.k, tau=self.tau, ridge_lambda=self.ridge_lambda, top_m=self.top_m, ) self._icl_core.fit(X, y, verbose=verbose, classification=False) return self
def _predict(self, X): return self._icl_core.predict(X)
[docs] class MoDirectRSICLClassifier(_ICLMoEInterpretMixin, ClassifierMixin, ModelBaseClassifier): """ICL-MoE post-processor for DirectRS classifiers. Takes a fitted ``MoDirectRSClassifier`` and builds a soft-gated mixture of experts using kNN attention in the DirectRS stretch embedding space. Operates in logit space internally; predictions use sigmoid. Parameters ---------- directrs_model : MoDirectRSClassifier A fitted DirectRS classifier (must have ``._core`` set). name : str, optional Model identifier. Default: "ICL-MoE-Cls". variant : str, default="hierarchical" One of "point_expert", "local_linear", "leaf_expert", "hierarchical". k : int, default=50 Number of nearest neighbors for kNN gating. tau : float, default=1.0 Temperature for softmax gating weights. ridge_lambda : float, default=1.0 Regularisation for local linear / residual ridge. top_m : int, default=5 Number of top leaf experts per tree (variant C / hierarchical). """ def __init__(self, directrs_model, name=None, variant="hierarchical", k=50, tau=1.0, ridge_lambda=1.0, top_m=5): self.directrs_model = directrs_model self.name = name or "ICL-MoE-Cls" self.variant = variant self.k = k self.tau = tau self.ridge_lambda = ridge_lambda self.top_m = top_m self._icl_core = None
[docs] def fit(self, X, y, sample_weight=None, verbose=False): """Fit ICL-MoE on training data for classification. Parameters ---------- X : np.ndarray of shape (n_samples, n_features) y : np.ndarray of shape (n_samples,) sample_weight : ignored verbose : bool, default=False Returns ------- self """ X, y, _ = self._validate_fit_inputs(X, y, sample_weight) if self.directrs_model._core is None: raise ValueError("directrs_model must be fitted before " "creating ICL-MoE. Call directrs_model.fit() first.") self._icl_core = ICLMoECore( self.directrs_model._core, variant=self.variant, k=self.k, tau=self.tau, ridge_lambda=self.ridge_lambda, top_m=self.top_m, ) self._icl_core.fit(X, y, verbose=verbose, classification=True) return self
def _predict_proba(self, X): """Predict probabilities via sigmoid on logit-space output.""" raw = self._icl_core.predict(X) prob_1 = _sigmoid(raw) return np.column_stack([1.0 - prob_1, prob_1])