astra.utils

This module contains utility functions used in the package.

Functions

build_model(model_class[, impute, ...])

Build a scikit-learn model with optional preprocessing steps.

get_data(data, features)

Reads a CSV, pickle, or parquet file and returns a pandas DataFrame.

get_estimator_name(model)

Get the name of the estimator from a scikit-learn model.

get_models(main_metric, sec_metrics[, ...])

Get models and their hyperparameters based on the main metric and secondary metrics.

get_optuna_grid(hparam_grid)

Convert a hyperparameter grid to an Optuna-compatible format.

get_scores(cv_results_df, main_metric, ...)

Get means and standard deviations of the main and secondary metrics from the CV results.

load_config(file_path)

Load configuration from a YAML file.

print_file_console(file, message[, mode, end])

Print a message to a file and the console.

print_final_results(final_model_name, ...[, ...])

Prints the final results to the console and a file.

print_performance(model_name, results_dict)

Print the performance of a model based on the results dictionary.

astra.utils.build_model(model_class: BaseEstimator, impute: str | float | int | None = None, remove_constant: float | None = None, remove_correlated: float | None = None, scaler: str | None = None) BaseEstimator | Pipeline[source]

Build a scikit-learn model with optional preprocessing steps.

Parameters:
  • model_class (BaseEstimator) -- A scikit-learn model class to be instantiated.

  • impute (str | float | int or None, default None) -- Imputation strategy to apply before the model. Valid options are 'mean', 'median', 'knn', or a numeric value for constant imputation. If None, no imputation is applied.

  • remove_constant (float or None, default None) -- Threshold for variance to remove constant features. If None, no features are removed.

  • remove_correlated (float or None, default None) -- Threshold for correlation to remove correlated features. If None, no features are removed.

  • scaler (str or None, default None) -- Type of scaler to apply before the model. Valid options are 'MinMax' or 'Standard'. If None, no scaling is applied.

Returns:

A scikit-learn model or a Pipeline with the specified preprocessing steps.

Return type:

BaseEstimator or Pipeline

Raises:

ValueError -- If an unknown scaler or imputation strategy is provided, or if remove_constant or remove_correlated are not numeric values.

astra.utils.get_data(data: str, features: str) DataFrame[source]

Reads a CSV, pickle, or parquet file and returns a pandas DataFrame.

Parameters:
  • data (str) -- Path to the data file.

  • features (str) -- Name of the column containing the features.

Returns:

A pandas DataFrame containing the data.

Return type:

pd.DataFrame

astra.utils.get_estimator_name(model: BaseEstimator) str[source]

Get the name of the estimator from a scikit-learn model.

Parameters:

model (BaseEstimator) -- A scikit-learn model. Can be a Pipeline or a direct estimator.

Returns:

The name of the estimator.

Return type:

str

astra.utils.get_models(main_metric: str, sec_metrics: list[str], scaler: str | None = None, custom_models: dict[str, None | dict[str, dict] | tuple[dict[str, dict], dict[str, dict]]] | None = None, use_optuna: bool = False) tuple[dict[str, BaseEstimator], dict[str, dict[str, list]], dict[str, dict[str, list]] | None][source]

Get models and their hyperparameters based on the main metric and secondary metrics.

Parameters:
  • main_metric (str) -- The main metric to determine the type of models.

  • sec_metrics (list of str) -- List of secondary metrics to validate against the main metric.

  • scaler (str or None, default None) -- The type of scaler used. If 'Standard', some models are excluded.

  • 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.

  • use_optuna (bool, default False) -- Whether to retrieve hyperparameter search dictionaries suitable for Optuna. Note that we do not currently support custom hyperparameter grids when using Optuna.

Returns:

A tuple containing: - A dictionary of models with their names as keys and scikit-learn estimators as values. - A dictionary of hyperparameters for each model. - A dictionary of custom hyperparameters if provided, otherwise None.

Return type:

tuple of dicts

Raises:

ValueError -- If the main metric or any secondary metric is not recognized.

astra.utils.get_optuna_grid(hparam_grid: dict[str, list]) dict[str, BaseDistribution][source]

Convert a hyperparameter grid to an Optuna-compatible format.

Parameters:

hparam_grid (dict[str, list]) -- A dictionary where keys are hyperparameter names and values are lists of possible values. For numeric values, the minimum and maximum will be used to create a uniform distribution. For categorical values, a categorical distribution will be created.

Returns:

A dictionary where keys are hyperparameter names and values are Optuna distribution objects.

Return type:

dict[str, BaseDistribution]

astra.utils.get_scores(cv_results_df: DataFrame, main_metric: str, sec_metrics: list[str], n_folds: int) tuple[dict[str, list[float]], float, float, float, dict[str, tuple[float, float, float]]][source]

Get means and standard deviations of the main and secondary metrics from the CV results.

Parameters:
  • cv_results_df (pd.DataFrame) -- DataFrame containing the CV results.

  • main_metric (str) -- The main metric to extract results for.

  • sec_metrics (list of str) -- Secondary metrics to extract results for.

  • n_folds (int) -- Number of folds used in the CV.

Returns:

A tuple containing: - A dictionary with metrics as keys and lists of scores as values. - Mean, standard deviation and median of the main metric, - and a dictionary with secondary metrics as keys and tuples of (mean, std, median) as values.

Return type:

tuple of (dict, float, float, float, dict)

astra.utils.load_config(file_path: str) tuple[source]

Load configuration from a YAML file.

Parameters:

file_path (str) -- Path to the YAML configuration file.

Returns:

Configuration loaded from the YAML file.

Return type:

dict

astra.utils.print_file_console(file: str, message: str, mode: str = 'a', end: str = '\n')[source]

Print a message to a file and the console.

Parameters:
  • file (str) -- Path to the file where the message will be written.

  • message (str) -- The message to print.

  • mode (str, default 'a') -- File mode for writing. Default is append mode.

  • end (str, default 'n') -- String appended after the message. Default is newline.

Return type:

None

astra.utils.print_final_results(final_model_name: str, final_hyperparameters: dict[str, int | float | str], main_metric: str, mean_score_main: float, std_score_main: float, median_score_main: float, sec_metrics_scores: dict[str, tuple[float, float, float]], file: str | None = None) None[source]

Prints the final results to the console and a file.

Parameters:
  • final_model_name (str) -- Name of the final model.

  • final_hyperparameters (dict[str, int | float | str]) -- Hyperparameters of the final model.

  • main_metric (str) -- Main metric.

  • mean_score_main (float) -- Mean score of the main metric.

  • std_score_main (float) -- Standard deviation of the main metric.

  • median_score_main (float) -- Median score of the main metric.

  • sec_metrics_scores (dict[str, tuple[float, float, float]]) -- Dictionary containing secondary metrics scores (mean, std, median).

  • file (str or None, default None) -- If provided, the output will additionally be written to this file.

Return type:

None

astra.utils.print_performance(model_name: str, results_dict: dict[str, list[float]], file: str | None = None)[source]

Print the performance of a model based on the results dictionary.

Parameters:
  • model_name (str) -- Name of the model.

  • results_dict (dict[str, list[float]]) -- Dictionary containing performance metrics.

  • file (str or None, default None) -- If provided, the output will additionally be written to this file.

Return type:

None