Coverage for /builds/ase/ase/ase/nomad.py: 91.23%

1import json 

2 

3import numpy as np 

4 

5import ase.units as units 

6from ase import Atoms 

7from ase.data import chemical_symbols 

8 

9 

10def read(fd, _includekeys=lambda key: True): 

11 """Read NomadEntry object from file.""" 

12 # _includekeys can be used to strip unnecessary keys out of a 

13 # downloaded nomad file so its size is suitable for inclusion 

14 # in the test suite. 

15 

16 def hook(dct): 

17 d = {k: dct[k] for k in dct if _includekeys(k)} 

18 return NomadEntry(d) 

19 

20 dct = json.load(fd, object_hook=hook) 

21 return dct 

22 

23 

24def section_system_to_atoms(section): 

25 """Covnert section_system into an Atoms object.""" 

26 assert section['name'] == 'section_system' 

27 numbers = section['atom_species'] 

28 numbers = np.array(numbers, int) 

29 numbers[numbers < 0] = 0 # We don't support Z < 0 

30 numbers[numbers >= len(chemical_symbols)] = 0 

31 positions = section['atom_positions']['flatData'] 

32 positions = np.array(positions).reshape(-1, 3) * units.m 

33 atoms = Atoms(numbers, positions=positions) 

34 atoms.info['nomad_uri'] = section['uri'] 

35 

36 pbc = section.get('configuration_periodic_dimensions') 

37 if pbc is not None: 

38 assert len(pbc) == 1 

39 pbc = pbc[0] # it's a list?? 

40 pbc = pbc['flatData'] 

41 assert len(pbc) == 3 

42 atoms.pbc = pbc 

43 

44 # celldisp? 

45 cell = section.get('lattice_vectors') 

46 if cell is not None: 

47 cell = cell['flatData'] 

48 cell = np.array(cell).reshape(3, 3) * units.m 

49 atoms.cell = cell 

50 

51 return atoms 

52 

53 

54class NomadEntry(dict): 

55 """An entry from the Nomad database. 

56 

57 The Nomad entry is represented as nested dictionaries and lists. 

58 

59 ASE converts each dictionary into a NomadEntry object which supports 

60 different actions. Some actions are only available when the NomadEntry 

61 represents a particular section.""" 

62 

63 def __init__(self, dct): 

64 # assert dct['type'] == 'nomad_calculation_2_0' 

65 # assert dct['name'] == 'calculation_context' 

66 # We could implement NomadEntries that represent sections. 

67 dict.__init__(self, dct) 

68 

69 @property 

70 def hash(self): 

71 # The hash is a string, so not __hash__ 

72 assert self['uri'].startswith('nmd://') 

73 return self['uri'][6:] 

74 

75 def toatoms(self): 

76 """Convert this NomadEntry into an Atoms object. 

77 

78 This NomadEntry must represent a section_system.""" 

79 return section_system_to_atoms(self) 

80 

81 def iterimages(self): 

82 """Yield Atoms object contained within this NomadEntry. 

83 

84 This NomadEntry must represent or contain a section_run.""" 

85 

86 if 'section_run' in self: 

87 run_sections = self['section_run'] 

88 else: 

89 assert self['name'] == 'section_run' 

90 run_sections = [self] # We assume that we are the section_run 

91 

92 for run in run_sections: 

93 systems = run['section_system'] 

94 for system in systems: 

95 atoms = section_system_to_atoms(system) 

96 atoms.info['nomad_run_gIndex'] = run['gIndex'] 

97 atoms.info['nomad_system_gIndex'] = system['gIndex'] 

98 

99 if self.get('name') == 'calculation_context': 

100 atoms.info['nomad_calculation_uri'] = self['uri'] 

101 yield atoms