Coverage for /builds/ase/ase/ase/visualize/plot.py: 88.89%

1# fmt: off 

2 

3from ase.io.utils import PlottingVariables, make_patch_list 

4 

5 

6class Matplotlib(PlottingVariables): 

7 def __init__(self, atoms, ax, 

8 rotation='', radii=None, 

9 colors=None, scale=1, offset=(0, 0), **parameters): 

10 PlottingVariables.__init__( 

11 self, atoms, rotation=rotation, 

12 radii=radii, colors=colors, scale=scale, 

13 extra_offset=offset, **parameters) 

14 

15 self.ax = ax 

16 self.figure = ax.figure 

17 self.ax.set_aspect('equal') 

18 

19 def write(self): 

20 self.write_body() 

21 self.ax.set_xlim(0, self.w) 

22 self.ax.set_ylim(0, self.h) 

23 

24 def write_body(self): 

25 patch_list = make_patch_list(self) 

26 for patch in patch_list: 

27 self.ax.add_patch(patch) 

28 

29 

30def animate(images, ax=None, 

31 interval=200, # in ms; same default value as in FuncAnimation 

32 save_count=None, # ignored as of 2023 with newer matplotlib 

33 **parameters): 

34 """Convert sequence of atoms objects into Matplotlib animation. 

35 

36 Each image is generated using plot_atoms(). Additional parameters 

37 are passed to this function.""" 

38 import matplotlib.pyplot as plt 

39 from matplotlib.animation import FuncAnimation 

40 

41 if ax is None: 

42 ax = plt.gca() 

43 

44 fig = ax.get_figure() 

45 

46 def drawimage(atoms): 

47 ax.clear() 

48 ax.axis('off') 

49 plot_atoms(atoms, ax=ax, **parameters) 

50 

51 animation = FuncAnimation(fig, drawimage, frames=images, 

52 init_func=lambda: None, 

53 interval=interval) 

54 return animation 

55 

56 

57def plot_atoms(atoms, ax=None, **parameters): 

58 """Plot an atoms object in a matplotlib subplot. 

59 

60 Parameters 

61 ---------- 

62 atoms : Atoms object 

63 ax : Matplotlib subplot object 

64 rotation : str, optional 

65 In degrees. In the form '10x,20y,30z' 

66 show_unit_cell : int, optional, default 2 

67 Draw the unit cell as dashed lines depending on value: 

68 0: Don't 

69 1: Do 

70 2: Do, making sure cell is visible 

71 radii : float, optional 

72 The radii of the atoms 

73 colors : list of strings, optional 

74 Color of the atoms, must be the same length as 

75 the number of atoms in the atoms object. 

76 scale : float, optional 

77 Scaling of the plotted atoms and lines. 

78 offset : tuple (float, float), optional 

79 Offset of the plotted atoms and lines. 

80 """ 

81 if isinstance(atoms, list): 

82 assert len(atoms) == 1 

83 atoms = atoms[0] 

84 

85 import matplotlib.pyplot as plt 

86 if ax is None: 

87 ax = plt.gca() 

88 Matplotlib(atoms, ax, **parameters).write() 

89 return ax