generate_chart_optimized.py

#!/usr/bin/env python3
"""
ClinVar Stacked Area Chart Generator - Optimized Version
Processes ClinVar variant data from 2015-2025
"""

import os
import pandas as pd
import matplotlib.pyplot as plt
from pathlib import Path
from datetime import datetime
import json

# Ensure we're in the correct directory
os.chdir(Path(__file__).parent)

# Configuration
DATA_DIR = "data"
YEARS = list(range(2015, datetime.now().year + 1))  # 2015 to current year

# Color scheme from original notebook
COLORS = {
    'Benign': '#00008B',              # Dark blue for benign
    'Likely benign': '#ADD8E6',       # Light blue for likely benign
    'Likely pathogenic': '#FF6347',   # Light red for likely pathogenic
    'Pathogenic': '#8B0000',          # Dark red for pathogenic
    'Uncertain significance': '#A9A9A9'  # Dark grey for VUS
}

def get_cached_data(year_month):
    """Check if cached JSON exists and load it"""
    json_cache_path = Path(DATA_DIR) / f"variant_summary_{year_month}.json"
    if json_cache_path.exists():
        try:
            with open(json_cache_path, 'r') as f:
                return json.load(f)
        except:
            return None
    return None

def save_cached_data(year_month, data):
    """Save data to JSON cache"""
    json_cache_path = Path(DATA_DIR) / f"variant_summary_{year_month}.json"
    # Convert numpy int64 to regular int for JSON serialization
    json_data = {k: int(v) for k, v in data.items()}
    with open(json_cache_path, 'w') as f:
        json.dump(json_data, f, indent=2)

def process_year(year):
    """Process a single year's data using chunked reading of compressed files"""
    current_month = datetime.now().strftime("%m")
    year_month = f"{year}-{current_month}"
    filename = f"variant_summary_{year_month}.txt.gz"
    filepath = Path(DATA_DIR) / filename

    # Check if cached results exist
    cached_data = get_cached_data(year_month)
    if cached_data:
        print(f"Using cached data for {year_month}: {sum(cached_data.values()):,} variants")
        return cached_data

    if not filepath.exists():
        print(f"Warning: {filepath} not found, skipping {year}")
        return None

    print(f"Processing {year_month}...", end=" ")

    # Categories to count
    categories = ['Pathogenic', 'Likely pathogenic', 'Uncertain significance', 'Likely benign', 'Benign']
    counts = {cat: 0 for cat in categories}

    try:
        # Read compressed files in chunks to manage memory
        chunk_size = 50000
        total_processed = 0

        for chunk in pd.read_csv(filepath, sep='\t', chunksize=chunk_size,
                                usecols=['Type', 'ClinicalSignificance', 'Stop', 'Start'],
                                low_memory=False, compression='gzip'):

            # Filter for indels and SNVs
            filtered = chunk[chunk['Type'].str.contains('indel|single nucleotide variant', case=False, na=False)]

            # Filter out variants > 50bp
            filtered = filtered[filtered['Stop'] - filtered['Start'] <= 50]

            # Count clinical significance categories
            for category in categories:
                counts[category] += filtered['ClinicalSignificance'].eq(category).sum()

            total_processed += len(chunk)

        total_variants = sum(counts.values())
        print(f"{total_variants:,} variants")

        # Save to cache
        save_cached_data(year_month, counts)

        return counts

    except Exception as e:
        print(f"Error processing {year}: {e}")
        return None

def create_chart(data_dict, output_filename):
    """Create the stacked area chart"""
    years = sorted(data_dict.keys())
    categories = ['Benign', 'Likely benign', 'Likely pathogenic', 'Pathogenic', 'Uncertain significance']

    # Prepare data
    y_data = {cat: [] for cat in categories}
    for year in years:
        year_data = data_dict[year]
        for cat in categories:
            y_data[cat].append(year_data.get(cat, 0))

    # Create plot
    plt.figure(figsize=(8, 6))

    plot_colors = [COLORS.get(cat, '#7f7f7f') for cat in categories]
    y_arrays = [y_data[cat] for cat in categories]

    plt.stackplot(years, *y_arrays, colors=plot_colors, labels=categories, alpha=0.7)

    # Formatting
    plt.legend(loc='upper left', frameon=True, fancybox=True, shadow=True)
    plt.xlabel('Year', fontsize=12)
    plt.ylabel('Number of Variants', fontsize=12)
    end_year = datetime.now().year
    plt.title(f'ClinVar Variants by Clinical Significance (2015-{end_year})', fontsize=14, fontweight='bold')

    # Format y-axis
    ax = plt.gca()
    max_val = max([sum(y_data[cat][i] for cat in categories) for i in range(len(years))])
    if max_val > 1000000:
        ax.yaxis.set_major_formatter(plt.FuncFormatter(lambda x, p: f'{x/1e6:.1f}M'))
        plt.ylabel('Number of Variants (Millions)', fontsize=12)

    plt.xticks(years, rotation=45)
    plt.grid(True, alpha=0.3)
    plt.tight_layout()

    # Save chart
    plt.savefig(output_filename, dpi=300, format='png', bbox_inches='tight')
    print(f"\nChart saved as: {output_filename}")

    return plt.gcf()

def print_summary(data_dict):
    """Print summary statistics"""
    print("\n" + "="*60)
    print("CLINVAR VARIANT SUMMARY")
    print("="*60)

    for year in sorted(data_dict.keys()):
        if data_dict[year]:
            total = sum(data_dict[year].values())
            print(f"{year}: {total:,} variants")

    # Growth calculation
    years = sorted([y for y in data_dict.keys() if data_dict[y]])
    if len(years) >= 2:
        first_total = sum(data_dict[years[0]].values())
        last_total = sum(data_dict[years[-1]].values())
        growth = last_total / first_total if first_total > 0 else 0
        print(f"\nGrowth from {years[0]} to {years[-1]}: {growth:.1f}x")


def main():
    """Main function"""

    print("Processing data files...")
    # Process all years
    all_data = {}
    for year in YEARS:
        counts = process_year(year)
        all_data[year] = counts

    # Generate final chart with all data
    output_file = "clinvar_chart.png"
    create_chart(all_data, output_file)
    
if __name__ == "__main__":
    main()