Update tools/visuals.py

tools/visuals.py

@@ -1,132 +1,163 @@
 import os
 import tempfile
 import pandas as pd
+import numpy as np
 import plotly.express as px
-from typing import Union, Tuple
+import plotly.figure_factory as ff
+import plotly.graph_objects as go
+from scipy.cluster.hierarchy import linkage, leaves_list
+from typing import Union, Tuple, List
 
-def _save_fig(fig, prefix: str, output_dir: str) -> str:
+
+def _save_fig(fig: go.Figure, prefix: str, output_dir: str) -> str:
     """
-    Save a Plotly figure as a PNG to a temp file and return its path.
+    Save a Plotly figure as a high-res PNG and return the file path.
     """
     os.makedirs(output_dir, exist_ok=True)
     tmp = tempfile.NamedTemporaryFile(suffix='.png', prefix=prefix, dir=output_dir, delete=False)
     path = tmp.name
     tmp.close()
-    try:
-        fig.write_image(path, scale=2)
-    except Exception as e:
-        raise
+    fig.write_image(path, scale=3)
     return path
 
 
 def histogram_tool(
     file_path: str,
     column: str,
-    output_dir: str = '/tmp',
-    bins: int = 30
-) -> Union[Tuple[px.Figure, str], str]:
+    bins: int = 30,
+    kde: bool = True,
+    output_dir: str = '/tmp'
+) -> Union[Tuple[ff.FigureFactory, str], str]:
     """
-    Build a histogram for a numeric column, return a Plotly Figure and PNG path,
-    or an error string starting with '❌'.
+    Create a histogram with optional KDE overlay for a given numeric column.
+
+    Returns (figure, png_path) or error string.
     """
-    # Load data
+    # Load
     ext = os.path.splitext(file_path)[1].lower()
-    try:
-        df = pd.read_excel(file_path) if ext in ('.xls', '.xlsx') else pd.read_csv(file_path)
-    except Exception as exc:
-        return f"❌ Failed to load file: {exc}"
+    df = pd.read_excel(file_path) if ext in ('.xls','.xlsx') else pd.read_csv(file_path)
 
-    # Validate column
+    # Validate
     if column not in df.columns:
         return f"❌ Column '{column}' not found."
-
-    # Coerce to numeric
-    df[column] = pd.to_numeric(df[column], errors='coerce')
-    series = df[column].dropna()
+    series = pd.to_numeric(df[column], errors='coerce').dropna()
     if series.empty:
-        return f"❌ No valid numeric data in '{column}'."
-
-    # Create figure
-    fig = px.histogram(
-        df,
-        x=column,
-        nbins=bins,
-        title=f"Histogram – {column}",
-        template='plotly_dark'
-    )
-    # Save PNG
+        return f"❌ No numeric data in '{column}'."
+
+    # Build histogram + KDE
+    if kde:
+        fig = ff.create_distplot([series], [column], bin_size=(series.max()-series.min())/bins)
+    else:
+        fig = px.histogram(series, nbins=bins, title=f"Histogram – {column}", template='plotly_dark')
+    fig.update_layout(template='plotly_dark')
+
+    # Save
     img_path = _save_fig(fig, f"hist_{column}_", output_dir)
     return fig, img_path
 
 
