Fix black formatting (#43)

Author: DiogoRibeiro7

examples/run_aft.py

@@ -1,6 +1,7 @@
-import sys
 import os
-sys.path.insert(0, os.path.abspath(os.path.join(os.path.dirname(__file__), '..')))
+import sys
+
+sys.path.insert(0, os.path.abspath(os.path.join(os.path.dirname(__file__), "..")))
 from gen_surv.interface import generate
 
 # Generate synthetic survival data using Log-Normal AFT model
@@ -11,7 +12,7 @@
     sigma=1.0,
     model_cens="exponential",
     cens_par=3.0,
-    seed=123
+    seed=123,
 )
 
 print(df.head())

examples/run_aft_weibull.py

@@ -2,20 +2,21 @@
 Example demonstrating Weibull AFT model and visualization capabilities.
 """
 
-import sys
 import os
+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__), '..')))
+sys.path.insert(0, os.path.abspath(os.path.join(os.path.dirname(__file__), "..")))
 
 from gen_surv import generate
 from gen_surv.visualization import (
-    plot_survival_curve,
-    plot_hazard_comparison,
+    describe_survival,
     plot_covariate_effect,
-    describe_survival
+    plot_hazard_comparison,
+    plot_survival_curve,
 )
 
 # 1. Generate data from different models for comparison
@@ -28,7 +29,7 @@
         scale=2.0,
         model_cens="uniform",
         cens_par=5.0,
-        seed=42
+        seed=42,
     ),
     "Weibull AFT (shape=1.0)": generate(
         model="aft_weibull",
@@ -38,7 +39,7 @@
         scale=2.0,
         model_cens="uniform",
         cens_par=5.0,
-        seed=42
+        seed=42,
     ),
     "Weibull AFT (shape=2.0)": generate(
         model="aft_weibull",
@@ -48,8 +49,8 @@
         scale=2.0,
         model_cens="uniform",
         cens_par=5.0,
-        seed=42
-    )
+        seed=42,
+    ),
 }
 
 # Print sample data
@@ -59,18 +60,15 @@
 
 # 2. Compare survival curves from different models
 fig1, ax1 = plot_survival_curve(
-    data=pd.concat(
-        [df.assign(_model=name) for name, df in models.items()]
-    ),
+    data=pd.concat([df.assign(_model=name) for name, df in models.items()]),
     group_col="_model",
-    title="Comparing Survival Curves with Different Weibull Shapes"
+    title="Comparing Survival Curves with Different Weibull Shapes",
 )
 plt.savefig("survival_curve_comparison.png", dpi=300, bbox_inches="tight")
 
 # 3. Compare hazard functions
 fig2, ax2 = plot_hazard_comparison(
-    models=models,
-    title="Comparing Hazard Functions with Different Weibull Shapes"
+    models=models, title="Comparing Hazard Functions with Different Weibull Shapes"
 )
 plt.savefig("hazard_comparison.png", dpi=300, bbox_inches="tight")
 
@@ -79,7 +77,7 @@
     data=models["Weibull AFT (shape=2.0)"],
     covariate_col="X0",
     n_groups=3,
-    title="Effect of X0 Covariate on Survival"
+    title="Effect of X0 Covariate on Survival",
 )
 plt.savefig("covariate_effect.png", dpi=300, bbox_inches="tight")
 

examples/run_cmm.py

@@ -1,6 +1,7 @@
-import sys
 import os
-sys.path.insert(0, os.path.abspath(os.path.join(os.path.dirname(__file__), '..')))
+import sys
+
+sys.path.insert(0, os.path.abspath(os.path.join(os.path.dirname(__file__), "..")))
 
 from gen_surv import generate
 
@@ -11,7 +12,7 @@
     cens_par=2.0,
     qmat=[[0, 0.1], [0.05, 0]],
     p0=[1.0, 0.0],
-    seed=42
+    seed=42,
 )
 
 print(df.head())

examples/run_competing_risks.py

@@ -2,53 +2,64 @@
 Example demonstrating the Competing Risks models and visualization.
 """
 
-import sys
 import os
