{
  "cells": [
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "\n# DirectRS Regression\n\nThis example demonstrates how to use DirectRS to post-process a pre-trained\nXGBoost regression model. DirectRS extracts a global stretch matrix from tree\ngeometry, fits per-tree Ridge regressions on stretched embeddings, and provides\nexact additive interpretability while maintaining or improving accuracy.\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. The ``ridge_alpha``\nparameter controls regularization of the per-tree Ridge regressions.\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": [
        "## Accuracy Comparison\nCompare R-squared scores between the base XGBoost model and DirectRS.\n\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "collapsed": false
      },
      "outputs": [],
      "source": [
        "from sklearn.metrics import r2_score\n\nbase_r2 = r2_score(ds.test_y, model.predict(ds.test_x))\ndrs_r2 = r2_score(ds.test_y, drs.predict(ds.test_x))\n\nprint(f\"Base XGBoost R\u00b2:  {base_r2:.4f}\")\nprint(f\"DirectRS R\u00b2:      {drs_r2:.4f}\")\nprint(f\"Improvement:      {drs_r2 - base_r2:+.4f}\")"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "## S' Stretch Analysis\nAnalyze the global stretch matrix S' extracted from tree geometry.\nThe eigenvalue spectrum shows the effective dimensionality, and\nfeature activity shows which features are most stretched.\n\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "collapsed": false
      },
      "outputs": [],
      "source": [
        "result = drs.get_global_stretch_analysis(ds.feature_names)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "Eigenvalue spectrum of S'.\n\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "collapsed": false
      },
      "outputs": [],
      "source": [
        "result.plot(\"eigenvalue_spectrum\")"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "Feature activity scores.\n\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "collapsed": false
      },
      "outputs": [],
      "source": [
        "result.plot(\"feature_activity\")"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "## Local Explanation\nDirectRS 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 = drs.explain_local(ds.test_x, feature_names=ds.feature_names)\nlocal = result.value\n\npred = drs.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.\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 the default ``slope`` mode\n(absolute Ridge coefficients aggregated across trees).\n\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "collapsed": false
      },
      "outputs": [],
      "source": [
        "result = drs.importance_global(feature_names=ds.feature_names)\nresult.plot()"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "## Main/Interaction Decomposition\nDecompose model variance into main effects and interactions using\neta-squared diagnostics. The main effect is computed from binned\npredictions f(x), not from individual coefficients.\n\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "collapsed": false
      },
      "outputs": [],
      "source": [
        "result = drs.importance_main_interaction(ds.test_x, feature_names=ds.feature_names)\nmi = result.value\n\nprint(f\"Orthogonalized variance split:\")\nprint(f\"  eta2_main = {mi['eta2_main']:.4f}  ({mi['eta2_main']*100:.1f}%)\")\nprint(f\"  eta2_int  = {mi['eta2_int']:.4f}  ({mi['eta2_int']*100:.1f}%)\")\nprint(f\"  rho(g, r) = {mi['rho']:.4f}\")"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "Main vs interaction importance bar chart.\n\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "collapsed": false
      },
      "outputs": [],
      "source": [
        "result.plot()"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "## Geometric Interaction Traces\nTrace feature interactions through the adjacency matrix A derived from\nthe log of the stretch matrix. Higher-order walks W^(k) = A^k capture\nk-step interaction paths.\n\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "collapsed": false
      },
      "outputs": [],
      "source": [
        "result = drs.geometric_interaction_traces(\n    feature_names=ds.feature_names, K=4, gamma=0.5\n)\ntraces = result.value\n\nprint(\"Interaction spectrum:\")\nfor k in range(len(traces['T'])):\n    print(f\"  k={k+1}: T_k = {traces['T'][k]:.6f},  E_k = {traces['E'][k]:.6f}\")"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "Adjacency matrix heatmap.\n\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "collapsed": false
      },
      "outputs": [],
      "source": [
        "result.plot(\"adjacency\")"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "Interaction energy spectrum.\n\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "collapsed": false
      },
      "outputs": [],
      "source": [
        "result.plot(\"spectrum\")"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "## Top Feature Interactions\nShow the strongest off-diagonal entries in the stretch matrix,\nrepresenting the most important pairwise feature interactions.\n\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "collapsed": false
      },
      "outputs": [],
      "source": [
        "result = drs.get_off_diagonal_analysis(ds.feature_names, top_k=15)\nresult.plot()"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "## FANOVA Comparison\nCompare DirectRS feature importance with FANOVA decomposition.\nFor depth-2 trees, both methods capture the same pairwise interactions.\n\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "collapsed": false
      },
      "outputs": [],
      "source": [
        "results = ts.interpret_fi()\nresults.plot()"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "collapsed": false
      },
      "outputs": [],
      "source": [
        "results = ts.interpret_ei()\nresults.plot()"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "Side-by-side comparison of importance scores.\n\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "collapsed": false
      },
      "outputs": [],
      "source": [
        "result_drs = drs.importance_global(feature_names=ds.feature_names, mode=\"slope\")\nresult_fanova = ts.interpret_fi()\n\ndrs_imp = result_drs.value['importance']\nfanova_table = result_fanova.table\nfanova_imp = dict(zip(fanova_table[\"Name\"], fanova_table[\"Score\"]))\n\nprint(f\"{'Feature':<12s} {'DirectRS':>10s} {'FANOVA':>10s}\")\nprint(\"-\" * 34)\nfor i, feat in enumerate(ds.feature_names):\n    print(f\"{feat:<12s} {drs_imp[i]:>10.4f} {fanova_imp.get(feat, 0.0):>10.4f}\")"
      ]
    }
  ],
  "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
}