metacausal.datasets.load_lalonde

metacausal.datasets.load_lalonde(raw=False, binarize_y=None)[source]

Load the Lalonde job training dataset.

The dataset contains 445 observations with a binary treatment indicator (job training program) and a continuous outcome (earnings in 1978).

Parameters:
  • raw (bool) – If True, return a DataFrame. If False (default), return (X, T, Y) numpy arrays.

  • binarize_y (Literal[None, 'median', 'positive']) – Optional binarization of the 1978 earnings outcome. None (default) keeps it continuous. "median" thresholds at the sample median of re78 (~50/50 split, useful as a balanced binary fixture). "positive" thresholds at re78 > 0 (~69/31 split, the natural “any 1978 earnings” indicator). Ignored when raw=True.

Returns:

If raw=False

tuple of (X, T, Y) where

  • X: covariate matrix, shape (n, 10)

  • T: binary treatment vector, shape (n,)

  • Y: outcome vector, shape (n,) — continuous if binarize_y is None, integer 0/1 otherwise.

If raw=True: DataFrame with all columns.

Return type:

DataFrame | tuple[ndarray, ndarray, ndarray]

Examples

>>> from metacausal.datasets import load_lalonde
>>> X, T, Y = load_lalonde()
>>> X.shape
(445, 10)
>>> sorted(set(T.tolist()))
[0, 1]
>>> X, T, Y_bin = load_lalonde(binarize_y="median")
>>> sorted(set(Y_bin.tolist()))
[0, 1]