From 3464e3e2f66b9fbb3223b767695b5d246652d55d Mon Sep 17 00:00:00 2001 From: Eva Janouskova Date: Mon, 7 Oct 2024 23:00:52 +0100 Subject: [PATCH 1/4] lm: predict() does not always return pd.Series --- src/tlo/lm.py | 1 + 1 file changed, 1 insertion(+) diff --git a/src/tlo/lm.py b/src/tlo/lm.py index e099714850..db0b10e716 100644 --- a/src/tlo/lm.py +++ b/src/tlo/lm.py @@ -456,6 +456,7 @@ def predict( # pop the boolean out of the series if we have a single row, # otherwise return the series if len(outcome) == 1 and squeeze_single_row_output: + assert isinstance(outcome.iloc[0], pd.Series) return outcome.iloc[0] else: return outcome From 03b018cc490b0bb933ab8b4924aef93d7be2f50b Mon Sep 17 00:00:00 2001 From: Eva Janouskova Date: Tue, 8 Oct 2024 19:09:05 +0100 Subject: [PATCH 2/4] Revert "lm: predict() does not always return pd.Series" This reverts commit 3464e3e2f66b9fbb3223b767695b5d246652d55d. --- src/tlo/lm.py | 1 - 1 file changed, 1 deletion(-) diff --git a/src/tlo/lm.py b/src/tlo/lm.py index db0b10e716..e099714850 100644 --- a/src/tlo/lm.py +++ b/src/tlo/lm.py @@ -456,7 +456,6 @@ def predict( # pop the boolean out of the series if we have a single row, # otherwise return the series if len(outcome) == 1 and squeeze_single_row_output: - assert isinstance(outcome.iloc[0], pd.Series) return outcome.iloc[0] else: return outcome From 71efb575a35bfb75232f41e3a9c89ce678c3ced5 Mon Sep 17 00:00:00 2001 From: Eva Janouskova Date: Tue, 8 Oct 2024 19:09:51 +0100 Subject: [PATCH 3/4] lm: precise output type of predict() --- src/tlo/lm.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/src/tlo/lm.py b/src/tlo/lm.py index e099714850..218b71e969 100644 --- a/src/tlo/lm.py +++ b/src/tlo/lm.py @@ -385,7 +385,7 @@ def predict( rng: Optional[np.random.RandomState] = None, squeeze_single_row_output=True, **kwargs - ) -> pd.Series: + ) -> Union[pd.Series, np.bool_]: """Evaluate linear model output for a given set of input data. :param df: The input ``DataFrame`` containing the input data to evaluate the From 789b1b5830572800c53744f603c494bcd8a23ce6 Mon Sep 17 00:00:00 2001 From: Eva Janouskova Date: Tue, 8 Oct 2024 19:17:19 +0100 Subject: [PATCH 4/4] lm: update description of squeeze_single_row_output by stating output type when False --- src/tlo/lm.py | 3 ++- 1 file changed, 2 insertions(+), 1 deletion(-) diff --git a/src/tlo/lm.py b/src/tlo/lm.py index 218b71e969..a7538a1a7a 100644 --- a/src/tlo/lm.py +++ b/src/tlo/lm.py @@ -396,7 +396,8 @@ def predict( output directly returned. :param squeeze_single_row_output: If ``rng`` argument is not ``None`` and this argument is set to ``True``, the output for a ``df`` input with a single-row - will be a scalar boolean value rather than a boolean ``Series``. + will be a scalar boolean value rather than a boolean ``Series``, if set to + ``False``, the output will always be a ``Series``. :param **kwargs: Values for any external variables included in model predictors. """