ds-agent-cli 0.1.0

package/src/tools/visualization_engine.py

@@ -0,0 +1,559 @@
+"""
+Comprehensive Visualization Engine (Matplotlib + Seaborn)
+Automatically generate all relevant plots for data analysis and model evaluation.
+
+All functions now return matplotlib Figure objects for Gradio compatibility.
+"""
+
+import matplotlib
+matplotlib.use('Agg')  # Non-interactive backend for Gradio
+
+import matplotlib.pyplot as plt
+import seaborn as sns
+import polars as pl
+import numpy as np
+import pandas as pd
+from typing import Dict, Any, List, Optional, Tuple
+from pathlib import Path
+import sys
+import os
+from sklearn.metrics import confusion_matrix, roc_curve, auc, precision_recall_curve
+
+# Add parent directory to path
+sys.path.append(os.path.dirname(os.path.dirname(os.path.abspath(__file__))))
+
+from ds_agent.utils.polars_helpers import load_dataframe
+from ds_agent.utils.validation import validate_file_exists
+
+# Import matplotlib visualization functions
+try:
+    from .matplotlib_visualizations import (
+        create_scatter_plot,
+        create_bar_chart,
+        create_histogram,
+        create_boxplot,
+        create_correlation_heatmap,
+        create_distribution_plot,
+        create_roc_curve,
+        create_confusion_matrix,
+        create_feature_importance,
+        create_residual_plot,
+        create_missing_values_heatmap,
+        create_missing_values_bar,
+        create_outlier_detection_boxplot,
+        save_figure,
+        close_figure
+    )
+except ImportError:
+    # Fallback for direct execution
+    from matplotlib_visualizations import (
+        create_scatter_plot,
+        create_bar_chart,
+        create_histogram,
+        create_boxplot,
+        create_correlation_heatmap,
+        create_distribution_plot,
+        create_roc_curve,
+        create_confusion_matrix,
+        create_feature_importance,
+        create_residual_plot,
+        create_missing_values_heatmap,
+        create_missing_values_bar,
+        create_outlier_detection_boxplot,
+        save_figure,
+        close_figure
+    )
+
+# Set global style
+sns.set_style('whitegrid')
+
+
+def generate_all_plots(file_path: str,
+                       target_col: Optional[str] = None,
+                       output_dir: str = "./outputs/plots") -> Dict[str, Any]:
+    """
+    Generate ALL plots for a dataset automatically.
+    
+    Generates:
+    - Data quality plots
+    - EDA plots
+    - Distribution plots
+    - Correlation plots
+    
+    Args:
+        file_path: Path to dataset
+        target_col: Optional target column
+        output_dir: Directory to save plots
+        
+    Returns:
+        Dictionary with Figure objects and saved file paths
+    """
+    validate_file_exists(file_path)
+    Path(output_dir).mkdir(parents=True, exist_ok=True)
+    
+    results = {
+        "output_directory": output_dir,
+        "plots_generated": [],
+        "figure_objects": [],  # Store Figure objects
+        "plot_categories": {}
+    }
+    
+    print(f"🎨 Generating comprehensive visualizations...")
+    
+    # 1. Data Quality Plots
+    quality_plots = generate_data_quality_plots(file_path, output_dir)
+    results["plot_categories"]["data_quality"] = quality_plots
+    results["plots_generated"].extend(quality_plots.get("plot_paths", []))
+    results["figure_objects"].extend(quality_plots.get("figures", []))
+    
+    # 2. EDA Plots
+    eda_plots = generate_eda_plots(file_path, target_col, output_dir)
+    results["plot_categories"]["eda"] = eda_plots
+    results["plots_generated"].extend(eda_plots.get("plot_paths", []))
+    results["figure_objects"].extend(eda_plots.get("figures", []))
+    
+    # 3. Distribution Plots
+    dist_plots = generate_distribution_plots(file_path, output_dir)
+    results["plot_categories"]["distributions"] = dist_plots
+    results["plots_generated"].extend(dist_plots.get("plot_paths", []))
+    results["figure_objects"].extend(dist_plots.get("figures", []))
+    
+    results["total_plots"] = len(results["plots_generated"])
+    print(f"✅ Generated {results['total_plots']} plots in {output_dir}")
+    
+    return results
+
+
+def generate_data_quality_plots(file_path: str, output_dir: str) -> Dict[str, Any]:
+    """Generate plots related to data quality using Matplotlib."""
+    df = load_dataframe(file_path).to_pandas()
+    plots = []
+    figures = []
+    
+    # 1. Missing values bar chart
+    missing_data = df.isnull().sum()
+    if missing_data.sum() > 0:
+        fig = create_missing_values_bar(
+            df=df,
+            title="Missing Values by Column",
+            figsize=(10, 6)
+        )
+        if fig is not None:
+            path = f"{output_dir}/missing_values.png"
+            save_figure(fig, path)
+            plots.append(path)
+            figures.append(fig)
+            print(f"   ✓ Missing values plot")
+    
+    # 2. Data types distribution (pie chart alternative - bar chart)
+    dtype_counts = df.dtypes.astype(str).value_counts()
+    fig = create_bar_chart(
+        categories=dtype_counts.index.tolist(),
+        values=dtype_counts.values,
+        title="Data Types Distribution",
+        xlabel="Data Type",
+        ylabel="Count",
+        figsize=(8, 6),
+        color='steelblue'
+    )
+    if fig is not None:
+        path = f"{output_dir}/data_types.png"
+        save_figure(fig, path)
+        plots.append(path)
+        figures.append(fig)
+        print(f"   ✓ Data types plot")
+    
+    # 3. Outlier detection (box plots)
+    numeric_cols = df.select_dtypes(include=[np.number]).columns[:6]  # Limit to 6
+    if len(numeric_cols) > 0:
+        fig = create_boxplot(
+            data=df[numeric_cols],
+            title="Outlier Detection (Box Plots)",
+            figsize=(12, 6)
+        )
+        if fig is not None:
+            path = f"{output_dir}/outliers_boxplot.png"
+            save_figure(fig, path)
+            plots.append(path)
+            figures.append(fig)
+            print(f"   ✓ Outlier detection plot")
+    
+    return {"plot_paths": plots, "figures": figures, "n_plots": len(plots)}
+
+
+def generate_eda_plots(file_path: str, target_col: Optional[str] = None, output_dir: str = "./outputs/plots/eda") -> Dict[str, Any]:
+    """Generate exploratory data analysis plots using Matplotlib."""
+    df = load_dataframe(file_path).to_pandas()
+    plots = []
+    figures = []
+    
+    numeric_cols = df.select_dtypes(include=[np.number]).columns.tolist()
+    if target_col and target_col in numeric_cols:
+        numeric_cols.remove(target_col)
+    
+    # 1. Correlation heatmap
+    if len(numeric_cols) > 1:
+        fig = create_correlation_heatmap(
+            data=df[numeric_cols[:15]],  # Limit to 15 features
+            title="Feature Correlation Matrix",
+            figsize=(12, 10),
+            annot=True,
+            cmap='RdBu_r'
+        )
+        if fig is not None:
+            path = f"{output_dir}/correlation_heatmap.png"
+            save_figure(fig, path)
+            plots.append(path)
+            figures.append(fig)
+            print(f"   ✓ Correlation heatmap")
+    
+    # 2. Feature relationships with target (scatter plots)
+    if target_col and target_col in df.columns and len(numeric_cols) > 0:
+        top_features = numeric_cols[:4]  # Top 4 features
+        
+        # Create multiple scatter plots
+        fig, axes = plt.subplots(2, 2, figsize=(14, 12))
+        axes = axes.flatten()
+        
+        for i, col in enumerate(top_features):
+            ax = axes[i]
+            ax.scatter(df[col], df[target_col], alpha=0.5, s=30, 
+                      c='steelblue', edgecolors='black', linewidth=0.5)
+            ax.set_xlabel(col, fontsize=11)
+            ax.set_ylabel(target_col, fontsize=11)
+            ax.set_title(f"{col} vs {target_col}", fontsize=12, fontweight='bold')
+            ax.grid(True, alpha=0.3, linestyle='--')
+        
+        fig.suptitle(f"Top Features vs {target_col}", fontsize=14, fontweight='bold', y=0.995)
+        plt.tight_layout()
+        
+        path = f"{output_dir}/feature_relationships.png"
+        save_figure(fig, path)
+        plots.append(path)
+        figures.append(fig)
+        print(f"   ✓ Feature relationships plot")
+    
+    # 3. Pairplot for top features (sample data for performance)
+    if len(numeric_cols) >= 3:
+        sample_size = min(1000, len(df))
+        sample_df = df[numeric_cols[:3]].sample(sample_size)
+        
+        # Create pairplot using seaborn
+        pair_grid = sns.pairplot(sample_df, corner=True, diag_kind='kde', 
+                                plot_kws={'alpha': 0.6, 's': 20})
+        fig = pair_grid.fig
+        fig.suptitle("Feature Pairplot (Top 3 Features)", fontsize=14, 
+                    fontweight='bold', y=1.01)
+        
+        path = f"{output_dir}/pairplot.png"
+        save_figure(fig, path)
+        plots.append(path)
+        figures.append(fig)
+        print(f"   ✓ Pairplot")
+    
+    return {"plot_paths": plots, "figures": figures, "n_plots": len(plots)}
+
+
+def generate_distribution_plots(file_path: str, output_dir: str) -> Dict[str, Any]:
+    """Generate distribution analysis plots using Matplotlib."""
+    df = load_dataframe(file_path).to_pandas()
+    plots = []
+    figures = []
+    
+    numeric_cols = df.select_dtypes(include=[np.number]).columns[:6]
+    
+    if len(numeric_cols) > 0:
+        # Histograms for numeric features in a grid
+        n_cols = min(3, len(numeric_cols))
+        n_rows = (len(numeric_cols) + n_cols - 1) // n_cols
+        
+        fig, axes = plt.subplots(n_rows, n_cols, figsize=(15, n_rows * 4))
+        axes = axes.flatten() if n_rows * n_cols > 1 else [axes]
+        
+        for i, col in enumerate(numeric_cols):
+            ax = axes[i]
+            data = df[col].dropna()
+            
+            # Create histogram with KDE
+            ax.hist(data, bins=30, color='steelblue', edgecolor='black', 
+                   alpha=0.7, density=True)
+            
+            # Add KDE
+            try:
+                sns.kdeplot(data, ax=ax, color='darkred', linewidth=2)
+            except:
+                pass  # Skip KDE if it fails
+            
+            ax.set_title(col[:25], fontsize=11, fontweight='bold')
+            ax.set_xlabel('Value', fontsize=10)
+            ax.set_ylabel('Density', fontsize=10)
+            ax.grid(True, alpha=0.3, linestyle='--')
+        
+        # Hide unused subplots
+        for i in range(len(numeric_cols), len(axes)):
+            axes[i].axis('off')
+        
+        fig.suptitle("Feature Distributions", fontsize=14, fontweight='bold', y=0.995)
+        plt.tight_layout()
+        
+        path = f"{output_dir}/distributions_histogram.png"
+        save_figure(fig, path)
+        plots.append(path)
+        figures.append(fig)
+        print(f"   ✓ Distribution histograms")
+    
+    return {"plot_paths": plots, "figures": figures, "n_plots": len(plots)}
+
+
+def generate_model_performance_plots(y_true, y_pred, y_pred_proba=None,
+                                     task_type="regression",
+                                     model_name="Model",
+                                     output_dir="./outputs/plots") -> Dict[str, Any]:
+    """
+    Generate model performance plots using Matplotlib.
+    
+    Args:
+        y_true: True labels
+        y_pred: Predicted labels
+        y_pred_proba: Predicted probabilities (for classification)
+        task_type: 'classification' or 'regression'
+        model_name: Name of the model
+        output_dir: Output directory
+        
+    Returns:
+        Dictionary with plot paths and figure objects
+    """
+    Path(output_dir).mkdir(parents=True, exist_ok=True)
+    plots = []
+    figures = []
+    
+    if task_type == "classification":
+        # 1. Confusion Matrix
+        cm = confusion_matrix(y_true, y_pred)
+        class_names = [f"Class {i}" for i in range(len(cm))]
+        
+        fig = create_confusion_matrix(
+            cm=cm,
+            class_names=class_names,
+            title=f"Confusion Matrix - {model_name}",
+            show_percentages=True,
+            figsize=(10, 8)
+        )
+        if fig is not None:
+            path = f"{output_dir}/confusion_matrix_{model_name}.png"
+            save_figure(fig, path)
+            plots.append(path)
+            figures.append(fig)
+            print(f"   ✓ Confusion matrix")
+        
+        # 2. ROC Curve (if probabilities provided)
+        if y_pred_proba is not None and len(np.unique(y_true)) == 2:
+            y_proba = y_pred_proba[:, 1] if y_pred_proba.ndim > 1 else y_pred_proba
+            fpr, tpr, _ = roc_curve(y_true, y_proba)
+            roc_auc = auc(fpr, tpr)
+            
+            models_data = {model_name: (fpr, tpr, roc_auc)}
+            fig = create_roc_curve(
+                models_data=models_data,
+                title=f"ROC Curve - {model_name}",
+                figsize=(10, 8)
+            )
+            if fig is not None:
+                path = f"{output_dir}/roc_curve_{model_name}.png"
+                save_figure(fig, path)
+                plots.append(path)
+                figures.append(fig)
+                print(f"   ✓ ROC curve")
+    
+    else:  # Regression
+        # 1. Residual plot (Predicted vs Actual + Residuals)
+        fig = create_residual_plot(
+            y_true=y_true,
+            y_pred=y_pred,
+            title=f"Residual Analysis - {model_name}",
+            figsize=(10, 6)
+        )
+        if fig is not None:
+            path = f"{output_dir}/residuals_{model_name}.png"
+            save_figure(fig, path)
+            plots.append(path)
+            figures.append(fig)
+            print(f"   ✓ Residual plot")
+        
+        # 2. Residuals distribution
+        residuals = y_true - y_pred
+        fig = create_histogram(
+            data=residuals,
+            title=f"Residuals Distribution - {model_name}",
+            xlabel="Residuals",
+            ylabel="Frequency",
+            bins=30,
+            kde=True,
+            figsize=(10, 6)
+        )
+        if fig is not None:
+            path = f"{output_dir}/residuals_dist_{model_name}.png"
+            save_figure(fig, path)
+            plots.append(path)
+            figures.append(fig)
+            print(f"   ✓ Residuals distribution")
+    
+    return {"plot_paths": plots, "figures": figures, "n_plots": len(plots)}
+
+
+def generate_feature_importance_plot(feature_importances: Dict[str, float],
+                                     output_path: str = "./outputs/plots/feature_importance.png",
+                                     top_n: int = 20) -> str:
+    """
+    Generate feature importance plot using Matplotlib.
+    
+    Args:
+        feature_importances: Dictionary of feature: importance
+        output_path: Where to save the plot
+        top_n: Number of top features to show
+        
+    Returns:
+        Path to saved plot
+    """
+    Path(output_path).parent.mkdir(parents=True, exist_ok=True)
+    
+    # Convert dict to lists
+    features = list(feature_importances.keys())
+    importances = np.array(list(feature_importances.values()))
+    
+    # Create plot
+    fig = create_feature_importance(
+        feature_names=features,
+        importances=importances,
+        title=f"Top {top_n} Feature Importances",
+        top_n=top_n,
+        figsize=(10, max(8, top_n * 0.4))
+    )
+    
+    if fig is not None:
+        save_figure(fig, output_path)
+        print(f"   ✓ Feature importance plot")
+        close_figure(fig)
+        return output_path
+    
+    return None
+
+
+def generate_learning_curve(train_sizes, train_scores, val_scores,
+                           model_name="Model",
+                           output_path="./outputs/plots/learning_curve.png") -> str:
+    """Generate learning curve plot using Matplotlib."""
+    Path(output_path).parent.mkdir(parents=True, exist_ok=True)
+    
+    # Calculate mean and std
+    if isinstance(train_scores, list):
+        train_scores = np.array(train_scores)
+        val_scores = np.array(val_scores)
+    
+    if train_scores.ndim > 1:
+        train_scores_mean = np.mean(train_scores, axis=1)
+        train_scores_std = np.std(train_scores, axis=1)
+        val_scores_mean = np.mean(val_scores, axis=1)
+        val_scores_std = np.std(val_scores, axis=1)
+    else:
+        train_scores_mean = train_scores
+        train_scores_std = np.zeros_like(train_scores)
+        val_scores_mean = val_scores
+        val_scores_std = np.zeros_like(val_scores)
+    
+    # Create plot
+    from .matplotlib_visualizations import create_learning_curve as mlp_learning_curve
+    
+    fig = mlp_learning_curve(
+        train_sizes=train_sizes,
+        train_scores_mean=train_scores_mean,
+        train_scores_std=train_scores_std,
+        val_scores_mean=val_scores_mean,
+        val_scores_std=val_scores_std,
+        title=f"Learning Curve - {model_name}",
+        figsize=(10, 6)
+    )
+    
+    if fig is not None:
+        save_figure(fig, output_path)
+        close_figure(fig)
+        return output_path
+    
+    return None
+
+
+def create_plot_gallery_html(plot_paths: List[str], output_path: str = "./outputs/plots/gallery.html") -> str:
+    """
+    Create an HTML gallery page showing all plots (now as PNG images).
+    
+    Args:
+        plot_paths: List of paths to plot files (now PNG instead of HTML)
+        output_path: Where to save the gallery
+        
+    Returns:
+        Path to gallery HTML
+    """
+    html_content = """
+    <!DOCTYPE html>
+    <html>
+    <head>
+        <title>Data Analysis Plot Gallery</title>
+        <style>
+            body {
+                font-family: Arial, sans-serif;
+                margin: 20px;
+                background-color: #f5f5f5;
+            }
+            h1 {
+                color: #333;
+                text-align: center;
+            }
+            .plot-container {
+                background: white;
+                margin: 20px 0;
+                padding: 20px;
+                border-radius: 8px;
+                box-shadow: 0 2px 4px rgba(0,0,0,0.1);
+            }
+            .plot-image {
+                width: 100%;
+                max-width: 1200px;
+                height: auto;
+                display: block;
+                margin: 0 auto;
+            }
+            .plot-title {
+                font-size: 18px;
+                font-weight: bold;
+                margin-bottom: 10px;
+                color: #555;
+            }
+        </style>
+    </head>
+    <body>
+        <h1>📊 Data Analysis Visualization Gallery</h1>
+        <p style="text-align: center; color: #666;">Total Plots: {}</p>
+    """.format(len(plot_paths))
+    
+    for i, plot_path in enumerate(plot_paths, 1):
+        plot_name = Path(plot_path).stem.replace('_', ' ').title()
+        rel_path = os.path.relpath(plot_path, os.path.dirname(output_path))
+        
+        html_content += f"""
+        <div class="plot-container">
+            <div class="plot-title">{i}. {plot_name}</div>
+            <img src="{rel_path}" alt="{plot_name}" class="plot-image">
+        </div>
+        """
+    
+    html_content += """
+    </body>
+    </html>
+    """
+    
+    Path(output_path).parent.mkdir(parents=True, exist_ok=True)
+    with open(output_path, 'w') as f:
+        f.write(html_content)
+    
+    print(f"✅ Created plot gallery: {output_path}")
+    return output_path

package/src/utils/__init__.py

@@ -0,0 +1,42 @@
+"""Utils module initialization."""
+
+from .polars_helpers import (
+    load_dataframe,
+    save_dataframe,
+    get_numeric_columns,
+    get_categorical_columns,
+    get_datetime_columns,
+    detect_id_columns,
+    get_column_info,
+    calculate_memory_usage,
+    split_features_target,
+)
+
+from .validation import (
+    ValidationError,
+    validate_file_exists,
+    validate_file_format,
+    validate_dataframe,
+    validate_column_exists,
+    validate_columns_exist,
+    validate_target_column,
+)
+
+__all__ = [
+    "load_dataframe",
+    "save_dataframe",
+    "get_numeric_columns",
+    "get_categorical_columns",
+    "get_datetime_columns",
+    "detect_id_columns",
+    "get_column_info",
+    "calculate_memory_usage",
+    "split_features_target",
+    "ValidationError",
+    "validate_file_exists",
+    "validate_file_format",
+    "validate_dataframe",
+    "validate_column_exists",
+    "validate_columns_exist",
+    "validate_target_column",
+]