1# coding: utf-8
2# Copyright (c) Pymatgen Development Team.
3# Distributed under the terms of the MIT License.
4
5"""
6This module implements core classes for calculation of diffraction patterns.
7"""
8
9import abc
10import collections
11
12import numpy as np
13
14from pymatgen.core.spectrum import Spectrum
15from pymatgen.util.plotting import add_fig_kwargs
16
17
18class DiffractionPattern(Spectrum):
19    """
20    A representation of a diffraction pattern
21    """
22
23    XLABEL = "$2\\Theta$"
24    YLABEL = "Intensity"
25
26    def __init__(self, x, y, hkls, d_hkls):
27        """
28        Args:
29            x: Two theta angles.
30            y: Intensities
31            hkls: [{"hkl": (h, k, l), "multiplicity": mult}],
32                where {"hkl": (h, k, l), "multiplicity": mult}
33                is a dict of Miller
34                indices for all diffracted lattice facets contributing to each
35                intensity.
36            d_hkls: List of interplanar spacings.
37        """
38        super().__init__(x, y, hkls, d_hkls)
39        self.hkls = hkls
40        self.d_hkls = d_hkls
41
42
43class AbstractDiffractionPatternCalculator(abc.ABC):
44    """
45    Abstract base class for computing the diffraction pattern of a crystal.
46    """
47
48    # Tolerance in which to treat two peaks as having the same two theta.
49    TWO_THETA_TOL = 1e-5
50
51    # Tolerance in which to treat a peak as effectively 0 if the scaled
52    # intensity is less than this number. Since the max intensity is 100,
53    # this means the peak must be less than 1e-5 of the peak intensity to be
54    # considered as zero. This deals with numerical issues where systematic
55    # absences do not cancel exactly to zero.
56    SCALED_INTENSITY_TOL = 1e-3
57
58    @abc.abstractmethod
59    def get_pattern(self, structure, scaled=True, two_theta_range=(0, 90)):
60        """
61        Calculates the diffraction pattern for a structure.
62
63        Args:
64            structure (Structure): Input structure
65            scaled (bool): Whether to return scaled intensities. The maximum
66                peak is set to a value of 100. Defaults to True. Use False if
67                you need the absolute values to combine XRD plots.
68            two_theta_range ([float of length 2]): Tuple for range of
69                two_thetas to calculate in degrees. Defaults to (0, 90). Set to
70                None if you want all diffracted beams within the limiting
71                sphere of radius 2 / wavelength.
72
73        Returns:
74            (DiffractionPattern)
75        """
76        pass
77
78    def get_plot(
79        self,
80        structure,
81        two_theta_range=(0, 90),
82        annotate_peaks="compact",
83        ax=None,
84        with_labels=True,
85        fontsize=16,
86    ):
87        """
88        Returns the diffraction plot as a matplotlib.pyplot.
89
90        Args:
91            structure: Input structure
92            two_theta_range ([float of length 2]): Tuple for range of
93                two_thetas to calculate in degrees. Defaults to (0, 90). Set to
94                None if you want all diffracted beams within the limiting
95                sphere of radius 2 / wavelength.
96            annotate_peaks (str or None): Whether and how to annotate the peaks
97                with hkl indices.  Default is 'compact', i.e. show short
98                version (oriented vertically), e.g. 100.  If 'full', show
99                long version, e.g. (1, 0, 0).  If None, do not show anything.
100            ax: matplotlib :class:`Axes` or None if a new figure should be
101                created.
102            with_labels: True to add xlabels and ylabels to the plot.
103            fontsize: (int) fontsize for peak labels.
104
105        Returns:
106            (matplotlib.pyplot)
107        """
108        if ax is None:
109            from pymatgen.util.plotting import pretty_plot
110
111            plt = pretty_plot(16, 10)
112            ax = plt.gca()
113        else:
114            # This to maintain the type of the return value.
115            import matplotlib.pyplot as plt
116
117        xrd = self.get_pattern(structure, two_theta_range=two_theta_range)
118        imax = max(xrd.y)
119
120        for two_theta, i, hkls, d_hkl in zip(xrd.x, xrd.y, xrd.hkls, xrd.d_hkls):
121            if two_theta_range[0] <= two_theta <= two_theta_range[1]:
122                hkl_tuples = [hkl["hkl"] for hkl in hkls]
123                label = ", ".join([str(hkl_tuple) for hkl_tuple in hkl_tuples])  # 'full' label
124                ax.plot([two_theta, two_theta], [0, i], color="k", linewidth=3, label=label)
125
126                if annotate_peaks == "full":
127                    ax.annotate(
128                        label,
129                        xy=[two_theta, i],
130                        xytext=[two_theta, i],
131                        fontsize=fontsize,
132                    )
133                elif annotate_peaks == "compact":
134                    if all(all(i < 10 for i in hkl_tuple) for hkl_tuple in hkl_tuples):
135                        label = ",".join(["".join([str(i) for i in hkl_tuple]) for hkl_tuple in hkl_tuples])
136                        # 'compact' label.  Would be unclear for indices >= 10
137                        # It would have more than 3 figures, e.g. 1031
138
139                    if i / imax > 0.5:  # Big peak: annotation on the side
140                        xytext = [-fontsize / 4, 0]
141                        ha = "right"
142                        va = "top"
143                    else:  # Small peak: annotation on top
144                        xytext = [0, 10]
145                        ha = "center"
146                        va = "bottom"
147
148                    ax.annotate(
149                        label,
150                        xy=[two_theta, i],
151                        xytext=xytext,
152                        textcoords="offset points",
153                        va=va,
154                        ha=ha,
155                        rotation=90,
156                        fontsize=fontsize,
157                    )
158
159        if with_labels:
160            ax.set_xlabel(r"$2\theta$ ($^\circ$)")
161            ax.set_ylabel("Intensities (scaled)")
162
163        if hasattr(ax, "tight_layout"):
164            ax.tight_layout()
165
166        return plt
167
168    def show_plot(self, structure, **kwargs):
169        """
170        Shows the diffraction plot.
171
172        Args:
173            structure (Structure): Input structure
174            two_theta_range ([float of length 2]): Tuple for range of
175                two_thetas to calculate in degrees. Defaults to (0, 90). Set to
176                None if you want all diffracted beams within the limiting
177                sphere of radius 2 / wavelength.
178            annotate_peaks (str or None): Whether and how to annotate the peaks
179                with hkl indices.  Default is 'compact', i.e. show short
180                version (oriented vertically), e.g. 100.  If 'full', show
181                long version, e.g. (1, 0, 0).  If None, do not show anything.
182        """
183        self.get_plot(structure, **kwargs).show()
184
185    @add_fig_kwargs
186    def plot_structures(self, structures, fontsize=6, **kwargs):
187        """
188        Plot diffraction patterns for multiple structures on the same figure.
189
190        Args:
191            structures (Structure): List of structures
192            two_theta_range ([float of length 2]): Tuple for range of
193                two_thetas to calculate in degrees. Defaults to (0, 90). Set to
194                None if you want all diffracted beams within the limiting
195                sphere of radius 2 / wavelength.
196            annotate_peaks (str or None): Whether and how to annotate the peaks
197                with hkl indices.  Default is 'compact', i.e. show short
198                version (oriented vertically), e.g. 100.  If 'full', show
199                long version, e.g. (1, 0, 0).  If None, do not show anything.
200            fontsize: (int) fontsize for peak labels.
201        """
202        import matplotlib.pyplot as plt
203
204        nrows = len(structures)
205        fig, axes = plt.subplots(nrows=nrows, ncols=1, sharex=True, squeeze=False)
206
207        for i, (ax, structure) in enumerate(zip(axes.ravel(), structures)):
208            self.get_plot(structure, fontsize=fontsize, ax=ax, with_labels=i == nrows - 1, **kwargs)
209            spg_symbol, spg_number = structure.get_space_group_info()
210            ax.set_title("{} {} ({}) ".format(structure.formula, spg_symbol, spg_number))
211
212        return fig
213
214
215def get_unique_families(hkls):
216    """
217    Returns unique families of Miller indices. Families must be permutations
218    of each other.
219
220    Args:
221        hkls ([h, k, l]): List of Miller indices.
222
223    Returns:
224        {hkl: multiplicity}: A dict with unique hkl and multiplicity.
225    """
226
227    # TODO: Definitely can be sped up.
228    def is_perm(hkl1, hkl2):
229        h1 = np.abs(hkl1)
230        h2 = np.abs(hkl2)
231        return all(i == j for i, j in zip(sorted(h1), sorted(h2)))
232
233    unique = collections.defaultdict(list)
234    for hkl1 in hkls:
235        found = False
236        for hkl2, v2 in unique.items():
237            if is_perm(hkl1, hkl2):
238                found = True
239                v2.append(hkl1)
240                break
241        if not found:
242            unique[hkl1].append(hkl1)
243
244    pretty_unique = {}
245    for k, v in unique.items():
246        pretty_unique[sorted(v)[-1]] = len(v)
247
248    return pretty_unique
249