dragon-ml-toolbox 13.0.0__py3-none-any.whl → 13.1.0__py3-none-any.whl

ml_tools/ML_optimization.py

@@ -17,9 +17,10 @@ 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 .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__ = [
@@ -40,66 +41,76 @@ class MLOptimizer:
     SNES and CEM algorithms do not accept bounds, the given bounds will be used as an initial starting point.
 
     Example:
-        >>> # 1. Get categorical info from preprocessing steps
-        >>> # e.g., from data_exploration.encode_categorical_features
-        >>> cat_mappings = {'feature_C': {'A': 0, 'B': 1}, 'feature_D': {'X': 0, 'Y': 1}}
-        >>> # e.g., from data_exploration.create_transformer_categorical_map
-        >>> # Assumes feature_C is at index 2 (cardinality 2) and feature_D is at index 3 (cardinality 2)
-        >>> cat_index_map = {2: 2, 3: 2}
+        >>> # 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)}
         >>>
-        >>> # 2. Initialize the optimizer
+        >>> # 3. Initialize the optimizer
         >>> optimizer = MLOptimizer(
         ...     inference_handler=my_handler,
-        ...     bounds=(lower_bounds, upper_bounds), # Bounds for ALL features
+        ...     schema=schema,
+        ...     continuous_bounds_map=cont_bounds,
         ...     task="max",
         ...     algorithm="Genetic",
-        ...     categorical_index_map=cat_index_map,
-        ...     categorical_mappings=cat_mappings,
         ... )
-        >>> # 3. Run the optimization
+        >>> # 4. Run the optimization
         >>> best_result = optimizer.run(
         ...     num_generations=100,
         ...     target_name="my_target",
-        ...     feature_names=my_feature_names,
         ...     save_dir="/path/to/results",
         ...     save_format="csv"
         ... )
     """
     def __init__(self,
                  inference_handler: PyTorchInferenceHandler,
-                 bounds: Tuple[List[float], List[float]],
+                 schema: FeatureSchema,
+                 continuous_bounds_map: Dict[str, Tuple[float, float]],
                  task: Literal["min", "max"],
                  algorithm: Literal["SNES", "CEM", "Genetic"] = "Genetic",
                  population_size: int = 200,
-                 categorical_index_map: Optional[Dict[int, int]] = None,
-                 categorical_mappings: Optional[Dict[str, Dict[str, int]]] = None,
                  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 and weights.
-            bounds (tuple[list[float], list[float]]): A tuple containing the lower and upper bounds for ALL solution features. 
-                Use the `optimization_tools.create_optimization_bounds()` helper to easily generate this and ensure unbiased categorical bounds.
+            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.
-            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]).
-            **searcher_kwargs: Additional keyword arguments for the selected search algorithm's constructor.
+            **searcher_kwargs: Additional keyword arguments for the selected 
+                               search algorithm's constructor.
         """
-        # Make a fitness function
+        # --- 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,
-            categorical_index_map=categorical_index_map,
+            # Get categorical info from the schema
+            categorical_index_map=schema.categorical_index_map,
             discretize_start_at_zero=discretize_start_at_zero
         )
         
-        # Call the existing factory function to get the problem and searcher factory
+        # --- 3. Create the problem and searcher factory ---
         self.problem, self.searcher_factory = create_pytorch_problem(
             evaluator=self.evaluator,
             bounds=bounds,
@@ -108,36 +119,36 @@ class MLOptimizer:
             population_size=population_size,
             **searcher_kwargs
         )
-        # Store categorical info to pass to the run function
-        self.categorical_map = categorical_index_map
-        self.categorical_mappings = categorical_mappings
+        
+        # --- 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],
-            feature_names: Optional[List[str]],
             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.
