dragon-ml-toolbox 10.2.0__py3-none-any.whl → 14.2.0__py3-none-any.whl

ml_tools/ML_optimization.py

@@ -5,7 +5,7 @@ import evotorch
 from evotorch.algorithms import SNES, CEM, GeneticAlgorithm
 from evotorch.logging import PandasLogger
 from evotorch.operators import SimulatedBinaryCrossOver, GaussianMutation
-from typing import Literal, Union, Tuple, List, Optional, Any, Callable
+from typing import Literal, Union, Tuple, List, Optional, Any, Callable, Dict
 from pathlib import Path
 from tqdm.auto import trange
 from contextlib import nullcontext
@@ -17,19 +17,216 @@ from ._script_info import _script_info
 from .ML_inference import PyTorchInferenceHandler
 from .keys import PyTorchInferenceKeys
 from .SQL import DatabaseManager
-from .optimization_tools import _save_result
-from .utilities import threshold_binary_values, save_dataframe
+from .optimization_tools import _save_result, create_optimization_bounds
+from .utilities import save_dataframe_filename
+from .math_utilities import discretize_categorical_values
+from ._schema import FeatureSchema
+
 
 __all__ = [
+    "MLOptimizer",
+    "FitnessEvaluator",
     "create_pytorch_problem",
     "run_optimization"
 ]
 
 
