Coverage for /builds/ase/ase/ase/io/amber.py: 86.52%

1# fmt: off 

2 

3import numpy as np 

4 

5import ase.units as units 

6 

7 

8def write_amber_coordinates(atoms, filename): 

9 from scipy.io import netcdf_file 

10 with netcdf_file(filename, 'w', mmap=False) as fout: 

11 _write_amber_coordinates(atoms, fout) 

12 

13 

14def read_amber_coordinates(filename): 

15 from scipy.io import netcdf_file 

16 with netcdf_file(filename, 'r', mmap=False) as fin: 

17 return _read_amber_coordinates(fin) 

18 

19 

20def _write_amber_coordinates(atoms, fout): 

21 fout.Conventions = 'AMBERRESTART' 

22 fout.ConventionVersion = "1.0" 

23 fout.title = 'Ase-generated-amber-restart-file' 

24 fout.application = "AMBER" 

25 fout.program = "ASE" 

26 fout.programVersion = "1.0" 

27 fout.createDimension('cell_spatial', 3) 

28 fout.createDimension('label', 5) 

29 fout.createDimension('cell_angular', 3) 

30 fout.createDimension('time', 1) 

31 time = fout.createVariable('time', 'd', ('time',)) 

32 time.units = 'picosecond' 

33 time[0] = 0 

34 fout.createDimension('spatial', 3) 

35 spatial = fout.createVariable('spatial', 'c', ('spatial',)) 

36 spatial[:] = np.asarray(list('xyz')) 

37 

38 natom = len(atoms) 

39 fout.createDimension('atom', natom) 

40 coordinates = fout.createVariable('coordinates', 'd', 

41 ('atom', 'spatial')) 

42 coordinates.units = 'angstrom' 

43 coordinates[:] = atoms.get_positions() 

44 

45 velocities = fout.createVariable('velocities', 'd', 

46 ('atom', 'spatial')) 

47 velocities.units = 'angstrom/picosecond' 

48 # Amber's units of time are 1/20.455 ps 

49 # Any other units are ignored in restart files, so these 

50 # are the only ones it is safe to print 

51 # See: http://ambermd.org/Questions/units.html 

52 # Apply conversion factor from ps: 

53 velocities.scale_factor = 20.455 

54 # get_velocities call returns velocities with units sqrt(eV/u) 

55 # so convert to Ang/ps 

56 factor = units.fs * 1000 / velocities.scale_factor 

57 velocities[:] = atoms.get_velocities() * factor 

58 

59 fout.createVariable('cell_angular', 'c', ('cell_angular', 'label')) 

60 

61 cell_spatial = fout.createVariable('cell_spatial', 'c', ('cell_spatial',)) 

62 cell_spatial[:] = ['a', 'b', 'c'] 

63 

64 cell_lengths = fout.createVariable('cell_lengths', 'd', ('cell_spatial',)) 

65 cell_lengths.units = 'angstrom' 

66 cell_lengths[:] = atoms.cell.lengths() 

67 

68 if not atoms.cell.orthorhombic: 

69 raise ValueError('Non-orthorhombic cell not supported with amber') 

70 

71 cell_angles = fout.createVariable('cell_angles', 'd', ('cell_angular',)) 

72 cell_angles[:3] = 90.0 

73 cell_angles.units = 'degree' 

74 

75 

76def _read_amber_coordinates(fin): 

77 from ase import Atoms 

78 

79 all_coordinates = fin.variables['coordinates'][:] 

80 get_last_frame = False 

81 if hasattr(all_coordinates, 'ndim'): 

82 if all_coordinates.ndim == 3: 

83 get_last_frame = True 

84 elif hasattr(all_coordinates, 'shape'): 

85 if len(all_coordinates.shape) == 3: 

86 get_last_frame = True 

87 if get_last_frame: 

88 all_coordinates = all_coordinates[-1] 

89 

90 atoms = Atoms(positions=all_coordinates) 

91 

92 if 'velocities' in fin.variables: 

93 all_velocities = fin.variables['velocities'] 

94 if hasattr(all_velocities, 'units'): 

95 if all_velocities.units != b'angstrom/picosecond': 

96 raise Exception( 

97 f'Unrecognised units {all_velocities.units}') 

98 if hasattr(all_velocities, 'scale_factor'): 

99 scale_factor = all_velocities.scale_factor 

100 else: 

101 scale_factor = 1.0 

102 all_velocities = all_velocities[:] * scale_factor 

103 all_velocities = all_velocities / (1000 * units.fs) 

104 if get_last_frame: 

105 all_velocities = all_velocities[-1] 

106 atoms.set_velocities(all_velocities) 

107 if 'cell_lengths' in fin.variables: 

108 all_abc = fin.variables['cell_lengths'] 

109 if get_last_frame: 

110 all_abc = all_abc[-1] 

111 a, b, c = all_abc 

112 all_angles = fin.variables['cell_angles'] 

113 if get_last_frame: 

114 all_angles = all_angles[-1] 

115 alpha, beta, gamma = all_angles 

116 

117 if (all(angle > 89.99 for angle in [alpha, beta, gamma]) and 

118 all(angle < 90.01 for angle in [alpha, beta, gamma])): 

119 atoms.set_cell( 

120 np.array([[a, 0, 0], 

121 [0, b, 0], 

122 [0, 0, c]])) 

123 atoms.set_pbc(True) 

124 else: 

125 raise NotImplementedError('only rectangular cells are' 

126 ' implemented in ASE-AMBER') 

127 

128 else: 

129 atoms.set_pbc(False) 

130 

131 return atoms