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

1# fmt: off 

2 

3import numpy as np 

4 

5from ase.data import atomic_numbers 

6from ase.units import Bohr 

7 

8 

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

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

11 if isinstance(fileobj, str): 

12 with open(fileobj) as fd: 

13 lines = fd.readlines() 

14 else: 

15 lines = fileobj 

16 

17 sep = '---------------' 

18 i = 0 # Counter for the lines 

19 k = 0 # Counter of sep 

20 assume6columns = False 

21 for line in lines: 

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

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

24 if i == 0: 

25 headings = line 

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

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

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

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

30 else: 

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

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

33 j = 4 

34 assume6columns = True 

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

36 if k == 1: 

37 break 

38 k += 1 

39 if i <= 1: 

40 pass 

41 else: 

42 words = line.split() 

43 if assume6columns is True: 

44 if len(words) != 6: 

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

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

47 

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

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

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

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

52 # ACF.dat units could be Bohr or Angstrom 

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

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

55 assert norm1 < displacement or norm2 < displacement 

56 i += 1