1# Copyright (C) 2011 Atsushi Togo
2# All rights reserved.
3#
4# This file is part of phonopy.
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 phonopy.phonon.tetrahedron_mesh import TetrahedronMesh
38from phonopy.structure.tetrahedron_method import TetrahedronMethod
39
40
41def get_pdos_indices(symmetry):
42    mapping = symmetry.get_map_atoms()
43    return [list(np.where(mapping == i)[0])
44            for i in symmetry.get_independent_atoms()]
45
46
47def write_total_dos(frequency_points,
48                    total_dos,
49                    comment=None,
50                    filename="total_dos.dat"):
51    with open(filename, 'w') as fp:
52        if comment is not None:
53            fp.write("# %s\n" % comment)
54
55        for freq, dos in zip(frequency_points, total_dos):
56            fp.write("%20.10f%20.10f\n" % (freq, dos))
57
58
59def write_partial_dos(frequency_points,
60                      partial_dos,
61                      comment=None,
62                      filename="partial_dos.dat"):
63    with open(filename, 'w') as fp:
64        if comment is not None:
65            fp.write("# %s\n" % comment)
66
67        for freq, pdos in zip(frequency_points, partial_dos.T):
68            fp.write("%20.10f" % freq)
69            fp.write(("%20.10f" * len(pdos)) % tuple(pdos))
70            fp.write("\n")
71
72
73def plot_total_dos(ax,
74                   frequency_points,
75                   total_dos,
76                   freq_Debye=None,
77                   Debye_fit_coef=None,
78                   xlabel=None,
79                   ylabel=None,
80                   draw_grid=True,
81                   flip_xy=False):
82    ax.xaxis.set_ticks_position('both')
83    ax.yaxis.set_ticks_position('both')
84    ax.xaxis.set_tick_params(which='both', direction='in')
85    ax.yaxis.set_tick_params(which='both', direction='in')
86
87    if freq_Debye is not None:
88        freq_pitch = frequency_points[1] - frequency_points[0]
89        num_points = int(freq_Debye / freq_pitch)
90        freqs = np.linspace(0, freq_Debye, num_points + 1)
91
92    if flip_xy:
93        ax.plot(total_dos, frequency_points, 'r-', linewidth=1)
94        if freq_Debye:
95            ax.plot(np.append(Debye_fit_coef * freqs**2, 0),
96                    np.append(freqs, freq_Debye), 'b-', linewidth=1)
97    else:
98        ax.plot(frequency_points, total_dos, 'r-', linewidth=1)
99        if freq_Debye:
100            ax.plot(np.append(freqs, freq_Debye),
101                    np.append(Debye_fit_coef * freqs**2, 0), 'b-', linewidth=1)
102
103    if xlabel:
104        ax.set_xlabel(xlabel)
105    if ylabel:
106        ax.set_ylabel(ylabel)
107
108    ax.grid(draw_grid)
109
110
111def plot_partial_dos(ax,
112                     frequency_points,
113                     partial_dos,
114                     indices=None,
115                     legend=None,
116                     xlabel=None,
117                     ylabel=None,
118                     draw_grid=True,
119                     flip_xy=False):
120    ax.xaxis.set_ticks_position('both')
121    ax.yaxis.set_ticks_position('both')
122    ax.xaxis.set_tick_params(which='both', direction='in')
123    ax.yaxis.set_tick_params(which='both', direction='in')
124
125    plots = []
126    num_pdos = len(partial_dos)
127
128    if indices is None:
129        indices = []
130        for i in range(num_pdos):
131            indices.append([i])
132
133    for set_for_sum in indices:
134        pdos_sum = np.zeros_like(frequency_points)
135        for i in set_for_sum:
136            if i > num_pdos - 1:
137                print("Index number \'%d\' is specified," % (i + 1))
138                print("but it is not allowed to be larger than the number of "
139                      "atoms.")
140                raise ValueError
141            if i < 0:
142                print("Index number \'%d\' is specified, but it must be "
143                      "positive." % (i + 1))
144                raise ValueError
145            pdos_sum += partial_dos[i]
146        if flip_xy:
147            plots.append(ax.plot(pdos_sum, frequency_points, linewidth=1))
148        else:
149            plots.append(ax.plot(frequency_points, pdos_sum, linewidth=1))
150
151    if legend is not None:
152        ax.legend(legend)
153
154    if xlabel:
155        ax.set_xlabel(xlabel)
156    if ylabel:
157        ax.set_ylabel(ylabel)
158
159    ax.grid(draw_grid)
160
161
162class NormalDistribution(object):
163    def __init__(self, sigma):
164        self._sigma = sigma
165
166    def calc(self, x):
167        return 1.0 / np.sqrt(2 * np.pi) / self._sigma * \
168            np.exp(-x**2 / 2.0 / self._sigma**2)
169
170
171class CauchyDistribution(object):
172    def __init__(self, gamma):
173        self._gamma = gamma
174
175    def calc(self, x):
176        return self._gamma / np.pi / (x**2 + self._gamma**2)
177
178
179def run_tetrahedron_method_dos(mesh,
180                               frequency_points,
181                               frequencies,
182                               grid_address,
183                               grid_mapping_table,
184                               relative_grid_address,
185                               coef=None):  # for each grid point
186    try:
187        import phonopy._phonopy as phonoc
188    except ImportError:
189        import sys
190        print("Phonopy C-extension has to be built properly.")
191        sys.exit(1)
192
193    if coef is None:
194        _coef = np.ones((frequencies.shape[0], 1, frequencies.shape[1]),
195                        dtype='double')
196    else:
197        _coef = np.array(coef, dtype='double', order='C')
198    arr_shape = frequencies.shape + (len(frequency_points), _coef.shape[1])
199    dos = np.zeros(arr_shape, dtype='double')
200
201    phonoc.tetrahedron_method_dos(
202        dos,
203        np.array(mesh, dtype='int_'),
204        frequency_points,
205        frequencies,
206        _coef,
207        np.array(grid_address, dtype='int_', order='C'),
208        np.array(grid_mapping_table, dtype='int_', order='C'),
209        relative_grid_address)
210    if coef is None:
211        return dos[:, :, :, 0].sum(axis=0).sum(axis=0) / np.prod(mesh)
212    else:
213        return dos.sum(axis=0).sum(axis=0) / np.prod(mesh)
214
215
216class Dos(object):
217    def __init__(self, mesh_object, sigma=None, use_tetrahedron_method=False):
218        self._mesh_object = mesh_object
219        self._frequencies = mesh_object.frequencies
220        self._weights = mesh_object.weights
221        if use_tetrahedron_method and sigma is None:
222            self._tetrahedron_mesh = TetrahedronMesh(
223                mesh_object.dynamical_matrix.primitive,
224                self._frequencies,
225                mesh_object.mesh_numbers,
226                np.array(mesh_object.grid_address, dtype='int_'),
227                np.array(mesh_object.grid_mapping_table, dtype='int_'),
228                mesh_object.ir_grid_points)
229        else:
230            self._tetrahedron_mesh = None
231
232        self._frequency_points = None
233        self._sigma = sigma
234        self.set_draw_area()
235        self.set_smearing_function('Normal')
236
237    @property
238    def frequency_points(self):
239        return self._frequency_points
240
241    def set_smearing_function(self, function_name):
242        """
243        function_name ==
244        'Normal': smearing is done by normal distribution.
245        'Cauchy': smearing is done by Cauchy distribution.
246        """
247        if function_name == 'Cauchy':
248            self._smearing_function = CauchyDistribution(self._sigma)
249        else:
250            self._smearing_function = NormalDistribution(self._sigma)
251
252    def set_sigma(self, sigma):
253        self._sigma = sigma
254
255    def set_draw_area(self,
256                      freq_min=None,
257                      freq_max=None,
258                      freq_pitch=None):
259
260        f_min = self._frequencies.min()
261        f_max = self._frequencies.max()
262
263        if self._sigma is None:
264            self._sigma = (f_max - f_min) / 100.0
265
266        if freq_min is None:
267            f_min -= self._sigma * 10
268        else:
269            f_min = freq_min
270
271        if freq_max is None:
272            f_max += self._sigma * 10
273        else:
274            f_max = freq_max
275
276        if freq_pitch is None:
277            f_delta = (f_max - f_min) / 200.0
278        else:
279            f_delta = freq_pitch
280        self._frequency_points = np.arange(f_min,
281                                           f_max + f_delta * 0.1,
282                                           f_delta)
283
284
285class TotalDos(Dos):
286    def __init__(self, mesh_object, sigma=None, use_tetrahedron_method=False):
287        Dos.__init__(self,
288                     mesh_object,
289                     sigma=sigma,
290                     use_tetrahedron_method=use_tetrahedron_method)
291        self._dos = None
292        self._freq_Debye = None
293        self._Debye_fit_coef = None
294        self._openmp_thm = True
295
296    def run(self):
297        if self._tetrahedron_mesh is None:
298            self._dos = np.array([self._get_density_of_states_at_freq(f)
299                                  for f in self._frequency_points])
300        else:
301            if self._openmp_thm:
302                self._run_tetrahedron_method_dos()
303            else:
304                self._dos = np.zeros_like(self._frequency_points)
305                thm = self._tetrahedron_mesh
306                thm.set(value='I', frequency_points=self._frequency_points)
307                for i, iw in enumerate(thm):
308                    self._dos += np.sum(iw * self._weights[i], axis=1)
309
310    @property
311    def dos(self):
312        return self._dos
313
314    def get_dos(self):
315        """
316        Return freqs and total dos
317        """
318        return self._frequency_points, self._dos
319
320    def get_Debye_frequency(self):
321        return self._freq_Debye
322
323    def set_Debye_frequency(self, num_atoms, freq_max_fit=None):
324        try:
325            from scipy.optimize import curve_fit
326        except ImportError:
327            print("You need to install python-scipy.")
328            sys.exit(1)
329
330        def Debye_dos(freq, a):
331            return a * freq**2
332
333        freq_min = self._frequency_points.min()
334        freq_max = self._frequency_points.max()
335
336        if freq_max_fit is None:
337            N_fit = int(len(self._frequency_points) / 4.0)  # Hard coded
338        else:
339            N_fit = int(freq_max_fit / (freq_max - freq_min) *
340                        len(self._frequency_points))
341        popt, pcov = curve_fit(Debye_dos,
342                               self._frequency_points[0:N_fit],
343                               self._dos[0:N_fit])
344        a2 = popt[0]
345        self._freq_Debye = (3 * 3 * num_atoms / a2)**(1.0 / 3)
346        self._Debye_fit_coef = a2
347
348    def plot(self,
349             ax,
350             xlabel=None,
351             ylabel=None,
352             draw_grid=True,
353             flip_xy=False):
354        if flip_xy:
355            _xlabel = 'Density of states'
356            _ylabel = 'Frequency'
357        else:
358            _xlabel = 'Frequency'
359            _ylabel = 'Density of states'
360
361        if xlabel is not None:
362            _xlabel = xlabel
363        if ylabel is not None:
364            _ylabel = ylabel
365
366        plot_total_dos(ax,
367                       self._frequency_points,
368                       self._dos,
369                       freq_Debye=self._freq_Debye,
370                       Debye_fit_coef=self._Debye_fit_coef,
371                       xlabel=_xlabel,
372                       ylabel=_ylabel,
373                       draw_grid=draw_grid,
374                       flip_xy=flip_xy)
375
376    def write(self, filename="total_dos.dat"):
377        if self._tetrahedron_mesh is None:
378            comment = "Sigma = %f" % self._sigma
379        else:
380            comment = "Tetrahedron method"
381
382        write_total_dos(self._frequency_points,
383                        self._dos,
384                        comment=comment,
385                        filename=filename)
386
387    def _run_tetrahedron_method_dos(self):
388        mesh_numbers = self._mesh_object.mesh_numbers
389        cell = self._mesh_object.dynamical_matrix.primitive
390        reciprocal_lattice = np.linalg.inv(cell.get_cell())
391        tm = TetrahedronMethod(reciprocal_lattice, mesh=mesh_numbers)
392        self._dos = run_tetrahedron_method_dos(
393            mesh_numbers,
394            self._frequency_points,
395            self._frequencies,
396            self._mesh_object.grid_address,
397            self._mesh_object.grid_mapping_table,
398            tm.get_tetrahedra())
399
400    def _get_density_of_states_at_freq(self, f):
401        return np.sum(np.dot(
402            self._weights, self._smearing_function.calc(self._frequencies - f))
403        ) / np.sum(self._weights)
404
405
406class PartialDos(Dos):
407    def __init__(self,
408                 mesh_object,
409                 sigma=None,
410                 use_tetrahedron_method=False,
411                 direction=None,
412                 xyz_projection=False):
413        Dos.__init__(self,
414                     mesh_object,
415                     sigma=sigma,
416                     use_tetrahedron_method=use_tetrahedron_method)
417        self._eigenvectors = self._mesh_object.eigenvectors
418        self._partial_dos = None
419
420        if xyz_projection:
421            self._eigvecs2 = np.abs(self._eigenvectors) ** 2
422        else:
423            num_atom = self._frequencies.shape[1] // 3
424            i_x = np.arange(num_atom, dtype='int') * 3
425            i_y = np.arange(num_atom, dtype='int') * 3 + 1
426            i_z = np.arange(num_atom, dtype='int') * 3 + 2
427            if direction is None:
428                self._eigvecs2 = np.abs(self._eigenvectors[:, i_x, :]) ** 2
429                self._eigvecs2 += np.abs(self._eigenvectors[:, i_y, :]) ** 2
430                self._eigvecs2 += np.abs(self._eigenvectors[:, i_z, :]) ** 2
431            else:
432                d = np.array(direction, dtype='double')
433                d /= np.linalg.norm(direction)
434                proj_eigvecs = self._eigenvectors[:, i_x, :] * d[0]
435                proj_eigvecs += self._eigenvectors[:, i_y, :] * d[1]
436                proj_eigvecs += self._eigenvectors[:, i_z, :] * d[2]
437                self._eigvecs2 = np.abs(proj_eigvecs) ** 2
438
439        self._openmp_thm = True
440
441    @property
442    def partial_dos(self):
443        return self._partial_dos
444
445    @property
446    def projected_dos(self):
447        return self._partial_dos
448
449    def run(self):
450        if self._tetrahedron_mesh is None:
451            self._run_smearing_method()
452        else:
453            if self._openmp_thm:
454                self._run_tetrahedron_method_dos()
455            else:
456                self._run_tetrahedron_method()
457
458    def get_partial_dos(self):
459        """
460        frequency_points: Sampling frequencies
461        partial_dos: [atom_index, frequency_points_index]
462        """
463        return self._frequency_points, self._partial_dos
464
465    def plot(self,
466             ax,
467             indices=None,
468             legend=None,
469             xlabel=None,
470             ylabel=None,
471             draw_grid=True,
472             flip_xy=False):
473
474        if flip_xy:
475            _xlabel = 'Partial density of states'
476            _ylabel = 'Frequency'
477        else:
478            _xlabel = 'Frequency'
479            _ylabel = 'Partial density of states'
480
481        if xlabel is not None:
482            _xlabel = xlabel
483        if ylabel is not None:
484            _ylabel = ylabel
485
486        plot_partial_dos(ax,
487                         self._frequency_points,
488                         self._partial_dos,
489                         indices=indices,
490                         legend=legend,
491                         xlabel=_xlabel,
492                         ylabel=_ylabel,
493                         draw_grid=draw_grid,
494                         flip_xy=flip_xy)
495
496    def write(self, filename="partial_dos.dat"):
497        if self._tetrahedron_mesh is None:
498            comment = "Sigma = %f" % self._sigma
499        else:
500            comment = "Tetrahedron method"
501
502        write_partial_dos(self._frequency_points,
503                          self._partial_dos,
504                          comment=comment,
505                          filename=filename)
506
507    def _run_smearing_method(self):
508        num_pdos = self._eigvecs2.shape[1]
509        num_freqs = len(self._frequency_points)
510        self._partial_dos = np.zeros((num_pdos, num_freqs), dtype='double')
511        weights = self._weights / float(np.sum(self._weights))
512        for i, freq in enumerate(self._frequency_points):
513            amplitudes = self._smearing_function.calc(self._frequencies - freq)
514            for j in range(self._partial_dos.shape[0]):
515                self._partial_dos[j, i] = np.dot(
516                    weights, self._eigvecs2[:, j, :] * amplitudes).sum()
517
518    def _run_tetrahedron_method(self):
519        num_pdos = self._eigvecs2.shape[1]
520        num_freqs = len(self._frequency_points)
521        self._partial_dos = np.zeros((num_pdos, num_freqs), dtype='double')
522        thm = self._tetrahedron_mesh
523        thm.set(value='I', frequency_points=self._frequency_points)
524        for i, iw in enumerate(thm):
525            w = self._weights[i]
526            self._partial_dos += np.dot(iw * w, self._eigvecs2[i].T).T
527
528    def _run_tetrahedron_method_dos(self):
529        mesh_numbers = self._mesh_object.mesh_numbers
530        cell = self._mesh_object.dynamical_matrix.primitive
531        reciprocal_lattice = np.linalg.inv(cell.get_cell())
532        tm = TetrahedronMethod(reciprocal_lattice, mesh=mesh_numbers)
533        pdos = run_tetrahedron_method_dos(
534            mesh_numbers,
535            self._frequency_points,
536            self._frequencies,
537            self._mesh_object.grid_address,
538            self._mesh_object.grid_mapping_table,
539            tm.get_tetrahedra(),
540            coef=self._eigvecs2)
541        self._partial_dos = pdos.T
542