Metadata-Version: 2.3
Name: mstts
Version: 0.1.1
Summary: Multi-class Stratified Train Test Spliter
Author: Mykhailo Vladymyrov
Author-email: Mykhailo Vladymyrov <neworldemancer@gmail.com>
Requires-Dist: numpy>=1.24.0
Requires-Dist: ortools>=9.0
Requires-Dist: ruff ; extra == 'dev'
Requires-Dist: mypy ; extra == 'dev'
Requires-Dist: pre-commit ; extra == 'dev'
Requires-Dist: rich ; extra == 'dev'
Requires-Dist: pytest ; extra == 'test'
Requires-Dist: pytest-cov ; extra == 'test'
Requires-Python: >=3.10
Project-URL: homepage, https://github.com/dsl-unibe-ch/mstts.git
Provides-Extra: dev
Provides-Extra: test
Description-Content-Type: text/markdown

# MSTTS

[![CI](https://github.com/dsl-unibe-ch/mstts/actions/workflows/ci.yml/badge.svg)](https://github.com/dsl-unibe-ch/mstts/actions/workflows/ci.yml)
[![PyPI version](https://img.shields.io/pypi/v/mstts.svg)](https://pypi.org/project/mstts/)
[![Python versions](https://img.shields.io/badge/python-3.10%2B-blue.svg)](https://pypi.org/project/mstts/)
[![License: Apache 2.0](https://img.shields.io/badge/License-Apache_2.0-blue.svg)](LICENSE)
[![Ruff](https://img.shields.io/endpoint?url=https://raw.githubusercontent.com/astral-sh/ruff/main/assets/badge/v2.json)](https://github.com/astral-sh/ruff)

**M**ulti-class **S**tratified **T**rain-**T**est **S**plitter — split a dataset into
train and test sets while preserving the per-column distribution of *many* target
classes at once.

Unlike a single-label stratified split, MSTTS balances **all** target columns
simultaneously by solving a constraint-optimization problem with Google OR-Tools
(CP-SAT). The split is deterministic: for the same data and parameters you always get
the same result. If no split satisfies the constraints, the solver says so instead of
returning a bad split.

## Installation

```bash
pip install mstts
```

Or install the latest development version straight from GitHub:

```bash
pip install git+ssh://git@github.com/dsl-unibe-ch/mstts.git
```

## Quick start

The package exposes two functions:

- `train_test_split(rows, ...)` — split row indices into train and test sets.
- `get_row_selection(rows, ...)` — lower level: pick a subset of rows that hits a target
  fraction per column. `train_test_split` is built on top of it.

`rows` is an `(n_samples, n_columns)` array-like of **non-negative integers** — usually
one-hot / multi-hot class indicators, but any non-negative counts work. Float arrays are
cast to `int`.

```python
import numpy as np
from mstts import train_test_split

# 8 samples, 4 binary class columns
rows = np.array([
    [1, 1, 0, 0],
    [1, 0, 1, 0],
    [1, 0, 0, 1],
    [0, 1, 1, 0],
    [0, 1, 0, 1],
    [0, 0, 1, 1],
    [1, 1, 1, 0],
    [0, 1, 1, 1],
])

train_idx, test_idx = train_test_split(rows, test_frac=0.25, test_prec=0.1)
# -> two lists of row indices; disjoint and together covering every row
```

Working from a `pandas.DataFrame`? Pass the class columns as an array and index back:

```python
class_cols = ["class1", "class2", "class3"]
train_idx, test_idx = train_test_split(df[class_cols].to_numpy(), test_frac=0.2)
train_df, test_df = df.iloc[train_idx], df.iloc[test_idx]
```

If no feasible split exists (or the 60 s solver time limit is hit),
`train_test_split` returns `(None, None)` and `get_row_selection` returns `None`.

### Every parameter, and what each returns

Both functions share `rows`, `v`, `n_strict`, and `non_strict_weight`; they differ only
in how the target fraction is named (`label_*` for the selection, `test_*` for the
split). The snippet below passes every argument explicitly at its default value:

```python
from mstts import get_row_selection, train_test_split

# get_row_selection — pick a subset of rows hitting a target fraction per column
selected = get_row_selection(
    rows,                 # (n_samples, n_columns) array-like of non-negative ints
    label_frac=0.8,       # target fraction of each column's total to include
    label_prec=0.15,      # allowance: included fraction may be label_frac ± label_prec
    v=1,                  # verbosity: 0 = silent, 1 = summary, 2 = debug
    n_strict=None,        # first n_strict columns are hard-constrained; None = all
    non_strict_weight=1,  # relative weight of priority (non-strict) columns
)
# -> list[int]: indices of the selected rows
#    or None     if no feasible selection is found / the time limit is reached

# train_test_split — split row indices into train and test sets
train_idx, test_idx = train_test_split(
    rows,                 # same as above
    test_frac=0.2,        # target fraction of each column's total for the TEST set
    test_prec=0.15,       # allowance: test fraction may be test_frac ± test_prec
    v=1,
    n_strict=None,
    non_strict_weight=1,
)
# -> (train_idx, test_idx): two lists of row indices — disjoint and together covering
#    every row (train_idx is the ~(1 - test_frac) portion)
#    or (None, None)         if no feasible split is found / the time limit is reached
```

## Thresholds & allowances

You specify a **target fraction** and an **allowance** (half-width of the tolerated
window) per split.

For `train_test_split`, they are expressed in terms of the **test** set:

| Parameter   | Default | Meaning                                                        |
| ----------- | ------- | -------------------------------------------------------------- |
| `test_frac` | `0.2`   | Target fraction of each column's total to place in the test set. |
| `test_prec` | `0.15`  | Allowance: the test fraction may land anywhere in `test_frac ± test_prec`. |

For `get_row_selection`, the same idea applies to the **selected** subset:

| Parameter    | Default | Meaning                                                       |
| ------------ | ------- | ------------------------------------------------------------- |
| `label_frac` | `0.8`   | Target fraction of each column's total to include.            |
| `label_prec` | `0.15`  | Allowance: the included fraction may land in `label_frac ± label_prec`. |

(`train_test_split` simply calls `get_row_selection` with `label_frac = 1 - test_frac`
and `label_prec = test_prec`, selecting the train set and returning its complement as
the test set.)

So the defaults keep **65 %–95 % of every column in the train set** (ideal 80 %), i.e.
5 %–35 % in the test set. Within that window the solver aims for the exact ideal
(`frac × column_total`), not just anywhere in range — see weighting below.

**Integer-count caveat.** Bounds are enforced on integer counts, and
`int(total × frac)` truncates. For a column with a small total the achievable fraction
is coarse — a column whose total is `2` can only be split `0 % / 50 % / 100 %`, so an
`80 % ± 15 %` request is infeasible for it. Give each column enough positive samples, or
widen the allowance, for tight fractions to be reachable.

## Strict vs. priority columns, and weighting

Columns come in two kinds, controlled by `n_strict`:

- **Strict columns** — the first `n_strict` columns. Their included count is a **hard
  constraint**: it *must* fall inside the `frac ± prec` window. They are *also* pulled
  toward the ideal target by the objective.
- **Priority columns** (a.k.a. non-strict) — the remaining columns. They are **not**
  range-constrained, but they still contribute to the objective, so the solver tries to
  bring them close to their target when it can.

`n_strict=None` (the default) treats **all** columns as strict.

The solver minimizes a weighted sum of per-column deviations from the ideal target:

```
minimize  Σ_j  weight_j · | included_count_j − round(frac × total_j) |
```

`non_strict_weight` sets the **relative importance of priority columns vs. strict
columns** in that objective. Because CP-SAT needs integer coefficients, the value is
read in two modes:

| `non_strict_weight` | Strict weight | Priority weight | Effect                                                    |
| ------------------- | ------------- | --------------- | --------------------------------------------------------- |
| `1` (default, int)  | 1             | 1               | Equal importance.                                         |
| int `> 1` (e.g. `2`, `5`) | 1       | that int        | Priority columns matter **more** than strict ones' objective term. |
| float `1.0`         | 10            | 10              | Equal importance.                                         |
| float `< 1.0` (e.g. `0.5`, `0.1`) | 10  | `int(10 × value)` | Priority columns matter **less** than strict ones (down-weighting; `~0.1` is the practical minimum — below `0.05` it rounds to `0` and priority columns are ignored). |

In short: pass an **int ≥ 1** to make priority columns weigh as much as or more than
strict columns, or a **float in `(0, 1]`** to weigh them less. Strict columns always
keep their hard window regardless of the weight.

### Verbosity

Both functions take `v`: `0` = silent, `1` = summary (objective, per-column
achieved-vs-target report), `2` = full debug (also prints the raw thresholds).

## Tests

Run the suite with:

```bash
uv run --extra test pytest        # or simply: pytest
```

15 tests across three files:

**`tests/test_get_row_selection.py`** — the core selector:
- *simple* / *fraction* / *precision* — the selection hits the requested per-column
  fraction within the allowance for the default (80 %), a 50 %, and a tight
  (75 % ± 5 %) target.
- *n_strict* — with only the first columns strict, those stay inside the window while
  the rest are left free.
- *edge_cases* — all-zeros, all-ones, single-column, and single-row inputs all return a
  valid selection.
- *empty* — empty input raises `IndexError`.

**`tests/test_train_test_split.py`** — the same matrix of cases for the train/test
wrapper (default / fraction / precision sizes, `n_strict`, edge cases, empty input), and
additionally checks that the train and test index sets are **disjoint and cover every
row**.

**`tests/test_integration.py`** — end-to-end on larger, realistic data:
- *integration_flow* — a 100×10 correlated dataset; verifies the train and test class
  distributions match within 15 %.
- *imbalanced_data* — columns with 10 % / 50 % / 90 % prevalence; verifies the rare
  class still appears in the test set and the common class lands near the target
  fraction.
- *non_binary_data* — integer count values rather than 0/1, confirming the solver
  handles non-binary columns.

> **Note on the toy datasets.** Because bounds apply to integer counts (see the caveat
> above), the small unit-test fixtures are scaled up with `np.tile` so the requested
> fractions are actually achievable — a column whose total is `2` cannot be split
> `80 / 20`.

## Releasing (for maintainers)

Publishing to PyPI is automated. Pushing a **version tag** (`v*`) triggers the
`Publish to PyPI` GitHub Action, which builds the package and uploads it via Trusted
Publishing (OIDC — no stored token). The published version is taken from the tag (a
leading `v` is stripped), so you never edit the version in `pyproject.toml` by hand:

```bash
git tag v0.1.0
git push origin v0.1.0     # -> publishes mstts 0.1.0
```

PyPI rejects re-uploading an existing version, so every release needs a new tag. Only
tags matching `v*` publish; other tags are ignored. A one-time PyPI trusted-publisher
setup is required before the first release — see
[`.github/workflows/README.md`](.github/workflows/README.md).

## License

Apache 2.0
