aiecs 1.5.1__py3-none-any.whl

aiecs/tools/task_tools/stats_tool.py

@@ -0,0 +1,688 @@
+import os
+import logging
+import tempfile
+from typing import Dict, Any, List, Optional, Union, Tuple
+from enum import Enum
+from dataclasses import dataclass
+
+import pandas as pd
+import numpy as np
+from pydantic import BaseModel, ConfigDict, Field
+
+from aiecs.tools.base_tool import BaseTool
+from aiecs.tools import register_tool
+
+# Enums for configuration options
+
+
+class ScalerType(str, Enum):
+    STANDARD = "standard"
+    MINMAX = "minmax"
+    ROBUST = "robust"
+    NONE = "none"
+
+
+# Exceptions
+class StatsToolError(Exception):
+    pass
+
+
+class FileOperationError(StatsToolError):
+    pass
+
+
+class AnalysisError(StatsToolError):
+    pass
+
+
+# Utility Dataclass for Statistical Results
+
+
+@dataclass
+class StatsResult:
+    """Structured statistical result."""
+
+    test_type: str
+    statistic: float
+    pvalue: float
+    significant: bool
+    additional_metrics: Dict[str, Any]
+
+    def to_dict(self) -> Dict[str, Any]:
+        return {
+            "test_type": self.test_type,
+            "statistic": self.statistic,
+            "pvalue": self.pvalue,
+            "significant": self.significant,
+            **self.additional_metrics,
+        }
+
+
+@register_tool("stats")
+class StatsTool(BaseTool):
+    """Enhanced statistical analysis tool for various data formats and operations."""
+
+    # Configuration schema
+    class Config(BaseModel):
+        """Configuration for the stats tool"""
+
+        model_config = ConfigDict(env_prefix="STATS_TOOL_")
+
+        max_file_size_mb: int = Field(default=200, description="Maximum file size in megabytes")
+        allowed_extensions: List[str] = Field(
+            default=[
+                ".sav",
+                ".sas7bdat",
+                ".por",
+                ".csv",
+                ".xlsx",
+                ".xls",
+                ".json",
+                ".parquet",
+                ".feather",
+            ],
+            description="Allowed file extensions",
+        )
+
+    def __init__(self, config: Dict[str, Any] = None):
+        super().__init__(config)
+
+        # Parse configuration
+        self.config = self.Config(**(config or {}))
+
+        self.logger = logging.getLogger(__name__)
+        if not self.logger.handlers:
+            h = logging.StreamHandler()
+            h.setFormatter(logging.Formatter("%(asctime)s %(levelname)s %(message)s"))
+            self.logger.addHandler(h)
+        self.logger.setLevel(logging.INFO)
+
+    def _load_data(
+        self,
+        file_path: str,
+        nrows: Optional[int] = None,
+        sheet_name: Optional[Union[str, int]] = 0,
+    ) -> pd.DataFrame:
+        """Load data from various file formats into a pandas DataFrame."""
+        try:
+            ext = os.path.splitext(file_path)[1].lower()
+            if ext in [".sav", ".sas7bdat", ".por"]:
+                import pyreadstat
+
+                if ext == ".sav":
+                    df, meta = pyreadstat.read_sav(file_path)
+                elif ext == ".sas7bdat":
+                    df, meta = pyreadstat.read_sas7bdat(file_path)
+                else:
+                    df, meta = pyreadstat.read_por(file_path)
+                return df
+            elif ext == ".csv":
+                return pd.read_csv(file_path, nrows=nrows)
+            elif ext in [".xlsx", ".xls"]:
+                return pd.read_excel(file_path, sheet_name=sheet_name, nrows=nrows)
+            elif ext == ".json":
+                return pd.read_json(file_path)
+            elif ext == ".parquet":
+                return pd.read_parquet(file_path)
+            elif ext == ".feather":
+                return pd.read_feather(file_path)
+            else:
+                raise FileOperationError(f"Unsupported file format: {ext}")
+        except Exception as e:
+            raise FileOperationError(f"Error reading file {file_path}: {str(e)}")
+
+    def _validate_variables(self, df: pd.DataFrame, vars_to_check: List[str]) -> None:
+        """Validate variables exist in the dataset."""
+        if not vars_to_check:
+            return
+        available_vars = df.columns.tolist()
+        missing_vars = [var for var in vars_to_check if var not in available_vars]
+        if missing_vars:
+            raise FileOperationError(f"Variables not found in dataset: {', '.join(missing_vars)}")
+
+    def _interpret_effect_size(self, d: float) -> str:
+        """Interpret Cohen's d or Cramer's V effect size."""
+        thresholds = [(0.2, "negligible"), (0.5, "small"), (0.8, "medium")]
+        for threshold, label in thresholds:
+            if abs(d) < threshold:
+                return label
+        return "large"
+
+    def read_data(
+        self,
+        file_path: str,
+        nrows: Optional[int] = None,
+        sheet_name: Optional[Union[str, int]] = 0,
+    ) -> Dict[str, Any]:
+        """Read data from various file formats."""
+        df = self._load_data(file_path, nrows, sheet_name)
+        return {
+            "variables": df.columns.tolist(),
+            "observations": len(df),
+            "dtypes": {col: str(dtype) for col, dtype in df.dtypes.items()},
+            "memory_usage": df.memory_usage(deep=True).sum() / (1024 * 1024),
+            "preview": df.head(5).to_dict(orient="records"),
+        }
+
+    def describe(
+        self,
+        file_path: str,
+        variables: Optional[List[str]] = None,
+        include_percentiles: bool = False,
+        percentiles: Optional[List[float]] = None,
+    ) -> Dict[str, Any]:
+        """Generate descriptive statistics for variables."""
+        df = self._load_data(file_path)
+        if variables:
+            self._validate_variables(df, variables)
+            df = df[variables]
+        desc = df.describe()
+        if include_percentiles and percentiles:
+            additional_percentiles = [p for p in percentiles if p not in [0.25, 0.5, 0.75]]
+            if additional_percentiles:
+                additional_desc = df.describe(percentiles=percentiles)
+                desc = pd.concat(
+                    [
+                        desc,
+                        additional_desc.loc[[f"{int(p*100)}%" for p in additional_percentiles]],
+                    ]
+                )
+        numeric_cols = df.select_dtypes(include=[np.number]).columns
+        if numeric_cols.any():
+            desc.loc["skew"] = df[numeric_cols].skew()
+            desc.loc["kurtosis"] = df[numeric_cols].kurt()
+        return {"statistics": desc.to_dict(), "summary": desc.to_string()}
+
+    def ttest(
+        self,
+        file_path: str,
+        var1: str,
+        var2: str,
+        equal_var: bool = True,
+        paired: bool = False,
+    ) -> Dict[str, Any]:
+        """Perform t-tests (independent or paired). Also handles legacy ttest_ind."""
+        df = self._load_data(file_path)
+        self._validate_variables(df, [var1, var2])
+        import scipy.stats as stats
+
+        a = df[var1].dropna().values
+        b = df[var2].dropna().values
+        if paired:
+            min_len = min(len(a), len(b))
+            stat, p = stats.ttest_rel(a[:min_len], b[:min_len])
+            test_type = "paired t-test"
+        else:
+            stat, p = stats.ttest_ind(a, b, equal_var=equal_var)
+            test_type = (
+                "independent t-test (equal variance)"
+                if equal_var
+                else "Welch's t-test (unequal variance)"
+            )
+        mean_a = np.mean(a)
+        mean_b = np.mean(b)
+        std_a = np.std(a, ddof=1)
+        std_b = np.std(b, ddof=1)
+        if equal_var:
+            pooled_std = np.sqrt(
+                ((len(a) - 1) * std_a**2 + (len(b) - 1) * std_b**2) / (len(a) + len(b) - 2)
+            )
+            cohens_d = (mean_a - mean_b) / pooled_std
+        else:
+            cohens_d = (mean_a - mean_b) / np.sqrt((std_a**2 + std_b**2) / 2)
+        return StatsResult(
+            test_type=test_type,
+            statistic=float(stat),
+            pvalue=float(p),
+            significant=p < 0.05,
+            additional_metrics={
+                "cohens_d": float(cohens_d),
+                "effect_size_interpretation": self._interpret_effect_size(cohens_d),
+                "group1_mean": float(mean_a),
+                "group2_mean": float(mean_b),
+                "group1_std": float(std_a),
+                "group2_std": float(std_b),
+                "group1_n": int(len(a)),
+                "group2_n": int(len(b)),
+            },
+        ).to_dict()
+
+    # Legacy method (now an alias)
+    ttest_ind = ttest
+
+    def correlation(
+        self,
+        file_path: str,
+        variables: Optional[List[str]] = None,
+        var1: Optional[str] = None,
+        var2: Optional[str] = None,
+        method: str = "pearson",
+    ) -> Dict[str, Any]:
+        """Perform correlation analysis."""
+        df = self._load_data(file_path)
+        if variables:
+            self._validate_variables(df, variables)
+        if var1 and var2:
+            self._validate_variables(df, [var1, var2])
+        import scipy.stats as stats
+
+        result = {}
+        if variables:
+            corr_matrix = df[variables].corr(method=method)
+            result["correlation_matrix"] = corr_matrix.to_dict()
+            flat_corrs = [
+                {
+                    "var1": v1,
+                    "var2": v2,
+                    "correlation": corr_matrix.loc[v1, v2],
+                    "abs_correlation": abs(corr_matrix.loc[v1, v2]),
+                }
+                for i, v1 in enumerate(variables)
+                for j, v2 in enumerate(variables)
+                if i < j
+            ]
+            flat_corrs.sort(key=lambda x: x["abs_correlation"], reverse=True)
+            result["pairs"] = flat_corrs
+        elif var1 and var2:
+            x = df[var1].dropna()
+            y = df[var2].dropna()
+            method_map = {
+                "pearson": (stats.pearsonr, "Pearson's r"),
+                "spearman": (stats.spearmanr, "Spearman's rho"),
+                "kendall": (stats.kendalltau, "Kendall's tau"),
+            }
+            func, method_name = method_map[method]
+            corr, p = func(x, y)
+            result = {
+                "method": method_name,
+                "correlation": float(corr),
+                "pvalue": float(p),
+                "significant": p < 0.05,
+                "n": len(x),
+            }
+        return result
+
+    def anova(
+        self,
+        file_path: str,
+        dependent: str,
+        factor: str,
+        post_hoc: bool = False,
+    ) -> Dict[str, Any]:
+        """Perform one-way ANOVA with optional post-hoc tests."""
+        df = self._load_data(file_path)
+        self._validate_variables(df, [dependent, factor])
+        import scipy.stats as stats
+        from statsmodels.stats.multicomp import pairwise_tukeyhsd
+
+        dependent_var = df[dependent].dropna()
+        factor_var = df[factor].dropna()
+        min_len = min(len(dependent_var), len(factor_var))
+        dependent_var = dependent_var[:min_len]
+        factor_var = factor_var[:min_len]
+        groups = {name: group[dependent].dropna().values for name, group in df.groupby(factor)}
+        stat, p = stats.f_oneway(*groups.values())
+        result = {
+            "F": float(stat),
+            "pvalue": float(p),
+            "significant": p < 0.05,
+            "groups": len(groups),
+            "group_sizes": {name: len(values) for name, values in groups.items()},
+            "group_means": {name: float(np.mean(values)) for name, values in groups.items()},
+            "group_std": {name: float(np.std(values, ddof=1)) for name, values in groups.items()},
+        }
+        if post_hoc:
+            post_hoc_df = pd.DataFrame({"value": dependent_var, "group": factor_var})
+            tukey = pairwise_tukeyhsd(post_hoc_df["value"], post_hoc_df["group"])
+            from itertools import combinations
+
+            group_pairs = list(combinations(tukey.groupsunique, 2))
+            tukey_results = [
+                {
+                    "group1": str(group1),
+                    "group2": str(group2),
+                    "mean_difference": float(mean_diff),
+                    "p_adjusted": float(p_adj),
+                    "significant": bool(reject),
+                    "conf_lower": float(lower),
+                    "conf_upper": float(upper),
+                }
+                for (
+                    group1,
+                    group2,
+                ), mean_diff, p_adj, lower, upper, reject in zip(
+                    group_pairs,
+                    tukey.meandiffs,
+                    tukey.pvalues,
+                    tukey.confint[:, 0],
+                    tukey.confint[:, 1],
+                    tukey.reject,
+                )
+            ]
+            result["post_hoc"] = {
+                "method": "Tukey HSD",
+                "alpha": 0.05,  # Standard significance level for Tukey HSD
+                "comparisons": tukey_results,
+            }
+        return result
+
+    def chi_square(
+        self, file_path: str, var1: str, var2: str, correction: bool = True
+    ) -> Dict[str, Any]:
+        """Perform chi-square test of independence."""
+        df = self._load_data(file_path)
+        self._validate_variables(df, [var1, var2])
+        import scipy.stats as stats
+
+        contingency = pd.crosstab(df[var1], df[var2])
+        chi2, p, dof, expected = stats.chi2_contingency(contingency, correction=correction)
+        n = contingency.sum().sum()
+        min_dim = min(contingency.shape) - 1
+        cramers_v = np.sqrt(chi2 / (n * min_dim))
+        return {
+            "chi2": float(chi2),
+            "pvalue": float(p),
+            "dof": int(dof),
+            "significant": p < 0.05,
+            "cramers_v": float(cramers_v),
+            "effect_size_interpretation": self._interpret_effect_size(cramers_v),
+            "contingency_table": contingency.to_dict(),
+            "expected_frequencies": pd.DataFrame(
+                expected, index=contingency.index, columns=contingency.columns
+            ).to_dict(),
+            "test_type": (
+                "Chi-square test with Yates correction" if correction else "Chi-square test"
+            ),
+        }
+
+    def non_parametric(
+        self,
+        file_path: str,
+        test_type: str,
+        variables: List[str],
+        grouping: Optional[str] = None,
+    ) -> Dict[str, Any]:
+        """Perform non-parametric statistical tests."""
+        df = self._load_data(file_path)
+        self._validate_variables(df, variables + ([grouping] if grouping else []))
+        import scipy.stats as stats
+
+        if test_type == "mann_whitney":
+            if len(variables) != 2:
+                raise AnalysisError("Mann-Whitney U test requires exactly 2 variables")
+            x = df[variables[0]].dropna().values
+            y = df[variables[1]].dropna().values
+            u_stat, p_value = stats.mannwhitneyu(x, y)
+            return StatsResult(
+                test_type="Mann-Whitney U test",
+                statistic=float(u_stat),
+                pvalue=float(p_value),
+                significant=p_value < 0.05,
+                additional_metrics={
+                    "n1": len(x),
+                    "n2": len(y),
+                    "median1": float(np.median(x)),
+                    "median2": float(np.median(y)),
+                },
+            ).to_dict()
+        elif test_type == "wilcoxon":
+            if len(variables) != 2:
+                raise AnalysisError("Wilcoxon signed-rank test requires exactly 2 variables")
+            x = df[variables[0]].dropna().values
+            y = df[variables[1]].dropna().values
+            min_len = min(len(x), len(y))
+            x = x[:min_len]
+            y = y[:min_len]
+            w_stat, p_value = stats.wilcoxon(x, y)
+            return StatsResult(
+                test_type="Wilcoxon signed-rank test",
+                statistic=float(w_stat),
+                pvalue=float(p_value),
+                significant=p_value < 0.05,
+                additional_metrics={
+                    "n_pairs": min_len,
+                    "median_difference": float(np.median(x - y)),
+                },
+            ).to_dict()
+        elif test_type == "kruskal":
+            if not grouping:
+                raise AnalysisError("Kruskal-Wallis test requires a grouping variable")
+            groups = {
+                f"{var}_{name}": group[var].dropna().values
+                for name, group in df.groupby(grouping)
+                for var in variables
+            }
+            h_stat, p_value = stats.kruskal(*groups.values())
+            return StatsResult(
+                test_type="Kruskal-Wallis H test",
+                statistic=float(h_stat),
+                pvalue=float(p_value),
+                significant=p_value < 0.05,
+                additional_metrics={
+                    "groups": len(groups),
+                    "group_sizes": {name: len(values) for name, values in groups.items()},
+                    "group_medians": {
+                        name: float(np.median(values)) for name, values in groups.items()
+                    },
+                },
+            ).to_dict()
+        elif test_type == "friedman":
+            if len(variables) < 2:
+                raise AnalysisError("Friedman test requires at least 2 variables")
+            data = df[variables].dropna()
+            chi2, p_value = stats.friedmanchisquare(*[data[var].values for var in variables])
+            return StatsResult(
+                test_type="Friedman test",
+                statistic=float(chi2),
+                pvalue=float(p_value),
+                significant=p_value < 0.05,
+                additional_metrics={
+                    "n_measures": len(variables),
+                    "n_samples": len(data),
+                    "variable_medians": {var: float(np.median(data[var])) for var in variables},
+                },
+            ).to_dict()
+        else:
+            raise AnalysisError(
+                f"Unsupported non-parametric test type: {test_type}. Supported types: mann_whitney, wilcoxon, kruskal, friedman"
+            )
+
+    def regression(
+        self,
+        file_path: str,
+        formula: str,
+        regression_type: str = "ols",
+        robust: bool = False,
+        structured_output: bool = True,
+    ) -> Dict[str, Any]:
+        """Perform regression analysis with various models."""
+        df = self._load_data(file_path)
+        import statsmodels.formula.api as smf
+
+        try:
+            model_map = {
+                "ols": smf.ols,
+                "logit": smf.logit,
+                "probit": smf.probit,
+                "poisson": smf.poisson,
+            }
+            model = model_map[regression_type](formula=formula, data=df)
+            fit = model.fit(cov_type="HC3" if robust else "nonrobust")
+            if structured_output:
+                result = {
+                    "model_type": regression_type,
+                    "formula": formula,
+                    "n_observations": int(fit.nobs),
+                    "r_squared": (float(fit.rsquared) if hasattr(fit, "rsquared") else None),
+                    "adj_r_squared": (
+                        float(fit.rsquared_adj) if hasattr(fit, "rsquared_adj") else None
+                    ),
+                    "aic": float(fit.aic) if hasattr(fit, "aic") else None,
+                    "bic": float(fit.bic) if hasattr(fit, "bic") else None,
+                    "f_statistic": (float(fit.fvalue) if hasattr(fit, "fvalue") else None),
+                    "f_pvalue": (float(fit.f_pvalue) if hasattr(fit, "f_pvalue") else None),
+                    "log_likelihood": (float(fit.llf) if hasattr(fit, "llf") else None),
+                    "coefficients": {
+                        var: {
+                            "coef": float(fit.params[var]),
+                            "std_err": float(fit.bse[var]),
+                            "t_value": (
+                                float(fit.tvalues[var]) if hasattr(fit, "tvalues") else None
+                            ),
+                            "p_value": float(fit.pvalues[var]),
+                            "significant": fit.pvalues[var] < 0.05,
+                            "conf_lower": float(fit.conf_int().loc[var, 0]),
+                            "conf_upper": float(fit.conf_int().loc[var, 1]),
+                        }
+                        for var in fit.params.index
+                    },
+                }
+                return {
+                    "summary_text": fit.summary().as_text(),
+                    "structured": result,
+                }
+            return {"summary": fit.summary().as_text()}
+        except Exception as e:
+            raise AnalysisError(f"Regression error: {str(e)}")
+
+    def time_series(
+        self,
+        file_path: str,
+        variable: str,
+        date_variable: Optional[str] = None,
+        model_type: str = "arima",
+        order: Optional[Tuple[int, int, int]] = (1, 1, 1),
+        seasonal_order: Optional[Tuple[int, int, int, int]] = None,
+        forecast_periods: int = 10,
+    ) -> Dict[str, Any]:
+        """Perform time series analysis."""
+        df = self._load_data(file_path)
+        self._validate_variables(df, [variable] + ([date_variable] if date_variable else []))
+        from statsmodels.tsa.arima.model import ARIMA
+        from statsmodels.tsa.statespace.sarimax import SARIMAX
+
+        try:
+            ts_data = df[variable].dropna()
+            if date_variable and date_variable in df.columns:
+                ts_data.index = df[date_variable]
+            if model_type == "arima":
+                model = ARIMA(ts_data, order=order)
+                fit = model.fit()
+                model_type_name = "ARIMA"
+            elif model_type == "sarima":
+                if not seasonal_order:
+                    raise AnalysisError("seasonal_order must be provided for SARIMA model")
+                model = SARIMAX(ts_data, order=order, seasonal_order=seasonal_order)
+                fit = model.fit(disp=False)
+                model_type_name = "SARIMA"
+            else:
+                raise AnalysisError(f"Unsupported time series model: {model_type}")
+            forecast = fit.forecast(steps=forecast_periods)
+            forecast_index = pd.date_range(
+                start=(
+                    ts_data.index[-1]
+                    if isinstance(ts_data.index, pd.DatetimeIndex)
+                    else len(ts_data)
+                ),
+                periods=forecast_periods + 1,
+                freq="D",
+            )[1:]
+            return {
+                "model_type": model_type_name,
+                "order": order,
+                "seasonal_order": (seasonal_order if model_type == "sarima" else None),
+                "aic": float(fit.aic),
+                "bic": float(fit.bic),
+                "forecast": {
+                    "values": (
+                        forecast.tolist()
+                        if isinstance(forecast, np.ndarray)
+                        else forecast.values.tolist()
+                    ),
+                    "index": (
+                        forecast_index.strftime("%Y-%m-%d").tolist()
+                        if isinstance(forecast_index, pd.DatetimeIndex)
+                        else list(range(len(forecast)))
+                    ),
+                },
+                "summary": str(fit.summary()),
+            }
+        except Exception as e:
+            raise AnalysisError(f"Time series analysis error: {str(e)}")
+
+    def preprocess(
+        self,
+        file_path: str,
+        variables: List[str],
+        operation: str,
+        scaler_type: ScalerType = ScalerType.STANDARD,
+        output_path: Optional[str] = None,
+    ) -> Dict[str, Any]:
+        """Preprocess data with various operations."""
+        df = self._load_data(file_path)
+        self._validate_variables(df, variables)
+        data = df[variables].copy()
+        result = {"operation": operation}
+        if operation == "scale":
+            from sklearn.preprocessing import (
+                StandardScaler,
+                MinMaxScaler,
+                RobustScaler,
+            )
+
+            scaler_map = {
+                ScalerType.STANDARD: (StandardScaler, "StandardScaler"),
+                ScalerType.MINMAX: (MinMaxScaler, "MinMaxScaler"),
+                ScalerType.ROBUST: (RobustScaler, "RobustScaler"),
+            }
+            scaler_cls, scaler_name = scaler_map[scaler_type]
+            scaler = scaler_cls()
+            scaled_data = scaler.fit_transform(data)
+            scaled_df = pd.DataFrame(
+                scaled_data,
+                columns=[f"{col}_scaled" for col in data.columns],
+                index=data.index,
+            )
+            result.update(
+                {
+                    "scaler": scaler_name,
+                    "original_stats": data.describe().to_dict(),
+                    "scaled_stats": scaled_df.describe().to_dict(),
+                    "preview": scaled_df.head(5).to_dict(orient="records"),
+                }
+            )
+            processed_df = scaled_df
+        elif operation == "impute":
+            import numpy as np
+
+            imputed_df = data.copy()
+            numeric_cols = data.select_dtypes(include=[np.number]).columns
+            for col in numeric_cols:
+                imputed_df[col] = data[col].fillna(data[col].mean())
+            cat_cols = data.select_dtypes(exclude=[np.number]).columns
+            for col in cat_cols:
+                imputed_df[col] = data[col].fillna(
+                    data[col].mode()[0] if not data[col].mode().empty else None
+                )
+            result.update(
+                {
+                    "imputation_method": {
+                        "numeric": "mean",
+                        "categorical": "mode",
+                    },
+                    "missing_counts_before": data.isna().sum().to_dict(),
+                    "missing_counts_after": imputed_df.isna().sum().to_dict(),
+                    "preview": imputed_df.head(5).to_dict(orient="records"),
+                }
+            )
+            processed_df = imputed_df
+        if output_path:
+            output_path = (
+                os.path.abspath(output_path)
+                if os.path.isabs(output_path)
+                else os.path.join(tempfile.gettempdir(), "stats_outputs", output_path)
+            )
+            os.makedirs(os.path.dirname(output_path), exist_ok=True)
+            processed_df.to_csv(output_path)
+            result["output_file"] = output_path
+        return result