Outcome: an offset for each response of a multi-response model

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

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`

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)

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]']