astra.benchmark

Functions

run(data[, name, features, target, ...])

Run the benchmark.
astra.benchmark.run(data: str, name: str | None = None, features: str = 'Features', target: str = 'Target', run_nested_CV: bool = False, use_optuna: bool = False, n_trials: int = 100, timeout: int = 3600, fold_col: str | list[str] = 'Fold', main_metric: str = 'R2', sec_metrics: list[str] = ['MSE', 'MAE'], parametric: str | bool = 'auto', impute: str | float | int | None = None, remove_constant: float | None = None, remove_correlated: float | None = None, scaler: str | None = None, custom_models: dict[str, None | dict[str, dict] | tuple[dict[str, dict], dict[str, dict]]] | None = None, n_jobs: int = 1, test_mode: bool = False) → None[source]

Run the benchmark.

Parameters:

  • data (str) -- Path to the dataset to train and evaluate models on. This should be a CSV, pickle, or parquet file.

  • name (str or None, default=None) -- Name of the experiment. Results will be saved in a folder with this name in the 'results' directory. Will be used to load cached results if they exist. If None, the name will be the file name of the data file without extension.

  • features (str, default='Features') -- Name of the column containing the features. Default: Features.

  • target (str, default='Target') -- Name of the column containing the target. Default: Target.

  • run_nested_CV (bool, default=False) -- Whether to run nested CV with hyperparameter tuning for the best models.

  • use_optuna (bool, default=False) -- Whether to use Optuna for hyperparameter optimization. If not set, GridSearchCV from scikit-learn will be used.

  • n_trials (int, default=100) -- Number of trials for Optuna hyperparameter search.

  • timeout (int, default=3600) -- Time limit (in seconds) for Optuna hyperparameter search. This is a soft timeout: when n_jobs > 1, trials already in flight when the timeout elapses will run to completion, so actual runtime may exceed this value.

  • fold_col (str or list[str], default='Fold') -- Name(s) of the column(s) containing the CV fold number(s). If a list is provided, models will be benchmarked in an nxk-fold CV, where n is the number of repeats and k is the number of folds. If a single string is provided, it will be treated as a single fold column. nxk-fold CV does not currently support nested CV and final hyperparameter tuning.

  • main_metric (str, default='R2') -- Main metric to use for model selection. This will be used to infer the prediction task (classification or regression).

  • sec_metrics (list[str], default=['MSE', 'MAE']) -- Secondary metrics to use for model selection.

  • parametric (str or bool, default='auto') -- Whether to use parametric tests. If 'auto', the assumptions of parametric tests will be checked, and parametric tests will be used if the assumptions are met.

  • impute (str or float or int or None, default=None) -- Method to use for imputing missing values. If None, no imputation will be performed. Valid choices are 'mean', 'median', 'knn', or a float or int value for constant imputation.

  • remove_constant (float or None, default=None) -- If specified, features with variance below this threshold will be removed. If None, no features are removed.

  • remove_correlated (float or None, default=None) -- If specified, features with correlation above this threshold will be removed. If None, no features are removed.

  • scaler (str or None, default=None) -- Type of scaler to use, if the data is to be scaled first. Valid choices are 'Standard' and 'MinMax'.

  • custom_models (dict[str, None | dict[str, dict] | tuple[dict[str, dict], dict[str, dict]]] or None, default=None) -- Dictionary of models to use for benchmarking. If None, default models will be used. The keys should be the model names, and the values should be dictionaries of starting hyperparameters for the model, and/or a dictionary of hyperparameter search grids. Default models are defined in astra.models.classification and astra.models.regression.

  • n_jobs (int, default=1) -- Number of jobs to run in parallel for hyperparameter tuning.

  • test_mode (bool, default=False) -- Exists for compatibility with the unit tests. If True, will run in test mode.

Raises:

ValueError -- If any of the arguments are invalid.

Return type:

None