sortscore 0.1.0__tar.gz

sortscore-0.1.0/LICENSE

@@ -0,0 +1,21 @@
+MIT License
+
+Copyright (c) 2025 dbaldridge-lab
+
+Permission is hereby granted, free of charge, to any person obtaining a copy
+of this software and associated documentation files (the "Software"), to deal
+in the Software without restriction, including without limitation the rights
+to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
+copies of the Software, and to permit persons to whom the Software is
+furnished to do so, subject to the following conditions:
+
+The above copyright notice and this permission notice shall be included in all
+copies or substantial portions of the Software.
+
+THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
+IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
+FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
+AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
+LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
+OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
+SOFTWARE.

sortscore-0.1.0/PKG-INFO

@@ -0,0 +1,75 @@
+Metadata-Version: 2.4
+Name: sortscore
+Version: 0.1.0
+Summary: A modular Python package for Sort-seq variant analysis
+Author-email: Caitlyn Chitwood <c.chitwood@wustl.edu>
+License: MIT
+Project-URL: Homepage, https://github.com/dbaldridge-lab/sortscore
+Keywords: bioinformatics,sequencing,variant-analysis,sort-seq
+Classifier: Development Status :: 5 - Production/Stable
+Classifier: Intended Audience :: Science/Research
+Classifier: Topic :: Scientific/Engineering :: Bio-Informatics
+Classifier: License :: OSI Approved :: MIT License
+Classifier: Programming Language :: Python :: 3
+Classifier: Programming Language :: Python :: 3.10
+Classifier: Programming Language :: Python :: 3.11
+Classifier: Programming Language :: Python :: 3.12
+Requires-Python: >=3.10
+Description-Content-Type: text/markdown
+License-File: LICENSE
+Requires-Dist: pandas>=2.0.0
+Requires-Dist: numpy>=1.24.0
+Requires-Dist: scipy>=1.10.0
+Requires-Dist: matplotlib>=3.6.0
+Requires-Dist: seaborn>=0.12.0
+Requires-Dist: biopython>=1.81
+Requires-Dist: mavehgvs>=0.7.0
+Dynamic: license-file
+
+# sortscore
+
+`sortscore` is a Python package for Sort-seq variant analysis, including scoring, normalization, and visualization.
+
+## Quick Start
+
+```bash
+python -m venv .venv
+source .venv/bin/activate
+pip install sortscore
+```
+Run a standard variant scoring analysis:
+
+```bash
+sortscore -n EXPERIMENT_NAME -e path/to/experiment_setup.csv -c path/to/config.json
+```
+If the `sortscore` command is not on your `PATH`, run:
+
+```bash
+python -m sortscore -n EXPERIMENT_NAME -e path/to/experiment_setup.csv -c path/to/config.json
+```
+
+## Example Heatmaps
+
+Single Experiment Scoring Heatmap
+
+![Single Experiment Scoring Heatmap](https://raw.githubusercontent.com/dbaldridge-lab/sortscore/main/docs/assets/single_experiment_heatmap.png)
+
+Batch Normalization Heatmap
+
+![Batch Normalization Heatmap](https://raw.githubusercontent.com/dbaldridge-lab/sortscore/main/docs/assets/batch_normalization_heatmap.png)
+
+## Documentation
+- [Installation](https://github.com/dbaldridge-lab/sortscore/blob/main/docs/installation.md)
+- [Usage](https://github.com/dbaldridge-lab/sortscore/blob/main/docs/usage.md)
+- [CLI Arguments](https://github.com/dbaldridge-lab/sortscore/blob/main/docs/cli_arguments.md)
+- [Visualization](https://github.com/dbaldridge-lab/sortscore/blob/main/docs/visualization.md)
+- [Batch Processing](https://github.com/dbaldridge-lab/sortscore/blob/main/docs/batch_processing.md)
+- [Troubleshooting](https://github.com/dbaldridge-lab/sortscore/blob/main/TROUBLESHOOTING.md)
+- [Contributing](https://github.com/dbaldridge-lab/sortscore/blob/main/CONTRIBUTING.md)
+
+## Demo
+
+- [Single Experiment Scoring Notebook Demo](https://github.com/dbaldridge-lab/sortscore/blob/main/demo_data/single_experiment_demo.ipynb)
+- [Batch Normalization Notebook Demo](https://github.com/dbaldridge-lab/sortscore/blob/main/demo_data/tiled_demo.ipynb)
+- [Example Config](https://github.com/dbaldridge-lab/sortscore/blob/main/demo_data/GLI2_oPool5b/config.json)
+- [Example Experiment Setup](https://github.com/dbaldridge-lab/sortscore/blob/main/demo_data/GLI2_oPool5b/experiment_setup.csv)

sortscore-0.1.0/README.md

@@ -0,0 +1,47 @@
+# sortscore
+
+`sortscore` is a Python package for Sort-seq variant analysis, including scoring, normalization, and visualization.
+
+## Quick Start
+
+```bash
+python -m venv .venv
+source .venv/bin/activate
+pip install sortscore
+```
+Run a standard variant scoring analysis:
+
+```bash
+sortscore -n EXPERIMENT_NAME -e path/to/experiment_setup.csv -c path/to/config.json
+```
+If the `sortscore` command is not on your `PATH`, run:
+
+```bash
+python -m sortscore -n EXPERIMENT_NAME -e path/to/experiment_setup.csv -c path/to/config.json
+```
+
+## Example Heatmaps
+
+Single Experiment Scoring Heatmap
+
+![Single Experiment Scoring Heatmap](https://raw.githubusercontent.com/dbaldridge-lab/sortscore/main/docs/assets/single_experiment_heatmap.png)
+
+Batch Normalization Heatmap
+
+![Batch Normalization Heatmap](https://raw.githubusercontent.com/dbaldridge-lab/sortscore/main/docs/assets/batch_normalization_heatmap.png)
+
+## Documentation
+- [Installation](https://github.com/dbaldridge-lab/sortscore/blob/main/docs/installation.md)
+- [Usage](https://github.com/dbaldridge-lab/sortscore/blob/main/docs/usage.md)
+- [CLI Arguments](https://github.com/dbaldridge-lab/sortscore/blob/main/docs/cli_arguments.md)
+- [Visualization](https://github.com/dbaldridge-lab/sortscore/blob/main/docs/visualization.md)
+- [Batch Processing](https://github.com/dbaldridge-lab/sortscore/blob/main/docs/batch_processing.md)
+- [Troubleshooting](https://github.com/dbaldridge-lab/sortscore/blob/main/TROUBLESHOOTING.md)
+- [Contributing](https://github.com/dbaldridge-lab/sortscore/blob/main/CONTRIBUTING.md)
+
+## Demo
+
+- [Single Experiment Scoring Notebook Demo](https://github.com/dbaldridge-lab/sortscore/blob/main/demo_data/single_experiment_demo.ipynb)
+- [Batch Normalization Notebook Demo](https://github.com/dbaldridge-lab/sortscore/blob/main/demo_data/tiled_demo.ipynb)
+- [Example Config](https://github.com/dbaldridge-lab/sortscore/blob/main/demo_data/GLI2_oPool5b/config.json)
+- [Example Experiment Setup](https://github.com/dbaldridge-lab/sortscore/blob/main/demo_data/GLI2_oPool5b/experiment_setup.csv)

sortscore-0.1.0/pyproject.toml

@@ -0,0 +1,43 @@
+[build-system]
+requires = ["setuptools>=68", "wheel"]
+build-backend = "setuptools.build_meta"
+
+[project]
+name = "sortscore"
+version = "0.1.0"
+description = "A modular Python package for Sort-seq variant analysis"
+readme = "README.md"
+requires-python = ">=3.10"
+license = { text = "MIT" }
+authors = [
+  { name = "Caitlyn Chitwood", email = "c.chitwood@wustl.edu" }
+]
+keywords = ["bioinformatics", "sequencing", "variant-analysis", "sort-seq"]
+classifiers = [
+  "Development Status :: 5 - Production/Stable",
+  "Intended Audience :: Science/Research",
+  "Topic :: Scientific/Engineering :: Bio-Informatics",
+  "License :: OSI Approved :: MIT License",
+  "Programming Language :: Python :: 3",
+  "Programming Language :: Python :: 3.10",
+  "Programming Language :: Python :: 3.11",
+  "Programming Language :: Python :: 3.12"
+]
+dependencies = [
+  "pandas>=2.0.0",
+  "numpy>=1.24.0",
+  "scipy>=1.10.0",
+  "matplotlib>=3.6.0",
+  "seaborn>=0.12.0",
+  "biopython>=1.81",
+  "mavehgvs>=0.7.0"
+]
+
+[project.urls]
+Homepage = "https://github.com/dbaldridge-lab/sortscore"
+
+[project.scripts]
+sortscore = "sortscore.cli:main"
+
+[tool.setuptools.packages.find]
+include = ["sortscore", "sortscore.*"]

sortscore-0.1.0/setup.cfg

@@ -0,0 +1,4 @@
+[egg_info]
+tag_build = 
+tag_date = 0
+

sortscore-0.1.0/sortscore/__init__.py

@@ -0,0 +1,19 @@
+"""
+sortscore: A modular Python package for Sort-seq variant analysis.
+
+This package provides tools for analyzing Sort-seq experimental data,
+calculating activity scores, and generating visualizations.
+"""
+
+__version__ = "0.1.0"
+__author__ = "Caitlyn Chitwood"
+__email__ = "c.chitwood@wustl.edu"
+
+# Import main classes for convenience
+from .utils.load_experiment import ExperimentConfig
+from .analysis.score import calculate_full_activity_scores
+
+__all__ = [
+    "ExperimentConfig", 
+    "calculate_full_activity_scores"
+]

sortscore-0.1.0/sortscore/__main__.py

@@ -0,0 +1,7 @@
+"""
+Allow running the sortscore package as a module with: python -m sortscore
+"""
+from sortscore.cli import main
+
+if __name__ == "__main__":
+    main()

sortscore-0.1.0/sortscore/analysis/__init__.py

@@ -0,0 +1,7 @@
+"""
+Analysis subpackage for Sort-seq variant analysis workflows.
+
+This subpackage contains modules for configuration, parameters, I/O, and utilities.
+"""
+
+from .score import calculate_activity_scores

sortscore-0.1.0/sortscore/analysis/aa_scores.py

@@ -0,0 +1,214 @@
+"""
+Amino acid scores processing and export for Sort-seq analysis.
+
+This module provides functions for processing amino acid scores from DNA variant data,
+including aggregation of synonymous codons, statistical analysis, and file export.
+"""
+import os
+import logging
+import pandas as pd
+import numpy as np
+from scipy import stats as scipy_stats
+from typing import List
+from sortscore.analysis.statistics import calculate_codon_and_replicate_variance
+
+
+def _get_score_column_from_avg_method(avg_method: str) -> str:
+    """Return the score column name for a given averaging method."""
+    if avg_method == 'simple-avg':
+        return 'avgscore'
+    return f"avgscore_{avg_method.replace('-', '_')}"
+
+
+def process_and_save_aa_scores(scores_df: pd.DataFrame, experiment, scores_dir: str, analysis_logger) -> None:
+    """
+    Process and save amino acid scores from variant data.
+    
+    This function handles the complete AA scores workflow including:
+    - Filtering out NaN values
+    - Checking if codon aggregation is needed
+    - Calculating appropriate statistics (with/without codon variance)
+    - Rounding score columns
+    - Saving to CSV file
+    - Logging output
+    
+    Parameters
+    ----------
+    scores_df : pd.DataFrame
+        DataFrame containing variant scores and annotations
+    experiment : ExperimentConfig
+        Experiment configuration containing metadata
+    scores_dir : str
+        Directory to save scores file
+    analysis_logger : AnalysisLogger
+        Logger instance for recording outputs
+        
+    Examples
+    --------
+    >>> process_and_save_aa_scores(scores_df, experiment, 'output/scores', 'suffix', logger)
+    """
+    if 'aa_seq_diff' not in scores_df.columns:
+        return
+
+    score_col = _get_score_column_from_avg_method(experiment.avg_method)
+    
+    aa_scores = build_aa_scores_table(scores_df, score_col)
+    
+    # Save to file
+    aa_scores_file = os.path.join(scores_dir, f"{experiment.experiment_name}_aa_scores.csv")
+    aa_scores.to_csv(aa_scores_file, index=False)
+    logging.info(f"Saved AA scores to {aa_scores_file} ({len(aa_scores)} unique AA variants)")
+    
+    # Log file output  
+    analysis_logger.log_output_file(
+        'aa_scores',
+        f"{experiment.experiment_name}_aa_scores.csv",
+        aa_scores_file,
+        variant_count=len(aa_scores)
+    )
+
+
+def build_aa_scores_table(scores_df: pd.DataFrame, score_col: str) -> pd.DataFrame:
+    """
+    Build the AA-level score table from a scored variant table.
+
+    Parameters
+    ----------
+    scores_df : pd.DataFrame
+        DataFrame containing variant scores and annotations.
+    score_col : str
+        Score column used to determine which rows are retained and how codon
+        variance is calculated.
+
+    Returns
+    -------
+    pd.DataFrame
+        AA-level score table with synonymous codons aggregated when needed.
+    """
+    if 'aa_seq_diff' not in scores_df.columns:
+        raise ValueError("scores_df must contain 'aa_seq_diff' column for AA score output")
+
+    scores_df_drop_nan = scores_df.dropna(subset=[score_col])
+    rep_score_columns = [
+        col for col in scores_df_drop_nan.columns
+        if col.startswith('Rep') and col.endswith('.score')
+    ]
+    return _check_codon_num(scores_df_drop_nan, score_col, rep_score_columns)
+
+
+def _check_codon_num(scores_df_drop_nan: pd.DataFrame, score_col: str, 
+                                  rep_score_columns: List[str]) -> pd.DataFrame:
+    """
+    Check if codon aggregation is needed and process AA scores accordingly.
+    
+    This function checks if there are multiple codons per AA variant and processes
+    the data using either DNA->AA aggregation (with codon variance) or AA-only
+    statistics (replicate variance only).
+    
+    Parameters
+    ----------
+    scores_df_drop_nan : pd.DataFrame
+        DataFrame with NaN values already filtered out
+    score_col : str
+        Name of the score column to use
+    rep_score_columns : List[str]
+        List of replicate score column names
+        
+    Returns
+    -------
+    pd.DataFrame
+        Processed AA scores with appropriate statistics
+    """
+    # Check if there are multiple codons per AA variant (DNA->AA case)
+    aa_variant_counts = scores_df_drop_nan.groupby(['aa_seq_diff', 'annotate_aa']).size()
+    needs_aggregation = (aa_variant_counts > 1).any()
+    
+    if needs_aggregation:
+        return _process_dna_to_aa_aggregation(scores_df_drop_nan, score_col, rep_score_columns)
+    else:
+        return _process_aa_only_scores(scores_df_drop_nan, rep_score_columns)
+
+
+def _process_dna_to_aa_aggregation(scores_df_drop_nan: pd.DataFrame, score_col: str, 
+                                  rep_score_columns: List[str]) -> pd.DataFrame:
+    """
+    Process DNA->AA aggregation case with codon variance decomposition.
+    
+    Parameters
+    ----------
+    scores_df_drop_nan : pd.DataFrame
+        DataFrame with variant scores (NaN filtered)
+    score_col : str
+        Name of the score column to use
+    rep_score_columns : List[str]
+        List of replicate score column names
+        
+    Returns
+    -------
+    pd.DataFrame
+        Aggregated AA scores with codon and replicate statistics
+    """
+    # DNA->AA aggregation case: aggregate synonymous variants
+    columns_to_average = ['avgscore', 'avgscore_rep_weighted'] + rep_score_columns
+    
+    # Calculate standard deviation and count of codon-level scores before AA aggregation
+    aa_scores_std = scores_df_drop_nan.groupby(['aa_seq_diff', 'annotate_aa'])[score_col].agg(['std', 'count']).reset_index()
+    aa_scores_std.columns = ['aa_seq_diff', 'annotate_aa', 'SD_codon', 'n_codons']
+    
+    # Calculate mean scores for aggregation
+    aa_scores = scores_df_drop_nan.groupby(['aa_seq_diff', 'annotate_aa'])[columns_to_average].mean().reset_index()
+    
+    # Merge the standard deviation and count of codon scores
+    aa_scores = aa_scores.merge(aa_scores_std, on=['aa_seq_diff', 'annotate_aa'], how='left')
+    
+    # Calculate statistics with codon and replicate variance decomposition
+    aa_scores = calculate_codon_and_replicate_variance(aa_scores, rep_score_columns)
+    
+    return aa_scores
+
+
+def _process_aa_only_scores(scores_df_drop_nan: pd.DataFrame, rep_score_columns: List[str]) -> pd.DataFrame:
+    """
+    Process AA-only case with simple replicate statistics.
+    
+    Parameters
+    ----------
+    scores_df_drop_nan : pd.DataFrame
+        DataFrame with variant scores (NaN filtered)
+    rep_score_columns : List[str]
+        List of replicate score column names
+        
+    Returns
+    -------
+    pd.DataFrame
+        AA scores with replicate statistics only
+    """
+    # AA-only case: no aggregation needed, just copy the data
+    columns_to_include = ['aa_seq_diff', 'annotate_aa', 'avgscore', 'avgscore_rep_weighted'] + rep_score_columns
+    
+    aa_scores = scores_df_drop_nan[columns_to_include].copy()
+    
+    # Calculate simple replicate statistics (no codon variance)
+    if len(rep_score_columns) >= 2:
+        aa_rep_mean = aa_scores[rep_score_columns].mean(axis=1)
+        aa_rep_std = aa_scores[rep_score_columns].std(axis=1, ddof=1)
+        
+        # Calculate n_measurements (just number of non-empty replicates)
+        n_measurements = aa_scores[rep_score_columns].notna().sum(axis=1)
+        
+        # Calculate SEM using only replicate variance
+        sem = aa_rep_std / np.sqrt(n_measurements)
+        
+        # Calculate 95% CI using t-distribution
+        df_actual = n_measurements - 1
+        t_critical = scipy_stats.t.ppf(0.975, df_actual)
+        aa_margin_of_error = t_critical * sem
+        
+        aa_scores['SD_rep'] = aa_rep_std
+        aa_scores['CV_rep'] = (aa_rep_std / aa_rep_mean).round(3)
+        aa_scores['n_measurements'] = n_measurements.astype('Int64')
+        aa_scores['SEM'] = sem
+        aa_scores['CI_lower'] = aa_rep_mean - aa_margin_of_error
+        aa_scores['CI_upper'] = aa_rep_mean + aa_margin_of_error
+    
+    return aa_scores

sortscore-0.1.0/sortscore/analysis/annotation.py

@@ -0,0 +1,188 @@
+"""
+Sequence annotation utilities for Sort-seq variant analysis.
+
+This module provides functions for annotating variant DataFrames with sequence differences,
+translations, and other derived sequence information.
+
+Examples
+--------
+>>> from sortscore.analysis.annotation import annotate_scores_dataframe
+>>> annotated_df = annotate_scores_dataframe(scores_df, experiment)
+"""
+import pandas as pd
+from sortscore.utils.sequence_parsing import compare_to_reference, compare_codon_lists, translate_dna
+
+NO_DIFF_MARKER = '='
+
+
+# TODO: #37 redundant, see if we can remove
+def annotate_scores_dataframe(
+    scores_df: pd.DataFrame, 
+    wt_dna_seq: str, 
+    mutagenesis_type: str = 'aa',
+) -> pd.DataFrame:
+    """
+    Add sequence annotation columns to a scores DataFrame.
+    
+    Parameters
+    ----------
+    scores_df : pd.DataFrame
+        DataFrame with variant sequences and scores.
+    wt_dna_seq : str
+        Wild-type DNA reference sequence.
+    mutagenesis_type : str, default 'aa'
+        Type of mutagenesis ('codon', 'snv', 'aa').
+    
+    Returns
+    -------
+    annotated_df : pd.DataFrame
+        DataFrame with added annotation columns.
+        
+    Examples
+    --------
+    >>> annotated_df = annotate_scores_dataframe(scores_df, wt_seq, 'dna')
+    """
+    df = scores_df.copy()
+    
+    # Check if aa_seq_diff already exists (from pre-annotated data)
+    has_pre_annotated_aa = 'aa_seq_diff' in df.columns
+    
+    # Treat 'dna' as a DNA-sequence variant type (full-length DNA sequences)
+    if mutagenesis_type in {'codon', 'snv'}:
+        # Add codon differences
+        df['codon_diff'] = df['variant_seq'].apply(
+            lambda x: compare_codon_lists(wt_dna_seq, x)
+        )
+        df['codon_diff'] = df['codon_diff'].fillna('')
+        
+        # Add DNA sequence differences
+        df['dna_seq_diff'] = df['variant_seq'].apply(
+            lambda x: compare_to_reference(wt_dna_seq, x, no_change_marker=NO_DIFF_MARKER)
+        )
+        
+        # Add AA sequence annotations only if not pre-annotated
+        if not has_pre_annotated_aa:
+            wt_aa_seq = translate_dna(wt_dna_seq)
+            df['aa_seq'] = df['variant_seq'].apply(translate_dna)
+            df['aa_seq_diff'] = df['aa_seq'].apply(
+                lambda x: compare_to_reference(wt_aa_seq, x, no_change_marker=NO_DIFF_MARKER)
+            )
+        
+    elif mutagenesis_type == 'aa':
+        # For AA variants, add sequence differences only if not pre-annotated
+        if not has_pre_annotated_aa:
+            wt_aa_seq = translate_dna(wt_dna_seq) if len(wt_dna_seq) % 3 == 0 else wt_dna_seq
+            df['aa_seq_diff'] = df['variant_seq'].apply(
+                lambda x: compare_to_reference(wt_aa_seq, x, no_change_marker=NO_DIFF_MARKER)
+            )
+    
+    # Map stop codon representations to * for standard notation in aa_seq_diff column
+    if 'aa_seq_diff' in df.columns:
+        df['aa_seq_diff'] = df['aa_seq_diff'].str.replace('X', '*', regex=False)
+        df['aa_seq_diff'] = df['aa_seq_diff'].str.replace('Ter', '*', regex=False)
+    
+    # Add functional annotations
+    df = add_variant_categories(df)
+    
+    return df
+
+# TODO: #37 isn't this redundant with similar functions
+def add_sequence_differences(
+    df: pd.DataFrame,
+    wt_dna_seq: str,
+    mutagenesis_type: str = 'aa',
+) -> pd.DataFrame:
+    """
+    Add sequence difference columns to a DataFrame.
+    
+    Parameters
+    ----------
+    df : pd.DataFrame
+        DataFrame with variant sequences.
+    wt_dna_seq : str
+        Wild-type DNA sequence.
+    mutagenesis_type : str, default 'aa'
+        Type of mutagenesis ('codon', 'snv', 'aa').
+        
+    Returns
+    -------
+    df : pd.DataFrame
+        DataFrame with sequence difference columns added.
+    """
+    df = df.copy()
+    
+    if mutagenesis_type in {'codon', 'snv'}:
+        # Add DNA sequence differences
+        df['dna_seq_diff'] = df['variant_seq'].apply(
+            lambda x: compare_to_reference(wt_dna_seq, x, no_change_marker=NO_DIFF_MARKER)
+        )
+        
+        # Add AA sequence differences
+        wt_aa_seq = translate_dna(wt_dna_seq)
+        df['aa_seq'] = df['variant_seq'].apply(translate_dna)
+        df['aa_seq_diff'] = df['aa_seq'].apply(
+            lambda x: compare_to_reference(wt_aa_seq, x, no_change_marker=NO_DIFF_MARKER)
+        )
+        
+    elif mutagenesis_type == 'aa':
+        # For AA variants, sequences are already amino acids
+        wt_aa_seq = translate_dna(wt_dna_seq) if len(wt_dna_seq) % 3 == 0 else wt_dna_seq
+        df['aa_seq_diff'] = df['variant_seq'].apply(
+            lambda x: compare_to_reference(wt_aa_seq, x, no_change_marker=NO_DIFF_MARKER)
+        )
+    
+    return df
+
+def classify_aa_variant(aa_diff, dna_diff=None):
+    if aa_diff == NO_DIFF_MARKER:
+        # Check if this is true WT (no DNA changes) or synonymous (DNA changes but same AA)
+        if dna_diff == NO_DIFF_MARKER:
+            return 'wt_dna'
+        else:
+            return 'synonymous'
+    elif '*' in aa_diff:
+        return 'nonsense'
+    else:
+        return 'missense_aa'
+
+def classify_dna_variant(dna_diff, aa_diff):
+    if dna_diff == NO_DIFF_MARKER:
+        return 'wt_dna'
+    elif aa_diff == NO_DIFF_MARKER:
+        return 'synonymous'
+    else:
+        return 'missense_dna'
+
+
+def add_variant_categories(df: pd.DataFrame) -> pd.DataFrame:
+    """
+    Add variant category annotations based on existing sequence difference columns.
+    
+    Parameters
+    ----------
+    df : pd.DataFrame
+        DataFrame with sequence difference columns (aa_seq_diff, dna_seq_diff).
+        
+    Returns
+    -------
+    df : pd.DataFrame
+        DataFrame with variant category columns added.
+    """
+    df = df.copy()
+    
+    # Classify DNA variants 
+    if 'dna_seq_diff' in df.columns:
+        if 'aa_seq_diff' in df.columns:
+            df['annotate_dna'] = df.apply(lambda row: classify_dna_variant(row['dna_seq_diff'], row['aa_seq_diff']), axis=1)
+        else:
+            df['annotate_dna'] = df['dna_seq_diff'].apply(lambda x: 'missense_dna' if x else 'wt_dna')
+    
+    
+    # Classify variants based on AA changes
+    if 'aa_seq_diff' in df.columns:
+        if 'dna_seq_diff' in df.columns:
+            df['annotate_aa'] = df.apply(lambda row: classify_aa_variant(row['aa_seq_diff'], row['dna_seq_diff']), axis=1)
+        else:
+            df['annotate_aa'] = df['aa_seq_diff'].apply(classify_aa_variant)
+
+    return df