-            feature_names (List[str] | None): Names of the solution features for labeling output.
-                If None, generic names like 'feature_0', 'feature_1', ... , will be created.
             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.
         """
-        # Call the existing run function with the stored problem, searcher, and categorical info
+        # Call the existing run function, passing info from the schema
         return run_optimization(
             problem=self.problem,
             searcher_factory=self.searcher_factory,
@@ -145,11 +156,13 @@ class MLOptimizer:
             target_name=target_name,
             save_dir=save_dir,
             save_format=save_format,
-            feature_names=feature_names,
+            # 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,
-            categorical_map=self.categorical_map,
-            categorical_mappings=self.categorical_mappings,
             discretize_start_at_zero=self.discretize_start_at_zero
         )
         

ml_tools/PSO_optimization.py

@@ -17,6 +17,10 @@ from ._script_info import _script_info
 from .SQL import DatabaseManager
 from .optimization_tools import _save_result
 
+"""
+DEPRECATED
+"""
+
 
 __all__ = [
     "ObjectiveFunction",
@@ -46,7 +50,7 @@ class ObjectiveFunction():
         self.binary_features = binary_features
         self.is_hybrid = False if binary_features <= 0 else True
         self.use_noise = add_noise
-        self._artifact = deserialize_object(trained_model_path, verbose=False, raise_on_error=True)
+        self._artifact = deserialize_object(trained_model_path, verbose=False)
         self.model = self._get_from_artifact(EnsembleKeys.MODEL)
         self.feature_names: Optional[list[str]] = self._get_from_artifact(EnsembleKeys.FEATURES) # type: ignore
         self.target_name: Optional[str] = self._get_from_artifact(EnsembleKeys.TARGET) # type: ignore

ml_tools/_schema.py

@@ -0,0 +1,19 @@
+from typing import NamedTuple, Tuple, Optional, Dict
+
+class FeatureSchema(NamedTuple):
+    """Holds the final, definitive schema for the model pipeline."""
+    
+    # The final, ordered list of all feature names
+    feature_names: Tuple[str, ...]
+    
+    # List of all continuous feature names
+    continuous_feature_names: Tuple[str, ...]
+    
+    # List of all categorical feature names
+    categorical_feature_names: Tuple[str, ...]
+    
+    # Map of {column_index: cardinality} for categorical features
+    categorical_index_map: Optional[Dict[int, int]]
+    
+    # The original string-to-int mappings (e.g., {'color': {'red': 0, 'blue': 1}})
+    categorical_mappings: Optional[Dict[str, Dict[str, int]]]

ml_tools/data_exploration.py

@@ -11,7 +11,7 @@ from .path_manager import sanitize_filename, make_fullpath
 from ._script_info import _script_info
 from ._logger import _LOGGER
 from .utilities import save_dataframe_filename
-
+from ._schema import FeatureSchema
 
 # Keep track of all available tools, show using `info()`
 __all__ = [
@@ -32,9 +32,9 @@ __all__ = [
     "drop_outlier_samples",
     "match_and_filter_columns_by_regex",
     "standardize_percentages",
-    "create_transformer_categorical_map",
     "reconstruct_one_hot",
-    "reconstruct_binary"
+    "reconstruct_binary",
+    "finalize_feature_schema"
 ]
 
 
@@ -977,49 +977,6 @@ def standardize_percentages(
     return df_copy
 
 
-def create_transformer_categorical_map(
-    df: pd.DataFrame,
-    mappings: Dict[str, Dict[str, int]],
-    verbose: bool = True
-) -> Dict[int, int]:
-    """
-    Creates the `categorical_map` required by a `TabularTransformer` model.
-
-    This function should be called late in the preprocessing pipeline, after all
-    column additions, deletions, or reordering have occurred. It uses the final
-    DataFrame's column order to map the correct column index to its cardinality.
-
-    Args:
-        df (pd.DataFrame): The final, processed DataFrame.
-        mappings (Dict[str, Dict[str, int]]): The mappings dictionary generated by
-          `encode_categorical_features`, containing the category-to-integer
-          mapping for each categorical column.
-        verbose (bool): If True, prints mapping progress.
-
-    Returns:
-        (Dict[int, int]): The final `categorical_map` for the transformer,
-          mapping each column's current index to its cardinality (e.g., {0: 3}).
-    """
-    transformer_map = {}
-    categorical_column_names = mappings.keys()
-
-    _LOGGER.info("Creating categorical map for TabularTransformer.")
-    for col_name in categorical_column_names:
-        if col_name in df.columns:
-            col_idx = df.columns.get_loc(col_name)
-            
-            # Get cardinality directly from the length of the mapping dictionary
-            cardinality = len(mappings[col_name])
-            
-            transformer_map[col_idx] = cardinality
-            if verbose:
-                print(f"  - Mapping column '{col_name}' at index {col_idx} with cardinality {cardinality}.")
-        else:
-            _LOGGER.warning(f"Categorical column '{col_name}' not found in the final DataFrame. Skipping.")
-            
-    return transformer_map
-
-
 def reconstruct_one_hot(
     df: pd.DataFrame,
     features_to_reconstruct: List[Union[str, Tuple[str, Optional[str]]]],
@@ -1274,6 +1231,78 @@ def reconstruct_binary(
     return new_df
 
 
+def finalize_feature_schema(
+    df_features: pd.DataFrame,
+    categorical_mappings: Optional[Dict[str, Dict[str, int]]]
+) -> FeatureSchema:
+    """
+    Analyzes the final features DataFrame to create a definitive schema.
+
+    This function is the "single source of truth" for column order
+    and type (categorical vs. continuous) for the entire ML pipeline.
+
+    It should be called at the end of the feature engineering process.
+
+    Args:
+        df_features (pd.DataFrame):
+            The final, processed DataFrame containing *only* feature columns
+            in the exact order they will be fed to the model.
+        categorical_mappings (Dict[str, Dict[str, int]] | None):
+            The mappings dictionary generated by
+            `encode_categorical_features`. Can be None if no
+            categorical features exist.
+
+    Returns:
+        FeatureSchema: A NamedTuple containing all necessary metadata for the pipeline.
+    """
+    feature_names: List[str] = df_features.columns.to_list()
+    
+    # Intermediate lists for building
+    continuous_feature_names_list: List[str] = []
+    categorical_feature_names_list: List[str] = []
+    categorical_index_map_dict: Dict[int, int] = {}
+
+    _LOGGER.info("Finalizing feature schema...")
+
+    if categorical_mappings:
+        # --- Categorical features are present ---
+        categorical_names_set = set(categorical_mappings.keys())
+        
+        for index, name in enumerate(feature_names):
+            if name in categorical_names_set:
+                # This is a categorical feature
+                cardinality = len(categorical_mappings[name])
+                categorical_index_map_dict[index] = cardinality
+                categorical_feature_names_list.append(name)
+            else:
+                # This is a continuous feature
+                continuous_feature_names_list.append(name)
+        
+        # Use the populated dict, or None if it's empty
+        final_index_map = categorical_index_map_dict if categorical_index_map_dict else None
+    
+    else:
+        # --- No categorical features ---
+        _LOGGER.info("No categorical mappings provided. Treating all features as continuous.")
+        continuous_feature_names_list = list(feature_names)
+        # categorical_feature_names_list remains empty
+        # categorical_index_map_dict remains empty
+        final_index_map = None # Explicitly set to None to match Optional type
+
+    _LOGGER.info(f"Schema created: {len(continuous_feature_names_list)} continuous, {len(categorical_feature_names_list)} categorical.")
+    
+    # Create the final immutable instance
+    schema_instance = FeatureSchema(
+        feature_names=tuple(feature_names),
+        continuous_feature_names=tuple(continuous_feature_names_list),
+        categorical_feature_names=tuple(categorical_feature_names_list),
+        categorical_index_map=final_index_map,
+        categorical_mappings=categorical_mappings
+    )
+    
+    return schema_instance
+
+
 def _validate_columns(df: pd.DataFrame, columns: list[str]):
     valid_columns = [column for column in columns if column in df.columns]
     return valid_columns

ml_tools/optimization_tools.py

@@ -9,6 +9,7 @@ from .utilities import yield_dataframes_from_dir
 from ._logger import _LOGGER
 from ._script_info import _script_info
 from .SQL import DatabaseManager
+from ._schema import FeatureSchema
 
 
 __all__ = [
@@ -19,35 +20,25 @@ __all__ = [
 
 
 def create_optimization_bounds(
-    csv_path: Union[str, Path],
+    schema: FeatureSchema,
     continuous_bounds_map: Dict[str, Tuple[float, float]],
-    categorical_map: Dict[int, int],
-    target_column: Optional[str] = None,
     start_at_zero: bool = True
 ) -> Tuple[List[float], List[float]]:
     """
