Coverage for /builds/debichem-team/python-ase/ase/io/amber.py: 86.52%

1import numpy as np 

2 

3import ase.units as units 

4 

5 

6def write_amber_coordinates(atoms, filename): 

7 from scipy.io import netcdf_file 

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

9 _write_amber_coordinates(atoms, fout) 

10 

11 

12def read_amber_coordinates(filename): 

13 from scipy.io import netcdf_file 

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

15 return _read_amber_coordinates(fin) 

16 

17 

18def _write_amber_coordinates(atoms, fout): 

19 fout.Conventions = 'AMBERRESTART' 

20 fout.ConventionVersion = "1.0" 

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

22 fout.application = "AMBER" 

23 fout.program = "ASE" 

24 fout.programVersion = "1.0" 

25 fout.createDimension('cell_spatial', 3) 

26 fout.createDimension('label', 5) 

27 fout.createDimension('cell_angular', 3) 

28 fout.createDimension('time', 1) 

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

30 time.units = 'picosecond' 

31 time[0] = 0 

32 fout.createDimension('spatial', 3) 

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

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

35 

36 natom = len(atoms) 

37 fout.createDimension('atom', natom) 

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

39 ('atom', 'spatial')) 

40 coordinates.units = 'angstrom' 

41 coordinates[:] = atoms.get_positions() 

42 

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

44 ('atom', 'spatial')) 

45 velocities.units = 'angstrom/picosecond' 

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

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

48 # are the only ones it is safe to print 

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

50 # Apply conversion factor from ps: 

51 velocities.scale_factor = 20.455 

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

53 # so convert to Ang/ps 

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

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

56 

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

58 

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

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

61 

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

63 cell_lengths.units = 'angstrom' 

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

65 

66 if not atoms.cell.orthorhombic: 

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

68 

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

70 cell_angles[:3] = 90.0 

71 cell_angles.units = 'degree' 

72 

73 

74def _read_amber_coordinates(fin): 

75 from ase import Atoms 

76 

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

78 get_last_frame = False 

79 if hasattr(all_coordinates, 'ndim'): 

80 if all_coordinates.ndim == 3: 

81 get_last_frame = True 

82 elif hasattr(all_coordinates, 'shape'): 

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

84 get_last_frame = True 

85 if get_last_frame: 

86 all_coordinates = all_coordinates[-1] 

87 

88 atoms = Atoms(positions=all_coordinates) 

89 

90 if 'velocities' in fin.variables: 

91 all_velocities = fin.variables['velocities'] 

92 if hasattr(all_velocities, 'units'): 

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

94 raise Exception( 

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

96 if hasattr(all_velocities, 'scale_factor'): 

97 scale_factor = all_velocities.scale_factor 

98 else: 

99 scale_factor = 1.0 

100 all_velocities = all_velocities[:] * scale_factor 

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

102 if get_last_frame: 

103 all_velocities = all_velocities[-1] 

104 atoms.set_velocities(all_velocities) 

105 if 'cell_lengths' in fin.variables: 

106 all_abc = fin.variables['cell_lengths'] 

107 if get_last_frame: 

108 all_abc = all_abc[-1] 

109 a, b, c = all_abc 

110 all_angles = fin.variables['cell_angles'] 

111 if get_last_frame: 

112 all_angles = all_angles[-1] 

113 alpha, beta, gamma = all_angles 

114 

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

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

117 atoms.set_cell( 

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

119 [0, b, 0], 

120 [0, 0, c]])) 

121 atoms.set_pbc(True) 

122 else: 

123 raise NotImplementedError('only rectangular cells are' 

124 ' implemented in ASE-AMBER') 

125 

126 else: 

127 atoms.set_pbc(False) 

128 

129 return atoms