Coverage for /builds/kinetik161/ase/ase/io/bader.py: 83.33%

1import numpy as np 

2 

3from ase.data import atomic_numbers 

4from ase.units import Bohr 

5 

6 

7def attach_charges(atoms, fileobj='ACF.dat', displacement=1e-4): 

8 """Attach the charges from the fileobj to the Atoms.""" 

9 if isinstance(fileobj, str): 

10 with open(fileobj) as fd: 

11 lines = fd.readlines() 

12 else: 

13 lines = fileobj 

14 

15 sep = '---------------' 

16 i = 0 # Counter for the lines 

17 k = 0 # Counter of sep 

18 assume6columns = False 

19 for line in lines: 

20 if line[0] == '\n': # check if there is an empty line in the 

21 i -= 1 # head of ACF.dat file 

22 if i == 0: 

23 headings = line 

24 if 'BADER' in headings.split(): 

25 j = headings.split().index('BADER') 

26 elif 'CHARGE' in headings.split(): 

27 j = headings.split().index('CHARGE') 

28 else: 

29 print('Can\'t find keyword "BADER" or "CHARGE".' 

30 ' Assuming the ACF.dat file has 6 columns.') 

31 j = 4 

32 assume6columns = True 

33 if sep in line: # Stop at last separator line 

34 if k == 1: 

35 break 

36 k += 1 

37 if not i > 1: 

38 pass 

39 else: 

40 words = line.split() 

41 if assume6columns is True: 

42 if len(words) != 6: 

43 raise OSError('Number of columns in ACF file incorrect!\n' 

44 'Check that Bader program version >= 0.25') 

45 

46 atom = atoms[int(words[0]) - 1] 

47 atom.charge = atomic_numbers[atom.symbol] - float(words[j]) 

48 if displacement is not None: # check if the atom positions match 

49 xyz = np.array([float(w) for w in words[1:4]]) 

50 # ACF.dat units could be Bohr or Angstrom 

51 norm1 = np.linalg.norm(atom.position - xyz) 

52 norm2 = np.linalg.norm(atom.position - xyz * Bohr) 

53 assert norm1 < displacement or norm2 < displacement 

54 i += 1