Merge pull request #1 from bscisel/dev-jwinter

Dev jwinter

Author: jwinter3

.gitignore

@@ -15,4 +15,5 @@ benchmark/bin/env.sh
 benchmark/src/results
 benchmark/src/results/overlap
 mprofile*dat
-*csv
+*csv
+!docs/notebooks/example.csv

Cargo.lock

@@ -1,5 +1,6 @@
 from polars_bio.polars_bio import InputFormat, ReadOptions, VcfReadOptions
 
+from .base_sequnce_quality_vis import visualize_base_sequence_quality
 from .context import ctx, set_option
 from .io import (
     describe_vcf,
@@ -16,6 +17,7 @@
 from .polars_ext import PolarsRangesOperations as LazyFrame
 from .range_op import FilterOp, count_overlaps, coverage, merge, nearest, overlap
 from .range_viz import visualize_intervals
+from .quality_stats import base_sequence_quality
 
 POLARS_BIO_MAX_THREADS = "datafusion.execution.target_partitions"
 
@@ -29,6 +31,7 @@
     "coverage",
     "ctx",
     "FilterOp",
+    "visualize_base_sequence_quality",
     "visualize_intervals",
     "read_bam",
     "read_vcf",
@@ -45,4 +48,5 @@
     "ReadOptions",
     "VcfReadOptions",
     "set_option",
+    "base_sequence_quality",
 ]

docs/notebooks/base_sequence_quality.ipynb

@@ -0,0 +1,51 @@
+from typing import Union
+
+import pandas as pd
+import polars as pl
+from matplotlib import pyplot as plt
+
+
+def visualize_base_sequence_quality(df: Union[pd.DataFrame, pl.DataFrame]) -> None:
+    """
+    Visualize the overlapping intervals.
+
+    Parameters:
+        df: Pandas DataFrame or Polars DataFrame. The DataFrame containing the base sequence quality results
+    """
+    assert isinstance(
+        df, (pd.DataFrame, pl.DataFrame)
+    ), "df must be a Pandas or Polars DataFrame"
+    df = df if isinstance(df, pd.DataFrame) else df.to_pandas()
+    df = df.sort_values(by="pos")
+
+    boxes = [
+        {
+            "label": int(row["pos"]),
+            "whislo": row["lower"],
+            "q1": row["q1"],
+            "med": row["median"],
+            "q3": row["q3"],
+            "whishi": row["upper"],
+        }
+        for _, row in df.iterrows()
+    ]
+
+    fig, ax = plt.subplots()
+    fig.set_size_inches(15, 5)
+
+    plot = ax.plot(df["pos"] + 1, df["avg"])
+    box_plot = ax.bxp(boxes, showfliers=False)
+
+    ax.set_title("base sequence quality")
+    ax.set_ylabel("Phred score")
+    ax.set_xlabel("Position in read (bp)")
+
+    ax.legend(
+        [plot[0], box_plot["medians"][0]],
+        ["Average of phred score", "Median of phred score"],
+    )
+
+    for label in ax.get_xticklabels():
+        label.set_fontsize(6)
+
+    plt.show()

docs/notebooks/example.csv

@@ -0,0 +1,64 @@
+from pathlib import Path
+from typing import Union
+import datafusion
+import polars as pl
+import pandas as pd
+import pyarrow as pa
+from .context import ctx
+from polars_bio.polars_bio import (
+    base_sequance_quality_scan,
+    base_sequance_quality_frame,
+)
+
+
+def base_sequence_quality(
+    df: Union[str, Path, pl.DataFrame, pl.LazyFrame, pd.DataFrame],
+    quality_scores_column: str = "quality_scores",
+    output_type: str = "polars.DataFrame",
+    target_partitions: int = 8,
+) -> Union[pl.DataFrame, pd.DataFrame]:
+    """
+    Compute base sequence quality statistics from various dataframe/file types.
+
+    Args:
+        df: Input data as a file path or dataframe.
+        quality_scores_column: Name of the column with quality scores.
+        output_type: Output type, either "polars.DataFrame" or "pandas.DataFrame".
+
+    Returns:
+        DataFrame with base sequence quality statistics.
+    """
+    ctx.set_option(
+        "datafusion.execution.target_partitions", str(target_partitions), False
+    )
+
+    if isinstance(df, (str, Path)):
+        df = str(df)
+        supported_exts = {".parquet", ".csv", ".bed", ".vcf", ".fastq"}
+        ext = set(Path(df).suffixes)
+        if not (supported_exts & ext or not ext):
+            raise ValueError(
+                "Input file must be a Parquet, CSV, BED, VCF, or FASTQ file."
+            )
+        result: datafusion.DataFrame = base_sequance_quality_scan(
+            ctx, df, quality_scores_column
+        )
+    else:
+        if isinstance(df, pl.LazyFrame):
+            arrow_table = df.collect().to_arrow()
+        elif isinstance(df, pl.DataFrame):
+            arrow_table = df.to_arrow()
+        elif isinstance(df, pd.DataFrame):
+            arrow_table = pa.Table.from_pandas(df)
+        else:
+            raise TypeError("Unsupported dataframe type.")
+        result: datafusion.DataFrame = base_sequance_quality_frame(
+            ctx, arrow_table, quality_scores_column
+        )
+
+    if output_type == "polars.DataFrame":
+        return result.to_polars()
+    elif output_type == "pandas.DataFrame":
+        return result.to_pandas()
+    else:
+        raise ValueError("output_type must be 'polars.DataFrame' or 'pandas.DataFrame'")