diff --git a/CHANGELOG.md b/CHANGELOG.md index 65f5bf4..b17f0c7 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -2,6 +2,17 @@ ## Unreleased +### New features + +- **Offsets of multi-response models**: in `Outcome`, a dict gives one + column for each response, keyed by the responses: the columns of a + list of columns (now a response of `Outcome` too), or the classes of a + `Multinomial`, whatever its output. + `Outcome(['y1', 'y2'], offset={'y1': 'O1', 'y2': 'O2'})` gives the + columns `y1`, `y2`, `offset[y1]` and `offset[y2]`, in the order of the + responses; a missing or unknown key raises. A single column, + `offset='O'`, is still one column, the same for every response. + ## 0.3.0 (2026-10-26) ### New features diff --git a/docs/responses.md b/docs/responses.md index 49fee65..6d0c440 100644 --- a/docs/responses.md +++ b/docs/responses.md @@ -202,6 +202,19 @@ y = outcome.fit_transform(BrainCancer) y.iloc[:5] ``` +A multi-response model, with a linear predictor for each response, +takes an offset for each: a dict, keyed by the responses (the columns +of a list of columns, or the classes of a `Multinomial`, whatever its +output), gives the columns `offset[...]` in the order of the responses. +A single column is the same offset for every response. + +```{code-cell} ipython3 +Carseats['O'] = 0.5 +outcome = Outcome(Multinomial('ShelveLoc'), + offset={'Bad': 'O', 'Good': 'Income', 'Medium': 'O'}) +outcome.fit_transform(Carseats).iloc[:3] +``` + 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` diff --git a/modelmatrix/outcome.py b/modelmatrix/outcome.py index 18881b1..c95628d 100644 --- a/modelmatrix/outcome.py +++ b/modelmatrix/outcome.py @@ -15,7 +15,7 @@ from ._native import _from_native from .model_spec import ModelSpec, Feature, _term_columns -from .response import Response, Binomial, _values, _output +from .response import Response, Binomial, Multinomial, _values, _output class _Column(Response): @@ -36,6 +36,38 @@ def __repr__(self): return repr(self.column) +class _Columns(Response): + + # columns of X as they are: the responses of a multi-response model + + def __init__(self, columns): + self.column_list = list(columns) + + @property + def columns(self): + return tuple(self.column_list) + + def _transform(self, X): + return pd.DataFrame({c: _values(X, c).to_numpy() for c in self.column_list}, + index=X.index if isinstance(X, pd.DataFrame) else None) + + def __repr__(self): + return repr(self.column_list) + + +def _components(response): + # the responses of a fitted multi-response model, each with its own + # linear predictor: the columns of a list of columns, the classes of + # a Multinomial (its levels, whatever its output), else None + if isinstance(response, _Columns): + return list(response.column_list) + if isinstance(response, Multinomial): + if isinstance(response.classes, list): + return list(response.classes) + return list(response.levels_) + return None + + class Outcome(Response): """ @@ -44,17 +76,26 @@ class Outcome(Response): Parameters ---------- - response : Response or column identifier + response : Response, column identifier, or list of them The response: a :py:class:`Surv`, :py:class:`Binomial` or - :py:class:`Multinomial`, or a column of `X` as it is. + :py:class:`Multinomial`, or a column of `X` as it is, or a list + of columns, the responses of a multi-response model. - **columns : column identifier or Feature + **columns : column identifier, Feature, or dict of them 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`. + For a multi-response model, a dict gives one column for each + response, e.g. an offset of each linear predictor: its keys are + the response's columns, or the classes of a + :py:class:`Multinomial` (whatever its output), and + `offset={'y1': 'O1', 'y2': 'O2'}` gives the columns + `offset[y1]` and `offset[y2]`, in the order of the responses. + A single column, `offset='O'`, is the same for every response. + Notes ----- @@ -77,21 +118,41 @@ def __init__(self, response, **columns): 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) + values = value.values() if isinstance(value, dict) else [value] + for v in values: + if isinstance(v, Feature): + used = _term_columns([v]) + 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(v) 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_ = {} + # the other columns, by their names: a dict gives one for each + # response, in their order + self.entries_ = {} for name, value in self.extra.items(): + if isinstance(value, dict): + components = _components(self.response_) + if components is None: + raise ValueError(f'{name!r} is a dict, one column for each response, ' + f'but {self.response!r} is not a multi-response') + missing = [c for c in components if c not in value] + unknown = [k for k in value if k not in components] + if missing or unknown: + raise ValueError(f'the keys of {name!r} should be the responses {components}: ' + f'missing {missing}, unknown {unknown}') + for c in components: + self.entries_[f'{name}[{c}]'] = value[c] + else: + self.entries_[name] = value + self.specs_ = {} + for name, value in self.entries_.items(): if isinstance(value, Feature): spec = ModelSpec([value]).fit(X) if len(spec.get_feature_names_out()) != 1: @@ -144,7 +205,7 @@ def _build(self, X, response): weight = y.pop(self.response_.trials) parts.append(y) extra = {} - for name, value in self.extra.items(): + for name, value in self.entries_.items(): if name in self.specs_: D = self.specs_[name].transform(X) extra[name] = np.asarray(D, dtype=float)[:, 0] @@ -162,12 +223,17 @@ def _build(self, X, response): 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()] + def show(v): + if isinstance(v, dict): + return '{' + ', '.join(f'{k!r}: {show(u)}' for k, u in v.items()) + '}' + return v.name if isinstance(v, Feature) else repr(v) + args = [repr(self.response)] + [f'{k}={show(v)}' for k, v in self.extra.items()] return f'Outcome({", ".join(args)})' def _as_response(response): if isinstance(response, Response): return response + if isinstance(response, list): + return _Columns(response) return _Column(response) diff --git a/tests/test_outcome.py b/tests/test_outcome.py index 63dce15..8f2a6ac 100644 --- a/tests/test_outcome.py +++ b/tests/test_outcome.py @@ -142,3 +142,71 @@ def test_native(D, backend): assert y_nw['strata'].to_list() == list(D['inst']) if backend == 'polars': assert y['strata'].dtype == lib.Enum(['a', 'b', 'c']) + + +# multi-response models: one column for each response + +@pytest.fixture +def M(): + rng = np.random.default_rng(3) + n = 30 + return pd.DataFrame({'y1': rng.standard_normal(n), 'y2': rng.standard_normal(n), + 'O1': rng.standard_normal(n), 'O2': rng.standard_normal(n), + 'O': rng.standard_normal(n), 'exposure': rng.uniform(1, 5, n), + 'kind': pd.Categorical(rng.choice(['lo', 'mid', 'hi'], n), + categories=['lo', 'mid', 'hi']), + 'x': rng.standard_normal(n)}) + + +def test_offset_for_each_response(M): + # in the order of the responses, whatever the order of the dict + Y = Outcome(['y1', 'y2'], offset={'y2': 'O2', 'y1': log_exposure}) + y = Y.fit_transform(M) + assert list(y.columns) == ['y1', 'y2', 'offset[y1]', 'offset[y2]'] + np.testing.assert_allclose(y[['y1', 'y2']], M[['y1', 'y2']]) + np.testing.assert_allclose(y['offset[y1]'], np.log(M['exposure'])) + np.testing.assert_allclose(y['offset[y2]'], M['O2']) + assert Y.columns == ('y1', 'y2', 'O2', 'exposure') + assert MS([wildcard(exclude=Y)]).fit(M).expanded_terms_ == ['O1', 'O', 'kind', 'x'] + assert repr(Y) == "Outcome(['y1', 'y2'], offset={'y2': 'O2', 'y1': log(exposure)})" + # new data, without the responses + O = Y.transform(M.drop(columns=['y1', 'y2']), response=False) + assert list(O.columns) == ['offset[y1]', 'offset[y2]'] + + +def test_one_offset_for_every_response(M): + y = Outcome(['y1', 'y2'], offset='O').fit_transform(M) + assert list(y.columns) == ['y1', 'y2', 'offset'] + + +@pytest.mark.parametrize('output', ['labels', 'onehot']) +def test_offset_for_each_class(M, output): + # keyed by the classes, whatever the output of the Multinomial + Y = Outcome(Multinomial('kind', output=output), + offset={'hi': 'O', 'lo': 'O1', 'mid': 'O2'}) + y = Y.fit_transform(M) + offsets = ['offset[lo]', 'offset[mid]', 'offset[hi]'] + assert list(y.columns)[-3:] == offsets + np.testing.assert_allclose(y[offsets], M[['O1', 'O2', 'O']]) + # the classes of counts are their columns + C = M.assign(n_lo=1, n_hi=2) + y = Outcome(Multinomial(['n_lo', 'n_hi']), offset={'n_lo': 'O1', 'n_hi': 'O2'}).fit_transform(C) + assert list(y.columns) == ['n_lo', 'n_hi', 'offset[n_lo]', 'offset[n_hi]'] + + +def test_offset_for_each_response_errors(M): + with pytest.raises(ValueError, match=r"missing \['y2'\], unknown \['y3'\]"): + Outcome(['y1', 'y2'], offset={'y1': 'O1', 'y3': 'O2'}).fit(M) + with pytest.raises(ValueError, match="missing \\['hi'\\]"): + Outcome(Multinomial('kind'), offset={'lo': 'O1', 'mid': 'O2'}).fit(M) + with pytest.raises(ValueError, match='not a multi-response'): + Outcome('y1', offset={'y1': 'O1'}).fit(M) + with pytest.raises(ValueError, match='single column'): + Outcome(['y1', 'y2'], offset={'y1': poly('x', 2), 'y2': 'O2'}).fit(M) + + +def test_offset_for_each_response_polars(M): + pl = pytest.importorskip('polars') + y = Outcome(['y1', 'y2'], offset={'y1': 'O1', 'y2': 'O2'}).fit_transform(pl.from_pandas(M)) + assert isinstance(y, pl.DataFrame) + assert y.columns == ['y1', 'y2', 'offset[y1]', 'offset[y2]']