Source code for variant_extractor.VariantExtractor
# Copyright 2022 - Barcelona Supercomputing Center
# Author: Rodrigo Martin
# MIT License
from typing import List, Optional
import warnings
import pysam
from .private._utils import compare_contigs, permute_breakend_sv, convert_inv_to_breakend, convert_del_to_ins
from .private._parser import parse_breakend_sv, parse_shorthand_sv, parse_sgl_sv, parse_standard_record
from .private._PendingBreakends import PendingBreakends
from .variants import VariantType
from .variants import VariantRecord
DATAFRAME_COLUMNS = ['start_chrom', 'start', 'end_chrom', 'end', 'ref',
'alt', 'length', 'brackets', 'type_inferred']
DATAFRAME_DTYPES = {'start_chrom': 'category', 'start': 'uint64', 'end_chrom': 'category', 'end': 'uint64', 'ref': 'category',
'alt': 'category', 'length': 'uint64', 'brackets': 'category', 'type_inferred': 'category'}
def _downcast(series):
series_max = series.max()
if series_max < 2 ** 8:
series = series.astype('uint8')
elif series_max < 2 ** 16:
series = series.astype('uint16')
elif series_max < 2 ** 32:
series = series.astype('uint32')
else:
series = series.astype('uint64')
return series
[docs]
class VariantExtractor:
"""
Reads and extracts variants from VCF files. This class is designed to be
used in a pipeline, where the variants are ingested from VCF files and then used in downstream analysis.
"""
[docs]
def __init__(self, vcf_file: str, pass_only=False, ensure_pairs=True, fasta_ref: Optional[str] = None):
"""
Parameters
----------
vcf_file : str
A VCF formatted file. The file is automatically opened.
pass_only : bool, optional
If :code:`True`, only records with PASS filter will be considered.
ensure_pairs : bool, optional
If :code:`True`, throws an exception if a breakend is missing a pair when all other were paired successfully.
fasta_ref : str, optional
A FASTA file with the reference genome. Must be indexed.
"""
self.__ensure_pairs = ensure_pairs
self.__pass_only = pass_only
self.__pairs_found = 0
self.__pending_breakends = PendingBreakends()
self.__fasta_ref = None
# Open FASTA file
if fasta_ref is not None:
self.__fasta_ref = pysam.FastaFile(fasta_ref)
# Open VCF file
vcf_handle = open(file=vcf_file, mode='r')
save = pysam.set_verbosity(0)
self.__variant_file = pysam.VariantFile(vcf_handle)
pysam.set_verbosity(save)
def __iter__(self):
# Read the next record from the VCF file
for rec in self.__variant_file:
yield from self.__handle_record(rec)
# Remove non-PASS records from the pending breakends if pass_only is True
if self.__pass_only:
vcf_records = list(self.__pending_breakends.values())
for vcf_record in vcf_records:
if 'PASS' not in vcf_record.filter:
self.__pending_breakends.remove(vcf_record)
# Only single-paired records or not ensuring pairs
if not self.__ensure_pairs or self.__pairs_found == 0:
for vcf_record in self.__pending_breakends.values():
yield from self.__handle_breakend_individual_sv(vcf_record)
# Found unpaired records
elif len(self.__pending_breakends) > 0:
exception_text = ''
for vcf_record in self.__pending_breakends.values():
exception_text += str(vcf_record)+'\n'
raise Exception(
(f'There are {len(self.__pending_breakends)} unpaired SV breakends. '
'Please, check the entires shown below in VCF file. '
f'Use ensure_pairs=False to ignore unpaired SV breakends.\n{exception_text}'))
def __handle_record(self, rec: pysam.VariantRecord) -> List[VariantRecord]:
if not rec.alts:
return []
if not rec.ref:
raise ValueError('Record does not have a REF field')
# Handle multiallelic records
if len(rec.alts) != 1:
return self.__handle_multiallelic_record(rec)
# Check if breakend SV record
vcf_record = parse_breakend_sv(rec)
if vcf_record:
return self.__handle_breakend_sv(vcf_record)
# Check PASS filter
if self.__pass_only and 'PASS' not in rec.filter:
return []
# Check if shorthand SV record
vcf_record = parse_shorthand_sv(rec)
if vcf_record:
return self.__handle_shorthand_sv(vcf_record)
# Check if single breakend SV record
vcf_record = parse_sgl_sv(rec)
if vcf_record:
return [vcf_record]
# Check if standard record
vcf_record = parse_standard_record(rec)
if vcf_record:
return self.__handle_standard_record(vcf_record)
else:
warnings.warn(f'Skipping unrecognized record:\n{rec}')
return []
def __handle_standard_record(self, vcf_record: VariantRecord) -> List[VariantRecord]:
record_list = []
if len(vcf_record.ref) == len(vcf_record.alt):
for i in range(len(vcf_record.ref)):
# Atomize SNVs
if vcf_record.ref[i] != vcf_record.alt[i]:
variant_id = f'{vcf_record.id}_{i}' if len(vcf_record.ref) > 1 and vcf_record.id is not None \
else vcf_record.id
new_vcf_record = vcf_record._replace(
ref=vcf_record.ref[i], pos=i+vcf_record.pos, end=i+vcf_record.pos, length=1, alt=vcf_record.alt[i], id=variant_id, variant_type=VariantType.SNV)
record_list.append(new_vcf_record)
elif len(vcf_record.ref) > len(vcf_record.alt):
# Deletion
vcf_record.variant_type=VariantType.DEL
record_list.append(vcf_record)
elif len(vcf_record.ref) < len(vcf_record.alt):
# Insertion
vcf_record.variant_type=VariantType.INS
record_list.append(vcf_record)
return record_list
def __handle_breakend_sv(self, vcf_record: VariantRecord) -> List[VariantRecord]:
# Check for pending breakends
previous_record = self.__pending_breakends.pop(vcf_record)
if previous_record is None:
self.__pending_breakends.push(vcf_record)
return []
# Mate breakend found, handle it
self.__pairs_found += 1
return self.__handle_braked_paired_sv(previous_record, vcf_record)
def __handle_braked_paired_sv(self, vcf_record_1: VariantRecord, vcf_record_2: VariantRecord) -> List[VariantRecord]:
# Check PASS filter
if self.__pass_only and ('PASS' not in vcf_record_1.filter or 'PASS' not in vcf_record_2.filter):
return []
# Unify filters
filters = set(vcf_record_1.filter) | set(vcf_record_2.filter)
filters.discard('PASS')
if len(filters) > 0:
vcf_record_1.filter = list(filters)
vcf_record_2.filter = list(filters)
contig_comparison = compare_contigs(vcf_record_1.contig, vcf_record_2.contig)
if contig_comparison == 0:
if vcf_record_1.pos < vcf_record_2.pos:
vcf_record = vcf_record_1
else:
vcf_record = vcf_record_2
elif contig_comparison == -1:
vcf_record = vcf_record_1
else:
vcf_record = vcf_record_2
# Keep the IDs of both VCF records represented by the paired BND.
vcf_record.ids = sorted(vcf_record_1.ids + vcf_record_2.ids)
return self.__handle_breakend_individual_sv(vcf_record)
def __handle_breakend_individual_sv(self, vcf_record: VariantRecord) -> List[VariantRecord]:
assert vcf_record.alt_sv_breakend is not None
contig_comparison = compare_contigs(vcf_record.contig, vcf_record.alt_sv_breakend.contig)
# Transform REF/ALT to equivalent notation so that REF contains the lowest contig and position
if contig_comparison == 1 or (contig_comparison == 0 and vcf_record.pos > vcf_record.alt_sv_breakend.pos):
vcf_record = permute_breakend_sv(vcf_record, self.__fasta_ref)
# Handle DEL notated variant as INS
if vcf_record.length == 1 and vcf_record.variant_type == VariantType.DEL and vcf_record.alt_sv_breakend is not None:
vcf_record = convert_del_to_ins(vcf_record)
if vcf_record.length == 0:
return []
return [vcf_record]
def __handle_shorthand_sv(self, vcf_record: VariantRecord) -> List[VariantRecord]:
if vcf_record.variant_type == VariantType.INV:
# Transform INV into breakend notation
vcf_record_1, vcf_record_2 = convert_inv_to_breakend(vcf_record, self.__fasta_ref)
return [vcf_record_1, vcf_record_2]
else:
return [vcf_record]
def __handle_multiallelic_record(self, rec: pysam.VariantRecord) -> List[VariantRecord]:
record_list = []
fake_rec = rec.copy()
assert fake_rec.alts is not None and len(fake_rec.alts) > 1
alts = fake_rec.alts
samples = dict()
for sample_name in rec.samples:
sample_dict = dict()
for key, value in rec.samples[sample_name].items():
sample_dict[key] = value
samples[sample_name] = sample_dict
original_id = fake_rec.id
for i, alt in enumerate(alts):
# WARNING: This overrides the record
fake_rec.alts = (alt,)
if original_id:
new_id = f'{original_id}_{i}'
fake_rec.id = new_id
new_records = self.__handle_record(fake_rec)
new_samples = dict()
for sample_name in samples:
new_samples[sample_name] = dict()
for key, value in samples[sample_name].items():
if not hasattr(value, '__iter__') or len(value) == self.__variant_file.header.formats[key].number:
new_samples[sample_name][key] = value
else:
if key == 'GT':
new_samples[sample_name][key] = (0, samples[sample_name][key][1])
elif len(value) == len(alts) + 1:
new_samples[sample_name][key] = (value[0], value[i + 1])
elif hasattr(value, '__iter__') and len(value) % len(alts) == 0:
new_samples[sample_name][key] = value[i::len(alts)]
else:
new_samples[sample_name][key] = value
for new_record in new_records:
new_record.samples = new_samples
record_list.extend(new_records)
return record_list
[docs]
@staticmethod
def empty_dataframe(extra_fields=[]):
"""Returns an empty pandas DataFrame with the columns used by this class.
"""
import pandas as pd
df = pd.DataFrame(columns=DATAFRAME_COLUMNS)
df = df.astype(DATAFRAME_DTYPES)
for field in extra_fields:
df[field] = None
return df
[docs]
def to_dataframe(self, extra_fields=[]):
"""Returns a pandas DataFrame with the variants extracted from the VCF file. The columns are:
- start_chrom: chromosome of the start position
- start: start position of the variant
- end_chrom: chromosome of the end position
- end: end position of the variant
- ref: reference allele
- alt: alternative allele
- length: length of the variant (0 for insertions)
- brackets: breakend brackets for breakend SVs (or equivalent for indels or shorthand SVs)
- type_inferred: inferred type of the variant (see VariantType)
Beware of the types of the columns, as they are optimized for memory usage. Use :code:`df.dtypes` to check the types of the columns.
For example, 'start' and 'end' are unsigned integers (of the minimum size possible!), so you must take this into account when performing
operations with these columns. If you want to calculate the distance between two variants, you should first convert the columns to a larger
(signed) integer type to avoid overflow issues.
The DataFrame can be extended with extra fields from the VariantRecord
by passing their names in the extra_fields parameter. For example, passing 'id' will add the id field to the DataFrame.
If :code:`variant_record_obj` is passed in extra_fields, the original VariantRecord object will be added to the DataFrame in a column named 'variant_record_obj'.
"""
import pandas as pd
variants = []
for variant_record in self:
start_chrom = variant_record.contig.replace('chr', '')
start = variant_record.pos
ref = variant_record.ref
alt = variant_record.alt
length = variant_record.length
end = variant_record.end
if variant_record.alt_sv_breakend:
end_chrom = variant_record.alt_sv_breakend.contig.replace('chr', '')
if start_chrom != end_chrom:
end = variant_record.alt_sv_breakend.pos
else:
end_chrom = start_chrom
# Inferred type
type_inferred = variant_record.variant_type.name
# breakends
breakends = ''
if type_inferred == VariantType.DEL.name:
breakends = 'N['
elif type_inferred == VariantType.DUP.name:
breakends = ']N'
elif type_inferred == VariantType.INV.name:
assert variant_record.alt_sv_breakend is not None
prefix = 'N' if variant_record.alt_sv_breakend.prefix else ''
suffix = 'N' if variant_record.alt_sv_breakend.suffix else ''
breakends = prefix + variant_record.alt_sv_breakend.bracket + suffix
elif type_inferred == VariantType.TRA.name:
assert variant_record.alt_sv_breakend is not None
prefix = 'N' if variant_record.alt_sv_breakend.prefix else ''
suffix = 'N' if variant_record.alt_sv_breakend.suffix else ''
breakends = prefix + variant_record.alt_sv_breakend.bracket + variant_record.alt_sv_breakend.bracket + suffix
extra_values = []
for field in extra_fields:
if field == 'variant_record_obj':
extra_values.append(variant_record)
elif hasattr(variant_record, field):
extra_values.append(getattr(variant_record, field))
else:
extra_values.append(None)
variants.append([start_chrom, start, end_chrom, end, ref, alt,
length, breakends, type_inferred] + extra_values)
df = pd.DataFrame(variants, columns=DATAFRAME_COLUMNS + extra_fields)
for col in DATAFRAME_COLUMNS:
df[col] = df[col].astype(DATAFRAME_DTYPES[col])
# Reduce memory usage by using the smallest possible data type for start, end and length
df['start'] = _downcast(df['start'])
df['end'] = _downcast(df['end'])
df['length'] = _downcast(df['length'])
return df