goad-toolkit 0.1.0__py3-none-any.whl

goad_toolkit/analytics.py

@@ -0,0 +1,325 @@
+from dataclasses import dataclass
+from typing import Any, List, Optional, Tuple, Union
+
+import numpy as np
+from loguru import logger
+from scipy import stats
+
+from goad_toolkit.distributions import DistributionRegistry
+
+
+@dataclass
+class KSTestResult:
+    """Result of a Kolmogorov-Smirnov test."""
+
+    statistic: float
+    p_value: float
+    error: Optional[str] = None
+
+
+@dataclass
+class FitResult:
+    """Result of a successful distribution fit."""
+
+    distribution: str  # Distribution name
+    dist_object: Any  # The actual distribution object
+    params: Tuple[float, ...]  # Fitted parameters
+    frozen_dist: Any  # Frozen distribution with parameters applied
+    success: bool = True
+    message: str = "Optimization successful"
+    log_likelihood: Optional[float] = None
+    kstest: Optional[KSTestResult] = None
+    best_likelihood: Optional[bool] = (
+        None  # Indicates if this is the best fit by likelihood
+    )
+    best_ks: Optional[bool] = None  # Indicates if this is the best fit by KS test
+
+    def __str__(self) -> str:
+        """String representation focused on distribution and parameters."""
+        return f"{self.distribution}: params={self.params}"
+
+    def __repr__(self) -> str:
+        """Detailed representation including fit quality metrics."""
+        ks_info = f", ks_pvalue={self.kstest.p_value:.4f}" if self.kstest else ""
+        ll_info = (
+            f", loglik={self.log_likelihood:.2f}"
+            if self.log_likelihood is not None
+            else ""
+        )
+        best_info = ""
+        if self.best_likelihood:
+            best_info += " (Best Likelihood)"
+        if self.best_ks:
+            best_info += " (Best KS)"
+
+        return f"Fit({self.distribution}: params={self.params}{ll_info}{ks_info}{best_info})"
+
+
+@dataclass
+class FailedFit:
+    """Result of a failed distribution fit."""
+
+    distribution: str
+    message: str
+    success: bool = False
+
+    def __str__(self) -> str:
+        """String representation with error message."""
+        return f"Failed {self.distribution}: {self.message}"
+
+
+# Type alias for either success or failure
+Result = Union[FitResult, FailedFit]
+
+
+class DistributionFitter:
+    """Class to fit distributions from a registry to data."""
+
+    def __init__(self) -> None:
+        """Initialize with a distribution registry."""
+        self._registry = DistributionRegistry()
+
+    @property
+    def registry(self) -> list[str]:
+        return self._registry.get_names()
+
+    def _get_bounds(self, data: np.ndarray, dist_obj) -> list[tuple]:
+        """Generate parameter bounds estimates based on data characteristics."""
+        # Calculate statistics for bounds
+        data_min = np.min(data)
+        data_max = np.max(data)
+        data_mean = np.mean(data)
+        data_std = np.std(data)
+        data_range = data_max - data_min
+
+        # Universal parameter bounds
+        loc_bound = (data_min - data_range, data_max + data_range)
+        std_lower_bound = float(data_std / 100)
+        scale_bound = (max(std_lower_bound, 0.001), data_std * 20)
+        shape_bound = (0.01, 10.0)
+
+        # Create bounds based on parameter count
+        if dist_obj.num_params == 1:
+            if dist_obj.is_discrete:
+                mean_lower_bound = max(0.1, float(data_mean / 10))
+                mean_upper_bound = float(data_mean * 10)
+                return [(mean_lower_bound, mean_upper_bound)]
+            else:
+                return [scale_bound]
+        elif dist_obj.num_params == 2:
+            return [loc_bound, scale_bound]
+        elif dist_obj.num_params == 3:
+            return [shape_bound, loc_bound, scale_bound]
+        else:
+            return [shape_bound] * (dist_obj.num_params - 2) + [loc_bound, scale_bound]
+
+    def _perform_kstest(self, data: np.ndarray, dist_obj, params) -> KSTestResult:
+        """Perform Kolmogorov-Smirnov test for goodness-of-fit."""
+        try:
+            # Create a frozen distribution with fitted parameters
+            fitted_dist = dist_obj.dist(*params)
+
+            # Run KS test - comparing data with the fitted distribution
+            ks_statistic, p_value = stats.kstest(data, fitted_dist.cdf)
+
+            return KSTestResult(statistic=ks_statistic, p_value=p_value)
+        except Exception as e:
+            return KSTestResult(
+                statistic=float("nan"), p_value=float("nan"), error=str(e)
+            )
+
+    def _calculate_loglikelihood(self, data: np.ndarray, dist_obj, params) -> float:
+        """Calculate log-likelihood of data given distribution and parameters.
+        For every datapoint, calculate the log probability density function (PDF)
+        We sum all the log PDFs to get the log-likelihood of the data.
+        If there are a lot of datapoints with very low probability, the log-likelihood
+        will be very negative (or -inf if probability is 0).
+
+        We will prefer distributions with higher log-likelihood values.
+        """
+        try:
+            return np.sum(dist_obj.dist.logpdf(data, *params))
+        except Exception as e:
+            logger.warning(f"Log-likelihood calculation failed: {str(e)}")
+            return -np.inf
+
+    def fit_distribution(
+        self, dist_name: str, data: np.ndarray, method: str = "mle"
+    ) -> Result:
+        """
+        Fit a specific distribution to data.
+
+        Args:
+            dist_name: Name of the distribution to fit
+            data: Data to fit the distribution to
+            method: Fitting method (default: 'mle')
+
+        Returns:
+            FitResult or FailedFit
+        """
+        try:
+            dist_obj = self._registry.get_distribution(dist_name)
+
+            # Get parameter bounds
+            bounds = self._get_bounds(data, dist_obj)
+            # logger.info(f"Bounds for {dist_name}: {bounds}")
+
+            # Perform the fit
+            result = stats.fit(dist_obj.dist, data, method=method, bounds=tuple(bounds))
+
+            # If fit was not successful, return failure
+            if not result.success:
+                logger.warning(f"Fitting failed: {result.message}")
+                return FailedFit(distribution=dist_name, message=str(result.message))
+
+            # Create frozen distribution with fitted parameters
+            frozen_dist = dist_obj.dist(*result.params)
+
+            # Run goodness-of-fit tests
+            kstest_result = self._perform_kstest(data, dist_obj, result.params)
+            log_likelihood = self._calculate_loglikelihood(
+                data, dist_obj, result.params
+            )
+
+            return FitResult(
+                distribution=dist_name,
+                dist_object=dist_obj.dist,
+                params=result.params,
+                frozen_dist=frozen_dist,
+                message=str(result.message),
+                log_likelihood=log_likelihood,
+                kstest=kstest_result,
+            )
+
+        except Exception as e:
+            logger.warning(f"Fitting failed: {str(e)}")
+            return FailedFit(
+                distribution=dist_name, message=f"Fitting failed: {str(e)}"
+            )
+
+    def _mark_best_fits(self, fits: List[Result], criterion: str) -> List[Result]:
+        """
+        Mark the best fits in a list of fit results.
+
+        Args:
+            fits: List of fit results to analyze
+            criterion: Selection criterion ('likelihood', 'ks', or 'combined')
+
+        Returns:
+            The same list with best_likelihood and best_ks attributes updated for successful fits
+        """
+        # Find best likelihood fit
+        best_likelihood_value = float("-inf")
+        best_likelihood_fit = None
+
+        # Find best KS test fit
+        best_ks_value = 0
+        best_ks_fit = None
+
+        # Find the best fits without creating a filtered list
+        for fit in fits:
+            if isinstance(fit, FitResult):
+                # Check for best likelihood
+                fit_likelihood = (
+                    fit.log_likelihood
+                    if fit.log_likelihood is not None
+                    else float("-inf")
+                )
+                if fit_likelihood > best_likelihood_value:
+                    best_likelihood_value = fit_likelihood
+                    best_likelihood_fit = fit
+
+                # Check for best KS test
+                fit_ks = fit.kstest.p_value if fit.kstest else 0
+                if fit_ks > best_ks_value:
+                    best_ks_value = fit_ks
+                    best_ks_fit = fit
+
+        # Only mark if we found best fits
+        if best_likelihood_fit is not None and best_ks_fit is not None:
+            # Mark based on criterion
+            if criterion == "likelihood":
+                # Mark only likelihood best
+                for fit in fits:
+                    if isinstance(fit, FitResult):
+                        fit.best_likelihood = (
+                            fit.distribution == best_likelihood_fit.distribution
+                        )
+
+            elif criterion == "ks":
+                # Mark only KS best
+                for fit in fits:
+                    if isinstance(fit, FitResult):
+                        fit.best_ks = fit.distribution == best_ks_fit.distribution
+
+            elif criterion == "combined":
+                # Mark both
+                for fit in fits:
+                    if isinstance(fit, FitResult):
+                        fit.best_likelihood = (
+                            fit.distribution == best_likelihood_fit.distribution
+                        )
+                        fit.best_ks = fit.distribution == best_ks_fit.distribution
+            else:
+                raise ValueError(f"Unknown criterion '{criterion}'")
+
+        # Return the original list
+        return fits
+
+    def fit(
+        self,
+        data: np.ndarray,
+        discrete: bool,
+        method: str = "mle",
+        criterion: str = "combined",
+    ) -> list[Result]:
+        """
+        Fit all registered distributions to data, mark the best fits, and return results.
+
+        Args:
+            data: Data to fit the distribution to
+            discrete: Whether to fit discrete (True) or continuous (False) distributions
+            method: Fitting method (default: 'mle')
+            criterion: Selection criterion ('likelihood', 'ks', or 'combined') for marking best fits
+
+        Returns:
+            List of Result objects (either FitResult or FailedFit) with best fits marked
+        """
+        if criterion not in ["likelihood", "ks", "combined"]:
+            raise ValueError(f"Unknown criterion '{criterion}'")
+
+        results = []
+        for dist_name in self._registry.get_names():
+            dist_obj = self._registry.get_distribution(dist_name)
+
+            if discrete and dist_obj.is_discrete:
+                results.append(self.fit_distribution(dist_name, data, method))
+
+            if not discrete and not dist_obj.is_discrete:
+                results.append(self.fit_distribution(dist_name, data, method))
+
+        # Mark the best fits based on criterion
+        self._mark_best_fits(results, criterion)
+
+        return results
+
+    @staticmethod
+    def best(results: list[Result], criterion: str = "combined") -> list[Result]:
+        if criterion == "ks":
+            return [
+                fit for fit in results if isinstance(fit, FitResult) and fit.best_ks
+            ]
+        elif criterion == "likelihood":
+            return [
+                fit
+                for fit in results
+                if isinstance(fit, FitResult) and fit.best_likelihood
+            ]
+        elif criterion == "combined":
+            return [
+                fit
+                for fit in results
+                if isinstance(fit, FitResult) and (fit.best_likelihood or fit.best_ks)
+            ]
+        else:
+            raise ValueError(f"Unknown criterion '{criterion}'")

