gen_optuna_study.py

#!/usr/bin/env python3
"""Generate a complete, resumable Optuna study script for a CSV dataset.

The script it writes:
  - loads the data, builds fixed CV folds ONCE (stratified / grouped / time-ordered) so every trial is
    scored on the same splits
  - defines an objective for the chosen model (lgbm | xgb | catboost | hgb | rf | logreg) with a
    search space sized to the data, early stopping inside each fold for boosting models, and
    trial.report(...) after every fold so the pruner can stop bad trials
  - creates the study with a seeded TPESampler (multivariate), a MedianPruner or HyperbandPruner,
    and an SQLite storage + study_name + load_if_exists=True so interrupted runs resume
  - stops early when the best value has not improved for --patience trials (callback)
  - saves trials to CSV, prints parameter importances, refits the best configuration on all data
    (boosting: lower learning rate, scaled iterations) and saves it with joblib
The generated script needs optuna >= 3, pandas, scikit-learn and the chosen library.

Usage:
    python gen_optuna_study.py data.csv --target churned --model lgbm --metric average_precision -o study.py
    python gen_optuna_study.py data.csv --target price --model xgb --metric neg_root_mean_squared_error --trials 100
    python gen_optuna_study.py data.csv --target churned --model catboost --group-col customer_id --pruner hyperband

Exit codes: 0 written, 2 usage or read error.
"""

import argparse
import csv
import math
import os
import sys
from itertools import islice

METRICS = {
    "roc_auc": ("roc_auc_score", "proba", "maximize"),
    "average_precision": ("average_precision_score", "proba", "maximize"),
    "neg_log_loss": ("log_loss", "proba", "minimize"),
    "accuracy": ("accuracy_score", "label", "maximize"),
    "neg_mean_absolute_error": ("mean_absolute_error", "value", "minimize"),
    "neg_root_mean_squared_error": ("root_mean_squared_error", "value", "minimize"),
    "r2": ("r2_score", "value", "maximize"),
}

SPACE = {
    "lgbm": '''        "num_leaves": trial.suggest_int("num_leaves", 7, {leaves_hi}, log=True),
        "min_child_samples": trial.suggest_int("min_child_samples", {mcs_lo}, {mcs_hi}, log=True),
        "subsample": trial.suggest_float("subsample", 0.5, 1.0),
        "subsample_freq": 1,
        "colsample_bytree": trial.suggest_float("colsample_bytree", 0.3, 1.0),
        "reg_lambda": trial.suggest_float("reg_lambda", 1e-3, 30.0, log=True),
        "reg_alpha": trial.suggest_float("reg_alpha", 1e-3, 10.0, log=True),''',
    "xgb": '''        "max_depth": trial.suggest_int("max_depth", 3, 10),
        "min_child_weight": trial.suggest_float("min_child_weight", 1.0, 50.0, log=True),
        "subsample": trial.suggest_float("subsample", 0.5, 1.0),
        "colsample_bytree": trial.suggest_float("colsample_bytree", 0.3, 1.0),
        "reg_lambda": trial.suggest_float("reg_lambda", 1e-3, 30.0, log=True),
        "gamma": trial.suggest_float("gamma", 1e-3, 5.0, log=True),''',
    "catboost": '''        "depth": trial.suggest_int("depth", 4, 9),
        "l2_leaf_reg": trial.suggest_float("l2_leaf_reg", 1.0, 20.0, log=True),
        "random_strength": trial.suggest_float("random_strength", 0.1, 10.0, log=True),
        "bootstrap_type": trial.suggest_categorical("bootstrap_type", ["Bayesian", "Bernoulli"]),''',
    "hgb": '''        "learning_rate": trial.suggest_float("learning_rate", 0.02, 0.3, log=True),
        "max_leaf_nodes": trial.suggest_int("max_leaf_nodes", 8, 128, log=True),
        "min_samples_leaf": trial.suggest_int("min_samples_leaf", {mcs_lo}, {mcs_hi}, log=True),
        "l2_regularization": trial.suggest_float("l2_regularization", 1e-4, 10.0, log=True),''',
    "rf": '''        "max_features": trial.suggest_float("max_features", 0.1, 0.8),
        "min_samples_leaf": trial.suggest_int("min_samples_leaf", 1, 50, log=True),''',
    "logreg": '''        "C": trial.suggest_float("C", 1e-4, 1e2, log=True),
        "penalty": trial.suggest_categorical("penalty", ["l1", "l2"]),''',
}

