Notebook Workflow

Quick In-Process Run

optimize() is the shortest supported notebook path. It does not start FastAPI or write to disk unless tracking is requested.

from sklearn.datasets import load_diabetes
from pspso import IntRange, OptimizationConfig, SearchSpace, optimize

X, y = load_diabetes(return_X_y=True)
search_space = SearchSpace({
    "n_estimators": IntRange(20, 80),
    "max_depth": IntRange(2, 10),
})
config = OptimizationConfig(
    task="regression",
    metric="rmse",
    strategy="pso",
    n_particles=4,
    n_iterations=3,
    random_state=42,
)
result = optimize(
    X,
    y,
    estimator="random_forest",
    search_space=search_space,
    config=config,
)
result

Use result.summary(), result.trials_frame(), result.predict(X), and result.evaluate(X, y) to inspect the outcome. predict_proba() is available only when the fitted estimator supports probabilities.

Tracked Notebook Run

from pspso import TrackingConfig

result = optimize(
    X,
    y,
    estimator="random_forest",
    search_space=search_space,
    config=config,
    tracking=TrackingConfig(
        experiment_name="Diabetes baselines",
        run_name="PSO notebook run",
        tags=("notebook", "baseline"),
        snapshot_data=True,
    ),
)
print(result.run_id, result.artifacts)

Tracked calls execute in the notebook process and write to .pspso/v1/. They appear in the same History page as GUI and CLI runs.

Full Example

"""Modern pspso API example for binary classification."""

from sklearn.datasets import load_breast_cancer
from sklearn.model_selection import train_test_split
from sklearn.preprocessing import StandardScaler

from pspso import Choice, FloatRange, OptimizationConfig, SearchSpace, optimize

X, y = load_breast_cancer(return_X_y=True)
X_train, X_val, y_train, y_val = train_test_split(X, y, test_size=0.2, random_state=42, stratify=y)
scaler = StandardScaler()
X_train = scaler.fit_transform(X_train)
X_val = scaler.transform(X_val)

result = optimize(
    X_train,
    y_train,
    estimator="svm",
    search_space=SearchSpace(
        {
            "kernel": Choice(["linear", "rbf"]),
            "C": FloatRange(0.1, 2.0, precision=1),
            "gamma": FloatRange(0.1, 1.0, precision=1),
        }
    ),
    config=OptimizationConfig(
        task="binary_classification",
        metric="roc_auc",
        strategy="random",
        max_trials=6,
        random_state=42,
    ),
    X_validation=X_val,
    y_validation=y_val,
)

print(result.best_params)
print(result.best_metric)