Add support for DataFrames and Numpy Arrays for log_plot()

src/dvclive/live.py

@@ -2,7 +2,9 @@
 import json
 import logging
 import math
+import numpy as np
 import os
+import pandas as pd
 import shutil
 import tempfile
 from pathlib import Path
@@ -35,6 +37,7 @@
     StrPath,
     catch_and_warn,
     clean_and_copy_into,
+    convert_datapoints_to_list_of_dicts,
     env2bool,
     inside_notebook,
     matplotlib_installed,
@@ -391,14 +394,20 @@ def log_image(self, name: str, val):
     def log_plot(
         self,
         name: str,
-        datapoints: List[Dict],
+        datapoints: pd.DataFrame | np.ndarray | List[Dict],
         x: str,
         y: str,
         template: Optional[str] = None,
         title: Optional[str] = None,
         x_label: Optional[str] = None,
         y_label: Optional[str] = None,
+        columns: Optional[List[str]] = None,
     ):
+        # Convert the given datapoints to List[Dict]
+        datapoints = convert_datapoints_to_list_of_dicts(
+            datapoints=datapoints, columns=columns
+        )
+
         if not CustomPlot.could_log(datapoints):
             raise InvalidDataTypeError(name, type(datapoints))
 

src/dvclive/utils.py

@@ -1,12 +1,14 @@
 import csv
 import json
+import numpy as np
 import os
+import pandas as pd
 import re
 import shutil
-import webbrowser
 from pathlib import Path
 from platform import uname
-from typing import Union
+from typing import Union, List, Dict, Optional
+import webbrowser
 
 StrPath = Union[str, Path]
 
@@ -194,3 +196,28 @@ def read_history(live, metric):
 def read_latest(live, metric_name):
     _, latest = parse_metrics(live)
     return latest["step"], latest[metric_name]
+
+
+def convert_datapoints_to_list_of_dicts(
+    datapoints: pd.DataFrame | np.ndarray | List[Dict],
+    columns: Optional[List[str]] = None,
+) -> List[Dict]:
+    """
+    Convert the given datapoints to a list of dictionaries.
+
+    Args:
+        datapoints: The input datapoints to be converted.
+        columns: The column columns for the datapoints. Applied only for np.ndarray inputs.
+
+    Returns:
+        A list of dictionaries representing the datapoints.
+
+    Raises:
+        None
+    """
+    if isinstance(datapoints, pd.DataFrame):
+        return datapoints.to_dict(orient="records")
+    elif isinstance(datapoints, np.ndarray):
+        return pd.DataFrame(datapoints, columns=columns).to_dict(orient="records")
+    else:
+        return datapoints

tests/test_utils.py

@@ -1,6 +1,8 @@
+import numpy as np
+import pandas as pd
 import pytest
 
-from dvclive.utils import standardize_metric_name
+from dvclive.utils import standardize_metric_name, convert_datapoints_to_list_of_dicts
 
 
 @pytest.mark.parametrize(
@@ -15,3 +17,25 @@
 )
 def test_standardize_metric_name(framework, logged, standardized):
     assert standardize_metric_name(logged, framework) == standardized
+
+
+class TestConvertDatapointsToListOfDicts:
+    def test_dataframe(self):
+        df = pd.DataFrame({"A": [1, 2], "B": [3, 4]})
+        expected_output = [{"A": 1, "B": 3}, {"A": 2, "B": 4}]
+        assert convert_datapoints_to_list_of_dicts(df) == expected_output
+
+    def test_ndarray_with_columns(self):
+        arr = np.array([[1, 3], [2, 4]])
+        columns = ["A", "B"]
+        expected_output = [{"A": 1, "B": 3}, {"A": 2, "B": 4}]
+        assert convert_datapoints_to_list_of_dicts(arr, columns) == expected_output
+
+    def test_ndarray_without_columns(self):
+        arr = np.array([[1, 3], [2, 4]])
+        expected_output = [{0: 1, 1: 3}, {0: 2, 1: 4}]
+        assert convert_datapoints_to_list_of_dicts(arr) == expected_output
+
+    def test_list_of_dicts(self):
+        list_of_dicts = [{"A": 1, "B": 3}, {"A": 2, "B": 4}]
+        assert convert_datapoints_to_list_of_dicts(list_of_dicts) == list_of_dicts