-    Generates the lower and upper bounds lists for the optimizer from a CSV header.
+    Generates the lower and upper bounds lists for the optimizer from a FeatureSchema.
 
     This helper function automates the creation of unbiased bounds for
     categorical features and combines them with user-defined bounds for
-    continuous features.
-
-    It reads *only* the header of the provided CSV to determine the full
-    list of feature columns and their order, excluding the specified target.
-    This is memory-efficient as the full dataset is not loaded.
+    continuous features, using the schema as the single source of truth
+    for feature order and type.
 
     Args:
-        csv_path (Union[str, Path]):
-            Path to the final, preprocessed CSV file. The column order in
-            this file must match the order expected by the model.
+        schema (FeatureSchema):
+            The definitive schema object created by 
+            `data_exploration.finalize_feature_schema()`.
         continuous_bounds_map (Dict[str, Tuple[float, float]]):
             A dictionary mapping the *name* of each **continuous** feature
             to its (min_bound, max_bound) tuple.
-        categorical_map (Dict[int, int]):
-            The map from the *index* of each **categorical** feature to its cardinality.
-            (e.g., {2: 4} for a feature at index 2 with 4 categories).
-        target_column (Optional[str], optional):
-            The name of the target column to exclude. If None (default), the *last column* in the CSV is assumed to be the target.
         start_at_zero (bool):
             - If True, assumes categorical encoding is [0, 1, ..., k-1].
               Bounds will be set as [-0.5, k - 0.5].
