dragon-ml-toolbox 20.2.0__py3-none-any.whl → 20.4.0__py3-none-any.whl

ml_tools/ML_optimization/_single_dragon.py

@@ -1,9 +1,10 @@
 from typing import Literal, Union, Optional
 from pathlib import Path
 
-from ..optimization_tools import create_optimization_bounds
+from ..optimization_tools import create_optimization_bounds, load_continuous_bounds_template
 from ..ML_inference import DragonInferenceHandler
 from ..schema import FeatureSchema
+from ..ML_configuration import DragonOptimizerConfig
 
 from .._core import get_logger
 from ..keys._keys import MLTaskKeys
@@ -29,35 +30,28 @@ class DragonOptimizer:
     SNES and CEM algorithms do not accept bounds, the given bounds will be used as an initial starting point.
 
     Example:
-        >>> # 1. Define bounds for continuous features
-        >>> cont_bounds = {'feature_A': (0, 100), 'feature_B': (-10, 10)}
+        >>> # 1. Define configuration
+        >>> config = DragonOptimizerConfig(
+        ...     target_name="my_target",
+        ...     task="max",
+        ...     continuous_bounds_map="path/to/bounds",
+        ...     save_directory="/path/to/results",
+        ...     algorithm="Genetic"
+        ... )
         >>>
         >>> # 2. Initialize the optimizer
         >>> optimizer = DragonOptimizer(
         ...     inference_handler=my_handler,
         ...     schema=schema,
-        ...     target_name="my_target",
-        ...     continuous_bounds_map=cont_bounds,
-        ...     task="max",
-        ...     algorithm="Genetic",
+        ...     config=config
         ... )
         >>> # 3. Run the optimization
-        >>> best_result = optimizer.run(
-        ...     num_generations=100,
-        ...     save_dir="/path/to/results",
-        ...     save_format="csv"
-        ... )
+        >>> best_result = optimizer.run()
     """
     def __init__(self,
                  inference_handler: DragonInferenceHandler,
                  schema: FeatureSchema,
-                 target_name: str,
-                 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):
+                 config: DragonOptimizerConfig):
         """
         Initializes the optimizer by creating the EvoTorch problem and searcher.
 
@@ -65,45 +59,43 @@ class DragonOptimizer:
             inference_handler (DragonInferenceHandler): 
                 An initialized inference handler containing the model.
             schema (FeatureSchema): 
