{
  "cells": [
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "\n# ICL-MoE Regression\n\nThis example demonstrates how to use ICL-MoE to post-process a fitted DirectRS\nregression model. ICL-MoE adds a kNN-based residual correction on top of\nDirectRS leaf experts, providing local adaptation in the stretched embedding\nspace while maintaining exact additive interpretability.\n\nWe use the CaliforniaHousing dataset with a depth-2 XGBoost base model.\n"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "## Setup\nImport libraries and suppress warnings.\n\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "collapsed": false
      },
      "outputs": [],
      "source": [
        "import warnings\nwarnings.filterwarnings(\"ignore\")\n\nimport numpy as np\nfrom modeva import DataSet, TestSuite\nfrom modeva.models import MoXGBRegressor"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "## Load Dataset\nLoad the CaliforniaHousing dataset and create a random train/test split.\n\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "collapsed": false
      },
      "outputs": [],
      "source": [
        "ds = DataSet()\nds.load(name=\"CaliforniaHousing\")\nds.set_random_split()"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "## Train Base Model\nTrain an XGBoost regressor with depth 2 (required for FANOVA comparison).\n\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "collapsed": false
      },
      "outputs": [],
      "source": [
        "model = MoXGBRegressor(\n    name=\"XGB-depth2\",\n    n_estimators=200, max_depth=2, learning_rate=0.1,\n    subsample=0.8, colsample_bytree=0.8,\n    random_state=42, verbosity=0\n)\nmodel.fit(ds.train_x, ds.train_y.ravel())\n\nts = TestSuite(ds, model)\nts.diagnose_accuracy_table().table"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "## Fit DirectRS\nPost-process the trained XGBoost model with DirectRS.\n\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "collapsed": false
      },
      "outputs": [],
      "source": [
        "from modeva.models import MoDirectRSRegressor\n\ndrs = MoDirectRSRegressor(\n    base_model=model, ridge_alpha=100.0, n_passes=1\n)\ndrs.fit(\n    ds.train_x, ds.train_y.ravel(),\n    X_val=ds.test_x, y_val=ds.test_y.ravel(),\n    verbose=True\n)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "## Fit ICL-MoE\nBuild ICL-MoE on top of the fitted DirectRS model. The hierarchical\nvariant combines leaf expert predictions with a kNN residual correction.\n\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "collapsed": false
      },
      "outputs": [],
      "source": [
        "from modeva.models import MoDirectRSICLRegressor\n\nicl = MoDirectRSICLRegressor(\n    directrs_model=drs, k=50, tau=1.0, ridge_lambda=1.0\n)\nicl.fit(ds.train_x, ds.train_y.ravel(), verbose=True)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "## Accuracy Comparison\nCompare R-squared and MAE between XGBoost, DirectRS, and ICL-MoE.\n\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "collapsed": false
      },
      "outputs": [],
      "source": [
        "from sklearn.metrics import r2_score, mean_absolute_error\n\ny_train = ds.train_y.ravel()\ny_test = ds.test_y.ravel()\n\nprint(f\"{'Model':<12s} {'Train R2':>10s} {'Test R2':>10s} {'Train MAE':>10s} {'Test MAE':>10s}\")\nprint(\"-\" * 54)\nfor name, m in [(\"XGBoost\", model), (\"DirectRS\", drs), (\"ICL-MoE\", icl)]:\n    tr = m.predict(ds.train_x)\n    te = m.predict(ds.test_x)\n    print(f\"{name:<12s} {r2_score(y_train, tr):>10.4f} {r2_score(y_test, te):>10.4f} \"\n          f\"{mean_absolute_error(y_train, tr):>10.4f} {mean_absolute_error(y_test, te):>10.4f}\")"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "## Local Explanation\nICL-MoE provides exact additive decomposition: f(x) = c_0 + sum_j c_j.\nWe verify the decomposition matches predictions to machine precision.\n\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "collapsed": false
      },
      "outputs": [],
      "source": [
        "result = icl.explain_local(ds.test_x, feature_names=ds.feature_names)\nlocal = result.value\n\npred = icl.predict(ds.test_x)\nrecon = local['intercept'] + local['contributions'].sum(axis=1)\nmax_err = np.max(np.abs(pred - recon))\n\nprint(f\"Max |predict - (intercept + sum contributions)|: {max_err:.2e}\")\nprint(f\"Decomposition exact to machine precision: {max_err < 1e-10}\")"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "Local explanation waterfall plot for sample 0.\n\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "collapsed": false
      },
      "outputs": [],
      "source": [
        "result.plot()"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "## Global Feature Importance\nCompute global feature importance using mean absolute contributions.\n\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "collapsed": false
      },
      "outputs": [],
      "source": [
        "result = icl.importance_global(\n    ds.test_x, feature_names=ds.feature_names, mode=\"contrib_abs\"\n)\nresult.plot()"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "## Importance Comparison\nCompare ICL-MoE, DirectRS, and FANOVA feature importance side-by-side.\n\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "collapsed": false
      },
      "outputs": [],
      "source": [
        "icl_imp = icl.importance_global(\n    ds.test_x, feature_names=ds.feature_names, mode=\"contrib_abs\"\n).value['importance']\n\ndrs_imp = drs.importance_global(\n    ds.test_x, feature_names=ds.feature_names, mode=\"contrib_abs\"\n).value['importance']\n\nresult_fanova = ts.interpret_fi()\nfanova_table = result_fanova.table\nfanova_imp = dict(zip(fanova_table[\"Name\"], fanova_table[\"Score\"]))\n\nprint(f\"{'Feature':<12s} {'ICL-MoE':>10s} {'DirectRS':>10s} {'FANOVA':>10s}\")\nprint(\"-\" * 44)\nfor i, feat in enumerate(ds.feature_names):\n    print(f\"{feat:<12s} {icl_imp[i]:>10.4f} {drs_imp[i]:>10.4f} {fanova_imp.get(feat, 0.0):>10.4f}\")"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "FANOVA feature importance plot.\n\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "collapsed": false
      },
      "outputs": [],
      "source": [
        "result_fanova.plot()"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "FANOVA interaction importance plot.\n\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "collapsed": false
      },
      "outputs": [],
      "source": [
        "result_ei = ts.interpret_ei()\nresult_ei.plot()"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "## Neighbour Analysis\nInspect the kNN neighbourhood used for residual correction.\n\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "collapsed": false
      },
      "outputs": [],
      "source": [
        "result = icl.get_neighbor_analysis(\n    ds.test_x, sample_index=0, feature_names=ds.feature_names\n)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "Neighbour weights.\n\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "collapsed": false
      },
      "outputs": [],
      "source": [
        "result.plot(\"weight_bar\")"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "PCA scatter of neighbours in z-space.\n\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "collapsed": false
      },
      "outputs": [],
      "source": [
        "result.plot(\"neighbor_scatter\")"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "## Leaf Gating Analysis\nExamine per-leaf routing for a specific tree.\n\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "collapsed": false
      },
      "outputs": [],
      "source": [
        "result = icl.get_leaf_gating_analysis(\n    ds.test_x, sample_index=0, tree_index=0,\n    feature_names=ds.feature_names\n)\n\nactive = result.value['active_leaf']\nprint(f\"Active leaf: {active}\")\nprint(f\"Tree 0 contribution: {result.value['prediction']:.4f}\")\nprint(f\"Number of leaves: {len(result.value['leaf_weights'])}\")"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "Leaf gating weights for tree 0.\n\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "collapsed": false
      },
      "outputs": [],
      "source": [
        "result.plot()"
      ]
    }
  ],
  "metadata": {
    "kernelspec": {
      "display_name": "Python 3",
      "language": "python",
      "name": "python3"
    },
    "language_info": {
      "codemirror_mode": {
        "name": "ipython",
        "version": 3
      },
      "file_extension": ".py",
      "mimetype": "text/x-python",
      "name": "python",
      "nbconvert_exporter": "python",
      "pygments_lexer": "ipython3",
      "version": "3.11.11"
    }
  },
  "nbformat": 4,
  "nbformat_minor": 0
}