1# Copyright (C) 2020 Atsushi Togo
2# All rights reserved.
3#
4# This file is part of phono3py.
5#
6# Redistribution and use in source and binary forms, with or without
7# modification, are permitted provided that the following conditions
8# are met:
9#
10# * Redistributions of source code must retain the above copyright
11#   notice, this list of conditions and the following disclaimer.
12#
13# * Redistributions in binary form must reproduce the above copyright
14#   notice, this list of conditions and the following disclaimer in
15#   the documentation and/or other materials provided with the
16#   distribution.
17#
18# * Neither the name of the phonopy project nor the names of its
19#   contributors may be used to endorse or promote products derived
20#   from this software without specific prior written permission.
21#
22# THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS
23# "AS IS" AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT
24# LIMITED TO, THE IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS
25# FOR A PARTICULAR PURPOSE ARE DISCLAIMED. IN NO EVENT SHALL THE
26# COPYRIGHT HOLDER OR CONTRIBUTORS BE LIABLE FOR ANY DIRECT, INDIRECT,
27# INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL DAMAGES (INCLUDING,
28# BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR SERVICES;
29# LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER
30# CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT
31# LIABILITY, OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN
32# ANY WAY OUT OF THE USE OF THIS SOFTWARE, EVEN IF ADVISED OF THE
33# POSSIBILITY OF SUCH DAMAGE.
34
35import sys
36import numpy as np
37from phono3py.phonon3.imag_self_energy import (
38    run_ise_at_frequency_points_batch, get_frequency_points, ImagSelfEnergy)
39from phono3py.phonon3.real_self_energy import imag_to_real
40from phono3py.file_IO import (
41    write_spectral_function_at_grid_point, write_spectral_function_to_hdf5)
42
43
44def run_spectral_function(interaction,
45                          grid_points,
46                          frequency_points=None,
47                          frequency_step=None,
48                          num_frequency_points=None,
49                          num_points_in_batch=None,
50                          sigmas=None,
51                          temperatures=None,
52                          band_indices=None,
53                          write_txt=False,
54                          write_hdf5=False,
55                          output_filename=None,
56                          log_level=0):
57    spf = SpectralFunction(interaction,
58                           grid_points,
59                           frequency_points=frequency_points,
60                           frequency_step=frequency_step,
61                           num_frequency_points=num_frequency_points,
62                           num_points_in_batch=num_points_in_batch,
63                           sigmas=sigmas,
64                           temperatures=temperatures,
65                           log_level=log_level)
66    for i, gp in enumerate(spf):
67        frequencies = interaction.get_phonons()[0]
68        for sigma_i, sigma in enumerate(spf.sigmas):
69            for t, spf_at_t in zip(
70                    temperatures, spf.spectral_functions[sigma_i, i]):
71                for j, bi in enumerate(band_indices):
72                    pos = 0
73                    for k in range(j):
74                        pos += len(band_indices[k])
75                    filename = write_spectral_function_at_grid_point(
76                        gp,
77                        bi,
78                        spf.frequency_points,
79                        spf_at_t[pos:(pos + len(bi))].sum(axis=0) / len(bi),
80                        interaction.mesh_numbers,
81                        t,
82                        sigma=sigma,
83                        filename=output_filename,
84                        is_mesh_symmetry=interaction.is_mesh_symmetry)
85                    if log_level:
86                        print("Spectral functions were written to")
87                        print("\"%s\"." % filename)
88
89            filename = write_spectral_function_to_hdf5(
90                gp,
91                band_indices,
92                temperatures,
93                spf.spectral_functions[sigma_i, i],
94                spf.shifts[sigma_i, i],
95                spf.half_linewidths[sigma_i, i],
96                interaction.mesh_numbers,
97                sigma=sigma,
98                frequency_points=spf.frequency_points,
99                frequencies=frequencies[gp],
100                filename=output_filename)
101
102            if log_level:
103                print("Spectral functions were stored in \"%s\"." % filename)
104                sys.stdout.flush()
105
106    return spf
107
108
109class SpectralFunction(object):
110    """Calculate spectral function"""
111
112    def __init__(self,
113                 interaction,
114                 grid_points,
115                 frequency_points=None,
116                 frequency_step=None,
117                 num_frequency_points=None,
118                 num_points_in_batch=None,
119                 sigmas=None,
120                 temperatures=None,
121                 log_level=0):
122        self._interaction = interaction
123        self._grid_points = grid_points
124        self._frequency_points_in = frequency_points
125        self._num_points_in_batch = num_points_in_batch
126        self._frequency_step = frequency_step
127        self._num_frequency_points = num_frequency_points
128        self._temperatures = temperatures
129        self._log_level = log_level
130
131        if sigmas is None:
132            self._sigmas = [None, ]
133        else:
134            self._sigmas = sigmas
135        self._frequency_points = None
136        self._gammas = None
137        self._deltas = None
138        self._spectral_functions = None
139        self._gp_index = None
140
141    def run(self):
142        for gp_index in self:
143            pass
144
145    def __iter__(self):
146        self._prepare()
147
148        return self
149
150    def next(self):
151        return self.__next__()
152
153    def __next__(self):
154        if self._gp_index >= len(self._grid_points):
155            if self._log_level:
156                print("-" * 74)
157            raise StopIteration
158
159        if self._log_level:
160            print(("-" * 24 + " Spectral function (%d/%d) " + "-" * 24)
161                  % (self._gp_index + 1, len(self._grid_points)))
162
163        gp = self._grid_points[self._gp_index]
164        ise = ImagSelfEnergy(self._interaction)
165        ise.set_grid_point(gp)
166
167        if self._log_level:
168            print("Running ph-ph interaction calculation...")
169            sys.stdout.flush()
170
171        ise.run_interaction()
172
173        for sigma_i, sigma in enumerate(self._sigmas):
174            self._run_gamma(ise, self._gp_index, sigma, sigma_i)
175            self._run_delta(self._gp_index, sigma_i)
176            self._run_spectral_function(self._gp_index, gp, sigma_i)
177
178        self._gp_index += 1
179        return gp
180
181    @property
182    def spectral_functions(self):
183        return self._spectral_functions
184
185    @property
186    def shifts(self):
187        return self._deltas
188
189    @property
190    def half_linewidths(self):
191        return self._gammas
192
193    @property
194    def frequency_points(self):
195        return self._frequency_points
196
197    @property
198    def grid_points(self):
199        return self._grid_points
200
201    @property
202    def sigmas(self):
203        return self._sigmas
204
205    def _prepare(self):
206        self._set_frequency_points()
207        self._gammas = np.zeros(
208            (len(self._sigmas), len(self._grid_points),
209             len(self._temperatures), len(self._interaction.band_indices),
210             len(self._frequency_points)),
211            dtype='double', order='C')
212        self._deltas = np.zeros_like(self._gammas)
213        self._spectral_functions = np.zeros_like(self._gammas)
214        self._gp_index = 0
215
216    def _run_gamma(self, ise, i, sigma, sigma_i):
217        if self._log_level:
218            print("* Imaginary part of self energy")
219
220        ise.set_sigma(sigma)
221        run_ise_at_frequency_points_batch(
222            self._frequency_points,
223            ise,
224            self._temperatures,
225            self._gammas[sigma_i, i],
226            nelems_in_batch=self._num_points_in_batch,
227            log_level=self._log_level)
228
229    def _run_delta(self, i, sigma_i):
230        if self._log_level:
231            print("* Real part of self energy")
232            print("Running Kramers-Kronig relation integration...")
233
234        for j, temp in enumerate(self._temperatures):
235            for k, bi in enumerate(self._interaction.band_indices):
236                re_part, fpoints = imag_to_real(
237                    self._gammas[sigma_i, i, j, k], self._frequency_points)
238                self._deltas[sigma_i, i, j, k] = -re_part
239        assert (np.abs(self._frequency_points - fpoints) < 1e-8).all()
240
241    def _run_spectral_function(self, i, grid_point, sigma_i):
242        """Compute spectral functions from self-energies
243
244        Note
245        ----
246        Spectral function A is defined by
247
248                -1        G_0
249            A = -- Im -----------
250                pi    1 - G_0 Pi
251
252        where pi = 3.14..., and Pi is the self energy, Pi = Delta - iGamma.
253        It is expected that the integral of A over frequency is
254        approximately 1 for each phonon mode.
255
256        """
257
258        if self._log_level:
259            print("* Spectral function")
260        frequencies = self._interaction.get_phonons()[0]
261        for j, temp in enumerate(self._temperatures):
262            for k, bi in enumerate(self._interaction.band_indices):
263                freq = frequencies[grid_point, bi]
264                gammas = self._gammas[sigma_i, i, j, k]
265                deltas = self._deltas[sigma_i, i, j, k]
266                vals = self._get_spectral_function(gammas, deltas, freq)
267                self._spectral_functions[sigma_i, i, j, k] = vals
268
269    def _get_spectral_function(self, gammas, deltas, freq):
270        fpoints = self._frequency_points
271        nums = 4 * freq ** 2 * gammas / np.pi
272        denoms = ((fpoints ** 2 - freq ** 2 - 2 * freq * deltas) ** 2
273                  + (2 * freq * gammas) ** 2)
274        vals = np.where(denoms > 0, nums / denoms, 0)
275        return vals
276
277    def _set_frequency_points(self):
278        if (self._interaction.get_phonons()[2] == 0).any():
279            if self._log_level:
280                print("Running harmonic phonon calculations...")
281            self._interaction.run_phonon_solver()
282        max_phonon_freq = np.amax(self._interaction.get_phonons()[0])
283        self._frequency_points = get_frequency_points(
284            max_phonon_freq=max_phonon_freq,
285            sigmas=self._sigmas,
286            frequency_points=self._frequency_points_in,
287            frequency_step=self._frequency_step,
288            num_frequency_points=self._num_frequency_points)
289