Switching the classification metric with partial_fit

Brush uses the scorer to evaluate programs: it drives selection, survival, the archive, and the final model choice. For classification, the available scorers are:

scorer

binary

multiclass

"log" / "multi_log"

log loss

multinomial log loss

"accuracy"

accuracy

accuracy

"balanced_accuracy"

balanced accuracy

balanced accuracy

"precision"

precision (threshold 0.5)

macro precision

"recall"

recall (threshold 0.5)

macro recall

"roc_auc"

AUROC

macro one-vs-rest AUROC

"average_precision_score"

average precision (AUPRC)

macro one-vs-rest average precision

The scorer is not used to fit parameters. Weights are always optimized with the log loss, and split thresholds with the gini impurity. This keeps parameter fitting smooth and well-behaved while letting you pick whichever metric you care about for model selection.

Because the scorer is just an estimator attribute, you can change it between calls to partial_fit. This notebook:

  1. Fits a model on an imbalanced binary problem using AUROC.

  2. Locks every internal node of the best program, leaving only the leaves and the weights free to change.

  3. Switches the scorer to average precision and calls partial_fit.

  4. Compares all metrics before and after the switch.

import numpy as np
import pandas as pd
import graphviz

from sklearn.datasets import make_classification
from sklearn.model_selection import train_test_split
from sklearn.metrics import (log_loss, accuracy_score, balanced_accuracy_score,
                             precision_score, recall_score, roc_auc_score,
                             average_precision_score)

from pybrush import BrushClassifier

1. An imbalanced binary problem

AUROC and average precision disagree the most when positives are rare: AUROC rewards ranking negatives correctly, while average precision focuses on how clean the top of the ranking is.

X, y = make_classification(
    n_samples=1000, n_features=6, n_informative=4, n_redundant=1,
    weights=[0.9, 0.1], class_sep=0.8, flip_y=0.02, random_state=42,
)

X_train, X_test, y_train, y_test = train_test_split(
    X, y, test_size=0.3, stratify=y, random_state=42)

print('train prevalence:', y_train.mean().round(3))
print('test prevalence: ', y_test.mean().round(3))

train prevalence: 0.109
test prevalence:  0.107

def all_metrics(est, X, y):
    proba = est.predict_proba(X)[:, 1]
    pred = est.predict(X)
    return {
        'log_loss':          log_loss(y, proba),
        'accuracy':          accuracy_score(y, pred),
        'balanced_accuracy': balanced_accuracy_score(y, pred),
        'precision':         precision_score(y, pred, zero_division=0),
        'recall':            recall_score(y, pred, zero_division=0),
        'roc_auc':           roc_auc_score(y, proba),
        'average_precision': average_precision_score(y, proba),
    }

def report(est, label):
    return pd.DataFrame({
        (label, 'train'): all_metrics(est, X_train, y_train),
        (label, 'test'):  all_metrics(est, X_test, y_test),
    })

2. Fit using AUROC as the scorer

est = BrushClassifier(
    functions=['SplitBest', 'Add', 'Sub', 'Mul', 'Div', 'Logabs', 'Exp'],
    scorer='roc_auc',
    pop_size=200,
    max_gens=30,
    max_depth=5,
    max_size=20,
    random_state=42,
    verbosity=0,
)
est.fit(X_train, y_train)

print('scorer:', est.parameters_.scorer)
print('internal scorer value (train):', round(est.best_estimator_.fitness.loss, 4))
print('model:', est.best_estimator_.get_model())

before = report(est, 'roc_auc')
before

scorer: roc_auc
internal scorer value (train): 0.9027
model: Logistic(Add(0.98,Add(-1.05*Mul(-1.32*x_4,-1.01*x_2),x_0)))
roc_auc
train test
log_loss 0.398210 0.387116
accuracy 0.810000 0.843333
balanced_accuracy 0.800987 0.871035
precision 0.338983 0.397260
recall 0.789474 0.906250
roc_auc 0.886956 0.938666
average_precision 0.678297 0.789859

3. Lock the internal nodes, switch to average precision, and refit

