Generators

The eleven models, grouped by module. Each function returns a pandas.DataFrame; see Output schemas for what the columns mean.

Cox proportional hazards

cphm

Cox Proportional Hazards Model (CPHM) data generation.

This module provides functions to generate survival data following the Cox Proportional Hazards Model with various censoring mechanisms.

generate_cphm_data

generate_cphm_data(
    n: int,
    rfunc: CensoringFunc,
    cens_par: float,
    beta: float,
    covariate_range: float,
    seed: int | None = None,
) -> NDArray[float64]

Generate data from a Cox Proportional Hazards Model (CPHM).

Parameters:

Name	Type	Description	Default
n	int	Number of samples to generate.	required
rfunc	callable	Function to generate censoring times, must accept (size, cens_par).	required
cens_par	float	Parameter passed to the censoring function.	required
beta	float	Coefficient for the covariate.	required
covariate_range	float	Range for the covariate (uniformly sampled from [0, covariate_range]).	required
seed	int	Random seed for reproducibility.	None

Returns:

Type	Description
NDArray[float64]	Array with shape (n, 3): [time, status, X0]

Source code in gen_surv/cphm.py

def generate_cphm_data(
    n: int,
    rfunc: CensoringFunc,
    cens_par: float,
    beta: float,
    covariate_range: float,
    seed: int | None = None,
) -> NDArray[np.float64]:
    """
    Generate data from a Cox Proportional Hazards Model (CPHM).

    Parameters
    ----------
    n : int
        Number of samples to generate.
    rfunc : callable
        Function to generate censoring times, must accept (size, cens_par).
    cens_par : float
        Parameter passed to the censoring function.
    beta : float
        Coefficient for the covariate.
    covariate_range : float
        Range for the covariate (uniformly sampled from [0, covariate_range]).
    seed : int, optional
        Random seed for reproducibility.

    Returns
    -------
    NDArray[np.float64]
        Array with shape ``(n, 3)``: ``[time, status, X0]``
    """
    rng = np.random.default_rng(seed)

    data: NDArray[np.float64] = np.zeros((n, 3), dtype=float)
    # Kept so the latent times can be reported as ground truth; they do not
    # affect any draw.
    event_times: NDArray[np.float64] = np.zeros(n, dtype=float)
    censoring_times: NDArray[np.float64] = np.zeros(n, dtype=float)

    for k in range(n):
        z = rng.uniform(0, covariate_range)
        c = rfunc(1, cens_par, rng)[0]
        x = rng.exponential(scale=1 / np.exp(beta * z))

        time = min(x, c)
        status = int(x <= c)

        data[k, :] = [time, status, z]
        event_times[k] = x
        censoring_times[k] = c

    record(
        beta=beta,
        covariates=data[:, 2].copy(),
        linear_predictor=beta * data[:, 2],
        event_time=event_times,
        censoring_time=censoring_times,
    )
    return data

gen_cphm

gen_cphm(
    n: int,
    model_cens: Literal["uniform", "exponential"],
    cens_par: float,
    beta: float,
    covariate_range: float,
    seed: int | None = None,
) -> DataFrame

Generate survival data following a Cox Proportional Hazards Model.

Parameters:

Name	Type	Description	Default
n	int	Number of observations.	required
model_cens	(uniform, exponential)	Type of censoring mechanism.	"uniform"
cens_par	float	Parameter for the censoring model.	required
beta	float	Coefficient for the covariate.	required
covariate_range	float	Upper bound for the covariate values (uniform between 0 and covariate_range).	required
seed	int	Random seed for reproducibility.	None

Returns:

Type	Description
DataFrame	DataFrame with columns ["time", "status", "X0"] - time: observed event or censoring time - status: event indicator (1=event, 0=censored) - X0: predictor variable

Examples:

>>> from gen_surv.cphm import gen_cphm
>>> df = gen_cphm(n=100, model_cens="uniform", cens_par=1.0, beta=0.5, covariate_range=2.0)
>>> df.head()
   time  status        X0
0  0.23     1.0       1.42
1  0.78     0.0       0.89
...

Source code in gen_surv/cphm.py

def gen_cphm(
    n: int,
    model_cens: Literal["uniform", "exponential"],
    cens_par: float,
    beta: float,
    covariate_range: float,
    seed: int | None = None,
) -> pd.DataFrame:
    """
    Generate survival data following a Cox Proportional Hazards Model.

    Parameters
    ----------
    n : int
        Number of observations.
    model_cens : {"uniform", "exponential"}
        Type of censoring mechanism.
    cens_par : float
        Parameter for the censoring model.
    beta : float
        Coefficient for the covariate.
    covariate_range : float
        Upper bound for the covariate values (uniform between 0 and covariate_range).
    seed : int, optional
        Random seed for reproducibility.

    Returns
    -------
    pd.DataFrame
        DataFrame with columns ["time", "status", "X0"]
        - time: observed event or censoring time
        - status: event indicator (1=event, 0=censored)
        - X0: predictor variable

    Examples
    --------
    >>> from gen_surv.cphm import gen_cphm
    >>> df = gen_cphm(n=100, model_cens="uniform", cens_par=1.0, beta=0.5, covariate_range=2.0)
    >>> df.head()
       time  status        X0
    0  0.23     1.0       1.42
    1  0.78     0.0       0.89
    ...
    """
    validate_gen_cphm_inputs(n, model_cens, cens_par, covariate_range, beta)

    rfunc = {"uniform": runifcens, "exponential": rexpocens}[model_cens]

    data = generate_cphm_data(n, rfunc, cens_par, beta, covariate_range, seed)
    return pd.DataFrame(data, columns=["time", "status", "X0"])

Accelerated failure time

aft

Accelerated Failure Time (AFT) models including Weibull, Log-Normal, and Log-Logistic distributions.

gen_aft_log_normal

gen_aft_log_normal(
    n: int,
    beta: List[float],
    sigma: float,
    model_cens: Literal["uniform", "exponential"],
    cens_par: float,
    seed: int | None = None,
) -> DataFrame

Simulate survival data under a Log-Normal Accelerated Failure Time (AFT) model.

Parameters:

Name	Type	Description	Default
n	int	Number of individuals	required
beta	list of float	Coefficients for covariates	required
sigma	float	Standard deviation of the log-error term	required
model_cens	(uniform, exponential)	Censoring mechanism	"uniform"
cens_par	float	Parameter for censoring distribution	required
seed	int	Random seed for reproducibility	None

Returns:

Type	Description
DataFrame	DataFrame with columns ['id', 'time', 'status', 'X0', ..., 'Xp']

Examples:

>>> from gen_surv.aft import gen_aft_log_normal
>>> df = gen_aft_log_normal(
...     n=100,
...     beta=[0.5, -0.3],
...     sigma=1.0,
...     model_cens="uniform",
...     cens_par=2.0,
...     seed=42,
... )
>>> df.head()

Source code in gen_surv/aft.py

def gen_aft_log_normal(
    n: int,
    beta: List[float],
    sigma: float,
    model_cens: Literal["uniform", "exponential"],
    cens_par: float,
    seed: int | None = None,
) -> pd.DataFrame:
    """
    Simulate survival data under a Log-Normal Accelerated Failure Time (AFT) model.

    Parameters
    ----------
    n : int
        Number of individuals
    beta : list of float
        Coefficients for covariates
    sigma : float
        Standard deviation of the log-error term
    model_cens : {"uniform", "exponential"}
        Censoring mechanism
    cens_par : float
        Parameter for censoring distribution
    seed : int, optional
        Random seed for reproducibility

    Returns
    -------
    pd.DataFrame
        DataFrame with columns ['id', 'time', 'status', 'X0', ..., 'Xp']

    Examples
    --------
    >>> from gen_surv.aft import gen_aft_log_normal
    >>> df = gen_aft_log_normal(
    ...     n=100,
    ...     beta=[0.5, -0.3],
    ...     sigma=1.0,
    ...     model_cens="uniform",
    ...     cens_par=2.0,
    ...     seed=42,
    ... )
    >>> df.head()
    """
    rng = np.random.default_rng(seed)
    validate_gen_aft_log_normal_inputs(n, beta, sigma, model_cens, cens_par)

    p = len(beta)
    X = rng.normal(size=(n, p))
    epsilon = rng.normal(loc=0.0, scale=sigma, size=n)
    log_T = X @ np.array(beta) + epsilon
    T = np.exp(log_T)

    rfunc = runifcens if model_cens == "uniform" else rexpocens
    C = rfunc(n, cens_par, rng)

    observed_time = np.minimum(T, C)
    status = (T <= C).astype(int)

    record(
        betas=np.asarray(beta, dtype=float),
        covariates=X,
        linear_predictor=X @ np.asarray(beta, dtype=float),
        event_time=T,
        censoring_time=C,
    )

    data = pd.DataFrame({"id": np.arange(n), "time": observed_time, "status": status})

    for j in range(p):
        data[f"X{j}"] = X[:, j]

    return data

gen_aft_weibull

gen_aft_weibull(
    n: int,
    beta: List[float],
    shape: float,
    scale: float,
    model_cens: Literal["uniform", "exponential"],
    cens_par: float,
    seed: int | None = None,
) -> DataFrame

Simulate survival data under a Weibull Accelerated Failure Time (AFT) model.

The Weibull AFT model has survival function: S(t|X) = exp(-(t/scale)^shape * exp(-X*beta))

Parameters:

Name	Type	Description	Default
n	int	Number of individuals	required
beta	list of float	Coefficients for covariates	required
shape	float	Weibull shape parameter (k > 0)	required
scale	float	Weibull scale parameter (λ > 0)	required
model_cens	(uniform, exponential)	Censoring mechanism	"uniform"
cens_par	float	Parameter for censoring distribution	required
seed	int	Random seed for reproducibility	None

Returns:

Type	Description
DataFrame	DataFrame with columns ['id', 'time', 'status', 'X0', ..., 'Xp']

Examples:

>>> from gen_surv.aft import gen_aft_weibull
>>> df = gen_aft_weibull(
...     n=100,
...     beta=[0.5, -0.3],
...     shape=1.2,
...     scale=2.0,
...     model_cens="uniform",
...     cens_par=2.0,
...     seed=42,
... )
>>> df.head()

Source code in gen_surv/aft.py

def gen_aft_weibull(
    n: int,
    beta: List[float],
    shape: float,
    scale: float,
    model_cens: Literal["uniform", "exponential"],
    cens_par: float,
    seed: int | None = None,
) -> pd.DataFrame:
    """
    Simulate survival data under a Weibull Accelerated Failure Time (AFT) model.

    The Weibull AFT model has survival function:
    S(t|X) = exp(-(t/scale)^shape * exp(-X*beta))

    Parameters
    ----------
    n : int
        Number of individuals
    beta : list of float
        Coefficients for covariates
    shape : float
        Weibull shape parameter (k > 0)
    scale : float
        Weibull scale parameter (λ > 0)
    model_cens : {"uniform", "exponential"}
        Censoring mechanism
    cens_par : float
        Parameter for censoring distribution
    seed : int, optional
        Random seed for reproducibility

    Returns
    -------
    pd.DataFrame
        DataFrame with columns ['id', 'time', 'status', 'X0', ..., 'Xp']

    Examples
    --------
    >>> from gen_surv.aft import gen_aft_weibull
    >>> df = gen_aft_weibull(
    ...     n=100,
    ...     beta=[0.5, -0.3],
    ...     shape=1.2,
    ...     scale=2.0,
    ...     model_cens="uniform",
    ...     cens_par=2.0,
    ...     seed=42,
    ... )
    >>> df.head()
    """
    rng = np.random.default_rng(seed)
    validate_gen_aft_weibull_inputs(n, beta, shape, scale, model_cens, cens_par)

    p = len(beta)
    X = rng.normal(size=(n, p))

    # Linear predictor
    eta = X @ np.array(beta)

    # Generate Weibull survival times
    U = rng.uniform(size=n)
    T = scale * (-np.log(U) * np.exp(-eta)) ** (1 / shape)

    # Generate censoring times
    rfunc = runifcens if model_cens == "uniform" else rexpocens
    C = rfunc(n, cens_par, rng)

    # Observed time is the minimum of event time and censoring time
    observed_time = np.minimum(T, C)
    status = (T <= C).astype(int)

    record(
        betas=np.asarray(beta, dtype=float),
        covariates=X,
        linear_predictor=X @ np.asarray(beta, dtype=float),
        event_time=T,
        censoring_time=C,
    )

    data = pd.DataFrame({"id": np.arange(n), "time": observed_time, "status": status})

    for j in range(p):
        data[f"X{j}"] = X[:, j]

    return data

gen_aft_log_logistic

gen_aft_log_logistic(
    n: int,
    beta: List[float],
    shape: float,
    scale: float,
    model_cens: Literal["uniform", "exponential"],
    cens_par: float,
    seed: int | None = None,
) -> DataFrame

Simulate survival data under a Log-Logistic Accelerated Failure Time (AFT) model.

The Log-Logistic AFT model has survival function: S(t|X) = 1 / (1 + (t/scale)^shape * exp(X*beta))

Log-logistic distribution is useful when the hazard rate first increases and then decreases.

Parameters:

Name	Type	Description	Default
n	int	Number of individuals	required
beta	list of float	Coefficients for covariates	required
shape	float	Log-logistic shape parameter (α > 0)	required
scale	float	Log-logistic scale parameter (β > 0)	required
model_cens	(uniform, exponential)	Censoring mechanism	"uniform"
cens_par	float	Parameter for censoring distribution	required
seed	int	Random seed for reproducibility	None

Returns:

Type	Description
DataFrame	DataFrame with columns ['id', 'time', 'status', 'X0', ..., 'Xp']

Examples:

