added documentation

GANDLF/utils/plot_utils.py

@@ -3,33 +3,36 @@
 import seaborn as sns
 import matplotlib.pyplot as plt
 from pathlib import Path
-from io import StringIO
 
 
-def plot_all(
-    df_training,
-    df_validation,
-    df_testing,
-    output_plot_dir,
-):
-    # Drop any columns that might have "_" in the values of their rows, this can be checked through the first row or data type
+def plot_all(df_training, df_validation, df_testing, output_plot_dir):
+    """
+    Plots training, validation, and testing data for loss and other metrics.
+
+    Args:
+        df_training (pd.DataFrame): DataFrame containing training data.
+        df_validation (pd.DataFrame): DataFrame containing validation data.
+        df_testing (pd.DataFrame): DataFrame containing testing data.
+        output_plot_dir (str): Directory to save the plots.
+
+    Returns:
+        tuple: Tuple containing the modified training, validation, and testing DataFrames.
+    """
+    # Drop any columns that might have "_" in the values of their rows
     banned_cols = [
         col
         for col in df_training.columns
         if any("_" in str(val) for val in df_training[col].values)
     ]
 
-    # Determine metrics from the column names by removing the train_ from the column names
+    # Determine metrics from the column names by removing the "train_" prefix
     metrics = [
         col.replace("train_", "")
         for col in df_training.columns
         if "train_" in col and col not in banned_cols
     ]
 
-    print("Metrics found: ", metrics)
-    print("Banned columns: ", banned_cols)
-
-    # Split the values of the banned columns into multiple columns  # Code for splitting output
+    # Split the values of the banned columns into multiple columns
     # for df in [df_training, df_validation, df_testing]:
     #     for col in banned_cols:
     #         if df[col].dtype == "object":
@@ -47,10 +50,7 @@ def plot_all(
         any(metric in col for col in df_training.columns) for metric in metrics
     ), "None of the specified metrics is in the dataframe."
 
-    required_cols = [
-        "epoch_no",
-        "train_loss",
-    ]
+    required_cols = ["epoch_no", "train_loss"]
 
     # Check if the required columns are in the dataframe
     assert all(
@@ -62,29 +62,15 @@ def plot_all(
     # Plot for loss
     plt.figure(figsize=(12, 6))
     if "train_loss" in df_training.columns:
-        sns.lineplot(
-            data=df_training,
-            x="epoch_no",
-            y="train_loss",
-            label="Training",
-        )
+        sns.lineplot(data=df_training, x="epoch_no", y="train_loss", label="Training")
 
     if "valid_loss" in df_validation.columns:
         sns.lineplot(
-            data=df_validation,
-            x="epoch_no",
-            y="valid_loss",
-            label="Validation",
+            data=df_validation, x="epoch_no", y="valid_loss", label="Validation"
         )
 
-    if df_testing is not None:
-        if "test_loss" in df_testing.columns:
-            sns.lineplot(
-                data=df_testing,
-                x="epoch_no",
-                y="test_loss",
-                label="Testing",
-            )
+    if df_testing is not None and "test_loss" in df_testing.columns:
+        sns.lineplot(data=df_testing, x="epoch_no", y="test_loss", label="Testing")
 
     plt.xlim(0, epochs - 1)
     plt.xlabel("Epoch")
@@ -114,14 +100,16 @@ def plot_all(
                     y=metric_col.replace("train", "valid"),
                     label=f"Validation {metric_col}",
                 )
-            if df_testing is not None:
-                if metric_col.replace("train", "test") in df_testing.columns:
-                    sns.lineplot(
-                        data=df_testing,
-                        x="epoch_no",
-                        y=metric_col.replace("train", "test"),
-                        label=f"Testing {metric_col}",
-                    )
+            if (
+                df_testing is not None
+                and metric_col.replace("train", "test") in df_testing.columns
+            ):
+                sns.lineplot(
+                    data=df_testing,
+                    x="epoch_no",
+                    y=metric_col.replace("train", "test"),
+                    label=f"Testing {metric_col}",
+                )
             plt.xlim(0, epochs - 1)
             plt.xlabel("Epoch")
             plt.ylabel(metric.capitalize())