A lock_nodes_depth larger than the tree depth locks every operator and split in the program. With keep_leaves_unlocked=True, the leaves (features and constants) stay free, so the search can still swap a leaf or grow it into a small subtree. With keep_current_weights=False, the weights are re-optimized.

We then set est.scorer = 'average_precision_score' before calling partial_fit. The new scorer is picked up by the engine, so every program is re-evaluated, selected, and archived by average precision from here on.

structure_before = est.best_estimator_.get_model()
dot_before = est.best_estimator_.get_model('dot')

est.scorer = 'average_precision_score'
est.partial_fit(
    X_train, y_train,
    lock_nodes_depth=est.max_depth + 1,  # deeper than any tree: lock every internal node
    keep_leaves_unlocked=True,           # ...but leave the leaves free to change
    keep_current_weights=False,          # weights can still be optimized
)

print('scorer:', est.parameters_.scorer)
print('internal scorer value (train):', round(est.best_estimator_.fitness.loss, 4))
print('model before:', structure_before)
print('model after: ', est.best_estimator_.get_model())

after = report(est, 'average_precision_score')

scorer: average_precision_score
internal scorer value (train): 0.917
model before: Logistic(Add(0.98,Add(-1.05*Mul(-1.32*x_4,-1.01*x_2),x_0)))
model after:  Logistic(Add(-2.77,Add(-1.53*Mul(x_4,x_2),3.87*Exp(0.42*x_0))))

The two programs side by side. The locked internal nodes are the same in both, and only the leaves (and the weights) differ.

from IPython.display import HTML

def svg(dot):
    return graphviz.Source(dot).pipe(format='svg').decode()

HTML(f"""
<div style="display: flex; gap: 2em; align-items: flex-start; flex-wrap: wrap;">
  <div><h4>Before (AUROC)</h4>{svg(dot_before)}</div>
  <div><h4>After (average precision)</h4>{svg(est.best_estimator_.get_model('dot'))}</div>
</div>
""")

Before (AUROC)

G ^ split feature fixed, * split threshold fixed 1473c7000 Logistic 105024bc0 Add 1473c7000->105024bc0 1473c70b0 Add 105024bc0->1473c70b0 105024bc0Offset 0.98 105024bc0->105024bc0Offset 1473884b0 Mul 1473c70b0->1473884b0 -1.05 x_0 x_0 1473c70b0->x_0 x_4 x_4 1473884b0->x_4 -1.32 x_2 x_2 1473884b0->x_2 -1.01

After (average precision)

G ^ split feature fixed, * split threshold fixed 15f54b570 Logistic 15f54ba00 Add 15f54b570->15f54ba00 15f5c09c0 Add 15f54ba00->15f5c09c0 15f54ba00Offset -2.77 15f54ba00->15f54ba00Offset 15f5c0a70 Mul 15f5c09c0->15f5c0a70 -1.53 15f5f0460 Exp 15f5c09c0->15f5f0460 3.87 x_4 x_4 15f5c0a70->x_4 x_2 x_2 15f5c0a70->x_2 x_0 x_0 15f5f0460->x_0 0.42

4. Compare every metric before and after the switch

comparison = pd.concat([before, after], axis=1)
comparison[('delta', 'train')] = comparison[('average_precision_score', 'train')] - comparison[('roc_auc', 'train')]
comparison[('delta', 'test')]  = comparison[('average_precision_score', 'test')]  - comparison[('roc_auc', 'test')]
comparison.round(4)
roc_auc average_precision_score delta
train test train test train test
log_loss 0.3982 0.3871 0.3819 0.3787 -0.0163 -0.0084
accuracy 0.8100 0.8433 0.8400 0.8633 0.0300 0.0200
balanced_accuracy 0.8010 0.8710 0.8178 0.8822 0.0168 0.0112
precision 0.3390 0.3973 0.3846 0.4328 0.0456 0.0356
recall 0.7895 0.9062 0.7895 0.9062 0.0000 0.0000
roc_auc 0.8870 0.9387 0.8892 0.9397 0.0022 0.0010
average_precision 0.6783 0.7899 0.7041 0.7983 0.0258 0.0085