TEMPLATE = '''"""Optuna study for {target!r} ({model}, {metric}). Generated by gen_optuna_study.py.

Run:    python {script} [--trials N] [--timeout SECONDS]
Resume: run the same command again - trials are stored in {storage} under study "{study_name}".
Needs:  optuna >= 3, pandas, scikit-learn{extra_req}
"""
import argparse

import joblib
import numpy as np
import optuna
import pandas as pd
from sklearn.metrics import {metric_fn}
from sklearn.model_selection import {splitter_import}
{model_import}
DATA, TARGET = {data!r}, {target!r}
DROP = {drop!r}
STORAGE, STUDY = {storage!r}, {study_name!r}
SEED = 0

df = pd.read_csv(DATA)
df = df[df[TARGET].notna()].reset_index(drop=True)
{order_line}{target_prep}{groups_line}X = df.drop(columns=[TARGET] + [c for c in DROP if c in df.columns])
CAT = [c for c in X.columns if not pd.api.types.is_numeric_dtype(X[c])]
{cat_prep}FOLDS = list({cv_expr})          # fixed once: every trial sees the same splits


def score(y_true, model, X_te):
{score_body}


def objective(trial):
    params = {{
{space}
    }}
{conditional}    values, iters = [], []
    for k, (tr, te) in enumerate(FOLDS):
        X_tr, X_te, y_tr, y_te = X.iloc[tr], X.iloc[te], y.iloc[tr], y.iloc[te]
{fit_block}
        values.append(score(y_te, model, X_te))
        trial.report(float(np.mean(values)), step=k)           # lets the pruner stop bad trials early
        if trial.should_prune():
            raise optuna.TrialPruned()
    trial.set_user_attr("std", float(np.std(values)))
    if iters:
        trial.set_user_attr("best_iteration_mean", float(np.mean(iters)))
    return float(np.mean(values))


class StopWhenStale:
    """Stop the study when the best value has not improved for `patience` completed trials."""

    def __init__(self, patience):
        self.patience = patience

    def __call__(self, study, trial):
        done = [t for t in study.trials if t.state == optuna.trial.TrialState.COMPLETE]
        if len(done) > self.patience and study.best_trial.number < done[-self.patience].number:
            study.stop()


def main():
    ap = argparse.ArgumentParser()
    ap.add_argument("--trials", type=int, default={trials})
    ap.add_argument("--timeout", type=int, default=None, help="seconds")
    args = ap.parse_args()
    optuna.logging.set_verbosity(optuna.logging.WARNING)
    study = optuna.create_study(
        study_name=STUDY, storage=STORAGE, load_if_exists=True, direction="{direction}",
        sampler=optuna.samplers.TPESampler(seed=SEED, multivariate=True, n_startup_trials=10),
        pruner={pruner})
    study.optimize(objective, n_trials=args.trials, timeout=args.timeout, callbacks=[StopWhenStale({patience})],
                   gc_after_trial=True)
    done = [t for t in study.trials if t.state == optuna.trial.TrialState.COMPLETE]
    pruned = [t for t in study.trials if t.state == optuna.trial.TrialState.PRUNED]
    print(f"trials: {{len(done)}} complete, {{len(pruned)}} pruned; best {{study.best_value:.5f}} +/- "
          f"{{study.best_trial.user_attrs.get('std', float('nan')):.5f}}")
    print("best params:", study.best_params)
    study.trials_dataframe().to_csv(f"{{STUDY}}_trials.csv", index=False)
    if len(done) >= 5:
        print("importances:", {{k: round(v, 3) for k, v in optuna.importance.get_param_importances(study).items()}})
{refit_block}
    joblib.dump(final, f"{{STUDY}}_best_model.joblib")
    print(f"saved {{STUDY}}_best_model.joblib and {{STUDY}}_trials.csv")


if __name__ == "__main__":
    main()
'''


