1import numpy as np
2from ase.units import Bohr, Hartree
3
4from gpaw.utilities import h2gpts
5from gpaw.fftw import get_efficient_fft_size
6from gpaw.wavefunctions.fd import FD
7
8
9def get_number_of_grid_points(cell_cv, h=None, mode=None, realspace=None,
10                              symmetry=None, log=None):
11    if mode is None:
12        mode = FD()
13
14    if realspace is None:
15        realspace = mode.name != 'pw'
16
17    if h is None:
18        if mode.name == 'pw':
19            h = np.pi / (4 * mode.ecut)**0.5
20        elif mode.name == 'lcao' and not realspace:
21            h = np.pi / (4 * 340 / Hartree)**0.5
22        else:
23            h = 0.2 / Bohr
24
25    if realspace or mode.name == 'fd':
26        N_c = h2gpts(h, cell_cv, 4)
27    else:
28        N_c = np.ceil((cell_cv**2).sum(1)**0.5 / h).astype(int)
29        if symmetry is None:
30            N_c = np.array([get_efficient_fft_size(N) for N in N_c])
31        else:
32            N_c = np.array([get_efficient_fft_size(N, n)
33                            for N, n in zip(N_c, symmetry.gcd_c)])
34
35    if symmetry is not None:
36        ok = symmetry.check_grid(N_c)
37        if not ok:
38            if log is not None:
39                log('Initial realspace grid '
40                    '({},{},{}) inconsistent with symmetries.'.format(*N_c))
41            # The grid is not symmetric enough. The essential problem
42            # is that we should start at some other Nmin_c and possibly with
43            # other gcd_c when getting the most efficient fft grids
44            gcd_c = symmetry.gcd_c.copy()
45            Nmin_c = N_c.copy()
46            for i in range(3):
47                for op_cc in symmetry.op_scc:
48                    for i, o in enumerate((op_cc.T).flat):
49                        i1, i2 = np.unravel_index(i, (3, 3))
50                        if i1 == i2:
51                            continue
52                        if o:
53                            # The axes are related and therefore they share
54                            # lowest common multiple of gcd
55                            gcd = np.lcm.reduce(gcd_c[[i1, i2]])
56                            gcd_c[[i1, i2]] = gcd
57                            # We just take the maximum of the two axes to make
58                            # sure that they are divisible always
59                            Nmin = np.max([Nmin_c[i1], Nmin_c[i2]])
60                            Nmin_c[i1] = Nmin
61                            Nmin_c[i2] = Nmin
62
63            N_c = np.array([get_efficient_fft_size(N, n)
64                            for N, n in zip(Nmin_c, gcd_c)])
65            log('Using symmetrized grid: ({},{},{}).\n'.format(*N_c))
66            ok = symmetry.check_grid(N_c)
67            assert ok, ('Grid still not constistent with symmetries '
68                        '({},{},{})'.format(*N_c))
69
70    return N_c
71