careamics 0.1.0rc5__py3-none-any.whl → 0.1.0rc7__py3-none-any.whl

careamics/callbacks/hyperparameters_callback.py

@@ -1,3 +1,5 @@
+"""Callback saving CAREamics configuration as hyperparameters in the model."""
+
 from pytorch_lightning import LightningModule, Trainer
 from pytorch_lightning.callbacks import Callback
 
@@ -11,13 +13,18 @@ class HyperParametersCallback(Callback):
     This allows saving the configuration as dictionnary in the checkpoints, and
     loading it subsequently in a CAREamist instance.
 
+    Parameters
+    ----------
+    config : Configuration
+        CAREamics configuration to be saved as hyperparameter in the model.
+
     Attributes
     ----------
     config : Configuration
         CAREamics configuration to be saved as hyperparameter in the model.
     """
 
-    def __init__(self, config: Configuration):
+    def __init__(self, config: Configuration) -> None:
         """
         Constructor.
 
@@ -28,14 +35,14 @@ class HyperParametersCallback(Callback):
         """
         self.config = config
 
-    def on_train_start(self, trainer: Trainer, pl_module: LightningModule):
+    def on_train_start(self, trainer: Trainer, pl_module: LightningModule) -> None:
         """
         Update the hyperparameters of the model with the configuration on train start.
 
         Parameters
         ----------
         trainer : Trainer
-            PyTorch Lightning trainer.
+            PyTorch Lightning trainer, unused.
         pl_module : LightningModule
             PyTorch Lightning module.
         """

careamics/callbacks/progress_bar_callback.py

@@ -1,3 +1,5 @@
+"""Progressbar callback."""
+
 import sys
 from typing import Dict, Union
 
@@ -10,7 +12,13 @@ class ProgressBarCallback(TQDMProgressBar):
     """Progress bar for training and validation steps."""
 
     def init_train_tqdm(self) -> tqdm:
-        """Override this to customize the tqdm bar for training."""
+        """Override this to customize the tqdm bar for training.
+
+        Returns
+        -------
+        tqdm
+            A tqdm bar.
+        """
         bar = tqdm(
             desc="Training",
             position=(2 * self.process_position),
@@ -23,7 +31,13 @@ class ProgressBarCallback(TQDMProgressBar):
         return bar
 
     def init_validation_tqdm(self) -> tqdm:
-        """Override this to customize the tqdm bar for validation."""
+        """Override this to customize the tqdm bar for validation.
+
+        Returns
+        -------
+        tqdm
+            A tqdm bar.
+        """
         # The main progress bar doesn't exist in `trainer.validate()`
         has_main_bar = self.train_progress_bar is not None
         bar = tqdm(
@@ -37,7 +51,13 @@ class ProgressBarCallback(TQDMProgressBar):
         return bar
 
     def init_test_tqdm(self) -> tqdm:
-        """Override this to customize the tqdm bar for testing."""
+        """Override this to customize the tqdm bar for testing.
+
+        Returns
+        -------
+        tqdm
+            A tqdm bar.
+        """
         bar = tqdm(
             desc="Testing",
             position=(2 * self.process_position),
@@ -52,6 +72,19 @@ class ProgressBarCallback(TQDMProgressBar):
     def get_metrics(
         self, trainer: Trainer, pl_module: LightningModule
     ) -> Dict[str, Union[int, str, float, Dict[str, float]]]:
-        """Override this to customize the metrics displayed in the progress bar."""
+        """Override this to customize the metrics displayed in the progress bar.
+
+        Parameters
+        ----------
+        trainer : Trainer
+            The trainer object.
+        pl_module : LightningModule
+            The LightningModule object, unused.
+
+        Returns
+        -------
+        dict
+            A dictionary with the metrics to display in the progress bar.
+        """
         pbar_metrics = trainer.progress_bar_metrics
         return {**pbar_metrics}