-def scatter_matrix_tool(
+def boxplot_tool(
     file_path: str,
-    cols: list[str],
+    column: str,
     output_dir: str = '/tmp'
 ) -> Union[Tuple[px.Figure, str], str]:
     """
-    Build a scatter-matrix for selected numeric columns, return figure and PNG path,
-    or an error string starting with '❌'.
+    Create a box plot with outliers for a numeric column.
+
+    Returns (figure, png_path) or error string.
     """
-    # Load data
     ext = os.path.splitext(file_path)[1].lower()
-    try:
-        df = pd.read_excel(file_path) if ext in ('.xls', '.xlsx') else pd.read_csv(file_path)
-    except Exception as exc:
-        return f"❌ Failed to load file: {exc}"
+    df = pd.read_excel(file_path) if ext in ('.xls','.xlsx') else pd.read_csv(file_path)
+    if column not in df.columns:
+        return f"❌ Column '{column}' not found."
+    series = pd.to_numeric(df[column], errors='coerce').dropna()
+    if series.empty:
+        return f"❌ No numeric data in '{column}'."
+
+    fig = px.box(series, points='outliers', title=f"Boxplot – {column}", template='plotly_dark')
+    img_path = _save_fig(fig, f"box_{column}_", output_dir)
+    return fig, img_path
 
-    # Validate columns
-    missing = [c for c in cols if c not in df.columns]
+
+def violin_tool(
+    file_path: str,
+    column: str,
+    output_dir: str = '/tmp'
+) -> Union[Tuple[px.Figure, str], str]:
+    """
+    Create a violin plot with inner box for a numeric column.
+
+    Returns (figure, png_path) or error string.
+    """
+    ext = os.path.splitext(file_path)[1].lower()
+    df = pd.read_excel(file_path) if ext in ('.xls','.xlsx') else pd.read_csv(file_path)
+    if column not in df.columns:
+        return f"❌ Column '{column}' not found."
+    series = pd.to_numeric(df[column], errors='coerce').dropna()
+    if series.empty:
+        return f"❌ No numeric data in '{column}'."
+
+    fig = px.violin(series, box=True, points='all', title=f"Violin – {column}", template='plotly_dark')
+    img_path = _save_fig(fig, f"violin_{column}_", output_dir)
+    return fig, img_path
+
+
+def scatter_matrix_tool(
+    file_path: str,
+    columns: List[str],
+    output_dir: str = '/tmp',
+    size: int = 5
+) -> Union[Tuple[px.Figure, str], str]:
+    """
+    Create an interactive scatter matrix for selected numeric columns.
+
+    Returns (figure, png_path) or error string.
+    """
+    ext = os.path.splitext(file_path)[1].lower()
+    df = pd.read_excel(file_path) if ext in ('.xls','.xlsx') else pd.read_csv(file_path)
+    missing = [c for c in columns if c not in df.columns]
     if missing:
         return f"❌ Missing columns: {', '.join(missing)}"
-
-    # Filter numeric
-    df_num = df[cols].apply(pd.to_numeric, errors='coerce').dropna()
+    df_num = df[columns].apply(pd.to_numeric, errors='coerce').dropna()
     if df_num.empty:
-        return f"❌ No valid numeric data in selected columns."
+        return "❌ No valid numeric data."
 
-    # Create figure
-    fig = px.scatter_matrix(
-        df_num,
-        dimensions=cols,
-        title="Scatter-Matrix",
-        template='plotly_dark'
-    )
-    # Save PNG
+    fig = px.scatter_matrix(df_num, dimensions=columns, title="Scatter Matrix", template='plotly_dark')
+    fig.update_traces(diagonal_visible=False, marker={'size': size})
     img_path = _save_fig(fig, "scatter_matrix_", output_dir)
     return fig, img_path
 
 
 def corr_heatmap_tool(
     file_path: str,
+    columns: List[str] = None,
     output_dir: str = '/tmp',
-    color_continuous_scale: str = 'RdBu'
+    cluster: bool = True
 ) -> Union[Tuple[px.Figure, str], str]:
     """
-    Build a correlation heatmap for numeric columns, return figure and PNG path,
-    or an error string starting with '❌'.
+    Create a correlation heatmap, with optional hierarchical clustering of variables.
+
+    Returns (figure, png_path) or error string.
     """
-    # Load data
     ext = os.path.splitext(file_path)[1].lower()
-    try:
-        df = pd.read_excel(file_path) if ext in ('.xls', '.xlsx') else pd.read_csv(file_path)
-    except Exception as exc:
-        return f"❌ Failed to load file: {exc}"
+    df = pd.read_excel(file_path) if ext in ('.xls','.xlsx') else pd.read_csv(file_path)
+    df_num = df.select_dtypes(include='number') if columns is None else df[columns]
+    df_num = df_num.apply(pd.to_numeric, errors='coerce').dropna(axis=1, how='all')
+    if df_num.shape[1] < 2:
+        return "❌ Need at least two numeric columns for correlation."
 
-    # Compute correlation
-    df_num = df.select_dtypes(include='number').apply(pd.to_numeric, errors='coerce')
-    if df_num.empty:
-        return "❌ No numeric columns available for correlation."
     corr = df_num.corr()
+    if cluster:
+        link = linkage(corr, method='average')
+        order = leaves_list(link)
+        corr = corr.iloc[order, order]
 
-    # Create figure
     fig = px.imshow(
         corr,
-        color_continuous_scale=color_continuous_scale,
+        color_continuous_scale='RdBu',
         title="Correlation Heatmap",
         labels=dict(color="Correlation"),
         template='plotly_dark'
     )
-    # Save PNG
     img_path = _save_fig(fig, "corr_heatmap_", output_dir)
     return fig, img_path