awe-ml — Analyze, Wrap, Explain

awe-ml is an end-to-end Python framework that carries a data set through the entire analytical workflow in one call, built on top of pandas and scikit-learn:

Explore -> Preprocess -> Select features -> Model -> Ensemble/Tune -> Wrap -> Explain -> Evaluate -> Report

Import name is awe_ml (import awe_ml); the distribution on PyPI is awe-ml. A different, unrelated package already owns the name awe on PyPI, so the import name is namespaced to avoid collisions.

Installation

pip install awe-ml

# with optional extras
pip install "awe-ml[xgboost,lightgbm,shap,report]"
# or everything at once
pip install "awe-ml[all]"

Quick start

import awe_ml as awe
from sklearn.datasets import fetch_openml

mtcars = fetch_openml(name="mtcars", version=1, as_frame=True).frame  # or use your own DataFrame

fit = awe.awe(mtcars, target="mpg")

fit.summary()                        # exploratory summary + leaderboard + evaluation + story
fit.plot(kind="importance")          # permutation importance
fit.plot(kind="pdp", var="wt")       # partial dependence (with ICE curves)
fit.plot(kind="leaderboard")         # cross-validated model comparison
fit.plot(kind="residuals")           # regression diagnostics (continuous targets)
fit.predict()                        # predictions on the training data
fit.explain()                        # the AWEExplainer object
fit.report(format="html")            # writes a report (HTML if `markdown` is installed)

Classification works the same way — the target type (continuous, binary, multiclass) is detected automatically, and binary targets additionally get ROC/AUC and a confusion-matrix heatmap:

from sklearn.datasets import load_iris

d = load_iris(as_frame=True)
iris = d.frame.assign(Species=d.target_names[d.target]).drop(columns=["target"])

fit = awe.awe(iris, target="Species")
fit.plot(kind="roc")         # binary targets only
fit.plot(kind="confusion")

Step-by-step (each stage is also its own function)

explore = awe.awe_explore(mtcars)
prep    = awe.awe_preprocess(mtcars, target="mpg")
model   = awe.awe_model(prep.train, target="mpg")
ens     = awe.awe_ensemble(model, top_n=3)                        # weighted ensemble of top candidates
tuned   = awe.awe_tune(prep.train, target="mpg", model="random_forest")  # grid-search tuning
wrapper = awe.awe_wrap(model)
explr   = awe.awe_explain(wrapper)
evalr   = awe.awe_evaluate(wrapper, newdata=prep.test)
story   = awe.awe_story(explore=explore, model=model, explainer=explr, evaluation=evalr)

Design

• Analyze: awe_explore(), eda_univariate(), eda_bivariate(), eda_report() (combined multi-panel overview), awe_preprocess() (with stratified train/test splitting), awe_select_features() (near-zero-variance, correlation, and univariate filters), awe_model() (automatic candidate model selection with k-fold cross-validation), awe_registry() (list every candidate available for a target type without fitting anything), and awe_tune() (grid-search hyperparameter tuning for decision_tree / random_forest / elastic_net / svm_radial).
• Wrap: awe_wrap() bundles any fitted candidate into a single AWEWrapper object with a consistent predict() / predict_proba() interface. Use which= to wrap any specific candidate, not just the overall best. awe_ensemble() combines the top-N cross-validated candidates into a performance-weighted ensemble that behaves like any other AWEWrapper. awe_save() / awe_load() persist any AWE object to disk via joblib.
• Explain: awe_explain() computes permutation importance and partial dependence / ICE curves directly against a wrapper's predict(), so it works identically for a plain model or a custom ensemble. awe_shap() adds a model-agnostic SHAP value summary when the optional shap package is installed.
• Compare: awe_compare() stacks the leaderboards of several runs (e.g. different feature sets) side by side, with a .plot() method.

Model bank

awe_model() always tries a solid bank of scikit-learn learners; XGBoost and LightGBM candidates are added automatically when those optional packages are installed. Call awe_registry("continuous") / awe_registry("binary") / awe_registry("multiclass") to list everything that would be tried, without fitting anything.

* only added when the optional package is installed.

Restrict the candidates tried with awe_model(data, target, models=[...]).

Plots

A consistent visual identity (theme_awe(), awe_palette()) is used throughout; every plotting function returns a matplotlib.figure.Figure so you can show, save, or embed it however you like.

• eda_univariate() / eda_bivariate() / eda_report() — distribution, relationship, and combined multi-panel overview plots
• plot_correlation() — correlation heatmap
• plot_missingness() — missing-value overview
• plot_feature_selection() — univariate association scores, selected vs. dropped
• plot_leaderboard() — cross-validated model comparison
• plot_tree() — visualise a fitted decision-tree candidate
• plot_roc() — ROC curve + rank-based AUC for binary targets (returns (fig, auc))
• plot_confusion() — confusion-matrix heatmap
• plot_residuals() — residuals-vs-fitted and normal Q-Q (continuous targets)
• explainer.plot(kind="importance" | "pdp") — variable importance / PDP+ICE
• compare.plot() — grouped-bar comparison across awe_compare() runs
• fit.plot(kind=...) — all of the above directly from a full awe() result

Status

Passes a 30-test pytest suite covering every public function, both target types (regression, binary and multiclass classification), model restriction, ensembling, tuning, comparison, explainability, evaluation, plotting, and reporting.