pywombat 0.1.0__py3-none-any.whl

pywombat/__init__.py

@@ -0,0 +1,2 @@
+def main() -> None:
+    print("Hello from pywombat!")

pywombat/cli.py

@@ -0,0 +1,935 @@
+"""CLI for wombat tool."""
+
+import re
+import warnings
+from pathlib import Path
+from typing import Optional
+
+import click
+import polars as pl
+import yaml
+
+
+@click.command()
+@click.argument("input_file", type=click.Path(exists=True, path_type=Path))
+@click.option(
+    "-o",
+    "--output",
+    type=str,
+    help="Output file prefix. If not specified, prints to stdout.",
+)
+@click.option(
+    "-f",
+    "--format",
+    "output_format",
+    type=click.Choice(["tsv", "parquet"], case_sensitive=False),
+    default="tsv",
+    help="Output format: tsv (default) or parquet.",
+)
+@click.option("-v", "--verbose", is_flag=True, help="Enable verbose output.")
+@click.option(
+    "-p",
+    "--pedigree",
+    type=click.Path(exists=True, path_type=Path),
+    help="Pedigree file to add father and mother genotype columns.",
+)
+@click.option(
+    "-F",
+    "--filter-config",
+    type=click.Path(exists=True, path_type=Path),
+    help="Filter configuration YAML file to apply quality and impact filters.",
+)
+def cli(
+    input_file: Path,
+    output: Optional[str],
+    output_format: str,
+    verbose: bool,
+    pedigree: Optional[Path],
+    filter_config: Optional[Path],
+):
+    """
+    Wombat: A tool for processing bcftools tabulated TSV files.
+
+    This command:
+
+    \b
+    1. Expands the '(null)' column containing NAME=value pairs separated by ';'
+    2. Preserves the CSQ (Consequence) column without melting
+    3. Melts sample columns into rows with sample names
+    4. Splits sample values (GT:DP:GQ:AD format) into separate columns:
+       - sample_gt: Genotype
+       - sample_dp: Read depth
+       - sample_gq: Genotype quality
+       - sample_ad: Allele depth (second value from comma-separated list)
+       - sample_vaf: Variant allele frequency (sample_ad / sample_dp)
+
+    \b
+    Examples:
+        wombat input.tsv -o output
+        wombat input.tsv -o output -f parquet
+        wombat input.tsv > output.tsv
+    """
+    try:
+        if verbose:
+            click.echo(f"Reading input file: {input_file}", err=True)
+
+        # Read the TSV file
+        df = pl.read_csv(input_file, separator="\t")
+
+        if verbose:
+            click.echo(
+                f"Input shape: {df.shape[0]} rows, {df.shape[1]} columns", err=True
+            )
+
+        # Read pedigree file if provided
+        pedigree_df = None
+        if pedigree:
+            if verbose:
+                click.echo(f"Reading pedigree file: {pedigree}", err=True)
+            pedigree_df = read_pedigree(pedigree)
+
+        # Process the dataframe
+        formatted_df = format_bcftools_tsv(df, pedigree_df)
+
+        if verbose:
+            click.echo(
+                f"Output shape: {formatted_df.shape[0]} rows, {formatted_df.shape[1]} columns",
+                err=True,
+            )
+
+        # Apply filters if provided
+        filter_config_data = None
+        if filter_config:
+            if verbose:
+                click.echo(f"Reading filter config: {filter_config}", err=True)
+            filter_config_data = load_filter_config(filter_config)
+
+        # Apply filters and write output
+        if filter_config_data:
+            apply_filters_and_write(
+                formatted_df,
+                filter_config_data,
+                output,
+                output_format,
+                verbose,
+            )
+        else:
+            # No filters - write single output file
+            if output:
+                # Construct output filename with prefix and format
+                output_path = Path(f"{output}.{output_format}")
+
+                if output_format == "tsv":
+                    formatted_df.write_csv(output_path, separator="\t")
+                elif output_format == "parquet":
+                    formatted_df.write_parquet(output_path)
+
+                click.echo(f"Formatted data written to {output_path}", err=True)
+            else:
+                # Write to stdout (only for TSV format)
+                if output_format != "tsv":
+                    click.echo(
+                        "Error: stdout output only supported for TSV format. Use -o to specify an output prefix for parquet.",
+                        err=True,
+                    )
+                    raise click.Abort()
+                click.echo(formatted_df.write_csv(separator="\t"), nl=False)
+
+    except Exception as e:
+        click.echo(f"Error: {e}", err=True)
+        raise click.Abort()
+
+
+def load_filter_config(config_path: Path) -> dict:
+    """Load and parse filter configuration from YAML file."""
+    with open(config_path, "r") as f:
+        config = yaml.safe_load(f)
+    return config
+
+
+def apply_quality_filters(
+    df: pl.DataFrame, quality_config: dict, verbose: bool = False
+) -> pl.DataFrame:
+    """Apply quality filters to the dataframe."""
+    if quality_config is None:
+        return df
+
+    original_rows = df.shape[0]
+
+    # Check for rare genotypes with '2' and warn
+    if "2" in str(df["sample_gt"].to_list()):
+        rare_gts = df.filter(pl.col("sample_gt").str.contains("2"))
+        if rare_gts.shape[0] > 0:
+            warnings.warn(
+                f"Found {rare_gts.shape[0]} rows with rare genotypes containing '2'. These will be kept."
+            )
+
+    # Filter: sample_gt must contain at least one '1' (default: true)
+    filter_no_alt = quality_config.get("filter_no_alt_allele", True)
+    if filter_no_alt:
+        df = df.filter(
+            pl.col("sample_gt").str.contains("1")
+            | pl.col("sample_gt").str.contains("2")
+        )
+
+    # Apply minimum depth filter
+    if "sample_dp_min" in quality_config:
+        min_dp = quality_config["sample_dp_min"]
+        df = df.filter(pl.col("sample_dp").cast(pl.Float64, strict=False) >= min_dp)
+
+    # Apply minimum GQ filter
+    if "sample_gq_min" in quality_config:
+        min_gq = quality_config["sample_gq_min"]
+        df = df.filter(pl.col("sample_gq").cast(pl.Float64, strict=False) >= min_gq)
+
+    # Determine genotype for VAF filters
+    # Het: contains exactly one '1' (0/1 or 1/0)
+    # HomAlt: 1/1
+    # HomRef: 0/0
+    is_het = (pl.col("sample_gt").str.count_matches("1") == 1) & ~pl.col(
+        "sample_gt"
+    ).str.contains("2")
+    is_hom_alt = pl.col("sample_gt") == "1/1"
+    is_hom_ref = pl.col("sample_gt") == "0/0"
+
+    # VAF filters for heterozygous
+    if "sample_vaf_het_min" in quality_config:
+        min_vaf_het = quality_config["sample_vaf_het_min"]
+        df = df.filter(~is_het | (pl.col("sample_vaf") >= min_vaf_het))
+
+    if "sample_vaf_het_max" in quality_config:
+        max_vaf_het = quality_config["sample_vaf_het_max"]
+        df = df.filter(~is_het | (pl.col("sample_vaf") <= max_vaf_het))
+
+    # VAF filters for homozygous alternate
+    if "sample_vaf_homalt_min" in quality_config:
+        min_vaf_homalt = quality_config["sample_vaf_homalt_min"]
+        df = df.filter(~is_hom_alt | (pl.col("sample_vaf") >= min_vaf_homalt))
+
+    # VAF filters for homozygous reference (wild type)
+    if "sample_vaf_hom_ref_max" in quality_config:
+        max_vaf_hom_ref = quality_config["sample_vaf_hom_ref_max"]
+        df = df.filter(~is_hom_ref | (pl.col("sample_vaf") <= max_vaf_hom_ref))
+
+    # Apply filters to parents if they exist and option is enabled
+    apply_to_parents = quality_config.get("apply_to_parents", False)
+    if apply_to_parents and "father" in df.columns:
+        # Apply same filters to father columns
+        if "sample_dp_min" in quality_config:
+            min_dp = quality_config["sample_dp_min"]
+            df = df.filter(
+                (pl.col("father_dp").is_null())
+                | (pl.col("father_dp").cast(pl.Float64, strict=False) >= min_dp)
+            )
+
+        if "sample_gq_min" in quality_config:
+            min_gq = quality_config["sample_gq_min"]
+            df = df.filter(
+                (pl.col("father_gq").is_null())
+                | (pl.col("father_gq").cast(pl.Float64, strict=False) >= min_gq)
+            )
+
+        # Father genotype checks
+        father_is_het = (pl.col("father_gt").str.count_matches("1") == 1) & ~pl.col(
+            "father_gt"
+        ).str.contains("2")
+        father_is_hom_alt = pl.col("father_gt") == "1/1"
+        father_is_hom_ref = pl.col("father_gt") == "0/0"
+
+        if "sample_vaf_het_min" in quality_config:
+            min_vaf_het = quality_config["sample_vaf_het_min"]
+            df = df.filter(
+                pl.col("father_vaf").is_null()
+                | ~father_is_het
+                | (pl.col("father_vaf") >= min_vaf_het)
+            )
+
+        if "sample_vaf_het_max" in quality_config:
+            max_vaf_het = quality_config["sample_vaf_het_max"]
+            df = df.filter(
+                pl.col("father_vaf").is_null()
+                | ~father_is_het
+                | (pl.col("father_vaf") <= max_vaf_het)
+            )
+
+        if "sample_vaf_homalt_min" in quality_config:
+            min_vaf_homalt = quality_config["sample_vaf_homalt_min"]
+            df = df.filter(
+                pl.col("father_vaf").is_null()
+                | ~father_is_hom_alt
+                | (pl.col("father_vaf") >= min_vaf_homalt)
+            )
+
+        if "sample_vaf_hom_ref_max" in quality_config:
+            max_vaf_hom_ref = quality_config["sample_vaf_hom_ref_max"]
+            df = df.filter(
+                pl.col("father_vaf").is_null()
+                | ~father_is_hom_ref
+                | (pl.col("father_vaf") <= max_vaf_hom_ref)
+            )
+
+        # Apply same filters to mother columns
+        if "sample_dp_min" in quality_config:
+            min_dp = quality_config["sample_dp_min"]
+            df = df.filter(
+                (pl.col("mother_dp").is_null())
+                | (pl.col("mother_dp").cast(pl.Float64, strict=False) >= min_dp)
+            )
+
+        if "sample_gq_min" in quality_config:
+            min_gq = quality_config["sample_gq_min"]
+            df = df.filter(
+                (pl.col("mother_gq").is_null())
+                | (pl.col("mother_gq").cast(pl.Float64, strict=False) >= min_gq)
+            )
+
+        # Mother genotype checks
+        mother_is_het = (pl.col("mother_gt").str.count_matches("1") == 1) & ~pl.col(
+            "mother_gt"
+        ).str.contains("2")
+        mother_is_hom_alt = pl.col("mother_gt") == "1/1"
+        mother_is_hom_ref = pl.col("mother_gt") == "0/0"
+
+        if "sample_vaf_het_min" in quality_config:
+            min_vaf_het = quality_config["sample_vaf_het_min"]
+            df = df.filter(
+                pl.col("mother_vaf").is_null()
+                | ~mother_is_het
+                | (pl.col("mother_vaf") >= min_vaf_het)
+            )
+
+        if "sample_vaf_het_max" in quality_config:
+            max_vaf_het = quality_config["sample_vaf_het_max"]
+            df = df.filter(
+                pl.col("mother_vaf").is_null()
+                | ~mother_is_het
+                | (pl.col("mother_vaf") <= max_vaf_het)
+            )
+
+        if "sample_vaf_homalt_min" in quality_config:
+            min_vaf_homalt = quality_config["sample_vaf_homalt_min"]
+            df = df.filter(
+                pl.col("mother_vaf").is_null()
+                | ~mother_is_hom_alt
+                | (pl.col("mother_vaf") >= min_vaf_homalt)
+            )
+
+        if "sample_vaf_hom_ref_max" in quality_config:
+            max_vaf_hom_ref = quality_config["sample_vaf_hom_ref_max"]
+            df = df.filter(
+                pl.col("mother_vaf").is_null()
+                | ~mother_is_hom_ref
+                | (pl.col("mother_vaf") <= max_vaf_hom_ref)
+            )
+
+    if verbose:
+        filtered_rows = df.shape[0]
+        click.echo(
+            f"Quality filters: {original_rows} -> {filtered_rows} rows ({original_rows - filtered_rows} filtered out)",
+            err=True,
+        )
+
+    return df
+
+
+def parse_impact_filter_expression(expression: str, df: pl.DataFrame) -> pl.Expr:
+    """Parse a filter expression string into a Polars expression."""
+    # Replace operators with Polars equivalents
+    # Support: =, !=, <=, >=, <, >, &, |, ()
+
+    expr_str = expression.strip()
+
+    # Split by logical operators while preserving them
+    tokens = re.split(r"(\s*[&|]\s*|\(|\))", expr_str)
+    tokens = [t.strip() for t in tokens if t.strip()]
+
+    def parse_condition(condition: str) -> pl.Expr:
+        """Parse a single condition into a Polars expression."""
+        condition = condition.strip()
+
+        # Try different operators in order of specificity
+        for op in ["<=", ">=", "!=", "=", "<", ">"]:
+            if op in condition:
+                parts = condition.split(op, 1)
+                if len(parts) == 2:
+                    col_name = parts[0].strip()
+                    value = parts[1].strip()
+
+                    # Check if column exists
+                    if col_name not in df.columns:
+                        raise ValueError(f"Column '{col_name}' not found in dataframe")
+
+                    # Try to convert value to number, otherwise treat as string
+                    try:
+                        value_num = float(value)
+                        col_expr = pl.col(col_name).cast(pl.Float64, strict=False)
+
+                        if op == "=":
+                            return col_expr == value_num
+                        elif op == "!=":
+                            return col_expr != value_num
+                        elif op == "<=":
+                            return col_expr <= value_num
+                        elif op == ">=":
+                            return col_expr >= value_num
+                        elif op == "<":
+                            return col_expr < value_num
+                        elif op == ">":
+                            return col_expr > value_num
+                    except ValueError:
+                        # String comparison (case-insensitive)
+                        value = value.strip("'\"")
+                        col_expr = pl.col(col_name).str.to_lowercase()
+                        value_lower = value.lower()
+
+                        if op == "=":
+                            return col_expr == value_lower
+                        elif op == "!=":
+                            return col_expr != value_lower
+                        else:
+                            raise ValueError(
+                                f"Operator '{op}' not supported for string comparison"
+                            )
+                break
+
+        raise ValueError(f"Could not parse condition: {condition}")
+
+    def build_expression(tokens: list, idx: int = 0) -> tuple[pl.Expr, int]:
+        """Recursively build expression from tokens."""
+        if idx >= len(tokens):
+            return None, idx
+
+        result = None
+        i = idx
+
+        while i < len(tokens):
+            token = tokens[i]
+
+            if token == "(":
+                # Parse sub-expression
+                sub_expr, new_i = build_expression(tokens, i + 1)
+                if result is None:
+                    result = sub_expr
+                i = new_i
+            elif token == ")":
+                # End of sub-expression
+                return result, i + 1
+            elif token == "&":
+                # AND operator
+                next_expr, new_i = build_expression(tokens, i + 1)
+                if next_expr is not None:
+                    result = result & next_expr if result is not None else next_expr
+                return result, new_i
+            elif token == "|":
+                # OR operator
+                next_expr, new_i = build_expression(tokens, i + 1)
+                if next_expr is not None:
+                    result = result | next_expr if result is not None else next_expr
+                return result, new_i
+            else:
+                # It's a condition
+                cond_expr = parse_condition(token)
+                if result is None:
+                    result = cond_expr
+                else:
+                    # If we have a result and encounter another condition without an operator,
+                    # assume AND
+                    result = result & cond_expr
+
+            i += 1
+
+        return result, i
+
+    expr, _ = build_expression(tokens)
+    return expr
+
+
+def apply_impact_filters(
+    df: pl.DataFrame,
+    impact_config: list,
+    output_prefix: str,
+    output_format: str,
+    verbose: bool,
+):
+    """Apply impact filters and create separate output files."""
+    if not impact_config:
+        return
+
+    # Sort impact filters by priority (lower number = higher priority)
+    impact_filters = sorted(impact_config, key=lambda x: x.get("priority", 999))
+
+    # Create a dict to store variants by impact filter
+    impact_variants = {}
+
+    # Apply each impact filter
+    for impact_filter in impact_filters:
+        name = impact_filter["name"]
+        priority = impact_filter.get("priority", 999)
+        expression = impact_filter["expression"]
+
+        if verbose:
+            click.echo(
+                f"Applying impact filter '{name}' (priority {priority})...", err=True
+            )
+
+        try:
+            # Parse and apply the filter expression
+            filter_expr = parse_impact_filter_expression(expression, df)
+            filtered_df = df.filter(filter_expr)
+
+            if verbose:
+                click.echo(
+                    f"  Impact filter '{name}': {filtered_df.shape[0]} variants",
+                    err=True,
+                )
+
+            # Store the filtered dataframe with its priority
+            impact_variants[name] = {
+                "df": filtered_df,
+                "priority": priority,
+            }
+        except Exception as e:
+            click.echo(
+                f"Error applying impact filter '{name}': {e}",
+                err=True,
+            )
+            raise
+
+    # Add flag_higher_impact column to each filtered dataframe
+    for name, data in impact_variants.items():
+        filtered_df = data["df"]
+        priority = data["priority"]
+
+        # Find variants that appear in higher priority filters
+        higher_priority_filters = []
+        for other_name, other_data in impact_variants.items():
+            if other_data["priority"] < priority:
+                higher_priority_filters.append(other_name)
+
+        if higher_priority_filters:
+            # Create a set of variant keys from higher priority filters
+            variant_keys_in_higher = set()
+            for other_name in higher_priority_filters:
+                other_df = impact_variants[other_name]["df"]
+                for row in other_df.select(["#CHROM", "POS", "REF", "ALT"]).iter_rows():
+                    variant_keys_in_higher.add(tuple(row))
+
+            # Add flag_higher_impact column
+            def check_higher_impact(chrom, pos, ref, alt):
+                key = (chrom, pos, ref, alt)
+                if key in variant_keys_in_higher:
+                    # Find which filters it appears in
+                    filters_with_variant = []
+                    for other_name in higher_priority_filters:
+                        other_df = impact_variants[other_name]["df"]
+                        match = other_df.filter(
+                            (pl.col("#CHROM") == chrom)
+                            & (pl.col("POS") == pos)
+                            & (pl.col("REF") == ref)
+                            & (pl.col("ALT") == alt)
+                        )
+                        if match.shape[0] > 0:
+                            filters_with_variant.append(other_name)
+                    return (
+                        ", ".join(filters_with_variant) if filters_with_variant else ""
+                    )
+                return ""
+
+            # Add the flag column
+            filtered_df = filtered_df.with_columns(
+                [
+                    pl.struct(["#CHROM", "POS", "REF", "ALT"])
+                    .map_elements(
+                        lambda x: check_higher_impact(
+                            x["#CHROM"], x["POS"], x["REF"], x["ALT"]
+                        ),
+                        return_dtype=pl.Utf8,
+                    )
+                    .alias("flag_higher_impact")
+                ]
+            )
+        else:
+            # No higher priority filters
+            filtered_df = filtered_df.with_columns(
+                [pl.lit("").alias("flag_higher_impact")]
+            )
+
+        # Write to file
+        output_filename = f"{output_prefix}_{name}.{output_format}"
+        output_path = Path(output_filename)
+
+        if output_format == "tsv":
+            filtered_df.write_csv(output_path, separator="\t")
+        elif output_format == "parquet":
+            filtered_df.write_parquet(output_path)
+
+        click.echo(
+            f"Written {filtered_df.shape[0]} variants to {output_path}", err=True
+        )
+
+
+def apply_filters_and_write(
+    df: pl.DataFrame,
+    filter_config: dict,
+    output_prefix: Optional[str],
+    output_format: str,
+    verbose: bool,
+):
+    """Apply filters and write output files."""
+    # Apply quality filters first
+    quality_config = filter_config.get("quality", {})
+    filtered_df = apply_quality_filters(df, quality_config, verbose)
+
+    # Get impact filters
+    impact_config = filter_config.get("impact", [])
+
+    if not impact_config:
+        # No impact filters - write single output file
+        if not output_prefix:
+            # Write to stdout
+            if output_format != "tsv":
+                click.echo(
+                    "Error: stdout output only supported for TSV format.",
+                    err=True,
+                )
+                raise click.Abort()
+            click.echo(filtered_df.write_csv(separator="\t"), nl=False)
+        else:
+            output_path = Path(f"{output_prefix}.{output_format}")
+
+            if output_format == "tsv":
+                filtered_df.write_csv(output_path, separator="\t")
+            elif output_format == "parquet":
+                filtered_df.write_parquet(output_path)
+
+            click.echo(f"Formatted data written to {output_path}", err=True)
+    else:
+        # Apply impact filters and create multiple output files
+        if not output_prefix:
+            click.echo(
+                "Error: Output prefix required when using impact filters.",
+                err=True,
+            )
+            raise click.Abort()
+
+        apply_impact_filters(
+            filtered_df,
+            impact_config,
+            output_prefix,
+            output_format,
+            verbose,
+        )
+
+
+def read_pedigree(pedigree_path: Path) -> pl.DataFrame:
+    """
+    Read a pedigree file and return a DataFrame with sample relationships.
+
+    Args:
+        pedigree_path: Path to the pedigree file
+
+    Returns:
+        DataFrame with columns: sample_id, father_id, mother_id
+    """
+    # Try reading with header first
+    df = pl.read_csv(pedigree_path, separator="\t")
+
+    # Check if first row has 'FID' in first column (indicates header)
+    if df.columns[0] == "FID" or "sample_id" in df.columns:
+        # Has header - use it as-is
+        pass
+    else:
+        # No header - assume standard pedigree format
+        # FID, sample_id, father_id, mother_id, sex, phenotype
+        df.columns = ["FID", "sample_id", "father_id", "mother_id", "sex", "phenotype"]
+
+    # Ensure we have the required columns (try different possible names)
+    if "sample_id" not in df.columns and len(df.columns) >= 4:
+        # Try to identify columns by position
+        df = df.rename(
+            {
+                df.columns[1]: "sample_id",
+                df.columns[2]: "father_id",
+                df.columns[3]: "mother_id",
+            }
+        )
+
+    # Handle different column names for father/mother
+    if "FatherBarcode" in df.columns:
+        df = df.rename({"FatherBarcode": "father_id", "MotherBarcode": "mother_id"})
+
+    # Select only the columns we need
+    pedigree_df = df.select(["sample_id", "father_id", "mother_id"])
+
+    # Replace 0 and -9 with null (indicating no parent)
+    pedigree_df = pedigree_df.with_columns(
+        [
+            pl.when(pl.col("father_id").cast(pl.Utf8).is_in(["0", "-9"]))
+            .then(None)
+            .otherwise(pl.col("father_id"))
+            .alias("father_id"),
+            pl.when(pl.col("mother_id").cast(pl.Utf8).is_in(["0", "-9"]))
+            .then(None)
+            .otherwise(pl.col("mother_id"))
+            .alias("mother_id"),
+        ]
+    )
+
+    return pedigree_df
+
+
+def add_parent_genotypes(df: pl.DataFrame, pedigree_df: pl.DataFrame) -> pl.DataFrame:
+    """
+    Add father and mother genotype columns to the DataFrame.
+
+    Args:
+        df: DataFrame with sample genotype information
+        pedigree_df: DataFrame with parent relationships
+
+    Returns:
+        DataFrame with added parent genotype columns
+    """
+    # Join with pedigree to get father and mother IDs for each sample
+    df = df.join(pedigree_df, left_on="sample", right_on="sample_id", how="left")
+
+    # Define the core variant-identifying columns for joining parent genotypes
+    # We only want to join on genomic position, not on annotation columns
+    # This ensures we match parents even if they have different VEP annotations
+    core_variant_cols = ["#CHROM", "POS", "REF", "ALT"]
+    # Check which columns actually exist in the dataframe
+    join_cols = [col for col in core_variant_cols if col in df.columns]
+
+    # Create a self-join friendly version of the data for looking up parent genotypes
+    # We select only the join columns + sample genotype information
+    parent_lookup = df.select(
+        join_cols
+        + [
+            pl.col("sample"),
+            pl.col("sample_gt"),
+            pl.col("sample_dp"),
+            pl.col("sample_gq"),
+            pl.col("sample_ad"),
+            pl.col("sample_vaf"),
+        ]
+    ).unique()
+
+    # Join for father's genotypes
+    # Match on genomic position AND father_id == sample
+    father_data = parent_lookup.rename(
+        {
+            "sample": "father_id",
+            "sample_gt": "father_gt",
+            "sample_dp": "father_dp",
+            "sample_gq": "father_gq",
+            "sample_ad": "father_ad",
+            "sample_vaf": "father_vaf",
+        }
+    )
+
+    df = df.join(father_data, on=join_cols + ["father_id"], how="left")
+
+    # Join for mother's genotypes
+    mother_data = parent_lookup.rename(
+        {
+            "sample": "mother_id",
+            "sample_gt": "mother_gt",
+            "sample_dp": "mother_dp",
+            "sample_gq": "mother_gq",
+            "sample_ad": "mother_ad",
+            "sample_vaf": "mother_vaf",
+        }
+    )
+
+    df = df.join(mother_data, on=join_cols + ["mother_id"], how="left")
+
+    # Rename father_id and mother_id to father and mother for debugging
+    df = df.rename({"father_id": "father", "mother_id": "mother"})
+
+    # Replace '.' with '0' for parent DP and GQ columns
+    df = df.with_columns(
+        [
+            pl.when(pl.col("father_dp") == ".")
+            .then(pl.lit("0"))
+            .otherwise(pl.col("father_dp"))
+            .alias("father_dp"),
+            pl.when(pl.col("father_gq") == ".")
+            .then(pl.lit("0"))
+            .otherwise(pl.col("father_gq"))
+            .alias("father_gq"),
+            pl.when(pl.col("mother_dp") == ".")
+            .then(pl.lit("0"))
+            .otherwise(pl.col("mother_dp"))
+            .alias("mother_dp"),
+            pl.when(pl.col("mother_gq") == ".")
+            .then(pl.lit("0"))
+            .otherwise(pl.col("mother_gq"))
+            .alias("mother_gq"),
+        ]
+    )
+
+    return df
+
+
+def format_bcftools_tsv(
+    df: pl.DataFrame, pedigree_df: Optional[pl.DataFrame] = None
+) -> pl.DataFrame:
+    """
+    Format a bcftools tabulated TSV DataFrame.
+
+    Args:
+        df: Input DataFrame from bcftools
+        pedigree_df: Optional pedigree DataFrame with parent information
+
+    Returns:
+        Formatted DataFrame with expanded fields and melted samples
+    """
+    # Find the (null) column
+    if "(null)" not in df.columns:
+        raise ValueError("Column '(null)' not found in the input file")
+
+    # Get column index of (null)
+    null_col_idx = df.columns.index("(null)")
+
+    # Split columns into: before (null), (null), and after (null)
+    cols_after = df.columns[null_col_idx + 1 :]
+
+    # Step 1: Expand the (null) column
+    # Split by semicolon and create new columns
+
+    # First, we need to extract all unique field names from the (null) column
+    # to know what columns to create
+    null_values = df.select("(null)").to_series()
+    all_fields = set()
+
+    for value in null_values:
+        if value and not (isinstance(value, float)):  # Skip null/NaN values
+            pairs = str(value).split(";")
+            for pair in pairs:
+                if "=" in pair:
+                    field_name = pair.split("=", 1)[0]
+                    all_fields.add(field_name)
+
+    # Create expressions to extract each field
+    for field in sorted(all_fields):
+        # Extract the field value from the (null) column
+        # Pattern: extract value after "field=" and before ";" or end of string
+        df = df.with_columns(
+            pl.col("(null)").str.extract(f"{field}=([^;]+)").alias(field)
+        )
+
+    # Drop the original (null) column
+    df = df.drop("(null)")
+
+    # Drop CSQ column if it exists (it was extracted from (null) column)
+    if "CSQ" in df.columns:
+        df = df.drop("CSQ")
+
+    # Step 2: Identify sample columns and extract sample names
+    # Sample columns have format "sample_name:..." in the header
+    # Skip the CSQ column as it should not be melted (handled above)
+    sample_cols = []
+    sample_names = []
+
+    for col in cols_after:
+        # Skip CSQ column
+        if col == "CSQ":
+            continue
+
+        if ":" in col:
+            sample_name = col.split(":", 1)[0]
+            sample_cols.append(col)
+            sample_names.append(sample_name)
+        else:
+            # If no colon, treat the whole column name as sample name
+            sample_cols.append(col)
+            sample_names.append(col)
+
+    if not sample_cols:
+        # No sample columns to melt, just return expanded data
+        return df
+
+    # Step 3: Melt the sample columns
+    # Keep all columns except sample columns as id_vars
+    id_vars = [col for col in df.columns if col not in sample_cols]
+
+    # Create a mapping of old column names to sample names
+    rename_map = {old: new for old, new in zip(sample_cols, sample_names)}
+
+    # Rename sample columns to just sample names before melting
+    df = df.rename(rename_map)
+
+    # Melt the dataframe
+    melted_df = df.melt(
+        id_vars=id_vars,
+        value_vars=sample_names,
+        variable_name="sample",
+        value_name="sample_value",
+    )
+
+    # Step 4: Split sample_value into GT:DP:GQ:AD format
+    # Split on ':' to get individual fields
+    # Use nullable=True to handle missing fields gracefully
+    melted_df = melted_df.with_columns(
+        [
+            # GT - first field (nullable for missing data)
+            pl.col("sample_value")
+            .str.split(":")
+            .list.get(0, null_on_oob=True)
+            .alias("sample_gt"),
+            # DP - second field (nullable for missing data)
+            pl.col("sample_value")
+            .str.split(":")
+            .list.get(1, null_on_oob=True)
+            .alias("sample_dp"),
+            # GQ - third field (nullable for missing data)
+            pl.col("sample_value")
+            .str.split(":")
+            .list.get(2, null_on_oob=True)
+            .alias("sample_gq"),
+            # AD - fourth field, split on ',' and keep second value (nullable)
+            pl.col("sample_value")
+            .str.split(":")
+            .list.get(3, null_on_oob=True)
+            .str.split(",")
+            .list.get(1, null_on_oob=True)
+            .alias("sample_ad"),
+        ]
+    )
+
+    # Replace '.' with '0' for DP and GQ columns
+    melted_df = melted_df.with_columns(
+        [
+            pl.when(pl.col("sample_dp") == ".")
+            .then(pl.lit("0"))
+            .otherwise(pl.col("sample_dp"))
+            .alias("sample_dp"),
+            pl.when(pl.col("sample_gq") == ".")
+            .then(pl.lit("0"))
+            .otherwise(pl.col("sample_gq"))
+            .alias("sample_gq"),
+        ]
+    )
+
+    # Step 5: Calculate sample_vaf as sample_ad / sample_dp
+    # Convert to numeric, calculate ratio, handle division by zero
+    melted_df = melted_df.with_columns(
+        [
+            (
+                pl.col("sample_ad").cast(pl.Float64, strict=False)
+                / pl.col("sample_dp").cast(pl.Float64, strict=False)
+            ).alias("sample_vaf")
+        ]
+    )
+
+    # Drop the original sample_value column
+    melted_df = melted_df.drop("sample_value")
+
+    # Step 6: Add parent genotype information if pedigree is provided
+    if pedigree_df is not None:
+        melted_df = add_parent_genotypes(melted_df, pedigree_df)
+
+    return melted_df
+
+
+if __name__ == "__main__":
+    cli()

