Datasets

v0.3.0

Built-in datasets for quick experimentation, plus synthetic data generation utilities. Load classic ML datasets with a single function call or generate custom data with configurable probability distributions.

Available Datasets

StatsLibX ships with six built-in datasets covering classification, regression, finance, and educational domains. All datasets are loaded as pandas.DataFrame (default) or can be returned as (X, y) numpy arrays.

Dataset	Description	Rows	Columns	Load
iris.csv	Fisher's Iris flower measurements (sepal & petal length/width) for three species: setosa, versicolor, virginica	150	5	`load_dataset("iris.csv")`
penguins.csv	Palmer Archipelago penguin morphometrics: bill dimensions, flipper length, body mass for Adelie, Chinstrap, Gentoo	344	7	`load_dataset("penguins.csv")`
titanic.csv	Passenger manifest from the Titanic disaster: survival status, class, age, sex, fare, embarkation port	418	12	`load_dataset("titanic.csv")`
sp500_companies.csv	S&P 500 constituent companies with financial metrics: market cap, EBITDA, revenue growth, sector, industry	503	21	`load_dataset("sp500_companies.csv")`
course_completion.csv	Student course completion records: demographics, engagement metrics, quiz scores, project grades, completion status	100,000	40	`load_dataset("course_completion.csv")`
Cocoa_Bubbles_Investment_Nigeria_Ghana_1980_2023.xlsx	Time-series economic data for cocoa production, pricing, and investment metrics across Nigeria and Ghana (1980–2023)	—	—	`load_dataset("Cocoa_Bubbles_Investment_Nigeria_Ghana_1980_2023.xlsx")`

Loading Datasets

The load_dataset() function is the primary interface for loading built-in data. It supports CSV, Excel, Parquet, and JSON formats.

load_dataset

load_dataset(name, backend='pandas', return_X_y=None, sep=',') -> pd.DataFrame | tuple[np.ndarray, np.ndarray]

pd.DataFrame | tuple[np.ndarray, np.ndarray]

Convenience Functions

Shortcut functions for the two most popular datasets. They wrapload_dataset() with the correct filename.

load_iris

load_iris(backend='pandas', return_X_y=None) -> pd.DataFrame | tuple[np.ndarray, np.ndarray]

pd.DataFrame | tuple[np.ndarray, np.ndarray]

load_penguins

load_penguins(backend='pandas', return_X_y=None) -> pd.DataFrame | tuple[np.ndarray, np.ndarray]

pd.DataFrame | tuple[np.ndarray, np.ndarray]

Synthetic Data Generation

The generate_dataset() function creates synthetic datasets with configurable probability distributions. Define a schema mapping column names to distribution configurations to produce data for testing, simulation, or benchmarking.

generate_dataset

generate_dataset(n_rows, schema, seed=None, save=False, filename=None) -> pd.DataFrame

pd.DataFrame

Schema Reference

Each column in the schema dict supports the following common parameters plus distribution-specific parameters:

Parameter	Required	Type	Description
dist	Yes	str	Distribution name: `normal`, `uniform`, `exponential`, `lognormal`, `poisson`, `binomial`, `categorical`
type	No	'int' \| 'float'	Output data type. Defaults to `float`. Not applicable for `categorical`.
round	No	int	Number of decimal places. Defaults to `2`. Ignored when `type='int'`.

Distribution-Specific Parameters

normal

Param	Default
mean	0
std	1

normal example

{"age": {"dist": "normal", "mean": 30, "std": 5, "type": "int"}}

uniform

Param	Default
low	0
high	1

uniform example

{"temperature": {"dist": "uniform", "low": 15, "high": 35, "type": "float", "round": 1}}

exponential

Param	Default
scale	1

exponential example

{"wait_time": {"dist": "exponential", "scale": 1.5, "type": "float", "round": 3}}

lognormal

Param	Default
mean	0
std	1

lognormal example

{"income": {"dist": "lognormal", "mean": 10, "std": 0.5, "type": "float", "round": 2}}

poisson

Param	Default
lam	1

poisson example

{"events_per_day": {"dist": "poisson", "lam": 3.5, "type": "int"}}

binomial

Param	Default
n	1
p	0.5

binomial example

{"successes": {"dist": "binomial", "n": 10, "p": 0.3, "type": "int"}}

categorical

Param	Required	Description
choices	Yes	List of categories to sample from

categorical example

{"department": {"dist": "categorical", "choices": ["Engineering", "Sales", "HR"]}}

Complete Example

Mixed dataset generation

from statslibx.datasets import generate_dataset

schema = {
    # Integer column: normally distributed ages
    "age": {"dist": "normal", "mean": 40, "std": 12, "type": "int"},

    # Float column: lognormal income distribution (right-skewed)
    "income": {"dist": "lognormal", "mean": 10.3, "std": 0.6, "type": "float", "round": 0},

    # Integer column: Poisson-distributed number of purchases
    "purchases": {"dist": "poisson", "lam": 2, "type": "int"},

    # Categorical column
    "segment": {"dist": "categorical", "choices": ["Premium", "Standard", "Budget"]},

    # Float column: uniform satisfaction score
    "satisfaction": {"dist": "uniform", "low": 0, "high": 10, "type": "float", "round": 1},

    # Integer column: binomial (10 trials, 70% probability)
    "conversions": {"dist": "binomial", "n": 10, "p": 0.7, "type": "int"},
}

df = generate_dataset(n_rows=5000, schema=schema, seed=123)
print(df.shape)  # (5000, 6)
print(df.head())
print(df.describe())

# Save to CSV
df = generate_dataset(n_rows=5000, schema=schema, seed=123, save=True, filename="customer_data")

Utility Functions

_X_y

_X_y(df, X_columns, y_column) -> tuple[np.ndarray, np.ndarray]

tuple[np.ndarray, np.ndarray]