stats-transformer 1.2.0__py3-none-any.whl

stats_transformer/__init__.py

@@ -0,0 +1,40 @@
+import importlib.metadata
+
+try:
+    __version__ = importlib.metadata.version("stats-transformer")
+except importlib.metadata.PackageNotFoundError:
+    __version__ = "unknown"
+
+from .featurization import FeatureEngineer, EventStudyBuilder
+from .models import RegressionModel
+from .pipeline import Pipeline
+from .visualization import (
+    BaseVisualizer, DataVisualizer, ModelVisualizer, RegressionVisualizer,
+    CoefficientBarChart, GroupedBarChart, StackedBarChart,
+    TimeSeriesPlot, IRFPlot, FacetedTimeSeries,
+    BinnedScatterPlot, ScatterWithRegression, CorrelationHeatmap
+)
+from .visualization.tables import TableGenerator
+from .models.timeseries.local_projections import LocalProjectionsModel
+
+__all__ = [
+    "FeatureEngineer",
+    "RegressionModel",
+    "Pipeline",
+    "BaseVisualizer",
+    "DataVisualizer",
+    "ModelVisualizer",
+    "RegressionVisualizer",
+    "CoefficientBarChart",
+    "GroupedBarChart",
+    "StackedBarChart",
+    "TimeSeriesPlot",
+    "IRFPlot",
+    "FacetedTimeSeries",
+    "BinnedScatterPlot",
+    "ScatterWithRegression",
+    "CorrelationHeatmap",
+    "TableGenerator",
+    "LocalProjectionsModel",
+    "EventStudyBuilder",
+]

stats_transformer/data/__init__.py

@@ -0,0 +1,94 @@
+"""Data utilities for the stats-transformer library."""
+
+import pandas as pd
+from importlib.resources import files
+
+
+EXAMPLE_DATASETS = {
+    "macrodb_gdp_inflation": {
+        "filename": "macrodb_gdp_inflation.parquet",
+        "description": "Macroeconomic dataset with GDP and inflation data for multiple countries",
+        "columns": {
+            "country": "Country code (ISO 3-letter)",
+            "date": "Year of observation",
+            "inflation": "Inflation rate (annual %)",
+            "gdp": "GDP per capita (current US$)",
+        },
+        "shape": (11490, 4),
+        "time_period": "Annual data",
+        "source": "Derived from World Bank and IMF data",
+        "notes": "Includes data for multiple countries across different time periods",
+    }
+}
+
+
+def list_examples():
+    """List packaged example datasets."""
+    return sorted(EXAMPLE_DATASETS)
+
+
+def describe_example(name="macrodb_gdp_inflation"):
+    """Return metadata for a packaged example dataset."""
+    if name not in EXAMPLE_DATASETS:
+        available = ", ".join(list_examples())
+        raise ValueError(f"Unknown example dataset '{name}'. Available examples: {available}")
+
+    metadata = EXAMPLE_DATASETS[name].copy()
+    metadata.pop("filename", None)
+    metadata["name"] = name
+    return metadata
+
+
+def load_example(name="macrodb_gdp_inflation"):
+    """Load a packaged example dataset."""
+    if name not in EXAMPLE_DATASETS:
+        available = ", ".join(list_examples())
+        raise ValueError(f"Unknown example dataset '{name}'. Available examples: {available}")
+
+    filename = EXAMPLE_DATASETS[name]["filename"]
+    try:
+        data_path = files("stats_transformer.data").joinpath(filename)
+        return pd.read_parquet(data_path)
+    except Exception as e:
+        # Fallback for development or testing
+        import os
+        fallback_path = os.path.join(os.path.dirname(__file__), filename)
+        if os.path.exists(fallback_path):
+            return pd.read_parquet(fallback_path)
+        else:
+            raise FileNotFoundError(
+                f"Example data file '{filename}' not found. Make sure the package was installed correctly."
+            ) from e
+
+
+def load_sample_data():
+    """Load the default sample macroeconomic dataset.
+
+    Returns:
+        DataFrame with columns: country, date, inflation, gdp
+
+    Example:
+        >>> from stats_transformer.data import load_sample_data
+        >>> df = load_sample_data()
+        >>> print(df.head())
+    """
+    return load_example("macrodb_gdp_inflation")
+
+
+def get_sample_data_description():
+    """Get description of the sample dataset.
+
+    Returns:
+        Dictionary with dataset description
+    """
+    return describe_example("macrodb_gdp_inflation")
+
+
+__all__ = [
+    "EXAMPLE_DATASETS",
+    "describe_example",
+    "get_sample_data_description",
+    "list_examples",
+    "load_example",
+    "load_sample_data",
+]

