First, thank you for updating this code!! Its a labor of love. I'm trying to make a script to refresh the HLA definitions, there has been quite a lot done in the last 12 years, could you include something like this (its not working, but I made it based on prev version):
from __future__ import division
from __future__ import print_function
from builtins import str
from builtins import map
import pandas as pd
import numpy as np
from collections import OrderedDict
from datetime import datetime
import os
import sys
from Bio import SeqIO
from Bio.SeqRecord import SeqRecord
from Bio.Seq import Seq
# 1. Add your exact Micromamba bin path to Python's searchable paths
#OPTITYPE_BIN_PATH = "/root/micromamba/envs/ngs_optitype/bin/"
#sys.path.append(OPTITYPE_BIN_PATH)
# Import OptiType's internal database generator
#try:
# from hlatyper import create_allele_dataframes
#except ImportError:
# raise ImportError("Cannot find 'hlatyper'. Ensure you are running this within the OptiType environment/container.")
# Configuration
HLA_DAT = "hla.dat"
DNA_OUT = "hla_reference_dna.fasta"
RNA_OUT = "hla_reference_rna.fasta"
HDF_OUT = "alleles.h5"
## optitype 1.5 compatability
FTRS_OUT = "alleles/features.parquet"
TABL_OUT = "alleles/table.parquet"
TARGET_LOCI = ["HLA-A", "HLA-B", "HLA-C"]
VERBOSE = False
def now(start=datetime.now()):
# function argument NOT to be ever set! It gets initialized at function definition time
# so we can calculate difference without further hassle.
return str(datetime.now() - start)[:-4]
def store_dataframes(out_hdf, **kwargs):
# DataFrames to serialize have to be passed by keyword arguments. An argument matrix1=DataFrame(...)
# will be written into table 'matrix1' in the HDF file.
complevel = kwargs.pop('complevel', 9) # default complevel & complib values if
complib = kwargs.pop('complib', 'zlib') # not explicitly asked for as arguments
if VERBOSE:
print(now(), 'Storing %d DataFrames in file %s with compression settings %d %s...' % (len(kwargs), out_hdf, complevel, complib))
# mode='w' so we always write a fresh file instead of appending to a stale one.
store = pd.HDFStore(out_hdf, mode='w', complevel=complevel, complib=complib)
for table_name, dataframe in kwargs.items():
store[table_name] = dataframe
store.close()
if VERBOSE:
print(now(), 'DataFrames stored in file.')
def feature_order(feature):
# to use it in sorted(..., key=feature_order) this returns:
# 0 for ('UTR', 1)
# 2 for exon, 1
# 3 for intron, 1
# 4 for exon, 2
# 5 for intron, 2
# ...
# 999 for UTR, 2
feat_type, feat_number = feature
assert feat_type in ('intron', 'exon', 'UTR'), 'Unknown feature in list (function accepts intron/exon/UTR)'
assert isinstance(feat_number, int), 'Feature number has to be integer'
return 0 if feature == ('UTR', 1) else 999 if feature == ('UTR', 2) else (feat_number * 2 + (feat_type == 'intron'))
def create_allele_dataframes(imgt_dat, fasta_gen, fasta_nuc):
from Bio import SeqIO
if VERBOSE:
print(now(), 'Loading IMGT allele dat file...')
alleles = OrderedDict()
with open(imgt_dat, 'r') as handle:
for i, record in enumerate(SeqIO.parse(handle, "imgt")):
# TODO: IMGT has changed the ID system. Now it's HLA00001.1 or HLA12345.2 showing versions. I'll get rid of this now though
record.id = record.id.split('.')[0]
alleles[record.id] = record
if VERBOSE:
print(now(), 'Initializing allele DataFrame...')
'''
id HLA000001
type A*01:01:01:01
locus A
class I --- I, II, TAP, MIC, other
flags 0 --- problems with allele stored here, see below
len_gen 3503
len_nuc 1098
full_gen 1
full_nuc 1
flagging: sum of the below codes
+1 if HLA type ends in a letter (N, Q, etc.)
+2 if CDS annotation != exon annotation
+4 if features don't add up to gen sequence
+8 if exons don't add up to nuc sequence
'''
allele_info = 'id type 4digit locus flags len_dat len_gen len_nuc full_gen full_nuc'
table = pd.DataFrame(index=list(alleles.keys()), columns=allele_info.split())
sequences = []
if VERBOSE:
print(now(), 'Filling DataFrame with allele data...')
all_features = [] # contains tuples: (HLA id, feature type, feature number, feature start, feature end)
for allele in alleles.values():
allele_type = allele.description.replace('HLA-', '').split(',')[0]
table.loc[allele.id, 'id'] = allele.id
table.loc[allele.id, 'type'] = allele_type
table.loc[allele.id, '4digit'] = ':'.join(allele_type.split(':')[:2])
table.loc[allele.id, 'locus'] = allele_type.split('*')[0]
table.loc[allele.id, 'flags'] = 0 if allele_type[-1].isdigit() else 1
table.loc[allele.id, 'len_dat'] = len(str(allele.seq))
table.loc[allele.id, 'len_gen'] = 0 # TODO: IT STILL DOESNT SOLVE PERFORMANCEWARNING!
table.loc[allele.id, 'len_nuc'] = 0 # we initialize these nulls so that we don't get a
table.loc[allele.id, 'full_gen'] = 0 # PerformanceWarning + pickling when storing HDF
table.loc[allele.id, 'full_nuc'] = 0 # because of NaNs (they don't map to ctypes)
sequences.append((allele.id, 'dat', str(allele.seq)))
# number of features in 2013-04-30 hla.dat:
# 9296 source (total # of alleles)
# 9291 CDS, 5 gene (HLA-P pseudogenes)
# 24493 exons, 3697 introns, 1027 UTRs
# CDS matches exons in 99+% of cases, some have a single base-pair extension at the end for some
# weird single-base exons or such (all on unimportant pseudogene loci)
# so we extract exons, introns and UTRs
features = [f for f in allele.features if f.type in ('exon', 'intron', 'UTR')]
for feature in features:
if feature.type in ('exon', 'intron'):
feature_num = int(feature.qualifiers['number'][0])
else: # UTR
# UTR either starts at the beginning or ends at the end
assert feature.location.start == 0 or feature.location.end == len(allele.seq)
feature_num = 1 if feature.location.start == 0 else 2 # 1 if 5' UTR, 2 if 3' UTR
all_features.append(
(allele.id, feature.type, feature_num,
int(feature.location.start), int(feature.location.end), len(feature), feature_order((feature.type, feature_num)))
)
# small sanity check, can be commented out
cds = [f for f in allele.features if f.type == 'CDS']
if cds:
if sum(map(len, [f for f in features if f.type == 'exon'])) != len(cds[0]):
if VERBOSE:
print("\tCDS length doesn't match sum of exons for", allele.id, allele_type)
table.loc[allele.id, 'flags'] += 2
else:
if VERBOSE:
print("\tNo CDS found for", allele.id, allele_type)
table.loc[allele.id, 'flags'] += 2
if VERBOSE:
print(now(), 'Loading gen and nuc files...')
# The gen/nuc FASTA records may be headed either by the OptiType accession scheme
# ("HLA:HLA00001" -> "HLA00001", which is the table index) or, as produced by
# parse_imgt_dat() here, by the allele type ("HLA-A*01:01:01:01"). Build a lookup
# from allele type to accession id so we can resolve the latter back to the index.
type_to_id = {t: i for i, t in zip(table.index, table['type'])}
def resolve_allele_id(record_id):
accession = record_id.replace('HLA:', '')
if accession in table.index:
return accession
return type_to_id.get(record_id.replace('HLA-', ''))
with open(fasta_gen, 'r') as fasta_gen:
for record in SeqIO.parse(fasta_gen, 'fasta'):
allele_id = resolve_allele_id(record.id)
if allele_id is None:
if VERBOSE:
print("\tNo matching allele for gen record", record.id)
continue
table.loc[allele_id, 'len_gen'] = len(record.seq)
sequences.append((allele_id, 'gen', str(record.seq)))
with open(fasta_nuc, 'r') as fasta_nuc:
for record in SeqIO.parse(fasta_nuc, 'fasta'):
allele_id = resolve_allele_id(record.id)
if allele_id is None:
if VERBOSE:
print("\tNo matching allele for nuc record", record.id)
continue
table.loc[allele_id, 'len_nuc'] = len(record.seq)
sequences.append((allele_id, 'nuc', str(record.seq)))
# convert list of tuples into DataFrame for features and sequences
all_features = pd.DataFrame(all_features, columns=['id', 'feature', 'number', 'start', 'end', 'length', 'order'])
sequences = pd.DataFrame(sequences, columns=['id', 'source', 'sequence'])
joined = pd.merge(table, all_features, how='inner', on='id')
exons_for_locus = {}
for i_locus, i_group in joined.groupby('locus'):
exons_for_locus[i_locus] = i_group[i_group['feature'] == 'exon']['number'].max()
if VERBOSE:
print(now(), 'Checking dat features vs gen/nuc sequences...')
for allele, features in joined.groupby('id'):
row = features.iloc[0] # first row of the features subtable. Contains all allele information because of the join
sum_features_length = features['length'].sum()
sum_exons_length = features.loc[features['feature'] == 'exon']['length'].sum()
if row['len_gen'] > 0 and row['len_gen'] != sum_features_length:
if VERBOSE:
print("\tFeature lengths don't add up to gen sequence length", allele, row['len_gen'], sum_features_length, row['type'])
table.loc[allele, 'flags'] += 4
if row['len_nuc'] > 0 and row['len_nuc'] != sum_exons_length:
if VERBOSE:
print("\tExon lengths don't add up to nuc sequence length", allele, row['len_nuc'], sum_exons_length, row['type'])
table.loc[allele, 'flags'] += 8
if VERBOSE:
print(now(), 'Sanity check finished. Computing feature sequences...')
ft_seq_lookup = OrderedDict()
ft_seq_lookup['---DUMMY---'] = 0 # it will be useful later on if 0 isn't used. lookup_id*boolean operation, etc.
ft_counter = 1
all_ft_counter = 0
all_features['seq_id'] = 0
for i_id, i_features in all_features.groupby('id'):
seq = sequences.loc[(sequences['id'] == i_id) & (sequences['source'] == 'dat')].iloc[0]['sequence']
for ft_idx, feature in i_features.iterrows():
ft_seq = seq[feature['start']:feature['end']]
all_ft_counter += 1
if ft_seq not in ft_seq_lookup:
ft_seq_lookup[ft_seq] = ft_counter
all_features.loc[ft_idx, 'seq_id'] = ft_counter
ft_counter += 1
else:
all_features.loc[ft_idx, 'seq_id'] = ft_seq_lookup[ft_seq]
feature_sequences = pd.DataFrame([seq for seq in list(ft_seq_lookup.keys())], columns=['sequence']) # , index=ft_seq_lookup.values() but it's 0,1,2,... anyway, as the default
return table, all_features, sequences, feature_sequences
def parse_imgt_dat():
dna_records = []
rna_records = []
print(f"Parsing {HLA_DAT}... This may take a moment.")
# Parse the EMBL formatted hla.dat file
for record in SeqIO.parse(HLA_DAT, "imgt"):
# The description usually looks like "HLA-A*01:01:01:01"
allele_name = record.description.split(",")[0].strip()
locus = allele_name.split("*")[0]
if locus not in TARGET_LOCI:
continue
exons = {}
introns = {}
# Map out the exact boundaries of exons and introns
for feature in record.features:
if feature.type == "exon":
if "number" in feature.qualifiers:
exon_num = feature.qualifiers["number"][0]
exons[exon_num] = feature.extract(record.seq)
elif feature.type == "intron":
if "number" in feature.qualifiers:
intron_num = feature.qualifiers["number"][0]
introns[intron_num] = feature.extract(record.seq)
# 1. Build RNA Sequence (Exon 2 + Exon 3)
if "2" in exons and "3" in exons:
rna_seq = exons["2"] + exons["3"]
rna_records.append(
SeqRecord(rna_seq, id=allele_name, description="Exon2+Exon3")
)
# 2. Build DNA Sequence (Intron 1 + Exon 2 + Intron 2 + Exon 3 + Intron 3)
# Note: Many alleles in IMGT lack full intronic data. We skip them if missing.
if all(k in introns for k in ["1", "2", "3"]) and all(k in exons for k in ["2", "3"]):
dna_seq = introns["1"] + exons["2"] + introns["2"] + exons["3"] + introns["3"]
dna_records.append(
SeqRecord(dna_seq, id=allele_name, description="Intron1-Intron3")
)
print(f"Extracted {len(rna_records)} RNA records and {len(dna_records)} DNA records.")
return dna_records, rna_records
def main():
if not os.path.exists(HLA_DAT):
print(f"Error: {HLA_DAT} not found. Please download it from the IPD-IMGT/HLA Github.")
return
# Extract sequences
dna_recs, rna_recs = parse_imgt_dat()
# Write trimmed sequences to FASTA
print(f"Writing {DNA_OUT}...")
SeqIO.write(dna_recs, DNA_OUT, "fasta")
print(f"Writing {RNA_OUT}...")
SeqIO.write(rna_recs, RNA_OUT, "fasta")
# Call OptiType's internal matrix generator
print("Generating new alleles.h5 hit matrix...")
# create_allele_dataframes expects: hla.dat path, dna fasta path, rna fasta path
table, all_features, sequences, feature_sequences = create_allele_dataframes(HLA_DAT, DNA_OUT, RNA_OUT)
# Serialize the four DataFrames into the HDF store OptiType loads at runtime.
print(f"Writing {HDF_OUT}...")
store_dataframes(
HDF_OUT,
table=table,
features=all_features,
sequences=sequences,
feature_sequences=feature_sequences,
)
## save parquet table and features (OptiType 1.5 compatibility)
os.makedirs(os.path.dirname(FTRS_OUT), exist_ok=True)
print(f"Writing {FTRS_OUT}...")
all_features.to_parquet(FTRS_OUT, index=False)
print(f"Writing {TABL_OUT}...")
table.to_parquet(TABL_OUT, index=True)
print("Update complete! You can now move alleles.h5, trimmed_gen.fasta, and trimmed_nuc.fasta into your OptiType/data/ folder.")
if __name__ == "__main__":
main()
First, thank you for updating this code!! Its a labor of love. I'm trying to make a script to refresh the HLA definitions, there has been quite a lot done in the last 12 years, could you include something like this (its not working, but I made it based on prev version):