aiecs 1.1.0__py3-none-any.whl → 1.2.0__py3-none-any.whl

aiecs/tools/statistics/data_visualizer_tool.py

@@ -0,0 +1,460 @@
+"""
+Data Visualizer Tool - Smart data visualization and chart generation
+
+This tool provides intelligent visualization capabilities with:
+- Auto chart type recommendation
+- Multiple chart types support
+- Interactive and static visualizations
+- Export in multiple formats
+"""
+
+import os
+import logging
+import tempfile
+from typing import Dict, Any, List, Optional, Union
+from enum import Enum
+
+import pandas as pd
+import numpy as np
+from pydantic import BaseModel, Field, ValidationError, ConfigDict
+
+from aiecs.tools.base_tool import BaseTool
+from aiecs.tools import register_tool
+
+
+class ChartType(str, Enum):
+    """Supported chart types"""
+    # Basic charts
+    LINE = "line"
+    BAR = "bar"
+    SCATTER = "scatter"
+    HISTOGRAM = "histogram"
+    BOX = "box"
+    VIOLIN = "violin"
+    
+    # Advanced charts
+    HEATMAP = "heatmap"
+    CORRELATION_MATRIX = "correlation_matrix"
+    PAIR_PLOT = "pair_plot"
+    PARALLEL_COORDINATES = "parallel_coordinates"
+    
+    # Statistical charts
+    DISTRIBUTION = "distribution"
+    QQ_PLOT = "qq_plot"
+    RESIDUAL_PLOT = "residual_plot"
+    
+    # Time series
+    TIME_SERIES = "time_series"
+    
+    # Auto-detect
+    AUTO = "auto"
+
+
+class VisualizationStyle(str, Enum):
+    """Visualization styles"""
+    STATIC = "static"
+    INTERACTIVE = "interactive"
+    ANIMATED = "animated"
+
+
+
+
+class DataVisualizerError(Exception):
+    """Base exception for DataVisualizer errors"""
+    pass
+
+
+class VisualizationError(DataVisualizerError):
+    """Raised when visualization fails"""
+    pass
+
+
+@register_tool('data_visualizer')
+class DataVisualizerTool(BaseTool):
+    """
+    Intelligent data visualization tool that can:
+    1. Auto-recommend appropriate chart types
+    2. Generate interactive visualizations
+    3. Create multi-dimensional plots
+    4. Export in multiple formats
+    
+    Integrates with chart_tool for core visualization operations.
+    """
+    
+    # Configuration schema
+    class Config(BaseModel):
+        """Configuration for the data visualizer tool"""
+        model_config = ConfigDict(env_prefix="DATA_VISUALIZER_")
+        
+        default_style: str = Field(
+            default="static",
+            description="Default visualization style"
+        )
+        default_output_dir: str = Field(
+            default=tempfile.gettempdir(),
+            description="Default directory for output files"
+        )
+        default_dpi: int = Field(
+            default=100,
+            description="Default DPI for image exports"
+        )
+        default_figsize: List[int] = Field(
+            default=[10, 6],
+            description="Default figure size in inches (width, height)"
+        )
+        enable_auto_recommendation: bool = Field(
+            default=True,
+            description="Whether to enable automatic chart type recommendation"
+        )
+    
+    def __init__(self, config: Optional[Dict[str, Any]] = None):
+        """
+        Initialize DataVisualizerTool with settings.
+        
+        Args:
+            config: Optional configuration overrides
+        """
+        super().__init__(config)
+        
+        # Parse configuration
+        self.config = self.Config(**(config or {}))
+        
+        self.logger = logging.getLogger(__name__)
+        if not self.logger.handlers:
+            handler = logging.StreamHandler()
+            handler.setFormatter(logging.Formatter('%(asctime)s %(levelname)s %(message)s'))
+            self.logger.addHandler(handler)
+        self.logger.setLevel(logging.INFO)
+        
+        # Initialize external tools
+        self._init_external_tools()
+    
+    def _init_external_tools(self):
+        """Initialize external task tools"""
+        self.external_tools = {}
+        
+        # Initialize ChartTool for visualization operations
+        try:
+            from aiecs.tools.task_tools.chart_tool import ChartTool
+            self.external_tools['chart'] = ChartTool()
+            self.logger.info("ChartTool initialized successfully")
+        except ImportError:
+            self.logger.warning("ChartTool not available")
+            self.external_tools['chart'] = None
+    
+    # Schema definitions
+    class VisualizeSchema(BaseModel):
+        """Schema for visualize operation"""
+        data: Union[Dict[str, Any], List[Dict[str, Any]]] = Field(description="Data to visualize")
+        chart_type: ChartType = Field(default=ChartType.AUTO, description="Type of chart")
+        x: Optional[str] = Field(default=None, description="X-axis column")
+        y: Optional[str] = Field(default=None, description="Y-axis column")
+        hue: Optional[str] = Field(default=None, description="Hue/color column")
+        style: VisualizationStyle = Field(default=VisualizationStyle.STATIC, description="Visualization style")
+        title: Optional[str] = Field(default=None, description="Chart title")
+        output_path: Optional[str] = Field(default=None, description="Output file path")
+    
+    class AutoVisualizeDatasetSchema(BaseModel):
+        """Schema for auto_visualize_dataset operation"""
+        data: Union[Dict[str, Any], List[Dict[str, Any]]] = Field(description="Dataset to visualize")
+        max_charts: int = Field(default=10, description="Maximum number of charts to generate")
+        focus_areas: Optional[List[str]] = Field(default=None, description="Areas to focus on")
+    
+    class RecommendChartTypeSchema(BaseModel):
+        """Schema for recommend_chart_type operation"""
+        data: Union[Dict[str, Any], List[Dict[str, Any]]] = Field(description="Data for recommendation")
+        x: Optional[str] = Field(default=None, description="X column")
+        y: Optional[str] = Field(default=None, description="Y column")
+    
+    def visualize(
+        self,
+        data: Union[Dict[str, Any], List[Dict[str, Any]], pd.DataFrame],
+        chart_type: ChartType = ChartType.AUTO,
+        x: Optional[str] = None,
+        y: Optional[str] = None,
+        hue: Optional[str] = None,
+        style: VisualizationStyle = VisualizationStyle.STATIC,
+        title: Optional[str] = None,
+        output_path: Optional[str] = None
+    ) -> Dict[str, Any]:
+        """
+        Create visualization with auto chart type recommendation.
+        
+        Args:
+            data: Data to visualize
+            chart_type: Type of chart (auto-recommended if AUTO)
+            x: X-axis column name
+            y: Y-axis column name
+            hue: Column for color encoding
+            style: Visualization style
+            title: Chart title
+            output_path: Path to save the chart
+            
+        Returns:
+            Dict containing:
+                - chart_info: Information about generated chart
+                - chart_type: Type of chart created
+                - recommendation_reason: Reason for chart type choice
+                - output_path: Path to saved chart (if saved)
+                
+        Raises:
+            VisualizationError: If visualization fails
+        """
+        try:
+            df = self._to_dataframe(data)
+            
+            # Auto-recommend chart type if needed
+            if chart_type == ChartType.AUTO:
+                chart_type, reason = self._recommend_chart_type(df, x, y)
+                self.logger.info(f"Auto-recommended chart type: {chart_type.value} - {reason}")
+            else:
+                reason = "User specified"
+            
+            # Generate output path if not provided
+            if output_path is None:
+                output_path = os.path.join(
+                    self.config.default_output_dir,
+                    f"chart_{chart_type.value}_{pd.Timestamp.now().strftime('%Y%m%d_%H%M%S')}.png"
+                )
+            
+            # Create visualization using chart_tool if available
+            if self.external_tools.get('chart'):
+                chart_result = self._create_chart_with_tool(df, chart_type, x, y, hue, title, output_path)
+            else:
+                chart_result = self._create_chart_matplotlib(df, chart_type, x, y, hue, title, output_path)
+            
+            return {
+                'chart_info': chart_result,
+                'chart_type': chart_type.value,
+                'recommendation_reason': reason,
+                'output_path': output_path,
+                'style': style.value
+            }
+            
+        except Exception as e:
+            self.logger.error(f"Error creating visualization: {e}")
+            raise VisualizationError(f"Visualization failed: {e}")
+    
+    def auto_visualize_dataset(
+        self,
+        data: Union[Dict[str, Any], List[Dict[str, Any]], pd.DataFrame],
+        max_charts: int = 10,
+        focus_areas: Optional[List[str]] = None
+    ) -> Dict[str, Any]:
+        """
+        Automatically generate a comprehensive visualization suite.
+        
+        Args:
+            data: Dataset to visualize
+            max_charts: Maximum number of charts to generate
+            focus_areas: Specific areas to focus on (distributions, correlations, outliers)
+            
+        Returns:
+            Dict containing information about all generated charts
+        """
+        try:
+            df = self._to_dataframe(data)
+            
+            generated_charts = []
+            chart_count = 0
+            
+            # Default focus areas
+            if focus_areas is None:
+                focus_areas = ['distributions', 'correlations', 'outliers']
+            
+            # Generate distribution charts
+            if 'distributions' in focus_areas and chart_count < max_charts:
+                numeric_cols = df.select_dtypes(include=[np.number]).columns[:5]
+                for col in numeric_cols:
+                    if chart_count >= max_charts:
+                        break
+                    chart_info = self.visualize(df, ChartType.HISTOGRAM, x=col, title=f"Distribution of {col}")
+                    generated_charts.append(chart_info)
+                    chart_count += 1
+            
+            # Generate correlation matrix
+            if 'correlations' in focus_areas and chart_count < max_charts:
+                numeric_cols = df.select_dtypes(include=[np.number]).columns
+                if len(numeric_cols) >= 2:
+                    chart_info = self.visualize(df, ChartType.CORRELATION_MATRIX, title="Correlation Matrix")
+                    generated_charts.append(chart_info)
+                    chart_count += 1
+            
+            # Generate box plots for outlier detection
+            if 'outliers' in focus_areas and chart_count < max_charts:
+                numeric_cols = df.select_dtypes(include=[np.number]).columns[:3]
+                for col in numeric_cols:
+                    if chart_count >= max_charts:
+                        break
+                    chart_info = self.visualize(df, ChartType.BOX, y=col, title=f"Box Plot of {col}")
+                    generated_charts.append(chart_info)
+                    chart_count += 1
+            
+            return {
+                'generated_charts': generated_charts,
+                'total_charts': len(generated_charts),
+                'focus_areas': focus_areas
+            }
+            
+        except Exception as e:
+            self.logger.error(f"Error in auto visualization: {e}")
+            raise VisualizationError(f"Auto visualization failed: {e}")
+    
+    def recommend_chart_type(
+        self,
+        data: Union[Dict[str, Any], List[Dict[str, Any]], pd.DataFrame],
+        x: Optional[str] = None,
+        y: Optional[str] = None
+    ) -> Dict[str, Any]:
+        """
+        Recommend appropriate chart type based on data characteristics.
+        
+        Args:
+            data: Data for analysis
+            x: X column name
+            y: Y column name
+            
+        Returns:
+            Dict containing recommended chart type and reasoning
+        """
+        try:
+            df = self._to_dataframe(data)
+            chart_type, reason = self._recommend_chart_type(df, x, y)
+            
+            return {
+                'recommended_chart': chart_type.value,
+                'reason': reason,
+                'confidence': 'high'
+            }
+            
+        except Exception as e:
+            self.logger.error(f"Error recommending chart type: {e}")
+            raise VisualizationError(f"Chart recommendation failed: {e}")
+    
+    # Internal helper methods
+    
+    def _to_dataframe(self, data: Union[Dict, List, pd.DataFrame]) -> pd.DataFrame:
+        """Convert data to DataFrame"""
+        if isinstance(data, pd.DataFrame):
+            return data
+        elif isinstance(data, list):
+            return pd.DataFrame(data)
+        elif isinstance(data, dict):
+            return pd.DataFrame([data])
+        else:
+            raise VisualizationError(f"Unsupported data type: {type(data)}")
+    
+    def _recommend_chart_type(self, df: pd.DataFrame, x: Optional[str], y: Optional[str]) -> tuple:
+        """Recommend chart type based on data characteristics"""
+        # If no columns specified, recommend based on data structure
+        if x is None and y is None:
+            numeric_cols = df.select_dtypes(include=[np.number]).columns
+            if len(numeric_cols) >= 2:
+                return ChartType.CORRELATION_MATRIX, "Multiple numeric columns detected"
+            elif len(numeric_cols) == 1:
+                return ChartType.HISTOGRAM, "Single numeric column detected"
+            else:
+                return ChartType.BAR, "Categorical data detected"
+        
+        # Determine column types
+        x_is_numeric = x and df[x].dtype in ['int64', 'float64'] if x in df.columns else False
+        y_is_numeric = y and df[y].dtype in ['int64', 'float64'] if y in df.columns else False
+        
+        # Both numeric: scatter or line
+        if x_is_numeric and y_is_numeric:
+            # Check if x looks like time series
+            if x and 'date' in x.lower() or 'time' in x.lower():
+                return ChartType.TIME_SERIES, "Time series data detected"
+            return ChartType.SCATTER, "Two numeric variables"
+        
+        # One numeric, one categorical: bar or box
+        if (x_is_numeric and not y_is_numeric) or (not x_is_numeric and y_is_numeric):
+            return ChartType.BAR, "Mix of numeric and categorical"
+        
+        # Both categorical: bar
+        if x and y and not x_is_numeric and not y_is_numeric:
+            return ChartType.BAR, "Two categorical variables"
+        
+        # Single numeric: histogram
+        if (x_is_numeric and y is None) or (y_is_numeric and x is None):
+            return ChartType.HISTOGRAM, "Single numeric variable distribution"
+        
+        # Default
+        return ChartType.BAR, "Default chart type"
+    
+    def _create_chart_with_tool(self, df: pd.DataFrame, chart_type: ChartType, x: Optional[str], 
+                                 y: Optional[str], hue: Optional[str], title: Optional[str], 
+                                 output_path: str) -> Dict[str, Any]:
+        """Create chart using chart_tool"""
+        chart_tool = self.external_tools['chart']
+        
+        # Convert chart type to chart_tool format
+        chart_config = {
+            'data': df.to_dict('records'),
+            'chart_type': chart_type.value,
+            'x': x,
+            'y': y,
+            'title': title or f"{chart_type.value.title()} Chart",
+            'output_path': output_path
+        }
+        
+        try:
+            result = chart_tool.run('create_chart', **chart_config)
+            return result
+        except Exception as e:
+            self.logger.warning(f"chart_tool failed, falling back to matplotlib: {e}")
+            return self._create_chart_matplotlib(df, chart_type, x, y, hue, title, output_path)
+    
+    def _create_chart_matplotlib(self, df: pd.DataFrame, chart_type: ChartType, x: Optional[str],
+                                  y: Optional[str], hue: Optional[str], title: Optional[str],
+                                  output_path: str) -> Dict[str, Any]:
+        """Create chart using matplotlib as fallback"""
+        import matplotlib.pyplot as plt
+        
+        fig, ax = plt.subplots(figsize=self.config.default_figsize)
+        
+        if chart_type == ChartType.HISTOGRAM and x:
+            ax.hist(df[x].dropna(), bins=30, edgecolor='black')
+            ax.set_xlabel(x)
+            ax.set_ylabel('Frequency')
+        elif chart_type == ChartType.SCATTER and x and y:
+            ax.scatter(df[x], df[y], alpha=0.6)
+            ax.set_xlabel(x)
+            ax.set_ylabel(y)
+        elif chart_type == ChartType.BAR and x and y:
+            df_grouped = df.groupby(x)[y].mean()
+            df_grouped.plot(kind='bar', ax=ax)
+            ax.set_xlabel(x)
+            ax.set_ylabel(y)
+        elif chart_type == ChartType.CORRELATION_MATRIX:
+            numeric_df = df.select_dtypes(include=[np.number])
+            if len(numeric_df.columns) >= 2:
+                corr = numeric_df.corr()
+                im = ax.imshow(corr, cmap='coolwarm', aspect='auto')
+                ax.set_xticks(range(len(corr.columns)))
+                ax.set_yticks(range(len(corr.columns)))
+                ax.set_xticklabels(corr.columns, rotation=45, ha='right')
+                ax.set_yticklabels(corr.columns)
+                plt.colorbar(im, ax=ax)
+        elif chart_type == ChartType.BOX and y:
+            df.boxplot(column=y, ax=ax)
+        else:
+            # Default: simple line plot
+            if x and y:
+                ax.plot(df[x], df[y])
+            elif x:
+                ax.plot(df[x])
+        
+        if title:
+            ax.set_title(title)
+        
+        plt.tight_layout()
+        plt.savefig(output_path, dpi=self.config.default_dpi, bbox_inches='tight')
+        plt.close()
+        
+        return {
+            'status': 'success',
+            'output_path': output_path,
+            'chart_type': chart_type.value
+        }
+