>>> from gen_surv.aft import gen_aft_log_logistic
>>> df = gen_aft_log_logistic(
...     n=100,
...     beta=[0.5, -0.3],
...     shape=1.2,
...     scale=2.0,
...     model_cens="uniform",
...     cens_par=2.0,
...     seed=42,
... )
>>> df.head()

Source code in gen_surv/aft.py

def gen_aft_log_logistic(
    n: int,
    beta: List[float],
    shape: float,
    scale: float,
    model_cens: Literal["uniform", "exponential"],
    cens_par: float,
    seed: int | None = None,
) -> pd.DataFrame:
    """
    Simulate survival data under a Log-Logistic Accelerated Failure Time (AFT) model.

    The Log-Logistic AFT model has survival function:
    S(t|X) = 1 / (1 + (t/scale)^shape * exp(X*beta))

    Log-logistic distribution is useful when the hazard rate first increases and then decreases.

    Parameters
    ----------
    n : int
        Number of individuals
    beta : list of float
        Coefficients for covariates
    shape : float
        Log-logistic shape parameter (α > 0)
    scale : float
        Log-logistic scale parameter (β > 0)
    model_cens : {"uniform", "exponential"}
        Censoring mechanism
    cens_par : float
        Parameter for censoring distribution
    seed : int, optional
        Random seed for reproducibility

    Returns
    -------
    pd.DataFrame
        DataFrame with columns ['id', 'time', 'status', 'X0', ..., 'Xp']

    Examples
    --------
    >>> from gen_surv.aft import gen_aft_log_logistic
    >>> df = gen_aft_log_logistic(
    ...     n=100,
    ...     beta=[0.5, -0.3],
    ...     shape=1.2,
    ...     scale=2.0,
    ...     model_cens="uniform",
    ...     cens_par=2.0,
    ...     seed=42,
    ... )
    >>> df.head()
    """
    rng = np.random.default_rng(seed)
    validate_gen_aft_log_logistic_inputs(n, beta, shape, scale, model_cens, cens_par)

    p = len(beta)
    X = rng.normal(size=(n, p))

    # Linear predictor
    eta = X @ np.array(beta)

    # Generate Log-Logistic survival times
    U = rng.uniform(size=n)

    # Inverse CDF method: S(t) = 1/(1 + (t/scale)^shape)
    # so t = scale * (1/S - 1)^(1/shape)
    # For random U ~ Uniform(0,1), we can use U as 1-S
    # t = scale * (1/(1-U) - 1)^(1/shape) * exp(-eta/shape)
    # simplifies to: t = scale * (U/(1-U))^(1/shape) * exp(-eta/shape)

    # Avoid numerical issues near 1
    U = np.clip(U, 0.001, 0.999)
    T = scale * (U / (1 - U)) ** (1 / shape) * np.exp(-eta / shape)

    # Generate censoring times
    rfunc = runifcens if model_cens == "uniform" else rexpocens
    C = rfunc(n, cens_par, rng)

    # Observed time is the minimum of event time and censoring time
    observed_time = np.minimum(T, C)
    status = (T <= C).astype(int)

    record(
        betas=np.asarray(beta, dtype=float),
        covariates=X,
        linear_predictor=X @ np.asarray(beta, dtype=float),
        event_time=T,
        censoring_time=C,
    )

    data = pd.DataFrame({"id": np.arange(n), "time": observed_time, "status": status})

    for j in range(p):
        data[f"X{j}"] = X[:, j]

    return data

Piecewise exponential

piecewise

Piecewise Exponential survival models.

This module provides functions for generating survival data from piecewise exponential distributions with time-dependent hazards.

gen_piecewise_exponential

gen_piecewise_exponential(
    n: int,
    breakpoints: list[float],
    hazard_rates: list[float],
    betas: list[float] | NDArray[float64] | None = None,
    n_covariates: int = 2,
    covariate_dist: Literal[
        "normal", "uniform", "binary"
    ] = "normal",
    covariate_params: dict[str, float] | None = None,
    model_cens: Literal[
        "uniform", "exponential"
    ] = "uniform",
    cens_par: float = 5.0,
    seed: int | None = None,
) -> DataFrame

Generate survival data using a piecewise exponential distribution.

Parameters:

Name	Type	Description	Default
n	int	Number of subjects.	required
breakpoints	list of float	Time points where hazard rates change. Must be in ascending order. The first interval is [0, breakpoints[0]), the second is [breakpoints[0], breakpoints[1]), etc.	required
hazard_rates	list of float	Hazard rates for each interval. Length should be len(breakpoints) + 1.	required
betas	list or array	Coefficients for covariates. If None, generates random coefficients.	None
n_covariates	int	Number of covariates to generate if betas is None.	2
covariate_dist	(normal, uniform, binary)	Distribution to generate covariates from.	"normal"
covariate_params	dict	Parameters for covariate distribution: - "normal": {"mean": float, "std": float} - "uniform": {"low": float, "high": float} - "binary": {"p": float} If None, uses defaults based on distribution.	None
model_cens	(uniform, exponential)	Censoring mechanism.	"uniform"
cens_par	float	Parameter for censoring distribution.	5.0
seed	int	Random seed for reproducibility.	None

Returns:

Type	Description
DataFrame	DataFrame with columns: - "id": Subject identifier - "time": Time to event or censoring - "status": Event indicator (1=event, 0=censored) - "X0", "X1", ...: Covariates

Examples:

>>> from gen_surv.piecewise import gen_piecewise_exponential
>>>
>>> # Generate data with 3 intervals (increasing hazard)
>>> df = gen_piecewise_exponential(
...     n=100,
...     breakpoints=[1.0, 3.0],
...     hazard_rates=[0.2, 0.5, 1.0],
...     betas=[0.8, -0.5],
...     seed=42
... )

Source code in gen_surv/piecewise.py

def gen_piecewise_exponential(
    n: int,
    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: dict[str, float] | None = None,
    model_cens: Literal["uniform", "exponential"] = "uniform",
    cens_par: float = 5.0,
    seed: int | None = None,
) -> pd.DataFrame:
    """
    Generate survival data using a piecewise exponential distribution.

    Parameters
    ----------
    n : int
        Number of subjects.
    breakpoints : list of float
        Time points where hazard rates change. Must be in ascending order.
        The first interval is [0, breakpoints[0]), the second is [breakpoints[0], breakpoints[1]), etc.
    hazard_rates : list of float
        Hazard rates for each interval. Length should be len(breakpoints) + 1.
    betas : list or array, optional
        Coefficients for covariates. If None, generates random coefficients.
    n_covariates : int, default=2
        Number of covariates to generate if betas is None.
    covariate_dist : {"normal", "uniform", "binary"}, default="normal"
        Distribution to generate covariates from.
    covariate_params : dict, optional
        Parameters for covariate distribution:
        - "normal": {"mean": float, "std": float}
        - "uniform": {"low": float, "high": float}
        - "binary": {"p": float}
        If None, uses defaults based on distribution.
    model_cens : {"uniform", "exponential"}, default="uniform"
        Censoring mechanism.
    cens_par : float, default=5.0
        Parameter for censoring distribution.
    seed : int, optional
        Random seed for reproducibility.

    Returns
    -------
    pd.DataFrame
        DataFrame with columns:
        - "id": Subject identifier
        - "time": Time to event or censoring
        - "status": Event indicator (1=event, 0=censored)
        - "X0", "X1", ...: Covariates

    Examples
    --------
    >>> from gen_surv.piecewise import gen_piecewise_exponential
    >>>
    >>> # Generate data with 3 intervals (increasing hazard)
    >>> df = gen_piecewise_exponential(
    ...     n=100,
    ...     breakpoints=[1.0, 3.0],
    ...     hazard_rates=[0.2, 0.5, 1.0],
    ...     betas=[0.8, -0.5],
    ...     seed=42
    ... )
    """
    rng = np.random.default_rng(seed)

    validate_gen_piecewise_inputs(
        n,
        breakpoints,
        hazard_rates,
        n_covariates,
        model_cens,
        cens_par,
        covariate_dist,
    )
    covariate_params = set_covariate_params(covariate_dist, covariate_params)

    # Set default betas if not provided
    betas, n_covariates = prepare_betas(betas, n_covariates, rng)

    # Generate covariates
    X = generate_covariates(n, n_covariates, covariate_dist, covariate_params, rng)

    # Calculate linear predictor
    linear_predictor = X @ betas

    # Generate survival times using piecewise exponential distribution
    survival_times = np.zeros(n)

    for i in range(n):
        # Adjust hazard rates by the covariate effect
        adjusted_hazard_rates = [h * np.exp(linear_predictor[i]) for h in hazard_rates]

        # Generate random uniform between 0 and 1
        u = rng.uniform(0, 1)

        # Calculate survival time using inverse CDF method for piecewise exponential
        remaining_time = -np.log(u)  # Initial time remaining (for standard exponential)
        total_time = 0.0

        # Start with the first interval [0, breakpoints[0])
        interval_width = breakpoints[0]
        hazard = adjusted_hazard_rates[0]
        time_to_consume = remaining_time / hazard

        if time_to_consume < interval_width:
            # Event occurs in first interval
            survival_times[i] = time_to_consume
            continue

        # Event occurs after first interval
        total_time += interval_width
        remaining_time -= hazard * interval_width

        # Go through middle intervals [breakpoints[j-1], breakpoints[j])
        #
        # The ``else`` belongs to the ``for``: it runs only when the loop was
        # not broken out of, meaning the subject outlived every bounded
        # interval. Releases up to 2.0.1 used a trailing ``if remaining_time >
        # 0`` here instead, which also ran after the ``break`` and overwrote
        # the time just computed with one derived from the *last* hazard rate.
        # Any event falling in a middle interval was therefore drawn at the
        # wrong rate.
        for j in range(1, len(breakpoints)):
            interval_width = breakpoints[j] - breakpoints[j - 1]
            hazard = adjusted_hazard_rates[j]
            time_to_consume = remaining_time / hazard

            if time_to_consume < interval_width:
                # Event occurs in this interval
                survival_times[i] = total_time + time_to_consume
                break

            # Event occurs after this interval
            total_time += interval_width
            remaining_time -= hazard * interval_width
        else:
            # Survived every bounded interval: the remainder is consumed at the
            # open-ended last rate.
            hazard = adjusted_hazard_rates[-1]
            survival_times[i] = total_time + remaining_time / hazard

    # Generate censoring times
    rfunc = runifcens if model_cens == "uniform" else rexpocens
    cens_times = rfunc(n, cens_par, rng)

    # Determine observed time and status
    observed_times = np.minimum(survival_times, cens_times)
    status = (survival_times <= cens_times).astype(int)

    record(
        betas=betas,
        covariates=X,
        linear_predictor=linear_predictor,
        event_time=survival_times,
        censoring_time=cens_times,
        breakpoints=np.asarray(breakpoints, dtype=float),
        hazard_rates=np.asarray(hazard_rates, dtype=float),
    )

    # Create DataFrame
    data = pd.DataFrame({"id": np.arange(n), "time": observed_times, "status": status})

    # Add covariates
    for j in range(n_covariates):
        data[f"X{j}"] = X[:, j]

    return data

piecewise_hazard_function

piecewise_hazard_function(
    t: float | NDArray[float64],
    breakpoints: list[float],
    hazard_rates: list[float],
) -> float | NDArray[float64]

Calculate the hazard function value at time t for a piecewise exponential distribution.

Parameters:

Name	Type	Description	Default
t	float or array	Time point(s) at which to evaluate the hazard function.	required
breakpoints	list of float	Time points where hazard rates change.	required
hazard_rates	list of float	Hazard rates for each interval.	required

Returns:

Type	Description
float or array	Hazard function value(s) at time t.

Source code in gen_surv/piecewise.py

def piecewise_hazard_function(
    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.

    Parameters
    ----------
    t : float or array
        Time point(s) at which to evaluate the hazard function.
    breakpoints : list of float
        Time points where hazard rates change.
    hazard_rates : list of float
        Hazard rates for each interval.

    Returns
    -------
    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)
    result = np.zeros_like(t_array)

    # Assign hazard rates based on time intervals
    result[t_array < 0] = 0  # Hazard is 0 for negative times

    # First interval: [0, breakpoints[0])
    mask = (t_array >= 0) & (t_array < breakpoints[0])
    result[mask] = hazard_rates[0]

    # Middle intervals: [breakpoints[j-1], breakpoints[j])
    for j in range(1, len(breakpoints)):
        mask = (t_array >= breakpoints[j - 1]) & (t_array < breakpoints[j])
        result[mask] = hazard_rates[j]

    # Last interval: [breakpoints[-1], infinity)
    mask = t_array >= breakpoints[-1]
    result[mask] = hazard_rates[-1]

    return result[0] if scalar_input else result

piecewise_survival_function

piecewise_survival_function(
    t: float | NDArray[float64],
    breakpoints: list[float],
    hazard_rates: list[float],
) -> float | NDArray[float64]

Calculate the survival function at time t for a piecewise exponential distribution.

Parameters:

Name	Type	Description	Default
t	float or array	Time point(s) at which to evaluate the survival function.	required
breakpoints	list of float	Time points where hazard rates change.	required
hazard_rates	list of float	Hazard rates for each interval.	required

Returns:

Type	Description
float or array	Survival function value(s) at time t.

Source code in gen_surv/piecewise.py