-                The definitive schema object from data_exploration.
-            target_name (str): 
-                target name to optimize.
-            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.
+                The definitive schema object.
+            config (DragonOptimizerConfig):
+                Configuration object containing optimization parameters.
         """
         # --- Store schema ---
         self.schema = schema
         # --- Store inference handler ---
         self.inference_handler = inference_handler
         
+        # --- Store config ---
+        self.config = config
+        
         # Ensure only Regression tasks are used
         allowed_tasks = [MLTaskKeys.REGRESSION, MLTaskKeys.MULTITARGET_REGRESSION]
         if self.inference_handler.task not in allowed_tasks:
             _LOGGER.error(f"DragonOptimizer only supports {allowed_tasks}. Got '{self.inference_handler.task}'.")
-            raise ValueError(f"Invalid Task: {self.inference_handler.task}")
+            raise ValueError()
         
         # --- store target name ---
-        self.target_name = target_name
+        self.target_name = config.target_name
         
         # --- flag to control single vs multi-target ---
         self.is_multi_target = False
         
         # --- 1. Create bounds from schema ---
-        # This is the robust way to get bounds
+        # Handle bounds loading if it's a path
+        raw_bounds_map = config.continuous_bounds_map
+        if isinstance(raw_bounds_map, (str, Path)):
+            continuous_bounds = load_continuous_bounds_template(raw_bounds_map)
+        else:
+            continuous_bounds = raw_bounds_map
+
+        # Robust way to get bounds
         bounds = create_optimization_bounds(
             schema=schema,
-            continuous_bounds_map=continuous_bounds_map,
-            start_at_zero=discretize_start_at_zero
+            continuous_bounds_map=continuous_bounds,
+            start_at_zero=config.discretize_start_at_zero
         )
         
         # Resolve target index if multi-target
@@ -114,26 +106,26 @@ class DragonOptimizer:
             _LOGGER.error("The provided inference handler does not have 'target_ids' defined.")
             raise ValueError()
 
-        if target_name not in self.inference_handler.target_ids:
-            _LOGGER.error(f"Target name '{target_name}' not found in the inference handler's 'target_ids': {self.inference_handler.target_ids}")
+        if self.target_name not in self.inference_handler.target_ids:
+            _LOGGER.error(f"Target name '{self.target_name}' not found in the inference handler's 'target_ids': {self.inference_handler.target_ids}")
             raise ValueError()
 
         if len(self.inference_handler.target_ids) == 1:
             # Single target regression
             target_index = None
-            _LOGGER.info(f"Optimization locked to single-target model '{target_name}'.")
+            _LOGGER.info(f"Optimization locked to single-target model '{self.target_name}'.")
         else:
             # Multi-target regression (optimizing one specific column)
-            target_index = self.inference_handler.target_ids.index(target_name)
+            target_index = self.inference_handler.target_ids.index(self.target_name)
             self.is_multi_target = True
-            _LOGGER.info(f"Optimization locked to target '{target_name}' (Index {target_index}) in a multi-target model.")
+            _LOGGER.info(f"Optimization locked to target '{self.target_name}' (Index {target_index}) in a multi-target model.")
         
         # --- 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,
+            discretize_start_at_zero=config.discretize_start_at_zero,
             target_index=target_index
         )
         
@@ -141,20 +133,13 @@ class DragonOptimizer:
         self.problem, self.searcher_factory = create_pytorch_problem(
             evaluator=self.evaluator,
             bounds=bounds,
-            task=task,
-            algorithm=algorithm,
-            population_size=population_size,
-            **searcher_kwargs
+            task=config.task, # type: ignore
+            algorithm=config.algorithm, # type: ignore
+            population_size=config.population_size,
+            **config.searcher_kwargs
         )
-        
-        # --- 4. Store other info needed by run() ---
-        self.discretize_start_at_zero = discretize_start_at_zero
 
     def run(self,
-            num_generations: int,
-            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.
@@ -163,15 +148,10 @@ class DragonOptimizer:
         provided during initialization.
 
         Args:
-            num_generations (int): The total number of generations for each repetition.
-            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.
+            Optional[dict]: A dictionary with the best result if repetitions is 1, otherwise None.
         """
         # Pass inference handler and target names for multi-target only
         if self.is_multi_target:
@@ -185,18 +165,18 @@ class DragonOptimizer:
         return run_optimization(
             problem=self.problem,
             searcher_factory=self.searcher_factory,
-            num_generations=num_generations,
+            num_generations=self.config.generations,
             target_name=self.target_name,
-            save_dir=save_dir,
-            save_format=save_format,
+            save_dir=self.config.save_directory,
+            save_format=self.config.save_format, # type: ignore
             # 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,
+            repetitions=self.config.repetitions,
             verbose=verbose,
-            discretize_start_at_zero=self.discretize_start_at_zero,
+            discretize_start_at_zero=self.config.discretize_start_at_zero,
             all_target_names=target_names_to_pass,
             inference_handler=inference_handler_to_pass
         )

ml_tools/ML_optimization/_single_manual.py