goad_toolkit/config.py

@@ -0,0 +1,18 @@
+from pathlib import Path
+
+from pydantic import BaseModel
+
+
+class FileConfig(BaseModel):
+    data_dir: Path = Path.home() / ".cache/mads_datasets/covid"
+    filename: Path = Path("covid.csv")
+    url: str = (
+        "https://raw.githubusercontent.com/mzelst/covid-19/master/data/rivm_by_day.csv"
+    )
+
+
+class DataConfig(BaseModel):
+    period: int = -14
+    window: int = 7
+    start_date: str = "2020-10-01"
+    end_date: str = "2021-06-01"

goad_toolkit/dataprocessor.py

@@ -0,0 +1,51 @@
+from abc import ABC, abstractmethod
+from pathlib import Path
+from typing import Optional
+
+import pandas as pd
+
+from goad_toolkit.config import DataConfig, FileConfig
+from goad_toolkit.datatransforms import (
+    DiffValues,
+    Pipeline,
+    RollingAvg,
+    SelectDataRange,
+    ShiftValues,
+    ZScaler,
+)
+from goad_toolkit.filehandler import FileHandler
+
+
+class DataProcessor(ABC):
+    def __init__(self, fileconfig: FileConfig, dataconfig: DataConfig) -> None:
+        self.pipeline = Pipeline()
+        self.filehandler = FileHandler(fileconfig)
+        self.config_pipeline(dataconfig)
+
+    @abstractmethod
+    def config_pipeline(self, dataconfig: DataConfig) -> None:
+        pass
+
+    def process(
+        self, filename: Optional[Path] = None, raw: bool = True
+    ) -> pd.DataFrame:
+        df = self.filehandler.load(filename, raw)
+        result = self.pipeline.apply(df)
+        return result
+
+
+class CovidDataProcessor(DataProcessor):
+    def config_pipeline(self, dataconfig: DataConfig) -> None:
+        self.pipeline.add(DiffValues, column="deaths")
+        self.pipeline.add(
+            ShiftValues, column="deaths", period=dataconfig.period, rename=True
+        )
+        self.pipeline.add(
+            SelectDataRange,
+            start_date=dataconfig.start_date,
+            end_date=dataconfig.end_date,
+        )
+        self.pipeline.add(RollingAvg, column="deaths_shifted", window=dataconfig.window)
+        self.pipeline.add(RollingAvg, column="positivetests", window=dataconfig.window)
+        self.pipeline.add(ZScaler, column="deaths_shifted", rename=True)
+        self.pipeline.add(ZScaler, column="positivetests", rename=True)

