1from ase import Atoms
2
3from gpaw.mpi import size, rank
4from gpaw import GPAW, FermiDirac
5from gpaw.analyse.simple_stm import SimpleStm
6from gpaw.test import equal
7
8
9def test_utilities_simple_stm(in_tmp_dir):
10    load = False
11    txt = '/dev/null'
12    txt = '-'
13
14    me = ''
15    if size > 1:
16        me += 'rank ' + str(rank) + ': '
17
18    BH = Atoms('BH', [[.0, .0, .41], [.0, .0, -1.23]],
19               cell=[5, 6, 6.5])
20    BH.center()
21
22    f3dname = 'stm3d.cube'
23
24    def testSTM(calc):
25        stm = SimpleStm(calc)
26        stm.write_3D([1, 0, 0], f3dname)  # single wf
27        wf = stm.gd.integrate(stm.ldos)
28
29        if size == 1:  # XXXX might have trouble reading in parallel
30            stm2 = SimpleStm(f3dname)
31            wf2 = stm2.gd.integrate(stm2.ldos)
32            print('Integrals: written, read=', wf, wf2)
33            equal(wf, wf2, 2.e-6)
34
35        stm.scan_const_current(0.02, 5)
36    #    print eigenvalue_string(calc)
37        stm.write_3D(3.1, f3dname)
38        wf2 = stm.gd.integrate(stm.ldos)
39    #    print "wf2=", wf2
40        equal(wf2, 2, 0.12)
41
42        return wf
43
44    # finite system without spin and width
45    fname = 'BH-nospin_wfs.gpw'
46    if not load:
47        BH.set_pbc(False)
48        cf = GPAW(nbands=3, h=.3, txt=txt)
49        BH.calc = cf
50        e1 = BH.get_potential_energy()
51        cf.write(fname, 'all')
52    else:
53        cf = GPAW(fname, txt=txt)
54    wf = testSTM(cf)
55
56    # finite system with spin
57    fname = 'BH-spin_Sz2_wfs.gpw'
58    BH.set_initial_magnetic_moments([1, 1])
59    if not load:
60        BH.set_pbc(False)
61        cf = GPAW(occupations=FermiDirac(0.1, fixmagmom=True),
62                  nbands=5,
63                  h=0.3,
64                  txt=txt)
65        BH.calc = cf
66        e2 = BH.get_potential_energy()
67        cf.write(fname, 'all')
68    else:
69        cf = GPAW(fname, txt=txt)
70    testSTM(cf)
71
72    # periodic system
73    if not load:
74        BH.set_pbc(True)
75        cp = GPAW(spinpol=True, nbands=3, h=.3,
76                  kpts=(2, 1, 1), txt=txt, symmetry='off')
77        BH.calc = cp
78        e3 = BH.get_potential_energy()
79        cp.write('BH-4kpts_wfs.gpw', 'all')
80    else:
81        cp = GPAW('BH-4kpts_wfs.gpw', txt=txt)
82
83    stmp = SimpleStm(cp)
84
85    stmp.write_3D(-4., f3dname)
86    print(me + 'Integrals(occ): 2 * wf, bias=', 2 * wf,
87          stmp.gd.integrate(stmp.ldos))
88    equal(2 * wf, stmp.gd.integrate(stmp.ldos), 0.02)
89
90    stmp.write_3D(+4., f3dname)
91    print(me + 'Integrals(unocc): 2 * wf, bias=', end=' ')
92    print(2 * wf, stmp.gd.integrate(stmp.ldos))
93    equal(2 * wf, stmp.gd.integrate(stmp.ldos), 0.02)
94
95    energy_tolerance = 0.007
96    equal(e1, -2.54026, energy_tolerance)
97    equal(e2, -1.51101, energy_tolerance)
98    equal(e3, -2.83573, energy_tolerance)
99