dragon-ml-toolbox 12.13.0__py3-none-any.whl → 14.3.0__py3-none-any.whl

{dragon_ml_toolbox-12.13.0.dist-info → dragon_ml_toolbox-14.3.0.dist-info}/METADATA

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: dragon-ml-toolbox
-Version: 12.13.0
+Version: 14.3.0
 Summary: A collection of tools for data science and machine learning projects.
 Author-email: "Karl L. Loza Vidaurre" <luigiloza@gmail.com>
 License-Expression: MIT
@@ -34,6 +34,10 @@ Requires-Dist: Pillow; extra == "ml"
 Requires-Dist: evotorch; extra == "ml"
 Requires-Dist: pyarrow; extra == "ml"
 Requires-Dist: colorlog; extra == "ml"
+Requires-Dist: torchmetrics; extra == "ml"
+Provides-Extra: py-tab
+Requires-Dist: pytorch_tabular; extra == "py-tab"
+Requires-Dist: omegaconf; extra == "py-tab"
 Provides-Extra: mice
 Requires-Dist: numpy<2.0; extra == "mice"
 Requires-Dist: pandas; extra == "mice"
@@ -142,10 +146,16 @@ ML_evaluation_multi
 ML_evaluation
 ML_inference
 ML_models
+ML_models_advanced # Requires the extra flag [py-tab]
 ML_optimization
 ML_scaler
 ML_trainer
 ML_utilities
+ML_vision_datasetmaster
+ML_vision_evaluation
+ML_vision_inference
+ML_vision_models
+ML_vision_transformers
 optimization_tools
 path_manager
 PSO_optimization