pywombat-0.1.0.dist-info/METADATA

@@ -0,0 +1,142 @@
+Metadata-Version: 2.4
+Name: pywombat
+Version: 0.1.0
+Summary: A CLI tool for processing and filtering bcftools tabulated TSV files with pedigree support
+Project-URL: Homepage, https://github.com/bourgeron-lab/pywombat
+Project-URL: Repository, https://github.com/bourgeron-lab/pywombat
+Project-URL: Issues, https://github.com/bourgeron-lab/pywombat/issues
+Author-email: Freddy Cliquet <fcliquet@pasteur.fr>
+License: MIT
+Keywords: bioinformatics,genomics,pedigree,variant-calling,vcf
+Classifier: Development Status :: 3 - Alpha
+Classifier: Intended Audience :: Science/Research
+Classifier: Programming Language :: Python :: 3
+Classifier: Programming Language :: Python :: 3.12
+Classifier: Topic :: Scientific/Engineering :: Bio-Informatics
+Requires-Python: >=3.12
+Requires-Dist: click>=8.1.0
+Requires-Dist: polars>=0.19.0
+Requires-Dist: pyyaml>=6.0
+Description-Content-Type: text/markdown
+
+# PyWombat
+
+A CLI tool for processing bcftools tabulated TSV files.
+
+## Installation
+
+This is a UV-managed Python package. To install:
+
+```bash
+uv sync
+```
+
+## Usage
+
+The `wombat` command processes bcftools tabulated TSV files:
+
+```bash
+# Format a bcftools TSV file and print to stdout
+wombat input.tsv
+
+# Format and save to output file (creates output.tsv by default)
+wombat input.tsv -o output
+
+# Format and save as parquet
+wombat input.tsv -o output -f parquet
+wombat input.tsv -o output --format parquet
+
+# Format with pedigree information to add parent genotypes
+wombat input.tsv --pedigree pedigree.tsv -o output
+```
+
+### What does `wombat` do?
+
+The `wombat` command processes bcftools tabulated TSV files by:
+
+1. **Expanding the `(null)` column**: This column contains multiple fields in the format `NAME=value` separated by semicolons (e.g., `DP=30;AF=0.5;AC=2`). Each field is extracted into its own column.
+
+2. **Preserving the `CSQ` column**: The CSQ (Consequence) column is preserved as-is and not melted, allowing VEP annotations to remain intact.
+
+3. **Melting and splitting sample columns**: After the `(null)` column, there are typically sample columns with values in `GT:DP:GQ:AD` format. The tool:
+   - Extracts the sample name (the part before the first `:` character)
+   - Transforms the wide format into long format
+   - Creates a `sample` column with the sample names
+   - Splits the sample values into separate columns:
+     - `sample_gt`: Genotype (e.g., 0/1, 1/1)
+     - `sample_dp`: Read depth
+     - `sample_gq`: Genotype quality
+     - `sample_ad`: Allele depth (takes the second value from comma-separated list)
+     - `sample_vaf`: Variant allele frequency (calculated as sample_ad / sample_dp)
+
+### Example
+
+**Input:**
+
+```tsv
+CHROM POS REF ALT (null) Sample1:GT:Sample1:DP:Sample1:GQ:Sample1:AD Sample2:GT:Sample2:DP:Sample2:GQ:Sample2:AD
+chr1 100 A T DP=30;AF=0.5;AC=2 0/1:15:99:5,10 1/1:18:99:0,18
+```
+
+**Output:**
+
+```tsv
+CHROM POS REF ALT AC AF DP sample sample_gt sample_dp sample_gq sample_ad sample_vaf
+chr1 100 A T 2 0.5 30 Sample1 0/1 15 99 10 0.6667
+chr1 100 A T 2 0.5 30 Sample2 1/1 18 99 18 1.0
+```
+
+Notes:
+
+- The `sample_ad` column contains the second value from the AD field (e.g., from `5,10` it extracts `10`)
+- The `sample_vaf` column is the variant allele frequency calculated as `sample_ad / sample_dp`
+- By default, output is in TSV format. Use `-f parquet` to output as Parquet files
+- The `-o` option specifies an output prefix (e.g., `-o output` creates `output.tsv` or `output.parquet`)
+
+### Pedigree Support
+
+You can provide a pedigree file with the `--pedigree` option to add parent genotype information to the output. This enables trio analysis by including the father's and mother's genotypes for each sample.
+
+**Pedigree File Format:**
+
+The pedigree file should be a tab-separated file with the following columns:
+
+- `FID`: Family ID
+- `sample_id`: Sample identifier (matches the sample names in the VCF)
+- `FatherBarcode`: Father's sample identifier (use `0` or `-9` if unknown)
+- `MotherBarcode`: Mother's sample identifier (use `0` or `-9` if unknown)
+- `Sex`: Sex of the sample (optional)
+- `Pheno`: Phenotype information (optional)
+
+Example pedigree file:
+
+```tsv
+FID sample_id FatherBarcode MotherBarcode Sex Pheno
+FAM1 Child1 Father1 Mother1 1 2
+FAM1 Father1 0 0 1 1
+FAM1 Mother1 0 0 2 1
+```
+
+**Output with Pedigree:**
+
+When using `--pedigree`, the output will include additional columns for each parent:
+
+- `father_gt`, `father_dp`, `father_gq`, `father_ad`, `father_vaf`: Father's genotype information
+- `mother_gt`, `mother_dp`, `mother_gq`, `mother_ad`, `mother_vaf`: Mother's genotype information
+
+These columns will contain the parent's genotype data for the same variant, allowing you to analyze inheritance patterns.
+
+## Development
+
+This project uses:
+
+- **UV** for package management
+- **Polars** for fast data processing
+- **Click** for CLI interface
+
+## Testing
+
+Test files are available in the `tests/` directory:
+
+- `test.tabulated.tsv` - Real bcftools output
+- `test_small.tsv` - Small example for quick testing

