alchemist-nrel 0.3.1__py3-none-any.whl → 0.3.2__py3-none-any.whl

alchemist_core/models/botorch_model.py

@@ -1,10 +1,12 @@
 import torch
 import numpy as np
 import pandas as pd
+from typing import Union, Tuple, Optional
 from botorch.models import SingleTaskGP
 from botorch.models.gp_regression_mixed import MixedSingleTaskGP
 from botorch.models.transforms import Normalize, Standardize
 from botorch.fit import fit_gpytorch_mll
+from botorch.exceptions import OptimizationWarning
 from gpytorch.mlls import ExactMarginalLogLikelihood
 from sklearn.metrics import mean_squared_error, mean_absolute_error, r2_score
 from sklearn.model_selection import KFold
@@ -15,6 +17,8 @@ from botorch.models.utils.assorted import InputDataWarning
 
 # Import necessary kernels from GPyTorch
 from gpytorch.kernels import MaternKernel, RBFKernel
+import gpytorch
+gpytorch.settings.cholesky_jitter(1e-2)
 
 logger = get_logger(__name__)
 
@@ -61,15 +65,25 @@ class BoTorchModel(BaseModel):
     
     def _get_cont_kernel_factory(self):
         """Returns a factory function for the continuous kernel."""
+        # Validate kernel type before creating factory
+        valid_kernels = ["matern", "rbf"]
+        kernel_lower = self.cont_kernel_type.lower()
+        
+        if kernel_lower not in valid_kernels:
+            raise ValueError(
+                f"Unknown kernel type: '{self.cont_kernel_type}'. "
+                f"Valid options for BoTorch are: {valid_kernels}"
+            )
+        
         def factory(batch_shape, ard_num_dims, active_dims):
-            if self.cont_kernel_type.lower() == "matern":
+            if kernel_lower == "matern":
                 return MaternKernel(
                     nu=self.matern_nu, 
                     ard_num_dims=ard_num_dims, 
                     active_dims=active_dims,
                     batch_shape=batch_shape
                 )
-            else:  # Default to RBF
+            else:  # RBF
                 return RBFKernel(
                     ard_num_dims=ard_num_dims, 
                     active_dims=active_dims,
@@ -154,12 +168,12 @@ class BoTorchModel(BaseModel):
         X_encoded = self._encode_categorical_data(X)
         
         # Convert to tensors
-        train_X = torch.tensor(X_encoded.values, dtype=torch.double)
-        train_Y = torch.tensor(y.values, dtype=torch.double).unsqueeze(-1)
+        train_X = torch.tensor(X_encoded.values, dtype=torch.float64)
+        train_Y = torch.tensor(y.values, dtype=torch.float64).unsqueeze(-1)
         
         # Convert noise values to tensor if available
         if noise is not None:
-            train_Yvar = torch.tensor(noise.values, dtype=torch.double).unsqueeze(-1)
+            train_Yvar = torch.tensor(noise.values, dtype=torch.float64).unsqueeze(-1)
             logger.info(f"Using provided noise values for BoTorch model regularization.")
         else:
             train_Yvar = None
@@ -208,12 +222,25 @@ class BoTorchModel(BaseModel):
                     outcome_transform=outcome_transform
                 )
         else:
-            # For continuous-only models
+            # For continuous-only models, we need to manually construct the covariance module
+            # SingleTaskGP doesn't accept cont_kernel_factory, so we create it and set it manually
+            from gpytorch.kernels import ScaleKernel
+            
+            # Get the kernel from our factory
+            num_dims = train_X.shape[-1]
+            base_kernel = cont_kernel_factory(
+                batch_shape=torch.Size([]),
+                ard_num_dims=num_dims,
+                active_dims=list(range(num_dims))
+            )
+            covar_module = ScaleKernel(base_kernel)
+            
             if noise is not None:
                 self.model = SingleTaskGP(
                     train_X=train_X, 
                     train_Y=train_Y, 
                     train_Yvar=train_Yvar,
+                    covar_module=covar_module,
                     input_transform=input_transform,
                     outcome_transform=outcome_transform
                 )
