diff --git a/autoemulate/compare.py b/autoemulate/compare.py index d74faf88..ea41e48e 100644 --- a/autoemulate/compare.py +++ b/autoemulate/compare.py @@ -1,9 +1,12 @@ +from typing import Optional + import matplotlib.pyplot as plt import numpy as np import pandas as pd from sklearn.decomposition import PCA from sklearn.metrics import make_scorer from sklearn.model_selection import cross_validate +from sklearn.model_selection import KFold from sklearn.model_selection import PredefinedSplit from sklearn.model_selection import train_test_split from sklearn.preprocessing import StandardScaler @@ -36,22 +39,22 @@ def __init__(self): def setup( self, - X, - y, - param_search=False, - param_search_type="random", - param_search_iters=20, - test_set_size=0.2, - scale=True, - scaler=StandardScaler(), - reduce_dim=False, - dim_reducer=PCA(), - fold_strategy="kfold", - folds=5, - n_jobs=None, - model_subset=None, - log_to_file=False, - ): + X: np.ndarray, + y: np.ndarray, + param_search: bool = False, + param_search_type: str = "random", + param_search_iters: int = 20, + test_set_size: float = 0.2, + scale: bool = True, + scaler: StandardScaler = StandardScaler(), + reduce_dim: bool = False, + dim_reducer: PCA = PCA(), + fold_strategy: str = "kfold", + folds: int = 5, + n_jobs: Optional[int] = None, + model_subset: Optional[list] = None, + log_to_file: bool = False, + ) -> None: """Sets up the automatic emulation. Parameters @@ -119,7 +122,9 @@ def setup( self.is_set_up = True self.cv_results = {} - def _check_input(self, X, y): + def _check_input( + self, X: np.ndarray, y: np.ndarray + ) -> tuple[np.ndarray, np.ndarray]: """Checks and possibly converts the input data. Parameters @@ -140,7 +145,7 @@ def _check_input(self, X, y): y = y.astype("float32") # needed for pytorch models return X, y - def _get_metrics(self, METRIC_REGISTRY): + def _get_metrics(self, METRIC_REGISTRY: dict) -> list: """ Get metrics from REGISTRY @@ -156,7 +161,7 @@ def _get_metrics(self, METRIC_REGISTRY): """ return [metric for metric in METRIC_REGISTRY.values()] - def _get_cv(self, CV_REGISTRY, fold_strategy, folds): + def _get_cv(self, CV_REGISTRY: dict, fold_strategy: str, folds: int) -> KFold: """Get cross-validation strategy from REGISTRY Parameters @@ -175,7 +180,7 @@ def _get_cv(self, CV_REGISTRY, fold_strategy, folds): """ return CV_REGISTRY[fold_strategy](folds=folds, shuffle=True) - def compare(self): + def compare(self) -> pd.DataFrame: """Compares the emulator models on the data. self.setup() must be run first. Returns @@ -190,7 +195,12 @@ def compare(self): self.scores_df = pd.DataFrame( columns=["model", "metric", "fold", "score"] ).astype( - {"model": "object", "metric": "object", "fold": "int64", "score": "float64"} + { + "model": "object", + "metric": "object", + "fold": "int64", + "score": "float64", + } ) for i in range(len(self.models)): @@ -246,7 +256,7 @@ def compare(self): return self.best_model - def get_model(self, rank=1, metric="r2"): + def get_model(self, rank: int = 1, metric: str = "r2"): # TODO: add return type """Get a fitted model based on it's rank in the comparison. Parameters @@ -284,7 +294,7 @@ def get_model(self, rank=1, metric="r2"): return chosen_model - def refit_model(self, model): + def refit_model(self, model): # TODO: add model type """Refits a model on the full data. Parameters @@ -303,8 +313,22 @@ def refit_model(self, model): model.fit(self.X, self.y) return model - def save_model(self, model=None, filepath=None): - """Saves the best model to disk.""" + def save_model( + self, model=None, filepath: str = None + ) -> None: # TODO add model type + """Saves the best model to disk. + + Parameters + ---------- + model : object + Fitted model. + filepath : str + Path to the model file. + + Returns + ------- + None + """ if not hasattr(self, "best_model"): raise RuntimeError("Must run compare() before save_model()") serialiser = ModelSerialiser() @@ -316,15 +340,27 @@ def save_model(self, model=None, filepath=None): serialiser.save_model(model, filepath) - def load_model(self, filepath=None): - """Loads a model from disk.""" + def load_model(self, filepath: str = None): # TODO add return type (model type) + """Loads a model from disk. + + Parameters + ---------- + filepath : str + Path to the model file. + + Returns + ------- + model : object + Loaded model. + """ serialiser = ModelSerialiser() if filepath is None: raise ValueError("Filepath must be provided") return serialiser.load_model(filepath) - def print_results(self, model=None, sort_by="r2"): + # TODO for print_results: suggestion, rename model to model_name here to not confuse with other references to the model object + def print_results(self, model: Optional[str] = None, sort_by: str = "r2") -> None: """Print cv results. Parameters @@ -342,13 +378,14 @@ def print_results(self, model=None, sort_by="r2"): sort_by=sort_by, ) + # TODO for plot_results: suggestion, rename model to model_name here to not confuse with other references to the model object def plot_results( self, - model=None, - plot_type="actual_vs_predicted", - n_cols=3, - figsize=None, - output_index=0, + model: Optional[str] = None, + plot_type: str = "actual_vs_predicted", + n_cols: int = 3, + figsize: Optional[tuple] = None, + output_index: int = 0, ): """Plots the results of the cross-validation. @@ -379,7 +416,7 @@ def plot_results( output_index=output_index, ) - def evaluate_model(self, model=None): + def evaluate_model(self, model=None) -> pd.DataFrame: # TODO add model type """ Evaluates the model on the hold-out set. @@ -411,7 +448,13 @@ def evaluate_model(self, model=None): return scores_df - def plot_model(self, model, plot="standard", n_cols=2, figsize=None): + def plot_model( + self, + model, + plot: str = "standard", + n_cols: int = 2, + figsize: Optional[tuple] = None, + ) -> None: # TODO add model type """Plots the model predictions vs. the true values. Parameters @@ -424,7 +467,14 @@ def plot_model(self, model, plot="standard", n_cols=2, figsize=None): “residual” draws the residuals, i.e. difference between observed and predicted values, (y-axis) vs. the predicted values (x-axis). n_cols : int, optional Number of columns in the plot grid for multi-output. Default is 2. + figsize : tuple, optional + Overrides the default figure size. """ _plot_model( - model, self.X[self.test_idxs], self.y[self.test_idxs], plot, n_cols, figsize + model, + self.X[self.test_idxs], + self.y[self.test_idxs], + plot, + n_cols, + figsize, ) diff --git a/autoemulate/cross_validate.py b/autoemulate/cross_validate.py index 328de066..2e6ad9be 100644 --- a/autoemulate/cross_validate.py +++ b/autoemulate/cross_validate.py @@ -1,3 +1,5 @@ +from typing import TYPE_CHECKING + import numpy as np import pandas as pd from sklearn.metrics import make_scorer @@ -5,10 +7,68 @@ from sklearn.model_selection import PredefinedSplit from sklearn.model_selection import train_test_split +from autoemulate.types import ArrayLike +from autoemulate.types import MatrixLike +from autoemulate.types import Union from autoemulate.utils import get_model_name +if TYPE_CHECKING: + from logging import Logger + from .types import Iterable + from sklearn.model_selection import BaseCrossValidator + from sklearn.model_selection import BaseShuffleSplit + from sklearn.pipeline import Pipeline + + +def run_cv( + X: MatrixLike, + y: Union[MatrixLike, ArrayLike], + cv: Union[int, BaseCrossValidator, Iterable, BaseShuffleSplit], + model: Pipeline, # TODO: Verify that this is correct + metrics: list, + n_jobs: int, + logger: Logger, +): + """Runs cross-validation on a model. + + Parameters + ---------- + X : {array-like, sparse matrix} of shape (n_samples, n_features) + The data to fit. Can be for example a list, or an array. + y : array-like of shape (n_samples,) or (n_samples, n_outputs), default=None + The target variable to try to predict in the case of supervised learning. + cv : int, cross-validation generator or an iterable, default=None + Determines the cross-validation splitting strategy. + Possible inputs for cv are: + + - None, to use the default 5-fold cross validation, + - int, to specify the number of folds in a `(Stratified)KFold`, + - CV splitter, + - An iterable yielding (train, test) splits as arrays of indices. + + For int/None inputs, if the estimator is a classifier and ``y`` is + either binary or multiclass, :class:`StratifiedKFold` is used. In all + other cases, :class:`KFold` is used. These splitters are instantiated + with `shuffle=False` so the splits will be the same across calls. + + Refer :ref:`User Guide ` for the various + cross-validation strategies that can be used here. + model : sklearn.pipeline.Pipeline + Model to cross-validate. + metrics : list + List of metrics to use for cross-validation. + n_jobs : int + Number of jobs to run in parallel. + logger : logging.Logger + Logger object. -def run_cv(X, y, cv, model, metrics, n_jobs, logger): + Returns + ------- + fitted_model : sklearn.pipeline.Pipeline + Fitted model. + cv_results : dict + Results of the cross-validation. + """ model_name = get_model_name(model) # The metrics we want to use for cross-validation @@ -39,22 +99,23 @@ def run_cv(X, y, cv, model, metrics, n_jobs, logger): return fitted_model, cv_results -def update_scores_df(scores_df, model, cv_results): +# TODO for update_scores_df: suggestion, rename model to model_name here to not confuse with other references to the model object +def update_scores_df(scores_df: pd.DataFrame, model: str, cv_results: dict) -> None: """Updates the scores dataframe with the results of the cross-validation. Parameters ---------- - scores_df : pandas.DataFrame - DataFrame with columns "model", "metric", "fold", "score". - model_name : str - Name of the model. - cv_results : dict - Results of the cross-validation. + scores_df : pandas.DataFrame + DataFrame with columns "model", "metric", "fold", "score". + model_name : str + Name of the model. + cv_results : dict + Results of the cross-validation. Returns ------- - None - Modifies the self.scores_df DataFrame in-place. + None + Modifies the self.scores_df DataFrame in-place. """ # Gather scores from each metric diff --git a/autoemulate/cv.py b/autoemulate/cv.py index 13307899..2476690a 100644 --- a/autoemulate/cv.py +++ b/autoemulate/cv.py @@ -1,21 +1,24 @@ from sklearn.model_selection import KFold from sklearn.model_selection import StratifiedKFold +from .types import Optional -def kfold(folds=None, shuffle=True): + +# TODO: KFold seems to only accept a n_splits parameter of Int type so we should also enforce that here to avoid bugs +def kfold(folds: Optional[int] = None, shuffle: bool = True) -> KFold: """scikit-learn class for k-fold cross validation. Parameters ---------- folds : int - Number of folds. + Number of folds. Must be at least 2. shuffle : bool - Whether or not to shuffle the data before splitting. + Whether or not to shuffle the data before splitting. Returns ------- kfold : sklearn.model_selection.KFold - An instance of the KFold class. + An instance of the KFold class. """ return KFold(n_splits=folds, shuffle=shuffle) diff --git a/autoemulate/data_splitting.py b/autoemulate/data_splitting.py index fdc8e610..5936da0e 100644 --- a/autoemulate/data_splitting.py +++ b/autoemulate/data_splitting.py @@ -1,8 +1,13 @@ import numpy as np from sklearn.model_selection import train_test_split +from .types import ArrayLike +from .types import Optional -def split_data(X, test_size=0.2, random_state=None): + +def split_data( + X: ArrayLike, test_size: float = 0.2, random_state: Optional[int] = None +) -> tuple[ArrayLike, ArrayLike]: """Splits the data into training and testing sets. Parameters diff --git a/autoemulate/datasets.py b/autoemulate/datasets.py index 6b32deca..80c1adfc 100644 --- a/autoemulate/datasets.py +++ b/autoemulate/datasets.py @@ -7,7 +7,13 @@ data_dir = Path(__file__).parent.parent / "data" -def fetch_data(dataset, split=False, test_size=0.2, random_state=42): +# TODO: Verify that this is a return of a tuple of np.ndarray elements +def fetch_data( + dataset: str, + split: bool = False, + test_size: float = 0.2, + random_state: int = 42, +) -> tuple[np.ndarray, np.ndarray]: """ Fetch a dataset by name. diff --git a/autoemulate/emulators/gaussian_process.py b/autoemulate/emulators/gaussian_process.py index 106ed452..121a9e9d 100644 --- a/autoemulate/emulators/gaussian_process.py +++ b/autoemulate/emulators/gaussian_process.py @@ -8,6 +8,10 @@ from skopt.space import Categorical from skopt.space import Real +from ..types import ArrayLike +from ..types import Literal +from ..types import Self + class GaussianProcess(BaseEstimator, RegressorMixin): """Gaussian Process Emulator. @@ -15,11 +19,11 @@ class GaussianProcess(BaseEstimator, RegressorMixin): Wraps Gaussian Process Regression from the mogp_emulator package. """ - def __init__(self, nugget="fit"): + def __init__(self, nugget: str = "fit"): """Initializes a GaussianProcess object.""" self.nugget = nugget - def fit(self, X, y): + def fit(self, X: ArrayLike, y: ArrayLike) -> Self: """Fits the emulator to the data. Parameters @@ -41,7 +45,7 @@ def fit(self, X, y): self.is_fitted_ = True return self - def predict(self, X, return_std=False): + def predict(self, X: ArrayLike, return_std: bool = False) -> ArrayLike: """Predicts the output of the simulator for a given input. Parameters @@ -66,8 +70,21 @@ def predict(self, X, return_std=False): else: return np.asarray(self.model_.predict(X).mean) - def get_grid_params(self, search_type="random"): - """Returns the grid parameters of the emulator.""" + def get_grid_params( + self, search_type: Literal["random", "bayes"] = "random" + ) -> dict[str, list]: + """Returns the grid parameters of the emulator. + + Parameters + ---------- + search_type : {"random", "bayes"} + The type of search to be used. + + Returns + ------- + dict[str, list] + The grid parameters of the emulator. + """ param_space_random = { "nugget": ["fit", "adaptive", "pivot"], } @@ -80,7 +97,16 @@ def get_grid_params(self, search_type="random"): elif search_type == "bayes": param_space = param_space_bayes + # TODO: Should this raise an error if the search type is not recognised? + return param_space - def _more_tags(self): + def _more_tags(self) -> dict: + """Returns more tags for the estimator. + + Returns + ------- + dict + The tags of the emulator. + """ return {"multioutput": False} diff --git a/autoemulate/emulators/gaussian_process_sk.py b/autoemulate/emulators/gaussian_process_sk.py index a9c46e47..9263599c 100644 --- a/autoemulate/emulators/gaussian_process_sk.py +++ b/autoemulate/emulators/gaussian_process_sk.py @@ -13,6 +13,11 @@ from skopt.space import Integer from skopt.space import Real +from ..types import ArrayLike +from ..types import Literal +from ..types import MatrixLike +from ..types import Optional +from ..types import Self from autoemulate.utils import suppress_convergence_warnings @@ -24,15 +29,34 @@ class GaussianProcessSk(BaseEstimator, RegressorMixin): def __init__( self, - kernel=RBF(), - alpha=1e-10, - optimizer="fmin_l_bfgs_b", - n_restarts_optimizer=20, - normalize_y=True, - copy_X_train=True, - random_state=None, + kernel: RBF = RBF(), # TODO: Is this the correct type then? + alpha: float = 1e-10, + optimizer: str = "fmin_l_bfgs_b", + n_restarts_optimizer: int = 20, + normalize_y: bool = True, + copy_X_train: bool = True, + random_state: Optional[int] = None, ): - """Initializes a GaussianProcess object.""" + """Initializes a GaussianProcess object. + + Parameters + ---------- + kernel : kernel object + The kernel specifying the covariance function of the GP. + alpha : float + Value added to the diagonal of the kernel matrix during fitting. + This can prevent a potential numerical issue during fitting. + optimizer : str + The optimizer to use for optimizing the kernel's parameters. + n_restarts_optimizer : int + The number of restarts of the optimizer for finding the kernel's parameters. + normalize_y : bool + Whether to normalize the target values. + copy_X_train : bool + Whether to make a copy of the training data. + random_state : int + The random state to use for the emulator. + """ self.kernel = kernel self.alpha = alpha self.optimizer = optimizer @@ -41,7 +65,8 @@ def __init__( self.copy_X_train = copy_X_train self.random_state = random_state - def fit(self, X, y): + # TODO: Should X be MatrixLike here? + def fit(self, X: ArrayLike, y: ArrayLike) -> Self: """Fits the emulator to the data. Parameters @@ -72,7 +97,7 @@ def fit(self, X, y): self.is_fitted_ = True return self - def predict(self, X, return_std=False): + def predict(self, X: MatrixLike, return_std: bool = False) -> ArrayLike: """Predicts the output of the simulator for a given input. Parameters @@ -93,7 +118,9 @@ def predict(self, X, return_std=False): check_is_fitted(self, "is_fitted_") return self.model_.predict(X, return_std=return_std) - def get_grid_params(self, search_type="random"): + def get_grid_params( + self, search_type: Literal["random", "bayes"] = "random" + ) -> dict[str, list]: """Returns the grid parameters of the emulator.""" param_space_random = { "kernel": [ @@ -118,7 +145,16 @@ def get_grid_params(self, search_type="random"): elif search_type == "bayes": param_space = param_space_bayes + # TODO: Should this raise an error if the search type is not recognised? + return param_space - def _more_tags(self): + def _more_tags(self) -> dict: + """Returns more tags for the estimator. + + Returns + ------- + dict + The multioutput tag. + """ return {"multioutput": True} diff --git a/autoemulate/emulators/gradient_boosting.py b/autoemulate/emulators/gradient_boosting.py index 266a76b3..add9c6d6 100644 --- a/autoemulate/emulators/gradient_boosting.py +++ b/autoemulate/emulators/gradient_boosting.py @@ -11,26 +11,93 @@ from skopt.space import Integer from skopt.space import Real +from ..types import Any +from ..types import ArrayLike +from ..types import Literal +from ..types import MatrixLike +from ..types import Optional +from ..types import Self +from ..types import Union + class GradientBoosting(BaseEstimator, RegressorMixin): """Gradient Boosting Emulator. Wraps Gradient Boosting regression from scikit-learn. + + Parameters + ---------- + loss : {'squared_error', 'ls', 'lad', 'huber', 'quantile'}, default='squared_error' + The loss function to be optimized. 'squared_error' refers to the + ordinary least squares fit. 'ls' refers to least squares fit. 'lad' + refers to least absolute deviation fit. 'huber' is a combination of + least squares and least absolute deviation. 'quantile' allows quantile + regression (use alpha to specify the quantile). + learning_rate : float, default=0.1 + The learning rate shrinks the contribution of each tree. There is a + trade-off between learning_rate and n_estimators. + n_estimators : int, default=100 + The number of boosting stages to be run. Gradient boosting is fairly + robust to over-fitting so a large number usually results in better + performance. + max_depth : int, default=3 + The maximum depth of the individual estimators. The maximum depth + limits the number of nodes in the tree. Tune this parameter for best + performance; the best value depends on the interaction of the input + variables. + min_samples_split : int, default=2 + The minimum number of samples required to split an internal node. + min_samples_leaf : int, default=1 + The minimum number of samples required to be at a leaf node. + subsample : float, default=1.0 + The fraction of samples to be used for fitting the individual base + learners. If smaller than 1.0 this results in Stochastic Gradient + Boosting. subsample interacts with the parameter n_estimators. Choosing + subsample < 1.0 leads to a reduction of variance and an increase in + bias. + max_features : {'auto', 'sqrt', 'log2'}, int or float, default=None + The number of features to consider when looking for the best split: + - If int, then consider max_features features at each split. + - If float, then max_features is a fraction and + int(max_features * n_features) features are considered at each split. + - If 'auto', then max_features=sqrt(n_features). + - If 'sqrt', then max_features=sqrt(n_features). + - If 'log2', then max_features=log2(n_features). + - If None, then max_features=n_features. + ccp_alpha : non-negative float, default=0.0 + Complexity parameter used for Minimal Cost-Complexity Pruning. The + subtree with the largest cost complexity that is smaller than + ccp_alpha will be chosen. By default, no pruning is performed. + n_iter_no_change : int, default=None + Number of iterations with no improvement to wait before stopping + fitting. Convergence is checked against the training loss or the + validation loss depending on the early_stopping parameter. + random_state : int, RandomState instance or None, default=None + Controls the random seed given at each base learner at each boosting + iteration. In addition, it controls the random permutation of the + features at each split. It also controls the random spliting of the + training data to obtain a validation set if n_iter_no_change is not + None. Pass an int for reproducible output across multiple function + calls. """ def __init__( self, - loss="squared_error", - learning_rate=0.1, - n_estimators=100, - max_depth=3, - min_samples_split=2, - min_samples_leaf=1, - subsample=1.0, - max_features=None, - ccp_alpha=0.0, - n_iter_no_change=None, - random_state=None, + loss: Literal[ + "squared_error", "ls", "lad", "huber", "quantile" + ] = "squared_error", + learning_rate: float = 0.1, + n_estimators: int = 100, + max_depth: int = 3, + min_samples_split: int = 2, + min_samples_leaf: int = 1, + subsample: float = 1.0, + max_features: Optional[ + Union[Literal["auto", "sqrt", "log2"], int, float] + ] = None, + ccp_alpha: float = 0.0, + n_iter_no_change: Optional[int] = None, + random_state: Optional[int] = None, ): """Initializes a GradientBoosting object.""" self.loss = loss @@ -45,7 +112,7 @@ def __init__( self.n_iter_no_change = n_iter_no_change self.random_state = random_state - def fit(self, X, y): + def fit(self, X: MatrixLike, y: ArrayLike) -> Self: """Fits the emulator to the data. Parameters @@ -82,7 +149,7 @@ def fit(self, X, y): self.is_fitted_ = True return self - def predict(self, X): + def predict(self, X: MatrixLike) -> ArrayLike: """Predicts the output of the emulator for a given input. Parameters @@ -99,7 +166,9 @@ def predict(self, X): check_is_fitted(self, "is_fitted_") return self.model_.predict(X) - def get_grid_params(self, search_type="random"): + def get_grid_params( + self, search_type: Literal["random", "bayes"] = "random" + ) -> dict[str, Any]: """Returns the grid parameters of the emulator.""" param_space_random = { "learning_rate": loguniform(0.01, 0.2), @@ -128,7 +197,16 @@ def get_grid_params(self, search_type="random"): elif search_type == "bayes": param_space = param_space_bayes + # TODO: Should this raise an error if the search type is not recognised? + return param_space - def _more_tags(self): + def _more_tags(self) -> dict[str, bool]: + """Returns more tags for the estimator. + + Returns + ------- + dict + Dictionary of tags. + """ return {"multioutput": False} diff --git a/autoemulate/emulators/neural_net_sk.py b/autoemulate/emulators/neural_net_sk.py index 43b540fe..199fbe8a 100644 --- a/autoemulate/emulators/neural_net_sk.py +++ b/autoemulate/emulators/neural_net_sk.py @@ -9,6 +9,11 @@ from skopt.space import Categorical from skopt.space import Real +from ..types import Any +from ..types import ArrayLike +from ..types import Literal +from ..types import MatrixLike +from ..types import Self from autoemulate.utils import suppress_convergence_warnings @@ -16,22 +21,45 @@ class NeuralNetSk(BaseEstimator, RegressorMixin): """Multi-layer perceptron Emulator. Wraps MLPRegressor from scikit-learn. + + Parameters + ---------- + hidden_layer_sizes : tuple, default=(100, 100) + The ith element represents the number of neurons in the ith hidden layer. + activation : {'identity', 'logistic', 'tanh', 'relu'}, default='relu' + Activation function for the hidden layer. + solver : {'lbfgs', 'sgd', 'adam'}, default='adam' + The solver for weight optimization. + alpha : float, default=0.0001 + L2 penalty (regularization term) parameter. + learning_rate : {'constant', 'invscaling', 'adaptive'}, default='constant' + Learning rate schedule for weight updates. + learning_rate_init : float, default=0.001 + The initial learning rate used. It controls the step-size in updating the weights. + max_iter : int, default=200 + Maximum number of iterations. + tol : float, default=1e-4 + Tolerance for the optimization. + random_state : int, RandomState instance or None, default=None + If int, random_state is the seed used by the random number generator; + If RandomState instance, random_state is the random number generator; + If None, the random number generator is the RandomState instance used by np.random. """ def __init__( self, - hidden_layer_sizes=( + hidden_layer_sizes: tuple[int, int] = ( 100, 100, ), - activation="relu", - solver="adam", - alpha=0.0001, - learning_rate="constant", - learning_rate_init=0.001, - max_iter=200, - tol=1e-4, - random_state=None, + activation: Literal["identity", "logistic", "tanh", "relu"] = "relu", + solver: Literal["lbfgs", "sgd", "adam"] = "adam", + alpha: float = 0.0001, + learning_rate: Literal["constant", "invscaling", "adaptive"] = "constant", + learning_rate_init: float = 0.001, + max_iter: int = 200, + tol: float = 1e-4, + random_state=None, # TODO add type ): """Initializes an MLPRegressor object.""" self.hidden_layer_sizes = hidden_layer_sizes @@ -44,7 +72,7 @@ def __init__( self.tol = tol self.random_state = random_state - def fit(self, X, y): + def fit(self, X: MatrixLike, y: ArrayLike) -> Self: """Fits the emulator to the data. Parameters @@ -79,7 +107,7 @@ def fit(self, X, y): self.is_fitted_ = True return self - def predict(self, X): + def predict(self, X: MatrixLike) -> ArrayLike: """Predicts the output of the simulator for a given input. Parameters @@ -96,8 +124,21 @@ def predict(self, X): check_is_fitted(self, "is_fitted_") return self.model_.predict(X) - def get_grid_params(self, search_type="random"): - """Returns the grid parameters of the emulator.""" + def get_grid_params( + self, search_type: Literal["random", "bayes"] = "random" + ) -> dict[str, Any]: + """Returns the grid parameters of the emulator. + + Parameters + ---------- + search_type : str + The type of search to use. Either "random" or "bayes". + + Returns + ------- + dict + The grid parameters of the emulator. + """ param_space_random = { "hidden_layer_sizes": [ (50,), @@ -132,12 +173,21 @@ def get_grid_params(self, search_type="random"): elif search_type == "bayes": param_space = param_space_bayes + # TODO: Should this raise an error if the search type is not recognised? + return param_space - def _more_tags(self): + def _more_tags(self) -> dict[str, bool]: + """Returns more tags for the estimator. + + Returns + ------- + dict + The multioutput tag. + """ return {"multioutput": True} - # def score(self, X, y, metric): + # def score(self, X: ArrayLike, y: ArrayLike, metric: Literal["rsme", "r2"]) -> float: # """Returns the score of the emulator. # Parameters diff --git a/autoemulate/emulators/neural_net_torch.py b/autoemulate/emulators/neural_net_torch.py index a0691505..4b7d5579 100644 --- a/autoemulate/emulators/neural_net_torch.py +++ b/autoemulate/emulators/neural_net_torch.py @@ -2,7 +2,6 @@ # to make it compatible with scikit-learn. Works with cross_validate and GridSearchCV, # but doesn't pass tests, because we're subclassing import warnings -from typing import List import numpy as np import torch @@ -12,18 +11,37 @@ from skorch import NeuralNetRegressor from skorch.callbacks import Callback +from ..types import ArrayLike +from ..types import List +from ..types import Literal +from ..types import MatrixLike +from ..types import Optional +from ..types import Self +from ..types import Union from autoemulate.emulators.neural_networks import get_module from autoemulate.utils import set_random_seed class InputShapeSetter(Callback): - """Callback to set input and output layer sizes dynamically.""" + """Callback to set input and output layer sizes dynamically. + + This is needed to support dynamic input size. + + Parameters + ---------- + net : NeuralNetRegressor + The neural network regressor. + X : Optional[Union[torch.Tensor, np.ndarray]], optional + The input samples, by default None. + y : Optional[Union[torch.Tensor, np.ndarray]], optional + The target values, by default None. + """ def on_train_begin( self, net: NeuralNetRegressor, - X: torch.Tensor | np.ndarray = None, - y: torch.Tensor | np.ndarray = None, + X: Optional[Union[torch.Tensor, np.ndarray]] = None, + y: Optional[Union[torch.Tensor, np.ndarray]] = None, **kwargs, ): if hasattr(net, "n_features_in_") and net.n_features_in_ != X.shape[-1]: @@ -44,13 +62,46 @@ class NeuralNetTorch(NeuralNetRegressor): module__input_size and module__output_size must be provided to define the input and output dimension of the data. + + Parameters + ---------- + module : str, optional + The module to use, by default "mlp" + criterion : torch.nn.Module, optional + The loss function, by default torch.nn.MSELoss + optimizer : torch.optim.Optimizer, optional + The optimizer to use, by default torch.optim.AdamW + lr : float, optional + The learning rate, by default 1e-3 + batch_size : int, optional + The batch size, by default 128 + max_epochs : int, optional + The maximum number of epochs, by default 1 + module__input_size : int, optional + The input size, by default 2 + module__output_size : int, optional + The output size, by default 1 + optimizer__weight_decay : float, optional + The weight decay, by default 0.0 + iterator_train__shuffle : bool, optional + Whether to shuffle the training data, by default True + callbacks : List[Callback], optional + The callbacks to use, by default [InputShapeSetter()] + train_split : bool, optional + Whether to split the data, by default False + verbose : int, optional + The verbosity level, by default 0 + random_state : int, optional + The random state, by default None + **kwargs + Additional keyword arguments to pass to the neural network regressor. """ def __init__( self, module: str = "mlp", - criterion=torch.nn.MSELoss, - optimizer=torch.optim.AdamW, + criterion: torch.nn.Module = torch.nn.MSELoss, # TODO: verify type here + optimizer: torch.optim.Optimizer = torch.optim.AdamW, # TODO: verify type here lr: float = 1e-3, batch_size: int = 128, max_epochs: int = 1, @@ -84,7 +135,20 @@ def __init__( ) self.initialize() - def set_params(self, **params): + def set_params(self, **params) -> Self: + """Set the parameters of the neural network regressor. + + Parameters + ---------- + **params + The parameters to set. If `random_state` is provided, it is set as + an attribute of the class. + + Returns + ------- + self + The neural network regressor with the new parameters. + """ if "random_state" in params: random_state = params.pop("random_state") if hasattr(self, "random_state"): @@ -95,7 +159,19 @@ def set_params(self, **params): self.initialize() return super().set_params(**params) - def initialize_module(self, reason=None): + def initialize_module(self, reason: Optional[str] = None) -> Self: + """Initializes the module. + + Parameters + ---------- + reason : str, optional + The reason for initializing the module, by default None + + Returns + ------- + self + The neural network regressor with the initialized module. + """ kwargs = self.get_params_for("module") if hasattr(self, "random_state"): kwargs["random_state"] = self.random_state @@ -103,13 +179,26 @@ def initialize_module(self, reason=None): self.module_ = module return self - def get_grid_params(self, search_type="random"): + def get_grid_params( + self, + search_type: Literal[ + "random", "bayes" + ] = "random", # TODO: Verify search_type types + ) -> dict[str, Union[bool, dict[str, str]]]: return self.module_.get_grid_params(search_type) - def __sklearn_is_fitted__(self): + def __sklearn_is_fitted__(self) -> bool: + """Private method to check if the model is fitted. This is used by scikit-learn.""" return hasattr(self, "n_features_in_") - def _more_tags(self): + def _more_tags(self) -> dict[str, Union[bool, dict[str, str]]]: + """Returns more tags for the estimator. + + Returns + ------- + dict + The tags of the neural network regressor. + """ return { "multioutput": True, "poor_score": True, @@ -122,7 +211,21 @@ def _more_tags(self): }, } - def check_data(self, X: np.ndarray, y: np.ndarray = None): + def check_data(self, X: ArrayLike, y: Optional[ArrayLike] = None) -> np.ndarray: + """Check the data for the neural network regressor. + + Parameters + ---------- + X : np.ndarray + The input samples. + y : np.ndarray, optional + The target values, by default None + + Returns + ------- + np.ndarray + The input samples. + """ if isinstance(y, np.ndarray): if np.iscomplex(X).any(): raise ValueError("Complex data not supported") @@ -158,20 +261,98 @@ def check_data(self, X: np.ndarray, y: np.ndarray = None): y = np.squeeze(y, axis=-1) return X, y - def fit_loop(self, X, y=None, epochs=None, **fit_params): + def fit_loop( + self, + X: ArrayLike, + y: Optional[ArrayLike] = None, + epochs: Optional[int] = None, + **fit_params, + ) -> Self: + """Loop to fit the neural network regressor. + + Parameters + ---------- + X : np.ndarray + The input samples. + y : np.ndarray, optional + The target values, by default None + epochs : int, optional + The number of epochs, by default None + **fit_params + Additional keyword arguments to pass to the fit method. + + Returns + ------- + self + The neural network regressor fitted to the data. + """ X, y = self.check_data(X, y) return super().fit_loop(X, y, epochs, **fit_params) - def partial_fit(self, X, y=None, classes=None, **fit_params): + def partial_fit( + self, + X: ArrayLike, + y: Optional[ArrayLike] = None, + classes: ArrayLike = None, + **fit_params, + ) -> Self: + """Partially fit the neural network regressor. + + Parameters + ---------- + X : np.ndarray + The input samples. + y : np.ndarray, optional + The target values, by default None + classes : np.ndarray, optional + The classes, by default None + **fit_params + Additional keyword arguments to pass to the fit method. + + Returns + ------- + self + The neural network regressor partially fitted to the data. + """ X, y = self.check_data(X, y) return super().partial_fit(X, y, classes, **fit_params) - def fit(self, X, y, **fit_params): + def fit(self, X: MatrixLike, y: Optional[ArrayLike], **fit_params) -> Self: + """ + Fits the emulator to the data. + + Parameters + ---------- + X : {array-like, sparse matrix}, shape (n_samples, n_features) + The training input samples. + y : array-like, shape (n_samples,) or (n_samples, n_outputs) + The target values (real numbers). + **fit_params + Additional keyword arguments to pass to the fit method. + + Returns + ------- + self : object + Returns self. + """ X, y = self.check_data(X, y) return super().fit(X, y, **fit_params) @torch.inference_mode() - def predict_proba(self, X): + def predict_proba(self, X: MatrixLike) -> np.ndarray: + """ + Predicts the output of the emulator for a given input. + + Parameters + ---------- + X : {array-like, sparse matrix}, shape (n_samples, n_features) + The input samples. + + Returns + ------- + y : ndarray of shape (n_samples, n_features) + The predicted values. + """ if not hasattr(self, "n_features_in_"): raise NotFittedError dtype = X.dtype if hasattr(X, "dtype") else None @@ -181,7 +362,21 @@ def predict_proba(self, X): y_pred = y_pred.astype(dtype) return y_pred - def infer(self, x: torch.Tensor, **fit_params): + def infer(self, x: torch.Tensor, **fit_params) -> np.ndarray: + """Predicts the output of the emulator for a given input. + + Parameters + ---------- + x : torch.Tensor + The input samples. + **fit_params + Additional keyword arguments to pass to the infer method. + + Returns + ------- + y_pred : np.ndarray + The predicted values. + """ if not hasattr(self, "n_features_in_"): setattr(self, "n_features_in_", x.size(1)) y_pred = super().infer(x, **fit_params) diff --git a/autoemulate/emulators/neural_networks/base.py b/autoemulate/emulators/neural_networks/base.py index 4583dcd3..34f97b60 100644 --- a/autoemulate/emulators/neural_networks/base.py +++ b/autoemulate/emulators/neural_networks/base.py @@ -1,6 +1,7 @@ import torch from torch import nn +from ...types import Optional from autoemulate.utils import set_random_seed @@ -12,9 +13,9 @@ class TorchModule(nn.Module): def __init__( self, module_name: str, - input_size: int = None, - output_size: int = None, - random_state: int = None, + input_size: Optional[int] = None, + output_size: Optional[int] = None, + random_state: Optional[int] = None, ): super(TorchModule, self).__init__() if random_state is not None: @@ -28,4 +29,5 @@ def get_grid_params(self, search_type: str = "random"): raise NotImplementedError("get_grid_params method not implemented.") def forward(self, X: torch.Tensor): + """Forward pass through the module""" raise NotImplementedError("forward method not implemented.") diff --git a/autoemulate/emulators/neural_networks/get_module.py b/autoemulate/emulators/neural_networks/get_module.py index 3b8a25ba..3054061c 100644 --- a/autoemulate/emulators/neural_networks/get_module.py +++ b/autoemulate/emulators/neural_networks/get_module.py @@ -1,11 +1,27 @@ +from ...types import Union from autoemulate.emulators.neural_networks import TorchModule from autoemulate.emulators.neural_networks.mlp import MLPModule -def get_module(module: str | TorchModule) -> TorchModule: +def get_module(module: Union[str, TorchModule]) -> TorchModule: """ Return the module class for NeuralNetRegressor. If `module` is already a TorchModule, then return it as is. + + Parameters + ---------- + module : str or TorchModule + The module class to use. + + Returns + ------- + TorchModule + The module class to use. + + Raises + ------ + NotImplementedError + If the module is not implemented. """ if not isinstance(module, str): return module diff --git a/autoemulate/emulators/neural_networks/mlp.py b/autoemulate/emulators/neural_networks/mlp.py index f1fc0f67..d709683d 100644 --- a/autoemulate/emulators/neural_networks/mlp.py +++ b/autoemulate/emulators/neural_networks/mlp.py @@ -7,6 +7,8 @@ from skopt.space import Real from torch import nn +from ...types import Literal +from ...types import Optional from autoemulate.emulators.neural_networks.base import TorchModule @@ -15,12 +17,14 @@ class MLPModule(TorchModule): def __init__( self, - input_size: int = None, - output_size: int = None, - random_state: int = None, + input_size: Optional[int] = None, + output_size: Optional[int] = None, + random_state: Optional[int] = None, hidden_layers: int = 1, hidden_size: int = 100, - hidden_activation: Tuple[callable] = nn.ReLU, + hidden_activation: tuple[ + callable + ] = nn.ReLU, # TODO: Is this really the correct type? ): super(MLPModule, self).__init__( module_name="mlp", @@ -37,7 +41,26 @@ def __init__( modules.append(nn.Linear(in_features=input_size, out_features=output_size)) self.model = nn.Sequential(*modules) - def get_grid_params(self, search_type: str = "random"): + def get_grid_params( + self, search_type: Literal["random", "bayes"] = "random" + ) -> dict[str, list]: + """Return the hyperparameter search space for the module. + + Parameters + ---------- + search_type : str, optional + The type of search space to return, by default "random" + + Returns + ------- + dict + The hyperparameter search space for the module + + Raises + ------ + ValueError + If the search type is not implemented + """ param_space = { "max_epochs": np.arange(10, 110, 10).tolist(), "batch_size": np.arange(2, 128, 2).tolist(), @@ -62,5 +85,17 @@ def get_grid_params(self, search_type: str = "random"): return param_space - def forward(self, X: torch.Tensor): + def forward(self, X: torch.Tensor) -> torch.Tensor: + """Forward pass through the module. + + Parameters + ---------- + X : torch.Tensor + The input data + + Returns + ------- + torch.Tensor + The output data + """ return self.model(X) diff --git a/autoemulate/emulators/polynomials.py b/autoemulate/emulators/polynomials.py index 468e7915..ae53a36d 100644 --- a/autoemulate/emulators/polynomials.py +++ b/autoemulate/emulators/polynomials.py @@ -10,19 +10,31 @@ from skopt.space import Categorical from skopt.space import Integer +from ..types import ArrayLike +from ..types import List +from ..types import Literal +from ..types import MatrixLike +from ..types import Self +from ..types import Union + class SecondOrderPolynomial(BaseEstimator, RegressorMixin): """Second order polynomial emulator. Creates a second order polynomial emulator. This is a linear model including all main effects, interactions and quadratic terms. + + Parameters + ---------- + degree : int, default=2 + The degree of the polynomial. """ - def __init__(self, degree=2): + def __init__(self, degree: int = 2): """Initializes a SecondOrderPolynomial object.""" self.degree = degree - def fit(self, X, y): + def fit(self, X: MatrixLike, y: ArrayLike) -> Self: """Fits the emulator to the data. Parameters @@ -49,7 +61,7 @@ def fit(self, X, y): self.is_fitted_ = True return self - def predict(self, X): + def predict(self, X: MatrixLike) -> ArrayLike: """Predicts the output of the emulator for a given input. Parameters @@ -67,7 +79,9 @@ def predict(self, X): predictions = self.model_.predict(X) return predictions - def get_grid_params(self, search_type="random"): + def get_grid_params( + self, search_type: Literal["random", "bayes"] = "random" + ) -> Union[dict, List[tuple[dict[str, Categorical], int]]]: """Get the parameter grid for the model. Parameters @@ -86,7 +100,16 @@ def get_grid_params(self, search_type="random"): elif search_type == "bayes": param_space = [({"degree": Categorical([2])}, 1)] + # TODO: Should this raise an error if the search type is not recognised? + return param_space def _more_tags(self): + """Returns more tags for the estimator. + + Returns + ------- + dict + The tags for the estimator. + """ return {"multioutput": True} diff --git a/autoemulate/emulators/random_forest.py b/autoemulate/emulators/random_forest.py index ea372473..625c9c5f 100644 --- a/autoemulate/emulators/random_forest.py +++ b/autoemulate/emulators/random_forest.py @@ -9,25 +9,74 @@ from skopt.space import Integer from skopt.space import Real +from ..types import Any +from ..types import ArrayLike +from ..types import Literal +from ..types import MatrixLike +from ..types import Optional +from ..types import Self +from ..types import Union + class RandomForest(BaseEstimator, RegressorMixin): """Random forest Emulator. Implements Random Forests regression from scikit-learn. + + Parameters + ---------- + n_estimators : int, default=100 + The number of trees in the forest. + criterion : {"squared_error", "mse", "mae"}, default="squared_error" + The function to measure the quality of a split. + max_depth : int, default=None + The maximum depth of the tree. If None, then nodes are expanded until + all leaves are pure or until all leaves contain less than + min_samples_split samples. + min_samples_split : int, default=2 + The minimum number of samples required to split an internal node. + min_samples_leaf : int, default=1 + The minimum number of samples required to be at a leaf node. + max_features : {"auto", "sqrt", "log2"}, int or float, default="auto" + The number of features to consider when looking for the best split: + - If int, then consider max_features features at each split. + - If float, then max_features is a fraction and + int(max_features * n_features) features are considered at each split. + - If "auto", then max_features=sqrt(n_features). + - If "sqrt", then max_features=sqrt(n_features). + - If "log2", then max_features=log2(n_features). + - If None, then max_features=n_features. + bootstrap : bool, default=True + Whether bootstrap samples are used when building trees. + oob_score : bool, default=False + Whether to use out-of-bag samples to estimate the R^2 on unseen data. + max_samples : int or float, default=None + If bootstrap is True, the number of samples to draw from X to train + each base estimator. If None (default), then draw X.shape[0] samples. + If int, then draw max_samples samples. + If float, then draw max_samples * X.shape[0] samples. Thus, max_samples + should be in the interval (0, 1). + random_state : int, RandomState instance or None, default=None + Controls both the randomness of the bootstrapping of the samples used + when building trees (if bootstrap=True) and the sampling of the features + to consider when looking for the best split at each node (if + max_features < n_features). """ def __init__( self, - n_estimators=100, - criterion="squared_error", - max_depth=None, - min_samples_split=2, - min_samples_leaf=1, - max_features=1.0, - bootstrap=True, - oob_score=False, - max_samples=None, - random_state=None, + n_estimators: int = 100, + criterion: Literal["squared_error", "mse", "mae"] = "squared_error", + max_depth: Optional[int] = None, + min_samples_split: int = 2, + min_samples_leaf: int = 1, + max_features: Optional[ + Union[Literal["auto", "sqrt", "log2"], int, float] + ] = 1.0, + bootstrap: bool = True, + oob_score: bool = False, + max_samples: Union[int, float] = None, + random_state=None, # TODO: set correct type ): """Initializes a RandomForest object.""" self.n_estimators = n_estimators @@ -41,7 +90,7 @@ def __init__( self.max_samples = max_samples self.random_state = random_state - def fit(self, X, y): + def fit(self, X: MatrixLike, y: ArrayLike) -> Self: """Fits the emulator to the data. Parameters @@ -72,7 +121,7 @@ def fit(self, X, y): self.is_fitted_ = True return self - def predict(self, X): + def predict(self, X: MatrixLike) -> ArrayLike: """Predicts the output of the simulator for a given input. Parameters @@ -80,10 +129,6 @@ def predict(self, X): X : {array-like, sparse matrix}, shape (n_samples, n_features) The training input samples. - return_std : bool - If True, returns a touple with two ndarrays, - one with the mean and one with the standard deviations of the prediction. - Returns ------- y : ndarray, shape (n_samples,) @@ -93,7 +138,9 @@ def predict(self, X): check_is_fitted(self, "is_fitted_") return self.model_.predict(X) - def get_grid_params(self, search_type="random"): + def get_grid_params( + self, search_type: Literal["random", "bayes"] = "random" + ) -> dict[str, Any]: """Returns the grid parameters of the emulator.""" param_space_random = { @@ -123,12 +170,21 @@ def get_grid_params(self, search_type="random"): elif search_type == "bayes": param_space = param_space_bayes + # TODO: Should this raise an error if the search type is not recognised? + return param_space def _more_tags(self): + """Returns more tags for the estimator. + + Returns + ------- + dict + The tags for the estimator. + """ return {"multioutput": True} - # def score(self, X, y, metric): + # def score(self, X: ArrayLike, y: ArrayLike, metric:Literal["rsme", "r2"]) -> float: # """Returns the score of the emulator. # Parameters @@ -149,5 +205,11 @@ def _more_tags(self): # return metric(y, predictions) # def _more_tags(self): + # """Returns more tags for the estimator. + # Returns + # ------- + # dict + # The tags for the estimator. + # """ # return {'non_deterministic': True, # 'multioutput': True} diff --git a/autoemulate/emulators/rbf.py b/autoemulate/emulators/rbf.py index 5fcaf73c..d36ad2da 100644 --- a/autoemulate/emulators/rbf.py +++ b/autoemulate/emulators/rbf.py @@ -11,19 +11,45 @@ from skopt.space import Integer from skopt.space import Real +from ..types import Any +from ..types import ArrayLike +from ..types import List +from ..types import Literal +from ..types import MatrixLike +from ..types import Self + class RBF(BaseEstimator, RegressorMixin): """Radial basis function Emulator. Wraps the RBF interpolator from scipy. + + Parameters + ---------- + smoothing : float, default=0.0 + The smoothing factor for the RBF interpolator. + kernel : {'linear', 'thin_plate_spline', 'cubic', 'quintic', 'multiquadric', 'inverse_multiquadric', 'gaussian'}, default='thin_plate_spline' + The kernel to be used in the RBF interpolator. + epsilon : float, default=1.0 + The epsilon parameter for the RBF interpolator. + degree : int, default=1 + The degree of the polynomial used in the RBF interpolator. """ def __init__( self, - smoothing=0.0, - kernel="thin_plate_spline", - epsilon=1.0, - degree=1, + smoothing: float = 0.0, + kernel: Literal[ + "linear", + "thin_plate_spline", + "cubic", + "quintic", + "multiquadric", + "inverse_multiquadric", + "gaussian", + ] = "thin_plate_spline", + epsilon: float = 1.0, + degree: int = 1, ): """Initializes an RBF object.""" self.smoothing = smoothing @@ -31,7 +57,7 @@ def __init__( self.epsilon = epsilon self.degree = degree - def fit(self, X, y): + def fit(self, X: MatrixLike, y: ArrayLike) -> Self: """Fits the emulator to the data. Parameters @@ -66,7 +92,7 @@ def fit(self, X, y): self.is_fitted_ = True return self - def predict(self, X): + def predict(self, X: MatrixLike) -> ArrayLike: """Predicts the output of the emulator for a given input. Parameters @@ -83,7 +109,9 @@ def predict(self, X): check_is_fitted(self, "is_fitted_") return self.model_(X) - def get_grid_params(self, search_type="random"): + def get_grid_params( + self, search_type: Literal["random", "bayes"] = "random" + ) -> List[dict[str, Any]]: """Returns the grid parameters of the emulator.""" # param_space_random = { # #"smoothing": uniform(0.0, 1.0), @@ -142,7 +170,16 @@ def get_grid_params(self, search_type="random"): elif search_type == "bayes": param_space = param_space_bayes + # TODO: Should this raise an error if the search type is not recognised? + return param_space def _more_tags(self): + """Returns more tags for the estimator. + + Returns + ------- + dict + The tags for the estimator. + """ return {"multioutput": True} diff --git a/autoemulate/emulators/support_vector_machines.py b/autoemulate/emulators/support_vector_machines.py index bca2216d..020c535b 100644 --- a/autoemulate/emulators/support_vector_machines.py +++ b/autoemulate/emulators/support_vector_machines.py @@ -11,6 +11,13 @@ from skopt.space import Integer from skopt.space import Real +from ..types import Any +from ..types import ArrayLike +from ..types import Literal +from ..types import MatrixLike +from ..types import Optional +from ..types import Self +from ..types import Union from autoemulate.utils import denormalise_y from autoemulate.utils import normalise_y @@ -19,22 +26,49 @@ class SupportVectorMachines(BaseEstimator, RegressorMixin): """Support Vector Machines Emulator. Wraps Support Vector Regressor from scikit-learn. + + Parameters + ---------- + kernel : {'rbf', 'linear', 'poly', 'sigmoid'}, default='rbf' + Specifies the kernel type to be used in the algorithm. + degree : int, default=3 + Degree of the polynomial kernel function ('poly'). + gamma : {'scale', 'auto'}, default='scale' + Kernel coefficient for 'rbf', 'poly' and 'sigmoid'. + coef0 : float, default=0.0 + Independent term in kernel function. It is only significant in 'poly' and 'sigmoid'. + tol : float, default=1e-3 + Tolerance for stopping criterion. + C : float, default=1.0 + Regularization parameter. The strength of the regularization is inversely proportional to C. + epsilon : float, default=0.1 + Epsilon in the epsilon-SVR model. It specifies the epsilon-tube within which no penalty is associated in the training loss function with points predicted within a distance epsilon from the actual value. + shrinking : bool, default=True + Whether to use the shrinking heuristic. + cache_size : float, default=200 + Specify the size of the kernel cache (in MB). + verbose : bool, default=False + Enable verbose output. + max_iter : int, default=-1 + Hard limit on iterations within solver, or -1 for no limit. + normalise_y : bool, default=True + Whether to normalise the target values before fitting the model. """ def __init__( self, - kernel="rbf", - degree=3, - gamma="scale", - coef0=0.0, - tol=1e-3, + kernel: Literal["rbf", "linear", "poly", "sigmoid"] = "rbf", + degree: int = 3, + gamma: Literal["scale", "auto"] = "scale", + coef0: float = 0.0, + tol: float = 1e-3, C=1.0, - epsilon=0.1, - shrinking=True, - cache_size=200, - verbose=False, - max_iter=-1, - normalise_y=True, + epsilon: float = 0.1, + shrinking: bool = True, + cache_size: int = 200, + verbose: bool = False, + max_iter: int = -1, + normalise_y: bool = True, ): """Initializes a SupportVectorMachines object.""" self.kernel = kernel @@ -50,7 +84,7 @@ def __init__( self.max_iter = max_iter self.normalise_y = normalise_y - def fit(self, X, y): + def fit(self, X: MatrixLike, y: ArrayLike) -> Self: """Fits the emulator to the data. Parameters @@ -101,7 +135,7 @@ def fit(self, X, y): self.is_fitted_ = True return self - def predict(self, X): + def predict(self, X: MatrixLike) -> ArrayLike: """Predicts the output for a given input. Parameters @@ -123,8 +157,22 @@ def predict(self, X): return y_pred - def get_grid_params(self, search_type="random"): - """Returns the grid paramaters for the emulator.""" + def get_grid_params( + self, search_type: Literal["random", "bayes"] = "random" + ) -> dict[str, Any]: + """Returns the grid paramaters for the emulator. + + Parameters + ---------- + search_type : str, optional + The type of parameter search to perform. Can be either 'random' or 'bayes'. + Defaults to 'random'. + + Returns + ------- + dict + The parameter grid for the model. + """ param_space_random = { "kernel": ["rbf", "linear", "poly", "sigmoid"], "degree": randint(2, 6), @@ -158,7 +206,16 @@ def get_grid_params(self, search_type="random"): elif search_type == "bayes": param_space = param_space_bayes + # TODO: Should this raise an error if the search type is not recognised? + return param_space def _more_tags(self): + """Returns more tags for the estimator. + + Returns + ------- + dict + The tags for the estimator. + """ return {"multioutput": False} diff --git a/autoemulate/emulators/xgboost.py b/autoemulate/emulators/xgboost.py index a0f4dd5c..44a7f3da 100644 --- a/autoemulate/emulators/xgboost.py +++ b/autoemulate/emulators/xgboost.py @@ -12,36 +12,88 @@ from skopt.space import Real from xgboost import XGBRegressor +from ..types import Any +from ..types import ArrayLike +from ..types import Literal +from ..types import MatrixLike +from ..types import Optional +from ..types import Self +from ..types import Union + class XGBoost(BaseEstimator, RegressorMixin): """XGBoost Emulator. Wraps XGBoost regression from scikit-learn. + + Parameters + ---------- + booster : {'gbtree', 'gblinear', 'dart'}, default='gbtree' + Which booster to use. + verbosity : int, default=0 + The degree of verbosity. + n_estimators : int, default=100 + The number of trees in the forest. + max_depth : int, default=6 + The maximum depth of the tree. + max_leaves : int, default=0 + The maximum number of leaves for base learners. + learning_rate : float, default=0.3 + The learning rate. + gamma : float, default=0 + Minimum loss reduction required to make a further partition on a leaf node of the tree. + min_child_weight : float, default=1 + Minimum sum of instance weight (hessian) needed in a child. + max_delta_step : float, default=0 + Maximum delta step we allow each tree's weight estimation to be. + subsample : float, default=1 + Subsample ratio of the training instance. + colsample_bytree : float, default=1 + Subsample ratio of columns when constructing each tree. + colsample_bylevel : float, default=1 + Subsample ratio of columns for each level. + colsample_bynode : float, default=1 + Subsample ratio of columns for each split. + reg_alpha : float, default=0 + L1 regularization term on weights. + reg_lambda : float, default=1 + L2 regularization term on weights. + objective : str, default='reg:squarederror' + Specify the learning task and the corresponding learning objective or a custom objective function. + tree_method : {'auto', 'exact', 'approx', 'hist', 'gpu_hist'}, default='auto' + The tree construction algorithm used in XGBoost. + random_state : int, RandomState instance or None, default=None + Controls both the randomness of the bootstrapping of the samples used + when building trees (if bootstrap=True) and the sampling of the features + to consider when looking for the best split at each node (if + max_features < n_features). + n_jobs : int, default=None + Number of parallel threads used to run XGBoost. """ def __init__( self, # general parameters - booster="gbtree", - verbosity=0, + booster: Literal["gbtree", "gblinear", "dart"] = "gbtree", + verbosity: int = 0, # tree booster parameters - n_estimators=100, - max_depth=6, - max_leaves=0, # no limit - learning_rate=0.3, - gamma=0, - min_child_weight=1, - max_delta_step=0, - subsample=1, - colsample_bytree=1, - colsample_bylevel=1, - colsample_bynode=1, - reg_alpha=0, - reg_lambda=1, - objective="reg:squarederror", - tree_method="auto", - random_state=None, - n_jobs=None, + n_estimators: int = 100, + max_depth: int = 6, + max_leaves: int = 0, # no limit + learning_rate: float = 0.3, + gamma: int = 0, + min_child_weight: int = 1, + max_delta_step: int = 0, + subsample: int = 1, + colsample_bytree: int = 1, + colsample_bylevel: int = 1, + colsample_bynode: int = 1, + reg_alpha: int = 0, + reg_lambda: int = 1, + objective: str = "reg:squarederror", + tree_method: Literal["auto", "exact", "approx", "hist", "gpu_hist"] = "auto", + random_state=None, # TODO: add correct type + n_jobs: Optional[int] = None, ): """Initializes a XGBoost object.""" self.booster = booster @@ -64,7 +116,7 @@ def __init__( self.random_state = random_state self.n_jobs = n_jobs - def fit(self, X, y): + def fit(self, X: MatrixLike, y: ArrayLike) -> Self: """Fits the emulator to the data. Parameters @@ -111,7 +163,7 @@ def fit(self, X, y): self.is_fitted_ = True return self - def predict(self, X): + def predict(self, X: MatrixLike) -> ArrayLike: """Predicts the output of the emulator for a given input. Parameters @@ -130,7 +182,9 @@ def predict(self, X): y_pred = self.model_.predict(X).astype(np.float64) return y_pred - def get_grid_params(self, search_type="random"): + def get_grid_params( + self, search_type: Literal["random", "bayes"] = "random" + ) -> dict[str, Any]: """Returns the grid parameters of the emulator.""" param_space_random = { "booster": ["gbtree", "dart"], @@ -169,7 +223,16 @@ def get_grid_params(self, search_type="random"): elif search_type == "bayes": param_space = param_space_bayes + # TODO: Should this raise an error if the search type is not recognised? + return param_space def _more_tags(self): + """Returns more tags for the estimator. + + Returns + ------- + dict + The tags for the estimator. + """ return {"multioutput": True} diff --git a/autoemulate/experimental_design.py b/autoemulate/experimental_design.py index 9ea0ecf3..66918aa5 100644 --- a/autoemulate/experimental_design.py +++ b/autoemulate/experimental_design.py @@ -2,10 +2,13 @@ from abc import abstractmethod import mogp_emulator +import numpy as np + +from .types import List class ExperimentalDesign(ABC): - def __init__(self, bounds_list): + def __init__(self, bounds_list: List[tuple[float, float]]): """Initializes a Sampler object. Parameters @@ -29,7 +32,7 @@ def sample(self, n: int): pass @abstractmethod - def get_n_parameters(self): + def get_n_parameters(self) -> int: """Returns the number of parameters in the sample space. Returns @@ -41,7 +44,7 @@ def get_n_parameters(self): class LatinHypercube(ExperimentalDesign): - def __init__(self, bounds_list): + def __init__(self, bounds_list: List[tuple[float, float]]): """Initializes a LatinHypercube object. Parameters @@ -53,7 +56,7 @@ def __init__(self, bounds_list): """ self.sampler = mogp_emulator.LatinHypercubeDesign(bounds_list) - def sample(self, n: int): + def sample(self, n: int) -> np.ndarray: """Samples n points from the sample space. Parameters @@ -68,7 +71,7 @@ def sample(self, n: int): """ return self.sampler.sample(n) - def get_n_parameters(self): + def get_n_parameters(self) -> int: """Returns the number of parameters in the sample space. Returns diff --git a/autoemulate/hyperparam_searching.py b/autoemulate/hyperparam_searching.py index b7507fc9..b0cc506e 100644 --- a/autoemulate/hyperparam_searching.py +++ b/autoemulate/hyperparam_searching.py @@ -3,23 +3,37 @@ from sklearn.model_selection import RandomizedSearchCV from skopt import BayesSearchCV -from autoemulate.utils import adjust_param_space -from autoemulate.utils import get_model_name -from autoemulate.utils import get_model_param_space -from autoemulate.utils import get_model_params +from .types import ArrayLike +from .types import MatrixLike +from .types import TYPE_CHECKING +from .utils import adjust_param_space +from .utils import get_model_name +from .utils import get_model_param_space +from .utils import get_model_params +if TYPE_CHECKING: + from logging import Logger + from .types import Iterable, Union, Optional, Literal + from sklearn.model_selection import BaseCrossValidator + from sklearn.model_selection import BaseShuffleSplit + from sklearn.pipeline import Pipeline + + SearchTypes = Literal["random", "bayes"] + + +# TODO: Note that BayesSearchCV takes a n_jobs Int parameter. We should enforce that here to avoid bugs. Also check that model and return type are correct. def optimize_params( - X, - y, - cv, - model, - search_type="random", - niter=20, - param_space=None, - n_jobs=None, - logger=None, -): + X: MatrixLike, + y: Union[MatrixLike, ArrayLike], + cv: Union[int, BaseCrossValidator, Iterable, BaseShuffleSplit], + model: Pipeline, # TODO: Verify that this is correct + search_type: SearchTypes = "random", + niter: int = 20, + param_space: Optional[dict] = None, + n_jobs: Optional[int] = None, + logger: Optional[Logger] = None, +) -> Pipeline: """Performs hyperparameter search for the provided model. Parameters @@ -35,16 +49,16 @@ def optimize_params( Type of search to perform. Can be "random" or "bayes", "grid" not yet implemented. niter : int, default=20 Number of parameter settings that are sampled. Trades off runtime vs quality of the solution. - param_space : dict, default=None + param_space : dict, default=None Dictionary with parameters names (string) as keys and lists of parameter settings to try as values, or a list of such dictionaries, in which case the grids spanned by each dictionary in the list are explored. This enables searching over any sequence of parameter settings. Parameters names should be prefixed with "model__" to indicate that they are parameters of the model. - n_jobs : int + n_jobs : int, optional Number of jobs to run in parallel. - logger : logging.Logger + logger : logging.Logger, optional Logger instance. Returns @@ -94,14 +108,19 @@ def optimize_params( return searcher.best_estimator_ -def process_param_space(model, search_type, param_space): - """Process parameter grid for hyperparameter search. +def process_param_space( + model: Pipeline, search_type: SearchTypes, param_space: dict +) -> dict: + """ + Process parameter grid for hyperparameter search. + Gets the parameter grid for the model and adjusts it to include prefixes for pipelines / multioutput estimators. Parameters ---------- - model : model instance to do hyperparameter search for. + model : Pipeline + Model instance to do hyperparameter search for. search_type : str, default="random" Type of search to perform. Can be "random" or "bayes", "grid" not yet implemented. param_space : dict, default=None @@ -126,7 +145,7 @@ def process_param_space(model, search_type, param_space): return param_space -def check_param_space(param_space, model): +def check_param_space(param_space: dict, model: Pipeline) -> dict: """Checks that the parameter grid is valid. Parameters @@ -145,12 +164,12 @@ def check_param_space(param_space, model): ------- param_space : dict """ - if type(param_space) != dict: + if not isinstance(param_space, dict): raise TypeError("param_space must be a dictionary") for key, value in param_space.items(): - if type(key) != str: + if not isinstance(key, str): raise TypeError("param_space keys must be strings") - if type(value) != list: + if not isinstance(value, list): raise TypeError("param_space values must be lists") inbuilt_grid = get_model_params(model) diff --git a/autoemulate/logging_config.py b/autoemulate/logging_config.py index d4712fdf..26758630 100644 --- a/autoemulate/logging_config.py +++ b/autoemulate/logging_config.py @@ -3,7 +3,7 @@ import warnings -def configure_logging(log_to_file=False): +def configure_logging(log_to_file: bool = False) -> logging.Logger: # Create a logger logger = logging.getLogger("autoemulate") logger.setLevel(logging.INFO) diff --git a/autoemulate/metrics.py b/autoemulate/metrics.py index c062019a..c6a295b1 100644 --- a/autoemulate/metrics.py +++ b/autoemulate/metrics.py @@ -2,8 +2,10 @@ from sklearn.metrics import mean_squared_error from sklearn.metrics import r2_score +from .types import ArrayLike -def rmse(y_true, y_pred): + +def rmse(y_true: ArrayLike, y_pred: ArrayLike) -> float: """Returns the root mean squared error. Parameters @@ -16,7 +18,7 @@ def rmse(y_true, y_pred): return mean_squared_error(y_true, y_pred, squared=False) -def r2(y_true, y_pred): +def r2(y_true: ArrayLike, y_pred: ArrayLike) -> float: """Returns the R^2 score. Parameters diff --git a/autoemulate/model_processing.py b/autoemulate/model_processing.py index f84813d8..1181668f 100644 --- a/autoemulate/model_processing.py +++ b/autoemulate/model_processing.py @@ -2,8 +2,11 @@ from sklearn.multioutput import MultiOutputRegressor from sklearn.pipeline import Pipeline +from .types import ArrayLike +from .types import Optional -def get_models(model_registry, model_subset=None): + +def get_models(model_registry: dict, model_subset: Optional[list] = None) -> list: """Get models from REGISTRY. Takes a subset of models if model_subset argument was used in setup(). @@ -27,7 +30,7 @@ def get_models(model_registry, model_subset=None): return models -def check_model_names(model_names, model_registry): +def check_model_names(model_names: list, model_registry: dict) -> None: """Check whether model_names are in MODEL_REGISTRY Parameters @@ -40,7 +43,11 @@ def check_model_names(model_names, model_registry): Returns ------- None - Raises ValueError if a model in chosen_models is not in MODEL_REGISTRY. + + Raises + ------ + ValueError + If a model in chosen_models is not in MODEL_REGISTRY. """ for model in model_names: if model not in model_registry: @@ -49,7 +56,7 @@ def check_model_names(model_names, model_registry): ) -def turn_models_into_multioutput(models, y): +def turn_models_into_multioutput(models: list, y: ArrayLike) -> list: """Turn single output models into multioutput models if y is 2D. Parameters @@ -65,15 +72,20 @@ def turn_models_into_multioutput(models, y): List of model instances, with single output models wrapped in MultiOutputRegressor. """ models_multi = [ - MultiOutputRegressor(model) - if not model._more_tags()["multioutput"] and (y.ndim > 1 and y.shape[1] > 1) - else model + ( + MultiOutputRegressor(model) + if not model._more_tags()["multioutput"] and (y.ndim > 1 and y.shape[1] > 1) + else model + ) for model in models ] return models_multi -def wrap_models_in_pipeline(models, scale, scaler, reduce_dim, dim_reducer): +# TODO: add types for scaler and dim_reducer +def wrap_models_in_pipeline( + models: list, scale: bool, scaler, reduce_dim: bool, dim_reducer +): """Wrap models in a pipeline if scale is True. Parameters @@ -110,8 +122,15 @@ def wrap_models_in_pipeline(models, scale, scaler, reduce_dim, dim_reducer): return models_piped +# TODO: add types for scaler and dim_reducer def get_and_process_models( - model_registry, model_subset, y, scale, scaler, reduce_dim, dim_reducer + model_registry: dict, + model_subset: list, + y: ArrayLike, + scale: bool, + scaler, + reduce_dim: bool, + dim_reducer, ): """Get and process models. @@ -127,6 +146,10 @@ def get_and_process_models( Whether to scale the data. scaler : sklearn.preprocessing object Scaler to use. + reduce_dim : bool + Whether to reduce the dimensionality of the data. + dim_reducer : sklearn.decomposition object + Dimensionality reduction method to use. Returns ------- diff --git a/autoemulate/plotting.py b/autoemulate/plotting.py index dd2aec8b..7973675d 100644 --- a/autoemulate/plotting.py +++ b/autoemulate/plotting.py @@ -2,10 +2,16 @@ import numpy as np from sklearn.metrics import PredictionErrorDisplay +from .types import TYPE_CHECKING from autoemulate.utils import get_model_name +if TYPE_CHECKING: + from .types import ArrayLike, Literal, Optional -def validate_inputs(cv_results, model_name): + PlotTypes = Literal["actual_vs_predicted", "residual_vs_predicted"] + + +def validate_inputs(cv_results: dict, model_name: str) -> None: """Validates cv_results and model_name for plotting. Parameters @@ -14,6 +20,16 @@ def validate_inputs(cv_results, model_name): A list of cross-validation results for each model. model_name : str The name of a model to plot. + + Returns + ------- + None + + Raises + ------ + ValueError + If cv_results is empty. + If model_name is not in cv_results. """ if not cv_results: raise ValueError("Run .compare() first.") @@ -25,8 +41,25 @@ def validate_inputs(cv_results, model_name): ) -def check_multioutput(y, output_index): - """Checks if y is multi-output and if the output_index is valid.""" +def check_multioutput(y: ArrayLike, output_index: int) -> None: + """Checks if y is multi-output and if the output_index is valid. + + Parameters + ---------- + y : array-like, shape (n_samples, n_outputs) + Simulation output. + output_index : int + The index of the output to plot. + + Returns + ------- + None + + Raises + ------ + ValueError + If output_index is out of range. + """ if y.ndim > 1: if (output_index > y.shape[1] - 1) | (output_index < 0): raise ValueError( @@ -38,17 +71,18 @@ def check_multioutput(y, output_index): ) +# TODO: Should X be MatrixLike? def plot_single_fold( - cv_results, - X, - y, - model_name, - fold_index, - ax, - plot_type="actual_vs_predicted", - annotation=" ", - output_index=0, -): + cv_results: dict, + X: ArrayLike, + y: ArrayLike, + model_name: str, + fold_index: int, + ax: plt.Axes, + plot_type: PlotTypes = "actual_vs_predicted", + annotation: str = " ", + output_index: int = 0, +) -> None: """Plots a single cv fold for a given model. Parameters @@ -74,6 +108,10 @@ def plot_single_fold( The annotation to add to the plot title. Default is an empty string. output_index : int, optional The index of the output to plot. Default is 0. + + Returns + ------- + None """ test_indices = cv_results[model_name]["indices"]["test"][fold_index] @@ -95,15 +133,16 @@ def plot_single_fold( ax.set_title(f"{model_name} - {title_suffix}") +# TODO: Should X be MatrixLike? def plot_best_fold_per_model( - cv_results, - X, - y, - n_cols=3, - plot_type="actual_vs_predicted", - figsize=None, - output_index=0, -): + cv_results: dict, + X: ArrayLike, + y: ArrayLike, + n_cols: int = 3, + plot_type: PlotTypes = "actual_vs_predicted", + figsize: Optional[tuple[int, int]] = None, + output_index: int = 0, +) -> None: """Plots results of the best (highest R^2) cv-fold for each model in cv_results. Parameters @@ -123,6 +162,10 @@ def plot_best_fold_per_model( Width, height in inches. Overrides the default figure size. output_index : int, optional The index of the output to plot. Default is 0. + + Returns + ------- + None """ n_models = len(cv_results) @@ -151,19 +194,19 @@ def plot_best_fold_per_model( plt.show() +# TODO: Should X be MatrixLike? def plot_model_folds( - cv_results, - X, - y, - model_name, - n_cols=3, - plot_type="actual_vs_predicted", - figsize=None, - output_index=0, -): + cv_results: dict, + X: ArrayLike, + y: ArrayLike, + model_name: str, + n_cols: int = 3, + plot_type: PlotTypes = "actual_vs_predicted", + figsize: Optional[tuple[int, int]] = None, + output_index: int = 0, +) -> None: """Plots all the folds for a given model. - Parameters ---------- cv_results : dict @@ -183,6 +226,10 @@ def plot_model_folds( Overrides the default figure size. output_index : int, optional The index of the output to plot. Default is 0. + + Returns + ------- + None """ n_folds = len(cv_results[model_name]["estimator"]) @@ -210,16 +257,17 @@ def plot_model_folds( plt.show() +# TODO: Should X be MatrixLike? def _plot_results( - cv_results, - X, - y, - model_name=None, - n_cols=3, - plot_type="actual_vs_predicted", - figsize=None, - output_index=0, -): + cv_results: dict, + X: ArrayLike, + y: ArrayLike, + model_name: Optional[str] = None, + n_cols: int = 3, + plot_type: PlotTypes = "actual_vs_predicted", + figsize: Optional[tuple[int, int]] = None, + output_index: int = 0, +) -> None: """Plots the results of cross-validation. Parameters @@ -264,7 +312,15 @@ def _plot_results( ) -def _plot_model(model, X, y, plot="standard", n_cols=2, figsize=None): +# TODO: add model type, is Pipeline correct? +def _plot_model( + model, + X: ArrayLike, + y: ArrayLike, + plot: Literal["standard", "residual"] = "standard", + n_cols: int = 2, + figsize: Optional[tuple[int, int]] = None, +): """Plots the model predictions vs. the true values. Parameters @@ -323,7 +379,10 @@ def _plot_model(model, X, y, plot="standard", n_cols=2, figsize=None): axs ): # Check to avoid index error if n_cols * n_rows > n_outputs display = PredictionErrorDisplay.from_predictions( - y_true=y[:, i], y_pred=y_pred[:, i], kind=plot_type, ax=axs[i] + y_true=y[:, i], + y_pred=y_pred[:, i], + kind=plot_type, + ax=axs[i], ) axs[i].set_title(f"{get_model_name(model)} - Test Set - Output {i+1}") diff --git a/autoemulate/printing.py b/autoemulate/printing.py index 2c65b9ee..27953877 100644 --- a/autoemulate/printing.py +++ b/autoemulate/printing.py @@ -1,9 +1,19 @@ +from .types import TYPE_CHECKING from autoemulate.utils import get_mean_scores from autoemulate.utils import get_model_name from autoemulate.utils import get_model_scores +if TYPE_CHECKING: + import pandas as pd -def _print_cv_results(models, scores_df, model=None, sort_by="r2"): + +# TODO: Suggestion to change model to model_name, to avoid confusion +def _print_cv_results( + models: list, + scores_df: pd.DataFrame, + model: str = None, + sort_by: str = "r2", +) -> None: """Print cv results. Parameters @@ -18,6 +28,14 @@ def _print_cv_results(models, scores_df, model=None, sort_by="r2"): sort_by : str, optional The metric to sort by. Default is "r2". + Returns + ------- + None + + Raises + ------ + ValueError + If model is not in self.models. """ # check if model is in self.models if model is not None: diff --git a/autoemulate/save.py b/autoemulate/save.py index b905ef85..1618740f 100644 --- a/autoemulate/save.py +++ b/autoemulate/save.py @@ -9,8 +9,21 @@ class ModelSerialiser: - def save_model(self, model, path): - """Saves a model + metadata to disk.""" + # TODO: Add model type, is Pipeline correct? + def save_model(self, model, path: str) -> None: + """Saves a model + metadata to disk. + + Parameters + ---------- + model : object + Model to save. + path : str + Path to save the model to. + + Returns + ------- + None + """ # model joblib.dump(model, path) @@ -24,8 +37,20 @@ def save_model(self, model, path): with open(self.get_meta_path(path), "w") as f: json.dump(meta, f) - def load_model(self, path): - """Loads a model from disk and checks version.""" + # TODO: Add return type, is Pipeline correct? + def load_model(self, path: str): + """Loads a model from disk and checks version. + + Parameters + ---------- + path : str + Path to the model file. + + Returns + ------- + + Model instance. + """ model = joblib.load(path) meta_path = self.get_meta_path(path) @@ -49,11 +74,21 @@ def load_model(self, path): return model - def get_meta_path(self, path): + def get_meta_path(self, path: str) -> str: """Returns the path to the metadata file. If the path has an extension, it is replaced with _meta.json. Otherwise, _meta.json is appended to the path. + + Parameters + ---------- + path : str + Path to the model file. + + Returns + ------- + str + Path to the metadata file. """ base, ext = os.path.splitext(path) meta_path = f"{base}_meta.json" if ext else f"{base}_meta{ext}.json" diff --git a/autoemulate/simulations/projectile.py b/autoemulate/simulations/projectile.py index ac21bb6b..edaaa0f7 100644 --- a/autoemulate/simulations/projectile.py +++ b/autoemulate/simulations/projectile.py @@ -3,6 +3,12 @@ import scipy from scipy.integrate import solve_ivp +from ..types import TYPE_CHECKING + +if TYPE_CHECKING: + from ..types import ArrayLike + from scipy.integrate._ivp.ivp import OdeResult + # Create our simulator, which solves a nonlinear differential equation describing projectile # motion with drag. A projectile is launched from an initial height of 2 meters at an # angle of 45 degrees and falls under the influence of gravity and air resistance. @@ -12,8 +18,23 @@ # define functions needed for simulator -def f(t, y, c): - "Compute RHS of system of differential equations, returning vector derivative" +def f(t: float, y: ArrayLike, c: float) -> ArrayLike: + """Compute RHS of system of differential equations, returning vector derivative. + + Parameters + ---------- + t : float + Time. + y : array-like, shape (4,) + State vector. + c : float + Drag coefficient. + + Returns + ------- + dydt : array-like, shape (4,) + Vector derivative. + """ # check inputs and extract @@ -35,8 +56,23 @@ def f(t, y, c): return dydt -def event(t, y, c): - "event to trigger end of integration" +def event(t: float, y: ArrayLike, c: float) -> float: + """Event to trigger end of integration + + Parameters + ---------- + t : float + Time. + y : array-like, shape (4,) + State vector. + c : float + Drag coefficient. + + Returns + ------- + float + Distance travelled by the projectile. + """ assert len(y) == 4 assert c >= 0.0 @@ -49,8 +85,19 @@ def event(t, y, c): # now can define simulator -def simulator_base(x): - "simulator to solve ODE system for projectile motion with drag. returns distance projectile travels" +def simulator_base(x: ArrayLike) -> OdeResult: + """Simulator to solve ODE system for projectile motion with drag. returns distance projectile travels + + Parameters + ---------- + x : array-like, shape (2,) + Drag coefficient and launch velocity. + + Returns + ------- + results : scipy.integrate._ivp.ivp.OdeResult + Results of the ODE solver. + """ # unpack values @@ -75,16 +122,38 @@ def simulator_base(x): return results -def simulator(x): - "simulator to solve ODE system for projectile motion with drag. returns distance projectile travels" +def simulator(x: ArrayLike) -> float: + """Simulator to solve ODE system for projectile motion with drag. returns distance projectile travels. + + Parameters + ---------- + x : array-like, shape (2,) + Drag coefficient and launch velocity. + + Returns + ------- + float + Distance travelled by the projectile. + """ results = simulator_base(x) return results.y_events[0][0][2] -def simulator_multioutput(x): - "simulator to solve ODE system with multiple outputs" +def simulator_multioutput(x: ArrayLike) -> tuple[float, float]: + """Simulator to solve ODE system with multiple outputs + + Parameters + ---------- + x : array-like, shape (2,) + Drag coefficient and launch velocity. + + Returns + ------- + tuple + Distance travelled by the projectile and the speed of the projectile. + """ results = simulator_base(x) @@ -97,8 +166,22 @@ def simulator_multioutput(x): # functions for printing out results -def print_results(inputs, arg, var): - "convenience function for printing out generic results" +def print_results(inputs, arg, var) -> None: + """Convenience function for printing out generic results + + Parameters + ---------- + inputs : array-like + Input values. + arg : array-like + Mean values. + var : array-like + Variance values. + + Returns + ------- + None + """ print( "---------------------------------------------------------------------------------" @@ -108,15 +191,43 @@ def print_results(inputs, arg, var): print("{} {} {}".format(pp, m, v)) -def print_predictions(inputs, pred, var): - "convenience function for printing predictions" +def print_predictions(inputs, pred, var) -> None: + """Convenience function for printing predictions + + Parameters + ---------- + inputs : array-like + Input values. + pred : array-like + Predicted mean values. + var : array-like + Predictive variance values. + + Returns + ------- + None + """ print("Target Point Predicted Mean Predictive Variance") print_results(inputs, pred, var) -def print_errors(inputs, errors, var): - "convenience function for printing out results and computing mean square error" +def print_errors(inputs, errors, var) -> None: + """Convenience function for printing out results and computing mean square error + + Parameters + ---------- + inputs : array-like + Input values. + errors : array-like + Error values. + var : array-like + Variance values. + + Returns + ------- + None + """ print("Target Point Standard Error Predictive Variance") print_results(inputs, errors, var) diff --git a/autoemulate/types.py b/autoemulate/types.py new file mode 100644 index 00000000..afe9e8b6 --- /dev/null +++ b/autoemulate/types.py @@ -0,0 +1,15 @@ +from collections.abc import Iterable +from typing import Any +from typing import Iterable +from typing import List +from typing import Literal +from typing import Optional +from typing import Self +from typing import TYPE_CHECKING +from typing import Union + +import numpy as np +from numpy.typing import ArrayLike + +Matrixlike = Union[np.ndarray, np.matrix, Iterable[Iterable[float]]] +ArrayLike diff --git a/autoemulate/utils.py b/autoemulate/utils.py index 190c1454..d5f6392b 100644 --- a/autoemulate/utils.py +++ b/autoemulate/utils.py @@ -10,6 +10,12 @@ from sklearn.multioutput import MultiOutputRegressor from sklearn.pipeline import Pipeline +from .types import TYPE_CHECKING + +if TYPE_CHECKING: + from .types import Literal, Union, ArrayLike, List + import pandas as pd + @contextmanager def suppress_convergence_warnings(): @@ -34,7 +40,8 @@ def suppress_convergence_warnings(): del os.environ["PYTHONWARNINGS"] -def get_model_name(model): +# TODO: Add model type, what is another model type than Pipeline? +def get_model_name(model) -> str: """Get the name of the base model. This function handles standalone models, models wrapped in a MultiOutputRegressor, @@ -70,6 +77,7 @@ def get_model_name(model): return type(model).__name__ +# TODO: add model type def get_model_params(model): """Get the parameters of the base model, which are not prefixed with `model__` or `estimator__`. @@ -104,7 +112,10 @@ def get_model_params(model): return model.get_params() -def get_model_param_space(model, search_type="random"): +# TODO: add model type +def get_model_param_space( + model, search_type: Literal["random", "bayes"] = "random" +) -> dict: """Get the parameter grid of the base model. This is used for hyperparameter search. This function handles standalone models, models wrapped in a MultiOutputRegressor, @@ -140,7 +151,8 @@ def get_model_param_space(model, search_type="random"): return model.get_grid_params(search_type) -def adjust_param_space(model, param_space): +# TODO: add model type +def adjust_param_space(model, param_space: dict) -> dict: """Adjusts param grid to be compatible with the model. Adds `model__` if model is a pipeline and `estimator__` if model is a MultiOutputRegressor. Or `model__estimator__` if both, @@ -176,7 +188,9 @@ def adjust_param_space(model, param_space): return add_prefix_to_param_space(param_space, prefix) -def add_prefix_to_param_space(param_space, prefix): +def add_prefix_to_param_space( + param_space: Union[dict, List[dict]], prefix: str +) -> Union[dict, List[dict]]: """Adds a prefix to all keys in a parameter grid. Works for three types of param_spaces: @@ -217,12 +231,25 @@ def add_prefix_to_param_space(param_space, prefix): return add_prefix_to_single_grid(param_space, prefix) -def add_prefix_to_single_grid(grid, prefix): - """Adds a prefix to all keys in a single parameter grid dictionary.""" +def add_prefix_to_single_grid(grid: dict, prefix: str) -> dict: + """Adds a prefix to all keys in a single parameter grid dictionary. + + Parameters + ---------- + grid : dict + The parameter grid to which the prefix will be added. + prefix : str + The prefix to be added to each parameter name in the grid. + + Returns + ------- + dict + The parameter grid with the prefix added to each key. + """ return {prefix + key: value for key, value in grid.items()} -def normalise_y(y): +def normalise_y(y: ArrayLike) -> tuple[ArrayLike, ArrayLike, ArrayLike]: """Normalize the target values y. Parameters @@ -238,14 +265,13 @@ def normalise_y(y): The mean of the target values. y_std : array-like, shape (n_outputs,) The standard deviation of the target values. - """ y_mean = np.mean(y, axis=0) y_std = np.std(y, axis=0) return (y - y_mean) / y_std, y_mean, y_std -def denormalise_y(y_pred, y_mean, y_std): +def denormalise_y(y_pred: ArrayLike, y_mean: ArrayLike, y_std: ArrayLike) -> ArrayLike: """Denormalize the predicted values. Parameters @@ -261,12 +287,13 @@ def denormalise_y(y_pred, y_mean, y_std): ------- y_pred_denorm : array-like, shape (n_samples,) or (n_samples, n_outputs) The denormalized predicted target values. - """ return y_pred * y_std + y_mean -def get_mean_scores(scores_df, metric): +def get_mean_scores( + scores_df: pd.DataFrame, metric: Literal["r2", "rmse"] +) -> pd.DataFrame: """Get the mean scores for each model and metric. Parameters @@ -308,7 +335,21 @@ def get_mean_scores(scores_df, metric): return means_df -def get_model_scores(scores_df, model_name): +def get_model_scores(scores_df: pd.DataFrame, model_name: str) -> pd.DataFrame: + """Get the scores for a specific model. + + Parameters + ---------- + scores_df : pandas.DataFrame + DataFrame with columns "model", "metric", "fold", "score". + model_name : str + The name of the model. + + Returns + ------- + model_scores : pandas.DataFrame + DataFrame with columns "fold", "metric", "score". + """ model_scores = scores_df[scores_df["model"] == model_name].pivot( index="fold", columns="metric", values="score" ) @@ -316,11 +357,19 @@ def get_model_scores(scores_df, model_name): return model_scores -def set_random_seed(seed: int, deterministic: bool = False): +def set_random_seed(seed: int, deterministic: bool = False) -> None: """Set random seed for Python, Numpy and PyTorch. - Args: - seed: int, the random seed to use. - deterministic: bool, use "deterministic" algorithms in PyTorch. + + Parameters + ---------- + seed : int + The random seed to use. + deterministic : bool + Use "deterministic" algorithms in PyTorch. + + Returns + ------- + None """ random.seed(seed) np.random.seed(seed)