modelmatrix 0.2.0__tar.gz → 0.3.0__tar.gz

{modelmatrix-0.2.0 → modelmatrix-0.3.0}/CHANGELOG.md

@@ -2,6 +2,29 @@
 
 ## Unreleased
 
+## 0.3.0 (2026-10-26)
+
+### 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

{modelmatrix-0.2.0/modelmatrix.egg-info → modelmatrix-0.3.0}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: modelmatrix
-Version: 0.2.0
+Version: 0.3.0
 Summary: Model matrices from tabular data, as scikit-learn transformers
 Author-email: Jonathan Taylor <jonathan.taylor@stanford.edu>
 License-Expression: BSD-3-Clause

{modelmatrix-0.2.0 → modelmatrix-0.3.0}/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

{modelmatrix-0.2.0 → modelmatrix-0.3.0}/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-0.2.0 → modelmatrix-0.3.0}/modelmatrix/__init__.py

@@ -18,6 +18,7 @@ from .model_spec import (ModelSpec,
                          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 @@ except ImportError:  # not installed
 __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-0.2.0 → modelmatrix-0.3.0}/modelmatrix/_native.py

@@ -12,6 +12,9 @@ from scipy import sparse as sp_sparse
 
 
 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 _from_native(X, columns=None):
         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-0.2.0 → modelmatrix-0.3.0}/modelmatrix/_version.py

@@ -18,7 +18,7 @@ version_tuple: tuple[int | str, ...]
 commit_id: str | None
 __commit_id__: str | None
 
-__version__ = version = '0.2.0'
-__version_tuple__ = version_tuple = (0, 2, 0)
+__version__ = version = '0.3.0'
+__version_tuple__ = version_tuple = (0, 3, 0)
 
-__commit_id__ = commit_id = 'g9dbe6c957'
+__commit_id__ = commit_id = 'g8c1f37d25'

modelmatrix-0.3.0/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-0.2.0 → modelmatrix-0.3.0}/modelmatrix/response.py

@@ -71,9 +71,7 @@ class Response(object):
             :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 @@ class Response(object):
         """
         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 @@ class Surv(Response):
             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)
 

{modelmatrix-0.2.0 → modelmatrix-0.3.0/modelmatrix.egg-info}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: modelmatrix
-Version: 0.2.0
+Version: 0.3.0
 Summary: Model matrices from tabular data, as scikit-learn transformers
 Author-email: Jonathan Taylor <jonathan.taylor@stanford.edu>
 License-Expression: BSD-3-Clause

{modelmatrix-0.2.0 → modelmatrix-0.3.0}/modelmatrix.egg-info/SOURCES.txt

@@ -26,6 +26,7 @@ modelmatrix/_native.py
 modelmatrix/_version.py
 modelmatrix/columns.py
 modelmatrix/model_spec.py
+modelmatrix/outcome.py
 modelmatrix/response.py
 modelmatrix/transforms.py
 modelmatrix.egg-info/PKG-INFO
@@ -41,6 +42,7 @@ tests/test_columns.py
 tests/test_model_matrix.py
 tests/test_model_spec.py
 tests/test_model_spec_hard.py
+tests/test_outcome.py
 tests/test_response.py
 tests/test_wildcard.py
 tests/data/Carseats.csv

{modelmatrix-0.2.0 → modelmatrix-0.3.0}/modelmatrix.egg-info/scm_file_list.json

@@ -26,6 +26,7 @@
     "modelmatrix/_native.py",
     "modelmatrix/columns.py",
     "modelmatrix/model_spec.py",
+    "modelmatrix/outcome.py",
     "modelmatrix/response.py",
     "modelmatrix/transforms.py",
     "pyproject.toml",
@@ -37,6 +38,7 @@
     "tests/test_model_matrix.py",
     "tests/test_model_spec.py",
     "tests/test_model_spec_hard.py",
+    "tests/test_outcome.py",
     "tests/test_response.py",
     "tests/test_wildcard.py"
   ]

{modelmatrix-0.2.0 → modelmatrix-0.3.0}/modelmatrix.egg-info/scm_version.json

@@ -1,7 +1,7 @@
 {
-  "tag": "0.2.0",
+  "tag": "0.3.0",
   "distance": 0,
-  "node": "g9dbe6c957fc56c4c60521bb8cabfc5191d2ae5c7",
+  "node": "g8c1f37d256a8f142a5270a220937e375a1ac01cf",
   "dirty": false,
   "branch": "HEAD",
   "node_date": "2026-10-26"

modelmatrix-0.3.0/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'])

{modelmatrix-0.2.0 → modelmatrix-0.3.0}/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)