Coverage for /builds/kinetik161/ase/ase/visualize/viewers.py: 82.93%

1""" 

2Module for managing viewers 

3 

4View plugins can be registered through the entrypoint system in with the 

5following in a module, such as a `viewer.py` file: 

6 

7```python3 

8VIEWER_ENTRYPOINT = ExternalViewer( 

9 desc="Visualization using <my package>", 

10 module="my_package.viewer" 

11) 

12``` 

13 

14Where module `my_package.viewer` contains a `view_my_viewer` function taking 

15and `ase.Atoms` object as the first argument, and also `**kwargs`. 

16 

17Then ones needs to register an entry point in `pyproject.toml` with 

18 

19```toml 

20[project.entry-points."ase.visualize"] 

21my_viewer = "my_package.viewer:VIEWER_ENTRYPOINT" 

22``` 

23 

24After this, call to `ase.visualize.view(atoms, viewer='my_viewer')` will be 

25forwarded to `my_package.viewer.view_my_viewer` function. 

26 

27""" 

28import pickle 

29import subprocess 

30import sys 

31import tempfile 

32import warnings 

33from contextlib import contextmanager 

34from io import BytesIO 

35from pathlib import Path 

36 

37from importlib.metadata import entry_points 

38 

39from importlib import import_module 

40 

41from ase.io import write 

42from ase.io.formats import ioformats 

43from ase.utils.plugins import ExternalViewer 

44 

45 

46class UnknownViewerError(Exception): 

47 """The view tyep is unknown""" 

48 

49 

50class AbstractViewer: 

51 def view(self, *args, **kwargss): 

52 raise NotImplementedError() 

53 

54 

55class PyViewer(AbstractViewer): 

56 def __init__(self, name: str, desc: str, module_name: str): 

57 """ 

58 Instantiate an viewer 

59 """ 

60 self.name = name 

61 self.desc = desc 

62 self.module_name = module_name 

63 

64 def _viewfunc(self): 

65 """Return the function used for viewing the atoms""" 

66 return getattr(self.module, "view_" + self.name, None) 

67 

68 @property 

69 def module(self): 

70 try: 

71 return import_module(self.module_name) 

72 except ImportError as err: 

73 raise UnknownViewerError( 

74 f"Viewer not recognized: {self.name}. Error: {err}" 

75 ) from err 

76 

77 def view(self, atoms, *args, **kwargs): 

78 return self._viewfunc()(atoms, *args, **kwargs) 

79 

80 

81class CLIViewer(AbstractViewer): 

82 """Generic viewer for""" 

83 

84 def __init__(self, name, fmt, argv): 

85 self.name = name 

86 self.fmt = fmt 

87 self.argv = argv 

88 

89 @property 

90 def ioformat(self): 

91 return ioformats[self.fmt] 

92 

93 @contextmanager 

94 def mktemp(self, atoms, data=None): 

95 ioformat = self.ioformat 

96 suffix = "." + ioformat.extensions[0] 

97 

98 if ioformat.isbinary: 

99 mode = "wb" 

100 else: 

101 mode = "w" 

102 

103 with tempfile.TemporaryDirectory(prefix="ase-view-") as dirname: 

104 # We use a tempdir rather than a tempfile because it's 

105 # less hassle to handle the cleanup on Windows (files 

106 # cannot be open on multiple processes). 

107 path = Path(dirname) / f"atoms{suffix}" 

108 with path.open(mode) as fd: 

109 if data is None: 

110 write(fd, atoms, format=self.fmt) 

111 else: 

112 write(fd, atoms, format=self.fmt, data=data) 

113 yield path 

114 

115 def view_blocking(self, atoms, data=None): 

116 with self.mktemp(atoms, data) as path: 

117 subprocess.check_call(self.argv + [str(path)]) 

118 

119 def view( 

120 self, 

121 atoms, 

122 data=None, 

123 repeat=None, 

124 **kwargs, 

125 ): 

126 """Spawn a new process in which to open the viewer.""" 

127 if repeat is not None: 

128 atoms = atoms.repeat(repeat) 

129 

130 proc = subprocess.Popen( 

131 [sys.executable, "-m", "ase.visualize.viewers"], 

132 stdin=subprocess.PIPE 

133 ) 

134 

135 pickle.dump((self, atoms, data), proc.stdin) 

136 proc.stdin.close() 

137 return proc 

138 

139 

140VIEWERS = {} 

141 

142 

143def _pipe_to_ase_gui(atoms, repeat, **kwargs): 