@@ -222,6 +249,7 @@ class BoTorchModel(BaseModel):
                 self.model = SingleTaskGP(
                     train_X=train_X, 
                     train_Y=train_Y,
+                    covar_module=covar_module,
                     input_transform=input_transform,
                     outcome_transform=outcome_transform
                 )
@@ -259,10 +287,10 @@ class BoTorchModel(BaseModel):
         
         # Convert to tensor - handle both DataFrame and numpy array inputs
         if isinstance(X_encoded, pd.DataFrame):
-            test_X = torch.tensor(X_encoded.values, dtype=torch.double)
+            test_X = torch.tensor(X_encoded.values, dtype=torch.float64)
         else:
             # If X_encoded is already a numpy array
-            test_X = torch.tensor(X_encoded, dtype=torch.double)
+            test_X = torch.tensor(X_encoded, dtype=torch.float64)
         
         # Set model to evaluation mode
         self.model.eval()
@@ -302,10 +330,10 @@ class BoTorchModel(BaseModel):
         
         # Convert to tensor
         if isinstance(X_encoded, pd.DataFrame):
-            X_tensor = torch.tensor(X_encoded.values, dtype=torch.double)
+            X_tensor = torch.tensor(X_encoded.values, dtype=torch.float64)
         else:
             # If X_encoded is already a numpy array
-            X_tensor = torch.tensor(X_encoded, dtype=torch.double)
+            X_tensor = torch.tensor(X_encoded, dtype=torch.float64)
         
         # Set model to evaluation mode
         self.model.eval()
@@ -415,20 +443,23 @@ class BoTorchModel(BaseModel):
         if self.model is None or self.fitted_state_dict is None:
             self.train(experiment_manager)
         
+        # Get target column name from experiment manager
+        target_col = experiment_manager.target_columns[0]
+        
         # Get data - handle noise column if present
         if 'Noise' in exp_df.columns:
-            X = exp_df.drop(columns=["Output", "Noise"])
+            X = exp_df.drop(columns=[target_col, "Noise"])
         else:
-            X = exp_df.drop(columns=["Output"])
+            X = exp_df.drop(columns=[target_col])
             
-        y = exp_df["Output"]
+        y = exp_df[target_col]
         
         # Encode categorical variables
         X_encoded = self._encode_categorical_data(X)
         
         # Convert to tensors
-        full_X = torch.tensor(X_encoded.values, dtype=torch.double)
-        full_Y = torch.tensor(y.values, dtype=torch.double).unsqueeze(-1)
+        full_X = torch.tensor(X_encoded.values, dtype=torch.float64)
+        full_Y = torch.tensor(y.values, dtype=torch.float64).unsqueeze(-1)
         
         # Metrics storage
         rmse_values = []
@@ -438,11 +469,11 @@ class BoTorchModel(BaseModel):
         n_obs = []
         
         # Calculate total steps for progress
-        total_steps = len(range(max(cv_splits+1, 5), len(full_X) + 1))
+        total_steps = len(range(5, len(full_X) + 1))
         current_step = 0
         
-        # Evaluate on increasing subsets of data
-        for i in range(max(cv_splits+1, 5), len(full_X) + 1):
+        # Evaluate on increasing subsets of data (starting at 5 for minimum CV size)
+        for i in range(5, len(full_X) + 1):
             if debug:
                 logger.info(f"Evaluating with {i} observations")
                 
@@ -459,48 +490,68 @@ class BoTorchModel(BaseModel):
             
             # Perform cross-validation for this subset size
             for train_idx, test_idx in kf.split(subset_np_X):