goad_toolkit/datatransforms.py

@@ -0,0 +1,184 @@
+from abc import ABC, abstractmethod
+from typing import Any, Dict, Optional, Type, TypeVar
+
+import pandas as pd
+from tqdm import tqdm
+
+T = TypeVar("T", bound="TransformBase")
+
+
+class TransformBase(ABC):
+    """Base class for all data transformations."""
+
+    def __init__(self, name: Optional[str] = None, **kwargs):
+        self._params = kwargs
+        self.name = name or self.__class__.__name__
+
+    def __call__(self, data: pd.DataFrame) -> pd.DataFrame:
+        """Checks columns, passes params to transform"""
+        self._validate_column(data)
+
+        # Call the transform method with unpacked parameters
+        return self.transform(data, **self._params)
+
+    def _validate_column(self, data: pd.DataFrame) -> None:
+        """Check if the column exists in the dataframe."""
+        column = self._params.get("column")
+        if column and column not in data.columns:
+            raise ValueError(f"Column '{column}' does not exist in the dataframe.")
+
+    @abstractmethod
+    def transform(self, data: pd.DataFrame, *args, **kwargs) -> pd.DataFrame:
+        """Transform the data."""
+        pass
+
+    def get_params(self) -> Dict[str, Any]:
+        """Get the current parameters."""
+        return self._params.copy()
+
+    def update_params(self, **kwargs) -> None:
+        """Update the parameters."""
+        self._params.update(kwargs)
+
+    def __repr__(self) -> str:
+        """String representation of the transform."""
+        params_str = ", ".join(f"{k}={v!r}" for k, v in self._params.items())
+        return f"{self.name}({params_str})"
+
+
+class ShiftValues(TransformBase):
+    """Shift values in a column by a specified period."""
+
+    def transform(
+        self,
+        data: pd.DataFrame,
+        column: str,
+        period: int,
+        rename: Optional[bool] = False,
+    ) -> pd.DataFrame:
+        """Shift values in a column by a specified period."""
+        if rename:
+            colname = f"{column}_shifted"
+        else:
+            colname = column
+        data[colname] = data[column].shift(period)
+        return data
+
+
+class DiffValues(TransformBase):
+    """Calculate the difference between consecutive values in a column."""
+
+    def transform(
+        self, data: pd.DataFrame, column: str, rename: bool = False
+    ) -> pd.DataFrame:
+        """Calculate the difference between consecutive values in a column."""
+        if rename:
+            colname = f"{column}_diff"
+        else:
+            colname = column
+        data[colname] = data[column].diff()
+        data.iloc[0, data.columns.get_loc(colname)] = 0
+        return data
+
+
+class SelectDataRange(TransformBase):
+    """Select rows within a specified date range."""
+
+    def transform(
+        self, data: pd.DataFrame, start_date: str, end_date: str
+    ) -> pd.DataFrame:
+        """Select rows within a specified date range."""
+        return data.loc[start_date:end_date]
+
+
+class RollingAvg(TransformBase):
+    """Calculate the rolling average of a column."""
+
+    def transform(
+        self, data: pd.DataFrame, column: str, window: int, rename: bool = False
+    ) -> pd.DataFrame:
+        """Calculate the rolling average of a column."""
+        if rename:
+            colname = f"{column}_rolling_avg"
+        else:
+            colname = column
+        data[colname] = data[column].rolling(window).mean()
+        data.dropna(subset=[colname], inplace=True)
+        return data
+
+
+class ZScaler(TransformBase):
+    """Standardize the values in a column."""
+
+    def transform(
+        self, data: pd.DataFrame, column: str, rename: bool = False
+    ) -> pd.DataFrame:
+        """Standardize the values in a column."""
+        if rename:
+            colname = f"{column}_zscore"
+        else:
+            colname = column
+        data[colname] = (data[column] - data[column].mean()) / data[column].std()
+        return data
+
+
+class Pipeline:
+    """Pipeline for chaining data transformations."""
+
+    def __init__(self):
+        self.transforms: Dict[str, Dict[str, Any]] = {}
+
+    def add(
+        self, transform_class: Type[T], name: Optional[str] = None, **kwargs
+    ) -> "Pipeline":
+        """Add a transformation to the pipeline without instantiating it yet."""
+        # Generate name if not provided
+        if name is None:
+            name = transform_class.__name__
+            counter = 1
+            while name in self.transforms:
+                name = f"{name}_{counter}"
+                counter += 1
+
+        self.transforms[name] = {"class": transform_class, "params": kwargs}
+        return self  # Allow method chaining
+
+    def apply(self, data: pd.DataFrame) -> pd.DataFrame:
+        """Apply all transformations in the pipeline."""
+        # Make a single copy at the pipeline level
+        result = data.copy() if len(self.transforms) > 0 else data
+        for name, transform_config in tqdm(
+            self.transforms.items(), desc="Applying transforms"
+        ):
+            transform_class = transform_config["class"]
+            params = transform_config["params"]
+            transform = transform_class(name=name, **params)
+            result = transform(result)
+        return result
+
+    def __getitem__(self, key: str) -> Dict[str, Any]:
+        """Get a specific transform configuration by name."""
+        if key in self.transforms:
+            return self.transforms[key].copy()
+        raise KeyError(f"Transform '{key}' not found in pipeline")
+
+    def __setitem__(self, key: str, params: Dict[str, Any]) -> None:
+        """Update parameters for a transform using dictionary-style assignment."""
+        if key in self.transforms:
+            self.transforms[key]["params"].update(params)
+        else:
+            raise KeyError(f"Transform '{key}' not found in pipeline")
+
+    def __repr__(self) -> str:
+        """String representation of the pipeline."""
+        if not self.transforms:
+            return "Pipeline(steps=[])"
+
+        steps = []
+        for name, config in self.transforms.items():
+            cls = config["class"].__name__
+            params = ", ".join(f"{k}={v!r}" for k, v in config["params"].items())
+            steps.append(f"  {name}: {cls}({params})")
+
+        steps_str = ",\n".join(steps)
+        return f"Pipeline(\n{steps_str}\n)"

