Statistical Context Engineering: Hierarchical feature enrichment for regression models
Project description
Statistical Context Engineering (SCE)
Hierarchical feature enrichment for regression with leakage-safe cross-fitting
Overview
Statistical Context Engineering (SCE) is a Python library for enriching tabular regression datasets with hierarchical statistical features. The method computes target-variable aggregations at multiple levels of categorical granularity (e.g., region → city → neighborhood) while preventing information leakage through k-fold cross-fitting.
The approach is grounded in Bayesian estimation principles and cooperative game theory, providing a principled framework for incorporating group-level context into predictive models.
Features
- Leakage prevention: Out-of-fold aggregation ensures no target information from the test set contaminates training features
- Automatic detection: Categorical columns are inferred from DataFrame dtypes when not explicitly specified
- Hierarchical aggregation: Captures statistical context at multiple granularity levels, including cross-column interactions
- Scikit-learn compatibility: Implements the transformer interface for seamless pipeline integration
- Uncertainty quantification: Cross-fitting provides fold-variance estimates for context features
Installation
From PyPI
pip install stat-context
From Source
git clone https://github.com/joint-hubs/sce.git
cd sce
pip install -e .
Optional Dependencies
pip install stat-context[dev] # Development and testing tools
pip install stat-context[data] # Remote dataset fetching
pip install stat-context[all] # All optional dependencies
Usage
Basic Example
from sce import StatisticalContextEngine, ContextConfig, AggregationMethod
from sklearn.pipeline import Pipeline
from sklearn.model_selection import train_test_split
from xgboost import XGBRegressor
import pandas as pd
# Load data
df = pd.read_csv("rental_data.csv")
X = df.drop(columns=["price"])
y = df["price"]
# Configure the context engine
config = ContextConfig(
target_col="price",
aggregations=[AggregationMethod.MEAN, AggregationMethod.STD],
use_cross_fitting=True,
n_folds=5
)
# Build pipeline
pipeline = Pipeline([
("sce", StatisticalContextEngine(config)),
("model", XGBRegressor(n_estimators=100))
])
# Train and predict
X_train, X_test, y_train, y_test = train_test_split(X, y, test_size=0.2)
pipeline.fit(X_train, y_train)
predictions = pipeline.predict(X_test)
Explicit Column Specification
config = ContextConfig(
target_col="price",
categorical_cols=["country", "region", "city", "neighborhood"],
aggregations=[
AggregationMethod.MEAN,
AggregationMethod.MEDIAN,
AggregationMethod.STD,
AggregationMethod.Q25,
AggregationMethod.Q75,
AggregationMethod.COUNT
],
use_cross_fitting=True,
n_folds=5,
min_group_size=5,
include_global_stats=True,
include_interactions=True
)
Method
SCE operates in four stages:
- Hierarchy identification: Categorical columns define grouping levels (e.g.,
city,neighborhood) - Statistical aggregation: For each group, compute configurable statistics (mean, std, quantiles, etc.)
- Cross-fitting: K-fold out-of-fold aggregation prevents target leakage during training
- Feature enrichment: Context features are appended to the original feature matrix
The transformation extends the feature space as follows:
Input: [city, beds, sqft]
Output: [city, beds, sqft, city_mean, city_std, city_count, ...]
During inference, aggregations are computed from the full training set without cross-fitting.
Datasets
Four benchmark datasets are provided for reproducibility:
| Dataset | Domain | Samples | Hier. Cols | Base Feats | +SCE Feats |
|---|---|---|---|---|---|
rental_poland_short |
Short-term rentals (Airbnb) | 1,185 | 4 | 9 | 504 |
rental_poland_long |
Long-term rentals (Otodom) | 1,016 | 5 | 4 | 724 |
rental_uae_contracts |
Dubai rental contracts | 50,000 | 6 | 8 | 323 |
sales_uae_transactions |
Dubai property sales | 50,000 | 7 | 10 | 370 |
from sce.io import load_dataset
df = load_dataset("rental_poland_short") # Bundled with package
df = load_dataset("rental_uae_contracts") # Downloaded on first use
Experimental Results
Experiments compare baseline models (XGBoost with original features) against SCE-enriched models using comprehensive hyperparameter search.
| Dataset | Baseline RMSE | +SCE RMSE | RMSE Reduction | R² Change |
|---|---|---|---|---|
| Airbnb Poland | 27,368 | 22,541 | 17.6% | +24.49 pp |
| Dubai Rentals | 465,037 | 360,267 | 22.5% | +3.83 pp |
| Dubai Transactions | 32,489,660 | 26,353,228 | 18.9% | +25.83 pp |
| Otodom Poland | 4,581 | 4,541 | 0.9% | +1.55 pp |
Results vary by dataset characteristics. The method shows larger improvements on datasets with informative categorical hierarchies and sufficient group sizes. See the results/ directory for detailed experimental outputs.
Reproducing Experiments
# Run experiments on all datasets
python scripts/run.py --all --search
# Run on a specific dataset
python scripts/run.py --dataset rental_uae_contracts --search
# Generate figures and tables
python scripts/generate_figures.py
python scripts/generate_paper_appendix_figures.py
API Reference
ContextConfig
| Parameter | Type | Default | Description |
|---|---|---|---|
target_col |
str |
Required | Name of the target variable column |
categorical_cols |
list[str] |
None |
Columns to aggregate over (auto-detected if None) |
aggregations |
list[AggregationMethod] |
[MEAN, STD] |
Statistics to compute |
use_cross_fitting |
bool |
True |
Enable out-of-fold aggregation |
n_folds |
int |
5 |
Number of cross-fitting folds |
min_group_size |
int |
5 |
Minimum samples required per group |
include_global_stats |
bool |
True |
Include dataset-wide statistics |
include_interactions |
bool |
False |
Include cross-column hierarchies |
StatisticalContextEngine
class StatisticalContextEngine(BaseEstimator, TransformerMixin):
def __init__(self, config: ContextConfig): ...
def fit(self, X: pd.DataFrame, y: pd.Series) -> Self: ...
def transform(self, X: pd.DataFrame) -> pd.DataFrame: ...
def fit_transform(self, X: pd.DataFrame, y: pd.Series) -> pd.DataFrame: ...
Package Structure
sce/
├── engine.py # StatisticalContextEngine transformer
├── config.py # ContextConfig and AggregationMethod
├── stats.py # Aggregation functions
├── pipeline.py # Pipeline utilities
├── cleanup.py # NaN handling and feature cleanup
├── importance.py # Feature importance analysis
└── io/ # Dataset loading
License
This software is released under the Creative Commons Attribution-NonCommercial 4.0 International License.
- Academic and research use: Permitted with attribution
- Commercial use: Requires a separate license agreement
Citation
If you use this software in your research, please cite:
@inproceedings{sce2026,
title={Statistical Context Engineering: Hierarchical Feature Enrichment for Regression},
author={[Authors]},
booktitle={Proceedings of the International Conference on Machine Learning},
year={2026}
}
Contributing
Contributions are welcome. Please see CONTRIBUTING.md for guidelines on code style, testing, and pull request procedures.
Links
Project details
Download files
Download the file for your platform. If you're not sure which to choose, learn more about installing packages.
Source Distribution
Built Distribution
Filter files by name, interpreter, ABI, and platform.
If you're not sure about the file name format, learn more about wheel file names.
Copy a direct link to the current filters
File details
Details for the file stat_context-0.3.2.tar.gz.
File metadata
- Download URL: stat_context-0.3.2.tar.gz
- Upload date:
- Size: 50.9 kB
- Tags: Source
- Uploaded using Trusted Publishing? Yes
- Uploaded via: twine/6.1.0 CPython/3.13.7
File hashes
| Algorithm | Hash digest | |
|---|---|---|
| SHA256 |
34b1da45a7f4f0f416dabf8962f1aa42f2aabce4b5b7dcbc69ff25aa12c83380
|
|
| MD5 |
6124a74201330e25d11e7f3c5cd8f3f9
|
|
| BLAKE2b-256 |
01afc2a16ede88f92ff49b59914e38f6a4325af6b78f99f86669bd5575f125eb
|
Provenance
The following attestation bundles were made for stat_context-0.3.2.tar.gz:
Publisher:
release.yml on joint-hubs/sce
-
Statement:
-
Statement type:
https://in-toto.io/Statement/v1 -
Predicate type:
https://docs.pypi.org/attestations/publish/v1 -
Subject name:
stat_context-0.3.2.tar.gz -
Subject digest:
34b1da45a7f4f0f416dabf8962f1aa42f2aabce4b5b7dcbc69ff25aa12c83380 - Sigstore transparency entry: 845197569
- Sigstore integration time:
-
Permalink:
joint-hubs/sce@97d2910300d7c6ff4ee82f08b27346ac834c44de -
Branch / Tag:
refs/tags/v0.3.2 - Owner: https://github.com/joint-hubs
-
Access:
public
-
Token Issuer:
https://token.actions.githubusercontent.com -
Runner Environment:
github-hosted -
Publication workflow:
release.yml@97d2910300d7c6ff4ee82f08b27346ac834c44de -
Trigger Event:
push
-
Statement type:
File details
Details for the file stat_context-0.3.2-py3-none-any.whl.
File metadata
- Download URL: stat_context-0.3.2-py3-none-any.whl
- Upload date:
- Size: 39.3 kB
- Tags: Python 3
- Uploaded using Trusted Publishing? Yes
- Uploaded via: twine/6.1.0 CPython/3.13.7
File hashes
| Algorithm | Hash digest | |
|---|---|---|
| SHA256 |
5c2fb2875471f98786140c156def35b1f1c011e538eab9b923c72132b737acbb
|
|
| MD5 |
66e2b31c8eef70d3640aa85594171bb9
|
|
| BLAKE2b-256 |
bb53a0aa6c149d5ec62b34a4ed697989bf8d84920e675f562596069e81a0e8e4
|
Provenance
The following attestation bundles were made for stat_context-0.3.2-py3-none-any.whl:
Publisher:
release.yml on joint-hubs/sce
-
Statement:
-
Statement type:
https://in-toto.io/Statement/v1 -
Predicate type:
https://docs.pypi.org/attestations/publish/v1 -
Subject name:
stat_context-0.3.2-py3-none-any.whl -
Subject digest:
5c2fb2875471f98786140c156def35b1f1c011e538eab9b923c72132b737acbb - Sigstore transparency entry: 845197573
- Sigstore integration time:
-
Permalink:
joint-hubs/sce@97d2910300d7c6ff4ee82f08b27346ac834c44de -
Branch / Tag:
refs/tags/v0.3.2 - Owner: https://github.com/joint-hubs
-
Access:
public
-
Token Issuer:
https://token.actions.githubusercontent.com -
Runner Environment:
github-hosted -
Publication workflow:
release.yml@97d2910300d7c6ff4ee82f08b27346ac834c44de -
Trigger Event:
push
-
Statement type: