Coverage for ase / _4 / atoms.py: 100.00%

1"""ASEv4 `Atoms` class.""" 

2 

3from ase._4.calculators.results import CalculationResults 

4from ase.atoms import Atoms as V3Atoms 

5from ase.atoms import _LimitedAtoms 

6from ase.outputs import ArrayProperty, all_outputs 

7 

8 

9class Atoms(_LimitedAtoms): 

10 """ASEv4 Atoms class.""" 

11 

12 @classmethod 

13 def from_v3atoms(cls, v3atoms: V3Atoms): 

14 """Make ASEv4 Atoms from ASEv3 Atoms.""" 

15 return cls( 

16 symbols=v3atoms.symbols, 

17 positions=v3atoms.positions, 

18 cell=v3atoms.cell, 

19 pbc=v3atoms.pbc, 

20 ) 

21 

22 def store( 

23 self, 

24 results: CalculationResults, 

25 prefix: str = '', 

26 suffix: str = '', 

27 ) -> None: 

28 """Stores `results` in `Atoms.info` and and `Atoms.arrays`. 

29 

30 Parameters 

31 ---------- 

32 results : `CalculationResults` 

33 Properties to be saved. 

34 prefix : str 

35 Prefix to the names of the properties. 

36 suffix : str 

37 Suffix to the names of the properties. 

38 

39 """ 

40 for prop_name, prop_val in results.properties.items(): 

41 output_type = all_outputs[prop_name] 

42 if ( 

43 isinstance(output_type, ArrayProperty) 

44 and output_type.shapespec[0] == 'natoms' 

45 ): 

46 self.arrays[prefix + prop_name + suffix] = prop_val 

47 else: 

48 self.info[prefix + prop_name + suffix] = prop_val