Locking the search during generation updates

This notebook shows how Brush’s locking mechanism works when you want to preserve part of a program while continuing evolution on new data.

The sequence below demonstrates three phases:

  1. Fit an initial model on a base dataset.

  2. Lock the upper part of the tree and continue with partial_fit(...).

  3. Unlock the model and run another partial_fit(...) to resume unrestricted search.

import numpy as np
import graphviz

from pybrush import BrushRegressor

1. Create a base problem and a shifted follow-up problem

The second dataset changes the coefficient pattern so that the model benefits from reuse, but still needs adaptation.

def mse(y_true, y_pred):
    y_true = np.asarray(y_true)
    y_pred = np.asarray(y_pred)
    return float(np.mean((y_true - y_pred) ** 2))

rng = np.random.default_rng(7)
n_samples = 500

X_base = rng.normal(size=(n_samples, 3))
y_base = np.array([
    4.0 * row[0] + 0.5 * row[1] if row[2] > 0 else -3.0 * row[1] + 0.2 * row[0]
    for row in X_base
])

# Shifts both in feature and target!
X_shifted = rng.normal(size=(n_samples, 3))
X_shifted[:, 0] = X_shifted[:, 0] * 1.4 + 1.0
X_shifted[:, 1] = X_shifted[:, 1] - 0.5

y_shifted = np.array([
    4.0 * row[0] - 0.25 * row[1] if row[2] > 0 else -5.0 * row[1] + 0.4 * row[0]
    for row in X_shifted
])

print('base shape:', X_base.shape, y_base.shape)
print('shifted shape:', X_shifted.shape, y_shifted.shape)

base shape: (500, 3) (500,)
shifted shape: (500, 3) (500,)

2. Fit the initial generation window

This first fit builds a population on the base data. The compact tree structure is the part we will selectively preserve in the next step.

est = BrushRegressor(
    functions=['SplitOn', 'SplitBest', 'Mul', 'Add', 'Sub'],
    pop_size=100,
    max_gens=25,
    max_depth=10,
    max_size=24,
    start_from_decision_trees=True,
    constants_simplification=True,
    inexact_simplification=False,
    verbosity=1,
    random_state=7,
)

est.fit(X_base, y_base)

print('stage 1 base mse:', mse(y_base, est.predict(X_base)))
print('stage 1 shifted mse:', mse(y_shifted, est.predict(X_shifted)))
print('stage 1 model:')
print(est.best_estimator_.get_model())

Completed 100% [====================]
stage 1 base mse: 2.024100915882261
stage 1 shifted mse: 6.257131345427466
stage 1 model:
Sub(Mul(Add(x_1,If(x_1>=-2.61,x_0,-4.43)),Sub(1.47*x_2,1.31)),-3.38*x_0)

3. Lock the upper part of the tree

Locking the top levels keeps the learned structure fixed while lower nodes and weights can still adapt. This is the key step for controlled transfer between generation windows.

est.best_estimator_.program.lock_nodes(
    3,
    keep_leaves_unlocked=True,
    keep_current_weights=True,
)

print('locked the top two levels of the current best estimator')
print('locked model:')
print(est.best_estimator_.get_model())
print('locked tree:')
print(est.best_estimator_.get_model('tree'))
graphviz.Source(est.best_estimator_.get_model('dot'))

locked the top two levels of the current best estimator
locked model:
Sub(Mul(Add(x_1,If(x_1>=-2.61,x_0,-4.43)),Sub(1.47*x_2,1.31)),-3.38*x_0)
locked tree:
Sub
|- Mul
|  |- Add
|  |  |- x_1
|  |  |- If(x_1>=-2.61)
|  |  |  |- x_0
|  |  |  |- -4.43
|  |- Sub
|  |  |- 1.47*x_2
|  |  |- 1.31
|- -3.38*x_0
../_images/572331e08f729e81422793029c6757b31e240d2432f1efd19cdbeb5ec1115616.svg

4. Continue evolution with the lock in place

partial_fit(...) reuses the current engine and population, then applies the requested lock depth before continuing the search on the new data.

est.partial_fit(
    X_shifted,
    y_shifted,
    lock_nodes_depth=3,
    keep_leaves_unlocked=True,
    keep_current_weights=True,
)

print('stage 2 base mse:', mse(y_base, est.predict(X_base)))
print('stage 2 shifted mse:', mse(y_shifted, est.predict(X_shifted)))
print('stage 2 model:')
print(est.best_estimator_.get_model())
print('stage 2 tree:')
print(est.best_estimator_.get_model('tree'))
graphviz.Source(est.best_estimator_.get_model('dot'))

Completed 100% [====================]
stage 2 base mse: 3.6734543201338834
stage 2 shifted mse: 3.8054243085275217
stage 2 model:
Sub(Mul(Add(0.85*Sub(1.81*x_0,-2.40*x_1),0.01),Sub(x_2,-1.40)),0.88*Sub(-0.22,-6.02*x_1))
stage 2 tree:
Sub
|- Mul
|  |- Add
|  |  |- 0.85*Sub
|  |  |  |- 1.81*x_0
|  |  |  |- -2.40*x_1
|  |  |- 0.01
|  |- Sub
|  |  |- x_2
|  |  |- -1.40
|- 0.88*Sub
|  |- -0.22
|  |- -6.02*x_1
../_images/4e2f56e6c46cb021716038d96f73bca64e2d04e0025d71c92de079e80b69b1f9.svg