def piecewise_survival_function(
    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.

    Parameters
    ----------
    t : float or array
        Time point(s) at which to evaluate the survival function.
    breakpoints : list of float
        Time points where hazard rates change.
    hazard_rates : list of float
        Hazard rates for each interval.

    Returns
    -------
    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)
    result = np.ones_like(t_array)

    # For each time point, calculate the survival function
    for i, ti in enumerate(t_array):
        if ti <= 0:
            continue  # Survival probability is 1 at time 0 or earlier

        cumulative_hazard = 0.0

        # First interval: [0, min(ti, breakpoints[0]))
        first_interval_end = min(ti, breakpoints[0]) if breakpoints else ti
        cumulative_hazard += hazard_rates[0] * first_interval_end

        if ti <= breakpoints[0]:
            result[i] = np.exp(-cumulative_hazard)
            continue

        # Middle intervals: [breakpoints[j-1], min(ti, breakpoints[j]))
        for j in range(1, len(breakpoints)):
            if ti <= breakpoints[j - 1]:
                break

            interval_start = breakpoints[j - 1]
            interval_end = min(ti, breakpoints[j])
            interval_width = interval_end - interval_start

            cumulative_hazard += hazard_rates[j] * interval_width

            if ti <= breakpoints[j]:
                break

        # Last interval: [breakpoints[-1], ti)
        if ti > breakpoints[-1]:
            last_interval_width = ti - breakpoints[-1]
            cumulative_hazard += hazard_rates[-1] * last_interval_width

        result[i] = np.exp(-cumulative_hazard)

    return result[0] if scalar_input else result

Competing risks

competing_risks

Competing Risks models for survival data simulation.

This module provides functions to generate survival data with competing risks under different hazard specifications.

gen_competing_risks

gen_competing_risks(
    n: int,
    n_risks: int = 2,
    baseline_hazards: (
        Union[List[float], ndarray] | None
    ) = None,
    betas: Union[List[List[float]], ndarray] | None = None,
    covariate_dist: Literal[
        "normal", "uniform", "binary"
    ] = "normal",
    covariate_params: Dict[str, float] | None = None,
    max_time: float | None = 10.0,
    model_cens: Literal[
        "uniform", "exponential"
    ] = "uniform",
    cens_par: float = 5.0,
    seed: int | None = None,
) -> DataFrame

Generate survival data with competing risks.

Parameters:

Name	Type	Description	Default
n	int	Number of subjects.	required
n_risks	int	Number of competing risks.	2
baseline_hazards	list of float or array	Baseline hazard rates for each risk. If None, uses [0.5, 0.3, ...] with decreasing values for subsequent risks.	None
betas	list of list of float or array	Coefficients for covariates, one list per risk. Shape should be (n_risks, n_covariates). If None, generates random coefficients.	None
covariate_dist	(normal, uniform, binary)	Distribution to generate covariates from.	"normal"
covariate_params	dict	Parameters for covariate distribution: - "normal": {"mean": float, "std": float} - "uniform": {"low": float, "high": float} - "binary": {"p": float} If None, uses defaults based on distribution.	None
max_time	float	Maximum simulation time. Set to None for no limit.	10.0
model_cens	(uniform, exponential)	Censoring mechanism.	"uniform"
cens_par	float	Parameter for censoring distribution.	5.0
seed	int	Random seed for reproducibility.	None

Returns:

Type	Description
DataFrame	DataFrame with columns: - "id": Subject identifier - "time": Time to event or censoring - "status": Event indicator (0=censored, 1,2,...=competing events) - "X0", "X1", ...: Covariates

Examples:

>>> from gen_surv.competing_risks import gen_competing_risks
>>>
>>> # Simple example with 2 competing risks
>>> df = gen_competing_risks(
...     n=100,
...     n_risks=2,
...     baseline_hazards=[0.5, 0.3],
...     betas=[[0.8, -0.5], [0.2, 0.7]],
...     seed=42
... )
>>>
>>> # Distribution of event types
>>> df["status"].value_counts()

Source code in gen_surv/competing_risks.py

def gen_competing_risks(
    n: int,
    n_risks: int = 2,
    baseline_hazards: Union[List[float], np.ndarray] | None = None,
    betas: Union[List[List[float]], np.ndarray] | None = None,
    covariate_dist: Literal["normal", "uniform", "binary"] = "normal",
    covariate_params: Dict[str, float] | None = None,
    max_time: float | None = 10.0,
    model_cens: Literal["uniform", "exponential"] = "uniform",
    cens_par: float = 5.0,
    seed: int | None = None,
) -> pd.DataFrame:
    """
    Generate survival data with competing risks.

    Parameters
    ----------
    n : int
        Number of subjects.
    n_risks : int, default=2
        Number of competing risks.
    baseline_hazards : list of float or array, optional
        Baseline hazard rates for each risk. If None, uses [0.5, 0.3, ...]
        with decreasing values for subsequent risks.
    betas : list of list of float or array, optional
        Coefficients for covariates, one list per risk.
        Shape should be (n_risks, n_covariates).
        If None, generates random coefficients.
    covariate_dist : {"normal", "uniform", "binary"}, default="normal"
        Distribution to generate covariates from.
    covariate_params : dict, optional
        Parameters for covariate distribution:
        - "normal": {"mean": float, "std": float}
        - "uniform": {"low": float, "high": float}
        - "binary": {"p": float}
        If None, uses defaults based on distribution.
    max_time : float, optional, default=10.0
        Maximum simulation time. Set to None for no limit.
    model_cens : {"uniform", "exponential"}, default="uniform"
        Censoring mechanism.
    cens_par : float, default=5.0
        Parameter for censoring distribution.
    seed : int, optional
        Random seed for reproducibility.

    Returns
    -------
    pd.DataFrame
        DataFrame with columns:
        - "id": Subject identifier
        - "time": Time to event or censoring
        - "status": Event indicator (0=censored, 1,2,...=competing events)
        - "X0", "X1", ...: Covariates

    Examples
    --------
    >>> from gen_surv.competing_risks import gen_competing_risks
    >>>
    >>> # Simple example with 2 competing risks
    >>> df = gen_competing_risks(
    ...     n=100,
    ...     n_risks=2,
    ...     baseline_hazards=[0.5, 0.3],
    ...     betas=[[0.8, -0.5], [0.2, 0.7]],
    ...     seed=42
    ... )
    >>>
    >>> # Distribution of event types
    >>> df["status"].value_counts()
    """
    rng = np.random.default_rng(seed)

    validate_competing_risks_inputs(
        n,
        n_risks,
        baseline_hazards,
        betas,
        covariate_dist,
        max_time,
        model_cens,
        cens_par,
    )

    # Set default baseline hazards if not provided
    if baseline_hazards is None:
        baseline_hazards = np.array([0.5 / (i + 1) for i in range(n_risks)])
    else:
        baseline_hazards = np.asarray(baseline_hazards, dtype=float)

    covariate_params = set_covariate_params(covariate_dist, covariate_params)
    n_covariates = 2
    betas, n_covariates = prepare_betas_matrix(betas, n_risks, n_covariates, rng)
    X = generate_covariates(n, n_covariates, covariate_dist, covariate_params, rng)

    # Calculate linear predictors for each risk
    linear_predictors = np.zeros((n, n_risks))
    for j in range(n_risks):
        linear_predictors[:, j] = X @ betas[j]

    # Calculate hazard rates
    hazard_rates = np.zeros_like(linear_predictors)
    for j in range(n_risks):
        hazard_rates[:, j] = baseline_hazards[j] * np.exp(linear_predictors[:, j])

    # Generate event times for each risk
    event_times = np.zeros((n, n_risks))
    for j in range(n_risks):
        # Use exponential distribution with rate = hazard
        event_times[:, j] = rng.exponential(1 / hazard_rates[:, j])

    # Generate censoring times
    rfunc = runifcens if model_cens == "uniform" else rexpocens
    cens_times = rfunc(n, cens_par, rng)

    # Find the minimum time for each subject (first event or censoring)
    min_event_times = np.min(event_times, axis=1)
    observed_times = np.minimum(min_event_times, cens_times)

    # Determine event type (0 = censored, 1...n_risks = event type)
    status = np.zeros(n, dtype=int)
    for i in range(n):
        if min_event_times[i] <= cens_times[i]:
            # Find which risk occurred first
            risk_index = np.argmin(event_times[i])
            status[i] = risk_index + 1  # 1-based indexing for event types

    # Cap times at max_time if specified
    if max_time is not None:
        over_max = observed_times > max_time
        observed_times[over_max] = max_time
        status[over_max] = 0  # Censored if beyond max_time

    record(
        betas=betas,
        covariates=X,
        cause_times=event_times,
        censoring_time=cens_times,
        first_event_time=min_event_times,
    )

    # Create DataFrame
    data = pd.DataFrame({"id": np.arange(n), "time": observed_times, "status": status})

    # Add covariates
    for j in range(n_covariates):
        data[f"X{j}"] = X[:, j]

    return data

gen_competing_risks_weibull

gen_competing_risks_weibull(
    n: int,
    n_risks: int = 2,
    shape_params: Union[List[float], ndarray] | None = None,
    scale_params: Union[List[float], ndarray] | None = None,
    betas: Union[List[List[float]], ndarray] | None = None,
    covariate_dist: Literal[
        "normal", "uniform", "binary"
    ] = "normal",
    covariate_params: Dict[str, float] | None = None,
    max_time: float | None = 10.0,
    model_cens: Literal[
        "uniform", "exponential"
    ] = "uniform",
    cens_par: float = 5.0,
    seed: int | None = None,
) -> DataFrame

Generate survival data with competing risks using Weibull hazards.

Parameters:

Name	Type	Description	Default
n	int	Number of subjects.	required
n_risks	int	Number of competing risks.	2
shape_params	list of float or array	Shape parameters for Weibull distribution, one per risk. If None, uses [1.2, 0.8, ...] alternating values.	None
scale_params	list of float or array	Scale parameters for Weibull distribution, one per risk. If None, uses [2.0, 3.0, ...] increasing values.	None
betas	list of list of float or array	Coefficients for covariates, one list per risk. Shape should be (n_risks, n_covariates). If None, generates random coefficients.	None
covariate_dist	(normal, uniform, binary)	Distribution to generate covariates from.	"normal"
covariate_params	dict	Parameters for covariate distribution: - "normal": {"mean": float, "std": float} - "uniform": {"low": float, "high": float} - "binary": {"p": float} If None, uses defaults based on distribution.	None
max_time	float	Maximum simulation time. Set to None for no limit.	10.0
model_cens	(uniform, exponential)	Censoring mechanism.	"uniform"
cens_par	float	Parameter for censoring distribution.	5.0
seed	int	Random seed for reproducibility.	None

Returns:

Type	Description
DataFrame	DataFrame with columns: - "id": Subject identifier - "time": Time to event or censoring - "status": Event indicator (0=censored, 1,2,...=competing events) - "X0", "X1", ...: Covariates

Examples:

>>> from gen_surv.competing_risks import gen_competing_risks_weibull
>>>
>>> # Example with 2 competing risks with different shapes
>>> df = gen_competing_risks_weibull(
...     n=100,
...     n_risks=2,
...     shape_params=[0.8, 1.5],  # Decreasing vs increasing hazard
...     scale_params=[2.0, 3.0],
...     betas=[[0.8, -0.5], [0.2, 0.7]],
...     seed=42
... )

Source code in gen_surv/competing_risks.py