-                # Split data
-                X_train = subset_X[train_idx]
-                y_train = subset_Y[train_idx]
-                X_test = subset_X[test_idx]
-                y_test = subset_Y[test_idx]
-                
-                # Create a new model with this fold's training data
-                # Need to recreate transforms with the same parameters as the main model
-                fold_input_transform, fold_outcome_transform = self._create_transforms(X_train, y_train)
-                
-                cont_kernel_factory = self._get_cont_kernel_factory()
-                if self.cat_dims and len(self.cat_dims) > 0:
-                    fold_model = MixedSingleTaskGP(
-                        X_train, y_train, 
-                        cat_dims=self.cat_dims,
-                        cont_kernel_factory=cont_kernel_factory,
-                        input_transform=fold_input_transform,
-                        outcome_transform=fold_outcome_transform
-                    )
-                else:
-                    fold_model = SingleTaskGP(
-                        X_train, y_train,
-                        input_transform=fold_input_transform,
-                        outcome_transform=fold_outcome_transform
-                    )
-                
-                # Train the fold model from scratch (don't load state_dict to avoid dimension mismatches)
-                # This is necessary because folds may have different categorical values or data shapes
-                mll = ExactMarginalLogLikelihood(fold_model.likelihood, fold_model)
-                fit_gpytorch_mll(mll)
-                
-                # Make predictions on test fold
-                fold_model.eval()
-                fold_model.likelihood.eval()
-                
-                with torch.no_grad():
-                    posterior = fold_model.posterior(X_test)
-                    preds = posterior.mean.squeeze(-1)
+                try:
+                    # Split data
+                    X_train = subset_X[train_idx]
+                    y_train = subset_Y[train_idx]
+                    X_test = subset_X[test_idx]
+                    y_test = subset_Y[test_idx]
+                    
+                    # Create a new model with this fold's training data
+                    # Need to recreate transforms with the same parameters as the main model
+                    fold_input_transform, fold_outcome_transform = self._create_transforms(X_train, y_train)
+                    
+                    cont_kernel_factory = self._get_cont_kernel_factory()
+                    if self.cat_dims and len(self.cat_dims) > 0:
+                        fold_model = MixedSingleTaskGP(
+                            X_train, y_train, 
+                            cat_dims=self.cat_dims,
+                            cont_kernel_factory=cont_kernel_factory,
+                            input_transform=fold_input_transform,
+                            outcome_transform=fold_outcome_transform
+                        )
+                    else:
+                        fold_model = SingleTaskGP(
+                            X_train, y_train,
+                            input_transform=fold_input_transform,
+                            outcome_transform=fold_outcome_transform
+                        )
+                    
+                    # Train the fold model from scratch (don't load state_dict to avoid dimension mismatches)
+                    # This is necessary because folds may have different categorical values or data shapes
+                    mll = ExactMarginalLogLikelihood(fold_model.likelihood, fold_model)
+                    
+                    # Suppress optimization warnings for small folds where convergence may be difficult
+                    import warnings
+                    with warnings.catch_warnings():
+                        warnings.filterwarnings('ignore', category=OptimizationWarning)
+                        # Use fit_gpytorch_mll with options that improve convergence for small datasets
+                        fit_gpytorch_mll(
+                            mll,
+                            options={
+                                "maxiter": 50,  # Reduce iterations for speed
+                                "ftol": 1e-6,   # Slightly relaxed tolerance
+                                "gtol": 1e-5,   # Slightly relaxed gradient tolerance
+                            }
+                        )
+                    
+                    # Make predictions on test fold
+                    fold_model.eval()
+                    fold_model.likelihood.eval()
                     
-                    # Store this fold's results
-                    fold_y_trues.append(y_test.squeeze(-1))
-                    fold_y_preds.append(preds)
+                    with torch.no_grad():
+                        posterior = fold_model.posterior(X_test)
+                        preds = posterior.mean.squeeze(-1)
+                        
+                        # Store this fold's results
+                        fold_y_trues.append(y_test.squeeze(-1))
+                        fold_y_preds.append(preds)
+                        
+                except Exception as e:
+                    # Skip this fold if optimization fails (can happen with small/difficult training sets)
+                    if debug:
+                        logger.warning(f"Skipping fold for subset size {i} due to error: {e}")
+                    continue
             
             # Combine all fold results for this subset size
             all_y_true = torch.cat(fold_y_trues).cpu().numpy()
@@ -795,14 +846,14 @@ class BoTorchModel(BaseModel):
         # Convert pandas/numpy data to tensors if needed
         if isinstance(X, pd.DataFrame):
             X_encoded = self._encode_categorical_data(X)
-            X_tensor = torch.tensor(X_encoded.values, dtype=torch.double)
+            X_tensor = torch.tensor(X_encoded.values, dtype=torch.float64)
         elif isinstance(X, np.ndarray):
