From 0b66b11d535cff14924a6f0ddb9ea17d9f43fd20 Mon Sep 17 00:00:00 2001
From: allegroai <>
Date: Tue, 13 Sep 2022 15:06:45 +0300
Subject: [PATCH] Fix interactive plots

---
 .../plotlympl/mplexporter/exporter.py         | 13 +++++----
 clearml/utilities/plotlympl/renderer.py       | 28 +++++++++++--------
 2 files changed, 24 insertions(+), 17 deletions(-)

diff --git a/clearml/utilities/plotlympl/mplexporter/exporter.py b/clearml/utilities/plotlympl/mplexporter/exporter.py
index 8927c141..b99ed965 100644
--- a/clearml/utilities/plotlympl/mplexporter/exporter.py
+++ b/clearml/utilities/plotlympl/mplexporter/exporter.py
@@ -133,11 +133,12 @@ class Exporter(object):
                 self.draw_line(ax, line)
             for text in ax.texts:
                 self.draw_text(ax, text)
-            for (text, ttp) in zip([ax.xaxis.label, ax.yaxis.label, ax.zaxis.label, ax.title],
-                                   ["xlabel", "ylabel", "zlabel", "title"]):
-                if(hasattr(text, "get_text") and text.get_text()):
-                    self.draw_text(ax, text, force_trans=ax.transAxes,
-                                   text_type=ttp)
+            for (text, ttp) in zip(
+                [ax.xaxis.label, ax.yaxis.label, ax.title] + ([ax.zaxis.label] if hasattr(ax, "zaxis") else []),
+                ["xlabel", "ylabel", "title"] + (["zlabel"] if hasattr(ax, "zaxis") else []),
+            ):
+                if hasattr(text, "get_text") and text.get_text():
+                    self.draw_text(ax, text, force_trans=ax.transAxes, text_type=ttp)
             for artist in ax.artists:
                 # TODO: process other artists
                 if isinstance(artist, matplotlib.text.Text):
@@ -279,7 +280,7 @@ class Exporter(object):
         # protected for removing  _offset_position() and default to "screen"
         offset_order = offset_dict[getattr(collection, '_offset_position', 'screen')]
 
-        self.renderer.draw_path_collection(paths=processed_paths,
+        self.renderer.draw_path_collection(ax, paths=processed_paths,
                                            path_coordinates=path_coords,
                                            path_transforms=path_transforms,
                                            offsets=offsets,
diff --git a/clearml/utilities/plotlympl/renderer.py b/clearml/utilities/plotlympl/renderer.py
index d86aa44f..c0d13d2c 100644
--- a/clearml/utilities/plotlympl/renderer.py
+++ b/clearml/utilities/plotlympl/renderer.py
@@ -399,22 +399,26 @@ class PlotlyRenderer(Renderer):
             marked_line = dict(
                 type="scatter",
                 mode=mode,
-                name=(
-                    str(props["label"])
-                    if isinstance(props["label"], six.string_types)
-                    else props["label"]
-                ),
-                x=props["data"][0],
-                y=props["data"][1],
+                name=(str(props["label"]) if isinstance(props["label"], six.string_types) else props["label"]),
+                x=props["data"][0]
+                if props.get("type") == "collection" and props.get("is_3d")
+                else [xy_pair[0] for xy_pair in props["data"]],
+                y=props["data"][1]
+                if props.get("type") == "collection" and props.get("is_3d")
+                else [xy_pair[1] for xy_pair in props["data"]],
                 xaxis="x{0}".format(self.axis_ct),
                 yaxis="y{0}".format(self.axis_ct),
                 line=line,
                 marker=marker,
             )
-            if len(props["data"]) >= 3:
-                marked_line["z"] = props["data"][2]
-                marked_line["zaxis"] = "z{0}".format(self.axis_ct)
+            if props.get("is_3d"):
+                marked_line["z"] = (
+                    props["data"][2]
+                    if props.get("type") == "collection"
+                    else [xyz_tuple[2] for xyz_tuple in props["data"]]
+                )
                 marked_line["type"] = "scatter3d"
+                marked_line["zaxis"] = "z{0}".format(self.axis_ct)
             if self.x_is_mpl_date:
                 formatter = (
                     self.current_mpl_ax.get_xaxis()
@@ -448,7 +452,7 @@ class PlotlyRenderer(Renderer):
             "images from matplotlib yet!"
         )
 
-    def draw_path_collection(self, **props):
+    def draw_path_collection(self, ax, **props):
         """Add a path collection to data list as a scatter plot.
 
         Current implementation defaults such collections as scatter plots.
@@ -487,6 +491,8 @@ class PlotlyRenderer(Renderer):
                 "label": None,
                 "markerstyle": markerstyle,
                 "linestyle": None,
+                "type": "collection",
+                "is_3d": "3d" in str(type(ax)).split(".")[-1].lower()
             }
             self.msg += "    Drawing path collection as markers\n"
             self.draw_marked_line(**scatter_props)