1"""
2This module finds diffusion paths through a structure based on a given
3potential field.
4
5If you use PathFinder algorithm for your research, please consider citing the
6following work::
7
8    Ziqin Rong, Daniil Kitchaev, Pieremanuele Canepa, Wenxuan Huang, Gerbrand
9    Ceder, The Journal of Chemical Physics 145 (7), 074112
10"""
11
12import logging
13import math
14from abc import ABCMeta
15import warnings
16
17import numpy as np
18import numpy.linalg as la
19import scipy.signal
20import scipy.stats
21from scipy.interpolate import interp1d
22
23from pymatgen.core.sites import PeriodicSite
24from pymatgen.core.structure import Structure
25from pymatgen.io.vasp.inputs import Poscar
26
27__author__ = "Daniil Kitchaev"
28__version__ = "1.0"
29__maintainer__ = "Daniil Kitchaev, Ziqin Rong"
30__email__ = "dkitch@mit.edu, rongzq08@mit.edu"
31__status__ = "Development"
32__date__ = "March 17, 2015"
33
34logger = logging.getLogger(__name__)
35
36warnings.warn(
37    "This code has been superseded by pymatgen.analysis.neb in the separate add-on package"
38    "pymatgen-diffusion. This module here is retained for backwards compatibility. It will be removed from"
39    "2022.1.1.",
40    FutureWarning,
41)
42
43
44class NEBPathfinder:
45    """
46    General pathfinder for interpolating between two structures, where the
47    interpolating path is calculated with the elastic band method with
48    respect to the given static potential for sites whose indices are given
49    in relax_sites, and is linear otherwise.
50    """
51
52    def __init__(self, start_struct, end_struct, relax_sites, v, n_images=20, mid_struct=None):
53        """
54        Args:
55            start_struct, end_struct: Endpoint structures to interpolate
56            relax_sites: List of site indices whose interpolation paths should
57                be relaxed
58            v: Static potential field to use for the elastic band relaxation
59            n_images: Number of interpolation images to generate
60            mid_struct: (optional) additional structure between the start and end structures to help
61        """
62        self.__s1 = start_struct
63        self.__s2 = end_struct
64        self.__mid = mid_struct
65        self.__relax_sites = relax_sites
66        self.__v = v
67        self.__n_images = n_images
68        self.__images = None
69        self.interpolate()
70
71    def interpolate(self):
72        """
73        Finds a set of n_images from self.s1 to self.s2, where all sites except
74        for the ones given in relax_sites, the interpolation is linear (as in
75        pymatgen.core.structure.interpolate), and for the site indices given
76        in relax_sites, the path is relaxed by the elastic band method within
77        the static potential V.
78
79        If a mid point is defined we will interpolate from s1--> mid -->s2
80        The final number of images will still be n_images.
81        """
82        if self.__mid is not None:
83            # to make arithmatic easier we will do the interpolation in two parts with n images each
84            # then just take every other image at the end, this results in exactly n images
85            images_0 = self.__s1.interpolate(self.__mid, nimages=self.__n_images, interpolate_lattices=False)[:-1]
86            images_1 = self.__mid.interpolate(self.__s2, nimages=self.__n_images, interpolate_lattices=False)
87            images = images_0 + images_1
88            images = images[::2]
89        else:
90            images = self.__s1.interpolate(self.__s2, nimages=self.__n_images, interpolate_lattices=False)
91        for site_i in self.__relax_sites:
92            start_f = images[0].sites[site_i].frac_coords
93            end_f = images[-1].sites[site_i].frac_coords
94
95            path = NEBPathfinder.string_relax(
96                NEBPathfinder.__f2d(start_f, self.__v),
97                NEBPathfinder.__f2d(end_f, self.__v),
98                self.__v,
99                n_images=(self.__n_images + 1),
100                dr=[
101                    self.__s1.lattice.a / self.__v.shape[0],
102                    self.__s1.lattice.b / self.__v.shape[1],
103                    self.__s1.lattice.c / self.__v.shape[2],
104                ],
105            )
106            for image_i, image in enumerate(images):
107                image.translate_sites(
108                    site_i,
109                    NEBPathfinder.__d2f(path[image_i], self.__v) - image.sites[site_i].frac_coords,
110                    frac_coords=True,
111                    to_unit_cell=True,
112                )
113        self.__images = images
114
115    @property
116    def images(self):
117        """
118        Returns a list of structures interpolating between the start and
119        endpoint structures.
120        """
121        return self.__images
122
123    def plot_images(self, outfile):
124        """
125        Generates a POSCAR with the calculated diffusion path with respect to the first endpoint.
126        :param outfile: Output file for the POSCAR
127        """
128        sum_struct = self.__images[0].sites
129        for image in self.__images:
130            for site_i in self.__relax_sites:
131                sum_struct.append(
132                    PeriodicSite(
133                        image.sites[site_i].specie,
134                        image.sites[site_i].frac_coords,
135                        self.__images[0].lattice,
136                        to_unit_cell=True,
137                        coords_are_cartesian=False,
138                    )
139                )
140        sum_struct = Structure.from_sites(sum_struct, validate_proximity=False)
141        p = Poscar(sum_struct)
142        p.write_file(outfile)
143
144    @staticmethod
145    def string_relax(
146        start,
147        end,
148        V,
149        n_images=25,
150        dr=None,
151        h=3.0,
152        k=0.17,
153        min_iter=100,
154        max_iter=10000,
155        max_tol=5e-6,
156    ):
157        """
158        Implements path relaxation via the elastic band method. In general, the
159        method is to define a path by a set of points (images) connected with
160        bands with some elasticity constant k. The images then relax along the
161        forces found in the potential field V, counterbalanced by the elastic
162        response of the elastic band. In general the endpoints of the band can
163        be allowed to relax also to their local minima, but in this calculation
164        they are kept fixed.
165
166        Args:
167            start, end: Endpoints of the path calculation given in discrete
168                coordinates with respect to the grid in V
169            V: potential field through which to calculate the path
170            n_images: number of images used to define the path. In general
171                anywhere from 20 to 40 seems to be good.
172            dr: Conversion ratio from discrete coordinates to real coordinates
173                for each of the three coordinate vectors
174            h: Step size for the relaxation. h = 0.1 works reliably, but is
175                slow. h=10 diverges with large gradients but for the types of
176                gradients seen in CHGCARs, works pretty reliably
177            k: Elastic constant for the band (in real units, not discrete)
178            min_iter, max_iter: Number of optimization steps the string will
179                take before exiting (even if unconverged)
180            max_tol: Convergence threshold such that if the string moves by
181                less than max_tol in a step, and at least min_iter steps have
182                passed, the algorithm will terminate. Depends strongly on the
183                size of the gradients in V, but 5e-6 works reasonably well for
184                CHGCARs.
185        """
186        #
187        # This code is based on the MATLAB example provided by
188        # Prof. Eric Vanden-Eijnden of NYU
189        # (http://www.cims.nyu.edu/~eve2/main.htm)
190        #
191
192        # logger.debug("Getting path from {} to {} (coords wrt V grid)".format(start, end))
193
194        # Set parameters
195        if not dr:
196            dr = np.array([1.0 / V.shape[0], 1.0 / V.shape[1], 1.0 / V.shape[2]])
197        else:
198            dr = np.array(dr, dtype=float)
199        keff = k * dr * n_images
200        h0 = h
201
202        # Initialize string
203        g1 = np.linspace(0, 1, n_images)
204        s0 = start
205        s1 = end
206        s = np.array([g * (s1 - s0) for g in g1]) + s0
207        ds = s - np.roll(s, 1, axis=0)
208        ds[0] = ds[0] - ds[0]
209        ls = np.cumsum(la.norm(ds, axis=1))
210        ls = ls / ls[-1]
211        fi = interp1d(ls, s, axis=0)
212        s = fi(g1)
213
214        # Evaluate initial distances (for elastic equilibrium)
215        ds0_plus = s - np.roll(s, 1, axis=0)
216        ds0_minus = s - np.roll(s, -1, axis=0)
217        ds0_plus[0] = ds0_plus[0] - ds0_plus[0]
218        ds0_minus[-1] = ds0_minus[-1] - ds0_minus[-1]
219
220        # Evaluate potential gradient outside the loop, as potential does not
221        # change per step in this approximation.
222        dV = np.gradient(V)
223
224        # Evolve string
225        for step in range(0, max_iter):
226            if step > min_iter:
227                # Gradually decay step size to prevent oscillations
228                h = h0 * np.exp(-2.0 * (step - min_iter) / max_iter)
229            else:
230                h = h0
231            # Calculate forces acting on string
232            d = V.shape
233            s0 = s
234            edV = np.array(
235                [
236                    [
237                        dV[0][int(pt[0]) % d[0]][int(pt[1]) % d[1]][int(pt[2]) % d[2]] / dr[0],
238                        dV[1][int(pt[0]) % d[0]][int(pt[1]) % d[1]][int(pt[2]) % d[2]] / dr[0],
239                        dV[2][int(pt[0]) % d[0]][int(pt[1]) % d[1]][int(pt[2]) % d[2]] / dr[0],
240                    ]
241                    for pt in s
242                ]
243            )
244            # if(step % 100 == 0):
245            #    logger.debug(edV)
246
247            # Update according to force due to potential and string elasticity
248            ds_plus = s - np.roll(s, 1, axis=0)
249            ds_minus = s - np.roll(s, -1, axis=0)
250            ds_plus[0] = ds_plus[0] - ds_plus[0]
251            ds_minus[-1] = ds_minus[-1] - ds_minus[-1]
252            Fpot = edV
253            Fel = keff * (la.norm(ds_plus) - la.norm(ds0_plus)) * (ds_plus / la.norm(ds_plus))
254            Fel += keff * (la.norm(ds_minus) - la.norm(ds0_minus)) * (ds_minus / la.norm(ds_minus))
255            s -= h * (Fpot + Fel)
256
257            # Fix endpoints
258            s[0] = s0[0]
259            s[-1] = s0[-1]
260
261            # Reparametrize string
262            ds = s - np.roll(s, 1, axis=0)
263            ds[0] = ds[0] - ds[0]
264            ls = np.cumsum(la.norm(ds, axis=1))
265            ls = ls / ls[-1]
266            fi = interp1d(ls, s, axis=0)
267            s = fi(g1)
268
269            tol = la.norm((s - s0) * dr) / n_images / h
270
271            if tol > 1e10:
272                raise ValueError("Pathfinding failed, path diverged! Consider reducing h to " "avoid divergence.")
273
274            if step > min_iter and tol < max_tol:
275                logger.debug("Converged at step {}".format(step))
276                break
277
278            if step % 100 == 0:
279                logger.debug("Step {} - ds = {}".format(step, tol))
280        return s
281
282    @staticmethod
283    def __f2d(frac_coords, v):
284        """
285        Converts fractional coordinates to discrete coordinates with respect to
286        the grid size of v
287        """
288        # frac_coords = frac_coords % 1
289        return np.array(
290            [
291                int(frac_coords[0] * v.shape[0]),
292                int(frac_coords[1] * v.shape[1]),
293                int(frac_coords[2] * v.shape[2]),
294            ]
295        )
296
297    @staticmethod
298    def __d2f(disc_coords, v):
299        """
300        Converts a point given in discrete coordinates withe respect to the
301        grid in v to fractional coordinates.
302        """
303        return np.array(
304            [
305                disc_coords[0] / v.shape[0],
306                disc_coords[1] / v.shape[1],
307                disc_coords[2] / v.shape[2],
308            ]
309        )
310
311
312class StaticPotential(metaclass=ABCMeta):
313    """
314    Defines a general static potential for diffusion calculations. Implements
315    grid-rescaling and smearing for the potential grid. Also provides a
316    function to normalize the potential from 0 to 1 (recommended).
317    """
318
319    def __init__(self, struct, pot):
320        """
321        :param struct: atomic structure of the potential
322        :param pot: volumentric data to be used as a potential
323        """
324        self.__v = pot
325        self.__s = struct
326
327    def get_v(self):
328        """
329        Returns the potential
330        """
331        return self.__v
332
333    def normalize(self):
334        """
335        Sets the potential range 0 to 1.
336        """
337        self.__v = self.__v - np.amin(self.__v)
338        self.__v = self.__v / np.amax(self.__v)
339
340    def rescale_field(self, new_dim):
341        """
342        Changes the discretization of the potential field by linear
343        interpolation. This is necessary if the potential field
344        obtained from DFT is strangely skewed, or is too fine or coarse. Obeys
345        periodic boundary conditions at the edges of
346        the cell. Alternatively useful for mixing potentials that originally
347        are on different grids.
348
349        :param new_dim: tuple giving the numpy shape of the new grid
350        """
351        v_dim = self.__v.shape
352        padded_v = np.lib.pad(self.__v, ((0, 1), (0, 1), (0, 1)), mode="wrap")
353        ogrid_list = np.array([list(c) for c in list(np.ndindex(v_dim[0] + 1, v_dim[1] + 1, v_dim[2] + 1))])
354        v_ogrid = padded_v.reshape(((v_dim[0] + 1) * (v_dim[1] + 1) * (v_dim[2] + 1), -1))
355        ngrid_a, ngrid_b, ngrid_c = np.mgrid[
356            0 : v_dim[0] : v_dim[0] / new_dim[0],
357            0 : v_dim[1] : v_dim[1] / new_dim[1],
358            0 : v_dim[2] : v_dim[2] / new_dim[2],
359        ]
360
361        v_ngrid = scipy.interpolate.griddata(ogrid_list, v_ogrid, (ngrid_a, ngrid_b, ngrid_c), method="linear").reshape(
362            (new_dim[0], new_dim[1], new_dim[2])
363        )
364        self.__v = v_ngrid
365
366    def gaussian_smear(self, r):
367        """
368        Applies an isotropic Gaussian smear of width (standard deviation) r to
369        the potential field. This is necessary to avoid finding paths through
370        narrow minima or nodes that may exist in the field (although any
371        potential or charge distribution generated from GGA should be
372        relatively smooth anyway). The smearing obeys periodic
373        boundary conditions at the edges of the cell.
374
375        :param r - Smearing width in cartesian coordinates, in the same units
376            as the structure lattice vectors
377        """
378        # Since scaling factor in fractional coords is not isotropic, have to
379        # have different radii in 3 directions
380        a_lat = self.__s.lattice.a
381        b_lat = self.__s.lattice.b
382        c_lat = self.__s.lattice.c
383
384        # Conversion factors for discretization of v
385        v_dim = self.__v.shape
386        r_frac = (r / a_lat, r / b_lat, r / c_lat)
387        r_disc = (
388            int(math.ceil(r_frac[0] * v_dim[0])),
389            int(math.ceil(r_frac[1] * v_dim[1])),
390            int(math.ceil(r_frac[2] * v_dim[2])),
391        )
392
393        # Apply smearing
394        # Gaussian filter
395        gauss_dist = np.zeros((r_disc[0] * 4 + 1, r_disc[1] * 4 + 1, r_disc[2] * 4 + 1))
396        for g_a in np.arange(-2.0 * r_disc[0], 2.0 * r_disc[0] + 1, 1.0):
397            for g_b in np.arange(-2.0 * r_disc[1], 2.0 * r_disc[1] + 1, 1.0):
398                for g_c in np.arange(-2.0 * r_disc[2], 2.0 * r_disc[2] + 1, 1.0):
399                    g = np.array([g_a / v_dim[0], g_b / v_dim[1], g_c / v_dim[2]]).T
400                    gauss_dist[int(g_a + r_disc[0])][int(g_b + r_disc[1])][int(g_c + r_disc[2])] = (
401                        la.norm(np.dot(self.__s.lattice.matrix, g)) / r
402                    )
403        gauss = scipy.stats.norm.pdf(gauss_dist)
404        gauss = gauss / np.sum(gauss, dtype=float)
405        padded_v = np.pad(
406            self.__v,
407            ((r_disc[0], r_disc[0]), (r_disc[1], r_disc[1]), (r_disc[2], r_disc[2])),
408            mode="wrap",
409        )
410        smeared_v = scipy.signal.convolve(padded_v, gauss, mode="valid")
411        self.__v = smeared_v
412
413
414class ChgcarPotential(StaticPotential):
415    """
416    Implements a potential field based on the charge density output from VASP.
417    """
418
419    def __init__(self, chgcar, smear=False, normalize=True):
420        """
421        :param chgcar: Chgcar object based on a VASP run of the structure of
422            interest (Chgcar.from_file("CHGCAR"))
423        :param smear: Whether or not to apply a Gaussian smearing to the
424            potential
425        :param normalize: Whether or not to normalize the potential to range
426            from 0 to 1
427        """
428        v = chgcar.data["total"]
429        v = v / (v.shape[0] * v.shape[1] * v.shape[2])
430        StaticPotential.__init__(self, chgcar.structure, v)
431        if smear:
432            self.gaussian_smear(2.0)
433        if normalize:
434            self.normalize()
435
436
437class FreeVolumePotential(StaticPotential):
438    """
439    Implements a potential field based on geometric distances from atoms in the
440    structure - basically, the potential
441    is lower at points farther away from any atoms in the structure.
442    """
443
444    def __init__(self, struct, dim, smear=False, normalize=True):
445        """
446        :param struct: Unit cell on which to base the potential
447        :param dim: Grid size for the potential
448        :param smear: Whether or not to apply a Gaussian smearing to the
449            potential
450        :param normalize: Whether or not to normalize the potential to range
451            from 0 to 1
452        """
453        self.__s = struct
454        v = FreeVolumePotential.__add_gaussians(struct, dim)
455        StaticPotential.__init__(self, struct, v)
456        if smear:
457            self.gaussian_smear(2.0)
458        if normalize:
459            self.normalize()
460
461    @staticmethod
462    def __add_gaussians(s, dim, r=1.5):
463        gauss_dist = np.zeros(dim)
464        for a_d in np.arange(0.0, dim[0], 1.0):
465            for b_d in np.arange(0.0, dim[1], 1.0):
466                for c_d in np.arange(0.0, dim[2], 1.0):
467                    coords_f = np.array([a_d / dim[0], b_d / dim[1], c_d / dim[2]])
468                    d_f = sorted(s.get_sites_in_sphere(coords_f, s.lattice.a), key=lambda x: x[1])[0][1]
469                    # logger.debug(d_f)
470                    gauss_dist[int(a_d)][int(b_d)][int(c_d)] = d_f / r
471        v = scipy.stats.norm.pdf(gauss_dist)
472        return v
473
474
475class MixedPotential(StaticPotential):
476    """
477    Implements a potential that is a weighted sum of some other potentials
478    """
479
480    def __init__(self, potentials, coefficients, smear=False, normalize=True):
481        """
482        Args:
483            potentials: List of objects extending the StaticPotential superclass
484            coefficients: Mixing weights for the elements of the potentials list
485            smear: Whether or not to apply a Gaussian smearing to the potential
486            normalize: Whether or not to normalize the potential to range from
487                0 to 1
488        """
489        v = potentials[0].get_v() * coefficients[0]
490        s = potentials[0].__s
491        for i in range(1, len(potentials)):
492            v += potentials[i].get_v() * coefficients[i]
493        StaticPotential.__init__(self, s, v)
494        if smear:
495            self.gaussian_smear(2.0)
496        if normalize:
497            self.normalize()
498