<jupyter_start><jupyter_text>Precursors<jupyter_code>gtf_file = '/home/drk/common/data/genomes/hg38/genes/gencode36/gencode.v36.basic.annotation.gtf.gz'
utr3_fasta = '/home/drk/rnaml/analysis/1-27/V36.3UTRs.selectedTranscript.fa'
utr5_fasta = '/home/drk/rnaml/analysis/1-27/V36.5UTRs.selectedTranscript.fa'
orf_fasta = '/home/drk/rnaml/analysis/1-27/V36.ORFs.selectedTranscript.fa'

stability_file = '/home/vagar/RNAbinding/transcript_isoform_analysis/Rinn_avg_K562.txt'
features_file = '/home/vagar/RNAbinding/half_life/human/half_life_features.txt'
def gtf_kv(s):
    """Convert the last gtf section of key/value pairs into a dict."""
    d = {}

    a = s.split(';')
    for key_val in a:
        if key_val.strip():
            eq_i = key_val.find('=')
            if eq_i != -1 and key_val[eq_i-1] != '"':
                kvs = key_val.split('=')
            else:
                kvs = key_val.split()

            key = kvs[0]
            if kvs[1][0] == '"' and kvs[-1][-1] == '"':
                val = (' '.join(kvs[1:]))[1:-1].strip()
            else:
                val = (' '.join(kvs[1:])).strip()

            d[key] = val

    return d
def trim_dot(gene_id):
    """Trim the .xyz suffix from a GENCODE gene_id."""
    dot_i = gene_id.find('.')
    if dot_i != -1:
        gene_id = gene_id[:dot_i]
    return gene_id
def read_gene_fasta(fasta_file, trim_geneid=False):
    """Read gene sequences from FASTA into a dict."""
    if fasta_file[-3:] == '.gz':
        fasta_open = gzip.open(fasta_file, 'rt')
    else:
        fasta_open = open(fasta_file, 'r')
        
    gene_seq = {}
    for line in fasta_open:
        if line[0] == '>':
            gene = line[1:].rstrip()
            if trim_geneid:
                gene = trim_dot(gene)
            gene_seq[gene] = ''
        else:
            gene_seq[gene] += line.rstrip()
            
    fasta_open.close()
        
    return gene_seq<jupyter_output><empty_output><jupyter_text>Read stability measurements<jupyter_code># read sample table
stability_df = pd.read_csv(stability_file, sep='\t', names=['Gene','Stability'])
# additional gene attributes

gene_type = {}
gene_name = {}
for line in gzip.open(gtf_file, 'rt'):
    a = line.split('\t')
    if len(a) > 8 and a[2] == 'gene':
        kv = gtf_kv(a[8])
        gene_id = trim_dot(kv['gene_id'])
        gene_type[gene_id] = kv['gene_type']
        gene_name[gene_id] = kv['gene_name']
    
stability_type = []
stability_name = []
for gene_id in stability_df.Gene:
    if gene_id in gene_type:
        stability_type.append(gene_type[gene_id])
        stability_name.append(gene_name[gene_id])
    else:
        stability_type.append('depracated')
        stability_name.append('NA')
    
stability_df['Name'] = stability_name
stability_df['Type'] = stability_type
stability_df = stability_df[['Gene','Name','Type','Stability']]
stability_df.head()
# count gene types
stability_df.value_counts('Type')
# plot stability for mRNA versus rest
sns.displot(x=np.sqrt(stability_df.Stability),
            hue=(stability_df.Type=='protein_coding'),
            kind='kde', aspect=2, height=4)
# keep only mRNAs
stability_df = stability_df[stability_df.Type=='protein_coding']
del stability_df['Type']<jupyter_output><empty_output><jupyter_text>Read UTRs<jupyter_code>gene_utr3 = read_gene_fasta(utr3_fasta)
print(len(gene_utr3))

gene_utr5 = read_gene_fasta(utr5_fasta)
print(len(gene_utr5))

gene_orf = read_gene_fasta(orf_fasta)
print(len(gene_orf))
# construct gene table, with stability and transcript segment lengths

stability_utr3 = []
stability_utr5 = []
stability_orf = []

gene_seq = {}

for gene in stability_df.Gene:
    if gene in gene_utr5:
        gene_seq[gene] = gene_utr5[gene]
        stability_utr5.append(len(gene_utr5[gene]))
    else:
        stability_utr5.append(0)
    if gene in gene_orf:
        gene_seq[gene] = gene_seq.get(gene,'') + gene_orf[gene]
        stability_orf.append(len(gene_orf[gene]))
    else:
        stability_orf.append(0)
    if gene in gene_utr3:
        gene_seq[gene] = gene_seq.get(gene,'') + gene_utr3[gene]
        stability_utr3.append(len(gene_utr3[gene]))
    else:
        stability_utr3.append(0)
        
