splitpilot 0.1.0__py3-none-any.whl

splitpilot/core/splitter.py

@@ -0,0 +1,441 @@
+from dataclasses import dataclass
+
+import pandas as pd
+from sklearn.model_selection import train_test_split
+
+
+@dataclass
+class SplitResult:
+    X_train: pd.DataFrame
+    X_test: pd.DataFrame
+    y_train: pd.Series
+    y_test: pd.Series
+    strategy: str
+
+    def __iter__(self):
+        yield self.X_train
+        yield self.X_test
+        yield self.y_train
+        yield self.y_test
+
+
+class DatasetSplitter:
+    """
+    Execute a train/test split using a recommended strategy.
+    """
+
+    def split(
+        self,
+        df: pd.DataFrame,
+        target: str,
+        strategy: str = "random",
+        group_column: str | None = None,
+        time_column: str | None = None,
+        test_size: float = 0.2,
+        random_state: int = 42,
+    ) -> SplitResult:
+
+        if target not in df.columns:
+            raise ValueError(
+                f"Target column '{target}' was not found in the dataset."
+            )
+
+        if not 0 < test_size < 1:
+            raise ValueError(
+                "test_size must be between 0 and 1."
+            )
+
+        if group_column is not None and group_column not in df.columns:
+            raise ValueError(
+                f"Group column '{group_column}' was not found in the dataset."
+            )
+
+        if time_column is not None and time_column not in df.columns:
+            raise ValueError(
+                f"Time column '{time_column}' was not found in the dataset."
+            )
+
+        X = df.drop(columns=[target])
+        y = df[target]
+
+        if strategy == "random":
+            return self._random_split(
+                X,
+                y,
+                test_size,
+                random_state,
+            )
+
+        if strategy == "group":
+            return self._group_split(
+                df,
+                target,
+                group_column,
+                test_size,
+                random_state,
+            )
+
+        if strategy == "group_stratified":
+            return self._group_stratified_split(
+                df,
+                target,
+                group_column,
+                test_size,
+                random_state,
+            )
+
+        if strategy == "time":
+            return self._time_split(
+                df,
+                target,
+                time_column,
+                test_size,
+            )
+
+        if strategy == "group_time":
+            return self._group_time_split(
+                df,
+                target,
+                group_column,
+                time_column,
+                test_size,
+            )
+
+        raise ValueError(
+            f"Unknown split strategy: '{strategy}'."
+        )
+
+    def _random_split(
+        self,
+        X,
+        y,
+        test_size,
+        random_state,
+    ):
+        X_train, X_test, y_train, y_test = train_test_split(
+            X,
+            y,
+            test_size=test_size,
+            random_state=random_state,
+        )
+
+        return SplitResult(
+            X_train=X_train,
+            X_test=X_test,
+            y_train=y_train,
+            y_test=y_test,
+            strategy="random",
+        )
+
+    def _group_split(
+        self,
+        df,
+        target,
+        group_column,
+        test_size,
+        random_state,
+    ):
+        if group_column is None:
+            raise ValueError(
+                "group_column is required for group splitting."
+            )
+
+        groups = df[group_column].drop_duplicates()
+
+        rng = pd.Series(groups).sample(
+            frac=1,
+            random_state=random_state,
+        )
+
+        test_count = max(
+            1,
+            round(len(groups) * test_size),
+        )
+
+        test_groups = set(
+            rng.iloc[:test_count]
+        )
+
+        train_groups = set(groups) - test_groups
+
+        train_mask = df[group_column].isin(train_groups)
+        test_mask = df[group_column].isin(test_groups)
+
+        train_df = df.loc[train_mask]
+        test_df = df.loc[test_mask]
+
+        return SplitResult(
+            X_train=train_df.drop(columns=[target]),
+            X_test=test_df.drop(columns=[target]),
+            y_train=train_df[target],
+            y_test=test_df[target],
+            strategy="group",
+        )
+
+    def _group_stratified_split(
+        self,
+        df,
+        target,
+        group_column,
+        test_size,
+        random_state,
+    ):
+        if group_column is None:
+            raise ValueError(
+                "group_column is required for group_stratified splitting."
+            )
+
+        group_target = (
+            df.groupby(group_column)[target]
+            .mean()
+            .reset_index()
+        )
+
+        group_target["_stratum"] = (
+            group_target[target]
+            >= group_target[target].median()
+        )
+
+        train_groups, test_groups = train_test_split(
+            group_target[group_column],
+            test_size=test_size,
+            random_state=random_state,
+            stratify=group_target["_stratum"],
+        )
+
+        train_groups = set(train_groups)
+        test_groups = set(test_groups)
+
+        train_mask = df[group_column].isin(train_groups)
+        test_mask = df[group_column].isin(test_groups)
+
+        train_df = df.loc[train_mask]
+        test_df = df.loc[test_mask]
+
+        return SplitResult(
+            X_train=train_df.drop(columns=[target]),
+            X_test=test_df.drop(columns=[target]),
+            y_train=train_df[target],
+            y_test=test_df[target],
+            strategy="group_stratified",
+        )
+
+    def _time_split(
+        self,
+        df,
+        target,
+        time_column,
+        test_size,
+    ):
+        if time_column is None:
+            raise ValueError(
+                "time_column is required for time splitting."
+            )
+
+        ordered = df.copy()
+
+        ordered[time_column] = pd.to_datetime(
+            ordered[time_column],
+            errors="coerce",
+        )
+
+        if ordered[time_column].isna().any():
+            raise ValueError(
+                f"Column '{time_column}' contains invalid dates."
+            )
+
+        ordered = ordered.sort_values(time_column)
+
+        split_index = int(
+            len(ordered) * (1 - test_size)
+        )
+
+        train_df = ordered.iloc[:split_index]
+        test_df = ordered.iloc[split_index:]
+
+        return SplitResult(
+            X_train=train_df.drop(columns=[target]),
+            X_test=test_df.drop(columns=[target]),
+            y_train=train_df[target],
+            y_test=test_df[target],
+            strategy="time",
+        )
+
+    def _group_time_split(
+        self,
+        df,
+        target,
+        group_column,
+        time_column,
+        test_size,
+    ):
+        """
+        Perform a leakage-safe chronological split while keeping
+        every group entirely inside either train or test.
+
+        A valid split must satisfy both:
+
+        1. No group appears in both train and test.
+        2. Every training observation occurs before every test observation.
+
+        Groups whose time ranges cross the selected temporal boundary
+        are excluded from the split rather than being divided between
+        train and test.
+        """
+
+        if group_column is None:
+            raise ValueError(
+                "group_column is required for group_time splitting."
+            )
+
+        if time_column is None:
+            raise ValueError(
+                "time_column is required for group_time splitting."
+            )
+
+        ordered = df.copy()
+
+        ordered[time_column] = pd.to_datetime(
+            ordered[time_column],
+            errors="coerce",
+        )
+
+        if ordered[time_column].isna().any():
+            raise ValueError(
+                f"Column '{time_column}' contains invalid dates."
+            )
+
+        group_ranges = (
+            ordered.groupby(group_column)[time_column]
+            .agg(["min", "max"])
+            .sort_values("max")
+        )
+
+        total_groups = len(group_ranges)
+
+        if total_groups < 2:
+            raise ValueError(
+                "group_time splitting requires at least two groups."
+            )
+
+        desired_test_groups = max(
+            1,
+            round(total_groups * test_size),
+        )
+
+        best_split = None
+
+        for boundary_index in range(
+            1,
+            total_groups,
+        ):
+            boundary_date = group_ranges["max"].iloc[
+                boundary_index - 1
+            ]
+
+            train_groups = set(
+                group_ranges.index[
+                    group_ranges["max"] <= boundary_date
+                ]
+            )
+
+            test_groups = set(
+                group_ranges.index[
+                    group_ranges["min"] > boundary_date
+                ]
+            )
+
+            if not train_groups or not test_groups:
+                continue
+
+            train_max = ordered.loc[
+                ordered[group_column].isin(train_groups),
+                time_column,
+            ].max()
+
+            test_min = ordered.loc[
+                ordered[group_column].isin(test_groups),
+                time_column,
+            ].min()
+
+            if train_max >= test_min:
+                continue
+
+            test_group_difference = abs(
+                len(test_groups) - desired_test_groups
+            )
+
+            train_group_difference = abs(
+                len(train_groups)
+                - (total_groups - desired_test_groups)
+            )
+
+            score = (
+                test_group_difference,
+                train_group_difference,
+            )
+
+            if (
+                best_split is None
+                or score < best_split["score"]
+            ):
+                best_split = {
+                    "score": score,
+                    "train_groups": train_groups,
+                    "test_groups": test_groups,
+                }
+
+        if best_split is None:
+            raise ValueError(
+                "Could not create a leakage-safe group_time split. "
+                "The groups have overlapping time ranges, so no "
+                "strict chronological boundary exists."
+            )
+
+        train_groups = best_split["train_groups"]
+        test_groups = best_split["test_groups"]
+
+        train_mask = ordered[group_column].isin(
+            train_groups
+        )
+
+        test_mask = ordered[group_column].isin(
+            test_groups
+        )
+
+        train_df = ordered.loc[train_mask]
+        test_df = ordered.loc[test_mask]
+
+        if train_df.empty or test_df.empty:
+            raise RuntimeError(
+                "group_time split produced an empty train or test set."
+            )
+
+        train_ids = set(
+            train_df[group_column]
+        )
+
+        test_ids = set(
+            test_df[group_column]
+        )
+
+        if train_ids.intersection(test_ids):
+            raise RuntimeError(
+                "group_time split produced overlapping groups."
+            )
+
+        train_max = train_df[time_column].max()
+        test_min = test_df[time_column].min()
+
+        if train_max >= test_min:
+            raise RuntimeError(
+                "group_time split failed chronological validation."
+            )
+
+        return SplitResult(
+            X_train=train_df.drop(columns=[target]),
+            X_test=test_df.drop(columns=[target]),
+            y_train=train_df[target],
+            y_test=test_df[target],
+            strategy="group_time",
+        )

splitpilot/models/__init__.py

@@ -0,0 +1 @@
+"""Data models used by SplitPilot."""

splitpilot/models/recommendation.py

@@ -0,0 +1,15 @@
+from dataclasses import dataclass
+
+
+@dataclass
+class SplitRecommendation:
+    """A recommended evaluation split."""
+
+    strategy: str
+    target: str
+    test_size: float
+    random_state: int | None
+    group_column: str | None
+    time_column: str | None
+    confidence: str
+    reasons: list[str]