@@ -506,5 +506,5 @@ def _save_result(
 
 def _handle_pandas_log(logger: PandasLogger, save_path: Path, target_name: str):
     log_dataframe = logger.to_dataframe()
-    save_dataframe_filename(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, verbose=2)
 

ml_tools/ML_scaler/_ML_scaler.py

@@ -33,7 +33,7 @@ class DragonScaler:
         self.continuous_feature_indices = continuous_feature_indices
 
     @classmethod
-    def fit(cls, dataset: Dataset, continuous_feature_indices: list[int], batch_size: int = 64) -> 'DragonScaler':
+    def fit(cls, dataset: Dataset, continuous_feature_indices: list[int], batch_size: int = 64, verbose: int = 3) -> 'DragonScaler':
         """
         Fits the scaler using a PyTorch Dataset (Method A) using Batched Welford's Algorithm.
         """
@@ -85,23 +85,25 @@ class DragonScaler:
                 n_total = new_n_total
 
         if n_total == 0:
-             _LOGGER.error("Dataset is empty. Scaler cannot be fitted.")
-             return cls(continuous_feature_indices=continuous_feature_indices)
+            _LOGGER.error("Dataset is empty. Scaler cannot be fitted.")
+            return cls(continuous_feature_indices=continuous_feature_indices)
 
         # Finalize Standard Deviation
         # Unbiased estimator (divide by n-1)
         if n_total < 2:
-            _LOGGER.warning(f"Only one sample found. Standard deviation set to 1.")
+            if verbose >= 1:
+                _LOGGER.warning(f"Only one sample found. Standard deviation set to 1.")
             std = torch.ones_like(mean_global) # type: ignore
         else:
             variance = m2_global / (n_total - 1)
             std = torch.sqrt(torch.clamp(variance, min=1e-8))
-
-        _LOGGER.info(f"Scaler fitted on {n_total} samples for {num_continuous_features} features (Welford's).")
+        
+        if verbose >= 2:
+            _LOGGER.info(f"Scaler fitted on {n_total} samples for {num_continuous_features} features (Welford's).")
         return cls(mean=mean_global, std=std, continuous_feature_indices=continuous_feature_indices)
 
     @classmethod
-    def fit_tensor(cls, data: torch.Tensor) -> 'DragonScaler':
+    def fit_tensor(cls, data: torch.Tensor, verbose: int = 3) -> 'DragonScaler':
         """
         Fits the scaler directly on a Tensor (Method B).
         Useful for targets or small datasets already in memory.
@@ -118,6 +120,9 @@ class DragonScaler:
         # Handle constant values (std=0) to prevent division by zero
         std = torch.where(std == 0, torch.tensor(1.0, device=data.device), std)
         
+        if verbose >= 2:
+            _LOGGER.info(f"Scaler fitted on tensor with {data.shape[0]} samples for {num_features} features.")
+        
         return cls(mean=mean, std=std, continuous_feature_indices=indices)
 
     def transform(self, data: torch.Tensor) -> torch.Tensor:

ml_tools/ML_scaler/__init__.py

@@ -2,9 +2,13 @@ from ._ML_scaler import (
     DragonScaler
 )
 
-from ._imprimir import info
+from .._core import _imprimir_disponibles
 
 
 __all__ = [
     "DragonScaler"
 ]
+
+
+def info():
+    _imprimir_disponibles(__all__)

ml_tools/ML_trainer/__init__.py

@@ -10,7 +10,7 @@ from ._dragon_detection_trainer import (
     DragonDetectionTrainer
 )
 
-from ._imprimir import info
+from .._core import _imprimir_disponibles
 
 
 __all__ = [
@@ -18,3 +18,7 @@ __all__ = [
     "DragonSequenceTrainer",
     "DragonDetectionTrainer",
 ]
+
+
+def info():
+    _imprimir_disponibles(__all__)

ml_tools/ML_trainer/_base_trainer.py

@@ -1,6 +1,6 @@
 from typing import Literal, Union, Optional, Any
 from pathlib import Path
-from torch.utils.data import DataLoader
+from torch.utils.data import DataLoader, Dataset
 import torch
 from torch import nn
 from abc import ABC, abstractmethod
@@ -10,6 +10,7 @@ from ..ML_callbacks._checkpoint import DragonModelCheckpoint
 from ..ML_callbacks._early_stop import _DragonEarlyStopping
 from ..ML_callbacks._scheduler import _DragonLRScheduler
 from ..ML_evaluation import plot_losses
+from ..ML_utilities import inspect_pth_file
 
 from ..path_manager import make_fullpath
 from ..keys._keys import PyTorchCheckpointKeys, MagicWords
@@ -89,11 +90,128 @@ class _BaseDragonTrainer(ABC):
         """Gives each callback a reference to this trainer instance."""
         for callback in self.callbacks:
             callback.set_trainer(self)
+    
+    def _make_dataloaders(self, 
+                          train_dataset: Any, 
+                          validation_dataset: Any, 
+                          batch_size: int, 
+                          shuffle: bool,
+                          collate_fn: Optional[Any] = None):
+        """
+        Shared logic to initialize standard DataLoaders. 
+        Subclasses can call this inside their _create_dataloaders implementation.
+        """
+        # Ensure stability on MPS devices by setting num_workers to 0
+        loader_workers = 0 if self.device.type == 'mps' else self.dataloader_workers
+        pin_memory = ("cuda" in self.device.type)
+        
+        self.train_loader = DataLoader(
+            dataset=train_dataset, 
+            batch_size=batch_size, 
+            shuffle=shuffle, 
+            num_workers=loader_workers, 
+            pin_memory=pin_memory,
+            drop_last=True,
+            collate_fn=collate_fn
+        )
+        
+        self.validation_loader = DataLoader(
+            dataset=validation_dataset, 
+            batch_size=batch_size, 
+            shuffle=False, 
+            num_workers=loader_workers, 
+            pin_memory=pin_memory,
+            collate_fn=collate_fn
+        )
+
+    def _validate_checkpoint_arg(self, model_checkpoint: Union[Path, str]) -> Union[Path, str]:
+        """Validates the model_checkpoint argument."""
+        if isinstance(model_checkpoint, Path):
+            return make_fullpath(model_checkpoint, enforce="file")
+        elif model_checkpoint in [MagicWords.BEST, MagicWords.CURRENT]:
+            return model_checkpoint
+        else:
+            _LOGGER.error(f"'model_checkpoint' must be a Path object, or the string '{MagicWords.BEST}', or the string '{MagicWords.CURRENT}'.")
+            raise ValueError()
+
+    def _validate_save_dir(self, save_dir: Union[str, Path]) -> Path:
+        """Validates and creates the save directory."""
+        return make_fullpath(save_dir, make=True, enforce="directory")
+
+    def _prepare_eval_data(self, 
+                           data: Optional[Union[DataLoader, Dataset]], 
+                           default_dataset: Optional[Dataset], 
+                           collate_fn: Optional[Any] = None) -> tuple[DataLoader, Any]:
+        """
+        Prepares the DataLoader and dataset artifact source for evaluation.
+        
+        Returns:
+            (eval_loader, dataset_for_artifacts)
+        """
+        eval_loader = None
+        dataset_for_artifacts = None
+        
+        # Loader workers config
+        loader_workers = 0 if self.device.type == 'mps' else self.dataloader_workers
+        pin_memory = (self.device.type == "cuda")
 
-    def _load_checkpoint(self, path: Union[str, Path]):
+        if isinstance(data, DataLoader):
+            eval_loader = data
+            if hasattr(data, 'dataset'):
+                dataset_for_artifacts = data.dataset
+        elif isinstance(data, Dataset):
+            eval_loader = DataLoader(data, 
+                                     batch_size=self._batch_size, 
+                                     shuffle=False, 
+                                     num_workers=loader_workers,
+                                     pin_memory=pin_memory,
+                                     collate_fn=collate_fn)
+            dataset_for_artifacts = data
+        else: # data is None
+            if default_dataset is None:
+                _LOGGER.error("Cannot evaluate. No data provided and no validation dataset available in the trainer.")
+                raise ValueError()
+            
+            eval_loader = DataLoader(default_dataset, 
+                                     batch_size=self._batch_size, 
+                                     shuffle=False, 
+                                     num_workers=loader_workers,
+                                     pin_memory=pin_memory,
+                                     collate_fn=collate_fn)
+            dataset_for_artifacts = default_dataset
+
+        if eval_loader is None:
+            _LOGGER.error("Cannot evaluate. No valid data was provided or found.")
+            raise ValueError()
+            
+        return eval_loader, dataset_for_artifacts
+
+    def _save_finalized_artifact(self, 
+                                 finalized_data: dict, 
+                                 save_dir: Union[str, Path], 
+                                 filename: str):
+        """
+        Handles the common logic for saving the finalized model dictionary to disk.
+        """
+        # handle save path
+        dir_path = self._validate_save_dir(save_dir)
+        full_path = dir_path / filename
+        
+        # checkpoint loading happens before dict creation. 
+        
+        torch.save(finalized_data, full_path)
+        
+        _LOGGER.info(f"Finalized model file saved to '{full_path}'")
+        
+        if full_path.is_file():
+            inspect_pth_file(pth_path=full_path, save_dir=dir_path, verbose=2)
+    
+    def _load_checkpoint(self, path: Union[str, Path], verbose: int = 3):
         """Loads a training checkpoint to resume training."""
         p = make_fullpath(path, enforce="file")
-        _LOGGER.info(f"Loading checkpoint from '{p.name}'...")
+        
+        if verbose >= 2:
+            _LOGGER.info(f"Loading checkpoint from '{p.name}'...")
         
         try:
             checkpoint = torch.load(p, map_location=self.device)
@@ -110,9 +228,11 @@ class _BaseDragonTrainer(ABC):
             # --- Load History ---
             if PyTorchCheckpointKeys.HISTORY in checkpoint:
                 self.history = checkpoint[PyTorchCheckpointKeys.HISTORY]
-                _LOGGER.info(f"Restored training history up to epoch {self.epoch}.")
+                if verbose >= 3:
+                    _LOGGER.info(f"Restored training history up to epoch {self.epoch}.")
             else:
-                _LOGGER.warning("No 'history' found in checkpoint. A new history will be started.")
+                if verbose >= 1:
+                    _LOGGER.warning("No 'history' found in checkpoint. A new history will be started.")
                 self.history = {} # Ensure it's at least an empty dict
             
             # --- Scheduler State Loading Logic ---
@@ -124,7 +244,8 @@ class _BaseDragonTrainer(ABC):
                 try:
                     self.scheduler.load_state_dict(checkpoint[PyTorchCheckpointKeys.SCHEDULER_STATE]) # type: ignore
                     scheduler_name = self.scheduler.__class__.__name__
-                    _LOGGER.info(f"Restored LR scheduler state for: {scheduler_name}")
+                    if verbose >= 3:
+                        _LOGGER.info(f"Restored LR scheduler state for: {scheduler_name}")
                 except Exception as e:
                     # Loading failed, likely a mismatch
                     scheduler_name = self.scheduler.__class__.__name__
@@ -134,7 +255,8 @@ class _BaseDragonTrainer(ABC):
             elif scheduler_object_exists and not scheduler_state_exists:
                 # Case 2: Scheduler provided, but no state in checkpoint.
                 scheduler_name = self.scheduler.__class__.__name__
-                _LOGGER.warning(f"'{scheduler_name}' was provided, but no scheduler state was found in the checkpoint. The scheduler will start from its initial state.")
+                if verbose >= 1:
+                    _LOGGER.warning(f"'{scheduler_name}' was provided, but no scheduler state was found in the checkpoint. The scheduler will start from its initial state.")
             
             elif not scheduler_object_exists and scheduler_state_exists:
                 # Case 3: State in checkpoint, but no scheduler provided.
@@ -145,9 +267,11 @@ class _BaseDragonTrainer(ABC):
             for cb in self.callbacks:
                 if isinstance(cb, DragonModelCheckpoint) and PyTorchCheckpointKeys.BEST_SCORE in checkpoint:
                     cb.best = checkpoint[PyTorchCheckpointKeys.BEST_SCORE]
-                    _LOGGER.info(f"Restored {cb.__class__.__name__} 'best' score to: {cb.best:.4f}")
+                    if verbose >= 3:
+                        _LOGGER.info(f"Restored {cb.__class__.__name__} 'best' score to: {cb.best:.4f}")
             
-            _LOGGER.info(f"Checkpoint loaded. Resuming training from epoch {self.start_epoch}.")
+            if verbose >= 2:
+                _LOGGER.info(f"Model restored to epoch {self.epoch}.")
             
         except Exception as e:
             _LOGGER.error(f"Failed to load checkpoint from '{p}': {e}")
@@ -243,16 +367,15 @@ class _BaseDragonTrainer(ABC):
         self.model.to(self.device)
         _LOGGER.info(f"Trainer and model moved to {self.device}.")
         
-    def _load_model_state_for_finalizing(self, model_checkpoint: Union[Path, Literal['best', 'current']]):
+    def _load_model_state_wrapper(self, model_checkpoint: Union[Path, Literal['best', 'current']], verbose: int = 2):
         """
         Private helper to load the correct model state_dict based on user's choice.
-        This is called by finalize_model_training() in subclasses.
         """
         if isinstance(model_checkpoint, Path):
-            self._load_checkpoint(path=model_checkpoint)
+            self._load_checkpoint(path=model_checkpoint, verbose=verbose)
         elif model_checkpoint == MagicWords.BEST and self._checkpoint_callback:
             path_to_latest = self._checkpoint_callback.best_checkpoint_path
-            self._load_checkpoint(path_to_latest)
+            self._load_checkpoint(path_to_latest, verbose=verbose)
         elif model_checkpoint == MagicWords.BEST and self._checkpoint_callback is None:
             _LOGGER.error(f"'model_checkpoint' set to '{MagicWords.BEST}' but no checkpoint callback was found.")
             raise ValueError()