def gen_competing_risks_weibull(
    n: int,
    n_risks: int = 2,
    shape_params: Union[List[float], np.ndarray] | None = None,
    scale_params: Union[List[float], np.ndarray] | None = None,
    betas: Union[List[List[float]], np.ndarray] | None = None,
    covariate_dist: Literal["normal", "uniform", "binary"] = "normal",
    covariate_params: Dict[str, float] | None = None,
    max_time: float | None = 10.0,
    model_cens: Literal["uniform", "exponential"] = "uniform",
    cens_par: float = 5.0,
    seed: int | None = None,
) -> pd.DataFrame:
    """
    Generate survival data with competing risks using Weibull hazards.

    Parameters
    ----------
    n : int
        Number of subjects.
    n_risks : int, default=2
        Number of competing risks.
    shape_params : list of float or array, optional
        Shape parameters for Weibull distribution, one per risk.
        If None, uses [1.2, 0.8, ...] alternating values.
    scale_params : list of float or array, optional
        Scale parameters for Weibull distribution, one per risk.
        If None, uses [2.0, 3.0, ...] increasing values.
    betas : list of list of float or array, optional
        Coefficients for covariates, one list per risk.
        Shape should be (n_risks, n_covariates).
        If None, generates random coefficients.
    covariate_dist : {"normal", "uniform", "binary"}, default="normal"
        Distribution to generate covariates from.
    covariate_params : dict, optional
        Parameters for covariate distribution:
        - "normal": {"mean": float, "std": float}
        - "uniform": {"low": float, "high": float}
        - "binary": {"p": float}
        If None, uses defaults based on distribution.
    max_time : float, optional, default=10.0
        Maximum simulation time. Set to None for no limit.
    model_cens : {"uniform", "exponential"}, default="uniform"
        Censoring mechanism.
    cens_par : float, default=5.0
        Parameter for censoring distribution.
    seed : int, optional
        Random seed for reproducibility.

    Returns
    -------
    pd.DataFrame
        DataFrame with columns:
        - "id": Subject identifier
        - "time": Time to event or censoring
        - "status": Event indicator (0=censored, 1,2,...=competing events)
        - "X0", "X1", ...: Covariates

    Examples
    --------
    >>> from gen_surv.competing_risks import gen_competing_risks_weibull
    >>>
    >>> # Example with 2 competing risks with different shapes
    >>> df = gen_competing_risks_weibull(
    ...     n=100,
    ...     n_risks=2,
    ...     shape_params=[0.8, 1.5],  # Decreasing vs increasing hazard
    ...     scale_params=[2.0, 3.0],
    ...     betas=[[0.8, -0.5], [0.2, 0.7]],
    ...     seed=42
    ... )
    """
    rng = np.random.default_rng(seed)

    validate_competing_risks_inputs(
        n,
        n_risks,
        None,
        betas,
        covariate_dist,
        max_time,
        model_cens,
        cens_par,
    )

    # Set default shape and scale parameters if not provided
    if shape_params is None:
        shape_params = np.array([1.2 if i % 2 == 0 else 0.8 for i in range(n_risks)])
    else:
        shape_params = np.asarray(shape_params, dtype=float)
        ensure_sequence_length(shape_params, n_risks, "shape_params")
        ensure_positive_sequence(shape_params, "shape_params")

    if scale_params is None:
        scale_params = np.array([2.0 + i for i in range(n_risks)])
    else:
        scale_params = np.asarray(scale_params, dtype=float)
        ensure_sequence_length(scale_params, n_risks, "scale_params")
        ensure_positive_sequence(scale_params, "scale_params")

    covariate_params = set_covariate_params(covariate_dist, covariate_params)
    n_covariates = 2
    betas, n_covariates = prepare_betas_matrix(betas, n_risks, n_covariates, rng)
    X = generate_covariates(n, n_covariates, covariate_dist, covariate_params, rng)

    # Calculate linear predictors for each risk
    linear_predictors = np.zeros((n, n_risks))
    for j in range(n_risks):
        linear_predictors[:, j] = X @ betas[j]

    # Generate event times for each risk using Weibull distribution
    event_times = np.zeros((n, n_risks))
    for j in range(n_risks):
        # Adjust the scale parameter using the linear predictor
        adjusted_scale = scale_params[j] * np.exp(
            -linear_predictors[:, j] / shape_params[j]
        )

        # Generate random uniform between 0 and 1
        u = rng.uniform(0, 1, size=n)

        # Convert to Weibull using inverse CDF: t = scale * (-log(1-u))^(1/shape)
        event_times[:, j] = adjusted_scale * (-np.log(1 - u)) ** (1 / shape_params[j])

    # Generate censoring times
    rfunc = runifcens if model_cens == "uniform" else rexpocens
    cens_times = rfunc(n, cens_par, rng)

    # Find the minimum time for each subject (first event or censoring)
    min_event_times = np.min(event_times, axis=1)
    observed_times = np.minimum(min_event_times, cens_times)

    # Determine event type (0 = censored, 1...n_risks = event type)
    status = np.zeros(n, dtype=int)
    for i in range(n):
        if min_event_times[i] <= cens_times[i]:
            # Find which risk occurred first
            risk_index = np.argmin(event_times[i])
            status[i] = risk_index + 1  # 1-based indexing for event types

    # Cap times at max_time if specified
    if max_time is not None:
        over_max = observed_times > max_time
        observed_times[over_max] = max_time
        status[over_max] = 0  # Censored if beyond max_time

    record(
        betas=betas,
        covariates=X,
        cause_times=event_times,
        censoring_time=cens_times,
        first_event_time=min_event_times,
    )

    # Create DataFrame
    data = pd.DataFrame({"id": np.arange(n), "time": observed_times, "status": status})

    # Add covariates
    for j in range(n_covariates):
        data[f"X{j}"] = X[:, j]

    return data

cause_specific_cumulative_incidence

cause_specific_cumulative_incidence(
    data: DataFrame,
    time_points: Union[List[float], ndarray],
    time_col: str = "time",
    status_col: str = "status",
    cause: int = 1,
) -> DataFrame

Calculate the cause-specific cumulative incidence function at specified time points.

Parameters:

Name	Type	Description	Default
data	DataFrame	DataFrame with competing risks data.	required
time_points	list of float or array	Time points at which to calculate the cumulative incidence.	required
time_col	str	Name of the time column.	"time"
status_col	str	Name of the status column (0=censored, 1,2,...=competing events).	"status"
cause	int	The cause/event type for which to calculate the incidence.	1

Returns:

Type	Description
DataFrame	DataFrame with time points and corresponding cumulative incidence values.

Notes

The cumulative incidence function for cause j is defined as: F_j(t) = P(T <= t, cause = j)

This is the probability of experiencing the event of type j before time t.

Source code in gen_surv/competing_risks.py

def cause_specific_cumulative_incidence(
    data: pd.DataFrame,
    time_points: Union[List[float], np.ndarray],
    time_col: str = "time",
    status_col: str = "status",
    cause: int = 1,
) -> pd.DataFrame:
    """
    Calculate the cause-specific cumulative incidence function at specified time points.

    Parameters
    ----------
    data : pd.DataFrame
        DataFrame with competing risks data.
    time_points : list of float or array
        Time points at which to calculate the cumulative incidence.
    time_col : str, default="time"
        Name of the time column.
    status_col : str, default="status"
        Name of the status column (0=censored, 1,2,...=competing events).
    cause : int, default=1
        The cause/event type for which to calculate the incidence.

    Returns
    -------
    pd.DataFrame
        DataFrame with time points and corresponding cumulative incidence values.

    Notes
    -----
    The cumulative incidence function for cause j is defined as:
    F_j(t) = P(T <= t, cause = j)

    This is the probability of experiencing the event of type j before time t.
    """
    # Validate the cause value
    unique_causes = set(data[status_col].unique()) - {0}  # Exclude censoring
    if cause not in unique_causes:
        raise ParameterError(
            "cause", cause, f"not found in the data. Available causes: {unique_causes}"
        )

    sorted_data = data.sort_values(by=time_col).copy()
    times = sorted_data[time_col].to_numpy()
    status = sorted_data[status_col].to_numpy()

    unique_times, idx = np.unique(times, return_index=True)
    counts = np.diff(np.append(idx, len(times)))
    at_risk = len(times) - idx

    d_all = np.zeros_like(unique_times, dtype=int)
    d_cause = np.zeros_like(unique_times, dtype=int)
    inverse = np.repeat(np.arange(len(unique_times)), counts)
    for i, s in enumerate(status):
        if s > 0:
            d_all[inverse[i]] += 1
            if s == cause:
                d_cause[inverse[i]] += 1

    surv = 1.0
    cif_vals = np.zeros_like(unique_times, dtype=float)
    ci = 0.0
    for i, t in enumerate(unique_times):
        prev_surv = surv
        surv *= 1 - d_all[i] / at_risk[i]
        ci += prev_surv * d_cause[i] / at_risk[i]
        cif_vals[i] = ci

    result = []
    for t in time_points:
        if t <= 0:
            result.append({"time": t, "incidence": 0.0})
        elif t >= unique_times[-1]:
            result.append({"time": t, "incidence": cif_vals[-1]})
        else:
            idx = np.searchsorted(unique_times, t, side="right") - 1
            result.append({"time": t, "incidence": cif_vals[idx]})

    return pd.DataFrame(result)

competing_risks_summary

competing_risks_summary(
    data: DataFrame,
    time_col: str = "time",
    status_col: str = "status",
    covariate_cols: list[str] | None = None,
) -> dict[str, Any]

Provide a summary of a competing risks dataset.

Parameters:

Name	Type	Description	Default
data	DataFrame	DataFrame with competing risks data.	required
time_col	str	Name of the time column.	"time"
status_col	str	Name of the status column (0=censored, 1,2,...=competing events).	"status"
covariate_cols	list of str	List of covariate columns to include in the summary. If None, all columns except time_col and status_col are considered.	None

Returns:

Type	Description
Dict[str, Any]	Dictionary with summary statistics.

Examples:

>>> from gen_surv.competing_risks import gen_competing_risks, competing_risks_summary
>>>
>>> # Generate data
>>> df = gen_competing_risks(n=100, n_risks=3, seed=42)
>>>
>>> # Get summary
>>> summary = competing_risks_summary(df)
>>> print(f"Number of events by cause: {summary['events_by_cause']}")
>>> print(f"Median time to first event: {summary['median_time']}")

Source code in gen_surv/competing_risks.py

def competing_risks_summary(
    data: pd.DataFrame,
    time_col: str = "time",
    status_col: str = "status",
    covariate_cols: list[str] | None = None,
) -> dict[str, Any]:
    """
    Provide a summary of a competing risks dataset.

    Parameters
    ----------
    data : pd.DataFrame
        DataFrame with competing risks data.
    time_col : str, default="time"
        Name of the time column.
    status_col : str, default="status"
        Name of the status column (0=censored, 1,2,...=competing events).
    covariate_cols : list of str, optional
        List of covariate columns to include in the summary.
        If None, all columns except time_col and status_col are considered.

    Returns
    -------
    Dict[str, Any]
        Dictionary with summary statistics.

    Examples
    --------
    >>> from gen_surv.competing_risks import gen_competing_risks, competing_risks_summary
    >>>
    >>> # Generate data
    >>> df = gen_competing_risks(n=100, n_risks=3, seed=42)
    >>>
    >>> # Get summary
    >>> summary = competing_risks_summary(df)
    >>> print(f"Number of events by cause: {summary['events_by_cause']}")
    >>> print(f"Median time to first event: {summary['median_time']}")
    """
    # Determine covariate columns if not provided
    if covariate_cols is None:
        covariate_cols = [
            col for col in data.columns if col not in [time_col, status_col, "id"]
        ]

    # Basic counts
    n_subjects = len(data)
    n_events = (data[status_col] > 0).sum()
    n_censored = n_subjects - n_events
    censoring_rate = n_censored / n_subjects

    # Events by cause
    causes = sorted(data[data[status_col] > 0][status_col].unique())
    events_by_cause = {}
    for cause in causes:
        n_cause = (data[status_col] == cause).sum()
        events_by_cause[int(cause)] = {
            "count": int(n_cause),
            "proportion": float(n_cause / n_subjects),
            "proportion_of_events": float(n_cause / n_events) if n_events > 0 else 0,
        }

    # Time statistics
    time_stats = {
        "min": float(data[time_col].min()),
        "max": float(data[time_col].max()),
        "median": float(data[time_col].median()),
        "mean": float(data[time_col].mean()),
    }

    # Median time to each type of event
    median_time_by_cause = {}
    for cause in causes:
        cause_times = data[data[status_col] == cause][time_col]
        if not cause_times.empty:
            median_time_by_cause[int(cause)] = float(cause_times.median())

    # Covariate statistics
    covariate_stats: dict[str, dict[str, float | int | dict[str, float]]] = {}
    for col in covariate_cols:
        col_data = data[col]

        # Check if numeric
        if pd.api.types.is_numeric_dtype(col_data):
            covariate_stats[col] = {
                "mean": float(col_data.mean()),
                "median": float(col_data.median()),
                "std": float(col_data.std()),
                "min": float(col_data.min()),
                "max": float(col_data.max()),
            }
        else:
            # Categorical statistics
            value_counts = col_data.value_counts(normalize=True).to_dict()
            covariate_stats[col] = {
                "categories": len(value_counts),
                "distribution": {str(k): float(v) for k, v in value_counts.items()},
            }

    # Compile final summary
    summary = {
        "n_subjects": n_subjects,
        "n_events": n_events,
        "n_censored": n_censored,
        "censoring_rate": censoring_rate,
        "n_causes": len(causes),
        "causes": list(map(int, causes)),
        "events_by_cause": events_by_cause,
        "time_stats": time_stats,
        "median_time_by_cause": median_time_by_cause,
        "covariate_stats": covariate_stats,
    }

    return summary

plot_cause_specific_hazards

plot_cause_specific_hazards(
    data: DataFrame,
    time_points: ndarray | None = None,
    time_col: str = "time",
    status_col: str = "status",
    bandwidth: float = 0.5,
    figsize: tuple[float, float] = (10, 6),
) -> tuple[Figure, Axes]

Plot cause-specific hazard functions.

Parameters:

Name	Type	Description	Default
data	DataFrame	DataFrame with competing risks data.	required
time_points	array	Time points at which to estimate hazards. If None, uses 100 equally spaced points from 0 to max time.	None
time_col	str	Name of the time column.	"time"
status_col	str	Name of the status column (0=censored, 1,2,...=competing events).	"status"
bandwidth	float	Bandwidth for kernel density estimation.	0.5
figsize	tuple	Figure size (width, height) in inches.	(10, 6)

Returns:

Type	Description
tuple	Figure and axes objects.

Notes

This function requires matplotlib and scipy.

Source code in gen_surv/competing_risks.py

def plot_cause_specific_hazards(
    data: pd.DataFrame,
    time_points: np.ndarray | None = None,
    time_col: str = "time",
    status_col: str = "status",
    bandwidth: float = 0.5,
    figsize: tuple[float, float] = (10, 6),
) -> tuple["Figure", "Axes"]:
    """
    Plot cause-specific hazard functions.

    Parameters
    ----------
    data : pd.DataFrame
        DataFrame with competing risks data.
    time_points : array, optional
        Time points at which to estimate hazards.
        If None, uses 100 equally spaced points from 0 to max time.
    time_col : str, default="time"
        Name of the time column.
    status_col : str, default="status"
        Name of the status column (0=censored, 1,2,...=competing events).
    bandwidth : float, default=0.5
        Bandwidth for kernel density estimation.
    figsize : tuple, default=(10, 6)
        Figure size (width, height) in inches.

    Returns
    -------
    tuple
        Figure and axes objects.

    Notes
    -----
    This function requires matplotlib and scipy.
    """
    try:
        import matplotlib.pyplot as plt
        from scipy.stats import gaussian_kde
    except ImportError:
        raise ImportError(
            "This function requires matplotlib and scipy. "
            "Install them with: pip install matplotlib scipy"
        )

    # Determine time points if not provided
    if time_points is None:
        max_time = data[time_col].max()
        time_points = np.linspace(0, max_time, 100)

    # Get unique causes (excluding censoring)
    causes = sorted([c for c in data[status_col].unique() if c > 0])

    # Create plot
    fig, ax = plt.subplots(figsize=figsize)

    times_sorted = np.sort(data[time_col].to_numpy())
    total = len(times_sorted)

    # Plot hazard for each cause
    for cause in causes:
        cause_data = data[data[status_col] == cause]
        if len(cause_data) < 5:
            continue

        kde = gaussian_kde(cause_data[time_col], bw_method=bandwidth)

        at_risk = total - np.searchsorted(times_sorted, time_points, side="left")
        at_risk = np.maximum(at_risk, 1)

        hazard = kde(time_points) * total / at_risk
        ax.plot(time_points, hazard, label=f"Cause {cause}")

    # Format plot
    ax.set_xlabel("Time")
    ax.set_ylabel("Hazard Rate")
    ax.set_title("Cause-Specific Hazard Functions")
    ax.legend()
    ax.grid(alpha=0.3)

    return fig, ax

