src/ai_dashboard/plot_plugin/heatmap_plot.py

"""Correlation Heatmap plugin"""

from typing import Any
import plotly.express as px
from dash import dcc, html
from dash.dependencies import Input, Output

from ..base import BasePlotPlugin


class CorrelationHeatmapPlugin(BasePlotPlugin):
    """Plugin to render a correlation heatmap of numeric columns in the dataframe."""

    name = "Correlation Heatmap"

    def controls(self) -> Any:
        return html.Div("No controls for this plot.")

    def render(self, **kwargs: Any) -> Any:  # pylint: disable=W0201
        """Render the correlation heatmap."""
        corr = self.dataframe[self.numeric_columns()].corr()
        fig = px.imshow(corr, text_auto=True, color_continuous_scale="RdBu_r")
        return dcc.Graph(figure=fig)

    def register_callbacks(self, app: Any) -> None:
        """Register callbacks for the plot."""

        @app.callback(  # type: ignore
            Output({"type": "plot-output", "plot": self.name}, "children"),
            Input("dynamic-plot-select", "value"),
        )
        def update(_) -> Any:  # type: ignore
            """Update the plot."""
            return self.render()