goad_toolkit/distributions.py

@@ -0,0 +1,80 @@
+from dataclasses import dataclass, field
+from typing import Any, Dict, List
+
+from scipy import stats
+
+
+@dataclass
+class Distribution:
+    """Dataclass to hold information about a statistical distribution."""
+
+    name: str
+    dist: Any  # The scipy.stats distribution object
+    is_discrete: bool
+    num_params: int
+
+    def __str__(self):
+        """String representation of the distribution."""
+        return f"Distribution({self.name},discrete={self.is_discrete},params={self.num_params})"
+
+    def __repr__(self):
+        """Detailed representation of the distribution."""
+        return f"Distribution(name='{self.name}',\n dist={self.dist.__class__.__name__},\n is_discrete={self.is_discrete},\n num_params={self.num_params})"
+
+
+@dataclass
+class DistributionRegistry:
+    """Registry for statistical distributions with metadata."""
+
+    distributions: Dict[str, Distribution] = field(default_factory=dict)
+
+    def __post_init__(self):
+        """Initialize the registry with common distributions."""
+        self.register_distribution("norm", stats.norm, is_discrete=False, num_params=2)
+        self.register_distribution(
+            "uniform", stats.uniform, is_discrete=False, num_params=2
+        )
+        self.register_distribution(
+            "lognorm", stats.lognorm, is_discrete=False, num_params=3
+        )
+        self.register_distribution(
+            "poisson", stats.poisson, is_discrete=True, num_params=1
+        )
+        self.register_distribution(
+            "exponential", stats.expon, is_discrete=False, num_params=2
+        )
+        self.register_distribution(
+            "skewnorm", stats.skewnorm, is_discrete=False, num_params=3
+        )
+        self.register_distribution(
+            "gamma", stats.gamma, is_discrete=False, num_params=3
+        )
+        self.register_distribution(
+            "weibull", stats.weibull_min, is_discrete=False, num_params=3
+        )
+
+    def __repr__(self) -> str:
+        """Detailed representation of the registry."""
+        return f"DistributionRegistry({self.get_names()})"
+
+    def register_distribution(
+        self, name: str, dist, is_discrete: bool, num_params: int
+    ):
+        """Register a new distribution with metadata."""
+        self.distributions[name] = Distribution(
+            name=name, dist=dist, is_discrete=is_discrete, num_params=num_params
+        )
+
+    def get_distribution(self, name: str) -> Distribution:
+        """Get a distribution by name."""
+        if name not in self.distributions:
+            raise ValueError(f"Distribution '{name}' not found in registry.")
+        return self.distributions[name]
+
+    def get_names(self) -> List[str]:
+        """Get all registered distribution names."""
+        return list(self.distributions.keys())
+
+    def is_discrete(self, name: str) -> bool:
+        """Check if a distribution is discrete."""
+        return self.get_distribution(name).is_discrete