Running astra benchmark on a synthetic dataset¶
In this tutorial, we show how to run astra benchmark on a synthetic classification dataset. We begin by creating the dataset and bringing it into an ASTRA-ready form. ASTRA expects a pd.DataFrame with one row per sample, a Features column containing the feature vector for each sample, a Target column containing the label, and a Fold column containing the 0-indexed CV fold assignment. The helper function below creates this format from a standard NumPy feature matrix and label array using stratified 5-fold splits.
import pickle
import numpy as np
import pandas as pd
from sklearn.datasets import make_classification
from sklearn.metrics import (
average_precision_score,
matthews_corrcoef,
roc_auc_score,
)
from sklearn.model_selection import StratifiedKFold, train_test_split
def prepare_for_ASTRA(X_train: np.ndarray, y_train: np.ndarray) -> pd.DataFrame:
"""
Prepare the training data for ASTRA by creating 5-fold cross-validation splits
and returning a pd.DataFrame with the following columns:
- 'Features': a list of feature vectors for each sample
- 'Target': the corresponding label for each sample
- 'Fold': the fold number (0-4) for cross-validation
Parameters
----------
X_train : array-like, shape (n_samples, n_features)
The training feature matrix.
y_train : array-like, shape (n_samples,)
The training labels.
Returns
-------
pd.DataFrame
A DataFrame containing the features, target labels, and fold assignments.
"""
kf = StratifiedKFold(n_splits=5, shuffle=True, random_state=42)
folds = np.empty(len(X_train), dtype=int)
for fold, (_, val_index) in enumerate(kf.split(X_train, y_train)):
folds[val_index] = fold
df = pd.DataFrame({"Features": list(X_train), "Target": y_train, "Fold": folds})
return df
X, y = make_classification(
n_samples=1000,
n_features=100,
n_informative=75,
n_redundant=10,
n_repeated=5,
random_state=42,
weights=[0.7, 0.3],
)
# set aside 10% of the data as test set
X_train, X_test, y_train, y_test = train_test_split(
X, y, test_size=0.1, random_state=42, stratify=y
)
df_train = prepare_for_ASTRA(X_train, y_train)
df_train.to_pickle("synthetic_data_train.pkl")
with open("synthetic_data_test.pkl", "wb") as f:
pickle.dump((X_test, y_test), f)
Now we can run astra benchmark. We use ROC-AUC as the main metric for model selection, with PR-AUC and MCC as secondary metrics. ASTRA will evaluate all built-in classifiers using 5-fold CV and automatically select the best one via statistical testing. Hyperparameter tuning of the final model is performed using Optuna (--use_optuna).
%%bash
astra benchmark synthetic_data_train.pkl --name synthetic_data_benchmark --use_optuna --main_metric roc_auc --sec_metrics pr_auc mcc
👋 Welcome to ASTRA - Automated model selection using statistical testing
----------------------------------------------
_
__ _ ___ | |_ _ __ __ _
/ _` |/ __|| __|| '__| / _` |
| (_| |\__ \| |_ | | | (_| |
\__,_||___/ \__||_| \__,_|
----------------------------------------------
🤔 For help, run: astra --help
🧪 To benchmark models, run: astra benchmark --help
🏆 To compare models, run: astra compare --help
15-03 10:17 - INFO: Starting benchmark for synthetic_data_benchmark.
15-03 10:17 - INFO: Loading data.
15-03 10:17 - INFO: Starting benchmarking.
15-03 10:17 - INFO: Features column: Features
15-03 10:17 - INFO: Target column: Target
15-03 10:17 - INFO: Running 5-fold CV.
15-03 10:17 - INFO: Fold column: Fold
15-03 10:17 - INFO: Using Optuna for hyperparameter optimization, with 100 trials and a timeout of 3600 seconds.
15-03 10:17 - INFO: Will check assumptions for parametric tests and use them if met.
15-03 10:17 - INFO: Getting models and parameters.
15-03 10:17 - INFO: Benchmarking classification models.
15-03 10:17 - INFO: Main metric: roc_auc
15-03 10:17 - INFO: Secondary metrics: ['pr_auc', 'mcc']
15-03 10:17 - INFO: Starting CV for all models using default hyperparameters.
15-03 10:17 - INFO: Running CV.
15-03 10:17 - INFO: Running LogisticRegression.
Performance for LogisticRegression:
roc_auc: 0.915 ± 0.018 (median: 0.918)
pr_auc: 0.854 ± 0.040 (median: 0.862)
mcc: 0.686 ± 0.056 (median: 0.674)
15-03 10:17 - INFO: Running GaussianProcessClassifier.
Performance for GaussianProcessClassifier:
roc_auc: 0.500 ± 0.000 (median: 0.500)
pr_auc: 0.299 ± 0.002 (median: 0.300)
mcc: 0.000 ± 0.000 (median: 0.000)
15-03 10:17 - INFO: Running BernoulliNB.
Performance for BernoulliNB:
roc_auc: 0.828 ± 0.032 (median: 0.822)
pr_auc: 0.686 ± 0.051 (median: 0.673)
mcc: 0.473 ± 0.095 (median: 0.421)
15-03 10:17 - INFO: Running GaussianNB.
Performance for GaussianNB:
roc_auc: 0.903 ± 0.022 (median: 0.893)
pr_auc: 0.823 ± 0.047 (median: 0.797)
mcc: 0.605 ± 0.100 (median: 0.574)
15-03 10:17 - INFO: Running DecisionTreeClassifier.
Performance for DecisionTreeClassifier:
roc_auc: 0.609 ± 0.021 (median: 0.620)
pr_auc: 0.369 ± 0.019 (median: 0.379)
mcc: 0.219 ± 0.046 (median: 0.242)
15-03 10:17 - INFO: Running ExtraTreeClassifier.
Performance for ExtraTreeClassifier:
roc_auc: 0.570 ± 0.038 (median: 0.593)
pr_auc: 0.341 ± 0.024 (median: 0.354)
mcc: 0.141 ± 0.076 (median: 0.180)
15-03 10:17 - INFO: Running ExtraTreesClassifier.
Performance for ExtraTreesClassifier:
roc_auc: 0.918 ± 0.014 (median: 0.913)
pr_auc: 0.852 ± 0.028 (median: 0.852)
mcc: 0.328 ± 0.059 (median: 0.329)
15-03 10:17 - INFO: Running RandomForestClassifier.
Performance for RandomForestClassifier:
roc_auc: 0.896 ± 0.015 (median: 0.896)
pr_auc: 0.805 ± 0.045 (median: 0.823)
mcc: 0.361 ± 0.049 (median: 0.370)
15-03 10:17 - INFO: Running GradientBoostingClassifier.
Performance for GradientBoostingClassifier:
roc_auc: 0.895 ± 0.025 (median: 0.901)
pr_auc: 0.820 ± 0.050 (median: 0.814)
mcc: 0.568 ± 0.053 (median: 0.559)
15-03 10:17 - INFO: Running BaggingClassifier.
Performance for BaggingClassifier:
roc_auc: 0.773 ± 0.049 (median: 0.794)
pr_auc: 0.587 ± 0.068 (median: 0.610)
mcc: 0.351 ± 0.070 (median: 0.311)
15-03 10:17 - INFO: Running HistGradientBoostingClassifier.
Performance for HistGradientBoostingClassifier:
roc_auc: 0.928 ± 0.009 (median: 0.925)
pr_auc: 0.874 ± 0.016 (median: 0.872)
mcc: 0.651 ± 0.069 (median: 0.673)
15-03 10:17 - INFO: Running AdaBoostClassifier.
Performance for AdaBoostClassifier:
roc_auc: 0.826 ± 0.026 (median: 0.828)
pr_auc: 0.691 ± 0.064 (median: 0.674)
mcc: 0.461 ± 0.066 (median: 0.472)
15-03 10:17 - INFO: Running KNeighborsClassifier.
Performance for KNeighborsClassifier:
roc_auc: 0.726 ± 0.038 (median: 0.730)
pr_auc: 0.507 ± 0.031 (median: 0.502)
mcc: 0.293 ± 0.076 (median: 0.309)
15-03 10:17 - INFO: Running NearestCentroid.
Performance for NearestCentroid:
roc_auc: 0.893 ± 0.029 (median: 0.888)
pr_auc: 0.803 ± 0.055 (median: 0.798)
mcc: 0.184 ± 0.030 (median: 0.200)
15-03 10:17 - INFO: Running LinearDiscriminantAnalysis.
Performance for LinearDiscriminantAnalysis:
roc_auc: 0.928 ± 0.015 (median: 0.922)
pr_auc: 0.881 ± 0.034 (median: 0.879)
mcc: 0.736 ± 0.061 (median: 0.742)
15-03 10:17 - INFO: Running SGDClassifier.
Performance for SGDClassifier:
roc_auc: 0.822 ± 0.047 (median: 0.828)
pr_auc: 0.660 ± 0.085 (median: 0.688)
mcc: 0.615 ± 0.105 (median: 0.668)
15-03 10:17 - INFO: Running MLPClassifier.
Performance for MLPClassifier:
roc_auc: 0.948 ± 0.015 (median: 0.949)
pr_auc: 0.897 ± 0.046 (median: 0.914)
mcc: 0.738 ± 0.062 (median: 0.712)
15-03 10:17 - INFO: Running LGBMClassifier.
Performance for LGBMClassifier:
roc_auc: 0.933 ± 0.014 (median: 0.924)
pr_auc: 0.885 ± 0.022 (median: 0.885)
mcc: 0.639 ± 0.030 (median: 0.644)
15-03 10:17 - INFO: Running XGBClassifier.
Performance for XGBClassifier:
roc_auc: 0.918 ± 0.015 (median: 0.915)
pr_auc: 0.857 ± 0.024 (median: 0.851)
mcc: 0.651 ± 0.033 (median: 0.667)
15-03 10:17 - INFO: Done!
15-03 10:17 - INFO: Finished CV for all models.
15-03 10:17 - INFO: Checking assumptions for parametric tests.
15-03 10:17 - INFO: Assumptions of parametric tests met: False.
15-03 10:17 - INFO: Finding best model.
15-03 10:17 - INFO: Best model: MLPClassifier. Reason: Conover post-hoc test.
15-03 10:17 - INFO: Starting final hyperparameter tuning.
15-03 10:18 - INFO: Done!
-------------
Final results
-------------
Final model: MLPClassifier
Hyperparameters:
hidden_layer_sizes: (256, 128, 64)
activation: relu
solver: adam
alpha: 0.00028665421960535865
learning_rate: constant
learning_rate_init: 0.004806645399496772
early_stopping: False
batch_size: 204
Mean roc_auc: 0.963 ± 0.007.
Median roc_auc: 0.965.
Mean pr_auc: 0.934 ± 0.018.
Median pr_auc: 0.929.
Mean mcc: 0.794 ± 0.056.
Median mcc: 0.803.
The output tells us that MLPClassifier was selected as the best model, which means that it performed significantly better than the other models. Parametric test assumptions were not met, so non-parametric tests were used. The output was saved as benchmark.log under results/synthetic_data_benchmark. The final model was saved as final_model.pkl.
ASTRA also saved the model's final performance (final_CV.pkl) and optimised hyperparameters (final_hyperparameters.pkl):
with open("results/synthetic_data_benchmark/final_CV.pkl", "rb") as f:
final_CV = pickle.load(f)
print({k: f"{np.mean(v):.3f} ± {np.std(v):.3f}" for k, v in final_CV.items()})
with open("results/synthetic_data_benchmark/final_hyperparameters.pkl", "rb") as f:
final_hyperparameters = pickle.load(f)
print(final_hyperparameters)
{'roc_auc': '0.963 ± 0.007', 'pr_auc': '0.934 ± 0.018', 'mcc': '0.794 ± 0.056'}
{'hidden_layer_sizes': (256, 128, 64), 'activation': 'relu', 'solver': 'adam', 'alpha': 0.00028665421960535865, 'learning_rate': 'constant', 'learning_rate_init': 0.004806645399496772, 'early_stopping': False, 'batch_size': 204}
We can also see the performance metrics of all models using default hyperparameters, which are saved in default_CV.pkl:
with open("results/synthetic_data_benchmark/default_CV.pkl", "rb") as f:
default_CV = pickle.load(f)
default_CV_summary = pd.DataFrame(
{
model: {
metric: f"{np.mean(scores):.3f} ± {np.std(scores):.3f}"
for metric, scores in metrics.items()
}
for model, metrics in default_CV.items()
}
).T
default_CV_summary.sort_values("roc_auc", ascending=False)
| roc_auc | pr_auc | mcc | |
|---|---|---|---|
| MLPClassifier | 0.948 ± 0.015 | 0.897 ± 0.046 | 0.738 ± 0.062 |
| LGBMClassifier | 0.933 ± 0.014 | 0.885 ± 0.022 | 0.639 ± 0.030 |
| LinearDiscriminantAnalysis | 0.928 ± 0.015 | 0.881 ± 0.034 | 0.736 ± 0.061 |
| HistGradientBoostingClassifier | 0.928 ± 0.009 | 0.874 ± 0.016 | 0.651 ± 0.069 |
| XGBClassifier | 0.918 ± 0.015 | 0.857 ± 0.024 | 0.651 ± 0.033 |
| ExtraTreesClassifier | 0.918 ± 0.014 | 0.852 ± 0.028 | 0.328 ± 0.059 |
| LogisticRegression | 0.915 ± 0.018 | 0.854 ± 0.040 | 0.686 ± 0.056 |
| GaussianNB | 0.903 ± 0.022 | 0.823 ± 0.047 | 0.605 ± 0.100 |
| RandomForestClassifier | 0.896 ± 0.015 | 0.805 ± 0.045 | 0.361 ± 0.049 |
| GradientBoostingClassifier | 0.895 ± 0.025 | 0.820 ± 0.050 | 0.568 ± 0.053 |
| NearestCentroid | 0.893 ± 0.029 | 0.803 ± 0.055 | 0.184 ± 0.030 |
| BernoulliNB | 0.828 ± 0.032 | 0.686 ± 0.051 | 0.473 ± 0.095 |
| AdaBoostClassifier | 0.826 ± 0.026 | 0.691 ± 0.064 | 0.461 ± 0.066 |
| SGDClassifier | 0.822 ± 0.047 | 0.660 ± 0.085 | 0.615 ± 0.105 |
| BaggingClassifier | 0.773 ± 0.049 | 0.587 ± 0.068 | 0.351 ± 0.070 |
| KNeighborsClassifier | 0.726 ± 0.038 | 0.507 ± 0.031 | 0.293 ± 0.076 |
| DecisionTreeClassifier | 0.609 ± 0.021 | 0.369 ± 0.019 | 0.219 ± 0.046 |
| ExtraTreeClassifier | 0.570 ± 0.038 | 0.341 ± 0.024 | 0.141 ± 0.076 |
| GaussianProcessClassifier | 0.500 ± 0.000 | 0.299 ± 0.002 | 0.000 ± 0.000 |
Finally we can evaluate the final model on the held-out test set:
with open("results/synthetic_data_benchmark/final_model.pkl", "rb") as f:
final_model = pickle.load(f)
with open("synthetic_data_test.pkl", "rb") as f:
X_test, y_test = pickle.load(f)
predictions_astra = final_model.predict_proba(X_test)[:, 1]
auc = roc_auc_score(y_test, predictions_astra)
prauc = average_precision_score(y_test, predictions_astra)
predictions_astra_binary = final_model.predict(X_test)
mcc = matthews_corrcoef(y_test, predictions_astra_binary)
print(f"Test AUC: {auc:.4f}")
print(f"Test PRAUC: {prauc:.4f}")
print(f"Test MCC: {mcc:.4f}")
Test AUC: 0.9600
Test PRAUC: 0.9226
Test MCC: 0.7739
In this tutorial we demonstrated how to use ASTRA to benchmark models by:
constructing an ASTRA-ready dataset,
running
astra benchmarkto evaluate all built-in classifiers via 5-fold CV, andletting ASTRA select the best model through statistical testing.
The selected model was then evaluated on a held-out test set. For more detail on the available options see the User Guide.