Add extensive summary tests (#45)

Author: DiogoRibeiro7

tests/test_export.py

@@ -1,6 +1,7 @@
 import os
 
 import pandas as pd
+import pytest
 
 from gen_surv import export_dataset, generate
 
@@ -35,3 +36,22 @@ def test_export_dataset_json(tmp_path):
     assert out_file.exists()
     loaded = pd.read_json(out_file, orient="table")
     pd.testing.assert_frame_equal(df.reset_index(drop=True), loaded)
+
+
+def test_export_dataset_feather_and_invalid(tmp_path):
+    df = generate(
+        model="cphm",
+        n=5,
+        model_cens="uniform",
+        cens_par=1.0,
+        beta=0.5,
+        covariate_range=1.0,
+    )
+    feather_file = tmp_path / "data.feather"
+    export_dataset(df, str(feather_file))
+    assert feather_file.exists()
+    loaded = pd.read_feather(feather_file)
+    pd.testing.assert_frame_equal(df.reset_index(drop=True), loaded)
+
+    with pytest.raises(ValueError):
+        export_dataset(df, str(tmp_path / "data.txt"), fmt="txt")

tests/test_piecewise.py

@@ -18,3 +18,19 @@ def test_piecewise_invalid_lengths():
         gen_piecewise_exponential(
             n=5, breakpoints=[1.0, 2.0], hazard_rates=[0.5], seed=42
         )
+
+def test_piecewise_invalid_hazard_and_breakpoints():
+    with pytest.raises(ValueError):
+        gen_piecewise_exponential(
+            n=5,
+            breakpoints=[2.0, 1.0],
+            hazard_rates=[0.5, 1.0, 1.5],
+            seed=42,
+        )
+    with pytest.raises(ValueError):
+        gen_piecewise_exponential(
+            n=5,
+            breakpoints=[1.0],
+            hazard_rates=[0.5, -1.0],
+            seed=42,
+        )

tests/test_piecewise_functions.py

@@ -0,0 +1,51 @@
+import numpy as np
+from gen_surv.piecewise import piecewise_hazard_function, piecewise_survival_function
+
+
+def test_piecewise_hazard_function_scalar_and_array():
+    breakpoints = [1.0, 2.0]
+    hazard_rates = [0.5, 1.0, 1.5]
+    # Scalar values
+    assert piecewise_hazard_function(0.5, breakpoints, hazard_rates) == 0.5
+    assert piecewise_hazard_function(1.5, breakpoints, hazard_rates) == 1.0
+    assert piecewise_hazard_function(3.0, breakpoints, hazard_rates) == 1.5
+    # Array values
+    arr = np.array([0.5, 1.5, 3.0])
+    np.testing.assert_allclose(
+        piecewise_hazard_function(arr, breakpoints, hazard_rates),
+        np.array([0.5, 1.0, 1.5]),
+    )
+
+
+def test_piecewise_hazard_function_negative_time():
+    """Hazard should be zero for negative times."""
+    breakpoints = [1.0, 2.0]
+    hazard_rates = [0.5, 1.0, 1.5]
+    assert piecewise_hazard_function(-1.0, breakpoints, hazard_rates) == 0
+    np.testing.assert_array_equal(
+        piecewise_hazard_function(np.array([-0.5, -2.0]), breakpoints, hazard_rates),
+        np.array([0.0, 0.0]),
+    )
+
+
+def test_piecewise_survival_function():
+    breakpoints = [1.0, 2.0]
+    hazard_rates = [0.5, 1.0, 1.5]
+    # Known survival probabilities
+    expected = np.exp(-np.array([0.0, 0.25, 1.0, 3.0]))
+    times = np.array([0.0, 0.5, 1.5, 3.0])
+    np.testing.assert_allclose(
+        piecewise_survival_function(times, breakpoints, hazard_rates),
+        expected,
+    )
+
+
+def test_piecewise_survival_function_scalar_and_negative():
+    breakpoints = [1.0, 2.0]
+    hazard_rates = [0.5, 1.0, 1.5]
+    # Scalar output should be a float
+    val = piecewise_survival_function(1.5, breakpoints, hazard_rates)
+    assert isinstance(val, float)
+    assert np.isclose(val, np.exp(-1.0))
+    # Negative times return survival of 1
+    assert piecewise_survival_function(-2.0, breakpoints, hazard_rates) == 1

tests/test_summary_extra.py

@@ -0,0 +1,51 @@
+import pandas as pd
+import pytest
+from gen_surv.summary import check_survival_data_quality, compare_survival_datasets
+from gen_surv import generate
+
+
+def test_check_survival_data_quality_fix_issues():
+    df = pd.DataFrame(
+        {
+            "time": [1.0, -0.5, None, 1.0],
+            "status": [1, 2, 0, 1],
+            "id": [1, 2, 3, 1],
+        }
+    )
+    fixed, issues = check_survival_data_quality(
+        df,
+        id_col="id",
+        max_time=2.0,
+        fix_issues=True,
+    )
+    assert issues["modifications"]["rows_dropped"] == 2
+    assert issues["modifications"]["values_fixed"] == 1
+    assert len(fixed) == 2
+
+
+def test_check_survival_data_quality_no_fix():
+    """Issues should be reported but data left unchanged when fix_issues=False."""
+    df = pd.DataFrame({"time": [-1.0, 2.0], "status": [3, 1]})
+    checked, issues = check_survival_data_quality(df, max_time=1.0, fix_issues=False)
+    # Data is returned unmodified
+    pd.testing.assert_frame_equal(df, checked)
+    assert issues["invalid_values"]["negative_time"] == 1
+    assert issues["invalid_values"]["excessive_time"] == 1
+    assert issues["invalid_values"]["invalid_status"] == 1
+
+
+def test_compare_survival_datasets_basic():
+    ds1 = generate(model="cphm", n=5, model_cens="uniform", cens_par=1.0, beta=0.5, covariate_range=1.0)
+    ds2 = generate(model="cphm", n=5, model_cens="uniform", cens_par=1.0, beta=1.0, covariate_range=1.0)
+    comparison = compare_survival_datasets({"A": ds1, "B": ds2})
+    assert set(["A", "B"]).issubset(comparison.columns)
+    assert "n_subjects" in comparison.index
+
+
+def test_compare_survival_datasets_with_covariates_and_empty_error():
+    ds = generate(model="cphm", n=3, model_cens="uniform", cens_par=1.0, beta=0.5, covariate_range=1.0)
+    comparison = compare_survival_datasets({"only": ds}, covariate_cols=["X0"])
+    assert "only" in comparison.columns
+    assert "X0_mean" in comparison.index
+    with pytest.raises(ValueError):
+        compare_survival_datasets({})

tests/test_summary_more.py

@@ -0,0 +1,58 @@
+import pandas as pd
+import pytest
+from gen_surv.summary import (
+    summarize_survival_dataset,
+    check_survival_data_quality,
+    _print_summary,
+)
+
+
+def test_summarize_survival_dataset_errors():
+    df = pd.DataFrame({"time": [1, 2], "status": [1, 0]})
+    # Missing time column
+    with pytest.raises(ValueError):
+        summarize_survival_dataset(df.drop(columns=["time"]))
+    # Missing ID column when specified
+    with pytest.raises(ValueError):
+        summarize_survival_dataset(df, id_col="id")
+    # Missing covariate columns
+    with pytest.raises(ValueError):
+        summarize_survival_dataset(df, covariate_cols=["bad"])
+
+
+def test_summarize_survival_dataset_verbose_output(capsys):
+    df = pd.DataFrame(
+        {
+            "time": [1.0, 2.0, 3.0],
+            "status": [1, 0, 1],
+            "id": [1, 2, 3],
+            "age": [30, 40, 50],
+            "group": ["A", "B", "A"],
+        }
+    )
+    summary = summarize_survival_dataset(
+        df, id_col="id", covariate_cols=["age", "group"]
+    )
+    _print_summary(summary, "time", "status", "id", ["age", "group"])
+    captured = capsys.readouterr().out
+    assert "SURVIVAL DATASET SUMMARY" in captured
+    assert "age:" in captured
+    assert "Categorical" in captured
+
+
+def test_check_survival_data_quality_duplicates_and_fix():
+    df = pd.DataFrame(
+        {
+            "time": [1.0, -1.0, 2.0, 1.0],
+            "status": [1, 1, 0, 1],
+            "id": [1, 1, 2, 1],
+        }
+    )
+    checked, issues = check_survival_data_quality(df, id_col="id", fix_issues=False)
+    assert issues["duplicates"]["duplicate_rows"] == 1
+    assert issues["duplicates"]["duplicate_ids"] == 2
+    fixed, issues_fixed = check_survival_data_quality(
+        df, id_col="id", max_time=2.0, fix_issues=True
+    )
+    assert len(fixed) < len(df)
+    assert issues_fixed["modifications"]["rows_dropped"] > 0

tests/test_validate.py

@@ -71,3 +71,16 @@ def test_validate_gen_aft_log_normal_inputs_valid():
 def test_validate_dg_biv_inputs_valid_weibull():
     """Valid parameters for a Weibull distribution should pass."""
     v.validate_dg_biv_inputs(5, "weibull", 0.1, [1.0, 1.0, 1.0, 1.0])
+
+
+def test_validate_gen_aft_weibull_inputs_and_log_logistic():
+    with pytest.raises(ValueError):
+        v.validate_gen_aft_weibull_inputs(0, [0.1], 1.0, 1.0, "uniform", 1.0)
+    with pytest.raises(ValueError):
+        v.validate_gen_aft_log_logistic_inputs(1, [0.1], -1.0, 1.0, "uniform", 1.0)
+
+
+def test_validate_competing_risks_inputs():
+    with pytest.raises(ValueError):
+        v.validate_competing_risks_inputs(1, 2, [0.1], None, "uniform", 1.0)
+    v.validate_competing_risks_inputs(1, 1, [0.5], [[0.1]], "uniform", 0.5)