autoemulate 0.3.0__tar.gz → 0.3.2__tar.gz

{autoemulate-0.3.0 → autoemulate-0.3.2}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.3
 Name: autoemulate
-Version: 0.3.0
+Version: 0.3.2
 Summary: A python package for semi-automated emulation
 License: MIT
 Author: AutoEmulate contributors (see our GitHub page)

{autoemulate-0.3.0 → autoemulate-0.3.2}/autoemulate/compare.py

@@ -24,7 +24,8 @@ from autoemulate.preprocess_target import get_dim_reducer
 from autoemulate.preprocess_target import NonTrainableTransformer
 from autoemulate.printing import _print_setup
 from autoemulate.save import ModelSerialiser
-from autoemulate.sensitivity_analysis import _plot_sensitivity_analysis
+from autoemulate.sensitivity_analysis import _plot_morris_analysis
+from autoemulate.sensitivity_analysis import _plot_sobol_analysis
 from autoemulate.sensitivity_analysis import _sensitivity_analysis
 from autoemulate.utils import _check_cv
 from autoemulate.utils import _ensure_2d
@@ -473,15 +474,15 @@ class AutoEmulate:
                     means = _sum_cvs({full_name: cv_results[full_name]}, metric)
                     current_score = means.iloc[0][f"{metric}"]
 
-                    # Compare with current best
-                    if (metric == "r2" and current_score > best_score) or (
-                        metric != "r2" and current_score < best_score
-                    ):
-                        best_score = current_score
-                        for model in prep_data["models"]:
-                            if get_model_name(model) == full_name:
+                    # Find the model with the correct name that has the best score
+                    for model in prep_data["models"]:
+                        if get_model_name(model) == full_name:
+                            if (metric == "r2" and current_score > best_score) or (
+                                metric != "r2" and current_score < best_score
+                            ):
+                                best_score = current_score
                                 best_model = model
