Coverage for gpaw/test/test_watermodel.py: 95%

1import time 

2 

3import ase.units as units 

4import numpy as np 

5import pytest 

6from ase import Atoms 

7from ase.calculators.tip3p import TIP3P, angleHOH, rOH 

8from ase.constraints import FixBondLengths 

9from ase.io import read 

10from ase.io.trajectory import Trajectory 

11from ase.md import Langevin as Langevin0 

12 

13from gpaw.utilities.watermodel import FixBondLengthsWaterModel, TIP3PWaterModel 

14 

15 

16def Langevin(*args, **kwargs): 

17 try: 

18 return Langevin0(*args, **kwargs) 

19 except TypeError: 

20 kT = kwargs.pop('temperature_K') * units.kB 

21 return Langevin0(*args, **kwargs, temperature=kT) 

22 

23 

24@pytest.mark.slow 

25def test_watermodel(in_tmp_dir): 

26 NSTEPS = 600 

27 SCALE = 200 

28 

29 cutoff = 4.0 

30 

31 # Set up water box at 20 deg C density 

32 x = angleHOH * np.pi / 180 / 2 

33 pos = [[0, 0, 0], 

34 [0, rOH * np.cos(x), rOH * np.sin(x)], 

35 [0, rOH * np.cos(x), -rOH * np.sin(x)]] 

36 atoms = Atoms('OH2', positions=pos) 

37 

38 vol = ((18.01528 / 6.022140857e23) / (0.9982 / 1e24))**(1 / 3.) 

39 atoms.set_cell((vol, vol, vol)) 

40 atoms.center() 

41 

42 N = 4 

43 atoms = atoms.repeat((N, N, N)) 

44 atoms.set_pbc(True) 

45 

46 pairs = [(3 * i + j, 3 * i + (j + 1) % 3) 

47 for i in range(len(atoms) // 3) 

48 for j in [0, 1, 2]] 

49 

50 # Create atoms object with old constraints for reference 

51 atoms_ref = atoms.copy() 

52 atoms_ref.constraints = FixBondLengths(pairs) 

53 

54 # RATTLE-type constraints on O-H1, O-H2, H1-H2. 

55 atoms.constraints = FixBondLengthsWaterModel(pairs) 

56 

57 atoms.calc = TIP3PWaterModel(rc=cutoff) 

58 atoms_ref.calc = TIP3P(rc=cutoff) 

59 

60 rng = np.random.RandomState(123) 

61 md = Langevin(atoms, 1 * units.fs, temperature_K=300, 

62 rng=rng, 

63 friction=0.01, logfile='C.log') 

64 traj = Trajectory('C.traj', 'w', atoms) 

65 md.attach(traj.write, interval=1) 

66 

67 start = time.time() 

68 with md: 

69 md.run(NSTEPS) 

70 end = time.time() 

71 Cversion = end - start 

72 print("%d steps of C-MD took %.3fs (%.0f ms/step)" % ( 

73 NSTEPS, Cversion, 

74 Cversion / NSTEPS * 1000)) 

75 traj.close() 

76 

77 rng = np.random.RandomState(123) 

78 md_ref = Langevin(atoms_ref, 1 * units.fs, temperature_K=300, 

79 rng=rng, 

80 friction=0.01, logfile='ref.log') 

81 traj_ref = Trajectory('ref.traj', 'w', atoms_ref) 

82 md_ref.attach(traj_ref.write, interval=1) 

83 start = time.time() 

84 with md_ref: 

85 md_ref.run(NSTEPS / SCALE) 

86 end = time.time() 

87 Pyversion = (end - start) * SCALE 

88 print("%d steps of Py-MD took %.3fs (%.0f ms/step)" % ( 

89 NSTEPS / SCALE, 

90 Pyversion / SCALE, Pyversion / NSTEPS * 1000)) 

91 traj_ref.close() 

92 

93 # Compare trajectories 

94 images = read('C.traj@:') 

95 images_ref = read('ref.traj@:') 

96 for img1, img2 in zip(images, images_ref): 

97 norm = np.linalg.norm(img1.get_positions() - img2.get_positions()) 

98 print(norm) 

99 assert norm < 1e-11 

100 

101 print("Speedup", Pyversion / Cversion)