gsppy 3.6.0__py3-none-any.whl → 4.1.0__py3-none-any.whl

gsppy/gsp.py

@@ -85,10 +85,12 @@ Version:
 - Current Version: 2.0
 """
 
+from __future__ import annotations
+
 import math
 import logging
 import multiprocessing as mp
-from typing import Dict, List, Tuple, Union, Optional, cast
+from typing import TYPE_CHECKING, Dict, List, Tuple, Union, Literal, Optional, cast, overload
 from itertools import chain
 from collections import Counter
 
@@ -100,8 +102,13 @@ from gsppy.utils import (
     is_subsequence_in_list_with_time_constraints,
 )
 from gsppy.pruning import PruningStrategy, create_default_pruning_strategy
+from gsppy.sequence import Sequence, dict_to_sequences
 from gsppy.accelerate import support_counts as support_counts_accel
 
+if TYPE_CHECKING:
+    import pandas as pd
+    import polars as pl
+
 logger: logging.Logger = logging.getLogger(__name__)
 
 
@@ -126,21 +133,37 @@ class GSP:
 
     def __init__(
         self,
-        raw_transactions: Union[List[List[str]], List[List[Tuple[str, float]]]],
+        raw_transactions: Union[
+            List[List[str]],
+            List[List[Tuple[str, float]]],
+            "pl.DataFrame",
+            "pl.LazyFrame",
+            "pd.DataFrame",
+        ],
         mingap: Optional[float] = None,
         maxgap: Optional[float] = None,
         maxspan: Optional[float] = None,
         verbose: bool = False,
         pruning_strategy: Optional[PruningStrategy] = None,
+        transaction_col: Optional[str] = None,
+        item_col: Optional[str] = None,
+        timestamp_col: Optional[str] = None,
+        sequence_col: Optional[str] = None,
     ):
         """
         Initialize the GSP algorithm with raw transactional data.
 
         Parameters:
-            raw_transactions (Union[List[List[str]], List[List[Tuple[str, float]]]]):
-                Input transaction dataset where each transaction is either:
-                - A list of items (e.g., [['A', 'B'], ['B', 'C', 'D']])
-                - A list of (item, timestamp) tuples (e.g., [[('A', 1.0), ('B', 2.0)]])
+            raw_transactions (Union[List[List[str]], List[List[Tuple[str, float]]], DataFrame]):
+                Input transaction dataset. Accepts:
+                - A list of transactions where each transaction is a list of items (e.g., [['A', 'B'], ['B', 'C', 'D']])
+                - A list of transactions with timestamps (e.g., [[('A', 1.0), ('B', 2.0)]])
+                - A Polars or Pandas DataFrame (requires 'gsppy[dataframe]' installation)
+
+                When using DataFrames, you must specify either:
+                - `sequence_col`: Column containing complete sequences (list format)
+                - `transaction_col` and `item_col`: Columns for grouped format
+
             mingap (Optional[float]): Minimum time gap required between consecutive items in patterns.
             maxgap (Optional[float]): Maximum time gap allowed between consecutive items in patterns.
             maxspan (Optional[float]): Maximum time span from first to last item in patterns.
@@ -149,6 +172,10 @@ class GSP:
             pruning_strategy (Optional[PruningStrategy]): Custom pruning strategy for candidate filtering.
                                                           If None, a default strategy is created based on
                                                           temporal constraints.
+            transaction_col (Optional[str]): DataFrame only - column name for transaction IDs (grouped format).
+            item_col (Optional[str]): DataFrame only - column name for items (grouped format).
+            timestamp_col (Optional[str]): DataFrame only - column name for timestamps.
+            sequence_col (Optional[str]): DataFrame only - column name containing sequences (sequence format).
 
         Attributes Initialized:
             - Processes the input raw transaction dataset.
@@ -161,6 +188,44 @@ class GSP:
             ValueError: If the input transaction dataset is empty, contains
                         fewer than two transactions, or is not properly formatted.
                         Also raised if temporal constraints are invalid.
