{
"cells": [
{
"cell_type": "markdown",
"id": "b76f1bb9",
"metadata": {},
"source": [
"# Locking the search during generation updates\n",
"\n",
"This notebook shows how Brush's locking mechanism works when you want to preserve part of a program while continuing evolution on new data.\n",
"\n",
"The sequence below demonstrates three phases:\n",
"\n",
"1. Fit an initial model on a base dataset.\n",
"2. Lock the upper part of the tree and continue with `partial_fit(...)`.\n",
"3. Unlock the model and run another `partial_fit(...)` to resume unrestricted search."
]
},
{
"cell_type": "code",
"execution_count": 1,
"id": "134fd133",
"metadata": {},
"outputs": [],
"source": [
"import numpy as np\n",
"import graphviz\n",
"\n",
"from pybrush import BrushRegressor"
]
},
{
"cell_type": "markdown",
"id": "06b8bfde",
"metadata": {},
"source": [
"## 1. Create a base problem and a shifted follow-up problem\n",
"\n",
"The second dataset changes the coefficient pattern so that the model benefits from reuse, but still needs adaptation."
]
},
{
"cell_type": "code",
"execution_count": 2,
"id": "7b69da78",
"metadata": {},
"outputs": [],
"source": [
"def mse(y_true, y_pred):\n",
" y_true = np.asarray(y_true)\n",
" y_pred = np.asarray(y_pred)\n",
" return float(np.mean((y_true - y_pred) ** 2))"
]
},
{
"cell_type": "code",
"execution_count": 3,
"id": "d4128f4b",
"metadata": {},
"outputs": [
{
"name": "stdout",
"output_type": "stream",
"text": [
"base shape: (500, 3) (500,)\n",
"shifted shape: (500, 3) (500,)\n"
]
}
],
"source": [
"rng = np.random.default_rng(7)\n",
"n_samples = 500\n",
"\n",
"X_base = rng.normal(size=(n_samples, 3))\n",
"y_base = np.array([\n",
" 4.0 * row[0] + 0.5 * row[1] if row[2] > 0 else -3.0 * row[1] + 0.2 * row[0]\n",
" for row in X_base\n",
"])\n",
"\n",
"# Shifts both in feature and target!\n",
"X_shifted = rng.normal(size=(n_samples, 3))\n",
"X_shifted[:, 0] = X_shifted[:, 0] * 1.4 + 1.0\n",
"X_shifted[:, 1] = X_shifted[:, 1] - 0.5\n",
"\n",
"y_shifted = np.array([\n",
" 4.0 * row[0] - 0.25 * row[1] if row[2] > 0 else -5.0 * row[1] + 0.4 * row[0]\n",
" for row in X_shifted\n",
"])\n",
"\n",
"print('base shape:', X_base.shape, y_base.shape)\n",
"print('shifted shape:', X_shifted.shape, y_shifted.shape)"
]
},
{
"cell_type": "markdown",
"id": "0b73890b",
"metadata": {},
"source": [
"## 2. Fit the initial generation window\n",
"\n",
"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."
]
},
{
"cell_type": "code",
"execution_count": 4,
"id": "2f8e59dd",
"metadata": {},
"outputs": [
{
"name": "stdout",
"output_type": "stream",
"text": [
"Completed 100% [====================]\n",
"stage 1 base mse: 2.024100915882261\n",
"stage 1 shifted mse: 6.257131345427466\n",
"stage 1 model:\n",
"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)\n"
]
}
],
"source": [
"est = BrushRegressor(\n",
" functions=['SplitOn', 'SplitBest', 'Mul', 'Add', 'Sub'],\n",
" pop_size=100,\n",
" max_gens=25,\n",
" max_depth=10,\n",
" max_size=24,\n",
" start_from_decision_trees=True,\n",
" constants_simplification=True,\n",
" inexact_simplification=False,\n",
" verbosity=1,\n",
" random_state=7,\n",
")\n",
"\n",
"est.fit(X_base, y_base)\n",
"\n",
"print('stage 1 base mse:', mse(y_base, est.predict(X_base)))\n",
"print('stage 1 shifted mse:', mse(y_shifted, est.predict(X_shifted)))\n",
"print('stage 1 model:')\n",
"print(est.best_estimator_.get_model())"
]
},
{
"cell_type": "markdown",
"id": "584e8ca5",
"metadata": {},
"source": [
"## 3. Lock the upper part of the tree\n",
"\n",
"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."
]
},
{
"cell_type": "code",
"execution_count": 5,
"id": "4a641739",
"metadata": {},
"outputs": [
{
"name": "stdout",
"output_type": "stream",
"text": [
"locked the top two levels of the current best estimator\n",
"locked model:\n",
"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)\n",
"locked tree:\n",
"Sub\n",
"|- Mul\n",
"| |- Add\n",
"| | |- x_1\n",
"| | |- If(x_1>=-2.61)\n",
"| | | |- x_0\n",
"| | | |- -4.43\n",
"| |- Sub\n",
"| | |- 1.47*x_2\n",
"| | |- 1.31\n",
"|- -3.38*x_0\n"
]
},
{
"data": {
"image/svg+xml": [
"\n",
"\n",
"\n",
"\n",
"\n"
],
"text/plain": [
""
]
},
"execution_count": 5,
"metadata": {},
"output_type": "execute_result"
}
],
"source": [
"est.best_estimator_.program.lock_nodes(\n",
" 3,\n",
" keep_leaves_unlocked=True,\n",
" keep_current_weights=True,\n",
")\n",
"\n",
"print('locked the top two levels of the current best estimator')\n",
"print('locked model:')\n",
"print(est.best_estimator_.get_model())\n",
"print('locked tree:')\n",
"print(est.best_estimator_.get_model('tree'))\n",
"graphviz.Source(est.best_estimator_.get_model('dot'))"
]
},
{
"cell_type": "markdown",
"id": "848deb6e",
"metadata": {},
"source": [
"## 4. Continue evolution with the lock in place\n",
"\n",
"`partial_fit(...)` reuses the current engine and population, then applies the requested lock depth before continuing the search on the new data."
]
},
{
"cell_type": "code",
"execution_count": 6,
"id": "31b06d59",
"metadata": {},
"outputs": [
{
"name": "stdout",
"output_type": "stream",
"text": [
"Completed 100% [====================]\n",
"stage 2 base mse: 3.6734543201338834\n",
"stage 2 shifted mse: 3.8054243085275217\n",
"stage 2 model:\n",
"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))\n",
"stage 2 tree:\n",
"Sub\n",
"|- Mul\n",
"| |- Add\n",
"| | |- 0.85*Sub\n",
"| | | |- 1.81*x_0\n",
"| | | |- -2.40*x_1\n",
"| | |- 0.01\n",
"| |- Sub\n",
"| | |- x_2\n",
"| | |- -1.40\n",
"|- 0.88*Sub\n",
"| |- -0.22\n",
"| |- -6.02*x_1\n"
]
},
{
"data": {
"image/svg+xml": [
"\n",
"\n",
"\n",
"\n",
"\n"
],
"text/plain": [
""
]
},
"execution_count": 6,
"metadata": {},
"output_type": "execute_result"
}
],
"source": [
"est.partial_fit(\n",
" X_shifted,\n",
" y_shifted,\n",
" lock_nodes_depth=3,\n",
" keep_leaves_unlocked=True,\n",
" keep_current_weights=True,\n",
")\n",
"\n",
"print('stage 2 base mse:', mse(y_base, est.predict(X_base)))\n",
"print('stage 2 shifted mse:', mse(y_shifted, est.predict(X_shifted)))\n",
"print('stage 2 model:')\n",
"print(est.best_estimator_.get_model())\n",
"print('stage 2 tree:')\n",
"print(est.best_estimator_.get_model('tree'))\n",
"graphviz.Source(est.best_estimator_.get_model('dot'))"
]
},
{
"cell_type": "markdown",
"id": "ff4751bb",
"metadata": {},
"source": [
"## 5. Stop locking and resume unrestricted search\n",
"\n",
"Passing `end_depth=0` unlocks the tree again. The last call shows that generation updates can resume without preserving the prior lock state."
]
},
{
"cell_type": "code",
"execution_count": 7,
"id": "0704cf48",
"metadata": {},
"outputs": [
{
"name": "stdout",
"output_type": "stream",
"text": [
"Completed 100% [====================]\n",
"stage 3 base mse: 3.648386877972349\n",
"stage 3 shifted mse: 3.7962339333011497\n",
"stage 3 unlocked model:\n",
"Sub(Mul(Add(Sub(1.57*x_0,-2.03*x_1),-0.05*x_2),Sub(x_2,-1.38)),0.93*Sub(-0.29,-5.58*x_1))\n",
"stage 3 unlocked tree:\n",
"Sub\n",
"|- Mul\n",
"| |- Add\n",
"| | |- Sub\n",
"| | | |- 1.57*x_0\n",
"| | | |- -2.03*x_1\n",
"| | |- -0.05*x_2\n",
"| |- Sub\n",
"| | |- x_2\n",
"| | |- -1.38\n",
"|- 0.93*Sub\n",
"| |- -0.29\n",
"| |- -5.58*x_1\n"
]
},
{
"data": {
"image/svg+xml": [
"\n",
"\n",
"\n",
"\n",
"\n"
],
"text/plain": [
""
]
},
"execution_count": 7,
"metadata": {},
"output_type": "execute_result"
}
],
"source": [
"est.best_estimator_.program.lock_nodes(0, keep_leaves_unlocked=True, keep_current_weights=False)\n",
"\n",
"est.max_gens = 50\n",
"est.partial_fit(\n",
" X_shifted,\n",
" y_shifted,\n",
" lock_nodes_depth=0,\n",
" keep_leaves_unlocked=True,\n",
" keep_current_weights=False,\n",
")\n",
"\n",
"print('stage 3 base mse:', mse(y_base, est.predict(X_base)))\n",
"print('stage 3 shifted mse:', mse(y_shifted, est.predict(X_shifted)))\n",
"print('stage 3 unlocked model:')\n",
"print(est.best_estimator_.get_model())\n",
"print('stage 3 unlocked tree:')\n",
"print(est.best_estimator_.get_model('tree'))\n",
"graphviz.Source(est.best_estimator_.get_model('dot'))"
]
}
],
"metadata": {
"kernelspec": {
"display_name": "brush",
"language": "python",
"name": "python3"
},
"language_info": {
"codemirror_mode": {
"name": "ipython",
"version": 3
},
"file_extension": ".py",
"mimetype": "text/x-python",
"name": "python",
"nbconvert_exporter": "python",
"pygments_lexer": "ipython3",
"version": "3.13.14"
}
},
"nbformat": 4,
"nbformat_minor": 5
}