ds-agent-cli 0.1.0

package/src/tools/plotly_visualizations.py

@@ -0,0 +1,497 @@
+"""
+Plotly Interactive Visualization Tools
+Create interactive, web-based visualizations that can be explored in browsers.
+"""
+
+import polars as pl
+import plotly.express as px
+import plotly.graph_objects as go
+from plotly.subplots import make_subplots
+import numpy as np
+from typing import Dict, Any, List, Optional
+from pathlib import Path
+import sys
+import os
+
+# Add parent directory to path for imports
+sys.path.append(os.path.dirname(os.path.dirname(os.path.abspath(__file__))))
+
+from ds_agent.utils.polars_helpers import (
+    load_dataframe,
+    get_numeric_columns,
+    get_categorical_columns,
+)
+from ds_agent.utils.validation import (
+    validate_file_exists,
+    validate_file_format,
+    validate_dataframe,
+    validate_column_exists,
+)
+
+
+def generate_interactive_scatter(
+    file_path: str,
+    x_col: str,
+    y_col: str,
+    color_col: Optional[str] = None,
+    size_col: Optional[str] = None,
+    output_path: str = "./outputs/plots/interactive/scatter.html"
+) -> Dict[str, Any]:
+    """
+    Create interactive scatter plot with Plotly.
+    
+    Args:
+        file_path: Path to dataset
+        x_col: Column for X-axis
+        y_col: Column for Y-axis
+        color_col: Optional column for color coding
+        size_col: Optional column for bubble size
+        output_path: Path to save HTML file
+        
+    Returns:
+        Dictionary with plot info and path
+    """
+    # Validation
+    validate_file_exists(file_path)
+    validate_file_format(file_path)
+    
+    # Load data
+    df = load_dataframe(file_path)
+    validate_dataframe(df)
+    validate_column_exists(df, x_col)
+    validate_column_exists(df, y_col)
+    
+    if color_col:
+        validate_column_exists(df, color_col)
+    if size_col:
+        validate_column_exists(df, size_col)
+    
+    # Convert to pandas for plotly
+    df_pd = df.to_pandas()
+    
+    # Create figure
+    fig = px.scatter(
+        df_pd,
+        x=x_col,
+        y=y_col,
+        color=color_col,
+        size=size_col,
+        hover_data=df_pd.columns.tolist(),
+        title=f"Interactive Scatter: {y_col} vs {x_col}",
+        template="plotly_white"
+    )
+    
+    # Update layout for better interactivity
+    fig.update_layout(
+        hovermode='closest',
+        height=600,
+        font=dict(size=12)
+    )
+    
+    # Save
+    Path(output_path).parent.mkdir(parents=True, exist_ok=True)
+    fig.write_html(output_path)
+    
+    return {
+        "status": "success",
+        "plot_type": "interactive_scatter",
+        "output_path": output_path,
+        "x_col": x_col,
+        "y_col": y_col,
+        "color_col": color_col,
+        "size_col": size_col,
+        "num_points": len(df)
+    }
+
+
+def generate_interactive_histogram(
+    file_path: str,
+    column: str,
+    bins: int = 30,
+    color_col: Optional[str] = None,
+    output_path: str = "./outputs/plots/interactive/histogram.html"
+) -> Dict[str, Any]:
+    """
+    Create interactive histogram with Plotly.
+    
+    Args:
+        file_path: Path to dataset
+        column: Column to plot
+        bins: Number of bins
+        color_col: Optional column for grouped histograms
+        output_path: Path to save HTML file
+        
+    Returns:
+        Dictionary with plot info
+    """
+    # Validation
+    validate_file_exists(file_path)
+    validate_file_format(file_path)
+    
+    df = load_dataframe(file_path)
+    validate_dataframe(df)
+    validate_column_exists(df, column)
+    
+    if color_col:
+        validate_column_exists(df, color_col)
+    
+    df_pd = df.to_pandas()
+    
+    # Create histogram
+    fig = px.histogram(
+        df_pd,
+        x=column,
+        nbins=bins,
+        color=color_col,
+        title=f"Distribution of {column}",
+        template="plotly_white",
+        marginal="box"  # Add box plot on top
+    )
+    
+    fig.update_layout(
+        bargap=0.1,
+        height=600,
+        showlegend=True if color_col else False
+    )
+    
+    # Save
+    Path(output_path).parent.mkdir(parents=True, exist_ok=True)
+    fig.write_html(output_path)
+    
+    return {
+        "status": "success",
+        "plot_type": "interactive_histogram",
+        "output_path": output_path,
+        "column": column,
+        "bins": bins,
+        "color_col": color_col
+    }
+
+
+def generate_interactive_correlation_heatmap(
+    file_path: str,
+    output_path: str = None
+) -> Dict[str, Any]:
+    """
+    Create interactive correlation heatmap with Plotly.
+    
+    Args:
+        file_path: Path to dataset
+        output_path: Path to save HTML file (auto-determined if None)
+        
+    Returns:
+        Dictionary with plot info
+    """
+    # Auto-determine output path based on environment
+    if output_path is None:
+        output_base = os.getenv("DS_AGENT_OUTPUT_DIR", "./outputs")
+        output_path = f"{output_base}/plots/interactive/correlation_heatmap.html"
+    
+    # Validation
+    validate_file_exists(file_path)
+    validate_file_format(file_path)
+    
+    df = load_dataframe(file_path)
+    validate_dataframe(df)
+    
+    # Get numeric columns
+    numeric_cols = get_numeric_columns(df)
+    
+    if len(numeric_cols) < 2:
+        return {
+            "status": "error",
+            "message": "Need at least 2 numeric columns for correlation"
+        }
+    
+    # Calculate correlation matrix
+    df_numeric = df.select(numeric_cols)
+    corr_matrix = df_numeric.to_pandas().corr()
+    
+    # Create heatmap
+    fig = go.Figure(data=go.Heatmap(
+        z=corr_matrix.values,
+        x=corr_matrix.columns,
+        y=corr_matrix.columns,
+        colorscale='RdBu',
+        zmid=0,
+        text=np.round(corr_matrix.values, 2),
+        texttemplate='%{text}',
+        textfont={"size": 10},
+        colorbar=dict(title="Correlation")
+    ))
+    
+    fig.update_layout(
+        title="Interactive Correlation Heatmap",
+        template="plotly_white",
+        height=max(600, len(numeric_cols) * 30),
+        width=max(600, len(numeric_cols) * 30),
+        xaxis={'side': 'bottom'},
+        yaxis={'side': 'left'}
+    )
+    
+    # Save
+    Path(output_path).parent.mkdir(parents=True, exist_ok=True)
+    fig.write_html(output_path)
+    
+    return {
+        "status": "success",
+        "plot_type": "interactive_correlation_heatmap",
+        "output_path": output_path,
+        "num_features": len(numeric_cols)
+    }
+
+
+def generate_interactive_box_plots(
+    file_path: str,
+    columns: Optional[List[str]] = None,
+    group_by: Optional[str] = None,
+    output_path: str = "./outputs/plots/interactive/box_plots.html"
+) -> Dict[str, Any]:
+    """
+    Create interactive box plots for outlier detection.
+    
+    Args:
+        file_path: Path to dataset
+        columns: Columns to plot (all numeric if None)
+        group_by: Optional categorical column for grouping
+        output_path: Path to save HTML file
+        
+    Returns:
+        Dictionary with plot info
+    """
+    # Validation
+    validate_file_exists(file_path)
+    validate_file_format(file_path)
+    
+    df = load_dataframe(file_path)
+    validate_dataframe(df)
+    
+    # Determine columns to plot
+    if columns is None:
+        columns = get_numeric_columns(df)
+    else:
+        for col in columns:
+            validate_column_exists(df, col)
+    
+    if len(columns) == 0:
+        return {
+            "status": "error",
+            "message": "No numeric columns to plot"
+        }
+    
+    if group_by:
+        validate_column_exists(df, group_by)
+    
+    df_pd = df.to_pandas()
+    
+    # Create subplots
+    rows = (len(columns) + 2) // 3  # 3 plots per row
+    cols = min(3, len(columns))
+    
+    fig = make_subplots(
+        rows=rows,
+        cols=cols,
+        subplot_titles=columns,
+        vertical_spacing=0.1
+    )
+    
+    for idx, col in enumerate(columns):
+        row = idx // 3 + 1
+        col_idx = idx % 3 + 1
+        
+        if group_by:
+            for group in df_pd[group_by].unique():
+                group_data = df_pd[df_pd[group_by] == group][col]
+                fig.add_trace(
+                    go.Box(y=group_data, name=str(group), showlegend=(idx == 0)),
+                    row=row,
+                    col=col_idx
+                )
+        else:
+            fig.add_trace(
+                go.Box(y=df_pd[col], name=col, showlegend=False),
+                row=row,
+                col=col_idx
+            )
+    
+    fig.update_layout(
+        title="Interactive Box Plots - Outlier Detection",
+        template="plotly_white",
+        height=400 * rows,
+        showlegend=bool(group_by)
+    )
+    
+    # Save
+    Path(output_path).parent.mkdir(parents=True, exist_ok=True)
+    fig.write_html(output_path)
+    
+    return {
+        "status": "success",
+        "plot_type": "interactive_box_plots",
+        "output_path": output_path,
+        "columns_plotted": columns,
+        "group_by": group_by
+    }
+
+
+def generate_interactive_time_series(
+    file_path: str,
+    time_col: str,
+    value_cols: List[str],
+    output_path: str = "./outputs/plots/interactive/time_series.html"
+) -> Dict[str, Any]:
+    """
+    Create interactive time series plot with Plotly.
+    
+    Args:
+        file_path: Path to dataset
+        time_col: Column with datetime values
+        value_cols: Columns to plot over time
+        output_path: Path to save HTML file
+        
+    Returns:
+        Dictionary with plot info
+    """
+    # Validation
+    validate_file_exists(file_path)
+    validate_file_format(file_path)
+    
+    df = load_dataframe(file_path)
+    validate_dataframe(df)
+    validate_column_exists(df, time_col)
+    
+    for col in value_cols:
+        validate_column_exists(df, col)
+    
+    # Parse datetime if needed
+    if df[time_col].dtype == pl.Utf8:
+        df = df.with_columns(
+            pl.col(time_col).str.strptime(pl.Datetime, strict=False).alias(time_col)
+        )
+    
+    df_pd = df.to_pandas()
+    
+    # Create figure
+    fig = go.Figure()
+    
+    for col in value_cols:
+        fig.add_trace(go.Scatter(
+            x=df_pd[time_col],
+            y=df_pd[col],
+            mode='lines+markers',
+            name=col,
+            hovertemplate=f'<b>{col}</b><br>Time: %{{x}}<br>Value: %{{y:.2f}}<extra></extra>'
+        ))
+    
+    fig.update_layout(
+        title="Interactive Time Series",
+        xaxis_title=time_col,
+        yaxis_title="Value",
+        template="plotly_white",
+        height=600,
+        hovermode='x unified',
+        legend=dict(
+            orientation="h",
+            yanchor="bottom",
+            y=1.02,
+            xanchor="right",
+            x=1
+        )
+    )
+    
+    # Add range slider
+    fig.update_xaxes(rangeslider_visible=True)
+    
+    # Save
+    Path(output_path).parent.mkdir(parents=True, exist_ok=True)
+    fig.write_html(output_path)
+    
+    return {
+        "status": "success",
+        "plot_type": "interactive_time_series",
+        "output_path": output_path,
+        "time_col": time_col,
+        "value_cols": value_cols
+    }
+
+
+def generate_plotly_dashboard(
+    file_path: str,
+    target_col: Optional[str] = None,
+    output_dir: str = "./outputs/plots/interactive"
+) -> Dict[str, Any]:
+    """
+    Generate a complete dashboard with multiple interactive plots.
+    
+    Args:
+        file_path: Path to dataset
+        target_col: Optional target column for supervised analysis
+        output_dir: Directory to save all plots
+        
+    Returns:
+        Dictionary with paths to all generated plots
+    """
+    # Validation
+    validate_file_exists(file_path)
+    validate_file_format(file_path)
+    
+    df = load_dataframe(file_path)
+    validate_dataframe(df)
+    
+    if target_col:
+        validate_column_exists(df, target_col)
+    
+    numeric_cols = get_numeric_columns(df)
+    categorical_cols = get_categorical_columns(df)
+    
+    plots_generated = []
+    
+    # 1. Correlation heatmap
+    if len(numeric_cols) >= 2:
+        result = generate_interactive_correlation_heatmap(
+            file_path,
+            output_path=f"{output_dir}/correlation_heatmap.html"
+        )
+        if result["status"] == "success":
+            plots_generated.append(result)
+    
+    # 2. Box plots for outliers
+    if len(numeric_cols) > 0:
+        result = generate_interactive_box_plots(
+            file_path,
+            columns=numeric_cols[:10],  # Limit to 10 for performance
+            output_path=f"{output_dir}/box_plots.html"
+        )
+        if result["status"] == "success":
+            plots_generated.append(result)
+    
+    # 3. Target variable analysis if provided
+    if target_col and target_col in numeric_cols:
+        # Scatter plots against target
+        for col in numeric_cols[:5]:  # Top 5 features
+            if col != target_col:
+                result = generate_interactive_scatter(
+                    file_path,
+                    x_col=col,
+                    y_col=target_col,
+                    output_path=f"{output_dir}/scatter_{col}_vs_{target_col}.html"
+                )
+                if result["status"] == "success":
+                    plots_generated.append(result)
+    
+    # 4. Distribution plots for numeric features
+    for col in numeric_cols[:5]:  # Top 5 features
+        result = generate_interactive_histogram(
+            file_path,
+            column=col,
+            output_path=f"{output_dir}/histogram_{col}.html"
+        )
+        if result["status"] == "success":
+            plots_generated.append(result)
+    
+    return {
+        "status": "success",
+        "plots_generated": len(plots_generated),
+        "plots": plots_generated,
+        "output_dir": output_dir
+    }