<commit_before>import sys
import numpy as np
from realtime import *

class Spectra(object):
    """Return a spectra object that can plot the total absorption spectrum or
    circular dichroism spectra for a given system. Can accept x,y,and z RT-TDDFT
    log files"""
    def __init__(self,x=None,y=None,z=None,s='abs',d=150,zp=None,auto=False):
        self.spectra = None
        self.frequency = None

        self.x = None
        self.y = None
        self.z = None

        self.spectra_type = s
        self.damp_const   = d
        self.zero_pad     = zp

        # Load all the RealTime objects
        self.directions = []
        if x:
            self.x = RealTime(x)
            self.directions.append('x')
        if y:
            self.y = RealTime(y)
            self.directions.append('y')
        if z:
            self.z = RealTime(z)
            self.directions.append('z')

        # Enforce consistent data lengths
        self.align_data()
 
        # Do the isotropic fourier transform 
        for q in self.directions:
            self.__dict__[q].fourier_tx(q,self.spectra_type,self.damp_const,
                self.zero_pad,auto=auto)

        self.spectra = np.zeros_like(self.__dict__[self.directions[0]].fourier)
        for q in self.directions:
            self.frequency = self.__dict__[q].frequency
            self.spectra += self.__dict__[q].fourier


    def plot(self,xlim=[0,15],ylim=None):
        toEV = 27.2114 
        import matplotlib.pyplot as plt
        ax = plt.subplot(111)
       
        ax.plot(self.frequency*toEV,self.spectra)
        ax.set_xlim(xlim)
        if not ylim:
            if self.spectra_type == 'abs':
                ax.set_ylim([0,4])
            elif self.spectra_type == 'ecd':
                ax.set_ylim([-400,400])
        plt.show()

    def align_data(self):
        lengths = []
        if self.x:
            lengths.append(self.x.min_length) 
        if self.y:
            lengths.append(self.y.min_length) 
        if self.z:
            lengths.append(self.z.min_length) 
        min_length = min(lengths)
        if self.x:
            self.x.truncate(self.x,min_length)
        if self.y:
            self.y.truncate(self.y,min_length)
        if self.z:
            self.z.truncate(self.z,min_length)

    def peaks(self,number=3,thresh=0.01):
        """ Return the peaks from the Fourier transform"""
        from scipy.signal import argrelextrema as pks
        # find all peak indices [idx], and remove those below thresh [jdx]
        idx = pks(self.spectra,np.greater,order=10)
        jdx = np.where((self.spectra[idx] >= thresh))
        kdx = idx[0][jdx[0]] # indices of peaks matching criteria
        print "First "+str(number)+" peaks (eV) found: "
        if number > len(idx[0]):
            number = len(idx[0])
        for i in xrange(number): 
            print "{0:.2f}".format(self.frequency[kdx][i]*27.2114)
        
            

         


if __name__ == '__main__':
    #spectra = Spectra(x='AuH-x2c_xx',
    #                  y='AuH-x2c_yy',
    #                  z='AuH-x2c_zz',
    #                  s='abs',auto=True)
    spectra = Spectra(x='mg-td',auto=True)
    spectra.x.test()
#    spectra.peaks()
    spectra.plot(xlim=[0,30])
<commit_msg>modify peak finder to pick out major peaks (fix indices)<commit_after>import sys
import numpy as np
from realtime import *

class Spectra(object):
    """Return a spectra object that can plot the total absorption spectrum or
    circular dichroism spectra for a given system. Can accept x,y,and z RT-TDDFT
    log files"""
    def __init__(self,x=None,y=None,z=None,s='abs',d=150,zp=None,auto=False):
        self.spectra = None
        self.frequency = None

        self.x = None
        self.y = None
        self.z = None

        self.spectra_type = s
        self.damp_const   = d
        self.zero_pad     = zp

        # Load all the RealTime objects
        self.directions = []
        if x:
            self.x = RealTime(x)
            self.directions.append('x')
        if y:
            self.y = RealTime(y)
            self.directions.append('y')
        if z:
            self.z = RealTime(z)
            self.directions.append('z')

        # Enforce consistent data lengths
        self.align_data()
 
        # Do the isotropic fourier transform 
        for q in self.directions:
            self.__dict__[q].fourier_tx(q,self.spectra_type,self.damp_const,
                self.zero_pad,auto=auto)

        self.spectra = np.zeros_like(self.__dict__[self.directions[0]].fourier)
        for q in self.directions:
            self.frequency = self.__dict__[q].frequency
            self.spectra += self.__dict__[q].fourier


    def plot(self,xlim=[0,15],ylim=None):
        toEV = 27.2114 
        import matplotlib.pyplot as plt
        ax = plt.subplot(111)
       
        ax.plot(self.frequency*toEV,self.spectra)
        ax.set_xlim(xlim)
        if not ylim:
            if self.spectra_type == 'abs':
                ax.set_ylim([0,4])
            elif self.spectra_type == 'ecd':
                ax.set_ylim([-400,400])
        plt.show()

    def align_data(self):
        lengths = []
        if self.x:
            lengths.append(self.x.min_length) 
        if self.y:
            lengths.append(self.y.min_length) 
        if self.z:
            lengths.append(self.z.min_length) 
        min_length = min(lengths)
        if self.x:
            self.x.truncate(self.x,min_length)
        if self.y:
            self.y.truncate(self.y,min_length)
        if self.z:
            self.z.truncate(self.z,min_length)

    def peaks(self,number=3,thresh=0.01):
        """ Return the peaks from the Fourier transform"""
        from scipy.signal import argrelextrema as pks
        # find all peak indices [idx], and remove those below thresh [jdx]
        idx = pks(self.spectra,np.greater,order=10)
        jdx = np.where((self.spectra[idx] >= thresh))
        kdx = idx[0][jdx[0]] # indices of peaks matching criteria
        print "First "+str(number)+" peaks (eV) found: "
        if number > len(kdx):
            number = len(kdx)
        for i in xrange(number): 
            print "{0:.2f}".format(self.frequency[kdx][i]*27.2114)
        
            

         


if __name__ == '__main__':
    #spectra = Spectra(x='AuH-x2c_xx',
    #                  y='AuH-x2c_yy',
    #                  z='AuH-x2c_zz',
    #                  s='abs',auto=True)
    spectra = Spectra(x='mg-td',auto=True)
    spectra.x.test()
    spectra.peaks()
    spectra.plot(xlim=[0,1500])

{'fraction_non_alphanumeric': 0.11921437271619975, 'fraction_numerical': 0.014616321559074299, 'mean_word_length': 1.8388072601555747, 'pattern_counts': {'":': 0, '<': 3, '<?xml version=': 0, '>': 7, 'https://': 0, 'lorem ipsum': 0, 'www.': 0, 'xml': 0}, 'pii_count': 0, 'substrings_counts': 0, 'word_list_counts': {'cursed_substrings.json': 0, 'profanity_word_list.json': 0, 'sexual_word_list.json': 0, 'zh_pornsignals.json': 0}}
{'dir': 'git-commits-cleaned', 'id': '5948535', 'n_tokens_mistral': 2256, 'n_tokens_neox': 2152, 'n_words': 443}