stats_transformer/data/panel_builder.py

@@ -0,0 +1,232 @@
+import logging
+import pandas as pd
+
+class PanelDataBuilder:
+    """
+    Utility class to standardize the assembly of panel data, 
+    especially handling mixed-frequency joins (e.g., merging monthly 
+    macroeconomic data onto a daily panel) and broadcasting global 
+    series (e.g., VIX, oil prices) onto country panels.
+    """
+    
+    def __init__(self, entity_col='country', time_col='date'):
+        self.entity_col = entity_col
+        self.time_col = time_col
+        self.logger = logging.getLogger(self.__class__.__name__)
+        
+    def join(self, panel, df, name, on_freq='M', cols=None, index_cols=None):
+        """
+        Merge a lower-frequency (or matched-frequency) dataframe onto a panel.
+        
+        Args:
+            panel (pd.DataFrame): The base panel, indexed by [entity_col, time_col].
+            df (pd.DataFrame): The dataset to merge.
+            name (str): A descriptive name for the dataset being joined (for logging).
+            on_freq (str): The frequency to align on before joining. 
+                           E.g., 'M' means both the panel and df will have their dates 
+                           converted to monthly periods to perform the match.
+            cols (list): Specific columns from `df` to keep. If None, keeps all.
+            index_cols (list): Columns in `df` that represent the entity and time. 
+                               Defaults to [self.entity_col, self.time_col].
+                               
+        Returns:
+            pd.DataFrame: The merged panel.
+        """
+        if df is None or df.empty:
+            self.logger.warning(f"Dataset '{name}' is empty or None. Returning original panel.")
+            return panel
+            
+        index_cols = index_cols or [self.entity_col, self.time_col]
+        
+        # Prepare the incoming dataframe
+        df_clean = df.copy()
+        
+        # Un-index if necessary to work with columns
+        if isinstance(df_clean.index, pd.MultiIndex):
+            df_clean = df_clean.reset_index()
+        elif df_clean.index.name in index_cols:
+            df_clean = df_clean.reset_index()
+            
+        # Ensure index columns are present
+        missing_cols = [c for c in index_cols if c not in df_clean.columns]
+        if missing_cols:
+            self.logger.warning(f"Cannot join '{name}': missing index columns {missing_cols}")
+            return panel
+            
+        # Keep only required columns
+        if cols:
+            keep_cols = index_cols + [c for c in cols if c in df_clean.columns and c not in index_cols]
+            df_clean = df_clean[keep_cols]
+            
+        # Drop overlapping columns that already exist in the panel (except the join keys)
+        panel_cols = panel.columns.tolist() if not isinstance(panel.index, pd.MultiIndex) else panel.columns.tolist() + list(panel.index.names)
+        overlap = [c for c in df_clean.columns if c in panel_cols and c not in index_cols]
+        if overlap:
+            df_clean = df_clean.drop(columns=overlap)
+            
+        if len(df_clean.columns) <= len(index_cols):
+            self.logger.info(f"No new columns to join for '{name}'. Returning original panel.")
+            return panel
+            
+        # Align frequencies
+        panel_reset = panel.copy()
+        is_multiindex_panel = isinstance(panel_reset.index, pd.MultiIndex)
+        if is_multiindex_panel:
+            panel_reset = panel_reset.reset_index()
+            
+        if self.time_col not in panel_reset.columns:
+            self.logger.warning(f"Panel is missing time column '{self.time_col}'.")
+            return panel
+            
+        # Create a temporary period column for merging
+        temp_period_col = f'_period_merge_{on_freq}'
+        
+        # Convert panel dates to period
+        panel_reset[self.time_col] = pd.to_datetime(panel_reset[self.time_col])
+        panel_reset[temp_period_col] = panel_reset[self.time_col].dt.to_period(on_freq)
+        
+        # Convert incoming df dates to period
+        # Assume the second element in index_cols is the time column
+        df_time_col = index_cols[1]
+        df_entity_col = index_cols[0]
+        
+        df_clean[df_time_col] = pd.to_datetime(df_clean[df_time_col])
+        df_clean[temp_period_col] = df_clean[df_time_col].dt.to_period(on_freq)
+        
+        # Deduplicate incoming df based on entity and period (keep last)
+        df_clean = df_clean.drop_duplicates(subset=[df_entity_col, temp_period_col], keep='last')
+        
+        # Prepare for merge
+        df_clean = df_clean.rename(columns={df_entity_col: self.entity_col})
+        df_clean = df_clean.drop(columns=[df_time_col])
+        
+        # Perform the merge
+        merged = panel_reset.merge(
+            df_clean,
+            on=[self.entity_col, temp_period_col],
+            how='left'
+        )
+        
+        # Cleanup
+        merged = merged.drop(columns=[temp_period_col])
+        
+        if is_multiindex_panel:
+            merged = merged.set_index([self.entity_col, self.time_col])
+            
+        # Log successful merge metrics
+        added_cols = [c for c in df_clean.columns if c not in [self.entity_col, temp_period_col]]
+        not_null_counts = [f"{c}={merged[c].notna().sum()}" for c in added_cols if c in merged.columns]
+        self.logger.info(f"Joined '{name}': {merged.shape}, non-null: {', '.join(not_null_counts)}")
+        
+        return merged
+
+    def join_global(self, panel, df, name, time_col=None, cols=None):
+        """
+        Broadcast a global time series (e.g., VIX, oil) onto a panel where 
+        the incoming dataset has no entity column, just dates.
+        
+        Args:
+            panel (pd.DataFrame): The base panel.
+            df (pd.DataFrame): The global dataset to merge.
+            name (str): A descriptive name for logging.
+            time_col (str): The column in `df` representing time. If None, defaults to `self.time_col`.
+            cols (list): Specific columns from `df` to keep.
+            
+        Returns:
+            pd.DataFrame: The merged panel.
+        """
+        if df is None or df.empty:
+            self.logger.warning(f"Global dataset '{name}' is empty or None. Returning original panel.")
+            return panel
+            
+        time_col = time_col or self.time_col
+        df_clean = df.copy()
+        
+        if df_clean.index.name == time_col:
+            df_clean = df_clean.reset_index()
+            
+        if time_col not in df_clean.columns:
+            self.logger.warning(f"Cannot join global '{name}': missing time column '{time_col}'")
+            return panel
+            
+        if cols:
+            keep_cols = [time_col] + [c for c in cols if c in df_clean.columns and c != time_col]
+            df_clean = df_clean[keep_cols]
+            
+        panel_cols = panel.columns.tolist() if not isinstance(panel.index, pd.MultiIndex) else panel.columns.tolist() + list(panel.index.names)
+        overlap = [c for c in df_clean.columns if c in panel_cols and c != time_col]
+        if overlap:
+            df_clean = df_clean.drop(columns=overlap)
+            
+        if len(df_clean.columns) <= 1:
+            self.logger.info(f"No new columns to join for global '{name}'. Returning original panel.")
+            return panel
+            
+        df_clean[time_col] = pd.to_datetime(df_clean[time_col])
+        df_clean = df_clean.drop_duplicates(subset=[time_col], keep='last')
+        
+        panel_reset = panel.copy()
+        is_multiindex_panel = isinstance(panel_reset.index, pd.MultiIndex)
+        if is_multiindex_panel:
+            panel_reset = panel_reset.reset_index()
+            
+        panel_reset[self.time_col] = pd.to_datetime(panel_reset[self.time_col])
+        
+        # If the incoming time_col is named differently than the panel's time_col, align them
+        if time_col != self.time_col:
+            df_clean = df_clean.rename(columns={time_col: self.time_col})
+            
+        # Left merge based on exact date match
+        # Note: If forward-filling (asof join) is desired instead of exact match, 
+        # that would require a separate method or flag (pd.merge_asof).
+        merged = panel_reset.merge(
+            df_clean,
+            on=self.time_col,
+            how='left'
+        )
+        
+        if is_multiindex_panel:
+            merged = merged.set_index([self.entity_col, self.time_col])
+            
+        added_cols = [c for c in df_clean.columns if c != self.time_col]
+        not_null_counts = [f"{c}={merged[c].notna().sum()}" for c in added_cols if c in merged.columns]
+        self.logger.info(f"Joined global '{name}': {merged.shape}, non-null: {', '.join(not_null_counts)}")
+        
+        return merged
+
+    def build(self, base_panel, join_specs):
+        """
+        Fluent-style builder to chain multiple joins.
+        
+        Args:
+            base_panel (pd.DataFrame): The initial panel.
+            join_specs (list of dicts): E.g.
+                [
+                    {"type": "entity", "df": cpi_df, "name": "CPI", "on_freq": "M", "cols": ["infl"]},
+                    {"type": "global", "df": vix_df, "name": "VIX", "cols": ["vix_close"]}
+                ]
+        Returns:
+            pd.DataFrame: Fully built panel.
+        """
+        panel = base_panel.copy()
+        
+        for spec in join_specs:
+            j_type = spec.get("type", "entity")
+            if j_type == "entity":
+                panel = self.join(
+                    panel,
+                    spec.get("df"),
+                    spec.get("name"),
+                    on_freq=spec.get("on_freq", "M"),
+                    cols=spec.get("cols"),
+                    index_cols=spec.get("index_cols")
+                )
+            elif j_type == "global":
+                panel = self.join_global(
+                    panel,
+                    spec.get("df"),
+                    spec.get("name"),
+                    time_col=spec.get("time_col"),
+                    cols=spec.get("cols")
+                )
+        return panel

stats_transformer/featurization/__init__.py

@@ -0,0 +1,9 @@
+import logging
+from .feature_engineering import FeatureEngineer
+from .data_merger import DataMerger
+from .event_study import EventStudyBuilder
+
+__all__ = ["FeatureEngineer", "DataMerger", "EventStudyBuilder", "suppress_logging"]
+
+def suppress_logging():
+    logging.getLogger("stats_transformer.featurization.feature_engineering").setLevel(logging.CRITICAL + 1)

stats_transformer/featurization/base.py

@@ -0,0 +1,69 @@
+"""
+Base feature engineering class.
+"""
+
+import logging
+from abc import ABC, abstractmethod
+from typing import Any, Dict, List, Optional, Union
+
+
+class BaseFeatureEngineer(ABC):
+    """
+    Abstract base class for feature engineering pipelines.
+
+    This class defines the interface for feature engineering transformations
+    that can be applied to time series data with entity-based grouping.
+    """
+
+    def __init__(self, params_path: Optional[str] = None, **kwargs):
+        """
+        Initialize the base feature engineer.
+
+        Args:
+            params_path: Optional path to YAML configuration file
+            **kwargs: Additional configuration parameters
+        """
+        self.params = {}
+        if params_path:
+            self.params = self._load_params(params_path)
+        
+        # Apply any kwargs to override params
+        for key, value in kwargs.items():
+            if key in self.params:
+                self.params[key] = value
+        
+        # Initialize logger
+        self.logger = logging.getLogger(self.__class__.__name__)
+        
+    def _load_params(self, params_path: str) -> Dict:
+        """
+        Load parameters from YAML file.
+        
+        Args:
+            params_path: Path to YAML file
+            
+        Returns:
+            Dictionary of parameters
+        """
+        import yaml
+        try:
+            with open(params_path, "r") as f:
+                return yaml.safe_load(f)
+        except Exception as e:
+            self.logger.error(f"Error loading parameters from {params_path}: {e}")
+            raise
+    
+    @abstractmethod
+    def fit(self, df):
+        """Fit the feature engineer to data."""
+        pass
+    
+    @abstractmethod
+    def transform(self, df):
+        """Transform the data."""
+        pass
+    
+    @abstractmethod
+    def fit_transform(self, df):
+        """Fit and transform the data."""
+        pass