Mixture cure

mixture

Mixture Cure Models for survival data simulation.

This module provides functions to generate survival data with a cure fraction, i.e., a proportion of subjects who are immune to the event of interest.

gen_mixture_cure

gen_mixture_cure(
    n: int,
    cure_fraction: float,
    baseline_hazard: float = 0.5,
    betas_survival: list[float] | None = None,
    betas_cure: list[float] | None = None,
    n_covariates: int = 2,
    covariate_dist: Literal[
        "normal", "uniform", "binary"
    ] = "normal",
    covariate_params: dict[str, float] | None = None,
    model_cens: Literal[
        "uniform", "exponential"
    ] = "uniform",
    cens_par: float = 5.0,
    max_time: float | None = 10.0,
    seed: int | None = None,
) -> DataFrame

Generate survival data with a cure fraction using a mixture cure model.

Parameters:

Name	Type	Description	Default
n	int	Number of subjects.	required
cure_fraction	float	Baseline probability of being cured (immune to the event). Should be between 0 and 1.	required
baseline_hazard	float	Baseline hazard rate for the non-cured population.	0.5
betas_survival	list of float	Coefficients for covariates in the survival component. If None, generates random coefficients.	None
betas_cure	list of float	Coefficients for covariates in the cure component. If None, generates random coefficients.	None
n_covariates	int	Number of covariates to generate if betas is None.	2
covariate_dist	(normal, uniform, binary)	Distribution to generate covariates from.	"normal"
covariate_params	dict	Parameters for covariate distribution: - "normal": {"mean": float, "std": float} - "uniform": {"low": float, "high": float} - "binary": {"p": float} If None, uses defaults based on distribution.	None
model_cens	(uniform, exponential)	Censoring mechanism.	"uniform"
cens_par	float	Parameter for censoring distribution.	5.0
max_time	float	Maximum simulation time. Set to None for no limit.	10.0
seed	int	Random seed for reproducibility.	None

Returns:

Type	Description
DataFrame	DataFrame with columns: - "id": Subject identifier - "time": Time to event or censoring - "status": Event indicator (1=event, 0=censored) - "cured": Indicator of cure status (1=cured, 0=not cured) - "X0", "X1", ...: Covariates

Examples:

>>> from gen_surv.mixture import gen_mixture_cure
>>>
>>> # Generate data with 30% baseline cure fraction
>>> df = gen_mixture_cure(
...     n=100,
...     cure_fraction=0.3,
...     betas_survival=[0.8, -0.5],
...     betas_cure=[-0.5, 0.8],
...     seed=42
... )
>>>
>>> # Check cure proportion
>>> print(f"Cured subjects: {df['cured'].mean():.2%}")

Source code in gen_surv/mixture.py

def gen_mixture_cure(
    n: int,
    cure_fraction: float,
    baseline_hazard: float = 0.5,
    betas_survival: list[float] | None = None,
    betas_cure: list[float] | None = None,
    n_covariates: int = 2,
    covariate_dist: Literal["normal", "uniform", "binary"] = "normal",
    covariate_params: dict[str, float] | None = None,
    model_cens: Literal["uniform", "exponential"] = "uniform",
    cens_par: float = 5.0,
    max_time: float | None = 10.0,
    seed: int | None = None,
) -> pd.DataFrame:
    """
    Generate survival data with a cure fraction using a mixture cure model.

    Parameters
    ----------
    n : int
        Number of subjects.
    cure_fraction : float
        Baseline probability of being cured (immune to the event).
        Should be between 0 and 1.
    baseline_hazard : float, default=0.5
        Baseline hazard rate for the non-cured population.
    betas_survival : list of float, optional
        Coefficients for covariates in the survival component.
        If None, generates random coefficients.
    betas_cure : list of float, optional
        Coefficients for covariates in the cure component.
        If None, generates random coefficients.
    n_covariates : int, default=2
        Number of covariates to generate if betas is None.
    covariate_dist : {"normal", "uniform", "binary"}, default="normal"
        Distribution to generate covariates from.
    covariate_params : dict, optional
        Parameters for covariate distribution:
        - "normal": {"mean": float, "std": float}
        - "uniform": {"low": float, "high": float}
        - "binary": {"p": float}
        If None, uses defaults based on distribution.
    model_cens : {"uniform", "exponential"}, default="uniform"
        Censoring mechanism.
    cens_par : float, default=5.0
        Parameter for censoring distribution.
    max_time : float, optional, default=10.0
        Maximum simulation time. Set to None for no limit.
    seed : int, optional
        Random seed for reproducibility.

    Returns
    -------
    pd.DataFrame
        DataFrame with columns:
        - "id": Subject identifier
        - "time": Time to event or censoring
        - "status": Event indicator (1=event, 0=censored)
        - "cured": Indicator of cure status (1=cured, 0=not cured)
        - "X0", "X1", ...: Covariates

    Examples
    --------
    >>> from gen_surv.mixture import gen_mixture_cure
    >>>
    >>> # Generate data with 30% baseline cure fraction
    >>> df = gen_mixture_cure(
    ...     n=100,
    ...     cure_fraction=0.3,
    ...     betas_survival=[0.8, -0.5],
    ...     betas_cure=[-0.5, 0.8],
    ...     seed=42
    ... )
    >>>
    >>> # Check cure proportion
    >>> print(f"Cured subjects: {df['cured'].mean():.2%}")
    """
    rng = np.random.default_rng(seed)
    validate_gen_mixture_inputs(
        n,
        cure_fraction,
        baseline_hazard,
        n_covariates,
        model_cens,
        cens_par,
        max_time,
        covariate_dist,
    )
    covariate_params = set_covariate_params(covariate_dist, covariate_params)
    betas_survival_arr, n_covariates = prepare_betas(
        betas_survival, n_covariates, rng, name="betas_survival"
    )
    betas_cure_arr, _ = prepare_betas(
        betas_cure, n_covariates, rng, name="betas_cure", enforce_length=True
    )
    X = generate_covariates(n, n_covariates, covariate_dist, covariate_params, rng)
    lp_survival = X @ betas_survival_arr
    lp_cure = X @ betas_cure_arr
    cured = _cure_status(lp_cure, cure_fraction, rng)
    survival_times = _survival_times(cured, lp_survival, baseline_hazard, max_time, rng)
    observed_times, status = _apply_censoring(
        survival_times, model_cens, cens_par, max_time, rng
    )

    record(
        betas_survival=betas_survival_arr,
        betas_cure=betas_cure_arr,
        covariates=X,
        linear_predictor=lp_survival,
        cure_linear_predictor=lp_cure,
        cured=cured,
        event_time=survival_times,
    )

    data = pd.DataFrame(
        {"id": np.arange(n), "time": observed_times, "status": status, "cured": cured}
    )

    for j in range(n_covariates):
        data[f"X{j}"] = X[:, j]

    return data

cure_fraction_estimate

cure_fraction_estimate(
    data: DataFrame,
    time_col: str = "time",
    status_col: str = "status",
    bandwidth: float = 0.1,
) -> float

Estimate the cure fraction from observed data using non-parametric methods.

Parameters:

Name	Type	Description	Default
data	DataFrame	DataFrame with survival data.	required
time_col	str	Name of the time column.	"time"
status_col	str	Name of the status column (1=event, 0=censored).	"status"
bandwidth	float	Bandwidth parameter for smoothing the tail of the survival curve.	0.1

Returns:

Type	Description
float	Estimated cure fraction.

Notes

This function uses a non-parametric approach to estimate the cure fraction based on the plateau of the survival curve. It may not be accurate for small sample sizes or heavy censoring.

Source code in gen_surv/mixture.py