+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__), '..')))
+sys.path.insert(0, os.path.abspath(os.path.join(os.path.dirname(__file__), "..")))
 
 from gen_surv import generate
-from gen_surv.competing_risks import gen_competing_risks, gen_competing_risks_weibull, cause_specific_cumulative_incidence
-from gen_surv.summary import summarize_survival_dataset, compare_survival_datasets
+from gen_surv.competing_risks import (
+    cause_specific_cumulative_incidence,
+    gen_competing_risks,
+    gen_competing_risks_weibull,
+)
+from gen_surv.summary import compare_survival_datasets, summarize_survival_dataset
 
 
 def plot_cause_specific_cumulative_incidence(df, time_points=None, figsize=(10, 6)):
     """Plot the cause-specific cumulative incidence functions."""
     if time_points is None:
         max_time = df["time"].max()
         time_points = np.linspace(0, max_time, 100)
-    
+
     # Get unique causes (excluding censoring)
     causes = sorted([c for c in df["status"].unique() if c > 0])
-    
+
     # Create the plot
     fig, ax = plt.subplots(figsize=figsize)
-    
+
     for cause in causes:
         cif = cause_specific_cumulative_incidence(df, time_points, cause=cause)
         ax.plot(cif["time"], cif["incidence"], label=f"Cause {cause}")
-    
+
     # Add overlay showing number of subjects at each time
     time_bins = np.linspace(0, df["time"].max(), 10)
     event_counts = np.histogram(df.loc[df["status"] > 0, "time"], bins=time_bins)[0]
-    
+
     # Add a secondary y-axis for event counts
     ax2 = ax.twinx()
-    ax2.bar(time_bins[:-1], event_counts, width=time_bins[1]-time_bins[0], 
-            alpha=0.2, color='gray', align='edge')
-    ax2.set_ylabel('Number of events')
+    ax2.bar(
+        time_bins[:-1],
+        event_counts,
+        width=time_bins[1] - time_bins[0],
+        alpha=0.2,
+        color="gray",
+        align="edge",
+    )
+    ax2.set_ylabel("Number of events")
     ax2.grid(False)
-    
+
     # Format the main plot
     ax.set_xlabel("Time")
     ax.set_ylabel("Cumulative Incidence")
     ax.set_title("Cause-Specific Cumulative Incidence Functions")
     ax.legend()
     ax.grid(alpha=0.3)
-    
+
     return fig, ax
 
 