-            X_tensor = torch.tensor(X, dtype=torch.double)
+            X_tensor = torch.tensor(X, dtype=torch.float64)
         else:
             X_tensor = X  # Assume it's already a tensor
             
         if isinstance(y, pd.Series) or isinstance(y, np.ndarray):
-            y_tensor = torch.tensor(y, dtype=torch.double).unsqueeze(-1)
+            y_tensor = torch.tensor(y, dtype=torch.float64).unsqueeze(-1)
         else:
             y_tensor = y  # Assume it's already a tensor
         
@@ -921,4 +972,182 @@ class BoTorchModel(BaseModel):
         else:
             logger.info("  ✓ Uncertainty appears well-calibrated")
         
-        logger.info(f"{'='*60}\n")
+        logger.info(f"{'='*60}\n")
+
+    def evaluate_acquisition(
+        self, 
+        X: Union[pd.DataFrame, np.ndarray], 
+        acq_func: str = 'ucb', 
+        acq_func_kwargs: Optional[dict] = None, 
+        maximize: bool = True
+    ) -> Tuple[np.ndarray, None]:
+        """
+        Evaluate acquisition function at given points using BoTorch functions.
+        
+        Args:
+            X: Points to evaluate (DataFrame or array with shape (n, d))
+            acq_func: Acquisition function name 
+                     Analytic: 'ei', 'logei', 'pi', 'logpi', 'ucb'
+                     Batch: 'qei', 'qucb', 'qnipv'
+            acq_func_kwargs: Additional parameters (e.g., {'beta': 0.5, 'mc_samples': 128})
+            maximize: Whether we're maximizing (True) or minimizing (False)
+            
+        Returns:
+            Tuple of (acq_values, None) - None because acq functions are deterministic
+            
+        Example:
+            >>> points = pd.DataFrame({'temp': [300, 350, 400], 'pressure': [1, 2, 3]})
+            >>> acq_vals, _ = model.evaluate_acquisition(points, acq_func='ei', maximize=True)
+        """
+        from botorch.acquisition.analytic import (
+            ExpectedImprovement,
+            LogExpectedImprovement,
+            ProbabilityOfImprovement,
+            LogProbabilityOfImprovement,
+            UpperConfidenceBound,
+        )
+        from botorch.acquisition.monte_carlo import (
+            qExpectedImprovement,
+            qUpperConfidenceBound,
+        )
+        from botorch.acquisition.active_learning import qNegIntegratedPosteriorVariance
+        from botorch.sampling import SobolQMCNormalSampler
+        
+        if not self.is_trained:
+            raise ValueError("Model must be trained before evaluating acquisition functions.")
+        
+        # Encode categorical variables (same preprocessing as predict())
+        X_encoded = self._encode_categorical_data(X)
+        
+        # Convert to torch tensor
+        if isinstance(X_encoded, pd.DataFrame):
+            X_tensor = torch.tensor(X_encoded.values, dtype=torch.float64)
+        else:
+            # If X_encoded is already a numpy array
+            X_tensor = torch.tensor(X_encoded, dtype=torch.float64)
+        
+        # Add q=1 dimension if not present (batch_size, d) -> (batch_size, 1, d)
+        if X_tensor.ndim == 2:
+            X_tensor = X_tensor.unsqueeze(-2)
+        
+        # Calculate best_f from ORIGINAL (untransformed) training data
+        # When using Standardize transform, best_f must be in original scale
+        if hasattr(self, 'Y_orig') and self.Y_orig is not None:
+            y_train_tensor = self.Y_orig
+        else:
+            # Fallback: use model's train_targets (may be transformed)
+            y_train_tensor = self.model.train_targets
+        
+        if maximize:
+            best_f = torch.max(y_train_tensor)
+        else:
+            best_f = torch.min(y_train_tensor)
+        
+        # Map acquisition function names
+        acq_func_lower = acq_func.lower()
+        
+        # Parse kwargs with defaults
+        if acq_func_kwargs is None:
+            acq_func_kwargs = {}
+        
+        beta = acq_func_kwargs.get('beta', 0.5)
+        mc_samples = acq_func_kwargs.get('mc_samples', 128)
+        
+        # Determine if this is a batch (q) acquisition function
+        is_batch_acq = acq_func_lower.startswith('q')
+        
+        # Create acquisition function
+        try:
+            if acq_func_lower in ['ei', 'expectedimprovement']:
+                acq_fn = ExpectedImprovement(
+                    model=self.model,
+                    best_f=best_f,
+                    maximize=maximize
+                )
+            elif acq_func_lower in ['logei', 'logexpectedimprovement']:
+                acq_fn = LogExpectedImprovement(
+                    model=self.model,
+                    best_f=best_f,
+                    maximize=maximize
+                )
+            elif acq_func_lower in ['pi', 'probabilityofimprovement']:
+                acq_fn = ProbabilityOfImprovement(
+                    model=self.model,
+                    best_f=best_f,
+                    maximize=maximize
+                )
+            elif acq_func_lower in ['logpi', 'logprobabilityofimprovement']:
+                acq_fn = LogProbabilityOfImprovement(
+                    model=self.model,
+                    best_f=best_f,
+                    maximize=maximize
+                )
+            elif acq_func_lower in ['ucb', 'upperconfidencebound']:
+                acq_fn = UpperConfidenceBound(
+                    model=self.model,
+                    beta=beta,
+                    maximize=maximize
+                )
+            elif acq_func_lower in ['qei', 'qexpectedimprovement']:
+                sampler = SobolQMCNormalSampler(sample_shape=torch.Size([mc_samples]))
+                acq_fn = qExpectedImprovement(
+                    model=self.model,
+                    best_f=best_f,
+                    sampler=sampler
+                )
+            elif acq_func_lower in ['qucb', 'qupperconfidencebound']:
+                sampler = SobolQMCNormalSampler(sample_shape=torch.Size([mc_samples]))
+                acq_fn = qUpperConfidenceBound(
+                    model=self.model,
+                    beta=beta,
+                    sampler=sampler
+                )
+            elif acq_func_lower in ['qnipv', 'qnegintegratedposteriorvariance', 'qipv']:
+                # qNIPV requires mc_points for integration over the search space
+                n_mc_points = acq_func_kwargs.get('n_mc_points', 500)
+                
+                # Generate MC points uniformly over the input space
+                # Get bounds from the encoded input space
+                if hasattr(self, 'X_train') and self.X_train is not None:
+                    # Use training data bounds
+                    lower_bounds = self.X_train.min(dim=0)[0]
+                    upper_bounds = self.X_train.max(dim=0)[0]
+                else:
+                    # Fallback: assume normalized space [0, 1]
+                    n_dims = X_tensor.shape[-1]
+                    lower_bounds = torch.zeros(n_dims, dtype=torch.float64)
+                    upper_bounds = torch.ones(n_dims, dtype=torch.float64)
+                
+                # Generate random points
+                mc_points = torch.rand(n_mc_points, len(lower_bounds), dtype=torch.float64)
+                mc_points = mc_points * (upper_bounds - lower_bounds) + lower_bounds
+                
+                acq_fn = qNegIntegratedPosteriorVariance(
+                    model=self.model,
+                    mc_points=mc_points
+                )
+            else:
+                raise ValueError(
+                    f"Unknown acquisition function '{acq_func}' for BoTorch backend. "
+                    f"Valid options are: 'ei', 'logei', 'pi', 'logpi', 'ucb', 'qei', 'qucb', 'qnipv'"
+                )
+            
+            # Evaluate acquisition function
+            with torch.no_grad():
+                if is_batch_acq:
+                    # For batch acquisitions, evaluate each point as q=1
+                    # X_tensor is already (batch_size, 1, d)
+                    acq_values = acq_fn(X_tensor).cpu().numpy()
+                else:
+                    # For analytic acquisitions
+                    acq_values = acq_fn(X_tensor).cpu().numpy()
+            
+            # Ensure output is 1D array
+            if acq_values.ndim > 1:
+                acq_values = acq_values.ravel()
+            
+            return acq_values, None
+            
+        except Exception as e:
+            logger.error(f"Error evaluating acquisition function: {e}")
+            raise