1"""Integration tests for structural relaxations."""
2
3import numpy as np
4import pytest
5import abipy.data as abidata
6import abipy.abilab as abilab
7import abipy.flowtk as flowtk
8
9from abipy.core.testing import has_matplotlib
10
11
12def ion_relaxation(tvars, ntime=50):
13    structure = abilab.Structure.from_file(abidata.cif_file("si.cif"))
14
15    # Perturb the structure (random perturbation of 0.1 Angstrom)
16    structure.perturb(distance=0.02)
17
18    global_vars = dict(
19        ecut=6,
20        ngkpt=[2,2,2],
21        shiftk=[0,0,0],
22        nshiftk=1,
23        chksymbreak=0,
24        paral_kgb=tvars.paral_kgb,
25    )
26
27    inp = abilab.AbinitInput(structure, pseudos=abidata.pseudos("14si.pspnc"))
28
29    # Global variables
30    inp.set_vars(global_vars)
31
32    # Dataset 1 (Atom Relaxation)
33    #inp[1].set_vars(
34    # FIXME here there's a bug
35    inp.set_vars(
36        optcell=0,
37        ionmov=2,
38        tolrff=0.02,
39        tolmxf=5.0e-5,
40        ntime=ntime,
41        #dilatmx=1.05, # FIXME: abinit crashes if I don't use this
42    )
43
44    return inp
45
46
47def itest_atomic_relaxation(fwp, tvars):
48    """Test atomic relaxation with automatic restart."""
49    # Build the flow
50    flow = flowtk.Flow(fwp.workdir, manager=fwp.manager)
51
52    ion_input = ion_relaxation(tvars, ntime=2)
53    ion_input["chksymtnons"] = 0
54    work = flow.register_task(ion_input, task_class=flowtk.RelaxTask)
55
56    flow.allocate()
57    flow.build_and_pickle_dump(abivalidate=True)
58
59    # Run t0, and check status
60    t0 = work[0]
61    t0.start_and_wait()
62    assert t0.returncode == 0
63    t0.check_status()
64    assert t0.status == t0.S_UNCONVERGED
65
66    assert t0.initial_structure == ion_input.structure
67    unconverged_structure = t0.get_final_structure()
68    assert unconverged_structure != t0.initial_structure
69
70    # Use the default value ntime=50 and we can converge the calculation.
71    t0.input.set_vars(ntime=50)
72
73    t0.build()
74    assert t0.restart()
75    t0.wait()
76    assert t0.num_restarts == 1
77
78    # At this point, we should have reached S_OK.
79    assert t0.status == t0.S_DONE
80    t0.check_status()
81    assert t0.status == t0.S_OK
82
83    final_structure = t0.get_final_structure()
84    assert final_structure != unconverged_structure
85
86    flow.show_status()
87    assert all(work.finalized for work in flow)
88    if not flow.all_ok:
89        flow.debug()
90        raise RuntimeError()
91
92    # post-processing tools
93    if has_matplotlib():
94        assert t0.inspect(show=False) is not None
95
96    with t0.open_hist() as hist:
97        hist.to_string(verbose=2)
98        # from_file accepts HIST files as well.
99        assert np.all(hist.structures[-1].frac_coords == abilab.Structure.from_file(hist.filepath).frac_coords)
100
101    with t0.open_gsr() as gsr:
102        gsr.to_string(verbose=2)
103        gsr.pressure == 1.8280
104
105    t0.get_results()
106
107
108def make_ion_ioncell_inputs(tvars, dilatmx, scalevol=1, ntime=50):
109    structure = abilab.Structure.from_file(abidata.cif_file("si.cif"))
110
111    # Perturb the structure (random perturbation of 0.1 Angstrom)
112    #structure.perturb(distance=0.01)
113
114    # Compress the lattice so that ABINIT complains about dilatmx
115    structure.scale_lattice(structure.volume * scalevol)
116
117    global_vars = dict(
118        ecut=6,
119        ecutsm=0.5,
120        ngkpt=[4,4,4],
121        shiftk=[0,0,0],
122        nshiftk=1,
123        chksymbreak=0,
124        paral_kgb=tvars.paral_kgb,
125    )
126
127    multi = abilab.MultiDataset(structure, pseudos=abidata.pseudos("14si.pspnc"), ndtset=2)
128
129    # Global variables
130    multi.set_vars(global_vars)
131
132    # Dataset 1 (Atom Relaxation)
133    multi[0].set_vars(
134        optcell=0,
135        ionmov=2,
136        tolrff=0.02,
137        tolmxf=5.0e-5,
138        ntime=ntime,
139    )
140
141    # Dataset 2 (Atom + Cell Relaxation)
142    multi[1].set_vars(
143        optcell=1,
144        ionmov=2,
145        dilatmx=dilatmx,
146        tolrff=0.02,
147        tolmxf=5.0e-5,
148        strfact=100,
149        ntime=ntime,
150    )
151
152    ion_inp, ioncell_inp = multi.split_datasets()
153    return ion_inp, ioncell_inp
154
155
156def itest_relaxation_with_restart_from_den(fwp, tvars):
157    """Test structural relaxations with automatic restart from DEN files."""
158    # Build the flow
159    flow = flowtk.Flow(fwp.workdir, manager=fwp.manager)
160
161    # Use small value for ntime to trigger restart, then disable the output of the WFK file.
162    ion_input, ioncell_input = make_ion_ioncell_inputs(tvars, dilatmx=1.1, ntime=3)
163    ion_input.set_vars(prtwf=0)
164    ioncell_input.set_vars(prtwf=0)
165
166    relax_work = flowtk.RelaxWork(ion_input, ioncell_input)
167    flow.register_work(relax_work)
168
169    assert flow.make_scheduler().start() == 0
170    flow.show_status()
171
172    assert all(work.finalized for work in flow)
173    if not flow.all_ok:
174        flow.debug()
175        raise RuntimeError()
176
177    # we should have (0, 1) restarts and no WFK file in outdir.
178    for i, task in enumerate(relax_work):
179        assert task.status == task.S_OK
180        assert task.num_restarts == i
181        assert task.num_corrections == 0
182        assert not task.outdir.has_abiext("WFK")
183
184    if has_matplotlib:
185        assert relax_work.plot_ion_relaxation(show=False) is not None
186        assert relax_work.plot_ioncell_relaxation(show=False) is not None
187
188    flow.rmtree()
189
190
191def itest_dilatmx_error_handler(fwp, tvars):
192    """
193    Test cell relaxation with automatic restart in the presence of dilatmx error.
194    """
195    pytest.xfail("dilatmxerror_handler is not portable and it's been disabled!")
196    # Build the flow
197    flow = flowtk.Flow(fwp.workdir, manager=fwp.manager)
198
199    # Decrease the volume to trigger DilatmxError
200    ion_input, ioncell_input = make_ion_ioncell_inputs(tvars, dilatmx=1.01, scalevol=0.8)
201
202    work = flowtk.Work()
203    work.register_relax_task(ioncell_input)
204
205    flow.register_work(work)
206    flow.allocate()
207    assert flow.make_scheduler().start() == 0
208    flow.show_status()
209
210    assert all(work.finalized for work in flow)
211    if not flow.all_ok:
212        flow.debug()
213        raise RuntimeError()
214
215    # t0 should have reached S_OK, and we should have DilatmxError in the corrections.
216    t0 = work[0]
217    assert t0.status == t0.S_OK
218    print(t0.corrections)
219    assert t0.num_corrections > 0
220    assert t0.corrections[0]["event"]["@class"] == "DilatmxError"
221
222
223def itest_relaxation_with_target_dilatmx(fwp, tvars):
224    """Test structural relaxations with automatic restart from DEN files."""
225    # Build the flow
226    flow = flowtk.Flow(fwp.workdir, manager=fwp.manager)
227
228    # Use small value for ntime to trigger restart, then disable the output of the WFK file.
229    ion_input, ioncell_input = make_ion_ioncell_inputs(tvars, dilatmx=1.05)
230
231    target_dilatmx = 1.03
232    relax_work = flowtk.RelaxWork(ion_input, ioncell_input, target_dilatmx=target_dilatmx)
233    flow.register_work(relax_work)
234
235    assert flow.make_scheduler().start() == 0
236    flow.show_status()
237
238    assert all(work.finalized for work in flow)
239    if not flow.all_ok:
240        flow.debug()
241        raise RuntimeError()
242    #assert relax_work.last_dilatmx <= target_dilatmx
243
244    # we should have (0, 1) restarts
245    for i, task in enumerate(relax_work):
246        assert task.status == task.S_OK
247        assert task.num_restarts == i
248        assert task.num_corrections == 0
249
250    assert relax_work[1].input["dilatmx"] == 1.03
251
252    # check that when decreasing the dilatmx it actually takes the previously relaxed
253    # structure and does not start from scratch again: the lattice should not be the same.
254    assert not np.allclose(relax_work.ion_task.get_final_structure().lattice_vectors(),
255                           relax_work.ioncell_task.input.structure.lattice_vectors())
256
257    flow.rmtree()
258