stability_df['UTR5'] = stability_utr5
stability_df['ORF'] = stability_orf
stability_df['UTR3'] = stability_utr3
stability_df['Length'] = stability_df['UTR5'] + stability_df['ORF'] + stability_df['UTR3']
stability_df.shape
# attach additional features

features_df = pd.read_csv(features_file, sep='\t')
features_df.set_index('ENSID',inplace=True)
features_df.head()
# plot exon distributions
sns.displot(np.sqrt(features_df.ORFEXONDENSITY), aspect=2)
# choose threshold
(np.sqrt(features_df.ORFEXONDENSITY) < 4.5).mean()
# clip and attach features

stability_exons = []

for gene in stability_df.Gene:
    if gene in features_df.index:
        stability_exons.append(features_df.loc[gene].ORFEXONDENSITY)
    else:
        stability_exons.append(np.nan)

stability_df['Exons'] = np.clip(np.sqrt(stability_exons), 0, 4.5)
# check
stability_df.loc[np.isnan(stability_df.Exons)].head()
# set NaNs to 0
stability_df.Exons.fillna(np.nanmean(stability_df.Exons), inplace=True)<jupyter_output><empty_output><jupyter_text>QC<jupyter_code># check for zero length ORF genes
stability_df[stability_df.ORF==0]<jupyter_output><empty_output><jupyter_text>PIGY/PYURF seems to be a case of two separate genes that should be considered isoforms of the same gene.<jupyter_code># remove zero length ORF genes
stability_df = stability_df[stability_df.ORF>0]
# check for zero length 3' UTR genes
print(stability_df[stability_df.UTR3==0].shape)
stability_df[stability_df.UTR3==0].sort_values('Stability', ascending=False).head(8)
# check for zero length 5' UTR genes
print(stability_df[stability_df.UTR5==0].shape)
stability_df[stability_df.UTR5==0].sort_values('Stability', ascending=False).head(8)<jupyter_output>(167, 8)<jupyter_text>Choose length thresholds<jupyter_code># plot 5'UTR distribution
sns.displot(np.log2(1+stability_df.UTR5))
# choose threshold
utr5_len_t = 2**11
utr5_len_t, (stability_df.UTR5<utr5_len_t).mean()
# plot 3'UTR distribution
sns.displot(np.log2(1+stability_df.UTR3))
# choose threshold
utr3_len_t = 3*2**12
utr3_len_t, np.log2(utr3_len_t), (stability_df.UTR3<utr3_len_t).mean()
# plot ORF distribution
sns.displot(np.log2(1+stability_df.ORF))
# choose threshold
cds_len_t = 3*2**12
cds_len_t, np.log2(cds_len_t), (stability_df.ORF<cds_len_t).mean()
# plot full length distribution
sns.displot(np.log2(1+stability_df.Length))
# choose threshold
full_len_t = 2**14
full_len_t, np.log2(full_len_t), (stability_df.Length<full_len_t).mean()<jupyter_output><empty_output><jupyter_text>Training data<jupyter_code># divide into folds

num_folds = 5
num_genes = stability_df.shape[0]
fold_genes = num_genes // num_folds

folds_num = [fold_genes]*(num_folds-1)
fold_final_genes = num_genes - fold_genes*(num_folds-1)
folds_num.append(fold_final_genes)

gene_folds = []
for fi in range(num_folds):
    gene_folds += [fi]*folds_num[fi]
gene_folds = np.array(gene_folds)
np.random.shuffle(gene_folds)

fold_masks = []
for fi in range(num_folds):
    fold_masks.append(gene_folds == fi)
def feature_bytes(values):
    """Convert numpy arrays to bytes features."""
    values = values.flatten().tostring()
    return tf.train.Feature(bytes_list=tf.train.BytesList(value=[values]))

def tfr_feature_ints(int_list):
    return tf.train.Feature(int64_list=tf.train.Int64List(value=int_list))
def tfr_feature_dna(seq_1hot):
    seq_indexes = np.where(seq_1hot)[1].astype('uint8')
    seq_indexes = seq_indexes.tostring()
    return tf.train.Feature(bytes_list=tf.train.BytesList(value=[seq_indexes]))
def tfr_feature_targets(targets):
    targets = targets.astype('float16')
    targets = targets.tostring()
    return tf.train.Feature(bytes_list=tf.train.BytesList(value=[targets]))