@@ -191,7 +201,6 @@ pip install "dragon-ml-toolbox[excel]"
 #### Modules:
 
 ```Bash
-constants
 custom_logger
 handle_excel
 path_manager

dragon_ml_toolbox-14.3.0.dist-info/RECORD

@@ -0,0 +1,48 @@
+dragon_ml_toolbox-14.3.0.dist-info/licenses/LICENSE,sha256=L35WDmmLZNTlJvxF6Vy7Uy4SYNi6rCfWUqlTHpoRMoU,1081
+dragon_ml_toolbox-14.3.0.dist-info/licenses/LICENSE-THIRD-PARTY.md,sha256=gkOdNDbKYpIJezwSo2CEnISkLeYfYHv9t8b5K2-P69A,2687
+ml_tools/ETL_cleaning.py,sha256=2VBRllV8F-ZiPylPp8Az2gwn5ztgazN0BH5OKnRUhV0,20402
+ml_tools/ETL_engineering.py,sha256=KfYqgsxupAx6e_TxwO1LZXeu5mFkIhVXJrNjP3CzIZc,54927
+ml_tools/GUI_tools.py,sha256=Va6ig-dHULPVRwQYYtH3fvY5XPIoqRcJpRW8oXC55Hw,45413
+ml_tools/MICE_imputation.py,sha256=KLJXGQLKJ6AuWWttAG-LCCaxpS-ygM4dXPiguHDaL6Y,20815
+ml_tools/ML_callbacks.py,sha256=elD2Yr030sv_6gX_m9GVd6HTyrbmt34nFS8lrgS4HtM,15808
+ml_tools/ML_datasetmaster.py,sha256=rsJgZEGBJmfeKF6cR8CQZzfEx4T7Y-p1wUnR15_nNw0,28400
+ml_tools/ML_evaluation.py,sha256=4GU86rUWMIGbkXrvN6PyjfGwKtWvXKE7pMlWpWeBq14,18988
+ml_tools/ML_evaluation_multi.py,sha256=rJKdgtq-9I7oaI7PRzq7aIZ84XdNV0xzlVePZW4nj0k,16095
+ml_tools/ML_inference.py,sha256=YJ953bhNWsdlPRtJQh3h2ACfMIgp8dQ9KtL9Azar-5s,23489
+ml_tools/ML_models.py,sha256=PqOcNlws7vCJMbiVCKqlPuktxvskZVUHG3VfU-Yshf8,31415
+ml_tools/ML_models_advanced.py,sha256=vk3PZBSu3DVso2S1rKTxxdS43XG8Q5FnasIL3-rMajc,12410
+ml_tools/ML_optimization.py,sha256=P0zkhKAwTpkorIBtR0AOIDcyexo5ngmvFUzo3DfNO-E,22692
+ml_tools/ML_scaler.py,sha256=tw6onj9o8_kk3FQYb930HUzvv1zsFZe2YZJdF3LtHkU,7538
+ml_tools/ML_trainer.py,sha256=ZWI4MbUcLeBxyfoUTL96l5tjHHMp9I64h4SdXnjYmBE,49795
+ml_tools/ML_utilities.py,sha256=z6LbpbZwhn8F__fWlKi-g-cAJQXSxwg1NHfC5FBoAyc,21139
+ml_tools/ML_vision_datasetmaster.py,sha256=feFNUBjybzVJJrdyqToQ_mLV1uDJXHkNL0tmn_zofSY,56034
+ml_tools/ML_vision_evaluation.py,sha256=t12R7i1RkOCt9zu1_lxSBr8OH6A6Get0k8ftDLctn6I,10486
+ml_tools/ML_vision_inference.py,sha256=He3KV3VJAm8PwO-fOq4b9VO8UXFr-GmpuCnoHXf4VZI,20588
+ml_tools/ML_vision_models.py,sha256=G3S4jB9AE9wMpU9ZygOgOx9q1K6t6LAXBYcJ-U2XQ1M,25600
+ml_tools/ML_vision_transformers.py,sha256=95e0aBkHY5VDGE8i5xy57COU7NvSNIgFknnhBubwE40,1832
+ml_tools/PSO_optimization.py,sha256=T-HWHMRJUnPvPwixdU5jif3_rnnI36TzcL8u3oSCwuA,22960
+ml_tools/RNN_forecast.py,sha256=Qa2KoZfdAvSjZ4yE78N4BFXtr3tTr0Gx7tQJZPotsh0,1967
+ml_tools/SQL.py,sha256=vXLPGfVVg8bfkbBE3HVfyEclVbdJy0TBhuQONtMwSCQ,11234
+ml_tools/VIF_factor.py,sha256=at5IVqPvicja2-DNSTSIIy3SkzDWCmLzo3qTG_qr5n8,10422
+ml_tools/_ML_vision_recipe.py,sha256=zrgxFUvTJqQVuwR7jWlbIC2FD29u6eNFPkTRoJ7yEZI,3178
+ml_tools/__init__.py,sha256=kJiankjz9_qXu7gU92mYqYg_anLvt-B6RtW0mMH8uGo,76
+ml_tools/_logger.py,sha256=dlp5cGbzooK9YSNSZYB4yjZrOaQUGW8PTrM411AOvL8,4717
+ml_tools/_schema.py,sha256=yu6aWmn_2Z4_AxAtJGDDCIa96y6JcUp-vgnCS013Qmw,3908
+ml_tools/_script_info.py,sha256=21r83LV3RubsNZ_RTEUON6RbDf7Mh4_udweNcvdF_Fk,212
+ml_tools/constants.py,sha256=3br5Rk9cL2IUo638eJuMOGdbGQaWssaUecYEvSeRBLM,3322
+ml_tools/custom_logger.py,sha256=TGc0Ww2Xlqj2XE3q4bP43hV7T3qnb5ci9f0pYHXF5TY,11226
+ml_tools/data_exploration.py,sha256=bwHzFJ-IAo5GN3T53F-1J_pXUg8VHS91sG_90utAsfg,69911
+ml_tools/ensemble_evaluation.py,sha256=FGHSe8LBI8_w8LjNeJWOcYQ1UK_mc6fVah8gmSvNVGg,26853
+ml_tools/ensemble_inference.py,sha256=0yLmLNj45RVVoSCLH1ZYJG9IoAhTkWUqEZmLOQTFGTY,9348
+ml_tools/ensemble_learning.py,sha256=vsIED7nlheYI4w2SBzP6SC1AnNeMfn-2A1Gqw5EfxsM,21964
+ml_tools/handle_excel.py,sha256=pfdAPb9ywegFkM9T54bRssDOsX-K7rSeV0RaMz7lEAo,14006
+ml_tools/keys.py,sha256=wZOBuEnnHc54vlOZiimnrxfk-sZh6f6suPppJW8rbPQ,3326
+ml_tools/math_utilities.py,sha256=xeKq1quR_3DYLgowcp4Uam_4s3JltUyOnqMOGuAiYWU,8802
+ml_tools/optimization_tools.py,sha256=TYFQ2nSnp7xxs-VyoZISWgnGJghFbsWasHjruegyJRs,12763
+ml_tools/path_manager.py,sha256=CyDU16pOKmC82jPubqJPT6EBt-u-3rGVbxyPIZCvDDY,18432
+ml_tools/serde.py,sha256=c8uDYjYry_VrLvoG4ixqDj5pij88lVn6Tu4NHcPkwDU,6943
+ml_tools/utilities.py,sha256=aWqvYzmxlD74PD5Yqu1VuTekDJeYLQrmPIU_VeVyRp0,22526
+dragon_ml_toolbox-14.3.0.dist-info/METADATA,sha256=TeVrfmCt4AVSweSN4Ai0yyZCJMQtSD1MHsUoEQHXLg4,6475
+dragon_ml_toolbox-14.3.0.dist-info/WHEEL,sha256=_zCd3N1l69ArxyTb8rzEoP9TpbYXkqRFSNOD5OuxnTs,91
+dragon_ml_toolbox-14.3.0.dist-info/top_level.txt,sha256=wm-oxax3ciyez6VoO4zsFd-gSok2VipYXnbg3TH9PtU,9
+dragon_ml_toolbox-14.3.0.dist-info/RECORD,,

{dragon_ml_toolbox-12.13.0.dist-info → dragon_ml_toolbox-14.3.0.dist-info}/licenses/LICENSE-THIRD-PARTY.md

@@ -27,3 +27,13 @@ This project depends on the following third-party packages. Each is governed by
 - [plotnine](https://github.com/has2k1/plotnine/blob/main/LICENSE)
 - [tqdm](https://github.com/tqdm/tqdm/blob/master/LICENSE)
 - [pyarrow](https://github.com/apache/arrow/blob/main/LICENSE.txt)
+- [colorlog](https://github.com/borntyping/python-colorlog/blob/main/LICENSE)
+- [evotorch](https://github.com/nnaisense/evotorch/blob/master/LICENSE)
+- [FreeSimpleGUI](https://github.com/spyoungtech/FreeSimpleGUI/blob/main/license.txt)
+- [nuitka](https://github.com/Nuitka/Nuitka/blob/main/LICENSE.txt)
+- [omegaconf](https://github.com/omry/omegaconf/blob/master/LICENSE)
+- [ordered-set](https://github.com/rspeer/ordered-set/blob/master/MIT-LICENSE)
+- [pyinstaller](https://github.com/pyinstaller/pyinstaller/blob/develop/COPYING.txt)
+- [pytorch_tabular](https://github.com/manujosephv/pytorch_tabular/blob/main/LICENSE)
+- [torchmetrics](https://github.com/Lightning-AI/torchmetrics/blob/master/LICENSE)
+- [zstandard](https://github.com/indygreg/python-zstandard/blob/main/LICENSE)

ml_tools/MICE_imputation.py

@@ -7,19 +7,20 @@ from plotnine import ggplot, labs, theme, element_blank # type: ignore
 from typing import Optional, Union
 
 from .utilities import load_dataframe, merge_dataframes, save_dataframe_filename
-from .math_utilities import threshold_binary_values
+from .math_utilities import threshold_binary_values, discretize_categorical_values
 from .path_manager import sanitize_filename, make_fullpath, list_csv_paths
 from ._logger import _LOGGER
 from ._script_info import _script_info
+from ._schema import FeatureSchema
 
 
 __all__ = [
+    "MiceImputer",
     "apply_mice",
     "save_imputed_datasets",
-    "get_na_column_names",
     "get_convergence_diagnostic",
     "get_imputed_distributions",
-    "run_mice_pipeline"
+    "run_mice_pipeline",
 ]
 
 
@@ -79,7 +80,7 @@ def save_imputed_datasets(save_dir: Union[str, Path], imputed_datasets: list, df
 
 
 #Get names of features that had missing values before imputation
-def get_na_column_names(df: pd.DataFrame):
+def _get_na_column_names(df: pd.DataFrame):
     return [col for col in df.columns if df[col].isna().any()]
 
 
@@ -264,7 +265,7 @@ def run_mice_pipeline(df_path_or_dir: Union[str,Path], target_columns: list[str]
         
         save_imputed_datasets(save_dir=save_datasets_path, imputed_datasets=imputed_datasets, df_targets=df_targets, imputed_dataset_names=imputed_dataset_names)
         
-        imputed_column_names = get_na_column_names(df=df)
+        imputed_column_names = _get_na_column_names(df=df)
         
         get_convergence_diagnostic(kernel=kernel, imputed_dataset_names=imputed_dataset_names, column_names=imputed_column_names, root_dir=save_metrics_path)
         
@@ -278,5 +279,206 @@ def _skip_targets(df: pd.DataFrame, target_cols: list[str]):
     return df_feats, df_targets
 
 
+# modern implementation
+class MiceImputer:
+    """
+    A modern MICE imputation pipeline that uses a FeatureSchema
+    to correctly discretize categorical features after imputation.
+    """
+    def __init__(self, 
+                 schema: FeatureSchema,
+                 iterations: int=20,
+                 resulting_datasets: int = 1,
+                 random_state: int = 101):
+        
+        self.schema = schema
+        self.random_state = random_state
+        self.iterations = iterations
+        self.resulting_datasets = resulting_datasets
+        
+        # --- Store schema info ---
+        
+        # 1. Categorical info
+        if not self.schema.categorical_index_map:
+            _LOGGER.warning("FeatureSchema has no 'categorical_index_map'. No discretization will be applied.")
+            self.cat_info = {}
+        else:
+            self.cat_info = self.schema.categorical_index_map
+            
+        # 2. Ordered feature names (critical for index mapping)
+        self.ordered_features = list(self.schema.feature_names)
+        
+        # 3. Names of categorical features
+        self.categorical_features = list(self.schema.categorical_feature_names)
+
+        _LOGGER.info(f"MiceImputer initialized. Found {len(self.cat_info)} categorical features to discretize.")
+        
+    def _post_process(self, imputed_df: pd.DataFrame) -> pd.DataFrame:
+        """
+        Applies schema-based discretization to a completed dataframe.
+        
+        This method works around the behavior of `discretize_categorical_values`
+        (which returns a full int32 array) by:
+        1. Calling it on the full, ordered feature array.
+        2. Extracting *only* the valid discretized categorical columns.
+        3. Updating the original float dataframe with these integer values.
+        """
+        # If no categorical features are defined, return the df as-is.
+        if not self.cat_info:
+            return imputed_df
+
+        try:
+            # 1. Ensure DataFrame columns match the schema order
+            # This is critical for the index-based categorical_info
+            df_ordered: pd.DataFrame = imputed_df[self.ordered_features] # type: ignore
+            
+            # 2. Convert to NumPy array
+            array_ordered = df_ordered.to_numpy()
+
+            # 3. Apply discretization utility (which returns a full int32 array)
+            # This array has *correct* categorical values but *truncated* continuous values.
+            discretized_array_int32 = discretize_categorical_values(
+                array_ordered,
+                self.cat_info,
+                start_at_zero=True  # Assuming 0-based indexing
+            )
+
+            # 4. Create a new DF from the int32 array, keeping the categorical columns.
+            df_discretized_cats = pd.DataFrame(
+                discretized_array_int32,
+                columns=self.ordered_features,
+                index=df_ordered.index  # <-- Critical: align index
+            )[self.categorical_features] # <-- Select only cat features
+
+            # 5. "Rejoin": Start with a fresh copy of the *original* imputed DF (which has correct continuous floats).
+            final_df = df_ordered.copy()
+            
+            # 6. Use .update() to "paste" the integer categorical values
+            # over the old float categorical values. Continuous floats are unaffected.
+            final_df.update(df_discretized_cats)
+            
+            return final_df
+
+        except Exception as e:
+            _LOGGER.error(f"Failed during post-processing discretization:\n\tInput DF shape: {imputed_df.shape}\n\tSchema features: {len(self.ordered_features)}\n\tCategorical info keys: {list(self.cat_info.keys())}\n{e}")
+            raise
+        
+    def _run_mice(self, 
+                  df: pd.DataFrame, 
+                  df_name: str) -> tuple[mf.ImputationKernel, list[pd.DataFrame], list[str]]:
+        """
+        Runs the MICE kernel and applies schema-based post-processing.
+        
+        Parameters:
+            df (pd.DataFrame): The input dataframe *with NaNs*. Should only contain feature columns.
+            df_name (str): The base name for the dataset.
+
+        Returns:
+            tuple[mf.ImputationKernel, list[pd.DataFrame], list[str]]:
+                - The trained MICE kernel
+                - A list of imputed and processed DataFrames
+                - A list of names for the new DataFrames
+        """
+        # Ensure input df only contains features from the schema and is in the correct order.
+        try:
+            df_feats = df[self.ordered_features]
+        except KeyError as e:
+            _LOGGER.error(f"Input DataFrame is missing required schema columns: {e}")
+            raise
+        
+        # 1. Initialize kernel
+        kernel = mf.ImputationKernel(
+            data=df_feats,
+            num_datasets=self.resulting_datasets,
+            random_state=self.random_state
+        )
+        
+        _LOGGER.info("➡️ Schema-based MICE imputation running...")
+        
+        # 2. Perform MICE
+        kernel.mice(self.iterations)
+        
+        # 3. Retrieve, process, and collect datasets
+        imputed_datasets = []
+        for i in range(self.resulting_datasets):
+            # complete_data returns a pd.DataFrame
+            completed_df = kernel.complete_data(dataset=i)
+            
+            # Apply our new discretization and ordering
+            processed_df = self._post_process(completed_df)
+            imputed_datasets.append(processed_df)
+
+        if not imputed_datasets:
+            _LOGGER.error("No imputed datasets were generated.")
+            raise ValueError()
+
+        # 4. Generate names
+        if self.resulting_datasets == 1:
+            imputed_dataset_names = [f"{df_name}_MICE"]
+        else:
+            imputed_dataset_names = [f"{df_name}_MICE_{i+1}" for i in range(self.resulting_datasets)]
+        
+        # 5. Validate indexes 
+        for imputed_df, subname in zip(imputed_datasets, imputed_dataset_names):
+            assert imputed_df.shape[0] == df.shape[0], f"❌ Row count mismatch in dataset {subname}"
+            assert all(imputed_df.index == df.index), f"❌ Index mismatch in dataset {subname}"
+        
+        _LOGGER.info("Schema-based MICE imputation complete.")
+        
+        return kernel, imputed_datasets, imputed_dataset_names
+        
+    def run_pipeline(self, 
+                     df_path_or_dir: Union[str,Path],
+                     save_datasets_dir: Union[str,Path], 
+                     save_metrics_dir: Union[str,Path],
+                     ):
+        """
+        Runs the complete MICE imputation pipeline.
+
+        This method automates the entire workflow:
+            1. Loads data from a CSV file path or a directory with CSV files.
+            2. Separates features and targets based on the `FeatureSchema`.
+            3. Runs the MICE algorithm on the feature set.
+            4. Applies schema-based post-processing to discretize categorical features.
+            5. Saves the final, processed, and imputed dataset(s) (re-joined with targets) to `save_datasets_dir`.
+            6. Generates and saves convergence and distribution plots for all imputed columns to `save_metrics_dir`.
+
+        Parameters
+        ----------
+        df_path_or_dir :[str,Path]
+            Path to a single CSV file or a directory containing multiple CSV files to impute.
+        save_datasets_dir : [str,Path]
+            Directory where the final imputed and processed dataset(s) will be saved as CSVs.
+        save_metrics_dir : [str,Path]
+            Directory where convergence and distribution plots will be saved.
+        """
+        # Check paths
+        save_datasets_path = make_fullpath(save_datasets_dir, make=True)
+        save_metrics_path = make_fullpath(save_metrics_dir, make=True)
+        
+        input_path = make_fullpath(df_path_or_dir)
+        if input_path.is_file():
+            all_file_paths = [input_path]
+        else:
+            all_file_paths = list(list_csv_paths(input_path).values())
+        
+        for df_path in all_file_paths:
+            
+            df, df_name = load_dataframe(df_path=df_path, kind="pandas")
+            
+            df_features: pd.DataFrame = df[self.schema.feature_names] # type: ignore
+            df_targets = df.drop(columns=self.schema.feature_names)
+            
+            imputed_column_names = _get_na_column_names(df=df_features)
+            
+            kernel, imputed_datasets, imputed_dataset_names = self._run_mice(df=df_features, df_name=df_name)
+            
+            save_imputed_datasets(save_dir=save_datasets_path, imputed_datasets=imputed_datasets, df_targets=df_targets, imputed_dataset_names=imputed_dataset_names)
+            
+            get_convergence_diagnostic(kernel=kernel, imputed_dataset_names=imputed_dataset_names, column_names=imputed_column_names, root_dir=save_metrics_path)
+            
+            get_imputed_distributions(kernel=kernel, df_name=df_name, root_dir=save_metrics_path, column_names=imputed_column_names)
+
+
 def info():
     _script_info(__all__)

ml_tools/ML_callbacks.py

@@ -5,7 +5,7 @@ from typing import Union, Literal, Optional
 from pathlib import Path
 
 from .path_manager import make_fullpath, sanitize_filename
-from .keys import PyTorchLogKeys
+from .keys import PyTorchLogKeys, PyTorchCheckpointKeys
 from ._logger import _LOGGER
 from ._script_info import _script_info
 
@@ -189,7 +189,7 @@ class EarlyStopping(Callback):
 
 class ModelCheckpoint(Callback):
     """
-    Saves the model weights to a directory with automated filename generation and rotation. 
+    Saves the model weights, optimizer state, LR scheduler state (if any), and epoch number to a directory with automated filename generation and rotation. 
     """
     def __init__(self, save_dir: Union[str,Path], checkpoint_name: Optional[str]=None, monitor: str = PyTorchLogKeys.VAL_LOSS,
                  save_best_only: bool = True, mode: Literal['auto', 'min', 'max']= 'auto', verbose: int = 0):
@@ -200,7 +200,7 @@ class ModelCheckpoint(Callback):
         Args:
             save_dir (str): Directory where checkpoint files will be saved.
             checkpoint_name (str| None): If None, the filename will include the epoch and score.
-            monitor (str): Metric to monitor for `save_best_only=True`.
+            monitor (str): Metric to monitor.
             save_best_only (bool): If true, save only the best model.
             mode (str): One of {'auto', 'min', 'max'}.
             verbose (int): Verbosity mode.
@@ -270,15 +270,29 @@ class ModelCheckpoint(Callback):
             if self.verbose > 0:
                 _LOGGER.info(f"Epoch {epoch}: {self.monitor} improved from {old_best_str} to {current:.4f}, saving model to {new_filepath}")
             
+            # Update best score *before* saving
+            self.best = current
+
+            # Create a comprehensive checkpoint dictionary
+            checkpoint_data = {
+                PyTorchCheckpointKeys.EPOCH: epoch,
+                PyTorchCheckpointKeys.MODEL_STATE: self.trainer.model.state_dict(), # type: ignore
+                PyTorchCheckpointKeys.OPTIMIZER_STATE: self.trainer.optimizer.state_dict(), # type: ignore
+                PyTorchCheckpointKeys.BEST_SCORE: self.best, 
+            }
+            
+            # Check for scheduler
+            if hasattr(self.trainer, 'scheduler') and self.trainer.scheduler is not None: # type: ignore
+                checkpoint_data[PyTorchCheckpointKeys.SCHEDULER_STATE] = self.trainer.scheduler.state_dict() # type: ignore
+            
             # Save the new best model
-            torch.save(self.trainer.model.state_dict(), new_filepath) # type: ignore
+            torch.save(checkpoint_data, new_filepath)
 
             # Delete the old best model file
             if self.last_best_filepath and self.last_best_filepath.exists():
                 self.last_best_filepath.unlink()
             
             # Update state
-            self.best = current
             self.last_best_filepath = new_filepath
 
     def _save_rolling_checkpoints(self, epoch, logs):
@@ -292,7 +306,19 @@ class ModelCheckpoint(Callback):
         
         if self.verbose > 0:
             _LOGGER.info(f'Epoch {epoch}: saving model to {filepath}')
-        torch.save(self.trainer.model.state_dict(), filepath) # type: ignore
+
+        # Create a comprehensive checkpoint dictionary
+        checkpoint_data = {
+            PyTorchCheckpointKeys.EPOCH: epoch,
+            PyTorchCheckpointKeys.MODEL_STATE: self.trainer.model.state_dict(), # type: ignore
+            PyTorchCheckpointKeys.OPTIMIZER_STATE: self.trainer.optimizer.state_dict(), # type: ignore
+            PyTorchCheckpointKeys.BEST_SCORE: self.best, # Save the current best score
+        }
+        
+        if hasattr(self.trainer, 'scheduler') and self.trainer.scheduler is not None: # type: ignore
+            checkpoint_data[PyTorchCheckpointKeys.SCHEDULER_STATE] = self.trainer.scheduler.state_dict() # type: ignore
+        
+        torch.save(checkpoint_data, filepath)
 
         self.saved_checkpoints.append(filepath)
 
@@ -309,19 +335,25 @@ class LRScheduler(Callback):
     """
     Callback to manage a PyTorch learning rate scheduler.
     """
-    def __init__(self, scheduler, monitor: Optional[str] = None):
+    def __init__(self, scheduler, monitor: Optional[str] = PyTorchLogKeys.VAL_LOSS):
         """
         This callback automatically calls the scheduler's `step()` method at the
         end of each epoch. It also logs a message when the learning rate changes.
 
         Args:
             scheduler: An initialized PyTorch learning rate scheduler.
-            monitor (str, optional): The metric to monitor for schedulers that require it, like `ReduceLROnPlateau`. Should match a key in the logs (e.g., 'val_loss').
+            monitor (str): The metric to monitor for schedulers that require it, like `ReduceLROnPlateau`. Should match a key in the logs (e.g., 'val_loss').
         """
         super().__init__()
         self.scheduler = scheduler
         self.monitor = monitor
         self.previous_lr = None
+        
+    def set_trainer(self, trainer):
+        """This is called by the Trainer to associate itself with the callback."""
+        super().set_trainer(trainer)
+        # Register the scheduler with the trainer so it can be added to the checkpoint
+        self.trainer.scheduler = self.scheduler # type: ignore
 
     def on_train_begin(self, logs=None):
         """Store the initial learning rate."""