+
+        Examples:
+            Basic usage with lists:
+
+            ```python
+            from gsppy.gsp import GSP
+
+            transactions = [["A", "B"], ["B", "C", "D"]]
+            gsp = GSP(transactions)
+            patterns = gsp.search(min_support=0.5)
+            ```
+
+            Using Polars DataFrame (grouped format):
+
+            ```python
+            import polars as pl
+            from gsppy.gsp import GSP
+
+            df = pl.DataFrame(
+                {
+                    "transaction_id": [1, 1, 2, 2, 2],
+                    "item": ["A", "B", "A", "C", "D"],
+                }
+            )
+            gsp = GSP(df, transaction_col="transaction_id", item_col="item")
+            patterns = gsp.search(min_support=0.5)
+            ```
+
+            Using Pandas DataFrame (sequence format):
+
+            ```python
+            import pandas as pd
+            from gsppy.gsp import GSP
+
+            df = pd.DataFrame({"sequence": [["A", "B"], ["A", "C", "D"]]})
+            gsp = GSP(df, sequence_col="sequence")
+            patterns = gsp.search(min_support=0.5)
+            ```
         """
         self.freq_patterns: List[Dict[Tuple[str, ...], int]] = []
         self.mingap = mingap
@@ -170,7 +235,13 @@ class GSP:
         self.pruning_strategy: PruningStrategy
         self._configure_logging()
         self._validate_temporal_constraints()
-        self._pre_processing(raw_transactions)
+
+        # Convert DataFrame to transaction list if necessary
+        transactions_to_process = self._convert_input_data(
+            raw_transactions, transaction_col, item_col, timestamp_col, sequence_col
+        )
+
+        self._pre_processing(transactions_to_process)
         # Initialize default pruning strategy if none provided
         if pruning_strategy is None:
             self.pruning_strategy = create_default_pruning_strategy(
@@ -180,6 +251,78 @@ class GSP:
         else:
             self.pruning_strategy = pruning_strategy
 
+    def _convert_input_data(
+        self,
+        raw_transactions: Union[
+            List[List[str]],
+            List[List[Tuple[str, float]]],
+            "pl.DataFrame",
+            "pl.LazyFrame",
+            "pd.DataFrame",
+        ],
+        transaction_col: Optional[str],
+        item_col: Optional[str],
+        timestamp_col: Optional[str],
+        sequence_col: Optional[str],
+    ) -> Union[List[List[str]], List[List[Tuple[str, float]]]]:
+        """
+        Convert input data to the expected transaction list format.
+
+        This method handles both traditional list inputs and DataFrame inputs
+        (Polars or Pandas). DataFrames are converted using the dataframe_adapters module.
+
+        Parameters:
+            raw_transactions: Input data (list or DataFrame)
+            transaction_col: Column name for transaction IDs (DataFrame grouped format)
+            item_col: Column name for items (DataFrame grouped format)
+            timestamp_col: Column name for timestamps (DataFrame)
+            sequence_col: Column name for sequences (DataFrame sequence format)
+
+        Returns:
+            Transaction list in the expected format
+
+        Raises:
+            ValueError: If DataFrame parameters are specified for non-DataFrame input
+                       or if DataFrame conversion fails
+        """
+        # Check if any DataFrame-specific parameters are provided
+        df_params_provided = any([transaction_col, item_col, timestamp_col, sequence_col])
+
+        # If it's a list, validate that no DataFrame parameters were provided
+        if isinstance(raw_transactions, list):
+            if df_params_provided:
+                raise ValueError(
+                    "DataFrame parameters (transaction_col, item_col, timestamp_col, sequence_col) "
+                    "cannot be used with list input"
+                )
+            return cast(Union[List[List[str]], List[List[Tuple[str, float]]]], raw_transactions)  # pyright: ignore[reportUnnecessaryCast]
+
+        # Otherwise, try to convert as DataFrame
+        from gsppy.dataframe_adapters import DataFrameAdapterError, dataframe_to_transactions
+
+        try:
+            logger.debug("Converting DataFrame input to transaction list")
+            transactions = dataframe_to_transactions(
+                raw_transactions,
+                transaction_col=transaction_col,
+                item_col=item_col,
+                timestamp_col=timestamp_col,
+                sequence_col=sequence_col,
+            )
+            logger.debug("Successfully converted DataFrame with %d transactions", len(transactions))
+            return transactions
+        except DataFrameAdapterError as e:
+            msg = f"Failed to convert DataFrame input: {e}"
+            logger.error(msg)
+            raise ValueError(msg) from e
+        except ImportError as e:
+            msg = (
+                "DataFrame input detected but dataframe_adapters module failed to import. "
+                "Install DataFrame support with: pip install 'gsppy[dataframe]'"
+            )
+            logger.error(msg)
+            raise ValueError(msg) from e
+
     def _configure_logging(self) -> None:
         """
         Configure logging for the GSP instance based on verbosity setting.
@@ -262,8 +405,7 @@ class GSP:
         # Validate temporal constraints are only used with timestamps
         if (self.mingap is not None or self.maxgap is not None or self.maxspan is not None) and not self.has_timestamps:
             logger.warning(
-                "Temporal constraints specified but transactions do not have timestamps. "
-                "Constraints will be ignored."
+                "Temporal constraints specified but transactions do not have timestamps. Constraints will be ignored."
             )
             # Clear temporal constraints since they cannot be applied
             self.mingap = None
