<commit_before># Version : nov. 18th, 11:30
""" python plot_BSE.py <JobName> <PathToJobFolder>
The Job Folder normally contains the SAVE folder as well as <JobName> folder. Path should be at least '.'.
Yambo must be able to run, because ypp is called multiple times.
"""

import matplotlib
matplotlib.use('Agg') # prevents crashes if no X server present
from yambopy import *
from qepy import *
import json
import matplotlib.pyplot as plt
import numpy as np
from math import sqrt
import sys

jobname = sys.argv[1]
path = sys.argv[2]

print 'Preparing JSON file. Calling ypp ...'
### Creating the 'absorptionspectra.json' file
y = YamboOut(folder=path,save_folder=path)
# Args : name of job, SAVE folder path, folder where job was run path
a = YamboBSEAbsorptionSpectra(jobname,path=path)
# Get excitons values
a.get_excitons(min_intensity=0.0005,max_energy=6,Degen_Step=0.01)
# Get excitonic WFs
a.get_wavefunctions(Degen_Step=0.01,repx=range(-1,2),repy=range(-1,2),repz=range(1))
# Write .json file
a.write_json(filename=path+'_'+jobname)

### Loading data from .json file
f = open(path+'_'+jobname+'.json')
data = json.load(f)
f.close()
print 'JSON file prepared and loaded.'

### Plotting the absorption spectra
print 'Absorption spectra ...'
# BSE spectra
plt.plot(data['E/ev[1]'], data['EPS-Im[2]'],label='BSE',lw=4)
# IP spectra
plt.plot(data['E/ev[1]'], data['EPSo-Im[4]'],label='IP',lw=4)
# Axes : lines for exciton energies, labels
for n,exciton in enumerate(data['excitons']):
    plt.axvline(exciton['energy'])
plt.xlabel('$\omega$ (eV)')
plt.gca().yaxis.set_major_locator(plt.NullLocator())
plt.legend()
#plt.draw()
#plt.show()
plt.savefig(path+'_'+jobname+'_abs.png', bbox_inches='tight')

### Lattice : necessary to determine bounds of exciton weights plot
print 'Reciprocal lattice ...'
lat = np.array(data['lattice'])
rlat = rec_lat(lat)
x,y,z = np.array(rlat )
cut=.65
xmin,ymin,_ = -(x+y)*cut
xmax,ymax,_ = +(x+y)*cut


### Plot excitons
print 'Excitons weights ...'
fig = plt.figure(figsize=(16,8))
# Issue if dvipng is not installed
#plt.rc('text', usetex=True)
#plt.rc('font', family='serif',serif="Computer Modern Roman",size=20)

sorted_excitons = sorted(data['excitons'],key=lambda x: x['energy'])

#colormap
cmap = plt.get_cmap("viridis")

for n,exciton in enumerate(sorted_excitons[:8]):
    ax = plt.subplot(2,4,n+1)
    ax.set_aspect('equal', 'datalim')
    w = np.array(exciton['weights'])
    qpt = np.array(exciton['qpts'])
    ax.scatter(qpt[:,0], qpt[:,1], marker='H', s=2.5, color=[cmap(sqrt(c)) for c in w], label="e: %lf"%exciton['energy'])
    plt.xlim([xmin,xmax])
    plt.ylim([ymin,ymax])
    plt.gca().yaxis.set_major_locator(plt.NullLocator())
    plt.gca().xaxis.set_major_locator(plt.NullLocator())
    ax.legend(prop={'size':12})

#plt.draw()
#plt.show()
plt.savefig(path+'_'+jobname+'_ex.png', bbox_inches='tight')

print 'Done.'<commit_msg>Buggy update before tests & cleanup<commit_after># Copyright (C) 2015 <NAME>
# All rights reserved.
#
# This file is part of yambopy

'''
Plot the excitonic weights on the brillouin zone using ypp
Plot the absorption spectrum with the excitonic energies
'''

import matplotlib
matplotlib.use('Agg') # prevents crashes if no X server present
from yambopy import *
from qepy import *
import json
import matplotlib.pyplot as plt
import numpy as np
from math import sqrt, ceil
import argparse

parser = argparse.ArgumentParser(description='Plot excitonic weights in BZ or excitonic eigenenergies on optical spectrum using ypp.')
parser.add_argument('folder', help='Folder with SAVE and run.')
parser.add_argument('jobname', help='Job from which excitons are of interest.')
parser.add_argument('-ie','--intexc' , help='Minimum intensity for excitons to be considered bright', default=0.005)
parser.add_argument('-de','--degenexc', help='Energy threshold under which different peaks are merged (eV)', default=0.01)
parser.add_argument('-me','--maxexc', help='Energy threshold after which excitons are not read anymore (eV)', default=4.0)
parser.add_argument('-c','--cut', help='Set the limits of the BZ plot.', default=0.2)
parser.add_argument('-ny','--noypp', help='Skips the ypp calls and the .json file building.', action='store_false')
parser.add_argument('-nw','--noweight', help='Skips the plot of the weights in the BZ.', action='store_false')
parser.add_argument('-ns','--nospectrum', help='Skips the plot of the spectrum with the excitons eigenenergies.', action='store_false')

