Coverage for gpaw/utilities/extend_grid.py: 88%

1import numpy as np 

2from ase.units import Bohr 

3from gpaw.grid_descriptor import GridDescriptor 

4from gpaw.utilities.grid import grid2grid 

5 

6 

7def extended_grid_descriptor(gd, 

8 extend_N_cd=None, 

9 N_c=None, extcomm=None): 

10 """Create grid descriptor for extended grid. 

11 

12 Provide only either extend_N_cd or N_c. 

13 

14 Parameters: 

15 

16 extend_N_cd: ndarray, int 

17 Number of extra grid points per axis (c) and 

18 direction (d, left or right) 

19 N_c: ndarray, int 

20 Number of grid points in extended grid 

21 extcomm: 

22 Communicator for the extended grid, defaults to gd.comm 

23 """ 

24 

25 if extcomm is None: 

26 extcomm = gd.comm 

27 

28 if extend_N_cd is None: 

29 assert N_c is not None, 'give only extend_N_cd or N_c' 

30 N_c = np.array(N_c, dtype=int) 

31 extend_N_cd = np.tile((N_c - gd.N_c) // 2, (2, 1)).T 

32 else: # extend_N_cd is not None: 

33 assert N_c is None, 'give only extend_N_cd or N_c' 

34 extend_N_cd = np.array(extend_N_cd, dtype=int) 

35 N_c = gd.N_c + extend_N_cd.sum(axis=1) 

36 

37 cell_cv = gd.h_cv * N_c 

38 move_c = gd.get_grid_spacings() * extend_N_cd[:, 0] 

39 

40 egd = GridDescriptor(N_c, cell_cv, gd.pbc_c, extcomm) 

41 egd.extend_N_cd = extend_N_cd 

42 

43 return egd, cell_cv * Bohr, move_c * Bohr 

44 

45 

46def extend_array(gd, ext_gd, a_g, aext_g): 

47 assert gd.comm.compare(ext_gd.comm) in ['ident', 'congruent'] 

48 offset_c = ext_gd.extend_N_cd[:, 0] 

49 grid2grid(gd.comm, gd, ext_gd, a_g, aext_g, offset1_c=offset_c) 

50 

51 

52def deextend_array(gd, ext_gd, a_g, aext_g): 

53 assert gd.comm.compare(ext_gd.comm) in ['ident', 'congruent'] 

54 offset_c = ext_gd.extend_N_cd[:, 0] 

55 grid2grid(gd.comm, ext_gd, gd, aext_g, a_g, offset2_c=offset_c) 

56 

57 

58def move_atoms(atoms, move_c): 

59 pos_a = atoms.get_positions() 

60 for pos in pos_a: 

61 pos += move_c 

62 atoms.set_positions(pos_a)