def main():
    ap = argparse.ArgumentParser(description="Generate a resumable Optuna study script.", epilog=__doc__,
                                 formatter_class=argparse.RawDescriptionHelpFormatter)
    ap.add_argument("csv")
    ap.add_argument("--target", required=True)
    ap.add_argument("--model", choices=sorted(SPACE), default="lgbm")
    ap.add_argument("--metric", choices=sorted(METRICS), default="roc_auc")
    ap.add_argument("--drop", nargs="*", default=[])
    ap.add_argument("--group-col")
    ap.add_argument("--time-col")
    ap.add_argument("--folds", type=int, default=5)
    ap.add_argument("--trials", type=int, default=100)
    ap.add_argument("--patience", type=int, default=30)
    ap.add_argument("--pruner", choices=["median", "hyperband", "none"], default="median")
    ap.add_argument("--study-name")
    ap.add_argument("-o", "--output")
    a = ap.parse_args()
    try:
        with open(a.csv, newline="", encoding="utf-8-sig") as fh:
            rd = csv.DictReader(fh)
            header = rd.fieldnames or []
            rows = list(islice(rd, 20000))
    except OSError as e:
        print(f"error: {e}", file=sys.stderr)
        return 2
    for c in (a.target, a.group_col, a.time_col):
        if c and c not in header:
            print(f"error: column {c!r} not found", file=sys.stderr)
            return 2
    metric_fn, kind, direction = METRICS[a.metric]
    classification = kind in ("proba", "label")
    n = len(rows)
    mcs = max(5, min(200, n // 100))
    leaves_hi = 255 if n > 100_000 else (127 if n > 20_000 else 63)
    drop = list(a.drop) + [c for c in (a.group_col, a.time_col) if c and c not in a.drop]
    if a.time_col:
        splitter, cv_expr = "TimeSeriesSplit", f"TimeSeriesSplit(n_splits={a.folds}).split(df)"
        order_line = f"df = df.sort_values({a.time_col!r}).reset_index(drop=True)\n"
    elif a.group_col:
        splitter = "StratifiedGroupKFold" if classification else "GroupKFold"
        cv_expr = (f"StratifiedGroupKFold(n_splits={a.folds}, shuffle=True, random_state=SEED).split(df, y, groups)"
                   if classification else f"GroupKFold(n_splits={a.folds}).split(df, y, groups)")
        order_line = ""
    else:
        splitter = "StratifiedKFold" if classification else "KFold"
        cv_expr = f"{splitter}(n_splits={a.folds}, shuffle=True, random_state=SEED).split(df, y)"
        order_line = ""
    groups_line = f"groups = df[{a.group_col!r}].to_numpy()\n" if a.group_col and not a.time_col else ""
    if classification:
        target_prep = ("positive = df[TARGET].value_counts().idxmin()\n"
                       "y = (df[TARGET] == positive).astype(int)\n")
    else:
        target_prep = "y = df[TARGET].astype(float)\n"
    if kind == "proba":
        clip = "np.clip(p, 1e-15, 1 - 1e-15)" if metric_fn == "log_loss" else "p"
        score_body = f"    p = model.predict_proba(X_te)[:, 1]\n    return {metric_fn}(y_true, {clip})"
    elif kind == "label":
        score_body = f"    return {metric_fn}(y_true, model.predict(X_te))"
    else:
        score_body = f"    return {metric_fn}(y_true, model.predict(X_te))"
    cls = {"lgbm": ("LGBMClassifier", "LGBMRegressor"), "xgb": ("XGBClassifier", "XGBRegressor"),
           "catboost": ("CatBoostClassifier", "CatBoostRegressor"),
           "hgb": ("HistGradientBoostingClassifier", "HistGradientBoostingRegressor"),
           "rf": ("RandomForestClassifier", "RandomForestRegressor"), "logreg": ("LogisticRegression", "Ridge")}[a.model]
    C = cls[0] if classification else cls[1]
    conditional = ""
    cat_prep = ""
    extra_req = {"lgbm": ", lightgbm", "xgb": ", xgboost", "catboost": ", catboost"}.get(a.model, "")
    if a.model == "lgbm":
        model_import = f"import lightgbm as lgb\nfrom lightgbm import {C}\n"
        cat_prep = "for c in CAT:\n    X[c] = X[c].astype(\"category\")\n"
        fit_block = (f"        model = {C}(**params, learning_rate=0.05, n_estimators=5000, random_state=SEED, verbose=-1)\n"
                     "        model.fit(X_tr, y_tr, eval_set=[(X_te, y_te)], callbacks=[lgb.early_stopping(100, verbose=False)])\n"
                     "        iters.append(model.best_iteration_)")
        refit = (f"    best = dict(study.best_params, subsample_freq=1)\n"
                 "    n_trees = int(study.best_trial.user_attrs.get('best_iteration_mean', 500) * 2 * 1.1)\n"
                 f"    final = {C}(**best, learning_rate=0.025, n_estimators=n_trees, random_state=SEED, verbose=-1).fit(X, y)")
    elif a.model == "xgb":
        model_import = f"from xgboost import {C}\n"
        cat_prep = "for c in CAT:\n    X[c] = X[c].astype(\"category\")\n"
        fit_block = (f"        model = {C}(**params, learning_rate=0.05, n_estimators=5000, tree_method=\"hist\",\n"
                     "                    enable_categorical=True, early_stopping_rounds=100, random_state=SEED)\n"
                     "        model.fit(X_tr, y_tr, eval_set=[(X_te, y_te)], verbose=False)\n"
                     "        iters.append(model.best_iteration)")
        refit = ("    n_trees = int(study.best_trial.user_attrs.get('best_iteration_mean', 500) * 2 * 1.1)\n"
                 f"    final = {C}(**study.best_params, learning_rate=0.025, n_estimators=n_trees, tree_method=\"hist\",\n"
                 "                 enable_categorical=True, random_state=SEED).fit(X, y)")
    elif a.model == "catboost":
        model_import = f"from catboost import {C}\n"
        cat_prep = "for c in CAT:\n    X[c] = X[c].astype(\"string\").fillna(\"__nan__\").astype(str)\n"
        conditional = ('    if params["bootstrap_type"] == "Bayesian":\n'
                       '        params["bagging_temperature"] = trial.suggest_float("bagging_temperature", 0.0, 1.0)\n'
                       '    else:\n'
                       '        params["subsample"] = trial.suggest_float("subsample", 0.6, 1.0)\n')
        fit_block = (f"        model = {C}(**params, learning_rate=0.05, iterations=3000, od_type=\"Iter\", od_wait=100,\n"
                     "                    cat_features=CAT, random_seed=SEED, verbose=0, allow_writing_files=False)\n"
                     "        model.fit(X_tr, y_tr, eval_set=(X_te, y_te), use_best_model=True)\n"
                     "        iters.append(model.get_best_iteration())")
        refit = ("    n_trees = int(study.best_trial.user_attrs.get('best_iteration_mean', 500) * 2 * 1.1)\n"
                 f"    final = {C}(**study.best_params, learning_rate=0.025, iterations=n_trees, cat_features=CAT,\n"
                 "                 random_seed=SEED, verbose=0, allow_writing_files=False).fit(X, y)")
    else:
        sk_mod = {"hgb": "sklearn.ensemble", "rf": "sklearn.ensemble", "logreg": "sklearn.linear_model"}[a.model]
        model_import = (f"from {sk_mod} import {C}\nfrom sklearn.compose import ColumnTransformer\n"
                        "from sklearn.pipeline import make_pipeline\n"
                        "from sklearn.preprocessing import OneHotEncoder, OrdinalEncoder, StandardScaler\n"
                        "from sklearn.impute import SimpleImputer\n")
        if a.model == "logreg":
            prep = ('ColumnTransformer([("num", make_pipeline(SimpleImputer(strategy="median"), StandardScaler()), '
                    '[c for c in X.columns if c not in CAT]), ("cat", OneHotEncoder(handle_unknown="ignore"), CAT)])')
            est = (f'{C}(C=params["C"], penalty=params["penalty"], solver="saga", max_iter=5000)' if classification
                   else f'{C}(alpha=1.0 / params["C"])')
        elif a.model == "hgb":
            prep = ('ColumnTransformer([("cat", OrdinalEncoder(handle_unknown="use_encoded_value", unknown_value=-1, '
                    'encoded_missing_value=-1), CAT)], remainder="passthrough")')
            est = f'{C}(**params, categorical_features=list(range(len(CAT))) or None, early_stopping=True, random_state=SEED)'
        else:
            prep = ('ColumnTransformer([("cat", OrdinalEncoder(handle_unknown="use_encoded_value", unknown_value=-1, '
                    'encoded_missing_value=-1), CAT)], remainder=SimpleImputer(strategy="median"))')
            est = f'{C}(**params, n_estimators=300, n_jobs=-1, random_state=SEED)'
        fit_block = (f"        model = make_pipeline({prep}, {est})\n"
                     "        model.fit(X_tr, y_tr)")
        refit = ("    trial = optuna.trial.FixedTrial(study.best_params)\n"
                 "    params = study.best_params\n"
                 f"    final = make_pipeline({prep}, {est}).fit(X, y)")
    space = SPACE[a.model].format(leaves_hi=leaves_hi, mcs_lo=max(2, mcs // 4), mcs_hi=mcs * 4)
    pruner = {"median": "optuna.pruners.MedianPruner(n_startup_trials=10, n_warmup_steps=1)",
              "hyperband": "optuna.pruners.HyperbandPruner(min_resource=1, max_resource=len(FOLDS))",
              "none": "optuna.pruners.NopPruner()"}[a.pruner]
    study_name = a.study_name or f"{a.model}_{a.target}"
    text = TEMPLATE.format(target=a.target, model=a.model, metric=a.metric, script=os.path.basename(a.output or "study.py"),
                           storage=f"sqlite:///{study_name}.db", study_name=study_name, extra_req=extra_req,
                           metric_fn=metric_fn, splitter_import=splitter, model_import=model_import, data=a.csv,
                           drop=drop, order_line=order_line, target_prep=target_prep, groups_line=groups_line,
                           cat_prep=cat_prep, cv_expr=cv_expr, score_body=score_body, space=space,
                           conditional=conditional, fit_block=fit_block, trials=a.trials, direction=direction,
                           pruner=pruner, patience=a.patience, refit_block=refit)
    if a.output:
        with open(a.output, "w", encoding="utf-8") as f:
            f.write(text)
        print(f"wrote {a.output}: {a.model}, {a.metric} ({direction}), {splitter}, {a.trials} trials, "
              f"{a.pruner} pruner, storage sqlite:///{study_name}.db", file=sys.stderr)
    else:
        sys.stdout.write(text)
    return 0


if __name__ == "__main__":
    sys.exit(main())