jobname = args.jobname
path    = args.folder 
cut     = args.cut
exc_int   = args.intexc
exc_degen = args.degenexc
exc_max_E = args.maxexc

def get_var(dictionary,variables):
    """
    To have compatibility with different version of yambo
    We provide a list of different possible tags
    """
    for var in variables:
        if var in dictionary:
            return dictionary[var]
    raise ValueError( 'Could not find the variables %s in the output file'%str(variables) )

if args.noypp:
    print 'Preparing JSON file. Calling ypp ...'
    ### Creating the 'absorptionspectra.json' file
    y = YamboOut(folder=path,save_folder=path)
    # Args : name of job, SAVE folder path, folder where job was run path
    a = YamboBSEAbsorptionSpectra(jobname,path=path)
    # Get excitons values
    a.get_excitons(min_intensity=exc_int,max_energy=exc_max_E,Degen_Step=exc_degen)
    # Get excitonic WFs
    a.get_wavefunctions(Degen_Step=exc_degen,repx=range(-1,2),repy=range(-1,2),repz=range(1))
    # Write .json file
    a.write_json(filename=path+'_'+jobname)
else:
    print 'YPP call disabled.'

### Loading data from .json file
f = open(path+'_'+jobname+'.json')
data = json.load(f)
f.close()
print 'JSON file loaded.'

### Plotting the absorption spectra
print 'Absorption spectra ...'
# BSE spectra
plt.plot(get_var(data,['E/ev','E/ev[1]']), get_var(data,['EPS-Im[2]' ]),label='BSE',lw=4)
# IP spectra
plt.plot(get_var(data,['E/ev','E/ev[1]']), get_var(data,['EPSo-Im[4]']),label='IP',lw=4)
# Axes : lines for exciton energies, labels
nexcitons = len(data['excitons'])
print 'Number of excitons: ', nexcitons 
for n,exciton in enumerate(data['excitons']):
    plt.axvline(exciton['energy'])
plt.xlabel('$\omega$ (eV)')
plt.ylabel('Intensity (arb. units)')
plt.gca().yaxis.set_major_locator(plt.NullLocator())
plt.legend()
#plt.draw()
#plt.show()
plt.savefig(path+'_'+jobname+'_abs.png', bbox_inches='tight')
print path+'_'+jobname+'_abs.png'

### Lattice : necessary to determine bounds of exciton weights plot
print 'Reciprocal lattice ...'
lat = np.array(data['lattice'])
rlat = rec_lat(lat)
x,y,z = np.array(rlat )
xmin,ymin,_ = -(x+y)*cut
xmax,ymax,_ = +(x+y)*cut


### Plot excitons
print 'Excitons weights ...'
nx = int(ceil(sqrt(nexcitons)))
ny = int(ceil(nexcitons*1.0/nx))
print "cols:",nx
print "rows:",ny
cmap = plt.get_cmap("gist_heat_r")

fig = plt.figure(figsize=(nx*4,ny*4))
# Issue if dvipng is not installed
#plt.rc('text', usetex=True)
#plt.rc('font', family='serif',serif="Computer Modern Roman",size=20)

sorted_excitons = sorted(data['excitons'],key=lambda x: x['energy'])

for n,exciton in enumerate(sorted_excitons):
    #get data
    w   = np.array(exciton['weights'])
    qpt = np.array(exciton['qpts'])

    #plot
    ax = plt.subplot(ny,nx,n+1)
    ax.scatter(qpt[:,0], qpt[:,1], s=20, c=w, marker='H', cmap=cmap, lw=0, label="e: %lf"%exciton['energy'])

    # axis
    plt.xlim([-cut,cut])
    plt.ylim([-cut,cut])
    plt.gca().yaxis.set_major_locator(plt.NullLocator())
    plt.gca().xaxis.set_major_locator(plt.NullLocator())
    ax.set_aspect('equal')

#plt.draw()
#plt.show()
plt.savefig(path+'_'+jobname+'_ex.png', bbox_inches='tight')
print 'path'+'_'+jobname+'_ex.png'
print 'Done.'

print 'Done.'
{'fraction_non_alphanumeric': 0.16569428238039674, 'fraction_numerical': 0.011150006482561908, 'mean_word_length': 5.271544715447154, 'pattern_counts': {'":': 0, '<': 7, '<?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': '2928693', 'n_tokens_mistral': 2937, 'n_tokens_neox': 2816, 'n_words': 671}