1import os
2from pyscf.pbc.gto import Cell
3from pyscf.pbc.scf import RHF, KRHF
4from pyscf.pbc.tdscf import KTDHF
5from pyscf.pbc.tdscf import krhf_slow_supercell as ktd, rhf_slow as td
6from pyscf.pbc.tools.pbc import super_cell
7from pyscf.tdscf.common_slow import eig
8
9from test_common import retrieve_m, retrieve_m_hf, adjust_mf_phase, ov_order, assert_vectors_close, tdhf_frozen_mask
10
11import unittest
12from numpy import testing
13import numpy
14
15
16class DiamondTestGamma(unittest.TestCase):
17    """Compare this (krhf_supercell_slow) @Gamma vs reference (pyscf)."""
18    @classmethod
19    def setUpClass(cls):
20        cls.cell = cell = Cell()
21        # Lift some degeneracies
22        cell.atom = '''
23        C 0.000000000000   0.000000000000   0.000000000000
24        C 1.67   1.68   1.69
25        '''
26        cell.basis = {'C': [[0, (0.8, 1.0)],
27                            [1, (1.0, 1.0)]]}
28        # cell.basis = 'gth-dzvp'
29        cell.pseudo = 'gth-pade'
30        cell.a = '''
31        0.000000000, 3.370137329, 3.370137329
32        3.370137329, 0.000000000, 3.370137329
33        3.370137329, 3.370137329, 0.000000000'''
34        cell.unit = 'B'
35        cell.verbose = 5
36        cell.build()
37
38        cls.model_krhf = model_krhf = KRHF(cell).density_fit()
39        model_krhf.kernel()
40
41        cls.td_model_krhf = td_model_krhf = KTDHF(model_krhf)
42        td_model_krhf.nroots = 5
43        td_model_krhf.kernel()
44
45        cls.ref_m_krhf = retrieve_m(td_model_krhf)
46
47    @classmethod
48    def tearDownClass(cls):
49        # These are here to remove temporary files
50        del cls.td_model_krhf
51        del cls.model_krhf
52        del cls.cell
53
54    def test_eri(self):
55        """Tests all ERI implementations: with and without symmetries."""
56        for eri in (ktd.PhysERI, ktd.PhysERI4, ktd.PhysERI8):
57            e = eri(self.model_krhf)
58            m = e.tdhf_full_form()
59            try:
60                testing.assert_allclose(self.ref_m_krhf, m, atol=1e-14)
61                vals, vecs = eig(m, nroots=self.td_model_krhf.nroots)
62                testing.assert_allclose(vals, self.td_model_krhf.e, atol=1e-5)
63            except Exception:
64                print("When testing {} the following exception occurred:".format(eri))
65                raise
66
67    def test_class(self):
68        """Tests container behavior."""
69        model = ktd.TDRHF(self.model_krhf)
70        model.nroots = self.td_model_krhf.nroots
71        assert model.fast
72        model.kernel()
73        e, xy = model.kernel()
74        model.fast = False
75        model.kernel()
76        # Slow vs fast
77        testing.assert_allclose(model.e, e)
78        assert_vectors_close(model.xy, xy)
79        # ... vs ref
80        testing.assert_allclose(model.e, self.td_model_krhf.e, atol=1e-5)
81        assert_vectors_close(model.xy.squeeze(), numpy.array(self.td_model_krhf.xy).squeeze(), atol=1e-12)
82
83
84class DiamondTestShiftedGamma(unittest.TestCase):
85    """Compare this (krhf_supercell_slow) @non-Gamma vs reference (rhf_slow)."""
86    @classmethod
87    def setUpClass(cls):
88        cls.cell = cell = Cell()
89        # Lift some degeneracies
90        cell.atom = '''
91        C 0.000000000000   0.000000000000   0.000000000000
92        C 1.67   1.68   1.69
93        '''
94        cell.basis = {'C': [[0, (0.8, 1.0)],
95                            [1, (1.0, 1.0)]]}
96        # cell.basis = 'gth-dzvp'
97        cell.pseudo = 'gth-pade'
98        cell.a = '''
99        0.000000000, 3.370137329, 3.370137329
100        3.370137329, 0.000000000, 3.370137329
101        3.370137329, 3.370137329, 0.000000000'''
102        cell.unit = 'B'
103        cell.verbose = 5
104        cell.build()
105
106        k = cell.get_abs_kpts((.1, .2, .3))
107
108        # The Gamma-point reference
109        cls.model_rhf = model_rhf = RHF(cell, k).density_fit()
110        model_rhf.conv_tol = 1e-14
111        model_rhf.kernel()
112
113        # K-points
114        cls.model_krhf = model_krhf = KRHF(cell, k).density_fit()
115        model_krhf.conv_tol = 1e-14
116        model_krhf.kernel()
117
118        adjust_mf_phase(model_rhf, model_krhf)
119
120        testing.assert_allclose(model_rhf.mo_energy, model_krhf.mo_energy[0])
121        testing.assert_allclose(model_rhf.mo_coeff, model_krhf.mo_coeff[0])
122
123        # The Gamma-point TD
124        cls.td_model_rhf = td_model_rhf = td.TDRHF(model_rhf)
125        td_model_rhf.kernel()
126        cls.ref_m = td_model_rhf.eri.tdhf_full_form()
127
128    @classmethod
129    def tearDownClass(cls):
130        # These are here to remove temporary files
131        del cls.td_model_rhf
132        del cls.model_krhf
133        del cls.model_rhf
134        del cls.cell
135
136    def test_eri(self):
137        """Tests all ERI implementations: with and without symmetries."""
138        for eri in (ktd.PhysERI, ktd.PhysERI4):
139            try:
140                e = eri(self.model_krhf)
141                m = e.tdhf_full_form()
142
143                # Test matrix vs ref
144                testing.assert_allclose(m, retrieve_m_hf(e), atol=1e-14)
145
146                # Test matrix vs pyscf
147                testing.assert_allclose(self.ref_m, m, atol=1e-10)
148            except Exception:
149                print("When testing {} the following exception occurred:".format(eri))
150                raise
151
152    def test_class(self):
153        """Tests container behavior."""
154        model = ktd.TDRHF(self.model_krhf)
155        model.nroots = self.td_model_rhf.nroots
156        assert not model.fast
157        model.kernel()
158        testing.assert_allclose(model.e, self.td_model_rhf.e, atol=1e-5)
159        nocc = nvirt = 4
160        testing.assert_equal(model.xy.shape, (len(model.e), 2, 1, 1, nocc, nvirt))
161        assert_vectors_close(model.xy.squeeze(), numpy.array(self.td_model_rhf.xy).squeeze(), atol=1e-9)
162
163
164class DiamondTestSupercell2(unittest.TestCase):
165    """Compare this (krhf_supercell_slow) @2kp vs supercell reference (rhf_slow)."""
166    k = 2
167    k_c = (0, 0, 0)
168    test8 = True
169
170    @classmethod
171    def setUpClass(cls):
172        cls.cell = cell = Cell()
173        # Lift some degeneracies
174        cell.atom = '''
175        C 0.000000000000   0.000000000000   0.000000000000
176        C 1.67   1.68   1.69
177        '''
178        cell.basis = {'C': [[0, (0.8, 1.0)],
179                            [1, (1.0, 1.0)]]}
180        # cell.basis = 'gth-dzvp'
181        cell.pseudo = 'gth-pade'
182        cell.a = '''
183        0.000000000, 3.370137329, 3.370137329
184        3.370137329, 0.000000000, 3.370137329
185        3.370137329, 3.370137329, 0.000000000'''
186        cell.unit = 'B'
187        cell.verbose = 5
188        cell.build()
189
190        k = cell.make_kpts([cls.k, 1, 1], scaled_center=cls.k_c)
191
192        # The Gamma-point reference
193        cls.model_rhf = model_rhf = RHF(super_cell(cell, [cls.k, 1, 1]), kpt=k[0]).density_fit()
194        model_rhf.conv_tol = 1e-14
195        model_rhf.kernel()
196
197        # K-points
198        cls.model_krhf = model_krhf = KRHF(cell, k).density_fit()
199        model_krhf.conv_tol = 1e-14
200        model_krhf.kernel()
201
202        adjust_mf_phase(model_rhf, model_krhf)
203
204        ke = numpy.concatenate(model_krhf.mo_energy)
205        ke.sort()
206
207        # Make sure mo energies are the same
208        testing.assert_allclose(model_rhf.mo_energy, ke)
209
210        # Make sure no degeneracies are present
211        testing.assert_array_less(1e-4, ke[1:] - ke[:-1])
212
213        cls.ov_order = ov_order(model_krhf)
214
215        # The Gamma-point TD
216        cls.td_model_rhf = td_model_rhf = td.TDRHF(model_rhf)
217        td_model_rhf.kernel()
218        cls.ref_m = td_model_rhf.eri.tdhf_full_form()
219
220    @classmethod
221    def tearDownClass(cls):
222        # These are here to remove temporary files
223        del cls.td_model_rhf
224        del cls.model_krhf
225        del cls.model_rhf
226        del cls.cell
227
228    def test_eri(self):
229        """Tests all ERI implementations: with and without symmetries."""
230        for eri in (ktd.PhysERI, ktd.PhysERI4, ktd.PhysERI8):
231            if not eri == ktd.PhysERI8 or self.test8:
232                try:
233                    e = eri(self.model_krhf)
234                    m = e.tdhf_full_form()
235
236                    # Test matrix vs ref
237                    testing.assert_allclose(m, retrieve_m_hf(e), atol=1e-11)
238
239                    # Test matrix vs pyscf
240                    testing.assert_allclose(self.ref_m, m[numpy.ix_(self.ov_order, self.ov_order)], atol=1e-5)
241                except Exception:
242                    print("When testing {} the following exception occurred:".format(eri))
243                    raise
244
245    def test_class(self):
246        """Tests container behavior."""
247        model = ktd.TDRHF(self.model_krhf)
248        model.nroots = self.td_model_rhf.nroots
249        assert not model.fast
250        model.kernel()
251        testing.assert_allclose(model.e, self.td_model_rhf.e, atol=1e-5)
252        nocc = nvirt = 4
253        testing.assert_equal(model.xy.shape, (len(model.e), 2, self.k, self.k, nocc, nvirt))
254        vecs = model.xy.reshape(len(model.xy), -1)[:, self.ov_order]
255        assert_vectors_close(vecs, numpy.array(self.td_model_rhf.xy).squeeze(), atol=1e-5)
256        # Test real
257        testing.assert_allclose(model.e.imag, 0, atol=1e-8)
258
259
260class DiamondTestSupercell3(DiamondTestSupercell2):
261    """Compare this (supercell_slow) @3kp vs supercell reference (rhf_slow)."""
262    k = 3
263    k_c = (.1, 0, 0)
264    test8 = False
265
266
267class FrozenTest(unittest.TestCase):
268    """Tests frozen behavior."""
269    k = 2
270    k_c = (0, 0, 0)
271    df_file = os.path.realpath(os.path.join(__file__, "..", "frozen_test_cderi.h5"))
272
273    @classmethod
274    def setUpClass(cls):
275        cls.cell = cell = Cell()
276        # Lift some degeneracies
277        cell.atom = '''
278        C 0.000000000000   0.000000000000   0.000000000000
279        C 1.67   1.68   1.69
280        '''
281        cell.basis = 'sto-3g'
282        cell.a = '''
283        0.000000000, 3.370137329, 3.370137329
284        3.370137329, 0.000000000, 3.370137329
285        3.370137329, 3.370137329, 0.000000000'''
286        cell.unit = 'B'
287        cell.verbose = 5
288        cell.build()
289
290        k = cell.make_kpts([cls.k, 1, 1], scaled_center=cls.k_c)
291
292        # K-points
293        cls.model_krhf = model_krhf = KRHF(cell, k).density_fit()
294        # model_krhf.with_df._cderi_to_save = cls.df_file
295        model_krhf.with_df._cderi = cls.df_file
296        model_krhf.conv_tol = 1e-14
297        model_krhf.kernel()
298
299        cls.td_model_krhf = model_ktd = ktd.TDRHF(model_krhf)
300        model_ktd.nroots = 5
301        model_ktd.kernel()
302
303    @classmethod
304    def tearDownClass(cls):
305        # These are here to remove temporary files
306        del cls.td_model_krhf
307        del cls.model_krhf
308        del cls.cell
309
310    def test_class(self):
311        """Tests container behavior (frozen vs non-frozen)."""
312        for frozen in (1, [0, 1]):
313            try:
314                model = ktd.TDRHF(self.model_krhf, frozen=frozen)
315                model.nroots = self.td_model_krhf.nroots
316                model.kernel()
317                mask_o, mask_v = tdhf_frozen_mask(model.eri, kind="o,v")
318                testing.assert_allclose(model.e, self.td_model_krhf.e, atol=1e-4)
319                assert_vectors_close(
320                    model.xy,
321                    numpy.array(self.td_model_krhf.xy)[..., mask_o, :][..., mask_v],
322                    atol=1e-3,
323                )
324
325            except Exception:
326                print("When testing class with frozen={} the following exception occurred:".format(repr(frozen)))
327                raise
328