@@ -449,13 +591,37 @@ class GSP:
         """
         logger.info("Run %d: %d candidates filtered to %d.", run, len(candidates), len(self.freq_patterns[run - 1]))
 
+    @overload
     def search(
         self,
         min_support: float = 0.2,
         max_k: Optional[int] = None,
         backend: Optional[str] = None,
         verbose: Optional[bool] = None,
-    ) -> List[Dict[Tuple[str, ...], int]]:
+        *,
+        return_sequences: Literal[False] = False,
+    ) -> List[Dict[Tuple[str, ...], int]]: ...
+
+    @overload
+    def search(
+        self,
+        min_support: float = 0.2,
+        max_k: Optional[int] = None,
+        backend: Optional[str] = None,
+        verbose: Optional[bool] = None,
+        *,
+        return_sequences: Literal[True],
+    ) -> List[List[Sequence]]: ...
+
+    def search(
+        self,
+        min_support: float = 0.2,
+        max_k: Optional[int] = None,
+        backend: Optional[str] = None,
+        verbose: Optional[bool] = None,
+        *,
+        return_sequences: bool = False,
+    ) -> Union[List[Dict[Tuple[str, ...], int]], List[List[Sequence]]]:
         """
         Execute the Generalized Sequential Pattern (GSP) mining algorithm.
 
@@ -476,11 +642,20 @@ class GSP:
                                     Note: temporal constraints always use Python backend.
             verbose (Optional[bool]): Override instance verbosity setting for this search.
                                      If None, uses the instance's verbose setting.
+            return_sequences (bool): If True, returns patterns as Sequence objects instead of
+                                    Dict[Tuple[str, ...], int]. Defaults to False for backward
+                                    compatibility. When True, returns List[List[Sequence]] where
+                                    each Sequence contains items, support count, and can be extended
+                                    with additional metadata.
 
         Returns:
-            List[Dict[Tuple[str, ...], int]]: A list of dictionaries containing frequent patterns
-                                              at each k-sequence level, with patterns as keys
-                                              and their support counts as values.
+            Union[List[Dict[Tuple[str, ...], int]], List[List[Sequence]]]:
+                If return_sequences is False (default):
+                    A list of dictionaries containing frequent patterns at each k-sequence level,
+                    with patterns as keys and their support counts as values.
+                If return_sequences is True:
+                    A list of lists containing Sequence objects at each k-sequence level,
+                    where each Sequence encapsulates the pattern items and support count.
 
         Raises:
             ValueError: If the minimum support threshold is not in the range `(0.0, 1.0]`.
@@ -491,7 +666,7 @@ class GSP:
             - Status updates for each iteration until the algorithm terminates.
 
         Examples:
-            Basic usage without temporal constraints:
+            Basic usage without temporal constraints (default tuple-based):
 
             ```python
             from gsppy.gsp import GSP
@@ -504,6 +679,28 @@ class GSP:
 
             gsp = GSP(transactions)
             patterns = gsp.search(min_support=0.3)
+            # Returns: [{('Bread',): 4, ('Milk',): 4, ...}, {('Bread', 'Milk'): 3, ...}, ...]
+            ```
+
+            Using Sequence objects for richer pattern representation:
+
+            ```python
+            from gsppy.gsp import GSP
+
+            transactions = [
+                ["Bread", "Milk"],
+                ["Bread", "Diaper", "Beer", "Eggs"],
+                ["Milk", "Diaper", "Beer", "Coke"],
+            ]
+
+            gsp = GSP(transactions)
+            patterns = gsp.search(min_support=0.3, return_sequences=True)
+            # Returns: [[Sequence(('Bread',), support=4), Sequence(('Milk',), support=4), ...], ...]
+
+            # Access pattern details
+            for level_patterns in patterns:
+                for seq in level_patterns:
+                    print(f"Pattern: {seq.items}, Support: {seq.support}")
             ```
 
             Usage with temporal constraints (requires timestamped transactions):
@@ -541,6 +738,9 @@ class GSP:
                 f"Using temporal constraints: mingap={self.mingap}, maxgap={self.maxgap}, maxspan={self.maxspan}"
             )
 
+        # Clear freq_patterns for this search (allow reusing the GSP instance)
+        self.freq_patterns = []
+
         # Convert fractional support to absolute count (ceil to preserve threshold semantics)
         abs_min_support = int(math.ceil(len(self.transactions) * float(min_support)))
 
@@ -588,4 +788,9 @@ class GSP:
             self.verbose = original_verbose
             self._configure_logging()
 
-        return self.freq_patterns[:-1]
+        # Return results in the requested format
+        result = self.freq_patterns[:-1]
+        if return_sequences:
+            # Convert Dict[Tuple[str, ...], int] to List[Sequence] for each level
+            return [dict_to_sequences(level_patterns) for level_patterns in result]
+        return result