feat: expose dev extras (#65)

Author: DiogoRibeiro7

.pre-commit-config.yaml

@@ -15,3 +15,5 @@ repos:
     rev: v1.15.0
     hooks:
       - id: mypy
+        pass_filenames: false
+        args: [--config-file=pyproject.toml, gen_surv]

examples/run_aft_weibull.py

@@ -6,7 +6,6 @@
 import sys
 
 import matplotlib.pyplot as plt
-import numpy as np
 import pandas as pd
 
 sys.path.insert(0, os.path.abspath(os.path.join(os.path.dirname(__file__), "..")))

examples/run_competing_risks.py

@@ -7,7 +7,6 @@
 
 import matplotlib.pyplot as plt
 import numpy as np
-import pandas as pd
 
 sys.path.insert(0, os.path.abspath(os.path.join(os.path.dirname(__file__), "..")))
 

gen_surv/__init__.py

@@ -11,14 +11,14 @@
 from .bivariate import sample_bivariate_distribution
 from .censoring import (
     CensoringModel,
+    GammaCensoring,
+    LogNormalCensoring,
+    WeibullCensoring,
     rexpocens,
+    rgammacens,
+    rlognormcens,
     runifcens,
     rweibcens,
-    rlognormcens,
-    rgammacens,
-    WeibullCensoring,
-    LogNormalCensoring,
-    GammaCensoring,
 )
 from .cmm import gen_cmm
 from .competing_risks import gen_competing_risks, gen_competing_risks_weibull
@@ -38,12 +38,10 @@
 
 # Visualization tools (requires matplotlib and lifelines)
 try:
-    from .visualization import (
-        describe_survival,
-        plot_covariate_effect,
-        plot_hazard_comparison,
-        plot_survival_curve,
-    )
+    from .visualization import describe_survival  # noqa: F401
+    from .visualization import plot_covariate_effect  # noqa: F401
+    from .visualization import plot_hazard_comparison  # noqa: F401
+    from .visualization import plot_survival_curve  # noqa: F401
 
     _has_visualization = True
 except ImportError:

gen_surv/_validation.py

@@ -55,7 +55,9 @@ class PositiveSequenceError(ValidationError):
     """Raised when a sequence contains non-positive elements."""
 
     def __init__(self, name: str, seq: Sequence[Any]) -> None:
-        super().__init__(f"All elements in '{name}' must be greater than 0; got {seq!r}")
+        super().__init__(
+            f"All elements in '{name}' must be greater than 0; got {seq!r}"
+        )
 
 
 class ListOfListsError(ValidationError):
@@ -69,9 +71,7 @@ class ParameterError(ValidationError):
     """Raised when a parameter falls outside its allowed range."""
 
     def __init__(self, name: str, value: Any, constraint: str) -> None:
-        super().__init__(
-            f"Invalid value for '{name}': {value!r}. {constraint}"
-        )
+        super().__init__(f"Invalid value for '{name}': {value!r}. {constraint}")
 
 
 _ALLOWED_CENSORING = {"uniform", "exponential"}
@@ -108,16 +108,19 @@ def _to_float_array(seq: Sequence[Any], name: str) -> NDArray[np.float64]:
     except (TypeError, ValueError) as exc:
         raise NumericSequenceError(name, seq) from exc
 
+
 def ensure_numeric_sequence(seq: Sequence[Any], name: str) -> None:
     """Ensure all elements of ``seq`` are numeric."""
     _to_float_array(seq, name)
 
+
 def ensure_positive_sequence(seq: Sequence[float], name: str) -> None:
     """Ensure all elements of ``seq`` are positive."""
     arr = _to_float_array(seq, name)
     if np.any(arr <= 0):
         raise PositiveSequenceError(name, seq)
 
+
 def ensure_censoring_model(model_cens: str) -> None:
     """Validate that the censoring model is supported."""
     ensure_in_choices(model_cens, "model_cens", _ALLOWED_CENSORING)

gen_surv/bivariate.py

@@ -1,10 +1,10 @@
+from typing import Sequence
+
 import numpy as np
 from numpy.typing import NDArray
-from typing import Sequence
 
 from .validate import validate_dg_biv_inputs
 
-
 _CHI2_SCALE = 0.5
 _CLIP_EPS = 1e-10
 
@@ -43,7 +43,9 @@ def sample_bivariate_distribution(
     mean = [0, 0]
     cov = [[1, corr], [corr, 1]]
     z = np.random.multivariate_normal(mean, cov, size=n)
-    u = 1 - np.exp(-_CHI2_SCALE * z**2)  # transform normals to uniform via chi-squared approx
+    u = 1 - np.exp(
+        -_CHI2_SCALE * z**2
+    )  # transform normals to uniform via chi-squared approx
     u = np.clip(u, _CLIP_EPS, 1 - _CLIP_EPS)  # avoid infs in tails
 
     # Step 2: Transform to marginals

gen_surv/censoring.py

@@ -1,6 +1,7 @@
+from typing import Protocol
+
 import numpy as np
 from numpy.typing import NDArray
-from typing import Protocol
 
 
 class CensoringFunc(Protocol):

gen_surv/cmm.py

@@ -1,6 +1,7 @@
+from typing import Sequence, TypedDict
+
 import numpy as np
 import pandas as pd
-from typing import Sequence, TypedDict
 
 from gen_surv.censoring import CensoringFunc, rexpocens, runifcens
 from gen_surv.validate import validate_gen_cmm_inputs

gen_surv/competing_risks.py

@@ -11,11 +11,11 @@
 import pandas as pd
 
 from ._validation import (
+    ParameterError,
     ensure_censoring_model,
     ensure_in_choices,
     ensure_positive_sequence,
     ensure_sequence_length,
-    ParameterError,
 )
 from .censoring import rexpocens, runifcens
 
@@ -109,9 +109,7 @@ def gen_competing_risks(
     n_covariates = 2  # Default number of covariates
 
     # Set default covariate parameters if not provided
-    ensure_in_choices(
-        covariate_dist, "covariate_dist", {"normal", "uniform", "binary"}
-    )
+    ensure_in_choices(covariate_dist, "covariate_dist", {"normal", "uniform", "binary"})
     if covariate_params is None:
         if covariate_dist == "normal":
             covariate_params = {"mean": 0.0, "std": 1.0}
@@ -309,9 +307,7 @@ def gen_competing_risks_weibull(
     n_covariates = 2  # Default number of covariates
 
     # Set default covariate parameters if not provided
-    ensure_in_choices(
-        covariate_dist, "covariate_dist", {"normal", "uniform", "binary"}
-    )
+    ensure_in_choices(covariate_dist, "covariate_dist", {"normal", "uniform", "binary"})
     if covariate_params is None:
         if covariate_dist == "normal":
             covariate_params = {"mean": 0.0, "std": 1.0}

gen_surv/interface.py

@@ -6,7 +6,7 @@
     >>> df = generate(model="cphm", n=100, model_cens="uniform", cens_par=1.0, beta=0.5, covariate_range=2.0)
 """
 
-from typing import Any, Literal, Protocol, Dict
+from typing import Any, Dict, Literal, Protocol
 
 import pandas as pd
 
@@ -18,6 +18,7 @@
 from gen_surv.piecewise import gen_piecewise_exponential
 from gen_surv.tdcm import gen_tdcm
 from gen_surv.thmm import gen_thmm
+
 from ._validation import ensure_in_choices
 
 # Type definitions for model names
@@ -35,6 +36,7 @@
     "piecewise_exponential",
 ]
 
+
 # Interface for generator callables
 class DataGenerator(Protocol):
     def __call__(self, **kwargs: Any) -> pd.DataFrame: ...