miniMDS/plotting.py at master · seqcode/miniMDS · 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
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
frommayaviimportmlab
importnumpyasnp
importlinear_algebraasla
importos
importsys
frommatplotlibimportpyplotasplt

#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) forcolorindefault_colors]	#convert to tuple

defplot_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))
ifradiusisNone:
radius=calculateRadius(structures)
fori, structureinenumerate(structures):
coords=np.array(structure.getCoords())
xs=coords[:,0]
ys=coords[:,1]
zs=coords[:,2]
ifcut:
midpoint=np.mean(xs)
indices=np.where(xs>midpoint)[0]
xs=xs[indices]
ys=ys[indices]
zs=zs[indices]
ifall_enrichmentsisnotNone:
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])
ifout_path:
mlab.savefig(out_path)
mlab.show()

fig, ax=plt.subplots()
ax.axis("off")
fori, colorinenumerate(colors):
plt.text(0.5, i, "structure {}".format(i+1), color=color)
plt.savefig("structures_legend")
plt.close()

defplot_structure_interactive(structure, enrichments=None, color=(1,0,0), radius=None, out_path=None):
ifradiusisNone:
radius=calculateRadius([structure])
coords=np.array(structure.getCoords())
xs=coords[:,0]
ys=coords[:,1]
zs=coords[:,2]
mlab.figure(bgcolor=(1,1,1))
ifenrichmentsisnotNone:
mlab.plot3d(xs, ys, zs, enrichments, tube_radius=radius, colormap="bwr")
else:
mlab.plot3d(xs, ys, zs, tube_radius=radius, color=color)
ifout_path:
mlab.savefig(out_path)
mlab.show()

defplot_structures_gif(structures, outname, all_enrichments=None, colors=default_colors, radius=None, increment=10):
if360%increment!=0:
print("Error. Increment must be factor of 360.")
sys.exit(0)
ifradiusisNone:
radius=calculateRadius(structures)
mlab.figure(bgcolor=(1,1,1))
foriinrange(0, 360, increment):
forj, structureinenumerate(structures):
coords=np.array(structure.getCoords())
ifall_enrichmentsisnotNone:
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))

fig, ax=plt.subplots()
ax.axis("off")
fori, colorinenumerate(colors):
plt.text(0.5, i, "structure {}".format(i+1), color=color)
plt.savefig("structures_legend")
plt.close()

defplot_structure_gif(structure, outname, enrichments=None, color=(1,0,0), radius=None, increment=10):
if360%increment!=0:
print("Error. Increment must be factor of 360.")
sys.exit(0)
ifradiusisNone:
radius=calculateRadius([structure])
coords=np.array(structure.getCoords())
mlab.figure(bgcolor=(1,1,1))
ifenrichmentsisnotNone:
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)
foriinrange(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))

defcalculateRadius(structures):
"""Calculate to-scale radius based on Kuhn length and diameter of chromatin"""
conversionFactors=np.zeros(len(structures))
forj, structureinenumerate(structures):
totDist=0
coords=structure.getCoords()
n=len(coords)
foriinrange(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)
returnchromatinDiameter/2*conversionFactor