<commit_before>import numpy as np
import matplotlib.pyplot as pl

from empca import *
from rotate_factor import ortho_rotation

class FAModel(object):
    """
    Factor Analysis deconvolution, using online learning.
    """
    def __init__(self,data,obsvar,latent_dim,
                 max_iter=100000,compute_total_nll=True,
                 learning_init_relsizes=[0.01,0.01,0.01],
                 Ninit_est=10,check_rate=100,Neff=10,
                 decay_t0_factor=1.,decay_pow=0.5):
        """

        """
        assert data.shape==obsvar.shape, 'Data and obs. ' + \
            'variance have different shapes'

        self.N = data.shape[0]
        self.D = data.shape[1]
        self.M = latent_dim
        self.avg_factor = np.exp(-1./Neff)
        self.jitter = np.zeros(self.D)
        self.decay_t0 = decay_t0_factor * self.N
        self.decay_pow = decay_pow
        self.running_nll = np.ones(max_iter/check_rate) * -np.Inf

        # Suffling the data
        ind = np.random.permutation(self.N)
        self.data = data[ind,:]
        self.obsvar = obsvar[ind,:]

        self.initialize()
        self.initial_mean   = self.mean.copy()
        self.initial_jitter = self.jitter.copy()
        self.initial_lambda = self.lam.copy()
        if compute_total_nll:
            self.initial_nll = self.total_negative_log_likelihood()
        self.calc_init_rates(Ninit_est,learning_init_relsizes)
        self.run_inference(max_iter,check_rate)

    def initialize(self):

        # inverse variance weighted mean
        if np.sum(self.obsvar)!=0.0:
            self.mean = np.sum(self.data / self.obsvar,axis=0) / \
                np.sum(1.0 / self.obsvar,axis=0)
        else:
            self.mean = np.mean(self.data,axis=0)

        # use EM PCA to initialize factor loadings
        if self.M==0:
            self.lam = np.zeros(1)
        else:
            empca = EMPCA(self.data.T,self.M)
            R = ortho_rotation(empca.lam,method='varimax')
            self.lam = np.dot(empca.lam,R)

    def run_inference(self,max_iter,check_rate):

        j = 0
        for ii in range(max_iter):

            # shuffle
            if (ii%self.N)==0:
                ind = np.random.permutation(self.N)
                self.data = self.data[ind,:]
                self.obsvar = self.obsvar[ind,:]

            # specify datum
            i = np.mod(ii,self.data.shape[0])
            self.datum = self.data[i,:]
            self.variance = self.obsvar[i,:]
            self.do_precalcs()

            # Neg. log likelihood
            if (ii%check_rate)==0:
                j = self.estimate_nll(j)
                if (j%10)==0: print ii,j,self.running_nll[j]
                # convergence test...

            # not converged, make a step
            self.calc_rates(j)
            self.make_gradient_step()

    def estimate_nll(self,j):
        est_nll = self.single_negative_log_likelihood()
        if j==0:
            self.running_nll[j] = est_nll
        else:
            self.running_nll[j] = self.running_nll[j-1] * self.avg_factor + \
                (1. - self.avg_factor) * est_nll
        return j+1

    def do_precalcs(self):
        """
        Calculate matrices used repeatedly below
        """
        self.invert_cov()
        self.dmm = self.datum - self.mean
        self.dmmi = np.dot(self.dmm,self.inv_cov)

    def invert_cov(self):
        """
        Make inverse covariance, using the inversion lemma.
        """
        if self.M==0:
            self.inv_cov = np.diag(1.0 / (self.jitter + self.variance))
        else:
            # repeated in lemma
            lamT = lam.T
            if np.sum(self.jitter+self.variance)!=0.0:
                psiI = np.diag(1.0 / (self.jitter + self.variance))
                psiIlam = np.dot(psiI,self.lam)

                # the lemma
                bar = np.dot(lamT,psiI)
                foo = np.linalg.inv(np.eye(self.M) + np.dot(lamT,psiIlam))
                self.inv_cov = psiI - np.dot(psiIlam,np.dot(foo,bar))
            else:
                # fix to lemma-like version
                self.inv_cov = np.linalg.inv(np.dot(self.lam,lamT))

    def make_gradient_step(self):

        self.mean -= self.mean_rate * self.mean_gradients()
        self.jitter -= self.jitter_rate * self.jitter_gradients()
        self.lam -= self.lambda_rate * self.lambda_gradients()

        ind = (self.jitter < 0.0)
        self.jitter[ind] = 0.0

    def calc_rates(self,t):

        t0 = self.decay_t0 * self.N
        p  = self.decay_pow

        self.mean_rate *= (t0 / (t0+t)) ** p
        self.jitter_rate *= (t0 / (t0+t)) ** p
        self.lambda_rate *= (t0 / (t0+t)) ** p

    def calc_init_rates(self,Ninit_est,learning_init_relsizes):

        self.mean_rate = np.Inf
        self.jitter_rate = np.Inf
        self.lambda_rate = np.Inf

        mf = learning_init_relsizes[0]
        jf = learning_init_relsizes[1]
        lf = learning_init_relsizes[2]

        ind = np.random.permutation(self.N)
        for i in range(Ninit_est):
            self.datum = self.data[ind[i],:]
            self.variance = self.obsvar[ind[i],:]
            self.do_precalcs()

            self.mean_rate = np.minimum(self.mean_rate,
                                        np.min(mf/np.abs(self.mean_gradients())))
            self.jitter_rate = np.minimum(self.jitter_rate,
                                        np.min(jf/np.abs(self.jitter_gradients())))
            self.lambda_rate = np.minimum(self.mean_rate,
                                        np.min(lf/np.abs(self.lambda_gradients())))

    def mean_gradients(self):
        return -2. * self.dmmi

    def jitter_gradients(self):
        pt = np.dot(self.inv_cov.T,self.dmm.T)
        return self.D * np.diag(self.inv_cov) - pt * self.dmmi

    def lambda_gradients(self):
        pt1 = np.dot(self.inv_cov,self.lam) # d by m
        pt2 = np.zeros(self.M)
        for m in range(self.M):
            tmp = np.dot(self.inv_cov,self.lam[:,m])
            pt2[m] = np.dot((self.dmm,tmp) # scalar
        v = self.dmm[None,:] * pt2[:,None]
        pt2 = np.zeros((self.D,self.M))
        for m in range(self.M):
            pt2[:,m] = np.dot(self.inv_cov,v[m,:])
        return 2. * (self.D * pt1 - pt2)

    def total_negative_log_likelihood(self):
        totnll = 0.0
        for i in range(self.N):
            self.datum = self.data[i,:]
            self.variance = self.obsvar[i,:]
            self.do_precalcs()
            totnll += self.single_negative_log_likelihood()
        return totnll

    def single_negative_log_likelihood(self):
        sgn, logdet = self.single_slogdet_cov()
        assert sgn>0
        pt1 = self.D * logdet
        pt2 = np.dot(self.dmm,np.dot(self.inv_cov,self.dmm.T))
        return (pt1 + pt2)

    def make_cov(self):
        psi = np.diag(self.variance+self.jitter)
        lamlamT = np.dot(self.lam,self.lam.T)
        return psi+lamlamT

    def single_slogdet_cov(self):
        """
        Return sign and value of log det of covariance,
        using the matrix determinant lemma
        """
        psid = self.variance+self.jitter
        pt1 = np.eye(self.M) + np.dot(self.lam.T,self.lam / psid[:,None])
        det = np.linalg.det(pt1) * np.prod(psid)
        sgn = 1
        if det<0:
            sgn = -1
        return sgn, np.log(np.abs(det))

    def _check_gradients(self,Ncheck = 10):

        mgrad = self.mean_gradients()
        jgrad = self.jitter_gradients()
        lgrad = self.lambda_gradients()
        Dsamp = np.random.permutation(self.D)
        Msamp = np.random.permutation(self.M)

        for ii in range(Ncheck):

            i = np.mod(ii,self.D)
            D = Dsamp[i]
            i = np.mod(ii,self.M)
            M = Msamp[i]

            estm = self._check_one_gradient('mean',D)
            estj = self._check_one_gradient('jitter',D)
            estl = self._check_one_gradient('lambda',(D,M))

            print '\n\nCheck #%d, D = %d M = %d' % (ii,D,M)
            print '(Used mean grad., ' + \
                'Est. mean grad) = %g %g' % (mgrad[D],estm)
            print '(Used jitter grad., ' + \
                'Est. jitter grad) = %g %g' % (jgrad[D],estj)
            print '(Used lambda grad., ' + \
                'Est. lamda grad) = %g %g' % (lgrad[D,M],estl)

    def _check_one_gradient(self,kind,ind,eps = 1.0e-3):

        if kind=='mean':
            parm = self.mean
        elif kind=='jitter':
            parm = self.jitter
        else:
            parm = self.lam

        h = np.abs(parm[ind]) * eps
        if h==0.0: h = eps

        parm[ind] += h
        self.invert_cov()
        l1 = self.single_negative_log_likelihood()
        parm[ind] -= h
        self.invert_cov()
        l2 = self.single_negative_log_likelihood()

        return (l1 - l2) / (h)

"""

    def run_inference(self,max_iter,check_factor,Nosc):

        check_iter = np.round(check_factor * self.N)

        count = 0
        nll_best = np.Inf
        nll = np.Inf
        for ii in range(max_iter):

            if (ii%check_iter)==0:
                if ii==0:
                    nll_new = self.initial_nll
                else:
                    nll_new = self.total_negative_log_likelihood()
                print 'Iteration %d has neg. log likelhood %g' % (ii,nll_new)
                dlt_nll = nll - nll_new
                if nll_new<nll_best:
                    lam = self.lam
                    jitter = self.jitter
                    mean = self.mean
                    nll_best = nll_new
                if (dlt_nll<0.):
                    count += 1
                    if count==Nosc:
                        self.mean = mean
                        self.lam = lam
                        self.jitter = jitter
                        self.nll = nll_best
                        break
                nll = nll_new
            if (ii%self.N)==0:
                ind = np.random.permutation(self.N)
                self.data = self.data[ind,:]
                self.obsvar = self.obsvar[ind,:]

            i = np.mod(ii,self.data.shape[0])
            self.datum = self.data[i,:]
            self.variance = self.obsvar[i,:]

            self.invert_cov()
            self.make_gradient_step()
"""
<commit_msg>slight speedup<commit_after>import numpy as np
import matplotlib.pyplot as pl
from time import time

from empca import *
from rotate_factor import ortho_rotation

class FAModel(object):
    """
    Factor Analysis deconvolution, using online learning.
    """
    def __init__(self,data,obsvar,latent_dim,
                 max_iter=100000,compute_total_nll=False,
                 learning_init_relsizes=[0.01,0.01,0.01],
                 Ninit_est=10,check_rate=100,Neff=10,
                 decay_t0_factor=1.,decay_pow=0.5):
        """

        """
        assert data.shape==obsvar.shape, 'Data and obs. ' + \
            'variance have different shapes'

        self.N = data.shape[0]
        self.D = data.shape[1]
        self.M = latent_dim
        self.psiI = np.zeros((self.D,self.D))
        self.eyeM = np.eye(self.M)
        self.avg_factor = np.exp(-1./Neff)
        self.jitter = np.zeros(self.D)
        self.decay_t0 = decay_t0_factor * self.N
        self.decay_pow = decay_pow
        self.running_nll = np.ones(max_iter/check_rate) * -np.Inf
        self.psi_diag_inds = np.diag_indices_from(self.psiI)

        # Suffling the data
        ind = np.random.permutation(self.N)
        self.data = data[ind,:]
        self.obsvar = obsvar[ind,:]

        self.initialize()
        self.initial_mean   = self.mean.copy()
        self.initial_jitter = self.jitter.copy()
        self.initial_lambda = self.lam.copy()
        if compute_total_nll:
            self.initial_nll = self.total_negative_log_likelihood()
        self.calc_init_rates(Ninit_est,learning_init_relsizes)
        self.run_inference(max_iter,check_rate)

    def initialize(self):

        # inverse variance weighted mean
        if np.sum(self.obsvar)!=0.0:
            self.mean = np.sum(self.data / self.obsvar,axis=0) / \
                np.sum(1.0 / self.obsvar,axis=0)
        else:
            self.mean = np.mean(self.data,axis=0)

        # use EM PCA to initialize factor loadings
        if self.M==0:
            self.lam = np.zeros(1)
        else:
            empca = EMPCA(self.data.T,self.M)
            R = ortho_rotation(empca.lam,method='varimax')
            self.lam = np.dot(empca.lam,R)

    def run_inference(self,max_iter,check_rate):

        j = 0
        for ii in range(max_iter):

            # shuffle
            if (ii%self.N)==0:
                ind = np.random.permutation(self.N)
                self.data = self.data[ind,:]
                self.obsvar = self.obsvar[ind,:]

            # specify datum
            i = np.mod(ii,self.data.shape[0])
            self.datum = self.data[i,:]
            self.variance = self.obsvar[i,:]
            t = time()
            self.do_precalcs()
            print time()-t

            # Neg. log likelihood
            if (ii%check_rate)==0:
                j = self.estimate_nll(j)
                if (j%10)==0: print ii,j,self.running_nll[j]
                # convergence test...

            # not converged, make a step
            t = time()
            self.calc_rates(j)
            print time()-t
            t = time()
            self.make_gradient_step()
            print time()-t
            print ii
            if ii==1:
                assert 0

    def estimate_nll(self,j):
        est_nll = self.single_negative_log_likelihood()
        if j==0:
            self.running_nll[j] = est_nll
        else:
            self.running_nll[j] = self.running_nll[j-1] * self.avg_factor + \
                (1. - self.avg_factor) * est_nll
        return j+1

    def do_precalcs(self):
        """
        Calculate matrices used repeatedly below
        """
        self.psid = self.jitter + self.variance
        self.invert_cov()
        self.dmm = self.datum - self.mean
        self.dmmi = np.dot(self.dmm,self.inv_cov)

    def invert_cov(self):
        """
        Make inverse covariance, using the inversion lemma.
        """
        if self.M==0:
            self.inv_cov = np.diag(1.0 / self.psid)
        else:
            if np.sum(self.psid)!=0.0:
                psiI = np.diag(1.0 / self.psid)
                psiIlam = np.zeros((self.D,self.M))
                psiIlam = self.crazy_dot(psiI,self.lam,psiIlam)

                # the lemma, last step is slowest
                foo = np.linalg.inv(self.eyeM + np.dot(self.lam.T,psiIlam))
                self.inv_cov = psiI
                self.inv_cov -= np.dot(psiIlam,np.dot(foo,psiIlam.T))
            else:
                # fix to lemma-like version
                self.inv_cov = np.linalg.inv(np.dot(self.lam,self.lam.T))

    def crazy_dot(self,a,b,result):
        """
        Somehow this is faster...
        """
        for m in range(self.M):
            result[:,m] = np.dot(a,b[:,m])
        return result

    def make_gradient_step(self):

        self.mean -= self.mean_rate * self.mean_gradients()
        self.jitter -= self.jitter_rate * self.jitter_gradients()
        self.lam -= self.lambda_rate * self.lambda_gradients()

        ind = (self.jitter < 0.0)
        self.jitter[ind] = 0.0

    def calc_rates(self,t):

        t0 = self.decay_t0 * self.N
        p  = self.decay_pow

        self.mean_rate *= (t0 / (t0+t)) ** p
        self.jitter_rate *= (t0 / (t0+t)) ** p
        self.lambda_rate *= (t0 / (t0+t)) ** p

    def calc_init_rates(self,Ninit_est,learning_init_relsizes):

        self.mean_rate = np.Inf
        self.jitter_rate = np.Inf
        self.lambda_rate = np.Inf

        mf = learning_init_relsizes[0]
        jf = learning_init_relsizes[1]
        lf = learning_init_relsizes[2]

        ind = np.random.permutation(self.N)
        for i in range(Ninit_est):
            self.datum = self.data[ind[i],:]
            self.variance = self.obsvar[ind[i],:]
            self.do_precalcs()

            self.mean_rate = np.minimum(self.mean_rate,
                                        np.min(mf/np.abs(self.mean_gradients())))
            self.jitter_rate = np.minimum(self.jitter_rate,
                                        np.min(jf/np.abs(self.jitter_gradients())))
            self.lambda_rate = np.minimum(self.mean_rate,
                                        np.min(lf/np.abs(self.lambda_gradients())))

    def mean_gradients(self):
        return -2. * self.dmmi

    def jitter_gradients(self):
        return self.D * self.inv_cov[self.psi_diag_inds] - self.dmmi * self.dmmi

    def lambda_gradients(self):
        pt1 = np.dot(self.inv_cov,self.lam)
        pt2 = np.dot(self.dmm,pt1)
        v = self.dmm[:,None] * pt2[None,:]
        pt2 = np.zeros((self.D,self.M)) # crazy, this is faster.
        for m in range(self.M):
            pt2[:,m] = np.dot(self.inv_cov,v[:,m])
        return 2. * (self.D * pt1 - pt2)

    def total_negative_log_likelihood(self):
        totnll = 0.0
        for i in range(self.N):
            self.datum = self.data[i,:]
            self.variance = self.obsvar[i,:]
            self.do_precalcs()
            totnll += self.single_negative_log_likelihood()
        return totnll

    def single_negative_log_likelihood(self):
        sgn, logdet = self.single_slogdet_cov()
        assert sgn>0
        pt1 = self.D * logdet
        pt2 = np.dot(self.dmm,np.dot(self.inv_cov,self.dmm.T))
        return (pt1 + pt2)

    def make_cov(self):
        psi = np.diag(self.psid)
        lamlamT = np.dot(self.lam,self.lam.T)
        return psi+lamlamT

    def single_slogdet_cov(self):
        """
        Return sign and value of log det of covariance,
        using the matrix determinant lemma
        """
        pt1 = self.eyeM + np.dot(self.lam.T,self.lam / self.psid[:,None])
        det = np.linalg.det(pt1) * np.prod(self.psid)
        sgn = 1
        if det<0:
            sgn = -1
        return sgn, np.log(np.abs(det))

    def _check_gradients(self,Ncheck = 10):

        mgrad = self.mean_gradients()
        jgrad = self.jitter_gradients()
        lgrad = self.lambda_gradients()
        Dsamp = np.random.permutation(self.D)
        Msamp = np.random.permutation(self.M)

        for ii in range(Ncheck):

            i = np.mod(ii,self.D)
            D = Dsamp[i]
            i = np.mod(ii,self.M)
            M = Msamp[i]

            estm = self._check_one_gradient('mean',D)
            estj = self._check_one_gradient('jitter',D)
            estl = self._check_one_gradient('lambda',(D,M))

            print '\n\nCheck #%d, D = %d M = %d' % (ii,D,M)
            print '(Used mean grad., ' + \
                'Est. mean grad) = %g %g' % (mgrad[D],estm)
            print '(Used jitter grad., ' + \
                'Est. jitter grad) = %g %g' % (jgrad[D],estj)
            print '(Used lambda grad., ' + \
                'Est. lamda grad) = %g %g' % (lgrad[D,M],estl)

    def _check_one_gradient(self,kind,ind,eps = 1.0e-3):

        if kind=='mean':
            parm = self.mean
        elif kind=='jitter':
            parm = self.jitter
        else:
            parm = self.lam

        h = np.abs(parm[ind]) * eps
        if h==0.0: h = eps

        parm[ind] += h
        self.invert_cov()
        l1 = self.single_negative_log_likelihood()
        parm[ind] -= h
        self.invert_cov()
        l2 = self.single_negative_log_likelihood()

        return (l1 - l2) / (h)

"""
in lambda gradients -
        pt2 = np.zeros(self.M)
        for m in range(self.M):
            tmp = np.dot(self.inv_cov,self.lam[:,m])
            pt2[m] = np.dot(self.dmm,tmp) # scalar

        pt2 = np.dot(self.inv_cov,v)

"""
<|endoftext|>


<commit_before># coding: utf-8

from __future__ import division, unicode_literals

"""
Some reimplementation of Henkelman's Transition State Analysis utilities,
which are originally in Perl. Additional features beyond those offered by
Henkelman's utilities will be added.

This allows the usage and customization in Python.
"""

__author__ = 'Shyue Ping Ong'
__copyright__ = 'Copyright 2013, The Materials Virtual Lab'
__version__ = '0.1'
__maintainer__ = 'Shyue Ping Ong'
__email__ = 'ongsp@ucsd.edu'
__date__ = '6/1/15'

import os
import glob

import numpy as np
from scipy.interpolate import PiecewisePolynomial

from pymatgen.util.plotting_utils import get_publication_quality_plot
from pymatgen.io.vaspio import Poscar, Outcar

def is_int(s):
    try:
        int(s)
        return True
    except ValueError:
        return False

class NEBAnalysis(object):
    """
    An NEBAnalysis class.
    """

    def __init__(self, outcars, structures):
        """
        Initializes an NEBAnalysis from Outcar and Structure objects. Use
        the static constructors, e.g., :class:`from_dir` instead if you
        prefer to have these automatically generated from a directory of NEB
        calculations.

        Args:
            outcars ([Outcar]): List of Outcar objects. Note that these have
                to be ordered from start to end along reaction coordinates.
            structures ([Structure]): List of Structures along reaction
                coordinate. Must be same length as outcar.
        """
        if len(outcars) != len(structures):
            raise ValueError("# of Outcars must be same as # of Structures")

        # Calculate cumulative root mean square distance between structures,
        # which serves as the reaction coordinate. Note that these are
        # calculated from the final relaxed structures as the coordinates may
        # have changed from the initial interpolation.
        r = [0]
        prev = structures[0]
        for st in structures[1:]:
            dists = np.array([s2.distance(s1) for s1, s2 in zip(prev, st)])
            r.append(np.sqrt(np.sum(dists ** 2)))
            prev = st
        r = np.cumsum(r)

        energies = []
        forces = []
        for i, o in enumerate(outcars):
            o.read_neb()
            energies.append(o.data["energy"])
            if i in [0, len(outcars) - 1]:
                forces.append(0)
            else:
                forces.append(o.data["tangent_force"])
        energies = np.array(energies)
        energies -= energies[0]
        forces = np.array(forces)
        self.r = np.array(r)
        self.energies = energies
        self.forces = forces
        self.spline = PiecewisePolynomial(self.r,
                                          np.array([self.energies,
                                                    -self.forces]).T,
                                          orders=3)

    def get_extrema(self, normalize_rxn_coordinate=True):
        """
        Returns the positions of the extrema along the MEP. Both local
        minimums and maximums are returned.

        Args:
            normalize_rxn_coordinate (bool): Whether to normalize the
                reaction coordinate to between 0 and 1. Defaults to True.

        Returns:
            (min_extrema, max_extrema), where the extrema are given as
            [(x1, y1), (x2, y2), ...].
        """
        x = []
        y = []
        scale = 1 if not normalize_rxn_coordinate else 1 / self.r[-1]
        for i in range(0, len(self.r) - 1):
            dr = self.r[i + 1] - self.r[i]
            a, b, c, d = self.spline_params[i]
            f = np.arange(0, 1.0, 0.01)
            x.extend(self.r[i] + f * dr)
            y.extend(a * f ** 3 + b * f ** 2 + c * f + d)
        min_extrema = []
        max_extrema = []
        for i in range(1, len(x) - 1):
            if y[i] < y[i-1] and y[i] < y[i+1]:
                min_extrema.append((x[i] * scale, y[i]))
            elif y[i] > y[i-1] and y[i] > y[i+1]:
                max_extrema.append((x[i] * scale, y[i]))
        return min_extrema, max_extrema

    def get_plot(self, normalize_rxn_coordinate=True, label_barrier=True):
        """
        Returns the NEB plot. Uses Henkelman's approach of spline fitting
        each section of the reaction path based on tangent force and energies.

        Args:
            normalize_rxn_coordinate (bool): Whether to normalize the
                reaction coordinate to between 0 and 1. Defaults to True.
            label_barrier (bool): Whether to label the maximum barrier.

        Returns:
            matplotlib.pyplot object.
        """
        plt = get_publication_quality_plot(12, 8)
        scale = 1 if not normalize_rxn_coordinate else 1 / self.r[-1]
        x = np.arange(0, np.max(self.r), 0.01)
        y = self.spline(x) * 1000
        plt.plot(x * scale, y, 'k-', linewidth=2)

        plt.plot(self.r * scale, self.energies * 1000, 'ro', markersize=10)
        plt.xlabel("Reaction coordinate")
        plt.ylabel("Energy (meV)")
        plt.ylim((np.min(y) - 10, np.max(y) * 1.02 + 20))
        if label_barrier:
            data = zip(x * scale, y)
            barrier = max(data, key=lambda d: d[1])
            plt.plot([0, barrier[0]], [barrier[1], barrier[1]], 'k--')
            plt.annotate('%.0f meV' % barrier[1],
                         xy=(barrier[0] / 2, barrier[1] * 1.02),
                         xytext=(barrier[0] / 2, barrier[1] * 1.02))
        plt.tight_layout()
        return plt

    @classmethod
    def from_dir(cls, root_dir, relaxation_dirs=None):
        """
        Initializes a NEBAnalysis object from a directory of a NEB run.
        Note that OUTCARs must be present in all image directories. For the
        terminal OUTCARs from relaxation calculations, you can specify the
        locations using relaxation_dir. If these are not specified, the code
        will attempt to look for the OUTCARs in 00 and 0n directories,
        followed by subdirs "start", "end" or "initial", "final" in the
        root_dir. These are just some typical conventions used
        preferentially in Shyue Ping's MAVRL research group. For the
        non-terminal points, the CONTCAR is read to obtain structures. For
        terminal points, the POSCAR is used. The image directories are
        assumed to be the only directories that can be resolved to integers.
        E.g., "00", "01", "02", "03", "04", "05", "06". The minimum
        sub-directory structure that can be parsed is of the following form (
        a 5-image example is shown):

        00:
        - POSCAR
        - OUTCAR
        01, 02, 03, 04, 05:
        - CONTCAR
        - OUTCAR
        06:
        - POSCAR
        - OUTCAR

        Args:
            root_dir (str): Path to the root directory of the NEB calculation.
            relaxation_dirs (tuple): This specifies the starting and ending
                relaxation directories from which the OUTCARs are read for the
                terminal points for the energies.

        Returns:
            NEBAnalysis object.
        """
        neb_dirs = []
        for d in os.listdir(root_dir):
            pth = os.path.join(root_dir, d)
            if os.path.isdir(pth) and is_int(d):
                i = int(d)
                neb_dirs.append((i, pth))
        neb_dirs = sorted(neb_dirs, key=lambda d: d[0])
        outcars = []
        structures = []

        # Setup the search sequence for the OUTCARs for the terminal
        # directories.
        terminal_dirs = []
        if relaxation_dirs is not None:
            terminal_dirs.append(relaxation_dirs)
        terminal_dirs.append((neb_dirs[0][1], neb_dirs[-1][1]))
        terminal_dirs.append([os.path.join(root_dir, d)
                              for d in ["start", "end"]])
        terminal_dirs.append([os.path.join(root_dir, d)
                              for d in ["initial", "final"]])

        for i, d in neb_dirs:
            outcar = glob.glob(os.path.join(d, "OUTCAR*"))
            contcar = glob.glob(os.path.join(d, "CONTCAR*"))
            poscar = glob.glob(os.path.join(d, "POSCAR*"))
            terminal = i == 0 or i == neb_dirs[-1][0]
            if terminal:
                found = False
                for ds in terminal_dirs:
                    od = ds[0] if i == 0 else ds[1]
                    outcar = glob.glob(os.path.join(od, "OUTCAR*"))
                    if outcar:
                        outcar = sorted(outcar)
                        outcars.append(Outcar(outcar[-1]))
                        found = True
                        break
                if not found:
                    raise ValueError("OUTCAR cannot be found for terminal "
                                     "point %s" % d)
                structures.append(Poscar.from_file(poscar[0]).structure)
            else:
                outcars.append(Outcar(outcar[0]))
                structures.append(Poscar.from_file(contcar[0]).structure)
        return NEBAnalysis(outcars, structures)<commit_msg>Fix get_extrema for new scipy fitting.<commit_after># coding: utf-8

from __future__ import division, unicode_literals

"""
Some reimplementation of Henkelman's Transition State Analysis utilities,
which are originally in Perl. Additional features beyond those offered by
Henkelman's utilities will be added.

This allows the usage and customization in Python.
"""

__author__ = 'Shyue Ping Ong'
__copyright__ = 'Copyright 2013, The Materials Virtual Lab'
__version__ = '0.1'
__maintainer__ = 'Shyue Ping Ong'
__email__ = 'ongsp@ucsd.edu'
__date__ = '6/1/15'

import os
import glob

import numpy as np
from scipy.interpolate import PiecewisePolynomial

from pymatgen.util.plotting_utils import get_publication_quality_plot
from pymatgen.io.vaspio import Poscar, Outcar

def is_int(s):
    try:
        int(s)
        return True
    except ValueError:
        return False

class NEBAnalysis(object):
    """
    An NEBAnalysis class.
    """

    def __init__(self, outcars, structures):
        """
        Initializes an NEBAnalysis from Outcar and Structure objects. Use
        the static constructors, e.g., :class:`from_dir` instead if you
        prefer to have these automatically generated from a directory of NEB
        calculations.

        Args:
            outcars ([Outcar]): List of Outcar objects. Note that these have
                to be ordered from start to end along reaction coordinates.
            structures ([Structure]): List of Structures along reaction
                coordinate. Must be same length as outcar.
        """
        if len(outcars) != len(structures):
            raise ValueError("# of Outcars must be same as # of Structures")

        # Calculate cumulative root mean square distance between structures,
        # which serves as the reaction coordinate. Note that these are
        # calculated from the final relaxed structures as the coordinates may
        # have changed from the initial interpolation.
        r = [0]
        prev = structures[0]
        for st in structures[1:]:
            dists = np.array([s2.distance(s1) for s1, s2 in zip(prev, st)])
            r.append(np.sqrt(np.sum(dists ** 2)))
            prev = st
        r = np.cumsum(r)

        energies = []
        forces = []
        for i, o in enumerate(outcars):
            o.read_neb()
            energies.append(o.data["energy"])
            if i in [0, len(outcars) - 1]:
                forces.append(0)
            else:
                forces.append(o.data["tangent_force"])
        energies = np.array(energies)
        energies -= energies[0]
        forces = np.array(forces)
        self.r = np.array(r)
        self.energies = energies
        self.forces = forces
        self.spline = PiecewisePolynomial(self.r,
                                          np.array([self.energies,
                                                    -self.forces]).T,
                                          orders=3)

    def get_extrema(self, normalize_rxn_coordinate=True):
        """
        Returns the positions of the extrema along the MEP. Both local
        minimums and maximums are returned.

        Args:
            normalize_rxn_coordinate (bool): Whether to normalize the
                reaction coordinate to between 0 and 1. Defaults to True.

        Returns:
            (min_extrema, max_extrema), where the extrema are given as
            [(x1, y1), (x2, y2), ...].
        """
        x = np.arange(0, np.max(self.r), 0.01)
        y = self.spline(x) * 1000

        scale = 1 if not normalize_rxn_coordinate else 1 / self.r[-1]
        min_extrema = []
        max_extrema = []
        for i in range(1, len(x) - 1):
            if y[i] < y[i-1] and y[i] < y[i+1]:
                min_extrema.append((x[i] * scale, y[i]))
            elif y[i] > y[i-1] and y[i] > y[i+1]:
                max_extrema.append((x[i] * scale, y[i]))
        return min_extrema, max_extrema

    def get_plot(self, normalize_rxn_coordinate=True, label_barrier=True):
        """
        Returns the NEB plot. Uses Henkelman's approach of spline fitting
        each section of the reaction path based on tangent force and energies.

        Args:
            normalize_rxn_coordinate (bool): Whether to normalize the
                reaction coordinate to between 0 and 1. Defaults to True.
            label_barrier (bool): Whether to label the maximum barrier.

        Returns:
            matplotlib.pyplot object.
        """
        plt = get_publication_quality_plot(12, 8)
        scale = 1 if not normalize_rxn_coordinate else 1 / self.r[-1]
        x = np.arange(0, np.max(self.r), 0.01)
        y = self.spline(x) * 1000
        plt.plot(x * scale, y, 'k-', linewidth=2)

        plt.plot(self.r * scale, self.energies * 1000, 'ro', markersize=10)
        plt.xlabel("Reaction coordinate")
        plt.ylabel("Energy (meV)")
        plt.ylim((np.min(y) - 10, np.max(y) * 1.02 + 20))
        if label_barrier:
            data = zip(x * scale, y)
            barrier = max(data, key=lambda d: d[1])
            plt.plot([0, barrier[0]], [barrier[1], barrier[1]], 'k--')
            plt.annotate('%.0f meV' % barrier[1],
                         xy=(barrier[0] / 2, barrier[1] * 1.02),
                         xytext=(barrier[0] / 2, barrier[1] * 1.02))
        plt.tight_layout()
        return plt

    @classmethod
    def from_dir(cls, root_dir, relaxation_dirs=None):
        """
        Initializes a NEBAnalysis object from a directory of a NEB run.
        Note that OUTCARs must be present in all image directories. For the
        terminal OUTCARs from relaxation calculations, you can specify the
        locations using relaxation_dir. If these are not specified, the code
        will attempt to look for the OUTCARs in 00 and 0n directories,
        followed by subdirs "start", "end" or "initial", "final" in the
        root_dir. These are just some typical conventions used
        preferentially in Shyue Ping's MAVRL research group. For the
        non-terminal points, the CONTCAR is read to obtain structures. For
        terminal points, the POSCAR is used. The image directories are
        assumed to be the only directories that can be resolved to integers.
        E.g., "00", "01", "02", "03", "04", "05", "06". The minimum
        sub-directory structure that can be parsed is of the following form (
        a 5-image example is shown):

        00:
        - POSCAR
        - OUTCAR
        01, 02, 03, 04, 05:
        - CONTCAR
        - OUTCAR
        06:
        - POSCAR
        - OUTCAR

        Args:
            root_dir (str): Path to the root directory of the NEB calculation.
            relaxation_dirs (tuple): This specifies the starting and ending
                relaxation directories from which the OUTCARs are read for the
                terminal points for the energies.

        Returns:
            NEBAnalysis object.
        """
        neb_dirs = []
        for d in os.listdir(root_dir):
            pth = os.path.join(root_dir, d)
            if os.path.isdir(pth) and is_int(d):
                i = int(d)
                neb_dirs.append((i, pth))
        neb_dirs = sorted(neb_dirs, key=lambda d: d[0])
        outcars = []
        structures = []

        # Setup the search sequence for the OUTCARs for the terminal
        # directories.
        terminal_dirs = []
        if relaxation_dirs is not None:
            terminal_dirs.append(relaxation_dirs)
        terminal_dirs.append((neb_dirs[0][1], neb_dirs[-1][1]))
        terminal_dirs.append([os.path.join(root_dir, d)
                              for d in ["start", "end"]])
        terminal_dirs.append([os.path.join(root_dir, d)
                              for d in ["initial", "final"]])

        for i, d in neb_dirs:
            outcar = glob.glob(os.path.join(d, "OUTCAR*"))
            contcar = glob.glob(os.path.join(d, "CONTCAR*"))
            poscar = glob.glob(os.path.join(d, "POSCAR*"))
            terminal = i == 0 or i == neb_dirs[-1][0]
            if terminal:
                found = False
                for ds in terminal_dirs:
                    od = ds[0] if i == 0 else ds[1]
                    outcar = glob.glob(os.path.join(od, "OUTCAR*"))
                    if outcar:
                        outcar = sorted(outcar)
                        outcars.append(Outcar(outcar[-1]))
                        found = True
                        break
                if not found:
                    raise ValueError("OUTCAR cannot be found for terminal "
                                     "point %s" % d)
                structures.append(Poscar.from_file(poscar[0]).structure)
            else:
                outcars.append(Outcar(outcar[0]))
                structures.append(Poscar.from_file(contcar[0]).structure)
        return NEBAnalysis(outcars, structures)<|endoftext|>


<commit_before>import yaml
from runners.models import Runner
from games.models import DEFAULT_INSTALLER
SUCCESS = (True, "")

def validate_installer(installer):
    errors = []
    is_valid = True
    rules = [
        script_is_not_the_default_one,
        doesnt_contain_useless_fields,
        files_is_an_array,
        scummvm_has_gameid,
    ]
    for rule in rules:
        success, message = rule(installer)
        if not success:
            errors.append(message)
            is_valid = False

    return (is_valid, errors)

def script_is_not_the_default_one(installer):
    script = yaml.safe_load(installer.content)
    if script == DEFAULT_INSTALLER:
        return (
            False,
            "Really? You haven't even modified the default installer"
        )
    return SUCCESS

def doesnt_contain_useless_fields(installer):
    script = yaml.safe_load(installer.content)
    for field in (
        'version', 'gogid', 'humbleid', 'game_slug', 'description',
        'installer_slug', 'name', 'notes', 'runner', 'slug', 'steamid', 'year'
    ):
        if field in script:
            return (False, "Don't put a '{}' field in the script.".format(field))
    return SUCCESS

def files_is_an_array(installer):
    script = yaml.safe_load(installer.content)
    if 'files' in script:
        if not isinstance(script['files'], list):
            return (False, "'files' section should be an array.")
    return SUCCESS

def scummvm_has_gameid(installer):
    try:
        runner = installer.runner.name
    except Runner.DoesNotExist:
        runner = ""
    script = yaml.safe_load(installer.content)
    if runner != 'scummvm':
        return SUCCESS
    if 'game' not in script:
        return (
            False,
            "Missing section 'game'"
        )
    if 'game_id' not in script['game']:
        return (
            False,
            "ScummVM game should have a game identifier in the 'game' section"
        )
    return SUCCESS
<commit_msg>Provide fallback if no script is found<commit_after>import yaml
from runners.models import Runner
from games.models import DEFAULT_INSTALLER
SUCCESS = (True, "")

def get_installer_script(installer):
    script = yaml.safe_load(installer.content)
    if not script:
        return {}
    return script

def validate_installer(installer):
    errors = []
    is_valid = True
    rules = [
        script_is_not_the_default_one,
        doesnt_contain_useless_fields,
        files_is_an_array,
        scummvm_has_gameid,
    ]
    for rule in rules:
        success, message = rule(installer)
        if not success:
            errors.append(message)
            is_valid = False

    return (is_valid, errors)

def script_is_not_the_default_one(installer):
    script = get_installer_script(installer)
    if script == DEFAULT_INSTALLER:
        return (
            False,
            "Really? You haven't even modified the default installer"
        )
    return SUCCESS

def doesnt_contain_useless_fields(installer):
    script = get_installer_script(installer)
    for field in (
        'version', 'gogid', 'humbleid', 'game_slug', 'description',
        'installer_slug', 'name', 'notes', 'runner', 'slug', 'steamid', 'year'
    ):
        if field in script:
            return (False, "Don't put a '{}' field in the script.".format(field))
    return SUCCESS

def files_is_an_array(installer):
    script = get_installer_script(installer)
    if 'files' in script:
        if not isinstance(script['files'], list):
            return (False, "'files' section should be an array.")
    return SUCCESS

def scummvm_has_gameid(installer):
    try:
        runner = installer.runner.name
    except Runner.DoesNotExist:
        runner = ""
    script = get_installer_script(installer)
    if runner != 'scummvm':
        return SUCCESS
    if 'game' not in script:
        return (
            False,
            "Missing section 'game'"
        )
    if 'game_id' not in script['game']:
        return (
            False,
            "ScummVM game should have a game identifier in the 'game' section"
        )
    return SUCCESS
<|endoftext|>


<commit_before>from rest_framework import serializers as ser

from api.base.serializers import (
    JSONAPISerializer,
    RelationshipField,
    RestrictedDictSerializer,
    LinksField,
    is_anonymized,
    DateByVersion,
)
from website.project.model import Node
from website.util import permissions as osf_permissions
from framework.auth.core import User
from website.preprints.model import PreprintService

class NodeLogIdentifiersSerializer(RestrictedDictSerializer):

    doi = ser.CharField(read_only=True)
    ark = ser.CharField(read_only=True)

class NodeLogInstitutionSerializer(RestrictedDictSerializer):

    id = ser.CharField(read_only=True)
    name = ser.CharField(read_only=True)

class NodeLogFileParamsSerializer(RestrictedDictSerializer):

    materialized = ser.CharField(read_only=True)
    url = ser.URLField(read_only=True)
    addon = ser.CharField(read_only=True)
    node_url = ser.URLField(read_only=True, source='node.url')
    node_title = ser.SerializerMethodField()

    def get_node_title(self, obj):
        user = self.context['request'].user
        node_title = obj['node']['title']
        node = Node.load(obj['node']['_id'])
        if not user.is_authenticated():
            if node.is_public:
                return node_title
        elif node.has_permission(user, osf_permissions.READ):
            return node_title
        return 'Private Component'

class NodeLogParamsSerializer(RestrictedDictSerializer):

    addon = ser.CharField(read_only=True)
    bucket = ser.CharField(read_only=True)
    citation_name = ser.CharField(read_only=True, source='citation.name')
    contributors = ser.SerializerMethodField(read_only=True)
    data_set = ser.CharField(read_only=True, source='dataset')
    destination = NodeLogFileParamsSerializer(read_only=True)
    figshare_title = ser.CharField(read_only=True, source='figshare.title')
    forward_url = ser.CharField(read_only=True)
    github_user = ser.CharField(read_only=True, source='github.user')
    github_repo = ser.CharField(read_only=True, source='github.repo')
    file = ser.DictField(read_only=True)
    filename = ser.CharField(read_only=True)
    kind = ser.CharField(read_only=True)
    folder = ser.CharField(read_only=True)
    folder_name = ser.CharField(read_only=True)
    license = ser.CharField(read_only=True, source='new_license')
    identifiers = NodeLogIdentifiersSerializer(read_only=True)
    institution = NodeLogInstitutionSerializer(read_only=True)
    old_page = ser.CharField(read_only=True)
    page = ser.CharField(read_only=True)
    page_id = ser.CharField(read_only=True)
    params_node = ser.SerializerMethodField(read_only=True)
    params_project = ser.SerializerMethodField(read_only=True)
    path = ser.CharField(read_only=True)
    pointer = ser.DictField(read_only=True)
    preprint = ser.CharField(read_only=True)
    preprint_provider = ser.SerializerMethodField(read_only=True)
    previous_institution = NodeLogInstitutionSerializer(read_only=True)
    source = NodeLogFileParamsSerializer(read_only=True)
    study = ser.CharField(read_only=True)
    tag = ser.CharField(read_only=True)
    tags = ser.CharField(read_only=True)
    target = NodeLogFileParamsSerializer(read_only=True)
    template_node = ser.DictField(read_only=True)
    title_new = ser.CharField(read_only=True)
    title_original = ser.CharField(read_only=True)
    updated_fields = ser.DictField(read_only=True)
    urls = ser.DictField(read_only=True)
    version = ser.CharField(read_only=True)
    wiki = ser.DictField(read_only=True)
    citation_name = ser.CharField(read_only=True, source='citation.name')
    institution = NodeLogInstitutionSerializer(read_only=True)

    def get_view_url(self, obj):
        urls = obj.get('urls', None)
        if urls:
            view = urls.get('view', None)
            if view:
                return view
        return None

    def get_params_node(self, obj):
        node_id = obj.get('node', None)
        if node_id:
            node = Node.load(node_id)
            return {'id': node_id, 'title': node.title}
        return None

    def get_params_project(self, obj):
        project_id = obj.get('project', None)
        if project_id:
            node = Node.load(project_id)
            return {'id': project_id, 'title': node.title}
        return None

    def get_contributors(self, obj):

        contributor_info = []

        if is_anonymized(self.context['request']):
            return contributor_info

        contributor_ids = obj.get('contributors', None)
        params_node = obj.get('node', None)

        if contributor_ids:
            for contrib_id in contributor_ids:
                user = User.load(contrib_id)
                unregistered_name = None
                if user.unclaimed_records.get(params_node):
                    unregistered_name = user.unclaimed_records[params_node].get('name', None)

                contributor_info.append({
                    'id': contrib_id,
                    'full_name': user.fullname,
                    'given_name': user.given_name,
                    'middle_names': user.middle_names,
                    'family_name': user.family_name,
                    'unregistered_name': unregistered_name,
                    'active': user.is_active
                })
        return contributor_info

    def get_preprint_provider(self, obj):
        preprint_id = obj.get('preprint', None)
        if preprint_id:
            preprint = PreprintService.load(preprint_id)
            if preprint:
                provider = preprint.provider
                return {'url': provider.external_url, 'name': provider.name}
        return None

class NodeLogSerializer(JSONAPISerializer):

    filterable_fields = frozenset(['action', 'date'])
    non_anonymized_fields = [
        'id',
        'date',
        'action',
    ]

    id = ser.CharField(read_only=True, source='_id')
    date = DateByVersion(read_only=True)
    action = ser.CharField(read_only=True)
    params = NodeLogParamsSerializer(read_only=True)
    links = LinksField({'self': 'get_absolute_url'})

    class Meta:
        type_ = 'logs'

    node = RelationshipField(
        related_view=lambda n: 'registrations:registration-detail' if getattr(n, 'is_registration', False) else 'nodes:node-detail',
        related_view_kwargs={'node_id': '<node._id>'},
    )

    original_node = RelationshipField(
        related_view=lambda n: 'registrations:registration-detail' if getattr(n, 'is_registration', False) else 'nodes:node-detail',
        related_view_kwargs={'node_id': '<original_node._id>'},
    )

    user = RelationshipField(
        related_view='users:user-detail',
        related_view_kwargs={'user_id': '<user._id>'},
    )

    # This would be a node_link, except that data isn't stored in the node log params
    linked_node = RelationshipField(
        related_view='nodes:node-detail',
        related_view_kwargs={'node_id': '<params.pointer.id>'}
    )
    template_node = RelationshipField(
        related_view='nodes:node-detail',
        related_view_kwargs={'node_id': '<params.template_node.id>'}
    )

    def get_absolute_url(self, obj):
        return obj.absolute_url
<commit_msg>Query only for title when serializing params_node and params_project<commit_after>from rest_framework import serializers as ser

from api.base.serializers import (
    JSONAPISerializer,
    RelationshipField,
    RestrictedDictSerializer,
    LinksField,
    is_anonymized,
    DateByVersion,
)
from website.project.model import Node
from website.util import permissions as osf_permissions
from framework.auth.core import User
from website.preprints.model import PreprintService

class NodeLogIdentifiersSerializer(RestrictedDictSerializer):

    doi = ser.CharField(read_only=True)
    ark = ser.CharField(read_only=True)

class NodeLogInstitutionSerializer(RestrictedDictSerializer):

    id = ser.CharField(read_only=True)
    name = ser.CharField(read_only=True)

class NodeLogFileParamsSerializer(RestrictedDictSerializer):

    materialized = ser.CharField(read_only=True)
    url = ser.URLField(read_only=True)
    addon = ser.CharField(read_only=True)
    node_url = ser.URLField(read_only=True, source='node.url')
    node_title = ser.SerializerMethodField()

    def get_node_title(self, obj):
        user = self.context['request'].user
        node_title = obj['node']['title']
        node = Node.load(obj['node']['_id'])
        if not user.is_authenticated():
            if node.is_public:
                return node_title
        elif node.has_permission(user, osf_permissions.READ):
            return node_title
        return 'Private Component'

class NodeLogParamsSerializer(RestrictedDictSerializer):

    addon = ser.CharField(read_only=True)
    bucket = ser.CharField(read_only=True)
    citation_name = ser.CharField(read_only=True, source='citation.name')
    contributors = ser.SerializerMethodField(read_only=True)
    data_set = ser.CharField(read_only=True, source='dataset')
    destination = NodeLogFileParamsSerializer(read_only=True)
    figshare_title = ser.CharField(read_only=True, source='figshare.title')
    forward_url = ser.CharField(read_only=True)
    github_user = ser.CharField(read_only=True, source='github.user')
    github_repo = ser.CharField(read_only=True, source='github.repo')
    file = ser.DictField(read_only=True)
    filename = ser.CharField(read_only=True)
    kind = ser.CharField(read_only=True)
    folder = ser.CharField(read_only=True)
    folder_name = ser.CharField(read_only=True)
    license = ser.CharField(read_only=True, source='new_license')
    identifiers = NodeLogIdentifiersSerializer(read_only=True)
    institution = NodeLogInstitutionSerializer(read_only=True)
    old_page = ser.CharField(read_only=True)
    page = ser.CharField(read_only=True)
    page_id = ser.CharField(read_only=True)
    params_node = ser.SerializerMethodField(read_only=True)
    params_project = ser.SerializerMethodField(read_only=True)
    path = ser.CharField(read_only=True)
    pointer = ser.DictField(read_only=True)
    preprint = ser.CharField(read_only=True)
    preprint_provider = ser.SerializerMethodField(read_only=True)
    previous_institution = NodeLogInstitutionSerializer(read_only=True)
    source = NodeLogFileParamsSerializer(read_only=True)
    study = ser.CharField(read_only=True)
    tag = ser.CharField(read_only=True)
    tags = ser.CharField(read_only=True)
    target = NodeLogFileParamsSerializer(read_only=True)
    template_node = ser.DictField(read_only=True)
    title_new = ser.CharField(read_only=True)
    title_original = ser.CharField(read_only=True)
    updated_fields = ser.DictField(read_only=True)
    urls = ser.DictField(read_only=True)
    version = ser.CharField(read_only=True)
    wiki = ser.DictField(read_only=True)
    citation_name = ser.CharField(read_only=True, source='citation.name')
    institution = NodeLogInstitutionSerializer(read_only=True)

    def get_view_url(self, obj):
        urls = obj.get('urls', None)
        if urls:
            view = urls.get('view', None)
            if view:
                return view
        return None

    def get_params_node(self, obj):
        node_id = obj.get('node', None)
        if node_id:
            node = Node.objects.filter(guids___id=node_id).values('title').get()
            return {'id': node_id, 'title': node['title']}
        return None

    def get_params_project(self, obj):
        project_id = obj.get('project', None)
        if project_id:
            node = Node.objects.filter(guids___id=project_id).values('title').get()
            return {'id': project_id, 'title': node['title']}
        return None

    def get_contributors(self, obj):

        contributor_info = []

        if is_anonymized(self.context['request']):
            return contributor_info

        contributor_ids = obj.get('contributors', None)
        params_node = obj.get('node', None)

        if contributor_ids:
            for contrib_id in contributor_ids:
                user = User.load(contrib_id)
                unregistered_name = None
                if user.unclaimed_records.get(params_node):
                    unregistered_name = user.unclaimed_records[params_node].get('name', None)

                contributor_info.append({
                    'id': contrib_id,
                    'full_name': user.fullname,
                    'given_name': user.given_name,
                    'middle_names': user.middle_names,
                    'family_name': user.family_name,
                    'unregistered_name': unregistered_name,
                    'active': user.is_active
                })
        return contributor_info

    def get_preprint_provider(self, obj):
        preprint_id = obj.get('preprint', None)
        if preprint_id:
            preprint = PreprintService.load(preprint_id)
            if preprint:
                provider = preprint.provider
                return {'url': provider.external_url, 'name': provider.name}
        return None

class NodeLogSerializer(JSONAPISerializer):

    filterable_fields = frozenset(['action', 'date'])
    non_anonymized_fields = [
        'id',
        'date',
        'action',
    ]

    id = ser.CharField(read_only=True, source='_id')
    date = DateByVersion(read_only=True)
    action = ser.CharField(read_only=True)
    params = NodeLogParamsSerializer(read_only=True)
    links = LinksField({'self': 'get_absolute_url'})

    class Meta:
        type_ = 'logs'

    node = RelationshipField(
        related_view=lambda n: 'registrations:registration-detail' if getattr(n, 'is_registration', False) else 'nodes:node-detail',
        related_view_kwargs={'node_id': '<node._id>'},
    )

    original_node = RelationshipField(
        related_view=lambda n: 'registrations:registration-detail' if getattr(n, 'is_registration', False) else 'nodes:node-detail',
        related_view_kwargs={'node_id': '<original_node._id>'},
    )

    user = RelationshipField(
        related_view='users:user-detail',
        related_view_kwargs={'user_id': '<user._id>'},
    )

    # This would be a node_link, except that data isn't stored in the node log params
    linked_node = RelationshipField(
        related_view='nodes:node-detail',
        related_view_kwargs={'node_id': '<params.pointer.id>'}
    )
    template_node = RelationshipField(
        related_view='nodes:node-detail',
        related_view_kwargs={'node_id': '<params.template_node.id>'}
    )

    def get_absolute_url(self, obj):
        return obj.absolute_url
<|endoftext|>


<commit_before>from ontobio.io.gpadparser import GpadParser
from ontobio.io.gafparser import GafParser
from ontobio import ecomap
import click
import pandas as pd
import datetime
from ontobio.io import qc
from ontobio.io.assocparser import Report

@click.command()
@click.option("--file1", "-file1", type=click.Path(), required=True)
@click.option("--file2", "-file2", type=click.Path(), required=True)
@click.option("--output", "-o", required=True)
@click.option("--count_by", "-cb", multiple=True, required=False)
@click.option("--file_type", "-file_type", required=True)
@click.option("--exclude_details", "-ed", type=click.BOOL, default=True, required=False)
def compare_files(file1, file2, output, count_by, exclude_details, file_type):
    # decide which parser to instantiate, GAF or GPAD
    pd.set_option('display.max_rows', 35000)
    df_file1, df_file2, assocs1, assocs2 = get_parser(file1, file2, file_type)

    # if count_by is passed in, generate a report of annotations grouped by count_by
    # count_by is a list.
    generate_group_report(df_file1, df_file2, count_by, file1, file2, output)

    # if exclude_details is False, try to compare files line by line and find matches.
    compare_associations(assocs1, assocs2, exclude_details, output)

def generate_group_report(df_file1, df_file2, count_by, file1, file2, output):
    if len(count_by) > 0:
        group_by_report_file = open(output + "_group_by_report", "w")
        file1_groups = group_by(df_file1, count_by, file1)
        file2_groups = group_by(df_file2, count_by, file2)

        s = "\n\n## GROUP BY SUMMARY \n\n"
        s += "This report generated on {}\n\n".format(datetime.date.today())
        s += "  * Group By Columns: " + str(count_by) + "\n"
        s += "  * Compared Files: " + file1 + ", " + file2 + "\n"

        print(s)
        for grouped_item in file1_groups['grouped_reports']:
            print(file1_groups['filename'])
            group_by_report_file.write(file1_groups['filename'])
            print(grouped_item)
            group_by_report_file.write(str(grouped_item))
            print("\n")

        for grouped_item in file2_groups['grouped_reports']:
            print(file2_groups['filename'])
            group_by_report_file.write(file2_groups['filename'])
            print(grouped_item)
            group_by_report_file.write(str(grouped_item))

        group_by_report_file.close()

def compare_associations(assocs1, assocs2, exclude_details, output):
    compare_report_file = open(output + "_compare_report", "w")
    processed_lines = 0
    exact_matches = 0
    close_matches = 0

    report = Report()
    for association in assocs1:
        max_match_score = 0
        processed_lines = processed_lines + 1
        if not exclude_details:
            for target in assocs2:
                match_score = 0
                if association.negated != target.negated:
                    continue
                if association.subject.id == target.subject.id and association.object.id == target.object.id:
                    match_score = 1
                    if sorted(str(q).upper() for q in association.qualifiers) == \
                            sorted(str(q).upper() for q in target.qualifiers):
                        match_score = 2
                        if association.evidence.type == target.evidence.type:
                            match_score = 3
                            if sorted(str(w).upper() for w in association.evidence.with_support_from) == \
                                    sorted(str(w).upper() for w in target.evidence.with_support_from):
                                match_score = 4
                                if sorted(
                                        str(r).upper() for r in association.evidence.has_supporting_reference) == \
                                        sorted(str(r).upper() for r in target.evidence.has_supporting_reference):
                                    match_score = 5
                if match_score > max_match_score:
                    max_match_score = match_score
            if max_match_score > 4:
                exact_matches = exact_matches + 1
            elif 1 > max_match_score < 5:
                close_matches = close_matches + 1
                report.add_association(association)
                report.n_lines = report.n_lines + 1
                report.warning(association.source_line, qc.ResultType.WARNING,
                               "line from file1 only has CLOSE match in file2", "")
            else:
                report.add_association(association)
                report.n_lines = report.n_lines + 1
                report.error(association.source_line, qc.ResultType.ERROR,
                             "line from file1 has NO match in file2", "")
    if not exclude_details:
        md_report = markdown_report(report, exact_matches, close_matches, processed_lines)
        print(md_report)
        compare_report_file.write(md_report)
        compare_report_file.close()

def markdown_report(report, exact_matches, close_matches, processed_lines):

    json = report.to_report_json()

    s = "\n\n## DIFF SUMMARY\n\n"
    s += "This report generated on {}\n\n".format(datetime.date.today())
    s += "  * Total Unmatched Associations: {}\n".format(json["associations"])
    s += "  * Total Lines Compared: " + str(processed_lines) + "\n"
    s += "  * Total Exact matches: " + str(exact_matches) + "\n"
    s += "  * Total Close matches: " + str(close_matches) + "\n\n"

    for (rule, messages) in sorted(json["messages"].items(), key=lambda t: t[0]):
        s += "### {rule}\n\n".format(rule=rule)
        s += "* total: {amount}\n".format(amount=len(messages))
        s += "\n"
        if len(messages) > 0:
            s += "#### Messages\n\n"
        for message in messages:
            obj = " ({})".format(message["obj"]) if message["obj"] else ""
            s += "* {level} - {type}: {message}{obj} -- `{line}`\n".format(level=message["level"],
                                                                           type=message["type"],
                                                                           message=message["message"],
                                                                           line=message["line"],
                                                                           obj=obj)

        return s

def get_parser(file1, file2, file_type):
    if file_type == 'gpad':
        gpad_parser_1 = GpadParser()
        gpad_parser_2 = GpadParser()
        assocs1 = gpad_parser_1.parse(file1, skipheader=True)
        assocs2 = gpad_parser_2.parse(file2, skipheader=True)
        df_file1 = read_gpad_csv(file1)
        df_file2 = read_gpad_csv(file2)
    else:
        gaf_parser_1 = GafParser()
        gaf_parser_2 = GafParser()
        assocs1 = gaf_parser_1.parse(file1, skipheader=True)
        assocs2 = gaf_parser_2.parse(file2, skipheader=True)

        df_file1 = read_gaf_csv(file1)
        df_file2 = read_gaf_csv(file2)
    return df_file1, df_file2, assocs1, assocs2

def read_gaf_csv(filename):
    ecomapping = ecomap.EcoMap()
    data_frame = pd.read_csv(filename,
                             comment='!',
                             sep='\t',
                             header=None,
                             na_filter=False,
                             names=["DB",
                                    "DB_Object_ID",
                                    "DB_Object_Symbol",
                                    "Qualifier",
                                    "GO_ID",
                                    "DB_Reference",
                                    "Evidence_code",
                                    "With_or_From",
                                    "Aspect",
                                    "DB_Object_Name",
                                    "DB_Object_Synonym",
                                    "DB_Object_Type,"
                                    "Taxon",
                                    "Date",
                                    "Assigned_By",
                                    "Annotation_Extension",
                                    "Gene_Product_Form_ID"]).fillna("")
    for eco_code in ecomapping.mappings():
        for ev in data_frame['Evidence_code']:
            if eco_code[2] == ev:
                data_frame['Evidence_code'] = data_frame['Evidence_code'].replace([eco_code[2]],
                                                                                ecomapping.ecoclass_to_coderef(
                                                                                eco_code[2])[0])
    return data_frame

def read_gpad_csv(filename):
    ecomapping = ecomap.EcoMap()
    data_frame = pd.read_csv(filename,
                             comment='!',
                             sep='\t',
                             header=None,
                             na_filter=False,
                             names=["DB",
                                    "DB_Object_ID",
                                    "Relation",
                                    "Ontology_Class_ID",
                                    "Reference",
                                    "Evidence_type",
                                    "With_or_From",
                                    "Interacting_taxon_ID",
                                    "Date",
                                    "Assigned_by",
                                    "Annotation_Extensions",
                                    "Annotation_Properties"]).fillna("")
    for eco_code in ecomapping.mappings():
        for ev in data_frame['Evidence_type']:
            if eco_code[2] == ev:
                data_frame['Evidence_type'] = data_frame['Evidence_type'].replace([eco_code[2]],
                                                                                ecomapping.ecoclass_to_coderef(eco_code[2])[0])
    return data_frame

def group_by(data_frame, count_by, file):
    stats = {}
    grouped_reports = []
    stats['filename'] = file
    stats['total_rows'] = data_frame.shape[0]
    if len(count_by) > 0:
        for grouper in count_by:
            stats['grouper'] = grouper
            grouped_reports.append(data_frame.groupby(grouper)[grouper].count())
        stats['grouped_reports'] = grouped_reports
    return stats

if __name__ == '__main__':
    compare_files()
<commit_msg>add counts per column<commit_after>from ontobio.io.gpadparser import GpadParser
from ontobio.io.gafparser import GafParser
from ontobio import ecomap
import click
import pandas as pd
import datetime
from ontobio.io import qc
from ontobio.io.assocparser import Report

@click.command()
@click.option("--file1", "-file1", type=click.Path(), required=True)
@click.option("--file2", "-file2", type=click.Path(), required=True)
@click.option("--output", "-o", type=click.STRING, required=True)
@click.option("--count_by", "-cb", type=click.STRING, multiple=True, required=False)
@click.option("--file_type", "-file_type", type=click.STRING, required=True)
@click.option("--exclude_details", "-ed", type=click.BOOL, default=True, required=False)
def compare_files(file1, file2, output, count_by, exclude_details, file_type):
    # decide which parser to instantiate, GAF or GPAD
    pd.set_option('display.max_rows', 35000)
    df_file1, df_file2, assocs1, assocs2 = get_parser(file1, file2, file_type)

    # if count_by is passed in, generate a report of annotations grouped by count_by
    # count_by is a list.
    generate_group_report(df_file1, df_file2, count_by, file1, file2, output)

    # if exclude_details is False, try to compare files line by line and find matches.
    compare_associations(assocs1, assocs2, exclude_details, output)

def generate_group_report(df_file1, df_file2, count_by, file1, file2, output):
    if len(count_by) > 0:
        group_by_report_file = open(output + "_group_by_report", "w")
        file1_groups = group_by(df_file1, count_by, file1)
        file2_groups = group_by(df_file2, count_by, file2)
        print(file1_groups.get("counts_by_column"))

        s = "\n\n## GROUP BY SUMMARY \n\n"
        s += "This report generated on {}\n\n".format(datetime.date.today())
        s += "  * Group By Columns: " + str(count_by) + "\n"
        s += "  * Compared Files: " + file1 + ", " + file2 + "\n"

        print(s)
        print(("Counts By Column: %s file: " % "\n" + file1) + "\n" + str(file1_groups.get("counts_per_column")) + "\n")
        print(("Counts By Column: %s file: " % "\n" + file2) + "\n" + str(file2_groups.get("counts_per_column")) + "\n")

        for grouped_item in file1_groups['grouped_reports']:
            print(file1_groups['filename'])
            group_by_report_file.write(file1_groups['filename'])
            print(grouped_item)
            group_by_report_file.write(str(grouped_item))
            print("\n")

        for grouped_item in file2_groups['grouped_reports']:
            print(file2_groups['filename'])
            group_by_report_file.write(file2_groups['filename'])
            print(grouped_item)
            group_by_report_file.write(str(grouped_item))

        group_by_report_file.close()

def compare_associations(assocs1, assocs2, exclude_details, output):
    compare_report_file = open(output + "_compare_report", "w")
    processed_lines = 0
    exact_matches = 0
    close_matches = 0

    report = Report()
    for association in assocs1:
        max_match_score = 0
        processed_lines = processed_lines + 1
        if not exclude_details:
            for target in assocs2:
                match_score = 0
                if association.negated != target.negated:
                    continue
                if association.subject.id == target.subject.id and association.object.id == target.object.id:
                    match_score = 1
                    if sorted(str(q).upper() for q in association.qualifiers) == \
                            sorted(str(q).upper() for q in target.qualifiers):
                        match_score = 2
                        if association.evidence.type == target.evidence.type:
                            match_score = 3
                            if sorted(str(w).upper() for w in association.evidence.with_support_from) == \
                                    sorted(str(w).upper() for w in target.evidence.with_support_from):
                                match_score = 4
                                if sorted(
                                        str(r).upper() for r in association.evidence.has_supporting_reference) == \
                                        sorted(str(r).upper() for r in target.evidence.has_supporting_reference):
                                    match_score = 5
                if match_score > max_match_score:
                    max_match_score = match_score
            if max_match_score > 4:
                exact_matches = exact_matches + 1
            elif 1 > max_match_score < 5:
                close_matches = close_matches + 1
                report.add_association(association)
                report.n_lines = report.n_lines + 1
                report.warning(association.source_line, qc.ResultType.WARNING,
                               "line from file1 only has CLOSE match in file2", "")
            else:
                report.add_association(association)
                report.n_lines = report.n_lines + 1
                report.error(association.source_line, qc.ResultType.ERROR,
                             "line from file1 has NO match in file2", "")
    if not exclude_details:
        md_report = markdown_report(report, exact_matches, close_matches, processed_lines)
        print(md_report)
        compare_report_file.write(md_report)
        compare_report_file.close()

def markdown_report(report, exact_matches, close_matches, processed_lines):

    json = report.to_report_json()

    s = "\n\n## DIFF SUMMARY\n\n"
    s += "This report generated on {}\n\n".format(datetime.date.today())
    s += "  * Total Unmatched Associations: {}\n".format(json["associations"])
    s += "  * Total Lines Compared: " + str(processed_lines) + "\n"
    s += "  * Total Exact matches: " + str(exact_matches) + "\n"
    s += "  * Total Close matches: " + str(close_matches) + "\n\n"

    for (rule, messages) in sorted(json["messages"].items(), key=lambda t: t[0]):
        s += "### {rule}\n\n".format(rule=rule)
        s += "* total: {amount}\n".format(amount=len(messages))
        s += "\n"
        if len(messages) > 0:
            s += "#### Messages\n\n"
        for message in messages:
            obj = " ({})".format(message["obj"]) if message["obj"] else ""
            s += "* {level} - {type}: {message}{obj} -- `{line}`\n".format(level=message["level"],
                                                                           type=message["type"],
                                                                           message=message["message"],
                                                                           line=message["line"],
                                                                           obj=obj)

        return s

def get_parser(file1, file2, file_type):
    if file_type == 'gpad':
        gpad_parser_1 = GpadParser()
        gpad_parser_2 = GpadParser()
        assocs1 = gpad_parser_1.parse(file1, skipheader=True)
        assocs2 = gpad_parser_2.parse(file2, skipheader=True)
        df_file1 = read_gpad_csv(file1)
        df_file2 = read_gpad_csv(file2)
    else:
        gaf_parser_1 = GafParser()
        gaf_parser_2 = GafParser()
        assocs1 = gaf_parser_1.parse(file1, skipheader=True)
        assocs2 = gaf_parser_2.parse(file2, skipheader=True)

        df_file1 = read_gaf_csv(file1)
        df_file2 = read_gaf_csv(file2)
    return df_file1, df_file2, assocs1, assocs2

def read_gaf_csv(filename):
    ecomapping = ecomap.EcoMap()
    data_frame = pd.read_csv(filename,
                             comment='!',
                             sep='\t',
                             header=None,
                             na_filter=False,
                             names=["DB",
                                    "DB_Object_ID",
                                    "DB_Object_Symbol",
                                    "Qualifier",
                                    "GO_ID",
                                    "DB_Reference",
                                    "Evidence_code",
                                    "With_or_From",
                                    "Aspect",
                                    "DB_Object_Name",
                                    "DB_Object_Synonym",
                                    "DB_Object_Type,"
                                    "Taxon",
                                    "Date",
                                    "Assigned_By",
                                    "Annotation_Extension",
                                    "Gene_Product_Form_ID"]).fillna("")
    for eco_code in ecomapping.mappings():
        for ev in data_frame['Evidence_code']:
            if eco_code[2] == ev:
                data_frame['Evidence_code'] = data_frame['Evidence_code'].replace([eco_code[2]],
                                                                                ecomapping.ecoclass_to_coderef(
                                                                                eco_code[2])[0])
    return data_frame

def read_gpad_csv(filename):
    ecomapping = ecomap.EcoMap()
    data_frame = pd.read_csv(filename,
                             comment='!',
                             sep='\t',
                             header=None,
                             na_filter=False,
                             names=["DB",
                                    "DB_Object_ID",
                                    "Relation",
                                    "Ontology_Class_ID",
                                    "Reference",
                                    "Evidence_type",
                                    "With_or_From",
                                    "Interacting_taxon_ID",
                                    "Date",
                                    "Assigned_by",
                                    "Annotation_Extensions",
                                    "Annotation_Properties"]).fillna("")
    for eco_code in ecomapping.mappings():
        for ev in data_frame['Evidence_type']:
            if eco_code[2] == ev:
                data_frame['Evidence_type'] = data_frame['Evidence_type'].replace([eco_code[2]],
                                                                                ecomapping.ecoclass_to_coderef(eco_code[2])[0])
    return data_frame

def group_by(data_frame, count_by, file):
    stats = {}
    grouped_reports = []
    stats['filename'] = file
    stats['total_rows'] = data_frame.shape[0]
    stats['counts_per_column'] = data_frame.nunique()
    if len(count_by) > 0:
        for grouper in count_by:
            stats['grouper'] = grouper
            grouped_reports.append(data_frame.groupby(grouper)[grouper].count())
        stats['grouped_reports'] = grouped_reports

    return stats

if __name__ == '__main__':
    compare_files()
<|endoftext|>


<commit_before># coding: utf-8

from __future__ import division, unicode_literals

"""
Some reimplementation of Henkelman's Transition State Analysis utilities,
which are originally in Perl. Additional features beyond those offered by
Henkelman's utilities will be added.

This allows the usage and customization in Python.
"""

__author__ = 'Shyue Ping Ong'
__copyright__ = 'Copyright 2013, The Materials Virtual Lab'
__version__ = '0.1'
__maintainer__ = 'Shyue Ping Ong'
__email__ = 'ongsp@ucsd.edu'
__date__ = '6/1/15'

from pymatgen.io.vaspio import Poscar, Outcar
import numpy as np
import os
import glob
from pymatgen.util.plotting_utils import get_publication_quality_plot

def is_int(s):
    try:
        int(s)
        return True
    except ValueError:
        return False

class NEBAnalysis(object):
    """
    An NEBAnalysis class.
    """

    def __init__(self, outcars, structures):
        """
        Initializes a CINEBAnalysis from Outcar and Structure objects. Use
        the static constructors if you prefer to have these automatically
        generated from a directory of NEB calculations.

        Args:
            outcars ([Outcar]): List of Outcar objects. Note that these have
                to be ordered from start to end along reaction coordinates.
            structures ([Structure]): List of Structures along reaction
                coordinate. Must be same length as outcar.
        """
        if len(outcars) != len(structures):
            raise ValueError("# of Outcars must be same as # of Structures")
        dists = [0]
        prev = structures[0]
        for s in structures[1:]:
            ss = 0
            for site1, site2 in zip(prev, s):
                dist = site2.distance(site1)
                ss += dist ** 2
            prev_dist = dists[-1]
            dists.append(np.sqrt(ss) + prev_dist)
            prev = s

        energies = []
        forces = []
        for i, o in enumerate(outcars):
            o.read_neb()
            energies.append(o.data["energy"])
            if i in [0, len(outcars) - 1]:
                forces.append(0)
            else:
                forces.append(o.data["tangent_force"])
        energies = np.array(energies)
        energies -= energies[0]
        forces = np.array(forces)
        self.r = np.array(dists)
        self.energies = energies
        self.forces = forces
        self._fit_splines()

    def _fit_splines(self):
        # Internal spline fitting procedure based on boundary conditions of
        # tangent force and energies.
        spline_params = {}
        for i in range(0, len(self.r) - 1):
            dr = self.r[i + 1] - self.r[i]
            f1 = self.forces[i] * dr
            f2 = self.forces[i + 1] * dr
            u1 = self.energies[i]
            u2 = self.energies[i + 1]
            fs = f1 + f2
            ud = u2 - u1
            d = u1
            c = -f1
            b = 3 * ud + f1 + fs
            a = -2 * ud - fs
            spline_params[i] = a, b, c, d
        self.spline_params = spline_params

    def get_extrema(self, normalize_rxn_coordinate=True):
        """
        Returns the positions of the extrema along the MEP. Both local
        minimums and maximums are returned.

        Args:
            normalize_rxn_coordinate (bool): Whether to normalize the
                reaction coordinate to between 0 and 1. Defaults to True.

        Returns:
            (min_extrema, max_extrema), where the extrema are given as
            [(x1, y1), (x2, y2), ...].
        """
        x = []
        y = []
        scale = 1 if not normalize_rxn_coordinate else 1 / self.r[-1]
        for i in range(0, len(self.r) - 1):
            dr = self.r[i + 1] - self.r[i]
            a, b, c, d = self.spline_params[i]
            f = np.arange(0, 1.0, 0.01)
            x.extend(self.r[i] + f * dr)
            y.extend(a * f ** 3 + b * f ** 2 + c * f + d)
        min_extrema = []
        max_extrema = []
        for i in range(1, len(x) - 1):
            if y[i] < y[i-1] and y[i] < y[i+1]:
                min_extrema.append((x[i] * scale, y[i]))
            elif y[i] > y[i-1] and y[i] > y[i+1]:
                max_extrema.append((x[i] * scale, y[i]))
        return min_extrema, max_extrema

    def get_plot(self, normalize_rxn_coordinate=True, label_barrier=True):
        """
        Returns the NEB plot. Uses Henkelman's approach of spline fitting
        each section of the reaction path based on tangent force and energies.

        Args:
            normalize_rxn_coordinate (bool): Whether to normalize the
                reaction coordinate to between 0 and 1. Defaults to True.
            label_barrier (bool): Whether to label the maximum barrier.

        Returns:
            matplotlib.pyplot object.
        """
        plt = get_publication_quality_plot(12, 8)
        scale = 1 if not normalize_rxn_coordinate else 1 / self.r[-1]
        all_x = []
        all_y = []
        for i in range(0, len(self.r) - 1):
            dr = self.r[i + 1] - self.r[i]
            a, b, c, d = self.spline_params[i]
            f = np.arange(0, 1.01, 0.01)
            x = self.r[i] + f * dr
            y = (a * f ** 3 + b * f ** 2 + c * f + d) * 1000
            all_x.extend(x * scale)
            all_y.extend(y)
            plt.plot(x * scale, y, 'k-', linewidth=2)

        plt.plot(self.r * scale, self.energies * 1000, 'ro', markersize=10)
        plt.xlabel("Reaction coordinate")
        plt.ylabel("Energy (meV)")
        plt.ylim((np.min(all_y) - 10, np.max(all_y) + 10))
        if label_barrier:
            data = zip(all_x, all_y)
            barrier = max(data, key=lambda d: d[1])
            plt.plot([0, barrier[0]], [barrier[1], barrier[1]], 'k--')
            plt.annotate('%.0f meV' % barrier[1],
                         xy=(barrier[0] / 2, barrier[1] + 1),
                         xytext=(barrier[0] / 2, barrier[1] + 1))
        return plt

    @classmethod
    def from_dir(cls, root_dir):
        """
        Initializes a NEBAnalysis object from a directory of a NEB run.
        Note that the directory must have OUTCARs in all image directories,
        including the terminal points. The termainal OUTCARs are usually
        obtained from relaxation calculations. For the non-terminal points,
        the CONTCAR is read to obtain structures. For terminal points,
        the POSCAR is used. The image directories are assumed to be the only
        directories that can be resolved to integers. E.g., "00", "01", "02",
        "03", "04", "05", "06". The minimum sub-directory structure that can be
        parsed is of the following form (a 5-image example is shown):

        00:
        - POSCAR
        - OUTCAR
        01, 02, 03, 04, 05:
        - CONTCAR
        - OUTCAR
        06:
        - POSCAR
        - OUTCAR

        Args:
            root_dir (str): Path to the root directory of the NEB calculation.

        Returns:
            NEBAnalysis object.
        """
        neb_dirs = []
        for d in os.listdir(root_dir):
            pth = os.path.join(root_dir, d)
            if os.path.isdir(pth) and is_int(d):
                i = int(d)
                neb_dirs.append((i, pth))
        neb_dirs = sorted(neb_dirs, key=lambda d: d[0])
        outcars = []
        structures = []
        for i, d in neb_dirs:
            outcar = glob.glob(os.path.join(d, "OUTCAR*"))
            contcar = glob.glob(os.path.join(d, "CONTCAR*"))
            poscar = glob.glob(os.path.join(d, "POSCAR*"))
            terminal = i == 0 or i == neb_dirs[-1][0]
            if terminal:
                outcars.append(Outcar(outcar[0]))
                structures.append(Poscar.from_file(poscar[0]).structure)
            else:
                outcars.append(Outcar(outcar[0]))
                structures.append(Poscar.from_file(contcar[0]).structure)
        return NEBAnalysis(outcars, structures)<commit_msg>Some efficiency speedups and also, allow specification of relaxation directories for NEB.<commit_after># coding: utf-8

from __future__ import division, unicode_literals

"""
Some reimplementation of Henkelman's Transition State Analysis utilities,
which are originally in Perl. Additional features beyond those offered by
Henkelman's utilities will be added.

This allows the usage and customization in Python.
"""

__author__ = 'Shyue Ping Ong'
__copyright__ = 'Copyright 2013, The Materials Virtual Lab'
__version__ = '0.1'
__maintainer__ = 'Shyue Ping Ong'
__email__ = 'ongsp@ucsd.edu'
__date__ = '6/1/15'

from pymatgen.io.vaspio import Poscar, Outcar
import numpy as np
import os
import glob
from pymatgen.util.plotting_utils import get_publication_quality_plot

def is_int(s):
    try:
        int(s)
        return True
    except ValueError:
        return False

class NEBAnalysis(object):
    """
    An NEBAnalysis class.
    """

    def __init__(self, outcars, structures):
        """
        Initializes a CINEBAnalysis from Outcar and Structure objects. Use
        the static constructors if you prefer to have these automatically
        generated from a directory of NEB calculations.

        Args:
            outcars ([Outcar]): List of Outcar objects. Note that these have
                to be ordered from start to end along reaction coordinates.
            structures ([Structure]): List of Structures along reaction
                coordinate. Must be same length as outcar.
        """
        if len(outcars) != len(structures):
            raise ValueError("# of Outcars must be same as # of Structures")

        # Calculate cumulative mean square distance between structures,
        # which serves as the reaction coordinate.
        r = [0]
        prev = structures[0]
        for st in structures[1:]:
            dists = np.array([s2.distance(s1) for s1, s2 in zip(prev, st)])
            r.append(np.sqrt(np.sum(dists ** 2)) + r[-1])
            prev = st

        energies = []
        forces = []
        for i, o in enumerate(outcars):
            o.read_neb()
            energies.append(o.data["energy"])
            if i in [0, len(outcars) - 1]:
                forces.append(0)
            else:
                forces.append(o.data["tangent_force"])
        energies = np.array(energies)
        energies -= energies[0]
        forces = np.array(forces)
        self.r = np.array(r)
        self.energies = energies
        self.forces = forces
        self._fit_splines()

    def _fit_splines(self):
        # Internal spline fitting procedure based on boundary conditions of
        # tangent force and energies.
        spline_params = {}
        for i in range(0, len(self.r) - 1):
            dr = self.r[i + 1] - self.r[i]
            f1 = self.forces[i] * dr
            f2 = self.forces[i + 1] * dr
            u1 = self.energies[i]
            u2 = self.energies[i + 1]
            fs = f1 + f2
            ud = u2 - u1
            d = u1
            c = -f1
            b = 3 * ud + f1 + fs
            a = -2 * ud - fs
            spline_params[i] = a, b, c, d
        self.spline_params = spline_params

    def get_extrema(self, normalize_rxn_coordinate=True):
        """
        Returns the positions of the extrema along the MEP. Both local
        minimums and maximums are returned.

        Args:
            normalize_rxn_coordinate (bool): Whether to normalize the
                reaction coordinate to between 0 and 1. Defaults to True.

        Returns:
            (min_extrema, max_extrema), where the extrema are given as
            [(x1, y1), (x2, y2), ...].
        """
        x = []
        y = []
        scale = 1 if not normalize_rxn_coordinate else 1 / self.r[-1]
        for i in range(0, len(self.r) - 1):
            dr = self.r[i + 1] - self.r[i]
            a, b, c, d = self.spline_params[i]
            f = np.arange(0, 1.0, 0.01)
            x.extend(self.r[i] + f * dr)
            y.extend(a * f ** 3 + b * f ** 2 + c * f + d)
        min_extrema = []
        max_extrema = []
        for i in range(1, len(x) - 1):
            if y[i] < y[i-1] and y[i] < y[i+1]:
                min_extrema.append((x[i] * scale, y[i]))
            elif y[i] > y[i-1] and y[i] > y[i+1]:
                max_extrema.append((x[i] * scale, y[i]))
        return min_extrema, max_extrema

    def get_plot(self, normalize_rxn_coordinate=True, label_barrier=True):
        """
        Returns the NEB plot. Uses Henkelman's approach of spline fitting
        each section of the reaction path based on tangent force and energies.

        Args:
            normalize_rxn_coordinate (bool): Whether to normalize the
                reaction coordinate to between 0 and 1. Defaults to True.
            label_barrier (bool): Whether to label the maximum barrier.

        Returns:
            matplotlib.pyplot object.
        """
        plt = get_publication_quality_plot(12, 8)
        scale = 1 if not normalize_rxn_coordinate else 1 / self.r[-1]
        all_x = []
        all_y = []
        for i in range(0, len(self.r) - 1):
            dr = self.r[i + 1] - self.r[i]
            a, b, c, d = self.spline_params[i]
            f = np.arange(0, 1.01, 0.01)
            x = self.r[i] + f * dr
            y = (a * f ** 3 + b * f ** 2 + c * f + d) * 1000
            all_x.extend(x * scale)
            all_y.extend(y)
            plt.plot(x * scale, y, 'k-', linewidth=2)

        plt.plot(self.r * scale, self.energies * 1000, 'ro', markersize=10)
        plt.xlabel("Reaction coordinate")
        plt.ylabel("Energy (meV)")
        plt.ylim((np.min(all_y) - 10, np.max(all_y) + 10))
        if label_barrier:
            data = zip(all_x, all_y)
            barrier = max(data, key=lambda d: d[1])
            plt.plot([0, barrier[0]], [barrier[1], barrier[1]], 'k--')
            plt.annotate('%.0f meV' % barrier[1],
                         xy=(barrier[0] / 2, barrier[1] + 1),
                         xytext=(barrier[0] / 2, barrier[1] + 1))
        return plt

    @classmethod
    def from_dir(cls, root_dir, relaxation_dirs=None):
        """
        Initializes a NEBAnalysis object from a directory of a NEB run.
        Note that OUTCARs must be present in all image directories. For the
        terminal OUTCARs from relaxation calculations, you can specify the
        locations using relaxation_dir. If these are not specified, the code
        will attempt to look for the OUTCARs in 00 and 0n directories,
        followed by subdirs "start", "end" or "initial", "final" in the
        root_dir. These are just some typical conventions used
        preferentially in Shyue Ping's MAVRL research group. For the
        non-terminal points, the CONTCAR is read to obtain structures. For
        terminal points, the POSCAR is used. The image directories are
        assumed to be the only directories that can be resolved to integers.
        E.g., "00", "01", "02", "03", "04", "05", "06". The minimum
        sub-directory structure that can be parsed is of the following form (
        a 5-image example is shown):

        00:
        - POSCAR
        - OUTCAR
        01, 02, 03, 04, 05:
        - CONTCAR
        - OUTCAR
        06:
        - POSCAR
        - OUTCAR

        Args:
            root_dir (str): Path to the root directory of the NEB calculation.
            relaxation_dirs (tuple): This specifies the starting and ending
                relaxation directories from which the OUTCARs are read for the
                terminal points for the energies.

        Returns:
            NEBAnalysis object.
        """
        neb_dirs = []
        for d in os.listdir(root_dir):
            pth = os.path.join(root_dir, d)
            if os.path.isdir(pth) and is_int(d):
                i = int(d)
                neb_dirs.append((i, pth))
        neb_dirs = sorted(neb_dirs, key=lambda d: d[0])
        outcars = []
        structures = []

        # Setup the search sequence for the OUTCARs for the terminal
        # directories.
        terminal_dirs = []
        if relaxation_dirs is not None:
            terminal_dirs.append(relaxation_dirs)
        terminal_dirs.append((neb_dirs[0][1], neb_dirs[-1][1]))
        terminal_dirs.append([os.path.join(root_dir, d)
                              for d in ["start", "end"]])
        terminal_dirs.append([os.path.join(root_dir, d)
                              for d in ["initial", "final"]])

        for i, d in neb_dirs:
            outcar = glob.glob(os.path.join(d, "OUTCAR*"))
            contcar = glob.glob(os.path.join(d, "CONTCAR*"))
            poscar = glob.glob(os.path.join(d, "POSCAR*"))
            terminal = i == 0 or i == neb_dirs[-1][0]
            if terminal:
                found = False
                for ds in terminal_dirs:
                    od = ds[0] if i == 0 else ds[1]
                    outcar = glob.glob(os.path.join(od, "OUTCAR*"))
                    if outcar:
                        outcar = sorted(outcar)
                        outcars.append(Outcar(outcar[-1]))
                        found = True
                        break
                if not found:
                    raise ValueError("OUTCAR cannot be found for terminal "
                                     "point %s" % d)
                structures.append(Poscar.from_file(poscar[0]).structure)
            else:
                outcars.append(Outcar(outcar[0]))
                structures.append(Poscar.from_file(contcar[0]).structure)
        return NEBAnalysis(outcars, structures)<|endoftext|>