DirectRS Regression

This example demonstrates how to use DirectRS to post-process a pre-trained XGBoost regression model. DirectRS extracts a global stretch matrix from tree geometry, fits per-tree Ri

# %%
# Setup
# -----
# Import libraries and suppress warnings.

import warnings
warnings.filterwarnings("ignore")

import numpy as np
from modeva import DataSet, TestSuite
from modeva.models import MoXGBRegressor

# %%
# Load Dataset
# ------------
# Load the CaliforniaHousing dataset and create a random train/test split.

ds = DataSet()
ds.load(name="CaliforniaHousing")
ds.set_random_split()

# %%
# Train Base Model
# ----------------
# Train an XGBoost regressor with depth 2 (required for FANOVA comparison).

model = MoXGBRegressor(
    name="XGB-depth2",
    n_estimators=200, max_depth=2, learning_rate=0.1,
    subsample=0.8, colsample_bytree=0.8,
    random_state=42, verbosity=0
)
model.fit(ds.train_x, ds.train_y.ravel())

ts = TestSuite(ds, model)
ts.diagnose_accuracy_table().table

# %%
# Fit DirectRS
# ------------
# Post-process the trained XGBoost model with DirectRS. The ``ridge_alpha``
# parameter controls regularization of the per-tree Ridge regressions.

from modeva.models import MoDirectRSRegressor

drs = MoDirectRSRegressor(
    base_model=model, ridge_alpha=100.0, n_passes=1
)
drs.fit(
    ds.train_x, ds.train_y.ravel(),
    X_val=ds.test_x, y_val=ds.test_y.ravel(),
    verbose=True
)

# %%
# Accuracy Comparison
# -------------------
# Compare R-squared scores between the base XGBoost model and DirectRS.

from sklearn.metrics import r2_score

base_r2 = r2_score(ds.test_y, model.predict(ds.test_x))
drs_r2 = r2_score(ds.test_y, drs.predict(ds.test_x))

print(f"Base XGBoost R²:  {base_r2:.4f}")
print(f"DirectRS R²:      {drs_r2:.4f}")
print(f"Improvement:      {drs_r2 - base_r2:+.4f}")

# %%
# S' Stretch Analysis
# -------------------
# Analyze the global stretch matrix S' extracted from tree geometry.
# The eigenvalue spectrum shows the effective dimensionality, and
# feature activity shows which features are most stretched.

result = drs.get_global_stretch_analysis(ds.feature_names)

# %%
# Eigenvalue spectrum of S'.
result.plot("eigenvalue_spectrum")

# %%
# Feature activity scores.
result.plot("feature_activity")

# %%
# Local Explanation
# -----------------
# DirectRS provides exact additive decomposition: f(x) = c_0 + sum_j c_j.
# We verify the decomposition matches predictions to machine precision.

result = drs.explain_local(ds.test_x, feature_names=ds.feature_names)
local = result.value

pred = drs.predict(ds.test_x)
recon = local['intercept'] + local['contributions'].sum(axis=1)
max_err = np.max(np.abs(pred - recon))

print(f"Max |predict - (intercept + sum contributions)|: {max_err:.2e}")
print(f"Decomposition exact to machine precision: {max_err < 1e-10}")

# %%
# Local explanation waterfall plot.
result.plot()

# %%
# Global Feature Importance
# -------------------------
# Compute global feature importance using the default ``slope`` mode
# (absolute Ridge coefficients aggregated across trees).

result = drs.importance_global(feature_names=ds.feature_names)
result.plot()

# %%
# Main/Interaction Decomposition
# ------------------------------
# Decompose model variance into main effects and interactions using
# eta-squared diagnostics. The main effect is computed from binned
# predictions f(x), not from individual coefficients.

result = drs.importance_main_interaction(ds.test_x, feature_names=ds.feature_names)
mi = result.value

print(f"Orthogonalized variance split:")
print(f"  eta2_main = {mi['eta2_main']:.4f}  ({mi['eta2_main']*100:.1f}%)")
print(f"  eta2_int  = {mi['eta2_int']:.4f}  ({mi['eta2_int']*100:.1f}%)")
print(f"  rho(g, r) = {mi['rho']:.4f}")

# %%
# Main vs interaction importance bar chart.
result.plot()

# %%
# Geometric Interaction Traces
# ----------------------------
# Trace feature interactions through the adjacency matrix A derived from
# the log of the stretch matrix. Higher-order walks W^(k) = A^k capture
# k-step interaction paths.

result = drs.geometric_interaction_traces(
    feature_names=ds.feature_names, K=4, gamma=0.5
)
traces = result.value

print("Interaction spectrum:")
for k in range(len(traces['T'])):
    print(f"  k={k+1}: T_k = {traces['T'][k]:.6f},  E_k = {traces['E'][k]:.6f}")

# %%
# Adjacency matrix heatmap.
result.plot("adjacency")

# %%
# Interaction energy spectrum.
result.plot("spectrum")

# %%
# Top Feature Interactions
# ------------------------
# Show the strongest off-diagonal entries in the stretch matrix,
# representing the most important pairwise feature interactions.

result = drs.get_off_diagonal_analysis(ds.feature_names, top_k=15)
result.plot()

# %%
# FANOVA Comparison
# -----------------
# Compare DirectRS feature importance with FANOVA decomposition.
# For depth-2 trees, both methods capture the same pairwise interactions.

results = ts.interpret_fi()
results.plot()

# %%
results = ts.interpret_ei()
results.plot()

# %%
# Side-by-side comparison of importance scores.

result_drs = drs.importance_global(feature_names=ds.feature_names, mode="slope")
result_fanova = ts.interpret_fi()

drs_imp = result_drs.value['importance']
fanova_table = result_fanova.table
fanova_imp = dict(zip(fanova_table["Name"], fanova_table["Score"]))

print(f"{'Feature':<12s} {'DirectRS':>10s} {'FANOVA':>10s}")
print("-" * 34)
for i, feat in enumerate(ds.feature_names):
    print(f"{feat:<12s} {drs_imp[i]:>10.4f} {fanova_imp.get(feat, 0.0):>10.4f}")