Reactivity_comparisons/heatmapcode.py at master · LucksLab/Reactivity_comparisons · GitHub

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
#!/usr/bin/env python
"""
Write heatmaps for F1-score of SHAPE slope-intercept combinations
"""

import SU
import OSU
import NAU
import re
#import LucksLabUtils_config
import PAU
from itertools import cycle
from collections import namedtuple
import cPickle as pickle
import os
import runscorefile
import numpy as np
import matplotlib
matplotlib.use('Agg')
import matplotlib.pyplot as plt

def plotf1(infile, outfile, title):
    #Make heatmap of F1-scores from RNA folding with SHAPE slope and intercept ranges on the axes
    #infile: string path to file containing data from scoring of RNA folding against crystal structure; contains SHAPE slope, intercept, sensitivity, PPV, and F1-score
    #outfile: string name for output file containing heatmap
    #title: string name for heatmap
    try:
        with open(infile, 'r') as f:
            f.readline()
            f.readline()
            x = []
            y = []
            z = []
            z_ind = -1
            for line in f:
                #Get SHAPE slope, intercept, and F1-score from each line
                #For slope and intercept, only one of each value is assigned to the x and y lists, rather than all the repetitions evident in the data file
                line_arr = line.split(',')
                outcheck = 0
                for i in x:
                    if i == line_arr[0]:
                        outcheck = 1
                if outcheck == 0:
                    x.append(line_arr[0])
                    startnewx = line_arr[0]
                    z_arr = []
                    if (len(line_arr)) == 5:
                        z_new = (line_arr[4])[:-1]
                        if len(z_new)>5:
                            z_new = z_new[:5]
                        z_new = float(z_new)
                        z_arr.append(z_new)
                    z.append(z_arr)
                    z_ind += 1
                else:
                    z_new = (line_arr[4])[:-1]
                    if len(z_new)>5:
                        z_new = z_new[:5]
                    z_new = float(z_new)
                    z[z_ind].append(z_new)
                outcheck = 0
                for i in y:
                    if i == line_arr[1]:
                        outcheck = 1
                if outcheck == 0:
                    y.append(line_arr[1])
    except EnvironmentError:
        print("Error opening txt file" + infile)

    z_in = np.array(z)
    fig, ax = plt.subplots()
    im = ax.imshow(z_in)

    cbar = ax.figure.colorbar(im, ax=ax, shrink=0.5)
    cbar.ax.set_ylabel("F score", rotation=-90, va="bottom", fontsize=24)
    cbar.ax.tick_params(labelsize=16)

    ax.set_xticks(np.arange(len(y)))
    ax.set_yticks(np.arange(len(x)))

    ax.set_xticklabels(y)
    ax.set_yticklabels(x)

    plt.setp(ax.get_xticklabels(), rotation=45, ha="right",rotation_mode="anchor", fontsize=14)
    plt.setp(ax.get_yticklabels(), fontsize=16)

    plt.xlabel('SHAPE Intercept')
    plt.ylabel('SHAPE Slope')

    for i in range(len(x)):
        for j in range(len(y)):
            text = ax.text(j, i, z_in[i, j], ha="center", va="center", color="w", fontsize=10)

    ax.set_title(title + " Heatmap", fontsize=24)
    fig.set_size_inches(16, 23)
    fig.tight_layout()
    fig.savefig(outfile, bbox_inches='tight')
    plt.close(fig)