-
Notifications
You must be signed in to change notification settings - Fork 0
/
plotting.py
108 lines (100 loc) · 4.14 KB
/
plotting.py
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
99
100
101
102
103
104
105
106
107
108
from mayavi import mlab
import numpy as np
import linear_algebra as la
import os
import sys
#from Rippe 2001
kl = 289 #Kuhn length (nm)
bpPerKL = 30000. #base pairs per Kuhn length
chromatinDiameter = 30 #diameter of heterochromatin (nm)
default_colors = np.array([[255,0,0], [0,255,238], [255,238,0], [0,102,255], [255,0,170], [255,102,0], [204,255,0], [0,238,255], [0,68,255], [255,0,102], [255,136,0], [0,255,34], [0,204,255], [34,0,255], [255,0,68], [255,170,0], [0,255,136], [0,170,255], [204,0,255], [255,204,0], [0,255,204], [0,136,255], [255,0,238]])/255.
default_colors = [tuple(color) for color in default_colors] #convert to tuple
def plot_structures_interactive(structures, all_enrichments=None, colors=default_colors, radius=None, cut=False, out_path=None):
mlab.close(all=True)
mlab.figure(bgcolor=(1,1,1))
if radius is None:
radius = calculateRadius(structures)
for i, structure in enumerate(structures):
coords = np.array(structure.getCoords())
xs = coords[:,0]
ys = coords[:,1]
zs = coords[:,2]
if cut:
midpoint = np.mean(xs)
indices = np.where(xs > midpoint)[0]
xs = xs[indices]
ys = ys[indices]
zs = zs[indices]
if all_enrichments is not None:
mlab.plot3d(xs, ys, zs, all_enrichments[i], tube_radius=radius, colormap="bwr")
else:
mlab.plot3d(xs, ys, zs, tube_radius=radius, color=colors[i])
if out_path:
mlab.savefig(out_path)
mlab.show()
def plot_structure_interactive(structure, enrichments=None, color=(1,0,0), radius=None, out_path=None):
if radius is None:
radius = calculateRadius([structure])
coords = np.array(structure.getCoords())
xs = coords[:,0]
ys = coords[:,1]
zs = coords[:,2]
mlab.figure(bgcolor=(1,1,1))
if enrichments is not None:
mlab.plot3d(xs, ys, zs, enrichments, tube_radius=radius, colormap="bwr")
else:
mlab.plot3d(xs, ys, zs, tube_radius=radius, color=color)
if out_path:
mlab.savefig(out_path)
mlab.show()
def plot_structures_gif(structures, outname, all_enrichments=None, colors=default_colors, radius=None, increment=10):
if 360%increment != 0:
print("Error. Increment must be factor of 360.")
sys.exit(0)
if radius is None:
radius = calculateRadius(structures)
mlab.figure(bgcolor=(1,1,1))
for i in range(0, 360, increment):
for j, structure in enumerate(structures):
coords = np.array(structure.getCoords())
if all_enrichments is not None:
s = mlab.plot3d(coords[:,0], coords[:,1], coords[:,2], all_enrichments[j], tube_radius=radius, colormap="bwr")
else:
s = mlab.plot3d(coords[:,0], coords[:,1], coords[:,2], tube_radius=radius, color=colors[j])
mlab.view(i)
mlab.savefig("{}_{:>03}.png".format(outname, i))
mlab.close()
os.system("convert {}_*.png {}.gif".format(outname, outname))
os.system("rm {}_*.png".format(outname))
def plot_structure_gif(structure, outname, enrichments=None, color=(1,0,0), radius=None, increment=10):
if 360%increment != 0:
print("Error. Increment must be factor of 360.")
sys.exit(0)
if radius is None:
radius = calculateRadius([structure])
coords = np.array(structure.getCoords())
mlab.figure(bgcolor=(1,1,1))
if enrichments is not None:
s = mlab.plot3d(coords[:,0], coords[:,1], coords[:,2], enrichments, tube_radius=radius, colormap="bwr")
else:
s = mlab.plot3d(coords[:,0], coords[:,1], coords[:,2], tube_radius=radius, color=color)
for i in range(0, 360, increment):
mlab.view(i)
mlab.savefig("{}_{:>03}.png".format(outname, i))
mlab.close()
os.system("convert {}_*.png {}.gif".format(outname, outname))
os.system("rm {}_*.png".format(outname))
def calculateRadius(structures):
"""Calculate to-scale radius based on Kuhn length and diameter of chromatin"""
conversionFactors = np.zeros(len(structures))
for j, structure in enumerate(structures):
totDist = 0
coords = structure.getCoords()
n = len(coords)
for i in range(1, n):
totDist += la.calcDistance(coords[i-1], coords[i])
avgDist = totDist/(n-1) #average distance between neighboring loci
physicalDist = kl * (structure.chrom.res/bpPerKL)**(1./2) #physical distance between neighboring loci (nm)
conversionFactors[j] = avgDist/physicalDist
conversionFactor = np.mean(conversionFactors)
return chromatinDiameter/2 * conversionFactor