@@ -61,7 +72,7 @@ def plot_cause_specific_cumulative_incidence(df, time_points=None, figsize=(10,
     betas=[[0.8, -0.5], [0.2, 0.7]],
     model_cens="uniform",
     cens_par=2.0,
-    seed=42
+    seed=42,
 )
 
 # 2. Generate data with Weibull hazards (different shapes)
@@ -74,7 +85,7 @@ def plot_cause_specific_cumulative_incidence(df, time_points=None, figsize=(10,
     betas=[[0.8, -0.5], [0.2, 0.7]],
     model_cens="uniform",
     cens_par=2.0,
-    seed=42
+    seed=42,
 )
 
 # 3. Print summary statistics for both datasets
@@ -96,17 +107,13 @@ def plot_cause_specific_cumulative_incidence(df, time_points=None, figsize=(10,
 time_points = np.linspace(0, 5, 100)
 
 fig1, ax1 = plot_cause_specific_cumulative_incidence(
-    data_exponential, 
-    time_points=time_points, 
-    figsize=(10, 6)
+    data_exponential, time_points=time_points, figsize=(10, 6)
 )
 plt.title("Cumulative Incidence Functions (Exponential Hazards)")
 plt.savefig("cr_exponential_cif.png", dpi=300, bbox_inches="tight")
 
 fig2, ax2 = plot_cause_specific_cumulative_incidence(
-    data_weibull, 
-    time_points=time_points, 
-    figsize=(10, 6)
+    data_weibull, time_points=time_points, figsize=(10, 6)
 )
 plt.title("Cumulative Incidence Functions (Weibull Hazards)")
 plt.savefig("cr_weibull_cif.png", dpi=300, bbox_inches="tight")
@@ -121,16 +128,15 @@ def plot_cause_specific_cumulative_incidence(df, time_points=None, figsize=(10,
     betas=[[0.8, -0.5], [0.2, 0.7]],
     model_cens="uniform",
     cens_par=2.0,
-    seed=42
+    seed=42,
 )
 print(data_unified.head())
 
 # 7. Compare datasets
 print("\nComparing datasets:")
-comparison = compare_survival_datasets({
-    "Exponential": data_exponential,
-    "Weibull": data_weibull
-})
+comparison = compare_survival_datasets(
+    {"Exponential": data_exponential, "Weibull": data_weibull}
+)
 print(comparison)
 
 # Show plots if running interactively

examples/run_cphm.py

@@ -1,6 +1,7 @@
-import sys
 import os
-sys.path.insert(0, os.path.abspath(os.path.join(os.path.dirname(__file__), '..')))
+import sys
+
+sys.path.insert(0, os.path.abspath(os.path.join(os.path.dirname(__file__), "..")))
 
 from gen_surv import generate
 
@@ -11,7 +12,7 @@
     cens_par=1.0,
     beta=0.5,
     covariate_range=2.0,
-    seed=42
+    seed=42,
 )
 
 print(df.head())

examples/run_tdcm.py

@@ -1,6 +1,7 @@
-import sys
 import os
-sys.path.insert(0, os.path.abspath(os.path.join(os.path.dirname(__file__), '..')))
+import sys
+
+sys.path.insert(0, os.path.abspath(os.path.join(os.path.dirname(__file__), "..")))
 
 from gen_surv import generate
 
@@ -14,7 +15,7 @@
     cens_par=1.0,
     beta=[0.1, 0.2, 0.3],
     lam=1.0,
-    seed=42
+    seed=42,
 )
 
 print(df.head())

examples/run_thmm.py

@@ -1,6 +1,7 @@
-import sys
 import os
-sys.path.insert(0, os.path.abspath(os.path.join(os.path.dirname(__file__), '..')))
+import sys
+
+sys.path.insert(0, os.path.abspath(os.path.join(os.path.dirname(__file__), "..")))
 
 from gen_surv import generate
 
@@ -12,7 +13,7 @@
     p0=[1.0, 0.0, 0.0],
     model_cens="exponential",
     cens_par=3.0,
-    seed=42
+    seed=42,
 )
 
 print(df.head())

gen_surv/__init__.py

@@ -5,31 +5,32 @@
 
 from importlib.metadata import PackageNotFoundError, version
 
-# Main interface
-from .interface import generate
+from .aft import gen_aft_log_logistic, gen_aft_log_normal, gen_aft_weibull
 
-# Individual generators
-from .cphm import gen_cphm
+# Helper functions
+from .bivariate import sample_bivariate_distribution
+from .censoring import rexpocens, runifcens
 from .cmm import gen_cmm
-from .tdcm import gen_tdcm
-from .thmm import gen_thmm
-from .aft import gen_aft_log_normal, gen_aft_weibull, gen_aft_log_logistic
 from .competing_risks import gen_competing_risks, gen_competing_risks_weibull
-from .mixture import gen_mixture_cure, cure_fraction_estimate
-from .piecewise import gen_piecewise_exponential
+
+# Individual generators
+from .cphm import gen_cphm
 from .export import export_dataset
 
-# Helper functions
-from .bivariate import sample_bivariate_distribution
-from .censoring import runifcens, rexpocens
+# Main interface
+from .interface import generate
+from .mixture import cure_fraction_estimate, gen_mixture_cure
+from .piecewise import gen_piecewise_exponential
+from .tdcm import gen_tdcm
+from .thmm import gen_thmm
 
 # Visualization tools (requires matplotlib and lifelines)
 try:
     from .visualization import (
-        plot_survival_curve,
-        plot_hazard_comparison,
-        plot_covariate_effect,
         describe_survival,
+        plot_covariate_effect,
+        plot_hazard_comparison,
+        plot_survival_curve,
     )
 
     _has_visualization = True