1import numpy as np 

2 

3 

4def write_py(fileobj, images): 

5 """Write to ASE-compatible python script.""" 

6 fileobj.write('import numpy as np\n\n') 

7 fileobj.write('from ase import Atoms\n\n') 

8 

9 if hasattr(images, 'get_positions'): 

10 images = [images] 

11 fileobj.write('images = [\n') 

12 

13 for image in images: 

14 fileobj.write(" Atoms(symbols='%s',\n" 

15 " pbc=np.array(%s),\n" 

16 " cell=np.array(\n%s),\n" 

17 " positions=np.array(\n%s)),\n" % ( 

18 image.get_chemical_formula(mode='reduce'), 

19 array_to_string(image.pbc, 0), 

20 array_to_string(image.cell), 

21 array_to_string(image.positions))) 

22 

23 fileobj.write(']\n') 

24 

25 

26def array_to_string(array, indent=14): 

27 """Converts given numpy array to a string, which when printed will pass 

28 flake8 tests.""" 

29 text = np.array2string(array, separator=', ', suppress_small=False, 

30 formatter={'float': '{:.8f}'.format, 

31 'bool': '{}'.format}) 

32 text = ' ' * indent + text.replace('\n', '\n' + ' ' * indent) 

33 return text