Coverage for gamdpy/misc/plot_molecule.py: 12%

1import matplotlib.pyplot as plt 

2import numpy as np 

3 

4def plot_molecule(top, positions, particle_types, filename="molecule.pdf", block=False): 

5 ''' 

6 This function write a pdf file with a drawing of the molecule  

7 

8 Parameters 

9 ---------- 

10 top : gamdpy topology object 

11 

12 positions : list or numpy array with positions of all atoms 

13 

14 particle_types : types of the molecule 

15 

16 filename : name of the output pdf file, default is molecule.pdf 

17 

18 block: boolean, default False. If True shows plot and blocks script until display window is closed 

19 

20 ''' 

21 fig = plt.figure() 

22 fig.suptitle("This is the molecule you're going to simulate", fontsize=16) 

23 ax = fig.add_subplot(projection='3d') 

24 

25 pos = np.array(positions) 

26 num_of_type = np.array([np.size(np.nonzero(particle_types==value)) for value in np.unique(particle_types)]) 

27 

28 D = pos.shape[1] 

29 

30 if D>2: 

31 ax.scatter(pos[:,0], pos[:,1], pos[:,2], c=particle_types) 

32 else: 

33 ax.scatter(pos[:,0], pos[:,1], pos[:,1]*0, c=particle_types) 

34 

35 cmap = plt.cm.plasma 

36 num_bond = np.size(np.unique(np.array(top.bonds)[:,2])) 

37 if num_bond == 0: 

38 print("There are no bonds in top.bonds") 

39 exit() 

40 for bond in top.bonds: 

41 part1 = pos[bond[0],:] 

42 part2 = pos[bond[1],:] 

43 bondt = bond[2] 

44 if D>2: 

45 ax.plot3D([part1[0], part2[0]], [part1[1], part2[1]], [part1[2], part2[2]], color=cmap(bondt/num_bond)) 

46 else: 

47 ax.plot3D([part1[0], part2[0]], [part1[1], part2[1]], [0, 0], color=cmap(bondt/num_bond)) 

48 

49 if block: plt.show(block=block) 

50 plt.savefig(filename, format="pdf", bbox_inches="tight")