Source code for denia.insights

"""Insights: what drives the target?

``drivers()`` answers the question every experimentalist asks right after
"what should I run next?": *which variables actually control my outcome?*

A random forest is fit on the campaign's featurized history and its impurity
importances are reported at two levels:

* **feature level** - individual species and settings
  (``Cation::Mn``, ``Preparation::Impregnation``, ``Temperature, K``), and
* **column level** - importances summed back to the original spreadsheet
  columns/groups, so the answer is phrased in the user's own vocabulary.

These are descriptive, not causal: they say what the model relies on given
the (possibly biased) history, which is exactly the right caveat to show a
lab before they reorganize their program around it.
"""

from __future__ import annotations

from typing import Optional, Tuple

import numpy as np
import pandas as pd
from sklearn.ensemble import RandomForestRegressor

from .features import Featurizer




def drivers(
    history: pd.DataFrame,
    schema,
    target: Optional[str] = None,
    n_estimators: int = 300,
    random_state: int = 0,
) -> Tuple[pd.DataFrame, pd.DataFrame]:
    """Return (feature_level, column_level) importance tables for a target.

    Both tables have columns ``name`` and ``importance`` and sum to 1.
    """
    target = target or schema.targets[0]
    feat = Featurizer(schema).fit(history)
    X = feat.transform(history)
    y = feat.target_vector(history, target)
    mask = np.isfinite(y)
    rf = RandomForestRegressor(
        n_estimators=n_estimators, random_state=random_state, n_jobs=-1, min_samples_leaf=2
    ).fit(X[mask], y[mask])

    names = feat.feature_names_
    imp = rf.feature_importances_
    feature_level = (
        pd.DataFrame({"name": names, "importance": imp})
        .sort_values("importance", ascending=False)
        .reset_index(drop=True)
    )

    def group_of(n: str) -> str:
        return n.split("::", 1)[0] if "::" in n else n

    column_level = (
        feature_level.assign(name=feature_level["name"].map(group_of))
        .groupby("name", as_index=False)["importance"].sum()
        .sort_values("importance", ascending=False)
        .reset_index(drop=True)
    )
    return feature_level, column_level