stats_transformer/featurization/data_merger.py

@@ -0,0 +1,128 @@
+import logging
+import os
+import pandas as pd
+import yaml
+
+ISO2_TO_ISO3 = {}
+
+class DataMerger:
+
+    def __init__(self, params_path="params.yaml"):
+        self.params_path = params_path
+        self.logger = logging.getLogger(__name__)
+        self._setup_logging()
+        self.config = self._load_config()
+
+    def _setup_logging(self):
+        self.logger.setLevel(logging.INFO)
+        if not self.logger.handlers:
+            handler = logging.StreamHandler()
+            handler.setFormatter(logging.Formatter("%(asctime)s - %(levelname)s - %(message)s"))
+            self.logger.addHandler(handler)
+
+    def _load_config(self):
+        if not self.params_path or not os.path.exists(self.params_path):
+            return {}
+        with open(self.params_path, "r") as f:
+            return yaml.safe_load(f)
+
+    def load_dataset(self, path):
+        if path.endswith(".parquet"):
+            return pd.read_parquet(path)
+        return pd.read_csv(path)
+
+    def _log_merge_diagnostics(self, df_left, df_right, merged_df, entity_col):
+        rows_left, rows_right, rows_merged = len(df_left), len(df_right), len(merged_df)
+        if entity_col in df_left.columns:
+            countries_left = set(df_left[entity_col].unique())
+        else:
+            countries_left = set(df_left.index.get_level_values(entity_col).unique())
+        if entity_col in df_right.columns:
+            countries_right = set(df_right[entity_col].unique())
+        else:
+            countries_right = set(df_right.index.get_level_values(entity_col).unique())
+        if entity_col in merged_df.columns:
+            countries_merged = set(merged_df[entity_col].unique())
+        elif merged_df.index.names and entity_col in merged_df.index.names:
+            countries_merged = set(merged_df.index.get_level_values(entity_col).unique())
+        else:
+            countries_merged = set()
+        dropped = countries_left - countries_merged
+        self.logger.info(f"Left:   {rows_left:>7} rows, {len(countries_left):>4} countries")
+        self.logger.info(f"Right:  {rows_right:>7} rows, {len(countries_right):>4} countries")
+        self.logger.info(f"Merged: {rows_merged:>7} rows, {len(countries_merged):>4} countries")
+        if dropped:
+            self.logger.warning(f"Dropped countries from left: {sorted(dropped)}")
+
+    def merge(self, df_left, df_right, on, how="inner"):
+        merged_df = pd.merge(df_left, df_right, on=on, how=how)
+        entity_col = on[0] if on else "country"
+        self._log_merge_diagnostics(df_left, df_right, merged_df, entity_col)
+        return merged_df
+
+    def merge_on_index(self, df_left, df_right, how="inner"):
+        merged_df = pd.merge(df_left, df_right, how=how, left_index=True, right_index=True)
+        entity_col = "country"
+        if "country" in df_left.index.names:
+            self._log_merge_diagnostics(df_left, df_right, merged_df, entity_col)
+        return merged_df
+
+    def merge_asof(self, df_left, df_right, left_on, right_on, by="country", direction="backward"):
+        df_left_sorted = df_left.sort_values(by=left_on)
+        df_right_sorted = df_right.sort_values(by=right_on)
+        merged_df = pd.merge_asof(df_left_sorted, df_right_sorted, left_on=left_on, right_on=right_on, by=by, direction=direction)
+        self._log_merge_diagnostics(df_left, df_right, merged_df, by)
+        return merged_df
+
+    def standardize_entity(self, df, entity_col):
+        df = df.rename(columns={entity_col: "country"})
+        df["country"] = df["country"].map(lambda x: ISO2_TO_ISO3.get(x, x))
+        return df
+
+    def run(self):
+        data_config = self.config.get("data", {})
+        datasets_config = data_config.get("datasets", [])
+
+        if not datasets_config:
+            self.logger.warning("No datasets found in config to merge")
+            return None
+
+        loaded_datasets = []
+        for ds_cfg in datasets_config:
+            name = ds_cfg.get("name")
+            self.logger.info(f"Loading dataset: {name}")
+            df = self.load_dataset(ds_cfg.get("path"))
+            entity_col = ds_cfg.get("entity_column", "country")
+            date_col = ds_cfg.get("date_column", "date")
+            if entity_col != "country":
+                df = self.standardize_entity(df, entity_col)
+            if date_col != "date":
+                df = df.rename(columns={date_col: "date"})
+            loaded_datasets.append(df)
+
+        if not loaded_datasets:
+            return None
+
+        merge_config = data_config.get("merge", {})
+        merge_on = merge_config.get("on", ["country", "date"])
+        merge_how = merge_config.get("how", "outer")
+        output_path = merge_config.get("output_path", "data/pipeline/resampled_merged.parquet")
+
+        merged_df = loaded_datasets[0]
+        for i, next_df in enumerate(loaded_datasets[1:], start=1):
+            self.logger.info(f"Merging dataset {i + 1} of {len(loaded_datasets)}")
+            merged_df = self.merge(merged_df, next_df, on=merge_on, how=merge_how)
+
+        os.makedirs(os.path.dirname(output_path), exist_ok=True)
+        merged_df.to_parquet(output_path)
+        self.logger.info(f"Merged data saved to {output_path} with shape {merged_df.shape}")
+        return merged_df
+
+if __name__ == "__main__":
+    import argparse
+    parser = argparse.ArgumentParser()
+    parser.add_argument("--config", default="params.yaml")
+    args = parser.parse_args()
+    dm = DataMerger(params_path=args.config)
+    dm.run()
+