144 buf = BytesIO() 

145 write(buf, atoms, format="traj") 

146 

147 args = [sys.executable, "-m", "ase", "gui", "-"] 

148 if repeat: 

149 args.append("--repeat={},{},{}".format(*repeat)) 

150 

151 proc = subprocess.Popen(args, stdin=subprocess.PIPE) 

152 proc.stdin.write(buf.getvalue()) 

153 proc.stdin.close() 

154 return proc 

155 

156 

157def define_viewer( 

158 name, desc, *, module=None, cli=False, fmt=None, argv=None, external=False 

159): 

160 if not external: 

161 if module is None: 

162 module = name 

163 module = "ase.visualize." + module 

164 if cli: 

165 fmt = CLIViewer(name, fmt, argv) 

166 else: 

167 if name == "ase": 

168 # Special case if the viewer is named `ase` then we use 

169 # the _pipe_to_ase_gui as the viewer method 

170 fmt = PyViewer(name, desc, module_name=None) 

171 fmt.view = _pipe_to_ase_gui 

172 else: 

173 fmt = PyViewer(name, desc, module_name=module) 

174 VIEWERS[name] = fmt 

175 return fmt 

176 

177 

178def define_external_viewer(entry_point): 

179 """Define external viewer""" 

180 

181 viewer_def = entry_point.load() 

182 if entry_point.name in VIEWERS: 

183 raise ValueError(f"Format {entry_point.name} already defined") 

184 if not isinstance(viewer_def, ExternalViewer): 

185 raise TypeError( 

186 "Wrong type for registering external IO formats " 

187 f"in format {entry_point.name}, expected " 

188 "ExternalViewer" 

189 ) 

190 define_viewer(entry_point.name, **viewer_def._asdict(), 

191 external=True) 

192 

193 

194def register_external_viewer_formats(group): 

195 if hasattr(entry_points(), "select"): 

196 viewer_entry_points = entry_points().select(group=group) 

197 else: 

198 viewer_entry_points = entry_points().get(group, ()) 

199 

200 for entry_point in viewer_entry_points: 

201 try: 

202 define_external_viewer(entry_point) 

203 except Exception as exc: 

204 warnings.warn( 

205 "Failed to register external " 

206 f"Viewer {entry_point.name}: {exc}" 

207 ) 

208 

209 

210define_viewer("ase", "View atoms using ase gui.") 

211define_viewer("ngl", "View atoms using nglview.") 

212define_viewer("mlab", "View atoms using matplotlib.") 

213define_viewer("sage", "View atoms using sage.") 

214define_viewer("x3d", "View atoms using x3d.") 

215 

216# CLI viweers that are internally supported 

217define_viewer( 

218 "avogadro", "View atoms using avogradro.", cli=True, fmt="cube", 

219 argv=["avogadro"] 

220) 

221define_viewer( 

222 "ase_gui_cli", "View atoms using ase gui.", cli=True, fmt="traj", 

223 argv=[sys.executable, '-m', 'ase.gui'], 

224) 

225define_viewer( 

226 "gopenmol", 

227 "View atoms using gopenmol.", 

228 cli=True, 

229 fmt="extxyz", 

230 argv=["runGOpenMol"], 

231) 

232define_viewer( 

233 "rasmol", 

234 "View atoms using rasmol.", 

235 cli=True, 

236 fmt="proteindatabank", 

237 argv=["rasmol", "-pdb"], 

238) 

239define_viewer("vmd", "View atoms using vmd.", cli=True, fmt="cube", 

240 argv=["vmd"]) 

241define_viewer( 

242 "xmakemol", 

243 "View atoms using xmakemol.", 

244 cli=True, 

245 fmt="extxyz", 

246 argv=["xmakemol", "-f"], 

247) 

248 

249register_external_viewer_formats("ase.visualize") 

250 

251CLI_VIEWERS = {key: value for key, value in VIEWERS.items() 

252 if isinstance(value, CLIViewer)} 

253PY_VIEWERS = {key: value for key, value in VIEWERS.items() 

254 if isinstance(value, PyViewer)} 

255 

256 

257def cli_main(): 

258 """ 

259 This is mainly to facilitate launching CLI viewer in a separate python 

260 process 

261 """ 

262 cli_viewer, atoms, data = pickle.load(sys.stdin.buffer) 

263 cli_viewer.view_blocking(atoms, data) 

264 

265 

266if __name__ == "__main__": 

267 cli_main()