Outcome, and Surv's columns named by role, for 0.3.0

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

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

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

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

modelmatrix/_native.py

@@ -12,16 +12,19 @@
 
 
 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
     if isinstance(S.dtype, pd.CategoricalDtype):
         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()
 

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)

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)