@@ -59,98 +50,86 @@ def create_optimization_bounds(
             A tuple containing two lists: (lower_bounds, upper_bounds).
 
     Raises:
-        ValueError: If a feature is defined in both maps, is missing from
-                    both maps, or if a name in `continuous_bounds_map`
-                    or `target_column` is not found in the CSV columns.
+        ValueError: If a feature is missing from `continuous_bounds_map`
+                    or if a feature name in the map is not a
+                    continuous feature according to the schema.
     """
-    # 1. Read header and determine feature names
-    full_csv_path = make_fullpath(csv_path, enforce="file")
-    try:
-        df_header = pd.read_csv(full_csv_path, nrows=0, encoding="utf-8")
-    except Exception as e:
-        _LOGGER.error(f"Failed to read header from CSV: {e}")
-        raise
-        
-    all_column_names = df_header.columns.to_list()
-    feature_names: List[str] = []
-
-    if target_column is None:
-        feature_names = all_column_names[:-1]
-        excluded_target = all_column_names[-1]
-        _LOGGER.info(f"No target_column provided. Assuming last column '{excluded_target}' is the target.")
-    else:
-        if target_column not in all_column_names:
-            _LOGGER.error(f"Target column '{target_column}' not found in CSV header.")
-            raise ValueError()
-        feature_names = [name for name in all_column_names if name != target_column]
-        _LOGGER.info(f"Excluding target column '{target_column}'.")
-    
-    # 2. Initialize bound lists
+    # 1. Get feature names and map from schema
+    feature_names = schema.feature_names
+    categorical_index_map = schema.categorical_index_map
     total_features = len(feature_names)
+
     if total_features <= 0:
-        _LOGGER.error("No feature columns remain after excluding the target.")
+        _LOGGER.error("Schema contains no features.")
         raise ValueError()
+        
+    _LOGGER.info(f"Generating bounds for {total_features} total features...")
 
+    # 2. Initialize bound lists
     lower_bounds: List[Optional[float]] = [None] * total_features
     upper_bounds: List[Optional[float]] = [None] * total_features
-    
-    _LOGGER.info(f"Generating bounds for {total_features} total features...")
 
     # 3. Populate categorical bounds (Index-based)
-    # The indices in categorical_map (e.g., {2: 4}) directly correspond
-    # to the indices in the `feature_names` list.
-    for index, cardinality in categorical_map.items():
-        if not (0 <= index < total_features):
-            _LOGGER.error(f"Categorical index {index} is out of range for the {total_features} features.")
-            raise ValueError()
-            
-        if start_at_zero:
-            # Rule for [0, k-1]: bounds are [-0.5, k - 0.5]
-            low = -0.5
-            high = float(cardinality) - 0.5
-        else:
-            # Rule for [1, k]: bounds are [0.5, k + 0.5]
-            low = 0.5
-            high = float(cardinality) + 0.5
-            
-        lower_bounds[index] = low
-        upper_bounds[index] = high
+    if categorical_index_map:
+        for index, cardinality in categorical_index_map.items():
+            if not (0 <= index < total_features):
+                _LOGGER.error(f"Categorical index {index} is out of range for the {total_features} features.")
+                raise ValueError()
+                
+            if start_at_zero:
+                # Rule for [0, k-1]: bounds are [-0.5, k - 0.5]
+                low = -0.5
+                high = float(cardinality) - 0.5
+            else:
+                # Rule for [1, k]: bounds are [0.5, k + 0.5]
+                low = 0.5
+                high = float(cardinality) + 0.5
+                
+            lower_bounds[index] = low
+            upper_bounds[index] = high
         
-    _LOGGER.info(f"Automatically set bounds for {len(categorical_map)} categorical features.")
+        _LOGGER.info(f"Automatically set bounds for {len(categorical_index_map)} categorical features.")
+    else:
+        _LOGGER.info("No categorical features found in schema.")
 
     # 4. Populate continuous bounds (Name-based)
+    # Use schema.continuous_feature_names for robust checking
+    continuous_names_set = set(schema.continuous_feature_names)
+    
+    if continuous_names_set != set(continuous_bounds_map.keys()):
+        missing_in_map = continuous_names_set - set(continuous_bounds_map.keys())
+        if missing_in_map:
+            _LOGGER.error(f"The following continuous features are missing from 'continuous_bounds_map': {list(missing_in_map)}")
+        
+        extra_in_map = set(continuous_bounds_map.keys()) - continuous_names_set
+        if extra_in_map:
+            _LOGGER.error(f"The following features in 'continuous_bounds_map' are not defined as continuous in the schema: {list(extra_in_map)}")
+            
+        raise ValueError("Mismatch between 'continuous_bounds_map' and schema's continuous features.")
+
     count_continuous = 0
     for name, (low, high) in continuous_bounds_map.items():
-        try:
-            # Map name to its index in the *feature-only* list
-            index = feature_names.index(name)
-        except ValueError:
-            _LOGGER.warning(f"Feature name '{name}' from 'continuous_bounds_map' not found in the CSV's feature columns.")
-            continue
-            
+        # Map name to its index in the *feature-only* list
+        # This is guaranteed to be correct by the schema
+        index = feature_names.index(name)
+
         if lower_bounds[index] is not None:
-            # This index was already set by the categorical map
-            _LOGGER.error(f"Feature '{name}' (at index {index}) is defined in both 'categorical_map' and 'continuous_bounds_map'.")
+            # This should be impossible if schema is correct, but good to check
+            _LOGGER.error(f"Schema conflict: Feature '{name}' (at index {index}) is defined as both continuous and categorical.")
             raise ValueError()
-            
+
         lower_bounds[index] = float(low)
         upper_bounds[index] = float(high)
         count_continuous += 1
         
     _LOGGER.info(f"Manually set bounds for {count_continuous} continuous features.")
 
-    # 5. Validation: Check for any remaining None values
-    missing_indices = []
-    for i in range(total_features):
-        if lower_bounds[i] is None:
-            missing_indices.append(i)
-            
-    if missing_indices:
+    # 5. Final Validation (all Nones should be filled)
+    if None in lower_bounds:
+        missing_indices = [i for i, b in enumerate(lower_bounds) if b is None]
         missing_names = [feature_names[i] for i in missing_indices]
-        _LOGGER.error(f"Bounds not defined for all features. Missing: {missing_names}")
-        raise ValueError()
-
-    # _LOGGER.info("All bounds successfully created.")
+        _LOGGER.error(f"Failed to create all bounds. This indicates an internal logic error. Missing: {missing_names}")
+        raise RuntimeError("Internal error: Not all bounds were populated.")
     
     # Cast to float lists, as 'None' sentinels are gone
     return (

ml_tools/serde.py

@@ -116,8 +116,7 @@ def deserialize_object(
             # Can't do an isinstance check on 'Any', skip it.
             if type_to_check is not Any and not isinstance(obj, type_to_check):
                 error_msg = (
-                    f"Type mismatch: Expected an instance of '{expected_type}', "
-                    f"but found '{type(obj)}' in '{true_filepath}'."
+                    f"Type mismatch: Expected an instance of '{expected_type}', but found '{type(obj)}' in '{true_filepath}'."
                 )
                 _LOGGER.error(error_msg)
                 raise TypeError()