1import numpy as np
2from collections import namedtuple
3from ase.geometry import find_mic
4
5
6def fit_raw(energies, forces, positions, cell=None, pbc=None):
7    """Calculates parameters for fitting images to a band, as for
8    a NEB plot."""
9    energies = np.array(energies) - energies[0]
10    n_images = len(energies)
11    fit_energies = np.empty((n_images - 1) * 20 + 1)
12    fit_path = np.empty((n_images - 1) * 20 + 1)
13
14    path = [0]
15    for i in range(n_images - 1):
16        dR = positions[i + 1] - positions[i]
17        if cell is not None and pbc is not None:
18            dR, _ = find_mic(dR, cell, pbc)
19        path.append(path[i] + np.sqrt((dR**2).sum()))
20
21    lines = []  # tangent lines
22    lastslope = None
23    for i in range(n_images):
24        if i == 0:
25            direction = positions[i + 1] - positions[i]
26            dpath = 0.5 * path[1]
27        elif i == n_images - 1:
28            direction = positions[-1] - positions[-2]
29            dpath = 0.5 * (path[-1] - path[-2])
30        else:
31            direction = positions[i + 1] - positions[i - 1]
32            dpath = 0.25 * (path[i + 1] - path[i - 1])
33
34        direction /= np.linalg.norm(direction)
35        slope = -(forces[i] * direction).sum()
36        x = np.linspace(path[i] - dpath, path[i] + dpath, 3)
37        y = energies[i] + slope * (x - path[i])
38        lines.append((x, y))
39
40        if i > 0:
41            s0 = path[i - 1]
42            s1 = path[i]
43            x = np.linspace(s0, s1, 20, endpoint=False)
44            c = np.linalg.solve(np.array([(1, s0, s0**2, s0**3),
45                                          (1, s1, s1**2, s1**3),
46                                          (0, 1, 2 * s0, 3 * s0**2),
47                                          (0, 1, 2 * s1, 3 * s1**2)]),
48                                np.array([energies[i - 1], energies[i],
49                                          lastslope, slope]))
50            y = c[0] + x * (c[1] + x * (c[2] + x * c[3]))
51            fit_path[(i - 1) * 20:i * 20] = x
52            fit_energies[(i - 1) * 20:i * 20] = y
53
54        lastslope = slope
55
56    fit_path[-1] = path[-1]
57    fit_energies[-1] = energies[-1]
58    return ForceFit(path, energies, fit_path, fit_energies, lines)
59
60
61class ForceFit(namedtuple('ForceFit', ['path', 'energies', 'fit_path',
62                                       'fit_energies', 'lines'])):
63    """Data container to hold fitting parameters for force curves."""
64
65    def plot(self, ax=None):
66        import matplotlib.pyplot as plt
67        if ax is None:
68            ax = plt.gca()
69
70        ax.plot(self.path, self.energies, 'o')
71        for x, y in self.lines:
72            ax.plot(x, y, '-g')
73        ax.plot(self.fit_path, self.fit_energies, 'k-')
74        ax.set_xlabel(r'path [Å]')
75        ax.set_ylabel('energy [eV]')
76        Ef = max(self.energies) - self.energies[0]
77        Er = max(self.energies) - self.energies[-1]
78        dE = self.energies[-1] - self.energies[0]
79        ax.set_title(r'$E_\mathrm{{f}} \approx$ {:.3f} eV; '
80                     r'$E_\mathrm{{r}} \approx$ {:.3f} eV; '
81                     r'$\Delta E$ = {:.3f} eV'.format(Ef, Er, dE))
82        return ax
83
84
85def fit_images(images):
86    """Fits a series of images with a smoothed line for producing a standard
87    NEB plot. Returns a `ForceFit` data structure; the plot can be produced
88    by calling the `plot` method of `ForceFit`."""
89    R = [atoms.positions for atoms in images]
90    E = [atoms.get_potential_energy() for atoms in images]
91    F = [atoms.get_forces() for atoms in images]  # XXX force consistent???
92    A = images[0].cell
93    pbc = images[0].pbc
94    return fit_raw(E, F, R, A, pbc)
95
96
97def force_curve(images, ax=None):
98    """Plot energies and forces as a function of accumulated displacements.
99
100    This is for testing whether a calculator's forces are consistent with
101    the energies on a set of geometries where energies and forces are
102    available."""
103
104    if ax is None:
105        import matplotlib.pyplot as plt
106        ax = plt.gca()
107
108    nim = len(images)
109
110    accumulated_distances = []
111    accumulated_distance = 0.0
112
113    # XXX force_consistent=True will work with some calculators,
114    # but won't work if images were loaded from a trajectory.
115    energies = [atoms.get_potential_energy()
116                for atoms in images]
117
118    for i in range(nim):
119        atoms = images[i]
120        f_ac = atoms.get_forces()
121
122        if i < nim - 1:
123            rightpos = images[i + 1].positions
124        else:
125            rightpos = atoms.positions
126
127        if i > 0:
128            leftpos = images[i - 1].positions
129        else:
130            leftpos = atoms.positions
131
132        disp_ac, _ = find_mic(rightpos - leftpos, cell=atoms.cell,
133                              pbc=atoms.pbc)
134
135        def total_displacement(disp):
136            disp_a = (disp**2).sum(axis=1)**.5
137            return sum(disp_a)
138
139        dE_fdotr = -0.5 * np.vdot(f_ac.ravel(), disp_ac.ravel())
140
141        linescale = 0.45
142
143        disp = 0.5 * total_displacement(disp_ac)
144
145        if i == 0 or i == nim - 1:
146            disp *= 2
147            dE_fdotr *= 2
148
149        x1 = accumulated_distance - disp * linescale
150        x2 = accumulated_distance + disp * linescale
151        y1 = energies[i] - dE_fdotr * linescale
152        y2 = energies[i] + dE_fdotr * linescale
153
154        ax.plot([x1, x2], [y1, y2], 'b-')
155        ax.plot(accumulated_distance, energies[i], 'bo')
156        ax.set_ylabel('Energy [eV]')
157        ax.set_xlabel('Accumulative distance [Å]')
158        accumulated_distances.append(accumulated_distance)
159        accumulated_distance += total_displacement(rightpos - atoms.positions)
160
161    ax.plot(accumulated_distances, energies, ':', zorder=-1, color='k')
162    return ax
163
164
165def plotfromfile(*fnames):
166    from ase.io import read
167    nplots = len(fnames)
168
169    for i, fname in enumerate(fnames):
170        images = read(fname, ':')
171        import matplotlib.pyplot as plt
172        plt.subplot(nplots, 1, 1 + i)
173        force_curve(images)
174    plt.show()
175
176
177if __name__ == '__main__':
178    import sys
179    fnames = sys.argv[1:]
180    plotfromfile(*fnames)
181