pywombat-0.1.0.dist-info/RECORD

@@ -0,0 +1,6 @@
+pywombat/__init__.py,sha256=iIPN9vJtsIUhl_DiKNnknxCamLinfayodLLFK8y-aJg,54
+pywombat/cli.py,sha256=kUokXfnaSCKLXiCu7jXbYOPlGTtL5wSzocM9gFtPy30,32801
+pywombat-0.1.0.dist-info/METADATA,sha256=3RlA_lLC7hKUxIrhQvvbBKEolYGOl_EVJgyDfDLI0sU,4982
+pywombat-0.1.0.dist-info/WHEEL,sha256=qtCwoSJWgHk21S1Kb4ihdzI2rlJ1ZKaIurTj_ngOhyQ,87
+pywombat-0.1.0.dist-info/entry_points.txt,sha256=Vt7U2ypbiEgCBlEV71ZPk287H5_HKmPBT4iBu6duEcE,44
+pywombat-0.1.0.dist-info/RECORD,,

pywombat-0.1.0.dist-info/WHEEL

@@ -0,0 +1,4 @@
+Wheel-Version: 1.0
+Generator: hatchling 1.27.0
+Root-Is-Purelib: true
+Tag: py3-none-any

pywombat-0.1.0.dist-info/entry_points.txt

@@ -0,0 +1,2 @@
+[console_scripts]
+wombat = pywombat.cli:cli