def cure_fraction_estimate(
    data: pd.DataFrame,
    time_col: str = "time",
    status_col: str = "status",
    bandwidth: float = 0.1,
) -> float:
    """
    Estimate the cure fraction from observed data using non-parametric methods.

    Parameters
    ----------
    data : pd.DataFrame
        DataFrame with survival data.
    time_col : str, default="time"
        Name of the time column.
    status_col : str, default="status"
        Name of the status column (1=event, 0=censored).
    bandwidth : float, default=0.1
        Bandwidth parameter for smoothing the tail of the survival curve.

    Returns
    -------
    float
        Estimated cure fraction.

    Notes
    -----
    This function uses a non-parametric approach to estimate the cure fraction
    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()

    # Calculate Kaplan-Meier estimate
    times = sorted_data[time_col].values
    status = sorted_data[status_col].values
    n = len(times)

    if n == 0:
        return 0.0

    # Calculate survival function
    survival = np.ones(n)

    for i in range(n):
        if i > 0:
            survival[i] = survival[i - 1]

        # Count subjects at risk at this time
        at_risk = n - i

        if status[i] == 1:  # Event
            survival[i] *= 1 - 1 / at_risk

    # Estimate cure fraction as the plateau of the survival curve
    # Use the last portion of the survival curve if enough data points
    tail_size = max(int(n * _TAIL_FRACTION), 1)
    tail_survival = survival[-tail_size:]

    # Apply smoothing if there are enough data points
    if tail_size > _SMOOTH_MIN_TAIL:
        # Use kernel smoothing
        weights = np.exp(
            -((np.arange(tail_size) - tail_size + 1) ** 2)
            / (2 * bandwidth * tail_size) ** 2
        )
        weights = weights / weights.sum()
        cure_fraction = float(np.sum(tail_survival * weights))
    else:
        # Just use the last survival probability
        cure_fraction = float(survival[-1])

    return cure_fraction

Illness-death, intervals

cmm

generate_event_times

generate_event_times(
    z1: float,
    beta: Sequence[float],
    rate: Sequence[float],
    rng: Generator | None = None,
) -> EventTimes

Generate event times for a continuous-time multi-state Markov model.

Parameters:

Name	Type	Description	Default
z1	float	Covariate value.	required
beta	Sequence[float]	List of 3 beta coefficients.	required
rate	Sequence[float]	List of 6 transition rate parameters.	required
rng	Generator	Random number generator to use. Defaults to None which creates a new generator.	None

Returns:

Type	Description
EventTimes	Dictionary with keys 't12', 't13', and 't23'.

Examples:

>>> from gen_surv.cmm import generate_event_times
>>> ev = generate_event_times(0.2, [0.1, -0.2, 0.3],
...                          [0.5, 1.0, 0.7, 1.2, 0.4, 1.5])
>>> sorted(ev.keys())
['t12', 't13', 't23']

Source code in gen_surv/cmm.py

def generate_event_times(
    z1: float,
    beta: Sequence[float],
    rate: Sequence[float],
    rng: np.random.Generator | None = None,
) -> EventTimes:
    """Generate event times for a continuous-time multi-state Markov model.

    Parameters
    ----------
    z1 : float
        Covariate value.
    beta : Sequence[float]
        List of 3 beta coefficients.
    rate : Sequence[float]
        List of 6 transition rate parameters.
    rng : np.random.Generator, optional
        Random number generator to use. Defaults to ``None`` which creates a new generator.

    Returns
    -------
    EventTimes
        Dictionary with keys ``'t12'``, ``'t13'``, and ``'t23'``.

    Examples
    --------
    >>> from gen_surv.cmm import generate_event_times
    >>> ev = generate_event_times(0.2, [0.1, -0.2, 0.3],
    ...                          [0.5, 1.0, 0.7, 1.2, 0.4, 1.5])
    >>> sorted(ev.keys())
    ['t12', 't13', 't23']
    """
    rng = np.random.default_rng() if rng is None else rng

    u = rng.uniform(size=3)
    rate_arr = np.asarray(rate).reshape(3, 2)
    beta_arr = np.asarray(beta)
    t = (-np.log(1 - u) / (rate_arr[:, 0] * np.exp(beta_arr * z1))) ** (
        1 / rate_arr[:, 1]
    )

    return {"t12": float(t[0]), "t13": float(t[1]), "t23": float(t[2])}

gen_cmm

gen_cmm(
    n: int,
    model_cens: str,
    cens_par: float,
    beta: Sequence[float],
    covariate_range: float,
    rate: Sequence[float],
    seed: RandomStateLike = None,
) -> DataFrame

Generate survival data using a continuous-time Markov model (CMM).

Parameters:

Name	Type	Description	Default
n	int	Number of individuals.	required
model_cens	str	"uniform" or "exponential".	required
cens_par	float	Parameter for censoring.	required
beta	Sequence[float]	Regression coefficients (length 3).	required
covariate_range	float	Upper bound for the covariate values.	required
rate	Sequence[float]	Transition rates (length 6).	required
seed	int	Random seed for reproducibility.	None

Returns:

Type

Description

DataFrame

Counting-process records with columns id, start, stop, from_state, to_state, status, X0, sorted by id, start then to_state. States are 1 (healthy), 2 (illness) and 3 (death). While a subject occupies state 1 it is simultaneously at risk of 1 -> 2 and 1 -> 3, so it contributes one row for each, both ending when it leaves state 1; status is 1 on the transition that occurred and 0 on the competing one. A subject that reaches state 2 contributes a further 2 -> 3 row. Subjects therefore contribute two or three rows each, not one.

Notes

Sojourn times are drawn on a reset clock, so the model is semi-Markov: the 2 -> 3 row spans t12 to t12 + t23 where t23 is an independent draw. This matches genCMM in the R package.

Examples:

>>> from gen_surv.cmm import gen_cmm
>>> df = gen_cmm(
...     n=50,
...     model_cens="uniform",
...     cens_par=2.0,
...     beta=[0.3, -0.2, 0.1],
...     covariate_range=1.0,
...     rate=[0.1, 1.0, 0.2, 1.2, 0.3, 1.5],
...     seed=42,
... )
>>> list(df.columns)
['id', 'start', 'stop', 'from_state', 'to_state', 'status', 'X0']

Source code in gen_surv/cmm.py

def gen_cmm(
    n: int,
    model_cens: str,
    cens_par: float,
    beta: Sequence[float],
    covariate_range: float,
    rate: Sequence[float],
    seed: RandomStateLike = None,
) -> pd.DataFrame:
    """Generate survival data using a continuous-time Markov model (CMM).

    Parameters
    ----------
    n : int
        Number of individuals.
    model_cens : str
        ``"uniform"`` or ``"exponential"``.
    cens_par : float
        Parameter for censoring.
    beta : Sequence[float]
        Regression coefficients (length 3).
    covariate_range : float
        Upper bound for the covariate values.
    rate : Sequence[float]
        Transition rates (length 6).
    seed : int, optional
        Random seed for reproducibility.

    Returns
    -------
    pd.DataFrame
        Counting-process records with columns ``id``, ``start``, ``stop``,
        ``from_state``, ``to_state``, ``status``, ``X0``, sorted by ``id``,
        ``start`` then ``to_state``.

        States are 1 (healthy), 2 (illness) and 3 (death). While a subject
        occupies state 1 it is simultaneously at risk of ``1 -> 2`` and
        ``1 -> 3``, so it contributes one row for each, both ending when it
        leaves state 1; ``status`` is 1 on the transition that occurred and 0 on
        the competing one. A subject that reaches state 2 contributes a further
        ``2 -> 3`` row. Subjects therefore contribute two or three rows each,
        not one.

    Notes
    -----
    Sojourn times are drawn on a reset clock, so the model is semi-Markov: the
    ``2 -> 3`` row spans ``t12`` to ``t12 + t23`` where ``t23`` is an
    independent draw. This matches ``genCMM`` in the R package.

    Examples
    --------
    >>> from gen_surv.cmm import gen_cmm
    >>> df = gen_cmm(
    ...     n=50,
    ...     model_cens="uniform",
    ...     cens_par=2.0,
    ...     beta=[0.3, -0.2, 0.1],
    ...     covariate_range=1.0,
    ...     rate=[0.1, 1.0, 0.2, 1.2, 0.3, 1.5],
    ...     seed=42,
    ... )
    >>> list(df.columns)
    ['id', 'start', 'stop', 'from_state', 'to_state', 'status', 'X0']
    """
    validate_gen_cmm_inputs(n, model_cens, cens_par, beta, covariate_range, rate)

    # `rate` is three (intensity, shape) pairs. The sojourn is drawn from
    # H(t) = lambda * t**rho, which is a Weibull cumulative hazard
    # (t / scale) ** shape with shape = rho and scale = lambda ** (-1 / rho).
    transitions = [
        Transition(
            origin,
            destination,
            WeibullBaseline(shape=shape, scale=float(intensity) ** (-1.0 / shape)),
            [float(coefficient)],
        )
        for (origin, destination), intensity, shape, coefficient in (
            ((1, 2), rate[0], float(rate[1]), beta[0]),
            ((1, 3), rate[2], float(rate[3]), beta[1]),
            ((2, 3), rate[4], float(rate[5]), beta[2]),
        )
    ]

    # A reset clock: the 2 -> 3 sojourn is measured from entry to state 2, which
    # is what makes this semi-Markov and what `genCMM` does in the R package.
    data = gen_multistate(
        n=n,
        transitions=transitions,
        clock="reset",
        initial_state=1,
        covariate_dist="uniform",
        covariate_params={"low": 0.0, "high": float(covariate_range)},
        # Validated above against the same two choices the engine
        # accepts; the cast tells the type checker what that check
        # already guarantees.
        model_cens=cast(Literal["uniform", "exponential"], model_cens),
        cens_par=cens_par,
        layout="intervals",
        seed=seed,
    )

    _record_transition_times(beta, rate)
    return data[_COLUMNS]

Illness-death, panel

thmm

calculate_transitions

calculate_transitions(
    z1: float,
    cens_par: float,
    beta: Sequence[float],
    rate: Sequence[float],
    rfunc: CensoringFunc,
    seed: RandomStateLike = None,
) -> TransitionTimes

Calculate transition and censoring times for THMM.

Parameters: - z1 (float): Covariate value. - cens_par (float): Censoring parameter. - beta (list of float): Coefficients for rate modification (length 3). - rate (list of float): Base rates (length 3). - rfunc (callable): Censoring function, e.g. runifcens or rexpocens. - seed (int, Generator or None): Seed or generator for reproducibility.

Returns: - dict with keys 'c', 't12', 't13', 't23'

Source code in gen_surv/thmm.py

def calculate_transitions(
    z1: float,
    cens_par: float,
    beta: Sequence[float],
    rate: Sequence[float],
    rfunc: CensoringFunc,
    seed: RandomStateLike = None,
) -> TransitionTimes:
    """
    Calculate transition and censoring times for THMM.

    Parameters:
    - z1 (float): Covariate value.
    - cens_par (float): Censoring parameter.
    - beta (list of float): Coefficients for rate modification (length 3).
    - rate (list of float): Base rates (length 3).
    - rfunc (callable): Censoring function, e.g. runifcens or rexpocens.
    - seed (int, Generator or None): Seed or generator for reproducibility.

    Returns:
    - dict with keys 'c', 't12', 't13', 't23'
    """
    rng = resolve_rng(seed)

    c = rfunc(1, cens_par, rng)[0]
    rate12 = rate[0] * np.exp(beta[0] * z1)
    rate13 = rate[1] * np.exp(beta[1] * z1)
    rate23 = rate[2] * np.exp(beta[2] * z1)

    t12 = rng.exponential(scale=1 / rate12)
    t13 = rng.exponential(scale=1 / rate13)
    t23 = rng.exponential(scale=1 / rate23)

    return {"c": c, "t12": t12, "t13": t13, "t23": t23}

gen_thmm

gen_thmm(
    n: int,
    model_cens: str,
    cens_par: float,
    beta: Sequence[float],
    covariate_range: float,
    rate: Sequence[float],
    seed: RandomStateLike = None,
) -> DataFrame

Generate THMM (Time-Homogeneous Markov Model) survival data.

Parameters:

Name	Type	Description	Default
n	int	Number of individuals.	required
model_cens	(uniform, exponential)	Censoring model.	"uniform"
cens_par	float	Censoring parameter.	required
beta	Sequence[float]	Length-3 regression coefficients.	required
covariate_range	float	Upper bound for the covariate values.	required
rate	Sequence[float]	Length-3 transition rates.	required
seed	int or Generator	Seed or generator for reproducibility.	None

Returns:

Type

Description

DataFrame

Columns = ["id", "time", "state", "X0"], one row per observation time giving the state occupied at that time. States are 1 (healthy), 2 (illness) and 3 (death). Every subject starts with an observation in state 1 at time 0, then contributes one or two further observations, so subjects yield two or three rows each rather than one. A subject still in state 1 or 2 when censoring occurs has a final observation in that state at the censoring time.

Notes

This panel layout -- a state recorded at each observation time -- matches genTHMM in the R package, and differs deliberately from :func:gen_surv.cmm.gen_cmm, which emits counting-process intervals.

All transition intensities are constant in time, so sojourn times are exponential and the reset and forward clocks coincide.

Examples:

>>> from gen_surv.thmm import gen_thmm
>>> df = gen_thmm(
...     n=3,
...     model_cens="uniform",
...     cens_par=5.0,
...     beta=[0.1, 0.2, 0.3],
...     covariate_range=1.0,
...     rate=[0.1, 0.1, 0.2],
...     seed=42,
... )

Source code in gen_surv/thmm.py

def gen_thmm(
    n: int,
    model_cens: str,
    cens_par: float,
    beta: Sequence[float],
    covariate_range: float,
    rate: Sequence[float],
    seed: RandomStateLike = None,
) -> pd.DataFrame:
    """Generate THMM (Time-Homogeneous Markov Model) survival data.

    Parameters
    ----------
    n : int
        Number of individuals.
    model_cens : {"uniform", "exponential"}
        Censoring model.
    cens_par : float
        Censoring parameter.
    beta : Sequence[float]
        Length-3 regression coefficients.
    covariate_range : float
        Upper bound for the covariate values.
    rate : Sequence[float]
        Length-3 transition rates.
    seed : int or numpy.random.Generator, optional
        Seed or generator for reproducibility.

    Returns
    -------
    pd.DataFrame
        Columns = ``["id", "time", "state", "X0"]``, one row per observation
        time giving the state occupied at that time.

        States are 1 (healthy), 2 (illness) and 3 (death). Every subject starts
        with an observation in state 1 at time 0, then contributes one or two
        further observations, so subjects yield two or three rows each rather
        than one. A subject still in state 1 or 2 when censoring occurs has a
        final observation in that state at the censoring time.

    Notes
    -----
    This panel layout -- a state recorded at each observation time -- matches
    ``genTHMM`` in the R package, and differs deliberately from
    :func:`gen_surv.cmm.gen_cmm`, which emits counting-process intervals.

    All transition intensities are constant in time, so sojourn times are
    exponential and the reset and forward clocks coincide.

    Examples
    --------
    >>> from gen_surv.thmm import gen_thmm
    >>> df = gen_thmm(
    ...     n=3,
    ...     model_cens="uniform",
    ...     cens_par=5.0,
    ...     beta=[0.1, 0.2, 0.3],
    ...     covariate_range=1.0,
    ...     rate=[0.1, 0.1, 0.2],
    ...     seed=42,
    ... )
    """
    validate_gen_thmm_inputs(n, model_cens, cens_par, beta, covariate_range, rate)

    # Every intensity is constant in time, which is what "time-homogeneous"
    # means, so the baseline is exponential and the clock makes no difference:
    # a constant hazard is memoryless.
    transitions = [
        Transition(
            origin,
            destination,
            ExponentialBaseline(rate=float(intensity)),
            [float(coefficient)],
        )
        for (origin, destination), intensity, coefficient in (
            ((1, 2), rate[0], beta[0]),
            ((1, 3), rate[1], beta[1]),
            ((2, 3), rate[2], beta[2]),
        )
    ]

    data = gen_multistate(
        n=n,
        transitions=transitions,
        clock="forward",
        initial_state=1,
        covariate_dist="uniform",
        covariate_params={"low": 0.0, "high": float(covariate_range)},
        # Validated above against the same two choices the engine
        # accepts; the cast tells the type checker what that check
        # already guarantees.
        model_cens=cast(Literal["uniform", "exponential"], model_cens),
        cens_par=cens_par,
        layout="panel",
        seed=seed,
    )

    # This model has numbered its subjects from 1 since it was ported.
    data["id"] = data["id"] + 1
    _record_transition_times(beta, rate)
    return data[["id", "time", "state", "X0"]]

Time-dependent covariates

tdcm

generate_censored_observations

generate_censored_observations(
    n: int,
    dist_par: Sequence[float],
    model_cens: str,
    cens_par: float,
    beta: Sequence[float],
    lam: float,
    b: NDArray[float64],
    seed: RandomStateLike = None,
) -> NDArray[float64]

Generate censored TDCM observations.

Parameters:

Name	Type	Description	Default
n	int	Number of individuals.	required
dist_par	Sequence[float]	Not directly used here (kept for API compatibility).	required
model_cens	(uniform, exponential)	Censoring model.	"uniform"
cens_par	float	Parameter for the censoring model.	required
beta	Sequence[float]	Length-2 list of regression coefficients.	required
lam	float	Rate parameter.	required
b	NDArray[float64]	Covariate matrix with two columns [., z1].	required
seed	int or Generator	Seed or generator for reproducibility.	None

Returns:

Type	Description
NDArray[float64]	Array of shape (n, 6) with columns [id, start, stop, status, covariate1 (z1), covariate2 (z2)].

Source code in gen_surv/tdcm.py

def generate_censored_observations(
    n: int,
    dist_par: Sequence[float],
    model_cens: str,
    cens_par: float,
    beta: Sequence[float],
    lam: float,
    b: NDArray[np.float64],
    seed: RandomStateLike = None,
) -> NDArray[np.float64]:
    """Generate censored TDCM observations.

    Parameters
    ----------
    n : int
        Number of individuals.
    dist_par : Sequence[float]
        Not directly used here (kept for API compatibility).
    model_cens : {"uniform", "exponential"}
        Censoring model.
    cens_par : float
        Parameter for the censoring model.
    beta : Sequence[float]
        Length-2 list of regression coefficients.
    lam : float
        Rate parameter.
    b : NDArray[np.float64]
        Covariate matrix with two columns ``[., z1]``.
    seed : int or numpy.random.Generator, optional
        Seed or generator for reproducibility.

    Returns
    -------
    NDArray[np.float64]
        Array of shape ``(n, 6)`` with columns
        ``[id, start, stop, status, covariate1 (z1), covariate2 (z2)]``.
    """
    rfunc: CensoringFunc = runifcens if model_cens == "uniform" else rexpocens
    rng = resolve_rng(seed)

    z1 = b[:, 1]

    # A heavy-tailed covariate distribution -- a Weibull `dist_par` shape well
    # below 1 -- puts `z1` in the tens of thousands, and `exp(beta[0] * z1)`
    # then leaves the range of a float. Neither outcome is survivable:
    # overflow to `inf` makes `t1 = log_term / inf` exactly 0.0, which
    # `status = (t <= c)` reports as an *observed event at time zero* in a
    # zero-length risk interval; underflow to 0.0 makes `t` infinite, silently
    # reported as censored. Both are frames of the right shape carrying data no
    # analysis should be handed, so this raises rather than returning one.
    # Errors are suppressed only because the result is inspected immediately
    # below and turned into a message that says what to change; NumPy's warning
    # names the expression, not the parameter behind it.
    with np.errstate(over="ignore", under="ignore"):
        exp_b0_z1 = np.exp(beta[0] * z1)
    if not np.all(np.isfinite(exp_b0_z1)) or np.any(exp_b0_z1 == 0.0):
        extreme = float(z1[np.argmax(np.abs(beta[0] * z1))])
        raise ParameterError(
            "beta",
            beta,
            f"combined with the covariate distribution, exp(beta[0] * z) left "
            f"the range of a float (largest covariate drawn: {extreme:.3g}). "
            f"Reduce beta[0], or widen the Weibull shape in dist_par -- a "
            f"shape below 1 is heavy-tailed and draws very large covariates",
        )

    x = lam * b[:, 0] * exp_b0_z1
    u = rng.uniform(size=n)
    c = rfunc(n, cens_par, rng)

    threshold = 1 - np.exp(-x)
    log_term = -np.log(1 - u)

    # Before the crossover the hazard is lam * exp(beta[0] * z1); after it, that
    # times exp(beta[1]). Inverting the cumulative hazard on each side:
    #
    #   before:  t = L / A
    #   after:   t = tau + (L - x) / (A * exp(beta[1]))
    #
    # with A = lam * exp(beta[0] * z1), tau = x / A the crossover time, x the
    # cumulative hazard accrued by then, and L = -log(1 - u). Expanding the
    # second gives the closed form below. Releases up to 2.0.2 had the sign of
    # the x term reversed, which placed "after the crossover" draws *before* it
    # and, for large beta[1], produced negative survival times.
    t1 = log_term / (lam * exp_b0_z1)
    t2 = (log_term + x * (np.exp(beta[1]) - 1)) / (lam * np.exp(beta[0] * z1 + beta[1]))
    mask = u < threshold
    t = np.where(mask, t1, t2)

    # The covariate's value over the interval actually observed: a subject
    # censored before its crossover never switched, whatever its latent event
    # time would have done.
    crossover = b[:, 0]
    z2 = (crossover <= np.minimum(t, c)).astype(float)

    time = np.minimum(t, c)
    status = (t <= c).astype(float)

    # The crossover time is what the returned frame cannot express: it records
    # only the covariate's value at exit, so a caller cannot split the risk
    # interval without this.
    record(
        beta=np.asarray(beta, dtype=float),
        covariates=z1,
        crossover_time=b[:, 0],
        event_time=t,
        censoring_time=c,
        switched_before_exit=z2,
    )

    ids = np.arange(1, n + 1, dtype=float)
    zeros = np.zeros(n, dtype=float)
    return np.column_stack((ids, zeros, time, status, z1, z2))

gen_tdcm

gen_tdcm(
    n: int,
    dist: str,
    corr: float,
    dist_par: Sequence[float],
    model_cens: str,
    cens_par: float,
    beta: Sequence[float],
    lam: float,
    seed: RandomStateLike = None,
) -> DataFrame

Generate TDCM (Time-Dependent Covariate Model) survival data.

Parameters:

Name	Type	Description	Default
n	int	Number of individuals.	required
dist	(weibull, exponential)	Type of marginal distributions.	"weibull"
corr	float	Correlation between the baseline covariate and the crossover time, on the latent normal scale. Must be in (0, 1) for dist='weibull' and (-1, 1) for dist='exponential'; the endpoints make the copula's covariance singular.	required
dist_par	Sequence[float]	Distribution parameters.	required
model_cens	(uniform, exponential)	Censoring model.	"uniform"
cens_par	float	Censoring parameter.	required
beta	Sequence[float]	Length-2 regression coefficients: the baseline covariate effect and the effect of the time-dependent covariate.	required
lam	float	Lambda rate parameter.	required
seed	int or Generator	Seed or generator for reproducibility.	None

Returns:

Type	Description
DataFrame	Columns are ["id", "start", "stop", "status", "covariate", "tdcov"].

Examples:

>>> from gen_surv.tdcm import gen_tdcm
>>> df = gen_tdcm(
...     n=5,
...     dist="exponential",
...     corr=0.3,
...     dist_par=[0.5, 1.0],
...     model_cens="uniform",
...     cens_par=2.0,
...     beta=[0.1, 0.2],
...     lam=0.5,
...     seed=42,
... )

Source code in gen_surv/tdcm.py

def gen_tdcm(
    n: int,
    dist: str,
    corr: float,
    dist_par: Sequence[float],
    model_cens: str,
    cens_par: float,
    beta: Sequence[float],
    lam: float,
    seed: RandomStateLike = None,
) -> pd.DataFrame:
    """Generate TDCM (Time-Dependent Covariate Model) survival data.

    Parameters
    ----------
    n : int
        Number of individuals.
    dist : {"weibull", "exponential"}
        Type of marginal distributions.
    corr : float
        Correlation between the baseline covariate and the crossover time, on
        the latent normal scale. Must be in ``(0, 1)`` for ``dist='weibull'``
        and ``(-1, 1)`` for ``dist='exponential'``; the endpoints make the
        copula's covariance singular.
    dist_par : Sequence[float]
        Distribution parameters.
    model_cens : {"uniform", "exponential"}
        Censoring model.
    cens_par : float
        Censoring parameter.
    beta : Sequence[float]
        Length-2 regression coefficients: the baseline covariate effect and the
        effect of the time-dependent covariate.
    lam : float
        Lambda rate parameter.
    seed : int or numpy.random.Generator, optional
        Seed or generator for reproducibility.

    Returns
    -------
    pd.DataFrame
        Columns are ``["id", "start", "stop", "status", "covariate", "tdcov"]``.

    Examples
    --------
    >>> from gen_surv.tdcm import gen_tdcm
    >>> df = gen_tdcm(
    ...     n=5,
    ...     dist="exponential",
    ...     corr=0.3,
    ...     dist_par=[0.5, 1.0],
    ...     model_cens="uniform",
    ...     cens_par=2.0,
    ...     beta=[0.1, 0.2],
    ...     lam=0.5,
    ...     seed=42,
    ... )
    """
    validate_gen_tdcm_inputs(n, dist, corr, dist_par, model_cens, cens_par, beta, lam)

    # One generator shared by both stages, so a single seed reproduces the
    # covariates and the event/censoring times together.
    rng = resolve_rng(seed)

    # Generate covariate matrix from bivariate distribution
    b = sample_bivariate_distribution(n, dist, corr, dist_par, rng)

    data = generate_censored_observations(
        n, dist_par, model_cens, cens_par, beta, lam, b, rng
    )

    return pd.DataFrame(
        data, columns=["id", "start", "stop", "status", "covariate", "tdcov"]
    )

Recurrent events

recurrent

Recurrent event data generation.

Subjects may experience the same event repeatedly during follow-up. The three processes here correspond to the models the data is usually analysed with:

ag Andersen-Gill. The intensity depends on the covariates but not on how many events have already happened, and the clock runs forward from entry. A non-homogeneous Poisson process. pwp_tt Prentice-Williams-Peterson in total time. As Andersen-Gill, but the intensity is scaled by a factor specific to the event number, so the risk of a second event may differ from the risk of a first. The clock still runs forward from entry. pwp_gt Prentice-Williams-Peterson in gap time. As pwp_tt, but the clock resets after every event, so the baseline hazard is a function of time since the previous event rather than time since entry.

All three return counting-process intervals, the canonical layout for transition data in this package.

gen_recurrent_events

gen_recurrent_events(
    n: int,
    process: Process = "ag",
    baseline: Baseline | BaselineHazard = "exponential",
    baseline_params: dict[str, float] | None = None,
    betas: Sequence[float] | None = None,
    n_covariates: int = 2,
    covariate_dist: Literal[
        "normal", "uniform", "binary"
    ] = "normal",
    covariate_params: dict[str, float] | None = None,
    stratum_effects: Sequence[float] | None = None,
    max_events: int | None = None,
    followup_time: float = 10.0,
    model_cens: Literal[
        "uniform", "exponential"
    ] = "uniform",
    cens_par: float = 20.0,
    seed: RandomStateLike = None,
) -> DataFrame

Generate recurrent event data in counting-process form.

Parameters:

Name	Type	Description	Default
n	int	Number of subjects. Each contributes one row per at-risk interval, so the frame is longer than n.	required
process	(ag, pwp_tt, pwp_gt)	Event process. ag is Andersen-Gill, whose intensity ignores the event history. pwp_tt and pwp_gt are Prentice-Williams-Peterson in total and gap time, whose intensity is scaled per event number by stratum_effects; pwp_gt additionally resets the clock after each event.	"ag"
baseline	(exponential, weibull, gompertz)	Baseline hazard family. Exponential is constant, Weibull is monotone, Gompertz is exponentially increasing or decreasing.	"exponential"
baseline_params	dict[str, float]	Parameters of the baseline. {"rate"} for exponential, {"shape", "scale"} for Weibull, {"rate", "shape"} for Gompertz. Defaults are filled in when omitted.	None
betas	Sequence[float]	Coefficients acting on the log intensity, one per covariate. Drawn at random when omitted, which is convenient for a smoke test and unusable for validation.	None
n_covariates	int	Number of covariates when betas is not supplied.	2
covariate_dist	(normal, uniform, binary)	Distribution the covariates are drawn from.	"normal"
covariate_params	dict[str, float]	Parameters of that distribution. Defaults are filled in when omitted.	None
stratum_effects	Sequence[float]	Multiplicative intensity factors by event number, for the two PWP processes. The final entry applies to all later events, so [1.0, 2.0] means "first event at the baseline rate, every subsequent one at twice that". Supplying it with process="ag" raises, because an Andersen-Gill intensity cannot depend on the event number.	None
max_events	int	Stop following a subject once it has this many events. None places no cap.	None
followup_time	float	Administrative end of follow-up, applied to every subject.	10.0
model_cens	(uniform, exponential)	Random dropout mechanism, applied on top of followup_time.	"uniform"
cens_par	float	Parameter of the dropout distribution: the upper bound for uniform, the mean for exponential.	20.0
seed	int or Generator	Seed or generator for reproducibility.	None

Returns:

Type	Description
DataFrame	Counting-process intervals with columns ["id", "start", "stop", "status", "enum", "X0", ..., "Xp"]. Each row is the interval over which a subject was at risk of its enum-th event; status is 1 if that event occurred at stop and 0 if follow-up ended first. Every subject contributes at least one row, and intervals are contiguous within a subject.

Raises:

Type	Description
ValidationError	If any parameter is outside its allowed range.

Examples:

>>> from gen_surv.recurrent import gen_recurrent_events
>>> df = gen_recurrent_events(
...     n=50,
...     process="ag",
...     baseline_params={"rate": 0.5},
...     betas=[0.4, -0.2],
...     followup_time=5.0,
...     seed=42,
... )
>>> list(df.columns)
['id', 'start', 'stop', 'status', 'enum', 'X0', 'X1']

Source code in gen_surv/recurrent.py

def gen_recurrent_events(
    n: int,
    process: Process = "ag",
    baseline: Baseline | BaselineHazard = "exponential",
    baseline_params: dict[str, float] | None = None,
    betas: Sequence[float] | None = None,
    n_covariates: int = 2,
    covariate_dist: Literal["normal", "uniform", "binary"] = "normal",
    covariate_params: dict[str, float] | None = None,
    stratum_effects: Sequence[float] | None = None,
    max_events: int | None = None,
    followup_time: float = 10.0,
    model_cens: Literal["uniform", "exponential"] = "uniform",
    cens_par: float = 20.0,
    seed: RandomStateLike = None,
) -> pd.DataFrame:
    """Generate recurrent event data in counting-process form.

    Parameters
    ----------
    n : int
        Number of subjects. Each contributes one row per at-risk interval, so
        the frame is longer than ``n``.
    process : {"ag", "pwp_tt", "pwp_gt"}
        Event process. ``ag`` is Andersen-Gill, whose intensity ignores the
        event history. ``pwp_tt`` and ``pwp_gt`` are Prentice-Williams-Peterson
        in total and gap time, whose intensity is scaled per event number by
        ``stratum_effects``; ``pwp_gt`` additionally resets the clock after each
        event.
    baseline : {"exponential", "weibull", "gompertz"}
        Baseline hazard family. Exponential is constant, Weibull is monotone,
        Gompertz is exponentially increasing or decreasing.
    baseline_params : dict[str, float], optional
        Parameters of the baseline. ``{"rate"}`` for exponential,
        ``{"shape", "scale"}`` for Weibull, ``{"rate", "shape"}`` for Gompertz.
        Defaults are filled in when omitted.
    betas : Sequence[float], optional
        Coefficients acting on the log intensity, one per covariate. Drawn at
        random when omitted, which is convenient for a smoke test and unusable
        for validation.
    n_covariates : int
        Number of covariates when ``betas`` is not supplied.
    covariate_dist : {"normal", "uniform", "binary"}
        Distribution the covariates are drawn from.
    covariate_params : dict[str, float], optional
        Parameters of that distribution. Defaults are filled in when omitted.
    stratum_effects : Sequence[float], optional
        Multiplicative intensity factors by event number, for the two PWP
        processes. The final entry applies to all later events, so
        ``[1.0, 2.0]`` means "first event at the baseline rate, every subsequent
        one at twice that". Supplying it with ``process="ag"`` raises, because
        an Andersen-Gill intensity cannot depend on the event number.
    max_events : int, optional
        Stop following a subject once it has this many events. ``None`` places
        no cap.
    followup_time : float
        Administrative end of follow-up, applied to every subject.
    model_cens : {"uniform", "exponential"}
        Random dropout mechanism, applied on top of ``followup_time``.
    cens_par : float
        Parameter of the dropout distribution: the upper bound for ``uniform``,
        the mean for ``exponential``.
    seed : int or numpy.random.Generator, optional
        Seed or generator for reproducibility.

    Returns
    -------
    pd.DataFrame
        Counting-process intervals with columns ``["id", "start", "stop",
        "status", "enum", "X0", ..., "Xp"]``. Each row is the interval over
        which a subject was at risk of its ``enum``-th event; ``status`` is 1 if
        that event occurred at ``stop`` and 0 if follow-up ended first. Every
        subject contributes at least one row, and intervals are contiguous
        within a subject.

    Raises
    ------
    ValidationError
        If any parameter is outside its allowed range.

    Examples
    --------
    >>> from gen_surv.recurrent import gen_recurrent_events
    >>> df = gen_recurrent_events(
    ...     n=50,
    ...     process="ag",
    ...     baseline_params={"rate": 0.5},
    ...     betas=[0.4, -0.2],
    ...     followup_time=5.0,
    ...     seed=42,
    ... )
    >>> list(df.columns)
    ['id', 'start', 'stop', 'status', 'enum', 'X0', 'X1']
    """
    validate_gen_recurrent_events_inputs(
        n=n,
        process=process,
        baseline=baseline,
        baseline_params=baseline_params,
        n_covariates=n_covariates,
        stratum_effects=stratum_effects,
        max_events=max_events,
        followup_time=followup_time,
        model_cens=model_cens,
        cens_par=cens_par,
    )

    rng = resolve_rng(seed)

    resolved_baseline = _resolve_baseline(baseline, baseline_params)

    covariate_params = set_covariate_params(covariate_dist, covariate_params)
    coefficients, n_covariates = prepare_betas(betas, n_covariates, rng, name="betas")
    covariates: NDArray[np.float64] = generate_covariates(
        n, n_covariates, covariate_dist, covariate_params, rng
    )
    eta = covariates @ coefficients

    rfunc: CensoringFunc = runifcens if model_cens == "uniform" else rexpocens
    dropout = rfunc(n, cens_par, rng)
    ends = np.minimum(dropout, followup_time)

    # The event history of one subject depends on its own previous draws, so
    # this loop is over subjects rather than over a vectorised array.
    records: list[tuple[int, float, float, int, int]] = []
    for subject in range(n):
        records.extend(
            _subject_rows(
                subject=subject,
                eta=float(eta[subject]),
                end=float(ends[subject]),
                process=process,
                baseline=resolved_baseline,
                stratum_effects=stratum_effects,
                max_events=max_events,
                rng=rng,
            )
        )

    record(
        betas=coefficients,
        covariates=covariates,
        linear_predictor=eta,
        baseline=resolved_baseline,
        followup_end=ends,
        dropout_time=dropout,
    )

    data = pd.DataFrame(records, columns=_COLUMNS)
    data = data.astype(
        {
            "id": "int64",
            "start": "float64",
            "stop": "float64",
            "status": "int64",
            "enum": "int64",
        }
    )

    for j in range(n_covariates):
        data[f"X{j}"] = covariates[data["id"].to_numpy(), j]

    return data.reset_index(drop=True)

The multistate engine

multistate

A general multistate engine.

A subject moves through a graph of states. Each edge carries its own baseline hazard and its own coefficients, so the intensity of the i -> j transition is

.. math::

\alpha_{ij}(t \mid X) = h_{0,ij}(t)\exp(X^\top\beta_{ij}).

Two clocks are supported, and the choice is what separates a Markov process from a semi-Markov one:

clock="forward" The hazard is a function of time since entry to the study. The process is Markov: where a subject has been does not matter, only where it is and how long the study has run. clock="reset" The hazard restarts at each entry to a state, so it is a function of time in the current state. The process is semi-Markov.

With an exponential baseline the two coincide, because a constant hazard is memoryless.

Both canonical layouts are available. layout="intervals" gives counting-process rows -- one per transition a subject was at risk of, over the interval it was at risk -- and layout="panel" gives one row per observation of the subject's state. See :doc:the output schemas page </getting-started/schemas>.

Examples:

An illness-death process with Weibull sojourns:

>>> from gen_surv import Transition, WeibullBaseline, gen_multistate
>>> transitions = [
...     Transition(1, 2, WeibullBaseline(shape=1.0, scale=3.0), [0.3]),
...     Transition(1, 3, WeibullBaseline(shape=1.0, scale=5.0), [0.1]),
...     Transition(2, 3, WeibullBaseline(shape=1.2, scale=2.0), [0.2]),
... ]
>>> frame = gen_multistate(n=100, transitions=transitions, clock="reset", seed=1)
>>> list(frame.columns)
['id', 'start', 'stop', 'from_state', 'to_state', 'status', 'X0']

Transition dataclass

Transition(
    origin: int,
    destination: int,
    baseline: BaselineHazard,
    coefficients: Sequence[float] = tuple(),
)

One edge of the transition graph.

Parameters:

Name	Type	Description	Default
origin	int	The state a subject moves from.	required
destination	int	The state it moves to. Must differ from origin.	required
baseline	BaselineHazard	The baseline hazard for this transition. Any object implementing the protocol works, so the shape is a parameter rather than a fork in the code.	required
coefficients	Sequence[float]	One coefficient per covariate, acting on the log intensity. Empty means the transition does not depend on the covariates.	tuple()

gen_multistate

gen_multistate(
    n: int,
    transitions: Sequence[Transition],
    clock: Clock = "forward",
    initial_state: int = 1,
    covariate_dist: Literal[
        "normal", "uniform", "binary"
    ] = "normal",
    covariate_params: dict[str, float] | None = None,
    model_cens: Literal[
        "uniform", "exponential"
    ] = "uniform",
    cens_par: float = 5.0,
    max_time: float | None = None,
    layout: Layout = "intervals",
    seed: RandomStateLike = None,
) -> DataFrame

Simulate a multistate process over an arbitrary transition graph.

Parameters:

Name	Type	Description	Default
n	int	Number of subjects. Each contributes several rows, so the frame is longer than n.	required
transitions	Sequence[Transition]	The graph. Every edge carries its own baseline hazard and coefficients. A state with no outgoing transition is absorbing.	required
clock	('forward', 'reset')	"forward" measures the hazard from entry to the study, giving a Markov process; "reset" measures it from entry to the current state, giving a semi-Markov one. They coincide for an exponential baseline.	"forward"
initial_state	int	The state every subject starts in, at time zero.	1
covariate_dist	('normal', 'uniform', 'binary')	Distribution the covariates are drawn from.	"normal"
covariate_params	dict[str, float]	Parameters of that distribution; defaults are filled in.	None
model_cens	('uniform', 'exponential')	Random censoring mechanism.	"uniform"
cens_par	float	Parameter of the censoring distribution.	5.0
max_time	float	Administrative end of follow-up, applied on top of random censoring.	None
layout	('intervals', 'panel')	"intervals" returns counting-process rows: one per transition a subject was at risk of, over the interval it was at risk, with status marking the one that occurred. "panel" returns one row per observation of the subject's state.	"intervals"
seed	int or Generator	Seed or generator for reproducibility.	None

Returns:

Type	Description
DataFrame	For layout="intervals": ["id", "start", "stop", "from_state", "to_state", "status", "X0", ...]. For layout="panel": ["id", "time", "state", "X0", ...].

Raises:

Type	Description
ValidationError	If the graph is malformed or any parameter is out of range.

Source code in gen_surv/multistate.py

def gen_multistate(
    n: int,
    transitions: Sequence[Transition],
    clock: Clock = "forward",
    initial_state: int = 1,
    covariate_dist: Literal["normal", "uniform", "binary"] = "normal",
    covariate_params: dict[str, float] | None = None,
    model_cens: Literal["uniform", "exponential"] = "uniform",
    cens_par: float = 5.0,
    max_time: float | None = None,
    layout: Layout = "intervals",
    seed: RandomStateLike = None,
) -> pd.DataFrame:
    """Simulate a multistate process over an arbitrary transition graph.

    Parameters
    ----------
    n : int
        Number of subjects. Each contributes several rows, so the frame is
        longer than ``n``.
    transitions : Sequence[Transition]
        The graph. Every edge carries its own baseline hazard and coefficients.
        A state with no outgoing transition is absorbing.
    clock : {"forward", "reset"}
        ``"forward"`` measures the hazard from entry to the study, giving a
        Markov process; ``"reset"`` measures it from entry to the current
        state, giving a semi-Markov one. They coincide for an exponential
        baseline.
    initial_state : int
        The state every subject starts in, at time zero.
    covariate_dist : {"normal", "uniform", "binary"}
        Distribution the covariates are drawn from.
    covariate_params : dict[str, float], optional
        Parameters of that distribution; defaults are filled in.
    model_cens : {"uniform", "exponential"}
        Random censoring mechanism.
    cens_par : float
        Parameter of the censoring distribution.
    max_time : float, optional
        Administrative end of follow-up, applied on top of random censoring.
    layout : {"intervals", "panel"}
        ``"intervals"`` returns counting-process rows: one per transition a
        subject was at risk of, over the interval it was at risk, with
        ``status`` marking the one that occurred. ``"panel"`` returns one row
        per observation of the subject's state.
    seed : int or numpy.random.Generator, optional
        Seed or generator for reproducibility.

    Returns
    -------
    pd.DataFrame
        For ``layout="intervals"``: ``["id", "start", "stop", "from_state",
        "to_state", "status", "X0", ...]``. For ``layout="panel"``:
        ``["id", "time", "state", "X0", ...]``.

    Raises
    ------
    ValidationError
        If the graph is malformed or any parameter is out of range.
    """
    ensure_positive_int(n, "n")
    ensure_in_choices(clock, "clock", ("forward", "reset"))
    ensure_in_choices(layout, "layout", ("intervals", "panel"))
    ensure_in_choices(model_cens, "model_cens", ("uniform", "exponential"))
    ensure_positive(cens_par, "cens_par")
    if max_time is not None:
        ensure_positive(max_time, "max_time")

    n_covariates = _validate_graph(transitions, initial_state)

    rng = resolve_rng(seed)
    covariate_params = set_covariate_params(covariate_dist, covariate_params)
    covariates = (
        generate_covariates(n, n_covariates, covariate_dist, covariate_params, rng)
        if n_covariates
        else np.zeros((n, 0))
    )

    coefficients = np.array(
        [list(t.coefficients) for t in transitions], dtype=float
    ).reshape(len(transitions), n_covariates)
    # One linear predictor per subject per transition.
    eta = covariates @ coefficients.T

    rfunc: CensoringFunc = runifcens if model_cens == "uniform" else rexpocens
    dropout = rfunc(n, cens_par, rng)
    ends = dropout if max_time is None else np.minimum(dropout, max_time)

    outgoing: dict[int, list[tuple[int, Transition]]] = {}
    for index, transition in enumerate(transitions):
        outgoing.setdefault(transition.origin, []).append((index, transition))

    occupancies, latent_times = _walk_cohort(
        eta=eta,
        ends=ends,
        outgoing=outgoing,
        initial_state=initial_state,
        clock=clock,
        rng=rng,
        n_transitions=len(transitions),
    )

    # Columns are accumulated as arrays and concatenated once. Building the
    # frame from Python tuples costs more than the sampling does.
    chunks: dict[str, list[NDArray[Any]]] = {}

    def add(**columns: NDArray[Any]) -> None:
        for key, values in columns.items():
            chunks.setdefault(key, []).append(values)

    if layout == "intervals":
        for subjects, state, entry, exit_time, destination in occupancies:
            width = len(subjects)
            for _, transition in outgoing[state]:
                add(
                    id=subjects,
                    start=entry,
                    stop=exit_time,
                    from_state=np.full(width, state, dtype=np.int64),
                    to_state=np.full(width, transition.destination, dtype=np.int64),
                    status=(destination == transition.destination).astype(np.int64),
                )
    else:
        add(
            id=np.arange(n, dtype=np.int64),
            time=np.zeros(n),
            state=np.full(n, initial_state, dtype=np.int64),
        )
        for subjects, state, _entry, exit_time, destination in occupancies:
            add(
                id=subjects,
                time=exit_time,
                # A transition is observed in its destination; the end of
                # follow-up is observed in the state still occupied.
                state=np.where(destination == -1, state, destination).astype(np.int64),
            )

    columns = INTERVAL_COLUMNS if layout == "intervals" else PANEL_COLUMNS
    data = pd.DataFrame({name: np.concatenate(chunks[name]) for name in columns})

    for j in range(n_covariates):
        data[f"X{j}"] = covariates[data["id"].to_numpy(), j]

    record(
        transitions=tuple(transitions),
        clock=clock,
        covariates=covariates,
        linear_predictor=eta,
        censoring_time=dropout,
        followup_end=ends,
        latent_times={
            (t.origin, t.destination): latent_times[:, i]
            for i, t in enumerate(transitions)
        },
    )

    # Waves put every subject's first occupancy before anyone's second, so the
    # frame is sorted back into per-subject order.
    order = ["id", "start"] if layout == "intervals" else ["id", "time"]
    if layout == "intervals":
        order.append("to_state")
    data = data.sort_values(order, kind="stable").reset_index(drop=True)

    return data