From da75587d910dc20cfc77093ebdf7e1d595ebcb79 Mon Sep 17 00:00:00 2001 From: Jonathan Taylor Date: Mon, 26 Oct 2026 15:58:57 -0700 Subject: [PATCH] Outcome: a response with an offset, weights, strata, named by role - Outcome(response, **columns): each keyword names an output column, taken from a column of X as it is (a categorical keeps its labels) or from a single-column feature fit as a term, e.g. log(exposure). The response's columns come first. - wildcard(exclude=outcome) leaves out every column it uses. - transform(X, response=False) builds only the other columns, e.g. the offset of new data. - With Binomial(output='proportion') the trials are the column weight, times any weights given. - Surv names its columns by role: start, stop, status. - Dataframes of other backends keep categorical columns (polars Enum). Co-Authored-By: Claude Opus 5.5 --- CHANGELOG.md | 21 +++++ docs/api/responses.rst | 2 + docs/responses.md | 48 ++++++++++- modelmatrix/__init__.py | 3 +- modelmatrix/_native.py | 10 ++- modelmatrix/outcome.py | 173 ++++++++++++++++++++++++++++++++++++++++ modelmatrix/response.py | 23 +++--- tests/test_outcome.py | 144 +++++++++++++++++++++++++++++++++ tests/test_response.py | 10 ++- 9 files changed, 413 insertions(+), 21 deletions(-) create mode 100644 modelmatrix/outcome.py create mode 100644 tests/test_outcome.py diff --git a/CHANGELOG.md b/CHANGELOG.md index 3d78833..706f038 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -2,6 +2,27 @@ ## Unreleased +### New features + +- **`Outcome`**: a response with the other columns a model takes with + it, each named by its keyword, so an estimator can find them by name: + `Outcome(Surv('time', 'status'), offset='O', weight='w', strata='inst')` + gives the columns `stop`, `status`, `offset`, `weight` and `strata`. A + column is taken as it is (a categorical keeps its labels); a feature of + a single column is fit as a term, e.g. `offset(log(exposure))`. + `wildcard(exclude=outcome)` leaves out all of its columns, and + `transform(X, response=False)` builds the offset of new data. With + `Binomial(..., output='proportion')`, the trials are the column + `weight`, times any weights given. +- Dataframes of other backends keep a categorical column (as a polars + `Enum`). + +### Breaking changes + +- `Surv` names its columns by their role, `start`, `stop` and `status` + (as R's and lifelines'), not as in the data: `Surv('time', 'status')` + gives the columns `stop` and `status`. + ## 0.2.0 (2026-10-26) ### New features diff --git a/docs/api/responses.rst b/docs/api/responses.rst index 5a42aff..8173ca0 100644 --- a/docs/api/responses.rst +++ b/docs/api/responses.rst @@ -7,5 +7,7 @@ Responses .. autoclass:: modelmatrix.Surv .. autoclass:: modelmatrix.Binomial .. autoclass:: modelmatrix.Multinomial +.. autoclass:: modelmatrix.Outcome + :members: transform .. autoclass:: modelmatrix.Response :members: fit, transform, fit_transform diff --git a/docs/responses.md b/docs/responses.md index 5331398..49fee65 100644 --- a/docs/responses.md +++ b/docs/responses.md @@ -112,9 +112,10 @@ design. ### Survival responses: `Surv` `Surv(time, status)` is R's survival response, and `Surv(start, stop, -status)` that of (start, stop] intervals. The response is a -`pd.DataFrame` with these columns, named as in the data, as glmstar's -`CoxFamily(event_id=..., status_id=..., start_id=...)` or lifelines take it. +status)` that of (start, stop] intervals. The response is a dataframe +whose columns are named by their role, whatever their names in the +data: `stop` (the times), `status` and, if given, `start`, as R's +`Surv` and lifelines' `CoxTimeVaryingFitter` name them. ```{code-cell} ipython3 BrainCancer = load_data('BrainCancer').dropna() @@ -125,7 +126,7 @@ y.iloc[:5] ``` ```{code-cell} ipython3 -sm.PHReg(y['time'], X, status=y['status']).fit().params +sm.PHReg(y['stop'], X, status=y['status']).fit().params ``` ### Binomial responses @@ -167,6 +168,45 @@ of counts for each class. Multinomial('ShelveLoc', output='onehot').fit_transform(Carseats).iloc[:5] ``` +### Offsets, weights and strata: `Outcome` + +`Outcome` puts the other columns a model takes with its response next +to it, each named by its keyword: an offset, weights, the strata of a +Cox model, or any other. A column is taken as it is (a categorical keeps +its labels), and a feature of a single column is fit as a term of a +`ModelSpec`, e.g. R's `offset(log(exposure))`: + +```{code-cell} ipython3 +from sklearn.preprocessing import FunctionTransformer +from modelmatrix import Outcome, derived_feature + +rng = np.random.default_rng(1) +P = pd.DataFrame({'x': rng.standard_normal(100), 'exposure': rng.uniform(1, 5, 100)}) +P['claims'] = rng.poisson(P['exposure'] * np.exp(0.3 * P['x'])) + +log_exposure = derived_feature(['exposure'], encoder=FunctionTransformer(np.log), + name='log(exposure)') +outcome = Outcome('claims', offset=log_exposure) +y = outcome.fit_transform(P) +X = MS([wildcard(exclude=outcome)], intercept=True).fit_transform(P) # leaves out exposure too +sm.GLM(y['claims'], X, family=sm.families.Poisson(), offset=y['offset']).fit().params +``` + +As the columns are named by their role, an estimator can find them by +name, e.g. a stratified Cox model with an offset: + +```{code-cell} ipython3 +BrainCancer['O'] = 0.1 +outcome = Outcome(Surv('time', 'status'), offset='O', strata='sex') +y = outcome.fit_transform(BrainCancer) +y.iloc[:5] +``` + +The offset of new data, whose response is not known, is +`outcome.transform(new, response=False)`. With a `Binomial` of +`output='proportion'`, the number of trials is the column `weight` +(times the weights given as `weight=`), as R's weights of a proportion. + ## Missing values As scikit-learn's transformers, a `ModelSpec` keeps every row: a row diff --git a/modelmatrix/__init__.py b/modelmatrix/__init__.py index c45089b..15032f9 100644 --- a/modelmatrix/__init__.py +++ b/modelmatrix/__init__.py @@ -18,6 +18,7 @@ Wildcard) from .columns import Column from .response import Response, Surv, Binomial, Multinomial +from .outcome import Outcome try: from ._version import version as __version__ @@ -27,4 +28,4 @@ __all__ = ['ModelSpec', 'Feature', 'Column', 'Contrast', 'Ordinal', 'poly', 'ns', 'bs', 'derived_feature', 'pca', 'contrast', 'ordinal', 'build_columns', 'wildcard', 'Wildcard', - 'Response', 'Surv', 'Binomial', 'Multinomial'] + 'Response', 'Surv', 'Binomial', 'Multinomial', 'Outcome'] diff --git a/modelmatrix/_native.py b/modelmatrix/_native.py index 0d1231e..b5bccfc 100644 --- a/modelmatrix/_native.py +++ b/modelmatrix/_native.py @@ -12,6 +12,9 @@ def _series_to_native(S, backend): + return _series_to_nw(S, backend).to_native() + +def _series_to_nw(S, backend): # a categorical is an Enum, if the backend has one, of its # categories as strings values, dtype = S.to_numpy(), None @@ -19,9 +22,9 @@ def _series_to_native(S, backend): values = np.asarray([None if pd.isna(v) else str(v) for v in S], dtype=object) dtype = nw.Enum([str(c) for c in S.cat.categories]) try: - return nw.new_series(str(S.name), values, dtype=dtype, backend=backend).to_native() + return nw.new_series(str(S.name), values, dtype=dtype, backend=backend) except NotImplementedError: # e.g. no Enum in pyarrow - return nw.new_series(str(S.name), values, backend=backend).to_native() + return nw.new_series(str(S.name), values, backend=backend) def _from_native(X, columns=None): """ @@ -62,6 +65,9 @@ def to_native(D): if duplicated: raise ValueError(f'the model matrix has duplicated column names {duplicated}, ' 'which this dataframe backend cannot hold') + if any(isinstance(t, pd.CategoricalDtype) for t in D.dtypes): # e.g. strata + return nw.from_dict({str(c): _series_to_nw(D[c].rename(str(c)), backend) + for c in D.columns}).to_native() return nw.from_dict({str(c): D[c].to_numpy() for c in D.columns}, backend=backend).to_native() diff --git a/modelmatrix/outcome.py b/modelmatrix/outcome.py new file mode 100644 index 0000000..18881b1 --- /dev/null +++ b/modelmatrix/outcome.py @@ -0,0 +1,173 @@ +""" +Outcome +======= + +A response with other columns a model takes with it, such as an offset, +weights or the strata of a Cox model, each named by its role: +`Outcome(Surv('time', 'status'), offset='O', weight='w')` builds a +dataframe with columns `stop`, `status`, `offset` and `weight`, so that +an estimator can find each column by its name. +""" + +from copy import deepcopy + +import numpy as np, pandas as pd + +from ._native import _from_native +from .model_spec import ModelSpec, Feature, _term_columns +from .response import Response, Binomial, _values, _output + + +class _Column(Response): + + # a column of X as it is + + def __init__(self, column): + self.column = column + + @property + def columns(self): + return (self.column,) + + def _transform(self, X): + return _values(X, self.column).copy() + + def __repr__(self): + return repr(self.column) + + +class Outcome(Response): + + """ + A response, with other columns named by their role. + + Parameters + ---------- + + response : Response or column identifier + The response: a :py:class:`Surv`, :py:class:`Binomial` or + :py:class:`Multinomial`, or a column of `X` as it is. + + **columns : column identifier or Feature + Other columns, each named by its keyword, e.g. `offset='O'` or + `weight='w'`: a column of `X`, as it is (a categorical, such as + the strata of a Cox model, keeps its labels), or a feature + of a single column, e.g. the log of a column, fit as a term of a + :py:class:`ModelSpec`. + + Notes + ----- + + The columns are those of the response, then the others in the order + of their keywords. With a :py:class:`Binomial` of + `output='proportion'`, its number of trials is the column `weight` + (times the weights given as `weight=`, if any), as R's prior weights + of a proportion. + + `wildcard(exclude=outcome)` leaves out every column the outcome uses, + so that none of them is a term of its own design. For an offset of + new data, without its response, use `transform(X, response=False)`. + """ + + def __init__(self, response, **columns): + self.response = response + self.extra = columns + + @property + def columns(self): + cols = list(_as_response(self.response).columns) + for name, value in self.extra.items(): + if isinstance(value, Feature): + used = _term_columns([value]) + if used is None: + raise ValueError(f'the column {name!r} of an Outcome cannot use a wildcard') + cols.extend(sorted(used, key=str)) + else: + cols.append(value) + return tuple(dict.fromkeys(cols)) + + def _fit(self, X): + self.response_ = deepcopy(_as_response(self.response)) + self.response_._fit(X) + self.response_.columns_ = self.response_.columns + self.specs_ = {} + for name, value in self.extra.items(): + if isinstance(value, Feature): + spec = ModelSpec([value]).fit(X) + if len(spec.get_feature_names_out()) != 1: + raise ValueError(f'the column {name!r} of an Outcome should be a single ' + f'column: {value.name!r} has ' + f'{len(spec.get_feature_names_out())}') + self.specs_[name] = spec + else: + _values(X, value) # raises if the column is not in X + self._build(X, response=True) # checks the names + + def transform(self, X, response=True): + """ + Build the outcome from `X`, after fitting. + + Parameters + ---------- + X : dataframe or np.ndarray + Data holding the columns of the outcome. + + response : bool, default=True + Build the response? If False, only the other columns, e.g. + the offset of new data, whose response is not known. + + Returns + ------- + y : dataframe or np.ndarray + Of the backend of *X*, an :py:class:`np.ndarray` for an array. + """ + self._check_fitted() + X, to_native = _from_native(X, set(self.columns)) + y = self._build(X, response) + return y if to_native is None else to_native(y) + + def _transform(self, X): + return self._build(X, response=True) + + def _build(self, X, response): + X_orig = X + if not isinstance(X, pd.DataFrame): # an array: its columns are integers + X = pd.DataFrame(np.asarray(X)) + parts = [] + weight = None + if response: + y = self.response_._transform(X) + y = y.to_frame() if isinstance(y, pd.Series) else y + if (isinstance(self.response_, Binomial) and self.response_.trials is not None + and self.response_.output == 'proportion'): + # the trials are the weights of a proportion + weight = y.pop(self.response_.trials) + parts.append(y) + extra = {} + for name, value in self.extra.items(): + if name in self.specs_: + D = self.specs_[name].transform(X) + extra[name] = np.asarray(D, dtype=float)[:, 0] + else: + extra[name] = _values(X, value).copy() # keeps a categorical + if weight is not None: + w = weight.to_numpy() + extra['weight'] = w if 'weight' not in extra else w * np.asarray(extra['weight'], float) + parts.append(pd.DataFrame(extra, index=X.index)) + D = pd.concat(parts, axis=1) + duplicated = list(D.columns[D.columns.duplicated()]) + if duplicated: + raise ValueError(f'the columns of the Outcome have the same names {duplicated}: ' + 'a column of the response has the name of another') + return _output(D, X_orig) + + def __repr__(self): + args = [repr(self.response)] + [f'{k}={v.name if isinstance(v, Feature) else repr(v)}' + for k, v in self.extra.items()] + return f'Outcome({", ".join(args)})' + + +def _as_response(response): + if isinstance(response, Response): + return response + return _Column(response) diff --git a/modelmatrix/response.py b/modelmatrix/response.py index 29e143e..560c3a5 100644 --- a/modelmatrix/response.py +++ b/modelmatrix/response.py @@ -71,9 +71,7 @@ def transform(self, X): :py:class:`pd.DataFrame`), an :py:class:`np.ndarray` for an array. """ - if not hasattr(self, 'columns_'): - raise NotFittedError(f'this {type(self).__name__} is not fitted yet: ' - 'call fit (or fit_transform) first') + self._check_fitted() X, to_native = _from_native(X, set(self.columns)) y = self._transform(X) return y if to_native is None else to_native(y) @@ -84,6 +82,11 @@ def fit_transform(self, X, y=None): """ return self.fit(X).transform(X) + def _check_fitted(self): + if not hasattr(self, 'columns_'): + raise NotFittedError(f'this {type(self).__name__} is not fitted yet: ' + 'call fit (or fit_transform) first') + def _fit(self, X): for col in self.columns: _values(X, col) # raises if the column is not in X @@ -147,10 +150,10 @@ class Surv(Response): Notes ----- - The response is a :py:class:`pd.DataFrame` with columns `start` - (if given), `stop` and `status`, named as in `X`, e.g. for glmstar's - `CoxFamily(event_id=stop, status_id=status, start_id=start)` or for - lifelines' `CoxPHFitter`. + The response is a dataframe with columns named by their role, + whatever their names in `X`: `start` (if given), `stop` and `status`, + the names of R's `Surv` and of lifelines' `CoxTimeVaryingFitter`. + (scikit-survival's `Surv` is another object, a structured array.) """ def __init__(self, *args, start=None, stop=None, status=None): @@ -190,9 +193,9 @@ def _transform(self, X): if (start >= stop).any(): raise ValueError(f'the start {self.start!r} of Surv should be less than ' f'its stop {self.stop!r}') - D[self.start] = start.to_numpy() - D[self.stop] = stop.to_numpy() - D[self.status] = status.to_numpy() + D['start'] = start.to_numpy() + D['stop'] = stop.to_numpy() + D['status'] = status.to_numpy() index = X.index if isinstance(X, pd.DataFrame) else None return _output(pd.DataFrame(D, index=index), X) diff --git a/tests/test_outcome.py b/tests/test_outcome.py new file mode 100644 index 0000000..63dce15 --- /dev/null +++ b/tests/test_outcome.py @@ -0,0 +1,144 @@ +# Outcome: a response with an offset, weights, strata, ..., named by role + +import numpy as np, pandas as pd +import pytest +import statsmodels.api as sm +from sklearn.exceptions import NotFittedError +from sklearn.preprocessing import FunctionTransformer + +from modelmatrix import (ModelSpec as MS, Outcome, Surv, Binomial, Multinomial, + wildcard, derived_feature, poly) + + +@pytest.fixture +def D(): + rng = np.random.default_rng(2) + n = 60 + exposure = rng.uniform(1, 5, n) + x = rng.standard_normal(n) + return pd.DataFrame({'time': rng.exponential(size=n) + 0.1, + 'status': rng.binomial(1, 0.7, n), + 'O': rng.standard_normal(n) / 4, + 'w': rng.uniform(0.5, 2, n), + 'inst': pd.Categorical(rng.choice(['a', 'b', 'c'], n)), + 'exposure': exposure, + 'x': x, + 'claims': rng.poisson(exposure * np.exp(0.3 * x)), + 'cases': rng.binomial(8, 0.4, n), + 'size': np.full(n, 8)}) + + +log_exposure = derived_feature(['exposure'], encoder=FunctionTransformer(np.log), + name='log(exposure)') + + +def test_columns_named_by_role(D): + Y = Outcome(Surv('time', 'status'), offset='O', weight='w', strata='inst') + y = Y.fit_transform(D) + assert list(y.columns) == ['stop', 'status', 'offset', 'weight', 'strata'] + np.testing.assert_allclose(y['offset'], D['O']) + np.testing.assert_allclose(y['weight'], D['w']) + assert isinstance(y['strata'].dtype, pd.CategoricalDtype) # labels, not contrasts + assert y.index.equals(D.index) + assert repr(Y) == "Outcome(Surv('time', 'status'), offset='O', weight='w', strata='inst')" + + +def test_wildcard_excludes_every_column(D): + Y = Outcome(Surv('time', 'status'), offset='O', weight='w', strata='inst') + assert Y.columns == ('time', 'status', 'O', 'w', 'inst') + spec = MS([wildcard(exclude=Y)]).fit(D) + assert spec.expanded_terms_ == ['exposure', 'x', 'claims', 'cases', 'size'] + Z = Outcome('claims', offset=log_exposure) + assert Z.columns == ('claims', 'exposure') + assert MS([wildcard(exclude=Z)]).fit(D).expanded_terms_ == [ + 'time', 'status', 'O', 'w', 'inst', 'x', 'cases', 'size'] + + +def test_poisson_offset_matches_statsmodels(D): + # R's glm(claims ~ x + offset(log(exposure)), family=poisson) + Y = Outcome('claims', offset=log_exposure) + y = Y.fit_transform(D) + assert list(y.columns) == ['claims', 'offset'] + X = MS(['x'], intercept=True).fit_transform(D) + family = sm.families.Poisson() + fit = sm.GLM(y['claims'], X, family=family, offset=y['offset']).fit() + ref = sm.GLM(D['claims'], X, family=family, offset=np.log(D['exposure'])).fit() + np.testing.assert_allclose(fit.params, ref.params) + + +def test_offset_of_new_data(D): + Y = Outcome('claims', offset=log_exposure, weight='w').fit(D) + new = D.drop(columns='claims').iloc[:5] + O = Y.transform(new, response=False) + assert list(O.columns) == ['offset', 'weight'] + np.testing.assert_allclose(O['offset'], np.log(new['exposure'])) + with pytest.raises(ValueError, match="'claims' is not in X"): + Y.transform(new) + + +def test_stratified_cox_matches_statsmodels(D): + Y = Outcome(Surv('time', 'status'), offset='O', strata='inst') + y = Y.fit_transform(D) + X = MS(['x']).fit_transform(D) + fit = sm.PHReg(y['stop'], X, status=y['status'], strata=y['strata'], + offset=y['offset']).fit() + ref = sm.PHReg(D['time'], D[['x']], status=D['status'], strata=D['inst'], + offset=D['O']).fit() + np.testing.assert_allclose(fit.params, ref.params) + + +def test_binomial_proportion_weights(D): + # the trials are the weights of a proportion, times the prior weights + y = Outcome(Binomial('cases', trials='size', output='proportion')).fit_transform(D) + assert list(y.columns) == ['cases/size', 'weight'] + np.testing.assert_allclose(y['weight'], D['size']) + y = Outcome(Binomial('cases', trials='size', output='proportion'), weight='w').fit_transform(D) + assert list(y.columns) == ['cases/size', 'weight'] + np.testing.assert_allclose(y['weight'], D['size'] * D['w']) + # counts are left as they are + y = Outcome(Binomial('cases', trials='size'), weight='w').fit_transform(D) + assert list(y.columns) == ['cases', 'size-cases', 'weight'] + + +def test_other_responses(D): + y = Outcome(Multinomial('inst'), weight='w').fit_transform(D) + assert list(y.columns) == ['inst', 'weight'] + assert isinstance(y['inst'].dtype, pd.CategoricalDtype) + # the response given is not modified by fit + response = Binomial('status') + Outcome(response, weight='w').fit(D) + assert not hasattr(response, 'levels_') + + +def test_errors(D): + with pytest.raises(NotFittedError): + Outcome('claims', offset='O').transform(D) + with pytest.raises(ValueError, match='single column'): + Outcome('claims', offset=poly('x', 2)).fit(D) + with pytest.raises(ValueError, match='same names'): + Outcome('claims', claims='O').fit(D) + with pytest.raises(ValueError, match="'nope' is not in X"): + Outcome('claims', offset='nope').fit(D) + with pytest.raises(ValueError, match='wildcard'): + Outcome('claims', offset=derived_feature(wildcard, name='all')).columns + + +def test_ndarray(D): + A = np.asarray(D[['time', 'status', 'O']], dtype=float) + y = Outcome(Surv(0, 1), offset=2).fit_transform(A) + assert isinstance(y, np.ndarray) + np.testing.assert_allclose(y, A) + + +@pytest.mark.parametrize('backend', ['polars', 'pyarrow']) +def test_native(D, backend): + lib = pytest.importorskip(backend) + nw = pytest.importorskip('narwhals') + P = lib.from_pandas(D) if backend == 'polars' else lib.Table.from_pandas(D) + y = Outcome(Surv('time', 'status'), offset='O', strata='inst').fit_transform(P) + y_nw = nw.from_native(y) + assert nw.get_native_namespace(y_nw) is lib + assert y_nw.columns == ['stop', 'status', 'offset', 'strata'] + assert y_nw['strata'].to_list() == list(D['inst']) + if backend == 'polars': + assert y['strata'].dtype == lib.Enum(['a', 'b', 'c']) diff --git a/tests/test_response.py b/tests/test_response.py index 0ceae5d..18ae3ba 100644 --- a/tests/test_response.py +++ b/tests/test_response.py @@ -68,15 +68,17 @@ def test_index_of_y(Carseats): def test_surv(S): y = Surv('time', 'status').fit_transform(S) - pd.testing.assert_frame_equal(y, S[['time', 'status']].astype(float)) + # named by role, whatever their names in the data + pd.testing.assert_frame_equal(y, S[['time', 'status']].astype(float) + .set_axis(['stop', 'status'], axis=1)) y3 = Surv('start', 'time', 'status').fit_transform(S) - assert list(y3.columns) == ['start', 'time', 'status'] + assert list(y3.columns) == ['start', 'stop', 'status'] y_kw = Surv(stop='time', status='status', start='start').fit_transform(S) pd.testing.assert_frame_equal(y3, y_kw) assert repr(Surv('start', 'time', 'status')) == "Surv('start', 'time', 'status')" # a boolean status yb = Surv('time', 'flag').fit_transform(S) - np.testing.assert_array_equal(yb['flag'], S['flag'].astype(float)) + np.testing.assert_array_equal(yb['status'], S['flag'].astype(float)) def test_surv_matches_statsmodels_PHReg(S): @@ -84,7 +86,7 @@ def test_surv_matches_statsmodels_PHReg(S): y = response.fit_transform(S) X = MS([wildcard(exclude=[response, 'start', 'cases', 'size', 'flag', 'kind'])]).fit_transform(S) assert list(X.columns) == ['x'] - fit = sm.PHReg(y['time'], X, status=y['status']).fit() + fit = sm.PHReg(y['stop'], X, status=y['status']).fit() ref = sm.PHReg(S['time'], S[['x']], status=S['status']).fit() np.testing.assert_allclose(fit.params, ref.params)