-                                break
+                            break
 
                 if best_model is None:
                     raise ValueError(
@@ -979,7 +980,14 @@ class AutoEmulate:
         return fig
 
     def sensitivity_analysis(
-        self, model=None, problem=None, N=1024, conf_level=0.95, as_df=True
+        self,
+        model=None,
+        method="sobol",
+        problem=None,
+        N=1024,
+        conf_level=0.95,
+        as_df=True,
+        **plot_kwargs,
     ):
         """Perform Sobol sensitivity analysis on a fitted emulator.
 
@@ -1038,6 +1046,10 @@ class AutoEmulate:
         The analysis requires N * (2D + 2) model evaluations, where D is the number of input
         parameters. For example, with N=1024 and 5 parameters, this requires 12,288 evaluations.
         """
+        self.method = method
+        if method not in ["sobol", "morris"]:
+            raise ValueError(f"Unknown method: {method}. Must be 'sobol' or 'morris'.")
+
         if model is None:
             if not hasattr(self, "best_model"):
                 raise RuntimeError("Must run compare() before sensitivity_analysis()")
@@ -1045,10 +1057,22 @@ class AutoEmulate:
             self.logger.info(
                 f"No model provided, using {get_model_name(model)}, which had the highest average cross-validation score, refitted on full data."
             )
-        Si = _sensitivity_analysis(model, problem, self.X, N, conf_level, as_df)
-        return Si
 
-    def plot_sensitivity_analysis(self, results, index="S1", n_cols=None, figsize=None):
+        df_results = _sensitivity_analysis(
+            model=model,
+            method=method,
+            problem=problem,
+            X=self.X,
+            N=N,
+            conf_level=conf_level,
+            as_df=as_df,
+        )
+
+        return df_results
+
+    def plot_sensitivity_analysis(
+        self, results, index="S1", param_groups=None, n_cols=None, figsize=None
+    ):
         """
         Plot the sensitivity analysis results.
 
@@ -1068,4 +1092,7 @@ class AutoEmulate:
             Figure size as (width, height) in inches.If None, automatically calculated.
 
         """
-        return _plot_sensitivity_analysis(results, index, n_cols, figsize)
+        if self.method == "sobol":
+            return _plot_sobol_analysis(results, index, n_cols, figsize)
+        elif self.method == "morris":
+            return _plot_morris_analysis(results, param_groups, n_cols, figsize)

{autoemulate-0.3.0 → autoemulate-0.3.2}/autoemulate/experimental/compare.py

@@ -5,23 +5,27 @@ from typing import Any
 import numpy as np
 from sklearn.model_selection import BaseCrossValidator, KFold
 
-from autoemulate.experimental.data.utils import InputTypeMixin
+from autoemulate.experimental.data.utils import ConversionMixin
+from autoemulate.experimental.device import TorchDeviceMixin
 from autoemulate.experimental.emulators import ALL_EMULATORS
 from autoemulate.experimental.emulators.base import Emulator
 from autoemulate.experimental.model_selection import cross_validate
 from autoemulate.experimental.tuner import Tuner
-from autoemulate.experimental.types import InputLike
+from autoemulate.experimental.types import DeviceLike, InputLike
 
 
-class AutoEmulate(InputTypeMixin):
+class AutoEmulate(ConversionMixin, TorchDeviceMixin):
     def __init__(
         self,
         x: InputLike,
         y: InputLike,
         models: list[type[Emulator]] | None = None,
+        device: DeviceLike | None = None,
     ):
+        TorchDeviceMixin.__init__(self, device=device)
         # TODO: refactor in https://github.com/alan-turing-institute/autoemulate/issues/400
         x, y = self._convert_to_tensors(x, y)
+        x, y = self._move_tensors_to_device(x, y)
 
         # Set default models if None
         updated_models = self.get_models(models)
@@ -73,7 +77,7 @@ class AutoEmulate(InputTypeMixin):
     def compare(
         self, n_iter: int = 10, cv: type[BaseCrossValidator] = KFold
     ) -> dict[str, dict[str, Any]]:
-        tuner = Tuner(self.train_val, y=None, n_iter=n_iter)
+        tuner = Tuner(self.train_val, y=None, n_iter=n_iter, device=self.device)
         models_evaluated = {}
         for model_cls in self.models:
             scores, configs = tuner.run(model_cls)

{autoemulate-0.3.0 → autoemulate-0.3.2}/autoemulate/experimental/data/preprocessors.py

@@ -1,7 +1,7 @@
 from abc import ABC, abstractmethod
 
 import torch
-from autoemulate.experimental.types import InputLike
+from autoemulate.experimental.types import TensorLike
 
 
 class Preprocessor(ABC):
@@ -9,14 +9,14 @@ class Preprocessor(ABC):
     def __init__(*args, **kwargs): ...
 
     @abstractmethod
-    def preprocess(self, x: InputLike) -> InputLike: ...
+    def preprocess(self, x: TensorLike) -> TensorLike: ...
 
 
 class Standardizer(Preprocessor):
-    mean: torch.Tensor
-    std: torch.Tensor
+    mean: TensorLike
+    std: TensorLike
 
-    def __init__(self, mean: torch.Tensor, std: torch.Tensor):
+    def __init__(self, mean: TensorLike, std: TensorLike):
         if len(mean.shape) != 2:
             raise ValueError(f"mean is expected to be 2D, shape passed ({mean.shape})")
         if len(std.shape) != 2:
@@ -28,19 +28,17 @@ class Standardizer(Preprocessor):
         self.mean = mean
         self.std = std
 
-    def preprocess(self, x):
+    def preprocess(self, x: TensorLike) -> TensorLike:
         """
         Parameters
         ----------
 
-        x : torch.Tensor
+        x : TensorLike
             The input tensor to be standardized.
 
         """
-        if not isinstance(x, torch.Tensor):
-            raise ValueError(f"Expected 2D torch.Tensor, actual type {type(x)}")
+        if not isinstance(x, TensorLike):
+            raise ValueError(f"Expected 2D TensorLike, actual type {type(x)}")
         if len(x.shape) != 2:
-            raise ValueError(
-                f"Expected 2D torch.Tensor, actual shape dim {len(x.shape)}"
-            )
+            raise ValueError(f"Expected 2D TensorLike, actual shape dim {len(x.shape)}")
         return (x - self.mean) / self.std

autoemulate-0.3.2/autoemulate/experimental/data/utils.py

@@ -0,0 +1,483 @@
+import numpy as np
+import torch
+from autoemulate.experimental.types import (
+    InputLike,
+    OutputLike,
+    TensorLike,
+    TorchScalarDType,
+)
+from sklearn.utils.validation import check_X_y
+from torch.utils.data import DataLoader, Dataset, Subset, TensorDataset, random_split
+
+
+class ConversionMixin:
+    """
+    Mixin class to convert input data to pytorch Datasets and DataLoaders.
+    """
+
+    def _convert_to_dataset(
+        self,
+        x: InputLike,
+        y: InputLike | None = None,
+    ) -> Dataset:
+        """
+        Convert input data to pytorch Dataset.
+        """
+        # Convert input to Dataset if not already
+        if isinstance(x, np.ndarray):
+            x = torch.tensor(x, dtype=torch.float32)
+        if isinstance(y, np.ndarray):
+            y = torch.tensor(y, dtype=torch.float32)
+
+        if isinstance(x, (torch.Tensor, np.ndarray)) and isinstance(
+            y, (torch.Tensor, np.ndarray)
+        ):
+            dataset = TensorDataset(x, y)
+        elif isinstance(x, (torch.Tensor, np.ndarray)) and y is None:
+            dataset = TensorDataset(x)
+        elif isinstance(x, Dataset) and y is None:
+            dataset = x
+        elif isinstance(x, DataLoader) and y is None:
+            dataset = x.dataset
+        else:
+            raise ValueError(
+                f"Unsupported type for x ({type(x)}). Must be numpy array or PyTorch "
+                "tensor."
+            )
+
+        return dataset
+
+    def _convert_to_dataloader(
+        self,
+        x: InputLike,
+        y: InputLike | None = None,
+        batch_size: int = 16,
+        shuffle: bool = True,
+    ) -> DataLoader:
+        """
+        Convert input data to pytorch DataLoaders.
+        """
+        if isinstance(x, DataLoader) and y is None:
+            dataloader = x
+        elif isinstance(x, DataLoader) and y is not None:
+            raise ValueError(
+                f"Since x is already a DataLoader, expect y to be None, not {type(y)}."
+            )
+        else:
+            dataset = self._convert_to_dataset(x, y)
+            dataloader = DataLoader(dataset, batch_size=batch_size, shuffle=shuffle)
+        return dataloader
+
+    def _convert_to_tensors(
+        self,
+        x: InputLike,
+        y: InputLike | None = None,
+        dtype: torch.dtype = torch.float32,
+    ) -> torch.Tensor | tuple[torch.Tensor, torch.Tensor]:
+        """
+        Convert InputLike x, y to Tensor or tuple of Tensors.
+        """
+        dataset = self._convert_to_dataset(x, y)
+
+        # Handle Subset of TensorDataset
+        if isinstance(dataset, Subset):
+            if isinstance(dataset.dataset, TensorDataset):
+                tensors = dataset.dataset.tensors
+                indices = dataset.indices
+
+                # Use indexing to get subset tensors
+                subset_tensors = tuple(tensor[indices] for tensor in tensors)
+                dataset = TensorDataset(*subset_tensors)
+            else:
+                raise ValueError(
+                    f"Subset must wrap a TensorDataset. Found {type(dataset.dataset)}."
+                )
+
+        if isinstance(dataset, TensorDataset):
+            if len(dataset.tensors) > 2:
+                raise ValueError(
+                    f"Dataset must have 2 or fewer tensors. Found "
+                    f"{len(dataset.tensors)}."
+                )
+            if len(dataset.tensors) == 2:
+                x, y = dataset.tensors
+                assert x.ndim == 2
+                assert y.ndim in (1, 2)
+                # Ensure always 2D tensors
+                if y.ndim == 1:
+                    y = y.unsqueeze(1)
+                return x.to(dtype), y.to(dtype)
+            if len(dataset.tensors) == 1:
+                (x,) = dataset.tensors
+                assert x.ndim == 2
+                return x.to(dtype)
+            msg = "Number of tensors returned must be greater than zero."
+            raise ValueError(msg)
+        # Note: this error will never be raised, the same error is raised in
+        # _convert_to_dataset
+        raise ValueError(
+            f"Unsupported type for x ({type(x)}). Must be numpy array or PyTorch "
+            "tensor."
+        )
+
+    def _convert_to_numpy(
+        self,
+        x: InputLike,
+        y: InputLike | None = None,
+    ) -> tuple[np.ndarray, np.ndarray | None]:
+        """
+        Convert InputLike x, y to tuple of numpy arrays.
+        """
+        if isinstance(x, np.ndarray) and (y is None or isinstance(y, np.ndarray)):
+            return x, y
+
+        result = self._convert_to_tensors(x, y)
+        if isinstance(result, tuple):
+            x_tensor, y_tensor = result
+            x_np, y_np = x_tensor.cpu().numpy(), y_tensor.cpu().numpy()
+            if (y_np.ndim == 2 and y_np.shape[1] == 1) or y_np.ndim == 1:
+                y_np = y_np.ravel()  # Ensure y is 1-dimensional
+                return check_X_y(x_np, y_np, multi_output=False, y_numeric=True)
+            return check_X_y(x_np, y_np, multi_output=True, y_numeric=True)
+
+        x_tensor = result
+        return x_tensor.cpu().numpy(), None
+
+    def _random_split(
+        self,
+        dataset: Dataset,
+        batch_size: int | None = None,
+        train_size: float = 0.8,
+        test_size: float = 0.2,
+    ) -> tuple[DataLoader, DataLoader]:
+        """
+        Split Dataset into train/test DataLoaders.
+
+        Parameters
+        ----------
+        dataset: Dataset
+            The data to split.
+        batch_size: int | None
+            The DataLoader batch_size. If None, sets batch_size to lenth of training
+            data. Defaults to None.
+        """
+        if train_size < 0.0 or train_size > 1.0 or test_size < 0.0 or test_size > 1.0:
+            raise ValueError(
+                f"Train size ({train_size}) and test size ({test_size}) must be "
+                "specified as a proportion between 0 and 1"
+            )
+        if test_size + train_size != 1.0:
+            raise ValueError(
+                f"Train size ({train_size}) and test size ({test_size}) must sum to 1"
+            )
+        train, test = tuple(random_split(dataset, [train_size, test_size]))
+        if batch_size is None:
+            batch_size = len(train)
+        train_loader = self._convert_to_dataloader(train, batch_size=batch_size)
+        test_loader = self._convert_to_dataloader(test, batch_size=batch_size)
+        return train_loader, test_loader
+
+    @staticmethod
+    def _normalize(x: TensorLike) -> tuple[TensorLike, TensorLike, TensorLike]:
+        x_mean = x.mean(0, keepdim=True)
+        x_std = x.std(0, keepdim=True)
+        return (x - x_mean) / x_std, x_mean, x_std
+
+    @staticmethod
+    def _denormalize(
+        x: TensorLike, x_mean: TensorLike, x_std: TensorLike
+    ) -> TensorLike:
+        return (x * x_std) + x_mean
+
+
+class ValidationMixin:
+    """
+    Mixin class for validation methods.
+    This class provides static methods for checking the types and shapes of
+    input and output data, as well as validating specific tensor shapes.
+    """
+
+    @staticmethod
+    def _check(x: TensorLike, y: TensorLike | None):
+        """
+        Check the types and shape are correct for the input data.
+        Checks are equivalent to sklearn's check_array.
+        """
+
+        if not isinstance(x, TensorLike):
+            raise ValueError(f"Expected x to be TensorLike, got {type(x)}")
+
+        if y is not None and not isinstance(y, TensorLike):
+            raise ValueError(f"Expected y to be TensorLike, got {type(y)}")
+
+        # Check x
+        if not torch.isfinite(x).all():
+            msg = "Input tensor x contains non-finite values"
+            raise ValueError(msg)
+        if x.dtype not in TorchScalarDType:
+            msg = (
+                f"Input tensor x has unsupported dtype {x.dtype}. "
+                "Expected float32, float64, int32, or int64."
+            )
+            raise ValueError(msg)
+
+        # Check y if not None
+        if y is not None:
+            if not torch.isfinite(y).all():
+                msg = "Input tensor y contains non-finite values"
+                raise ValueError(msg)
+            if y.dtype not in TorchScalarDType:
+                msg = (
+                    f"Input tensor y has unsupported dtype {y.dtype}. "
+                    "Expected float32, float64, int32, or int64."
+                )
+                raise ValueError(msg)
+
+        return x, y
+
+    @staticmethod
+    def _check_output(output: OutputLike):
+        """
+        Check the types and shape are correct
+        for the output data.
+        """
+        if not isinstance(output, OutputLike):
+            raise ValueError(f"Expected OutputLike, got {type(output)}")
+
+        if isinstance(output, TensorLike) and output.ndim != 2:
+            raise ValueError(f"Expected output to be 2D tensor, got {output.ndim}D")
+
+    @staticmethod
+    def check_vector(X: TensorLike) -> TensorLike:
+        """
+        Validate that the input is a 1D TensorLike.
+
+        Parameters
+        ----------
+        X : TensorLike
+            Input tensor to validate.
+
+        Returns
+        -------
+        TensorLike
+            Validated 1D tensor.
+
+        Raises
+        ------
+        ValueError
+            If X is not a TensorLike or is not 1-dimensional.
+        """
+        if not isinstance(X, TensorLike):
+            raise ValueError(f"Expected TensorLike, got {type(X)}")
+        if X.ndim != 1:
+            raise ValueError(f"Expected 1D tensor, got {X.ndim}D")
+        return X
+
+    @staticmethod
+    def check_matrix(X: TensorLike) -> TensorLike:
+        """
+        Validate that the input is a 2D TensorLike.
+
+        Parameters
+        ----------
+        X : TensorLike
+            Input tensor to validate.
+
+        Returns
+        -------
+        TensorLike
+            Validated 2D tensor.
+
+        Raises
+        ------
+        ValueError
+            If X is not a TensorLike or is not 2-dimensional.
+        """
+        if not isinstance(X, TensorLike):
+            raise ValueError(f"Expected TensorLike, got {type(X)}")
+        if X.ndim != 2:
+            raise ValueError(f"Expected 2D tensor, got {X.ndim}D")
+        return X
+
+    @staticmethod
+    def check_pair(X: TensorLike, Y: TensorLike) -> tuple[TensorLike, TensorLike]:
+        """
+        Validate that two tensors have the same number of rows.
+
+        Parameters
+        ----------
+        X : TensorLike
+            First tensor.
+        Y : TensorLike
+            Second tensor.
+
+        Returns
+        -------
+        tuple[TensorLike, TensorLike]
+            The validated pair of tensors.
+
+        Raises
+        ------
+        ValueError
+            If X and Y do not have the same number of rows.
+        """
+        if X.shape[0] != Y.shape[0]:
+            msg = "X and Y must have the same number of rows"
+            raise ValueError(msg)
+        return X, Y
+
+    @staticmethod
+    def check_covariance(Y: TensorLike, Sigma: TensorLike) -> TensorLike:
+        """
+        Validate and return the covariance matrix.
+
+        Parameters
+        ----------
+        Y : TensorLike
+            Output tensor.
+        Sigma : TensorLike
+            Covariance matrix, which may be full, diagonal, or a scalar per sample.
+
+        Returns
+        -------
+        TensorLike
+            Validated covariance matrix.
+
+        Raises
+        ------
+        ValueError
+            If Sigma does not have a valid shape relative to Y.
+        """
+        if (
+            Sigma.shape == (Y.shape[0], Y.shape[1], Y.shape[1])
+            or Sigma.shape == (Y.shape[0], Y.shape[1])
+            or Sigma.shape == (Y.shape[0],)
+        ):
+            return Sigma
+        msg = "Invalid covariance matrix shape"
+        raise ValueError(msg)
+
+    @staticmethod
+    def trace(Sigma: TensorLike, d: int) -> TensorLike:
+        """
+        Compute the trace of the covariance matrix (A-optimal design criterion).
+
+        Parameters
+        ----------
+        Sigma : TensorLike
+            Covariance matrix (full, diagonal, or scalar).
+        d : int
+            Dimension of the output.
+
+        Returns
+        -------
+        TensorLike
+            The computed trace value.
+
+        Raises
+        ------
+        ValueError
+            If Sigma does not have a valid shape.
+        """
+        if Sigma.dim() == 3 and Sigma.shape[1:] == (d, d):
+            return torch.diagonal(Sigma, dim1=1, dim2=2).sum(dim=1).mean()
+        if Sigma.dim() == 2 and Sigma.shape[1] == d:
+            return Sigma.sum(dim=1).mean()
+        if Sigma.dim() == 1:
+            return d * Sigma.mean()
+        raise ValueError(f"Invalid covariance matrix shape: {Sigma.shape}")
+
+    @staticmethod
+    def logdet(Sigma: TensorLike, dim: int) -> TensorLike:
+        """
+        Compute the log-determinant of the covariance matrix (D-optimal design
+        criterion).
+
+        Parameters
+        ----------
+        Sigma : TensorLike
+            Covariance matrix (full, diagonal, or scalar).
+        dim : int
+            Dimension of the output.
+
+        Returns
+        -------
+        TensorLike
+            The computed log-determinant value.
+
+        Raises
+        ------
+        ValueError
+            If Sigma does not have a valid shape.
+        """
+        if len(Sigma.shape) == 3 and Sigma.shape[1:] == (dim, dim):
+            return torch.logdet(Sigma).mean()
+        if len(Sigma.shape) == 2 and Sigma.shape[1] == dim:
+            return torch.sum(torch.log(Sigma), dim=1).mean()
+        if len(Sigma.shape) == 1:
+            return dim * torch.log(Sigma).mean()
+        raise ValueError(f"Invalid covariance matrix shape: {Sigma.shape}")
+
+    @staticmethod
+    def max_eigval(Sigma: TensorLike) -> TensorLike:
+        """
+        Compute the maximum eigenvalue of the covariance matrix (E-optimal design
+        criterion).
+
+        Parameters
+        ----------
+        Sigma : TensorLike
+            Covariance matrix (full, diagonal, or scalar).
+
+        Returns
+        -------
+        TensorLike
+            The average maximum eigenvalue.
+
+        Raises
+        ------
+        ValueError
+            If Sigma does not have a valid shape.
+        """
+        if Sigma.dim() == 3 and Sigma.shape[1:] == (Sigma.shape[1], Sigma.shape[1]):
+            eigvals = torch.linalg.eigvalsh(Sigma)
+            return eigvals[:, -1].mean()  # Eigenvalues are sorted in ascending order
+        if Sigma.dim() == 2:
+            return Sigma.max(dim=1).values.mean()
+        if Sigma.dim() == 1:
+            return Sigma.mean()
+        raise ValueError(f"Invalid covariance matrix shape: {Sigma.shape}")
+
+    ### Validation methods from old utils.py ###
+    ### Leave here in case want to restore/refactor later ###
+
+    # def _ensure_2d(arr):
+    #     """Ensure that arr is a 2D."""
+    #     if arr.ndim == 1:
+    #         arr = arr.reshape(-1, 1)
+    #     return arr
+
+    # def _ensure_1d_if_column_vec(arr):
+    #     """Ensure that arr is 1D if shape is (n, 1)."""
+    #     if arr.ndim == 2 and arr.shape[1] == 1:
+    #         arr = arr.ravel()
+    #     if arr.ndim > 2 or arr.ndim < 1:
+    #         raise ValueError(
+    #             f"arr should be 1D or 2D. Found {arr.ndim}D array with shape {arr.shape}"  # noqa: E501
+    #         )
+    #     return arr
+
+    # def _check_cv(cv):
+    #     """Ensure that cross-validation method is valid"""
+    #     if cv is None:
+    #         msg = "cross_validator cannot be None"
+    #         raise ValueError(msg)
+    #     if not isinstance(cv, KFold):
+    #         msg = (
+    #             "cross_validator should be an instance of KFold cross-validation. We do not "  # noqa: E501
+    #             "currently support other cross-validation methods."
+    #         )
+    #         raise ValueError(
+    #             msg
+    #         )
+    #     return cv