Skip to content
← Back to PyStatsBio overview on sgcx.org
Back to top

# Source code for pystatsbio.power._survival

    """Power calculations for survival (time-to-event) endpoints.

    Validates against: R gsDesign::nSurv(), TrialSize
    """

    from __future__ import annotations

    import math

    from scipy.stats import norm

    from pystatsbio.power._common import PowerResult, _check_power_args, _solve_parameter

    _VALID_METHODS = ("schoenfeld", "freedman", "lachin_foulkes")
    _VALID_ALTERNATIVES = ("two.sided", "one.sided")

    # ---------------------------------------------------------------------------
    # Internal power computations — one per method
    # ---------------------------------------------------------------------------

    def _z_alpha(alpha: float, alternative: str) -> float:
        """Critical z-value for the given alpha and sidedness."""
        if alternative == "two.sided":
            return norm.ppf(1.0 - alpha / 2.0)
        return norm.ppf(1.0 - alpha)

    def _logrank_power_schoenfeld(
        n: float,
        hr: float,
        alpha: float,
        alternative: str,
        p_event: float,
        alloc_ratio: float,
    ) -> float:
        """Schoenfeld (1981) formula.

        Number of events:  d = (z_alpha + z_beta)^2 * (r+1)^2 / (r * log(HR)^2)
        Rearranged to solve for power given n:
            d = n * p_event
            z_beta = sqrt(d * r / (r+1)^2) * |log(HR)| - z_alpha
            power = Phi(z_beta)
        """
        r = alloc_ratio
        d = n * p_event  # expected number of events
        log_hr = math.log(hr)
        za = _z_alpha(alpha, alternative)

        z_beta = math.sqrt(d * r / (r + 1.0) ** 2) * abs(log_hr) - za

        pwr = float(norm.cdf(z_beta))
        return pwr

    def _logrank_power_freedman(
        n: float,
        hr: float,
        alpha: float,
        alternative: str,
        p_event: float,
        alloc_ratio: float,
    ) -> float:
        """Freedman (1982) formula.

        Uses (HR-1)/(HR+1) instead of log(HR) in the variance estimate.
        d = (z_alpha + z_beta)^2 * (HR+1)^2 / ((HR-1)^2) * (r+1)^2 / (r)
        Rearranged: z_beta = sqrt(d * r / (r+1)^2) * |HR-1| / (HR+1) - z_alpha
        """
        r = alloc_ratio
        d = n * p_event
        za = _z_alpha(alpha, alternative)

        # Freedman uses (HR-1)/(HR+1) instead of log(HR)
        z_beta = math.sqrt(d * r / (r + 1.0) ** 2) * abs(hr - 1.0) / (hr + 1.0) - za

        pwr = float(norm.cdf(z_beta))
        return pwr

    def _logrank_power_lachin_foulkes(
        n: float,
        hr: float,
        alpha: float,
        alternative: str,
        p_event: float,
        alloc_ratio: float,
    ) -> float:
        """Lachin & Foulkes (1986) formula.

        Adds a correction for the difference in survival distributions between
        the two groups, yielding a tighter variance estimate. For equal
        allocation (r=1) and no censoring (p_event=1), reduces to:
            z_beta = sqrt(n * p_event / 4) * |log(HR)| - z_alpha
        with a variance correction term.

        This implementation uses the simplified Lachin-Foulkes formula
        appropriate for exponential survival.
        """
        r = alloc_ratio
        d = n * p_event
        log_hr = math.log(hr)
        za = _z_alpha(alpha, alternative)

        # Proportion in each group
        p1 = r / (r + 1.0)  # treatment
        p2 = 1.0 / (r + 1.0)  # control

        # Under exponential survival, the variance of log-HR is approximately
        # 1/(d * p1 * p2). The Lachin-Foulkes correction adjusts for unequal
        # censoring patterns, but in the simplified version (uniform censoring),
        # the formula is:
        z_beta = abs(log_hr) * math.sqrt(d * p1 * p2) - za

        pwr = float(norm.cdf(z_beta))
        return pwr

    # ---------------------------------------------------------------------------
    # Public API
    # ---------------------------------------------------------------------------

    def power_logrank(
        n: int | None = None,
        hr: float | None = None,
        alpha: float = 0.05,
        power: float | None = None,
        alternative: str = "two.sided",
        p_event: float = 1.0,
        alloc_ratio: float = 1.0,
        method: str = "schoenfeld",
    ) -> PowerResult:
        """Power calculation for the log-rank test.

        Exactly one of ``n``, ``hr``, ``power`` must be ``None``.

        Parameters
        ----------
        n : int or None
            Total sample size (both groups combined).
        hr : float or None
            Hazard ratio under the alternative hypothesis. Must be != 1.
        alpha : float
            Significance level (default 0.05).
        power : float or None
            Desired power.
        alternative : str
            ``'two.sided'`` or ``'one.sided'``.
        p_event : float
            Probability of observing an event (1.0 = no censoring).
        alloc_ratio : float
            Allocation ratio (n_treatment / n_control). Default 1:1.
        method : str
            ``'schoenfeld'``, ``'freedman'``, or ``'lachin_foulkes'``.

        Returns
        -------
        PowerResult

        Examples
        --------
        >>> r = power_logrank(hr=0.7, alpha=0.05, power=0.80)
        >>> r.n  # total N for both groups
        186

        Validates against: R gsDesign::nSurv(), TrialSize
        """
        if alternative not in _VALID_ALTERNATIVES:
            raise ValueError(
                f"alternative must be one of {_VALID_ALTERNATIVES}, got {alternative!r}"
            )
        if method not in _VALID_METHODS:
            raise ValueError(f"method must be one of {_VALID_METHODS}, got {method!r}")
        if not (0.0 < p_event <= 1.0):
            raise ValueError(f"p_event must be in (0, 1], got {p_event}")
        if alloc_ratio <= 0.0:
            raise ValueError(f"alloc_ratio must be > 0, got {alloc_ratio}")

        solve_for = _check_power_args(
            n=n, effect=hr, power=power, alpha=alpha, effect_name="hr",
        )

        # HR must not be 1.0 when solving for n or power
        if hr is not None and hr == 1.0:
            raise ValueError("hr must be != 1.0 (no effect)")
        if hr is not None and hr <= 0.0:
            raise ValueError(f"hr must be > 0, got {hr}")

        # Select the method
        power_funcs = {
            "schoenfeld": _logrank_power_schoenfeld,
            "freedman": _logrank_power_freedman,
            "lachin_foulkes": _logrank_power_lachin_foulkes,
        }
        _power_func = power_funcs[method]

        def _compute(n_val: float, hr_val: float) -> float:
            return _power_func(n_val, hr_val, alpha, alternative, p_event, alloc_ratio)

        if solve_for == "power":
            assert n is not None and hr is not None
            result_power = _compute(float(n), hr)
            result_n = n
            result_hr = hr

        elif solve_for == "n":
            assert hr is not None and power is not None
            raw_n = _solve_parameter(
                func=lambda x: _compute(x, hr),
                target=power,
                bracket=(2.0, 1e7),
            )
            result_n = math.ceil(raw_n)
            result_power = power
            result_hr = hr

        else:  # solve_for == "effect" (hr)
            assert n is not None and power is not None
            # HR can be < 1 or > 1. For two-sided, solve for HR < 1 by convention.
            result_hr = _solve_parameter(
                func=lambda x: _compute(float(n), x),
                target=power,
                bracket=(0.01, 0.999),
            )
            result_n = n
            result_power = power

        method_labels = {
            "schoenfeld": "Schoenfeld",
            "freedman": "Freedman",
            "lachin_foulkes": "Lachin-Foulkes",
        }

        return PowerResult(
            n=result_n,
            power=result_power,
            effect_size=result_hr,
            alpha=alpha,
            alternative=alternative,
            method=f"Log-rank test power calculation ({method_labels[method]})",
            note="n is total sample size (both groups combined)",
        )