def write_fold(fold_index, out_dir, seqs_per_tfr=128, pseudocount=0.125):
    """Write a cross-validation fold's TFRecords."""
    tf_opts = tf.io.TFRecordOptions(compression_type='ZLIB')
    
    # slice fold's genes
    fold_df = stability_df[fold_masks[fold_index]]
    
    gi = 0
    num_tfr = int(np.ceil(fold_df.shape[0] / seqs_per_tfr))
    for fi in range(num_tfr):
        tfr_file = '%s/fold%d-%d.tfr' % (out_dir,fold_index,fi)
        with tf.io.TFRecordWriter(tfr_file, tf_opts) as writer:
            for ti in range(seqs_per_tfr):
                if gi < fold_df.shape[0]:
                    sgene = fold_df.iloc[gi]
                    gene_id = sgene.Gene

                    # utr5
                    if gene_id in gene_utr5:        
                        gene_utr5_1hot = dna_io.dna_1hot(gene_utr5[gene_id])
                        if gene_utr5_1hot.shape[0] > utr5_len_t:
                            gene_utr5_1hot = gene_utr5_1hot[-utr5_len_t:]
                    else:
                        gene_utr5_1hot = np.zeros((0,4), dtype='bool')

                    # cds
                    if gene_id in gene_orf:        
                        gene_cds_1hot = dna_io.dna_1hot(gene_orf[gene_id])
                        if gene_cds_1hot.shape[0] > cds_len_t:
                            li = (gene_cds_1hot.shape[0] - cds_len_t) // 2
                            gene_cds_1hot = gene_cds_1hot[li:li+cds_len_t]
                    else:
                        gene_cds_1hot = np.zeros((0,4), dtype='bool')

                    # utr3
                    if gene_id in gene_utr3:        
                        gene_utr3_1hot = dna_io.dna_1hot(gene_utr3[gene_id])
                        if gene_utr3_1hot.shape[0] > utr3_len_t:
                            li = (gene_utr3_1hot.shape[0] - utr3_len_t)
                            gene_utr3_1hot = gene_utr3_1hot[:utr3_len_t]
                    else:
                        gene_utr3_1hot = np.zeros((0,4), dtype='bool')
                        
                    # features
                    gene_features = np.array([sgene.Exons])

                    # targets
                    targets = np.log2(sgene.Stability + pseudocount)

                    # hash to bytes
                    seq_lengths = [gene_utr5_1hot.shape[0], gene_cds_1hot.shape[0], gene_utr3_1hot.shape[0]]
                    features_dict = {
                        'lengths': tfr_feature_ints(seq_lengths),
                        'utr3': tfr_feature_dna(gene_utr3_1hot),
                        'cds': tfr_feature_dna(gene_cds_1hot),
                        'utr5': tfr_feature_dna(gene_utr5_1hot),
                        'features': tfr_feature_targets(gene_features),
                        'targets': tfr_feature_targets(targets)
                    }

                    # write example
                    example = tf.train.Example(features=tf.train.Features(feature=features_dict))
                    writer.write(example.SerializeToString())
                
                    # next gene
                    gi += 1
################################################################
# output directory

out_dir = 'data_cv'
if os.path.isdir(out_dir):
    shutil.rmtree(out_dir)
os.mkdir(out_dir)

################################################################
# stats
stats_dict = {}
stats_dict['num_targets'] = 1
stats_dict['target_length'] = 1
stats_dict['num_features'] = 1

stats_dict['length_utr5'] = utr5_len_t
stats_dict['length_cds'] = cds_len_t
stats_dict['length_utr3'] = utr3_len_t
stats_dict['length_full'] = full_len_t

stats_dict['num_folds'] = num_folds
for fi in range(num_folds):
    stats_dict['fold%d_seqs' % fi] = folds_num[fi]

with open('%s/statistics.json' % out_dir, 'w') as stats_json_out:
    json.dump(stats_dict, stats_json_out, indent=4)

################################################################
# write tfrecords

tfrecords_dir = '%s/tfrecords' % out_dir
os.mkdir(tfrecords_dir)
for fi in range(num_folds):
    write_fold(fi, tfrecords_dir)<jupyter_output><ipython-input-34-13d201adb886>:10: DeprecationWarning: tostring() is deprecated. Use tobytes() instead.
  seq_indexes = seq_indexes.tostring()
<ipython-input-34-13d201adb886>:14: DeprecationWarning: tostring() is deprecated. Use tobytes() instead.
  targets = targets.tostring()
{'fraction_non_alphanumeric': 0.1384590190780316, 'fraction_numerical': 0.021698190452796202, 'mean_word_length': 3.0835840855901036, 'pattern_counts': {'":': 0, '<': 35, '<?xml version=': 0, '>': 35, 'https://': 0, 'lorem ipsum': 0, 'www.': 0, 'xml': 0}, 'pii_count': 0, 'substrings_counts': 0, 'word_list_counts': {'cursed_substrings.json': 0, 'profanity_word_list.json': 0, 'sexual_word_list.json': 0, 'zh_pornsignals.json': 0}}
{'chain_length': '7', 'dir': 'jupyter-structured-clean-dedup', 'id': '615574', 'license': 'Apache-2.0', 'n_tokens_mistral': 4657, 'n_tokens_neox': 4391, 'n_words': 685, 'path': 'jupyter/saluki_data.ipynb', 'repo_name': 'prasathlab/basenji', 'size': '12213'}