datacarve 0.1.0__tar.gz

This diff represents the content of publicly available package versions that have been released to one of the supported registries. The information contained in this diff is provided for informational purposes only and reflects changes between package versions as they appear in their respective public registries.
@@ -0,0 +1,21 @@
1
+ MIT License
2
+
3
+ Copyright (c) 2019 Vassilios Vonikakis
4
+
5
+ Permission is hereby granted, free of charge, to any person obtaining a copy
6
+ of this software and associated documentation files (the "Software"), to deal
7
+ in the Software without restriction, including without limitation the rights
8
+ to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
9
+ copies of the Software, and to permit persons to whom the Software is
10
+ furnished to do so, subject to the following conditions:
11
+
12
+ The above copyright notice and this permission notice shall be included in all
13
+ copies or substantial portions of the Software.
14
+
15
+ THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
16
+ IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
17
+ FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
18
+ AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
19
+ LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
20
+ OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
21
+ SOFTWARE.
@@ -0,0 +1,298 @@
1
+ Metadata-Version: 2.4
2
+ Name: datacarve
3
+ Version: 0.1.0
4
+ Summary: Distributional dataset undersampling: carve a balanced subset out of a large dataset with MILP optimization, enforcing target distributions across all dimensions simultaneously.
5
+ Author: Vasileios Vonikakis
6
+ License: MIT
7
+ Project-URL: Homepage, https://github.com/bbonik/datacarve
8
+ Project-URL: Repository, https://github.com/bbonik/datacarve
9
+ Keywords: distributional dataset undersampling,undersampling,dataset balancing,subset selection,MILP,integer programming,data shaping,fairness,imbalanced data,sampling
10
+ Classifier: Development Status :: 4 - Beta
11
+ Classifier: Intended Audience :: Science/Research
12
+ Classifier: Intended Audience :: Developers
13
+ Classifier: License :: OSI Approved :: MIT License
14
+ Classifier: Programming Language :: Python :: 3
15
+ Classifier: Programming Language :: Python :: 3.10
16
+ Classifier: Programming Language :: Python :: 3.11
17
+ Classifier: Programming Language :: Python :: 3.12
18
+ Classifier: Topic :: Scientific/Engineering :: Artificial Intelligence
19
+ Classifier: Topic :: Scientific/Engineering :: Mathematics
20
+ Requires-Python: >=3.10
21
+ Description-Content-Type: text/markdown
22
+ License-File: LICENSE
23
+ Requires-Dist: ortools>=9.8
24
+ Requires-Dist: numpy>=1.26
25
+ Requires-Dist: scipy>=1.11
26
+ Provides-Extra: plot
27
+ Requires-Dist: pandas>=2.0; extra == "plot"
28
+ Requires-Dist: matplotlib>=3.8; extra == "plot"
29
+ Provides-Extra: dev
30
+ Requires-Dist: pytest>=8.0; extra == "dev"
31
+ Requires-Dist: pandas>=2.0; extra == "dev"
32
+ Requires-Dist: matplotlib>=3.8; extra == "dev"
33
+ Dynamic: license-file
34
+
35
+ # datacarve
36
+
37
+ [![CI](https://github.com/bbonik/datacarve/actions/workflows/ci.yml/badge.svg)](https://github.com/bbonik/datacarve/actions/workflows/ci.yml)
38
+ [![Python 3.10+](https://img.shields.io/badge/python-3.10%2B-blue.svg)](https://www.python.org/downloads/)
39
+ [![License: MIT](https://img.shields.io/badge/License-MIT-yellow.svg)](LICENSE)
40
+
41
+ **Carve a balanced subset out of a large dataset — distributional dataset undersampling via MILP optimization.**
42
+
43
+ ```bash
44
+ pip install datacarve
45
+ ```
46
+ ```python
47
+ from datacarve import undersample_dataset
48
+
49
+ mask = undersample_dataset(data, data_to_keep=1000) # balanced across ALL dimensions
50
+ subset = data[mask]
51
+ ```
52
+
53
+ `datacarve` selects the **provably optimal subset** of a dataset whose attributes jointly follow the distributions you specify — balanced across any number of dimensions, numeric or categorical, *all at once* (e.g. gender *and* age *and* race *and* label). Typical uses: **fair evaluation sets** for bias audits and Responsible AI compliance, **LLM data mixtures** (eval suites, SFT subsets, red-teaming pools), quota samples, and matched cohorts. Under the hood it is a Mixed Integer Linear Programming (**MILP**) formulation that exploits the redundancies of a large dataset to carve a compact, distribution-shaped version of it, while also minimizing cross-attribute correlations. Formerly known as `distributional_dataset_undersampling`.
54
+
55
+ <img src="https://github.com/bbonik/datacarve/raw/master/assets/example.png" width="900">
56
+
57
+ ## The problem: your dataset is imbalanced in several ways at once
58
+
59
+ Real datasets are rarely skewed along just one attribute. Take the classic Adult census dataset (48,842 rows): **two-thirds male, 85% White, 76% low-income, ages bunched between 25 and 45** — four imbalances at the same time. Train or evaluate on it as-is, and your metrics are quietly dominated by the majority groups.
60
+
61
+ Fixing **one** attribute is easy: group by it, sample equally per group. Fixing **all of them at once** is a different kind of problem, and this is the part few people appreciate until they try:
62
+
63
+ - **Every row you keep counts toward every histogram simultaneously.** A row that improves your gender balance may worsen your age balance. There is no "safe" row to drop.
64
+ - **Stratifying on the combination of attributes explodes.** 2 sexes × 5 races × 2 income classes × 10 age bins = 200 strata — most of which are nearly or completely empty in the original data. You cannot sample equally from empty strata.
65
+ - **Greedy selection has no guarantee.** Picking whichever row locally improves balance routinely paints itself into corners where every remaining candidate makes some attribute worse.
66
+
67
+ Selecting the best possible subset under joint distributional constraints is a **combinatorial optimization problem**. Treating it like one — instead of approximating with heuristics — is the whole point of this package.
68
+
69
+ ## How datacarve solves it
70
+
71
+ `datacarve` formulates the selection as a **Mixed Integer Linear Program (MILP)**: one binary keep/drop decision per row, constraints that tie the selected counts in every (attribute, bin) cell to your target distribution, and an objective that minimizes total deviation from the targets while also suppressing cross-attribute correlations.
72
+
73
+ ```mermaid
74
+ flowchart LR
75
+ A["Large skewed dataset<br/>(N rows, M attributes)"] --> Q["Quantize each attribute<br/>into bins / categories"]
76
+ T["Target distribution<br/>per attribute<br/>(uniform, gaussian, custom)"] --> S
77
+ K["Subset size K"] --> S
78
+ Q --> S{"MILP solver<br/>one binary variable per row:<br/>keep or drop"}
79
+ S --> O["Optimal subset of<br/>K real rows"]
80
+ O --> R["All M marginals match<br/>their targets jointly<br/>+ minimal cross-correlations"]
81
+ ```
82
+
83
+ The solver either **proves it found the optimal subset**, or — given a time budget — returns the best subset found with a quality bound. Three properties fall out of this that heuristics cannot offer:
84
+
85
+ 1. **Exactness.** The selected counts per group are guaranteed, not approximate: you can state "200 rows per race, 500 per sex" in a datasheet and mean it.
86
+ 2. **Jointness.** All attributes are satisfied *simultaneously* — numeric ones shaped to any distribution (uniform, gaussian, custom histogram), categorical ones to exact per-category counts.
87
+ 3. **Real data only.** The subset is made of your actual rows. Nothing is synthesized, duplicated, or reweighted.
88
+
89
+ The technique is *complementary to dimensionality reduction*: instead of reducing feature dimensions while keeping all observations, it reduces observations while imposing distributional constraints on the dimensions.
90
+
91
+ The figure above shows it in action on a 6-dimensional dataset (11K datapoints), where dimension D5 is a linear combination of D0 and D3. Three 1K subsets are carved with Uniform, Gaussian and Triangular targets: every histogram takes the target shape, and the D0–D5 correlation visible in the original is broken in the subsets.
92
+
93
+ ## Fairness and Responsible AI
94
+
95
+ This is the use case the package was built around, and it has only become more urgent since the original papers (ICIP 2016 / IEEE TMM 2017). Today, model cards, datasheets, bias audits, and regulations such as the **EU AI Act** all expect evidence that systems were evaluated on **representative, balanced data across sensitive attributes** — and "we randomly sampled and hoped" does not qualify.
96
+
97
+ `datacarve` turns that requirement into a one-liner with a provable result:
98
+
99
+ - **Balanced evaluation sets.** In a random 1,000-row sample of Adult, the smallest racial group gets ~8 rows — its accuracy estimate is statistical noise that swings on a couple of lucky predictions. A carved set gives *every* group the same 200-row evidence base, making per-group metrics comparable and equally trustworthy. See the [worked notebook](notebooks/balanced_evaluation_sets.ipynb): sex 500/500, race 5×200, income 500/500, age flat — simultaneously, in seconds.
100
+ - **Auditable by construction.** Because the constraints are explicit and the solver's result status is reported, the composition of your eval set is a *documented guarantee*, not a post-hoc observation — exactly what a datasheet or compliance review wants to see.
101
+ - **Realistic, not just uniform, targets.** Fairness rarely means "make everything equal". Per-attribute targets let you balance sensitive attributes exactly while keeping, say, a realistic 3:1 label ratio: `target_distribution=["uniform", "uniform", [3, 1]]`.
102
+
103
+ ## Applications
104
+
105
+ Any situation where you need a **subset of fixed size whose attributes follow prescribed distributions, jointly across several attributes**, is a candidate.
106
+
107
+ ### In the LLM era
108
+
109
+ Modern LLM work is largely *data curation under a budget* — which is exactly this problem. Attributes don't need to be raw columns: task labels, topic clusters from embeddings, difficulty scores, or length buckets all work.
110
+
111
+ - **Balanced benchmark & eval suites.** Carve an evaluation set that is balanced across task type × domain × difficulty × language × prompt length, so a model's headline score isn't dominated by whichever category the benchmark over-collected. Same for regression-testing suites that must stay small enough to run on every checkpoint.
112
+ - **Fine-tuning mixtures (SFT).** Instruction datasets skew heavily by source, topic and length. Carve a compact training subset that hits an exact target mixture (e.g. 30% coding, 30% reasoning, 20% writing, 20% multilingual — with a target length distribution) instead of eyeballing sampling ratios.
113
+ - **Safety & red-teaming sets.** Balance adversarial prompts across harm categories × attack styles × targeted demographics, so safety metrics cover the space instead of over-testing the most common attack type.
114
+ - **Human evaluation & preference data.** Annotator time is the scarcest resource in RLHF pipelines; carve the candidate pool so every scenario type gets equal annotation coverage.
115
+
116
+ ### Classical ML and beyond
117
+
118
+ - **Dataset debiasing / data-centric AI.** Reshape a skewed training set toward a target distribution instead of collecting new data, leveraging redundancy already present in the dataset.
119
+ - **Causal inference & epidemiology.** Select a control cohort whose covariate distributions match a treatment group (or any reference population). This generalizes matching approaches such as cardinality matching: the target can be *any* distribution, not just another group's.
120
+ - **Survey statistics & market research.** Quota sampling and panel calibration: pick respondents so that the sample matches census demographics across several attributes simultaneously ([worked notebook](notebooks/survey_quota_sampling.ipynb)).
121
+ - **A/B testing.** Assign experiment groups that are balanced across multiple covariates, rather than relying on randomization alone for small samples.
122
+ - **Drug discovery / cheminformatics.** Select compound libraries with desired property distributions (molecular weight, logP, solubility, ...) while minimizing redundancy between correlated properties.
123
+ - **Simulation & testing.** Choose a representative, affordable subset of test scenarios (e.g., driving scenarios spanning weather × traffic × speed distributions) when running all of them is too expensive.
124
+
125
+ ## How it compares to other approaches
126
+
127
+ | Approach | What it does | Limitation this method addresses |
128
+ |---|---|---|
129
+ | **Random / stratified sampling** | Samples uniformly, or balances strata of *one* attribute. | Cannot jointly balance several attributes; multi-attribute stratification explodes combinatorially and leaves many empty strata. |
130
+ | **Class balancing** (e.g., random undersampling, SMOTE in `imbalanced-learn`) | Balances a single categorical label, possibly by synthesizing points. | Single-label only; synthetic points may be unrealistic. This method handles multiple *continuous or categorical* dimensions and only ever selects real datapoints. |
131
+ | **Reweighting / calibration** (importance weights, raking) | Keeps all data but assigns weights so that weighted statistics match targets. | The dataset stays large and individual high-weight points dominate; many ML pipelines and human-evaluation settings need an *actual subset*, not weights. |
132
+ | **Matching methods** (propensity score, cardinality matching) | Selects a control group whose covariates match a treatment group. | Matches to *another sample's* distribution; here the target is arbitrary (uniform, gaussian, custom), and correlation between attributes is minimized explicitly. |
133
+ | **Coreset selection / data pruning** | Selects a subset that preserves model loss or gradient information. | Optimizes for a *model's* training objective, not for interpretable distributional guarantees; typically gives no control over per-attribute histograms. |
134
+ | **Greedy / heuristic subset selection** | Iteratively picks points that locally improve balance. | No global guarantee: a point that helps attribute A may hurt attribute B. The MILP reasons about all attributes and all points jointly, and returns a certified optimal (or bounded) solution. |
135
+
136
+ In short: this method occupies a niche none of the standard tools cover — **exact, jointly multi-attribute, distribution-targeted subset selection of real datapoints**. One binary decision variable per datapoint solves comfortably up to hundreds of thousands of rows on a laptop; for larger datasets the built-in [pre-reduction stage](#very-large-datasets) extends it to tens of millions.
137
+
138
+ ## Installation
139
+
140
+ Requires Python 3.10+ (tested with Python 3.12).
141
+
142
+ ```bash
143
+ pip install datacarve # core (solver only)
144
+ pip install "datacarve[plot]" # core + scatterplot matrices
145
+ ```
146
+
147
+ Or from source:
148
+
149
+ ```bash
150
+ git clone https://github.com/bbonik/datacarve.git
151
+ cd datacarve
152
+
153
+ # create and activate a virtual environment
154
+ python3 -m venv .venv
155
+ source .venv/bin/activate # on Windows: .venv\Scripts\activate
156
+
157
+ pip install -e ".[plot]"
158
+ ```
159
+
160
+ The MILP solver is [Google OR-Tools](https://developers.google.com/optimization) (CBC backend), which is installed automatically — no separate solver installation is needed.
161
+
162
+ ## Quick start
163
+
164
+ ```python
165
+ import numpy as np
166
+ from datacarve import undersample_dataset
167
+
168
+ rng = np.random.default_rng(0)
169
+ data = rng.random((5000, 4)) # [N observations, M dimensions]
170
+
171
+ mask = undersample_dataset(
172
+ data=data,
173
+ data_to_keep=500, # size of the undersampled subset
174
+ target_distribution="uniform", # 'uniform', 'gaussian', 'weibull', 'triangular'
175
+ bins=10, # quantization bins per dimension
176
+ lamda=0.5, # weight of the correlation-minimization objective
177
+ )
178
+
179
+ subset = data[mask] # boolean mask over the original observations
180
+ ```
181
+
182
+ You can also pass a **custom target distribution** as an array of bin weights (one weight per bin, automatically normalized):
183
+
184
+ ```python
185
+ # triangular-ish custom target over 10 bins
186
+ mask = undersample_dataset(
187
+ data=data,
188
+ data_to_keep=500,
189
+ target_distribution=[1, 2, 3, 4, 5, 5, 4, 3, 2, 1],
190
+ bins=10,
191
+ )
192
+ ```
193
+
194
+ Useful options:
195
+
196
+ | Parameter | Default | Description |
197
+ |---|---|---|
198
+ | `data_to_keep` | `1000` | Number of datapoints to keep. |
199
+ | `data_scaling` | `'minmax'` | Per-feature scaling to [0, 1]. Use `None` if the data is already scaled. |
200
+ | `target_distribution` | `'uniform'` | Built-in name, a custom array of bin weights, or a list with one spec per dimension. See [Per-dimension targets](#per-dimension-targets-and-categorical-attributes). |
201
+ | `bins` | `10` | Quantization bins per numeric dimension. Categorical dimensions use one bin per unique value. |
202
+ | `categorical_dims` | `None` | Column indices to treat as categorical (one bin per unique value). |
203
+ | `lamda` | `0.5` | Balance between distribution matching (`0`) and correlation minimization (`>0`). |
204
+ | `prereduce` | `None` | Pre-reduce huge datasets before solving: `'auto'`, or an int cap per joint cell. See [Very large datasets](#very-large-datasets). |
205
+ | `solver` | `'CBC'` | MILP solver backend: `'CBC'`, `'SCIP'`, or `'SAT'`. See [Choosing a solver](#choosing-a-solver). |
206
+ | `max_solver_time_sec` | `10.0` | Time budget for the MILP solver. Increase for large datasets. |
207
+ | `verbose` | `True` | Print progress and solver statistics. |
208
+ | `scatterplot_matrix` | `'auto'` | Show scatterplot matrices (auto-disabled for >10 dimensions). |
209
+
210
+ ## Per-dimension targets and categorical attributes
211
+
212
+ Each dimension can get its **own target distribution** — pass a list with one spec per dimension, mixing built-in names and custom weight arrays:
213
+
214
+ ```python
215
+ mask = undersample_dataset(
216
+ data=data, # shape (N, 3)
217
+ data_to_keep=500,
218
+ target_distribution=["uniform", "gaussian", [1, 2, 3, 4, 5, 5, 4, 3, 2, 1]],
219
+ )
220
+ ```
221
+
222
+ **Categorical attributes** (e.g. gender, race, class labels) should not be quantized into equal-width bins. Mark them with `categorical_dims` and each unique value becomes its own bin, so `'uniform'` means "equal counts per category":
223
+
224
+ ```python
225
+ # column 2 holds a label-encoded category (e.g. 0=A, 1=B, 2=C)
226
+ mask = undersample_dataset(
227
+ data=data,
228
+ data_to_keep=300,
229
+ target_distribution=["uniform", "uniform", [3, 2, 1]], # 3:2:1 over categories
230
+ categorical_dims=[2],
231
+ )
232
+ ```
233
+
234
+ This is the typical recipe for fairness-style curation: balance the categorical attributes exactly (equal counts per gender/race/label) while shaping the continuous attributes (age, pose, brightness) to a target distribution — all jointly, in one optimization.
235
+
236
+ ## Very large datasets
237
+
238
+ The MILP uses one binary variable per row, which is comfortable up to several hundred thousand rows. Beyond that, use the built-in **pre-reduction** stage:
239
+
240
+ ```python
241
+ mask = undersample_dataset(
242
+ data=huge_data, # e.g. 10 million rows
243
+ data_to_keep=1000,
244
+ prereduce="auto", # or an explicit per-cell cap, e.g. prereduce=50
245
+ )
246
+ ```
247
+
248
+ Pre-reduction groups rows by their joint quantization cell (the combination of bin indices across all attributes). Rows in the same cell are interchangeable with respect to every histogram constraint, so overcrowded cells are randomly downsampled to a cap while **rare cells are always kept in full** — unlike naive random subsampling, which preserves the skew you are trying to fix and can wipe out rare categories entirely. With `'auto'`, the cap is chosen adaptively so the reduced pool stays at a size the solver handles in seconds. The returned mask always refers to the original rows.
249
+
250
+ Measured on a laptop: a 10-million-row dataset is carved into a perfectly balanced 1,000-row subset (proven optimal) in about 10 seconds end-to-end.
251
+
252
+ The standalone `prereduce_dataset()` function exposes the same stage with control over the cap, grouping granularity, and random seed.
253
+
254
+ ## Choosing a solver
255
+
256
+ All three backends are free, open source, and bundled with OR-Tools — no extra installation needed. They solve the exact same model; they differ in *how* they search, which matters once problems get hard.
257
+
258
+ | Solver | Best for | Character |
259
+ |---|---|---|
260
+ | `'CBC'` (default) | Easy to moderate problems | Classic branch-and-bound. Fastest when the problem is not too constrained; if it reports `optimal` within the time budget, stay with it. |
261
+ | `'SAT'` (CP-SAT) | Hard instances that hit the time limit | Clause-learning search, multi-core. When the status is `feasible` (time ran out before optimality was proven), it typically finds noticeably *better* subsets than CBC in the same time budget. |
262
+ | `'SCIP'` | Medium-hard instances | Modern branch-and-cut. Worth trying when CBC finds a solution quickly but struggles to prove it optimal. |
263
+
264
+ **Rule of thumb:**
265
+
266
+ 1. Start with the default (`'CBC'`).
267
+ 2. Check the reported result status (printed when `verbose=True`).
268
+ 3. If the status is `optimal` — done, no reason to switch.
269
+ 4. If the status is `feasible` (the time budget ran out), re-run with `solver='SAT'` and/or a larger `max_solver_time_sec`. This is where CP-SAT shines: on a hard 11K-point benchmark instance, all solvers hit a 60s budget, but CP-SAT returned the best subset found.
270
+ 5. If no solution is found at all, the constraints may be too tight for your data: increase `max_solver_time_sec`, reduce `bins`, or reduce `data_to_keep`.
271
+
272
+ What makes an instance "hard"? More datapoints, more dimensions, strongly imbalanced data relative to the target (little redundancy to exploit), and correlated dimensions all increase difficulty.
273
+
274
+ ## Examples
275
+
276
+ Two executed walkthrough notebooks in [`notebooks/`](notebooks/):
277
+
278
+ - **[Building fair, balanced evaluation sets](notebooks/balanced_evaluation_sets.ipynb)** — carves a 1,000-row eval set from the Adult census data, balanced across sex, race, income and age *simultaneously*, and shows why per-group accuracy numbers become trustworthy.
279
+ - **[Survey quota sampling](notebooks/survey_quota_sampling.ipynb)** — selects a quota sample from a skewed respondent panel, hitting census-style age/gender/region targets exactly (fully offline, synthetic data).
280
+
281
+ Runnable scripts in the [`examples/`](examples/) folder:
282
+
283
+ - **`example_6d_dataset.py`** — undersamples the bundled 6-dimensional dataset (11K datapoints) down to a uniform 1K subset.
284
+ - **`example_random_data.py`** — generates a random N-dimensional dataset (a different random distribution per dimension) and undersamples it.
285
+ - **`example_sklearn_datasets.py`** — applies the technique to classic scikit-learn datasets (diabetes, iris, breast cancer). Requires `scikit-learn`.
286
+
287
+ ```bash
288
+ python examples/example_6d_dataset.py
289
+ ```
290
+
291
+ Solver benchmarks live in [`benchmarks/`](benchmarks/).
292
+
293
+ ## Citations
294
+
295
+ If you use this code in your research please cite the following papers:
296
+
297
+ 1. [Vonikakis, V., Subramanian, R., Arnfred, J., & Winkler, S. A Probabilistic Approach to People-Centric Photo Selection and Sequencing. IEEE Transactions in Multimedia, 11(19), pp.2609-2624, 2017.](https://www.researchgate.net/publication/316569587_A_Probabilistic_Approach_to_People-Centric_Photo_Selection_and_Sequencing)
298
+ 2. [V. Vonikakis, R. Subramanian, S. Winkler. Shaping Datasets: Optimal Data Selection for Specific Target Distributions. Proc. ICIP2016, Phoenix, USA, Sept. 25-28, 2016.](http://vintage.winklerbros.net/Publications/icip2016a.pdf)
@@ -0,0 +1,264 @@
1
+ # datacarve
2
+
3
+ [![CI](https://github.com/bbonik/datacarve/actions/workflows/ci.yml/badge.svg)](https://github.com/bbonik/datacarve/actions/workflows/ci.yml)
4
+ [![Python 3.10+](https://img.shields.io/badge/python-3.10%2B-blue.svg)](https://www.python.org/downloads/)
5
+ [![License: MIT](https://img.shields.io/badge/License-MIT-yellow.svg)](LICENSE)
6
+
7
+ **Carve a balanced subset out of a large dataset — distributional dataset undersampling via MILP optimization.**
8
+
9
+ ```bash
10
+ pip install datacarve
11
+ ```
12
+ ```python
13
+ from datacarve import undersample_dataset
14
+
15
+ mask = undersample_dataset(data, data_to_keep=1000) # balanced across ALL dimensions
16
+ subset = data[mask]
17
+ ```
18
+
19
+ `datacarve` selects the **provably optimal subset** of a dataset whose attributes jointly follow the distributions you specify — balanced across any number of dimensions, numeric or categorical, *all at once* (e.g. gender *and* age *and* race *and* label). Typical uses: **fair evaluation sets** for bias audits and Responsible AI compliance, **LLM data mixtures** (eval suites, SFT subsets, red-teaming pools), quota samples, and matched cohorts. Under the hood it is a Mixed Integer Linear Programming (**MILP**) formulation that exploits the redundancies of a large dataset to carve a compact, distribution-shaped version of it, while also minimizing cross-attribute correlations. Formerly known as `distributional_dataset_undersampling`.
20
+
21
+ <img src="https://github.com/bbonik/datacarve/raw/master/assets/example.png" width="900">
22
+
23
+ ## The problem: your dataset is imbalanced in several ways at once
24
+
25
+ Real datasets are rarely skewed along just one attribute. Take the classic Adult census dataset (48,842 rows): **two-thirds male, 85% White, 76% low-income, ages bunched between 25 and 45** — four imbalances at the same time. Train or evaluate on it as-is, and your metrics are quietly dominated by the majority groups.
26
+
27
+ Fixing **one** attribute is easy: group by it, sample equally per group. Fixing **all of them at once** is a different kind of problem, and this is the part few people appreciate until they try:
28
+
29
+ - **Every row you keep counts toward every histogram simultaneously.** A row that improves your gender balance may worsen your age balance. There is no "safe" row to drop.
30
+ - **Stratifying on the combination of attributes explodes.** 2 sexes × 5 races × 2 income classes × 10 age bins = 200 strata — most of which are nearly or completely empty in the original data. You cannot sample equally from empty strata.
31
+ - **Greedy selection has no guarantee.** Picking whichever row locally improves balance routinely paints itself into corners where every remaining candidate makes some attribute worse.
32
+
33
+ Selecting the best possible subset under joint distributional constraints is a **combinatorial optimization problem**. Treating it like one — instead of approximating with heuristics — is the whole point of this package.
34
+
35
+ ## How datacarve solves it
36
+
37
+ `datacarve` formulates the selection as a **Mixed Integer Linear Program (MILP)**: one binary keep/drop decision per row, constraints that tie the selected counts in every (attribute, bin) cell to your target distribution, and an objective that minimizes total deviation from the targets while also suppressing cross-attribute correlations.
38
+
39
+ ```mermaid
40
+ flowchart LR
41
+ A["Large skewed dataset<br/>(N rows, M attributes)"] --> Q["Quantize each attribute<br/>into bins / categories"]
42
+ T["Target distribution<br/>per attribute<br/>(uniform, gaussian, custom)"] --> S
43
+ K["Subset size K"] --> S
44
+ Q --> S{"MILP solver<br/>one binary variable per row:<br/>keep or drop"}
45
+ S --> O["Optimal subset of<br/>K real rows"]
46
+ O --> R["All M marginals match<br/>their targets jointly<br/>+ minimal cross-correlations"]
47
+ ```
48
+
49
+ The solver either **proves it found the optimal subset**, or — given a time budget — returns the best subset found with a quality bound. Three properties fall out of this that heuristics cannot offer:
50
+
51
+ 1. **Exactness.** The selected counts per group are guaranteed, not approximate: you can state "200 rows per race, 500 per sex" in a datasheet and mean it.
52
+ 2. **Jointness.** All attributes are satisfied *simultaneously* — numeric ones shaped to any distribution (uniform, gaussian, custom histogram), categorical ones to exact per-category counts.
53
+ 3. **Real data only.** The subset is made of your actual rows. Nothing is synthesized, duplicated, or reweighted.
54
+
55
+ The technique is *complementary to dimensionality reduction*: instead of reducing feature dimensions while keeping all observations, it reduces observations while imposing distributional constraints on the dimensions.
56
+
57
+ The figure above shows it in action on a 6-dimensional dataset (11K datapoints), where dimension D5 is a linear combination of D0 and D3. Three 1K subsets are carved with Uniform, Gaussian and Triangular targets: every histogram takes the target shape, and the D0–D5 correlation visible in the original is broken in the subsets.
58
+
59
+ ## Fairness and Responsible AI
60
+
61
+ This is the use case the package was built around, and it has only become more urgent since the original papers (ICIP 2016 / IEEE TMM 2017). Today, model cards, datasheets, bias audits, and regulations such as the **EU AI Act** all expect evidence that systems were evaluated on **representative, balanced data across sensitive attributes** — and "we randomly sampled and hoped" does not qualify.
62
+
63
+ `datacarve` turns that requirement into a one-liner with a provable result:
64
+
65
+ - **Balanced evaluation sets.** In a random 1,000-row sample of Adult, the smallest racial group gets ~8 rows — its accuracy estimate is statistical noise that swings on a couple of lucky predictions. A carved set gives *every* group the same 200-row evidence base, making per-group metrics comparable and equally trustworthy. See the [worked notebook](notebooks/balanced_evaluation_sets.ipynb): sex 500/500, race 5×200, income 500/500, age flat — simultaneously, in seconds.
66
+ - **Auditable by construction.** Because the constraints are explicit and the solver's result status is reported, the composition of your eval set is a *documented guarantee*, not a post-hoc observation — exactly what a datasheet or compliance review wants to see.
67
+ - **Realistic, not just uniform, targets.** Fairness rarely means "make everything equal". Per-attribute targets let you balance sensitive attributes exactly while keeping, say, a realistic 3:1 label ratio: `target_distribution=["uniform", "uniform", [3, 1]]`.
68
+
69
+ ## Applications
70
+
71
+ Any situation where you need a **subset of fixed size whose attributes follow prescribed distributions, jointly across several attributes**, is a candidate.
72
+
73
+ ### In the LLM era
74
+
75
+ Modern LLM work is largely *data curation under a budget* — which is exactly this problem. Attributes don't need to be raw columns: task labels, topic clusters from embeddings, difficulty scores, or length buckets all work.
76
+
77
+ - **Balanced benchmark & eval suites.** Carve an evaluation set that is balanced across task type × domain × difficulty × language × prompt length, so a model's headline score isn't dominated by whichever category the benchmark over-collected. Same for regression-testing suites that must stay small enough to run on every checkpoint.
78
+ - **Fine-tuning mixtures (SFT).** Instruction datasets skew heavily by source, topic and length. Carve a compact training subset that hits an exact target mixture (e.g. 30% coding, 30% reasoning, 20% writing, 20% multilingual — with a target length distribution) instead of eyeballing sampling ratios.
79
+ - **Safety & red-teaming sets.** Balance adversarial prompts across harm categories × attack styles × targeted demographics, so safety metrics cover the space instead of over-testing the most common attack type.
80
+ - **Human evaluation & preference data.** Annotator time is the scarcest resource in RLHF pipelines; carve the candidate pool so every scenario type gets equal annotation coverage.
81
+
82
+ ### Classical ML and beyond
83
+
84
+ - **Dataset debiasing / data-centric AI.** Reshape a skewed training set toward a target distribution instead of collecting new data, leveraging redundancy already present in the dataset.
85
+ - **Causal inference & epidemiology.** Select a control cohort whose covariate distributions match a treatment group (or any reference population). This generalizes matching approaches such as cardinality matching: the target can be *any* distribution, not just another group's.
86
+ - **Survey statistics & market research.** Quota sampling and panel calibration: pick respondents so that the sample matches census demographics across several attributes simultaneously ([worked notebook](notebooks/survey_quota_sampling.ipynb)).
87
+ - **A/B testing.** Assign experiment groups that are balanced across multiple covariates, rather than relying on randomization alone for small samples.
88
+ - **Drug discovery / cheminformatics.** Select compound libraries with desired property distributions (molecular weight, logP, solubility, ...) while minimizing redundancy between correlated properties.
89
+ - **Simulation & testing.** Choose a representative, affordable subset of test scenarios (e.g., driving scenarios spanning weather × traffic × speed distributions) when running all of them is too expensive.
90
+
91
+ ## How it compares to other approaches
92
+
93
+ | Approach | What it does | Limitation this method addresses |
94
+ |---|---|---|
95
+ | **Random / stratified sampling** | Samples uniformly, or balances strata of *one* attribute. | Cannot jointly balance several attributes; multi-attribute stratification explodes combinatorially and leaves many empty strata. |
96
+ | **Class balancing** (e.g., random undersampling, SMOTE in `imbalanced-learn`) | Balances a single categorical label, possibly by synthesizing points. | Single-label only; synthetic points may be unrealistic. This method handles multiple *continuous or categorical* dimensions and only ever selects real datapoints. |
97
+ | **Reweighting / calibration** (importance weights, raking) | Keeps all data but assigns weights so that weighted statistics match targets. | The dataset stays large and individual high-weight points dominate; many ML pipelines and human-evaluation settings need an *actual subset*, not weights. |
98
+ | **Matching methods** (propensity score, cardinality matching) | Selects a control group whose covariates match a treatment group. | Matches to *another sample's* distribution; here the target is arbitrary (uniform, gaussian, custom), and correlation between attributes is minimized explicitly. |
99
+ | **Coreset selection / data pruning** | Selects a subset that preserves model loss or gradient information. | Optimizes for a *model's* training objective, not for interpretable distributional guarantees; typically gives no control over per-attribute histograms. |
100
+ | **Greedy / heuristic subset selection** | Iteratively picks points that locally improve balance. | No global guarantee: a point that helps attribute A may hurt attribute B. The MILP reasons about all attributes and all points jointly, and returns a certified optimal (or bounded) solution. |
101
+
102
+ In short: this method occupies a niche none of the standard tools cover — **exact, jointly multi-attribute, distribution-targeted subset selection of real datapoints**. One binary decision variable per datapoint solves comfortably up to hundreds of thousands of rows on a laptop; for larger datasets the built-in [pre-reduction stage](#very-large-datasets) extends it to tens of millions.
103
+
104
+ ## Installation
105
+
106
+ Requires Python 3.10+ (tested with Python 3.12).
107
+
108
+ ```bash
109
+ pip install datacarve # core (solver only)
110
+ pip install "datacarve[plot]" # core + scatterplot matrices
111
+ ```
112
+
113
+ Or from source:
114
+
115
+ ```bash
116
+ git clone https://github.com/bbonik/datacarve.git
117
+ cd datacarve
118
+
119
+ # create and activate a virtual environment
120
+ python3 -m venv .venv
121
+ source .venv/bin/activate # on Windows: .venv\Scripts\activate
122
+
123
+ pip install -e ".[plot]"
124
+ ```
125
+
126
+ The MILP solver is [Google OR-Tools](https://developers.google.com/optimization) (CBC backend), which is installed automatically — no separate solver installation is needed.
127
+
128
+ ## Quick start
129
+
130
+ ```python
131
+ import numpy as np
132
+ from datacarve import undersample_dataset
133
+
134
+ rng = np.random.default_rng(0)
135
+ data = rng.random((5000, 4)) # [N observations, M dimensions]
136
+
137
+ mask = undersample_dataset(
138
+ data=data,
139
+ data_to_keep=500, # size of the undersampled subset
140
+ target_distribution="uniform", # 'uniform', 'gaussian', 'weibull', 'triangular'
141
+ bins=10, # quantization bins per dimension
142
+ lamda=0.5, # weight of the correlation-minimization objective
143
+ )
144
+
145
+ subset = data[mask] # boolean mask over the original observations
146
+ ```
147
+
148
+ You can also pass a **custom target distribution** as an array of bin weights (one weight per bin, automatically normalized):
149
+
150
+ ```python
151
+ # triangular-ish custom target over 10 bins
152
+ mask = undersample_dataset(
153
+ data=data,
154
+ data_to_keep=500,
155
+ target_distribution=[1, 2, 3, 4, 5, 5, 4, 3, 2, 1],
156
+ bins=10,
157
+ )
158
+ ```
159
+
160
+ Useful options:
161
+
162
+ | Parameter | Default | Description |
163
+ |---|---|---|
164
+ | `data_to_keep` | `1000` | Number of datapoints to keep. |
165
+ | `data_scaling` | `'minmax'` | Per-feature scaling to [0, 1]. Use `None` if the data is already scaled. |
166
+ | `target_distribution` | `'uniform'` | Built-in name, a custom array of bin weights, or a list with one spec per dimension. See [Per-dimension targets](#per-dimension-targets-and-categorical-attributes). |
167
+ | `bins` | `10` | Quantization bins per numeric dimension. Categorical dimensions use one bin per unique value. |
168
+ | `categorical_dims` | `None` | Column indices to treat as categorical (one bin per unique value). |
169
+ | `lamda` | `0.5` | Balance between distribution matching (`0`) and correlation minimization (`>0`). |
170
+ | `prereduce` | `None` | Pre-reduce huge datasets before solving: `'auto'`, or an int cap per joint cell. See [Very large datasets](#very-large-datasets). |
171
+ | `solver` | `'CBC'` | MILP solver backend: `'CBC'`, `'SCIP'`, or `'SAT'`. See [Choosing a solver](#choosing-a-solver). |
172
+ | `max_solver_time_sec` | `10.0` | Time budget for the MILP solver. Increase for large datasets. |
173
+ | `verbose` | `True` | Print progress and solver statistics. |
174
+ | `scatterplot_matrix` | `'auto'` | Show scatterplot matrices (auto-disabled for >10 dimensions). |
175
+
176
+ ## Per-dimension targets and categorical attributes
177
+
178
+ Each dimension can get its **own target distribution** — pass a list with one spec per dimension, mixing built-in names and custom weight arrays:
179
+
180
+ ```python
181
+ mask = undersample_dataset(
182
+ data=data, # shape (N, 3)
183
+ data_to_keep=500,
184
+ target_distribution=["uniform", "gaussian", [1, 2, 3, 4, 5, 5, 4, 3, 2, 1]],
185
+ )
186
+ ```
187
+
188
+ **Categorical attributes** (e.g. gender, race, class labels) should not be quantized into equal-width bins. Mark them with `categorical_dims` and each unique value becomes its own bin, so `'uniform'` means "equal counts per category":
189
+
190
+ ```python
191
+ # column 2 holds a label-encoded category (e.g. 0=A, 1=B, 2=C)
192
+ mask = undersample_dataset(
193
+ data=data,
194
+ data_to_keep=300,
195
+ target_distribution=["uniform", "uniform", [3, 2, 1]], # 3:2:1 over categories
196
+ categorical_dims=[2],
197
+ )
198
+ ```
199
+
200
+ This is the typical recipe for fairness-style curation: balance the categorical attributes exactly (equal counts per gender/race/label) while shaping the continuous attributes (age, pose, brightness) to a target distribution — all jointly, in one optimization.
201
+
202
+ ## Very large datasets
203
+
204
+ The MILP uses one binary variable per row, which is comfortable up to several hundred thousand rows. Beyond that, use the built-in **pre-reduction** stage:
205
+
206
+ ```python
207
+ mask = undersample_dataset(
208
+ data=huge_data, # e.g. 10 million rows
209
+ data_to_keep=1000,
210
+ prereduce="auto", # or an explicit per-cell cap, e.g. prereduce=50
211
+ )
212
+ ```
213
+
214
+ Pre-reduction groups rows by their joint quantization cell (the combination of bin indices across all attributes). Rows in the same cell are interchangeable with respect to every histogram constraint, so overcrowded cells are randomly downsampled to a cap while **rare cells are always kept in full** — unlike naive random subsampling, which preserves the skew you are trying to fix and can wipe out rare categories entirely. With `'auto'`, the cap is chosen adaptively so the reduced pool stays at a size the solver handles in seconds. The returned mask always refers to the original rows.
215
+
216
+ Measured on a laptop: a 10-million-row dataset is carved into a perfectly balanced 1,000-row subset (proven optimal) in about 10 seconds end-to-end.
217
+
218
+ The standalone `prereduce_dataset()` function exposes the same stage with control over the cap, grouping granularity, and random seed.
219
+
220
+ ## Choosing a solver
221
+
222
+ All three backends are free, open source, and bundled with OR-Tools — no extra installation needed. They solve the exact same model; they differ in *how* they search, which matters once problems get hard.
223
+
224
+ | Solver | Best for | Character |
225
+ |---|---|---|
226
+ | `'CBC'` (default) | Easy to moderate problems | Classic branch-and-bound. Fastest when the problem is not too constrained; if it reports `optimal` within the time budget, stay with it. |
227
+ | `'SAT'` (CP-SAT) | Hard instances that hit the time limit | Clause-learning search, multi-core. When the status is `feasible` (time ran out before optimality was proven), it typically finds noticeably *better* subsets than CBC in the same time budget. |
228
+ | `'SCIP'` | Medium-hard instances | Modern branch-and-cut. Worth trying when CBC finds a solution quickly but struggles to prove it optimal. |
229
+
230
+ **Rule of thumb:**
231
+
232
+ 1. Start with the default (`'CBC'`).
233
+ 2. Check the reported result status (printed when `verbose=True`).
234
+ 3. If the status is `optimal` — done, no reason to switch.
235
+ 4. If the status is `feasible` (the time budget ran out), re-run with `solver='SAT'` and/or a larger `max_solver_time_sec`. This is where CP-SAT shines: on a hard 11K-point benchmark instance, all solvers hit a 60s budget, but CP-SAT returned the best subset found.
236
+ 5. If no solution is found at all, the constraints may be too tight for your data: increase `max_solver_time_sec`, reduce `bins`, or reduce `data_to_keep`.
237
+
238
+ What makes an instance "hard"? More datapoints, more dimensions, strongly imbalanced data relative to the target (little redundancy to exploit), and correlated dimensions all increase difficulty.
239
+
240
+ ## Examples
241
+
242
+ Two executed walkthrough notebooks in [`notebooks/`](notebooks/):
243
+
244
+ - **[Building fair, balanced evaluation sets](notebooks/balanced_evaluation_sets.ipynb)** — carves a 1,000-row eval set from the Adult census data, balanced across sex, race, income and age *simultaneously*, and shows why per-group accuracy numbers become trustworthy.
245
+ - **[Survey quota sampling](notebooks/survey_quota_sampling.ipynb)** — selects a quota sample from a skewed respondent panel, hitting census-style age/gender/region targets exactly (fully offline, synthetic data).
246
+
247
+ Runnable scripts in the [`examples/`](examples/) folder:
248
+
249
+ - **`example_6d_dataset.py`** — undersamples the bundled 6-dimensional dataset (11K datapoints) down to a uniform 1K subset.
250
+ - **`example_random_data.py`** — generates a random N-dimensional dataset (a different random distribution per dimension) and undersamples it.
251
+ - **`example_sklearn_datasets.py`** — applies the technique to classic scikit-learn datasets (diabetes, iris, breast cancer). Requires `scikit-learn`.
252
+
253
+ ```bash
254
+ python examples/example_6d_dataset.py
255
+ ```
256
+
257
+ Solver benchmarks live in [`benchmarks/`](benchmarks/).
258
+
259
+ ## Citations
260
+
261
+ If you use this code in your research please cite the following papers:
262
+
263
+ 1. [Vonikakis, V., Subramanian, R., Arnfred, J., & Winkler, S. A Probabilistic Approach to People-Centric Photo Selection and Sequencing. IEEE Transactions in Multimedia, 11(19), pp.2609-2624, 2017.](https://www.researchgate.net/publication/316569587_A_Probabilistic_Approach_to_People-Centric_Photo_Selection_and_Sequencing)
264
+ 2. [V. Vonikakis, R. Subramanian, S. Winkler. Shaping Datasets: Optimal Data Selection for Specific Target Distributions. Proc. ICIP2016, Phoenix, USA, Sept. 25-28, 2016.](http://vintage.winklerbros.net/Publications/icip2016a.pdf)