#!/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())