+class MLOptimizer:
+    """
+    A wrapper class for setting up and running EvoTorch optimization tasks.
+
+    This class combines the functionality of `FitnessEvaluator`, `create_pytorch_problem`, and
+    `run_optimization` into a single, streamlined workflow. 
+    
+    SNES and CEM algorithms do not accept bounds, the given bounds will be used as an initial starting point.
+
+    Example:
+        >>> # 1. Get the final schema from data exploration
+        >>> schema = data_exploration.finalize_feature_schema(...)
+        >>> # 2. Define bounds for continuous features
+        >>> cont_bounds = {'feature_A': (0, 100), 'feature_B': (-10, 10)}
+        >>>
+        >>> # 3. Initialize the optimizer
+        >>> optimizer = MLOptimizer(
+        ...     inference_handler=my_handler,
+        ...     schema=schema,
+        ...     continuous_bounds_map=cont_bounds,
+        ...     task="max",
+        ...     algorithm="Genetic",
+        ... )
+        >>> # 4. Run the optimization
+        >>> best_result = optimizer.run(
+        ...     num_generations=100,
+        ...     target_name="my_target",
+        ...     save_dir="/path/to/results",
+        ...     save_format="csv"
+        ... )
+    """
+    def __init__(self,
+                 inference_handler: PyTorchInferenceHandler,
+                 schema: FeatureSchema,
+                 continuous_bounds_map: Dict[str, Tuple[float, float]],
+                 task: Literal["min", "max"],
+                 algorithm: Literal["SNES", "CEM", "Genetic"] = "Genetic",
+                 population_size: int = 200,
+                 discretize_start_at_zero: bool = True,
+                 **searcher_kwargs):
+        """
+        Initializes the optimizer by creating the EvoTorch problem and searcher.
+
+        Args:
+            inference_handler (PyTorchInferenceHandler): 
+                An initialized inference handler containing the model.
+            schema (FeatureSchema): 
+                The definitive schema object from data_exploration.
+            continuous_bounds_map (Dict[str, Tuple[float, float]]):
+                A dictionary mapping the *name* of each **continuous** feature
+                to its (min_bound, max_bound) tuple.
+            task (str): The optimization goal, either "min" or "max".
+            algorithm (str): The search algorithm to use ("SNES", "CEM", "Genetic").
+            population_size (int): Population size for CEM and GeneticAlgorithm.
+            discretize_start_at_zero (bool): 
+                True if the discrete encoding starts at 0 (e.g., [0, 1, 2]).
+                False if it starts at 1 (e.g., [1, 2, 3]).
+            **searcher_kwargs: Additional keyword arguments for the selected 
+                               search algorithm's constructor.
+        """
+        # --- Store schema ---
+        self.schema = schema
+        
+        # --- 1. Create bounds from schema ---
+        # This is the new, robust way to get bounds
+        bounds = create_optimization_bounds(
+            schema=schema,
+            continuous_bounds_map=continuous_bounds_map,
+            start_at_zero=discretize_start_at_zero
+        )
+
+        # --- 2. Make a fitness function ---
+        self.evaluator = FitnessEvaluator(
+            inference_handler=inference_handler,
+            # Get categorical info from the schema
+            categorical_index_map=schema.categorical_index_map,
+            discretize_start_at_zero=discretize_start_at_zero
+        )
+        
+        # --- 3. Create the problem and searcher factory ---
+        self.problem, self.searcher_factory = create_pytorch_problem(
+            evaluator=self.evaluator,
+            bounds=bounds,
+            task=task,
+            algorithm=algorithm,
+            population_size=population_size,
+            **searcher_kwargs
+        )
+        
+        # --- 4. Store other info needed by run() ---
+        self.discretize_start_at_zero = discretize_start_at_zero
+
+    def run(self,
+            num_generations: int,
+            target_name: str,
+            save_dir: Union[str, Path],
+            save_format: Literal['csv', 'sqlite', 'both'],
+            repetitions: int = 1,
+            verbose: bool = True) -> Optional[dict]:
+        """
+        Runs the evolutionary optimization process using the pre-configured settings.
+
+        The `feature_names` are automatically pulled from the `FeatureSchema`
+        provided during initialization.
+
+        Args:
+            num_generations (int): The total number of generations for each repetition.
+            target_name (str): Target name used for the CSV filename and/or SQL table.
+            save_dir (str | Path): The directory where result files will be saved.
+            save_format (Literal['csv', 'sqlite', 'both']): The format for saving results.
+            repetitions (int): The number of independent times to run the optimization.
+            verbose (bool): If True, enables detailed logging.
+
+        Returns:
+            Optional[dict]: A dictionary with the best result if repetitions is 1, 
+                            otherwise None.
+        """
+        # Call the existing run function, passing info from the schema
+        return run_optimization(
+            problem=self.problem,
+            searcher_factory=self.searcher_factory,
+            num_generations=num_generations,
+            target_name=target_name,
+            save_dir=save_dir,
+            save_format=save_format,
+            # Get the definitive feature names (as a list) from the schema
+            feature_names=list(self.schema.feature_names),
+            # Get categorical info from the schema
+            categorical_map=self.schema.categorical_index_map,
+            categorical_mappings=self.schema.categorical_mappings,
+            repetitions=repetitions,
+            verbose=verbose,
+            discretize_start_at_zero=self.discretize_start_at_zero
+        )
+        
+        
+class FitnessEvaluator:
+    """
+    A callable class that wraps the PyTorch model inference handler and performs
+    on-the-fly discretization for the EvoTorch fitness function.
+
+    This class is automatically instantiated by MLOptimizer and passed to
+    create_pytorch_problem, encapsulating the evaluation logic.
+    """
+    def __init__(self,
+                 inference_handler: PyTorchInferenceHandler,
+                 categorical_index_map: Optional[Dict[int, int]] = None,
+                 discretize_start_at_zero: bool = True):
+        """
+        Initializes the fitness evaluator.
+
+        Args:
+            inference_handler (PyTorchInferenceHandler): 
+                An initialized inference handler containing the model.
+            categorical_index_map (Dict[int, int] | None): 
+                Maps {column_index: cardinality} for discretization.
+            discretize_start_at_zero (bool): 
+                True if discrete encoding starts at 0.
+        """
+        self.inference_handler = inference_handler
+        self.categorical_index_map = categorical_index_map
+        self.discretize_start_at_zero = discretize_start_at_zero
+        
+        # Expose the device
+        self.device = self.inference_handler.device
+
+    def __call__(self, solution_tensor: torch.Tensor) -> torch.Tensor:
+        """
+        This is the fitness function EvoTorch will call.
+
+        It receives a batch of continuous solutions, discretizes the
+        categorical ones, and returns the model's predictions.
+        """
+        # Clone to avoid modifying the optimizer's internal state (SNES, CEM, GA)
+        processed_tensor = solution_tensor.clone()
+        
+        if self.categorical_index_map:
+            for col_idx, cardinality in self.categorical_index_map.items():
+                # 1. Round (using torch.floor(x + 0.5) for "round half up" behavior)
+                rounded_col = torch.floor(processed_tensor[:, col_idx] + 0.5)
+                
+                # 2. Determine clamping bounds
+                min_bound = 0 if self.discretize_start_at_zero else 1
+                max_bound = cardinality - 1 if self.discretize_start_at_zero else cardinality
+                
+                # 3. Clamp the values and update the processed tensor
+                processed_tensor[:, col_idx] = torch.clamp(rounded_col, min_bound, max_bound)
+
+        # Use the *processed_tensor* for prediction
+        predictions = self.inference_handler.predict_batch(processed_tensor)[PyTorchInferenceKeys.PREDICTIONS]
+        return predictions.flatten()
+
+
 def create_pytorch_problem(
-    inference_handler: PyTorchInferenceHandler,
+    evaluator: FitnessEvaluator,
     bounds: Tuple[List[float], List[float]],
-    binary_features: int,
     task: Literal["min", "max"],
     algorithm: Literal["SNES", "CEM", "Genetic"] = "Genetic",
     population_size: int = 200,
@@ -38,14 +235,14 @@ def create_pytorch_problem(
     """
     Creates and configures an EvoTorch Problem and a Searcher factory class for a PyTorch model.
     
-    SNES and CEM do not accept bounds, the given bounds will be used as initial bounds only.
+    SNES and CEM do not accept bounds, the given bounds will be used as an initial starting point.
     
     The Genetic Algorithm works directly with the bounds, and operators such as SimulatedBinaryCrossOver and GaussianMutation.
     
     Args:
-        inference_handler (PyTorchInferenceHandler): An initialized inference handler containing the model and weights.
+        evaluator (FitnessEvaluator): A callable class that wraps the model inference and handles on-the-fly discretization.
         bounds (tuple[list[float], list[float]]): A tuple containing the lower and upper bounds for the solution features.
-        binary_features (int): Number of binary features located at the END of the feature vector. Will be automatically added to the bounds.
+            Use the `optimization_tools.create_optimization_bounds()` helper to easily generate this and ensure unbiased categorical bounds.
         task (str): The optimization goal, either "minimize" or "maximize".
         algorithm (str): The search algorithm to use.
         population_size (int): Used for CEM and GeneticAlgorithm.
@@ -60,26 +257,14 @@ def create_pytorch_problem(
     lower_bounds = list(bounds[0])
     upper_bounds = list(bounds[1])
     
-    # add binary bounds
-    if binary_features > 0:
-        lower_bounds.extend([0.45] * binary_features)
-        upper_bounds.extend([0.55] * binary_features)
-    
     solution_length = len(lower_bounds)
-    device = inference_handler.device
+    device = evaluator.device
 
-    # Define the fitness function that EvoTorch will call.
-    def fitness_func(solution_tensor: torch.Tensor) -> torch.Tensor:
-        # Directly use the continuous-valued tensor from the optimizer for prediction
-        predictions = inference_handler.predict_batch(solution_tensor)[PyTorchInferenceKeys.PREDICTIONS]
-        return predictions.flatten()
-    
-    
     # Create the Problem instance.
     if algorithm == "CEM" or algorithm == "SNES":
         problem = evotorch.Problem(
             objective_sense=task,
-            objective_func=fitness_func,
+            objective_func=evaluator,
             solution_length=solution_length,
             initial_bounds=(lower_bounds, upper_bounds),
             device=device,
@@ -105,7 +290,7 @@ def create_pytorch_problem(
     elif algorithm == "Genetic":
         problem = evotorch.Problem(
             objective_sense=task,
-            objective_func=fitness_func,
+            objective_func=evaluator,
             solution_length=solution_length,
             bounds=(lower_bounds, upper_bounds),
             device=device,
@@ -141,12 +326,14 @@ def run_optimization(
     searcher_factory: Callable[[],Any],
     num_generations: int,
     target_name: str,
-    binary_features: int,
     save_dir: Union[str, Path],
     save_format: Literal['csv', 'sqlite', 'both'],
     feature_names: Optional[List[str]],
     repetitions: int = 1,
-    verbose: bool = True
+    verbose: bool = True,
+    categorical_map: Optional[Dict[int, int]] = None,
+    categorical_mappings: Optional[Dict[str, Dict[str, int]]] = None,
+    discretize_start_at_zero: bool = True
 ) -> Optional[dict]:
     """
     Runs the evolutionary optimization process, with support for multiple repetitions.
@@ -169,7 +356,6 @@ def run_optimization(
         searcher_factory (Callable): The searcher factory to generate fresh evolutionary algorithms.
         num_generations (int): The total number of generations to run the search algorithm for in each repetition.
         target_name (str): Target name that will also be used for the CSV filename and SQL table.
-        binary_features (int): Number of binary features located at the END of the feature vector.
         save_dir (str | Path): The directory where the result file(s) will be saved.
         save_format (Literal['csv', 'sqlite', 'both'], optional): The format for
             saving results during iterative analysis.
@@ -179,13 +365,18 @@ def run_optimization(
         repetitions (int, optional): The number of independent times to run the
             entire optimization process.
         verbose (bool): Add an Evotorch Pandas logger saved as a csv. Only for the first repetition.
+        categorical_index_map (Dict[int, int] | None): Used to discretize values after optimization. Maps {column_index: cardinality}.
+        categorical_mappings (Dict[str, Dict[str, int]] | None): Used to map discrete integer values back to strings (e.g., {0: 'Category_A'}) before saving.
+        discretize_start_at_zero (bool): 
+            True if the discrete encoding starts at 0 (e.g., [0, 1, 2]).
+            False if it starts at 1 (e.g., [1, 2, 3]).
 
     Returns:
         Optional[dict]: A dictionary containing the best feature values and the
         fitness score if `repetitions` is 1. Returns `None` if `repetitions`
         is greater than 1, as results are streamed to files instead.
     """
-    # preprocess paths
+    # --- 1. Setup Paths and Feature Names ---
     save_path = make_fullpath(save_dir, make=True, enforce="directory")
     
     sanitized_target_name = sanitize_filename(target_name)
@@ -193,54 +384,38 @@ def run_optimization(
         sanitized_target_name = sanitized_target_name + ".csv"
     
     csv_path = save_path / sanitized_target_name
-    
     db_path = save_path / "Optimization.db"
     db_table_name = target_name
     
-    # preprocess feature names
+    # Use problem's solution_length to create default names if none provided
     if feature_names is None:
-        feature_names = [f"feature_{i}" for i in range(problem.solution_length)] # type: ignore
+        feat_len = problem.solution_length
+        feature_names = [f"feature_{i}" for i in range(feat_len)] # type: ignore
     
+    # --- 2. Run Optimization ---
     # --- SINGLE RUN LOGIC ---
     if repetitions <= 1:
-        searcher = searcher_factory()
-        _LOGGER.info(f"🤖 Starting optimization with {searcher.__class__.__name__} Algorithm for {num_generations} generations...")
-        # for _ in trange(num_generations, desc="Optimizing"):
-        #     searcher.step()
+        _LOGGER.info(f"🤖 Starting optimization for {num_generations} generations...")
         
-        # Attach logger if requested
-        if verbose:
-            pandas_logger = PandasLogger(searcher)
-            
-        searcher.run(num_generations) # Use the built-in run method for simplicity
-            
-        # # DEBUG new searcher objects
-        # for status_key in searcher.iter_status_keys():
-        #     print("===", status_key, "===")
-        #     print(searcher.status[status_key])
-        #     print()
-        
-        # Get results from the .status dictionary 
-        # SNES and CEM use the key 'center' to get mean values if needed    best_solution_tensor = searcher.status["center"]
-        best_solution_container = searcher.status["pop_best"]
-        best_solution_tensor = best_solution_container.values
-        best_fitness = best_solution_container.evals
-
-        best_solution_np = best_solution_tensor.cpu().numpy()
-        
-        # threshold binary features
-        if binary_features > 0:
-            best_solution_thresholded = threshold_binary_values(input_array=best_solution_np, binary_values=binary_features)
-        else:
-            best_solution_thresholded = best_solution_np
-
-        result_dict = {name: value for name, value in zip(feature_names, best_solution_thresholded)}
-        result_dict[target_name] = best_fitness.item()
+        result_dict, pandas_logger = _run_single_optimization_rep(
+            searcher_factory=searcher_factory,
+            num_generations=num_generations,
+            feature_names=feature_names,
+            target_name=target_name,
+            categorical_map=categorical_map,
+            discretize_start_at_zero=discretize_start_at_zero,
+            attach_logger=verbose
+        )
         
-        _save_result(result_dict, 'csv', csv_path) # Single run defaults to CSV
+        # Single run defaults to CSV, pass mappings for reverse mapping
+        _save_result(
+            result_dict=result_dict, 
+            save_format='csv', 
+            csv_path=csv_path,
+            categorical_mappings=categorical_mappings
+        )
         
-        # Process logger
-        if verbose:
+        if pandas_logger:
             _handle_pandas_log(pandas_logger, save_path=save_path, target_name=target_name)
         
         _LOGGER.info(f"Optimization complete. Best solution saved to '{csv_path.name}'")
@@ -249,60 +424,109 @@ def run_optimization(
     # --- MULTIPLE REPETITIONS LOGIC ---
     else:
         _LOGGER.info(f"🏁 Starting optimal solution space analysis with {repetitions} repetitions...")
-
+        
+        first_run_logger = None # To store the logger from the first rep
         db_context = DatabaseManager(db_path) if save_format in ['sqlite', 'both'] else nullcontext()
         
         with db_context as db_manager:
+            # --- Setup Database Schema (if applicable) ---
             if db_manager:
-                schema = {name: "REAL" for name in feature_names}
+                schema = {}
+                categorical_cols = set(categorical_mappings.keys()) if categorical_mappings else set()
+                
+                for name in feature_names:
+                    schema[name] = "TEXT" if name in categorical_cols else "REAL"
                 schema[target_name] = "REAL"
+                
                 db_manager.create_table(db_table_name, schema)
             
+            # --- Repetitions Loop ---
             print("")
-            # Repetitions loop
-            pandas_logger = None
             for i in trange(repetitions, desc="Repetitions"):
-                # CRITICAL: Create a fresh searcher for each run using the factory
-                searcher = searcher_factory()
-                
-                # Attach logger if requested
-                if verbose and i==0:
-                    pandas_logger = PandasLogger(searcher)
                 
-                searcher.run(num_generations) # Use the built-in run method for simplicity
+                # Only attach a logger for the first repetition if verbose
+                attach_logger = verbose and (i == 0)
                 
-                # Get results from the .status dictionary 
-                # SNES and CEM use the key 'center' to get mean values if needed    best_solution_tensor = searcher.status["center"]
-                best_solution_container = searcher.status["pop_best"]
-                best_solution_tensor = best_solution_container.values
-                best_fitness = best_solution_container.evals
-             
-                best_solution_np = best_solution_tensor.cpu().numpy()
+                result_dict, pandas_logger = _run_single_optimization_rep(
+                    searcher_factory=searcher_factory,
+                    num_generations=num_generations,
+                    feature_names=feature_names,
+                    target_name=target_name,
+                    categorical_map=categorical_map,
+                    discretize_start_at_zero=discretize_start_at_zero,
+                    attach_logger=attach_logger
+                )
                 
-                # threshold binary features
-                if binary_features > 0:
-                    best_solution_thresholded = threshold_binary_values(input_array=best_solution_np, binary_values=binary_features)
-                else:
-                    best_solution_thresholded = best_solution_np
-                
-                # make results dictionary
-                result_dict = {name: value for name, value in zip(feature_names, best_solution_thresholded)}
-                result_dict[target_name] = best_fitness.item()
+                if pandas_logger:
+                    first_run_logger = pandas_logger
                 
                 # Save each result incrementally
-                _save_result(result_dict, save_format, csv_path, db_manager, db_table_name)
+                _save_result(
+                    result_dict=result_dict, 
+                    save_format=save_format, 
+                    csv_path=csv_path, 
+                    db_manager=db_manager, 
+                    db_table_name=db_table_name, 
+                    categorical_mappings=categorical_mappings
+                )
             
-        # Process logger
-        if pandas_logger is not None:
-            _handle_pandas_log(pandas_logger, save_path=save_path, target_name=target_name)      
+        if first_run_logger:
+            _handle_pandas_log(first_run_logger, save_path=save_path, target_name=target_name)      
         
         _LOGGER.info(f"Optimal solution space complete. Results saved to '{save_path}'")
         return None
 
 
+def _run_single_optimization_rep(
+    searcher_factory: Callable[[],Any],
+    num_generations: int,
+    feature_names: List[str],
+    target_name: str,
+    categorical_map: Optional[Dict[int, int]],
+    discretize_start_at_zero: bool,
+    attach_logger: bool
+) -> Tuple[dict, Optional[PandasLogger]]:
+    """
+    Internal helper to run one full optimization repetition.
+    
+    Handles searcher creation, logging, running, and result post-processing.
+    """
+    # CRITICAL: Create a fresh searcher for each run using the factory
+    searcher = searcher_factory()
+    
+    # Attach logger if requested
+    pandas_logger = PandasLogger(searcher) if attach_logger else None
+    
+    # Run the optimization
+    searcher.run(num_generations)
+    
+    # Get the best result
+    best_solution_container = searcher.status["pop_best"]
+    best_solution_tensor = best_solution_container.values
+    best_fitness = best_solution_container.evals
+ 
+    best_solution_np = best_solution_tensor.cpu().numpy()
+    
+    # Discretize categorical/binary features
+    if categorical_map:
+        best_solution_thresholded = discretize_categorical_values(
+            input_array=best_solution_np,
+            categorical_info=categorical_map,
+            start_at_zero=discretize_start_at_zero
+        )
+    else:
+        best_solution_thresholded = best_solution_np
+    
+    # Format results into a dictionary
+    result_dict = {name: value for name, value in zip(feature_names, best_solution_thresholded)}
+    result_dict[target_name] = best_fitness.item()
+    
+    return result_dict, pandas_logger
+
+
 def _handle_pandas_log(logger: PandasLogger, save_path: Path, target_name: str):
     log_dataframe = logger.to_dataframe()
-    save_dataframe(df=log_dataframe, save_dir=save_path / "EvolutionLogs", filename=target_name)
+    save_dataframe_filename(df=log_dataframe, save_dir=save_path / "EvolutionLogs", filename=target_name)
 
 
 def info():

ml_tools/ML_scaler.py

@@ -2,14 +2,17 @@ import torch
 from torch.utils.data import Dataset, DataLoader
 from pathlib import Path
 from typing import Union, List, Optional
+
 from ._logger import _LOGGER
 from ._script_info import _script_info
 from .path_manager import make_fullpath
 
+
 __all__ = [
     "PytorchScaler"
 ]
 
+
 class PytorchScaler:
     """
     Standardizes continuous features in a PyTorch dataset by subtracting the
@@ -149,24 +152,25 @@ class PytorchScaler:
         
         return data_clone
 
-    def save(self, filepath: Union[str, Path]):
+    def save(self, filepath: Union[str, Path], verbose: bool=True):
         """
         Saves the scaler's state (mean, std, indices) to a .pth file.
 
         Args:
             filepath (str | Path): The path to save the file.
         """
-        path_obj = make_fullpath(filepath)
+        path_obj = make_fullpath(filepath, make=True, enforce="file")
         state = {
             'mean': self.mean_,
             'std': self.std_,
             'continuous_feature_indices': self.continuous_feature_indices
         }
         torch.save(state, path_obj)
-        _LOGGER.info(f"PytorchScaler state saved to '{path_obj.name}'.")
+        if verbose:
+            _LOGGER.info(f"PytorchScaler state saved as '{path_obj.name}'.")
 
     @staticmethod
-    def load(filepath: Union[str, Path]) -> 'PytorchScaler':
+    def load(filepath: Union[str, Path], verbose: bool=True) -> 'PytorchScaler':
         """
         Loads a scaler's state from a .pth file.
 
@@ -178,7 +182,8 @@ class PytorchScaler:
         """
         path_obj = make_fullpath(filepath, enforce="file")
         state = torch.load(path_obj)
-        _LOGGER.info(f"PytorchScaler state loaded from '{path_obj.name}'.")
+        if verbose:
+            _LOGGER.info(f"PytorchScaler state loaded from '{path_obj.name}'.")
         return PytorchScaler(
             mean=state['mean'],
             std=state['std'],