feat: share covariate utilities (#73)

Author: DiogoRibeiro7

gen_surv/_covariates.py

@@ -0,0 +1,68 @@
+"""Utilities for generating covariate matrices with validation."""
+
+from typing import Literal, cast
+
+import numpy as np
+from numpy.typing import NDArray
+
+from ._validation import ParameterError, ensure_positive
+
+_CovParams = dict[str, float | tuple[float, float]]
+
+
+def set_covariate_params(
+    covariate_dist: Literal["normal", "uniform", "binary"],
+    covariate_params: _CovParams | None,
+) -> _CovParams:
+    """Return covariate distribution parameters with defaults filled in."""
+    if covariate_params is not None:
+        return covariate_params
+    if covariate_dist == "normal":
+        return {"mean": 0.0, "std": 1.0}
+    if covariate_dist == "uniform":
+        return {"low": 0.0, "high": 1.0}
+    if covariate_dist == "binary":
+        return {"p": 0.5}
+    raise ParameterError(
+        "covariate_dist",
+        covariate_dist,
+        "unsupported covariate distribution; choose from 'normal', 'uniform', or 'binary'",
+    )
+
+
+def generate_covariates(
+    n: int,
+    n_covariates: int,
+    covariate_dist: Literal["normal", "uniform", "binary"],
+    covariate_params: _CovParams,
+) -> NDArray[np.float64]:
+    """Generate covariate matrix according to the specified distribution."""
+    if covariate_dist == "normal":
+        std = cast(float, covariate_params.get("std", 1.0))
+        ensure_positive(std, "covariate_params['std']")
+        mean = cast(float, covariate_params.get("mean", 0.0))
+        return np.random.normal(mean, std, size=(n, n_covariates))
+    if covariate_dist == "uniform":
+        low = cast(float, covariate_params.get("low", 0.0))
+        high = cast(float, covariate_params.get("high", 1.0))
+        if high <= low:
+            raise ParameterError(
+                "covariate_params['high']",
+                high,
+                "must be greater than 'low'",
+            )
+        return np.random.uniform(low, high, size=(n, n_covariates))
+    if covariate_dist == "binary":
+        p = cast(float, covariate_params.get("p", 0.5))
+        if not 0 <= p <= 1:
+            raise ParameterError(
+                "covariate_params['p']",
+                p,
+                "must be between 0 and 1",
+            )
+        return np.random.binomial(1, p, size=(n, n_covariates)).astype(float)
+    raise ParameterError(
+        "covariate_dist",
+        covariate_dist,
+        "unsupported covariate distribution; choose from 'normal', 'uniform', or 'binary'",
+    )

gen_surv/mixture.py

@@ -11,36 +11,22 @@
 import pandas as pd
 from numpy.typing import NDArray
 
-_TAIL_FRACTION = 0.1
-_SMOOTH_MIN_TAIL = 3
+_TAIL_FRACTION: float = 0.1
+_SMOOTH_MIN_TAIL: int = 3
 
+from ._covariates import generate_covariates, set_covariate_params
 from ._validation import (
     LengthError,
     ParameterError,
     ensure_censoring_model,
     ensure_in_choices,
+    ensure_numeric_sequence,
     ensure_positive,
+    ensure_positive_int,
 )
 from .censoring import rexpocens, runifcens
 
 
-def _set_covariate_params(
-    covariate_dist: str,
-    covariate_params: dict[str, float | tuple[float, float]] | None,
-) -> dict[str, float | tuple[float, float]]:
-    if covariate_params is not None:
-        return covariate_params
-    if covariate_dist == "normal":
-        return {"mean": 0.0, "std": 1.0}
-    if covariate_dist == "uniform":
-        return {"low": 0.0, "high": 1.0}
-    if covariate_dist == "binary":
-        return {"p": 0.5}
-    raise ParameterError(
-        "covariate_dist", covariate_dist, "must be one of {'normal','uniform','binary'}"
-    )
-
-
 def _prepare_betas(
     betas_survival: list[float] | None,
     betas_cure: list[float] | None,
@@ -49,49 +35,26 @@ def _prepare_betas(
     if betas_survival is None:
         betas_survival_arr = np.random.normal(0, 0.5, size=n_covariates)
     else:
+        ensure_numeric_sequence(betas_survival, "betas_survival")
         betas_survival_arr = np.asarray(betas_survival, dtype=float)
         n_covariates = len(betas_survival_arr)
 
     if betas_cure is None:
         betas_cure_arr = np.random.normal(0, 0.5, size=n_covariates)
     else:
+        ensure_numeric_sequence(betas_cure, "betas_cure")
         betas_cure_arr = np.asarray(betas_cure, dtype=float)
         if len(betas_cure_arr) != n_covariates:
             raise LengthError("betas_cure", len(betas_cure_arr), n_covariates)
 
     return betas_survival_arr, betas_cure_arr, n_covariates
 
 
-def _generate_covariates(
-    n: int,
-    n_covariates: int,
-    covariate_dist: str,
-    covariate_params: dict[str, float | tuple[float, float]],
-) -> NDArray[np.float64]:
-    if covariate_dist == "normal":
-        return np.random.normal(
-            covariate_params.get("mean", 0.0),
-            covariate_params.get("std", 1.0),
-            size=(n, n_covariates),
-        )
-    if covariate_dist == "uniform":
-        return np.random.uniform(
-            covariate_params.get("low", 0.0),
-            covariate_params.get("high", 1.0),
-            size=(n, n_covariates),
-        )
-    if covariate_dist == "binary":
-        return np.random.binomial(
-            1, covariate_params.get("p", 0.5), size=(n, n_covariates)
-        ).astype(float)
-    raise ParameterError(
-        "covariate_dist", covariate_dist, "must be one of {'normal','uniform','binary'}"
-    )
-
-
 def _cure_status(
     lp_cure: NDArray[np.float64], cure_fraction: float
 ) -> NDArray[np.int64]:
+    if not 0 < cure_fraction < 1:
+        raise ParameterError("cure_fraction", cure_fraction, "must be between 0 and 1")
     cure_probs = 1 / (
         1 + np.exp(-(np.log(cure_fraction / (1 - cure_fraction)) + lp_cure))
     )
@@ -104,6 +67,9 @@ def _survival_times(
     baseline_hazard: float,
     max_time: float | None,
 ) -> NDArray[np.float64]:
+    ensure_positive(baseline_hazard, "baseline_hazard")
+    if max_time is not None:
+        ensure_positive(max_time, "max_time")
     n = cured.size
     times = np.zeros(n, dtype=float)
     non_cured = cured == 0
@@ -122,6 +88,10 @@ def _apply_censoring(
     cens_par: float,
     max_time: float | None,
 ) -> tuple[NDArray[np.float64], NDArray[np.int64]]:
+    ensure_censoring_model(model_cens)
+    ensure_positive(cens_par, "cens_par")
+    if max_time is not None:
+        ensure_positive(max_time, "max_time")
     rfunc = runifcens if model_cens == "uniform" else rexpocens
     cens_times = rfunc(len(survival_times), cens_par)
     observed = np.minimum(survival_times, cens_times)
@@ -213,16 +183,21 @@ def gen_mixture_cure(
     if seed is not None:
         np.random.seed(seed)
 
+    ensure_positive_int(n, "n")
+    ensure_positive_int(n_covariates, "n_covariates")
+    ensure_positive(baseline_hazard, "baseline_hazard")
+    ensure_positive(cens_par, "cens_par")
+    if max_time is not None:
+        ensure_positive(max_time, "max_time")
     if not 0 <= cure_fraction <= 1:
         raise ParameterError("cure_fraction", cure_fraction, "must be between 0 and 1")
-    ensure_positive(baseline_hazard, "baseline_hazard")
 
     ensure_in_choices(covariate_dist, "covariate_dist", {"normal", "uniform", "binary"})
-    covariate_params = _set_covariate_params(covariate_dist, covariate_params)
+    covariate_params = set_covariate_params(covariate_dist, covariate_params)
     betas_survival_arr, betas_cure_arr, n_covariates = _prepare_betas(
         betas_survival, betas_cure, n_covariates
     )
-    X = _generate_covariates(n, n_covariates, covariate_dist, covariate_params)
+    X = generate_covariates(n, n_covariates, covariate_dist, covariate_params)
     lp_survival = X @ betas_survival_arr
     lp_cure = X @ betas_cure_arr
     cured = _cure_status(lp_cure, cure_fraction)
@@ -274,6 +249,14 @@ def cure_fraction_estimate(
     based on the plateau of the survival curve. It may not be accurate for
     small sample sizes or heavy censoring.
     """
+    if time_col not in data.columns or status_col not in data.columns:
+        missing = [c for c in (time_col, status_col) if c not in data.columns]
+        raise ParameterError(
+            "data",
+            data.columns.tolist(),
+            f"missing required column(s): {', '.join(missing)}",
+        )
+    ensure_positive(bandwidth, "bandwidth")
     # Sort data by time
     sorted_data = data.sort_values(by=time_col).copy()
 

gen_surv/piecewise.py

@@ -5,32 +5,52 @@
 exponential distributions with time-dependent hazards.
 """
 
-from typing import Dict, List, Literal, Optional, Tuple, Union
+from typing import Literal
 
 import numpy as np
 import pandas as pd
+from numpy.typing import NDArray
 
+from ._covariates import generate_covariates, set_covariate_params
 from ._validation import (
     ParameterError,
     ensure_censoring_model,
     ensure_in_choices,
+    ensure_numeric_sequence,
+    ensure_positive,
+    ensure_positive_int,
     ensure_positive_sequence,
     ensure_sequence_length,
 )
 from .censoring import rexpocens, runifcens
 
 
+def _validate_piecewise_params(
+    breakpoints: list[float], hazard_rates: list[float]
+) -> None:
+    """Validate breakpoint and hazard rate sequences."""
+    ensure_sequence_length(hazard_rates, len(breakpoints) + 1, "hazard_rates")
+    ensure_positive_sequence(breakpoints, "breakpoints")
+    ensure_positive_sequence(hazard_rates, "hazard_rates")
+    if np.any(np.diff(breakpoints) <= 0):
+        raise ParameterError(
+            "breakpoints",
+            breakpoints,
+            "must be a strictly increasing sequence",
+        )
+
+
 def gen_piecewise_exponential(
     n: int,
-    breakpoints: List[float],
-    hazard_rates: List[float],
-    betas: Optional[Union[List[float], np.ndarray]] = None,
+    breakpoints: list[float],
+    hazard_rates: list[float],
+    betas: list[float] | NDArray[np.float64] | None = None,
     n_covariates: int = 2,
     covariate_dist: Literal["normal", "uniform", "binary"] = "normal",
-    covariate_params: Optional[Dict[str, Union[float, Tuple[float, float]]]] = None,
+    covariate_params: dict[str, float | tuple[float, float]] | None = None,
     model_cens: Literal["uniform", "exponential"] = "uniform",
     cens_par: float = 5.0,
-    seed: Optional[int] = None,
+    seed: int | None = None,
 ) -> pd.DataFrame:
     """
     Generate survival data using a piecewise exponential distribution.
@@ -88,55 +108,27 @@ def gen_piecewise_exponential(
     if seed is not None:
         np.random.seed(seed)
 
+    ensure_positive_int(n, "n")
+    ensure_positive_int(n_covariates, "n_covariates")
+    ensure_positive(cens_par, "cens_par")
+
     # Validate inputs
-    ensure_sequence_length(hazard_rates, len(breakpoints) + 1, "hazard_rates")
-    ensure_positive_sequence(breakpoints, "breakpoints")
-    ensure_positive_sequence(hazard_rates, "hazard_rates")
-    if np.any(np.diff(breakpoints) <= 0):
-        raise ParameterError("breakpoints", breakpoints, "must be in ascending order")
+    _validate_piecewise_params(breakpoints, hazard_rates)
 
     ensure_censoring_model(model_cens)
     ensure_in_choices(covariate_dist, "covariate_dist", {"normal", "uniform", "binary"})
-
-    # Set default covariate parameters if not provided
-    if covariate_params is None:
-        if covariate_dist == "normal":
-            covariate_params = {"mean": 0.0, "std": 1.0}
-        elif covariate_dist == "uniform":
-            covariate_params = {"low": 0.0, "high": 1.0}
-        elif covariate_dist == "binary":
-            covariate_params = {"p": 0.5}
+    covariate_params = set_covariate_params(covariate_dist, covariate_params)
 
     # Set default betas if not provided
     if betas is None:
         betas = np.random.normal(0, 0.5, size=n_covariates)
     else:
-        betas = np.array(betas)
+        ensure_numeric_sequence(betas, "betas")
+        betas = np.array(betas, dtype=float)
         n_covariates = len(betas)
 
     # Generate covariates
-    if covariate_dist == "normal":
-        X = np.random.normal(
-            covariate_params.get("mean", 0.0),
-            covariate_params.get("std", 1.0),
-            size=(n, n_covariates),
-        )
-    elif covariate_dist == "uniform":
-        X = np.random.uniform(
-            covariate_params.get("low", 0.0),
-            covariate_params.get("high", 1.0),
-            size=(n, n_covariates),
-        )
-    elif covariate_dist == "binary":
-        X = np.random.binomial(
-            1, covariate_params.get("p", 0.5), size=(n, n_covariates)
-        )
-    else:  # pragma: no cover - validated above
-        raise ParameterError(
-            "covariate_dist",
-            covariate_dist,
-            "must be one of {'normal', 'uniform', 'binary'}",
-        )
+    X = generate_covariates(n, n_covariates, covariate_dist, covariate_params)
 
     # Calculate linear predictor
     linear_predictor = X @ betas
@@ -209,8 +201,10 @@ def gen_piecewise_exponential(
 
 
 def piecewise_hazard_function(
-    t: Union[float, np.ndarray], breakpoints: List[float], hazard_rates: List[float]
-) -> Union[float, np.ndarray]:
+    t: float | NDArray[np.float64],
+    breakpoints: list[float],
+    hazard_rates: list[float],
+) -> float | NDArray[np.float64]:
     """
     Calculate the hazard function value at time t for a piecewise exponential distribution.
 
@@ -228,6 +222,8 @@ def piecewise_hazard_function(
     float or array
         Hazard function value(s) at time t.
     """
+    _validate_piecewise_params(breakpoints, hazard_rates)
+
     # Convert scalar input to array for consistent processing
     scalar_input = np.isscalar(t)
     t_array = np.atleast_1d(t)
@@ -253,8 +249,10 @@ def piecewise_hazard_function(
 
 
 def piecewise_survival_function(
-    t: Union[float, np.ndarray], breakpoints: List[float], hazard_rates: List[float]
-) -> Union[float, np.ndarray]:
+    t: float | NDArray[np.float64],
+    breakpoints: list[float],
+    hazard_rates: list[float],
+) -> float | NDArray[np.float64]:
     """
     Calculate the survival function at time t for a piecewise exponential distribution.
 
@@ -272,6 +270,8 @@ def piecewise_survival_function(
     float or array
         Survival function value(s) at time t.
     """
+    _validate_piecewise_params(breakpoints, hazard_rates)
+
     # Convert scalar input to array for consistent processing
     scalar_input = np.isscalar(t)
     t_array = np.atleast_1d(t)