Add scatter plot

src/ai_dashboard/plot_plugin/scatter_plot.py

@@ -0,0 +1,91 @@
+"""This is the scatter plot module for AI Dashboard."""
+
+from typing import Any
+import logging
+
+import plotly.express as px
+from dash import html, dcc
+from dash.dependencies import Input, Output
+
+from ..base import BasePlotPlugin
+
+logger = logging.getLogger(__name__)
+
+
+class ScatterPlotPlugin(BasePlotPlugin):
+    """Scatter plot plugin for AI Dashboard."""
+
+    name = "Scatter Plot"
+
+    def dropdown(self, id_suffix: str, label: str) -> Any:
+        """Generate a dropdown for plot controls."""
+        cols = self.dataframe.columns.tolist()
+        return html.Div(
+            [
+                html.Label(label),
+                dcc.Dropdown(
+                    id={"type": "control", "plot": self.name, "axis": id_suffix},
+                    options=[{"label": c, "value": c} for c in cols],  # type: ignore
+                    value=cols[0],
+                    clearable=False,
+                    persistence=True,
+                    persistence_type="memory",
+                ),
+            ],
+            style={"width": "170px", "marginRight": "12px"},
+        )
+
+    def controls(self) -> Any:
+        """Connect the scatter plot plugin to the Dash app."""
+        return html.Div(
+            [
+                self.dropdown("x", "X-Axis"),
+                self.dropdown("y", "Y-Axis"),
+                self.dropdown("size", "Size"),
+                self.dropdown("color", "Color"),
+                self.dropdown("hover", "Hover"),
+            ],
+            style={"display": "flex", "flexWrap": "wrap", "gap": "12px"},
+        )
+
+    def render(self, **kwargs: Any) -> Any:
+        """Render the scatter plot."""
+        x_axis = kwargs.get("x_axis")
+        y_axis = kwargs.get("y_axis")
+        size_axis = kwargs.get("size_axis")
+        color_axis = kwargs.get("color_axis")
+        hover_axis = kwargs.get("hover_axis")
+        fig = px.scatter(
+            self.dataframe,
+            x=x_axis,
+            y=y_axis,
+            size=size_axis if size_axis in self.dataframe.columns else None,
+            color=color_axis if color_axis in self.dataframe.columns else None,
+            hover_name=hover_axis if hover_axis in self.dataframe.columns else None,
+            log_x=True,
+            size_max=60,
+        )
+        return dcc.Graph(figure=fig)
+
+    def register_callbacks(self, app: Any) -> None:
+        """Register the callbacks for the scatter plot plugin."""
+
+        @app.callback(  # type: ignore
+            Output({"type": "plot-output", "plot": self.name}, "children"),
+            Input({"type": "control", "plot": self.name, "axis": "x"}, "value"),
+            Input({"type": "control", "plot": self.name, "axis": "y"}, "value"),
+            Input({"type": "control", "plot": self.name, "axis": "size"}, "value"),
+            Input({"type": "control", "plot": self.name, "axis": "color"}, "value"),
+            Input({"type": "control", "plot": self.name, "axis": "hover"}, "value"),
+        )
+        def update_scatter(
+            x_axis: str, y_axis: str, size_axis: str, color_axis: str, hover_axis: str
+        ) -> Any:
+            """Update the scatter plot based on user input."""
+            return self.render(
+                x_axis=x_axis,
+                y_axis=y_axis,
+                size_axis=size_axis,
+                color_axis=color_axis,
+                hover_axis=hover_axis,
+            )