move modules into scripts to avoid having to import entire local package

Author: poldrack

src/bettercode/simple_workflow/snakemake_workflow/scripts/compute_correlation.py

@@ -4,9 +4,13 @@
 
 import pandas as pd
 
-from bettercode.simple_workflow.correlation import (
-    compute_correlation_matrix,
-)
+
+def compute_correlation_matrix(
+    df: pd.DataFrame,
+    method: str = "spearman",
+) -> pd.DataFrame:
+    """Compute correlation matrix using the specified method."""
+    return df.corr(method=method)
 
 
 def main():

src/bettercode/simple_workflow/snakemake_workflow/scripts/filter_data.py

@@ -4,9 +4,10 @@
 
 import pandas as pd
 
-from bettercode.simple_workflow.filter_data import (
-    filter_numerical_columns,
-)
+
+def filter_numerical_columns(df: pd.DataFrame) -> pd.DataFrame:
+    """Filter a dataframe to keep only numerical columns."""
+    return df.select_dtypes(include=["number"])
 
 
 def main():

src/bettercode/simple_workflow/snakemake_workflow/scripts/generate_heatmap.py

@@ -2,11 +2,45 @@
 
 from pathlib import Path
 
+import matplotlib.pyplot as plt
 import pandas as pd
+import seaborn as sns
 
-from bettercode.simple_workflow.visualization import (
-    generate_clustered_heatmap,
-)
+
+def generate_clustered_heatmap(
+    corr_matrix: pd.DataFrame,
+    output_path: Path | None = None,
+    figsize: tuple[int, int] = (8, 10),
+    cmap: str = "coolwarm",
+    vmin: float = -1.0,
+    vmax: float = 1.0,
+) -> sns.matrix.ClusterGrid:
+    """Generate a clustered heatmap from a correlation matrix."""
+    # Create clustered heatmap
+    g = sns.clustermap(
+        corr_matrix,
+        cmap=cmap,
+        vmin=vmin,
+        vmax=vmax,
+        figsize=figsize,
+        dendrogram_ratio=(0.1, 0.1),
+        cbar_pos=(0.02, 0.8, 0.03, 0.15),
+        xticklabels=False,
+        yticklabels=True,
+    )
+
+    # Set y-axis label font size
+    plt.setp(g.ax_heatmap.get_yticklabels(), rotation=0, fontsize=3)
+
+    # Set title
+    g.fig.suptitle("Clustered Correlation Heatmap (Spearman)", y=1.02, fontsize=14)
+
+    # Save if output path provided
+    if output_path is not None:
+        output_path.parent.mkdir(parents=True, exist_ok=True)
+        g.savefig(output_path, dpi=300, bbox_inches="tight")
+
+    return g
 
 
 def main():

src/bettercode/simple_workflow/snakemake_workflow/scripts/join_data.py

@@ -4,7 +4,14 @@
 
 import pandas as pd
 
-from bettercode.simple_workflow.join_data import join_dataframes
+
+def join_dataframes(
+    df1: pd.DataFrame,
+    df2: pd.DataFrame,
+    how: str = "inner",
+) -> pd.DataFrame:
+    """Join two dataframes based on their index."""
+    return df1.join(df2, how=how, lsuffix="_mv", rsuffix="_demo")
 
 
 def main():