#!/usr/bin/env python
# coding: utf-8
# In[1]:
"""Task that generates and plots uv/vis Spectra for solutes in solvent.
Also contains routines to calculate spectral warp parameters and RGB colours from spectra"""
# # Main Routine
# In[1]:
import numpy as np
from matplotlib import pyplot
wavelength_eV_conv=1239.84193
[docs]class SpectraTask:
def __init__(self,**kwargs):
self.wrapper = None
self.script_settings = None
args = self.make_parser().parse_args("")
for arg in vars(args):
setattr(self,arg,None)
#from amp import Amp
[docs] def main(args,fig=None,ax=None,plotlabel=None,rgb=np.array((0.0,0.0,0.0))):
"""
Main routine for plotting a spectrum. Capable of applying Gaussian broadening to a stick spectrum to
produce a broadened spectrum, and also of applying spectral warping to shift/scale the stick spectrum.
*Arguments*:
args: namespace or class
All arguments for the job - see documentation below.
fig, ax: matplotlib objects
Figure and axis objects for matplotlib.pyplot. Initialised anew if they have value None on entry.
plotlabel:
Label to add the key for this dataset.
rgb:
Colour of the line/points for this dataset. If set to (-1.0,-1.0,-1.0), the RGB colour will be
calculated from the spectrum.
*Returns*:
broad_spectrum,spec_plot,fig,ax,all_transition_origins
*Output*:
Plots the spectrum to 'args.output' as a png file if requested.
"""
from ase.calculators.onetep import Onetep;
from ase.calculators.nwchem import NWChem;
from esteem.wrappers import onetep
from esteem.wrappers import nwchem
from ase.io import Trajectory
import os
#validate_args(args)
if args.inputformat.lower()=="nwchem":
calc=NWChem(label="temp")
wrapper=nwchem.NWChemWrapper()
read_excitations = wrapper.read_excitations
elif args.inputformat.lower()=="onetep":
calc=Onetep(label="temp")
wrapper=onetep.OnetepWrapper()
read_excitations = wrapper.read_excitations
#elif args.inputformat.lower()=="amp":
# calc = []
# for f in args.files:
# calc.append(Amp.load(f))
# args.files = []
else:
raise Exception('Unrecognised input format. The value of inputformat was: {}'.format(args.inputformat))
all_excitations = []
all_transition_origins = []
if args.trajectory is not None:
for trajfile in args.trajectory:
traj = Trajectory(trajfile)
for frame in traj:
E = []
i=0
excitations = []
for c in calc:
i = i + 1
frame.set_calculator(c)
E.append(frame.get_potential_energy())
if len(E)>1:
excitations.append(np.array((i,E[-1]-E[0],1.0)))
all_excitations.append(excitations)
for f in args.files:
if not os.path.isfile(f):
raise OSError(f'Could not read file {f}')
#try:
label = f.replace(".nwo","")
label = label.replace(".out","")
calc.label = label
#calc.read(label)
read_excitations(calc)
all_excitations.append(calc.results['excitations'])
if args.inputformat.lower()=="nwchem":
all_transition_origins.append(calc.results['transition_origins'])
#calc.read(label)
#except:
# print(f'Reading excitations failed for: {f}')
stick_spectrum = np.array(all_excitations)
wv = np.arange(args.wavelength[0],args.wavelength[1],args.wavelength[2])
broad_spectrum = np.zeros((len(wv),3))
for i,x in enumerate(wv):
broad_spectrum[i,:] = np.array((x,spectral_value(x,stick_spectrum,args),wavelength_eV_conv/x))
if args.renorm:
if isinstance(args.renorm,float):
broad_spectrum[:,1] = broad_spectrum[:,1] * args.renorm
else: # renormalise so highest peak has maximum value 1
broad_spectrum[:,1] = broad_spectrum[:,1] / np.amax(broad_spectrum[:,1])
if (rgb == np.array((-1.0,-1.0,-1.0))).all():
rgb = RGB_colour(stick_spectrum,args)
# Set up and make plot
if fig is None and ax is None:
fig, ax = pyplot.subplots()
ax.set_xlabel('Wavelength (nm)')
ax.set_ylabel('Absorbtion (arb)')
#x_ticks = np.arange(args.wavelength[0], args.wavelength[1], 25)
#pyplot.xticks(x_ticks)
#spec_plot = plot_sticks(stick_spectrum,rgb,fig,ax,all_transition_origins)
spec_plot = plot_spectrum(broad_spectrum,rgb,fig,ax,plotlabel)
if args.output is not None:
fig.savefig(args.output)
return broad_spectrum,spec_plot,fig,ax,all_transition_origins
def make_parser(self):
# Parse command line values
main_help = ('Spectra.py: Generates optical spectra from pre-existing excited state \n'+
'calculation output files')
epi_help = ('')
from argparse import ArgumentParser
parser = ArgumentParser(description=main_help,epilog=epi_help)
parser.add_argument('--files','-f',nargs='*',default=[],type=str,help='List of output files from which to generate spectrum')
parser.add_argument('--trajectory','-t',nargs='*',default=None,type=str,help='Trajectory file if using a calculator')
parser.add_argument('--wavelength','-w',default=[200.0,800.0,1.0],nargs=3,type=float,help='List of minimum, maximum wavelength and step')
parser.add_argument('--warp_params','-s',default=[0.0],nargs='?',type=float,help='Parameters for spectral warp - meaning differs depending on warp_scheme argument')
parser.add_argument('--warp_scheme',default="beta",nargs='?',type=str,help='Scheme for spectral warp (beta,alphabeta or betabeta)')
parser.add_argument('--broad','-b',default=0.05,type=float,help='Broadening for spectrum, in eV')
parser.add_argument('--renorm','-R',default=True,type=bool,help='Renormalise spectrum so that highest peak is at 1.0')
parser.add_argument('--inputformat','-i',default="onetep",type=str,help='Expected format of output files')
parser.add_argument('--output','-o',default=None,type=str,help='File to write final plots to')
parser.add_argument('--illuminant','-I',default='D65_illuminant.txt',type=str,help='Spectrum of illuminant for calculating RGB colour')
parser.add_argument('--XYZresponse','-X',default='XYZ_response.txt',type=str,help='Response spectrum X, Y and Z functions for calculating RGB colour')
# For use in Drivers only, for setting up spectral warping and colouring plots
parser.add_argument('--exc_suffix','-e',default="exc",nargs='?',type=str,help='Suffix of excitation calculation directories')
parser.add_argument('--warp_broad',default=None,nargs='?',type=float,help='Broadening to apply when calculating warp parameters')
parser.add_argument('--warp_origin_ref_peak_range',default=None,nargs='?',type=str,help='Wavelength range in origin spectrum to be included in finding peaks for spectral warping')
parser.add_argument('--warp_origin_prefix',default="is_tddft",nargs='?',type=str,help='Prefix where spectral warp origin output files will be found')
parser.add_argument('--warp_dest_ref_peak_range',default=None,nargs='?',type=str,help='Wavelength range in destination spectrum to be included in finding peaks for spectral warping')
parser.add_argument('--warp_dest_prefix',default="is_tddft",nargs='?',type=str,help='Prefix where spectral warp destination output files will be found')
parser.add_argument('--warp_inputformat',default="nwchem",nargs='?',type=str,help='Expected format of output files for calculating spectral warp parameters')
parser.add_argument('--warp_files',default="tddft.nwo",nargs='?',type=str,help='Files for calculating spectral warp parameters')
parser.add_argument('--line_colours',default=None,nargs='?',type=str,help='Line colours for final plot')
parser.add_argument('--merge_solutes',default={},type=dict,help='Dictionary of solutes that should be merged into each key')
return parser
def validate_args(args):
default_args = make_parser().parse_args("")
for arg in vars(args):
if arg not in default_args:
raise Exception(f"Unrecognised argument '{arg}'")
# # Helper Routines
# In[ ]:
[docs] def find_spectral_warp_params(args,dest_spectrum,origin_spectrum,arrow1_pos,arrow2_pos):
"""
Finds spectral warping parameters via a range of schemes. See spectra documentation page
for more detail.
*Arguments*
dest_spectrum: numpy array of floats
Contains the spectral warp 'destination' spectrum (usually a high level of theory that can only be afforded for the solute molecule)
origin_spectrum: numpy array of floats
Contains the spectral warp 'origin' spectrum (usually a cheap level of theory, same as in the Clusters job)
args.warp_scheme: str
The scheme used for the warping. Allowed values: 'beta', 'alphabeta', 'betabeta'
args.warp_origin_ref_peak_range: list of 2 floats
Peak range searched when looking for 'reference' peaks in the origin spectrum for spectral warping.
args.warp_dest_ref_peak_range: list of 2 floats
Peak range searched when looking for 'reference' peaks in the destination spectrum for spectral warping.
args.broad:
Energy gaussian broadening applied to the stick spectra, useful to merge peaks that you want to treat
as one peak in spectral warping.
*Returns*
[beta], [alpha, beta], [beta1, beta2, omega1o, omega2o]: floats describing spectral warp parameters
Also sets arrow1_pos, arrow2_pos for use in spectral warp plots.
"""
if args.warp_scheme == None:
return
hc = wavelength_eV_conv
# Mask anything outside the allowed range
dest_spectrum_masked = mask_spectrum(dest_spectrum,
args.warp_dest_ref_peak_range[0],args.warp_dest_ref_peak_range[1])
origin_spectrum_masked = mask_spectrum(origin_spectrum,
args.warp_origin_ref_peak_range[0],args.warp_origin_ref_peak_range[1])
# Find maximum value of remaining spectra
dest_maxloc = np.argmax(dest_spectrum_masked[:,1])
dest_lambdamax = dest_spectrum[dest_maxloc,0]
origin_maxloc = np.argmax(origin_spectrum_masked[:,1])
origin_lambdamax = origin_spectrum[origin_maxloc,0]
# Set up position of (first) arrow on plot
arrow1_pos[0] = origin_lambdamax
arrow1_pos[1] = origin_spectrum[origin_maxloc,1]
arrow1_pos[2] = dest_lambdamax - origin_lambdamax
arrow1_pos[3] = 0
# Convert to energy
omega1d = hc / dest_lambdamax
omega1o = hc / origin_lambdamax
print(f'lambda(max) of origin spectrum = {origin_lambdamax}')
print(f'lambda(max) of dest spectrum = {dest_lambdamax}')
# Set acceptable ranges of result
betamin = -1.50
betamax = 2.4
alphamin = 0.3
alphamax = 1.8
if args.warp_scheme == 'beta':
beta = omega1d - omega1o
# Report and check results
print(f'Spectral warp beta param from origin to dest = {beta} eV')
if beta<betamin or beta>betamax:
raise Exception(f'beta parameter {beta} out of allowed range ({betamin} to {betamax}). Stopping.')
return [beta]
# Get spectrum masked again
dest_spectrum_tmp = mask_spectrum(dest_spectrum,
args.warp_dest_ref_peak_range[0],args.warp_dest_ref_peak_range[1])
origin_spectrum_tmp = mask_spectrum(origin_spectrum,
args.warp_origin_ref_peak_range[0],args.warp_origin_ref_peak_range[1])
# Mask section +- 4*broad from omega1d of dest spectrum
lower = hc / (omega1d+4*args.broad)
upper = hc / (omega1d-4*args.broad)
dest_spectrum_masked = mask_spectrum(dest_spectrum_tmp,lower,upper,swap=True)
# Mask section +- 4*broad from omega1o of origin spectrum
lower = hc / (omega1o+4*args.broad)
upper = hc / (omega1o-4*args.broad)
origin_spectrum_masked = mask_spectrum(origin_spectrum_tmp,lower,upper,swap=True)
# Find maximum value of remaining spectra
dest_maxloc = np.argmax(dest_spectrum_masked[:,1])
dest_lambdamax = dest_spectrum[dest_maxloc,0]
origin_maxloc = np.argmax(origin_spectrum_masked[:,1])
origin_lambdamax = origin_spectrum[origin_maxloc,0]
# Set up position of second arrow on plot
arrow2_pos[0] = origin_lambdamax
arrow2_pos[1] = origin_spectrum[origin_maxloc,1]
arrow2_pos[2] = dest_lambdamax - origin_lambdamax
arrow2_pos[3] = 0
omega2d = hc / dest_lambdamax
omega2o = hc / origin_lambdamax
print(f'second peak lambda(max) of origin spectrum = {origin_lambdamax}')
print(f'second peak lambda(max) of dest spectrum = {dest_lambdamax}')
if args.warp_scheme == 'avgebeta':
beta = ((omega1d - omega1o) + (omega2d - omega2o))*0.5
print(f'Spectral warp beta param from origin to dest = {beta} eV')
if beta<betamin or beta>betamax:
raise Exception(f'beta parameter {beta} out of allowed range ({betamin} to {betamax}). Stopping.')
return [beta]
if args.warp_scheme == 'alphabeta':
alpha = (omega1d - omega2d) / (omega1o - omega2o)
beta = omega2d - alpha*omega2o
print(f'Spectral warp alpha, beta params from origin to dest = {alpha} {beta} eV')
if beta<betamin or beta>betamax:
raise Exception(f'beta parameter {beta} out of allowed range ({betamin} to {betamax}). Stopping.')
if alpha<alphamin or alpha>alphamax:
raise Exception(f'alpha parameter {alpha} out of allowed range ({alphamin} to {alphamax}). Stopping.')
return [alpha, beta]
if args.warp_scheme == 'betabeta':
beta1 = omega1d - omega1o
beta2 = omega2d - omega2o
# Ensure omega1o < omega2o, swap if not.
if (omega1o > omega2o):
tmp = omega2o; omega2o = omega1o; omega1o = tmp
tmp = beta2; beta2 = beta1; beta1 = tmp
print(f'Spectral warp beta1, beta2, omega1, omega2 params from origin to dest = {beta1} {beta2} {omega1o} {omega2o} eV')
if beta1<betamin or beta1>betamax:
raise Exception(f'beta1 parameter {beta1} out of allowed range ({betamin} to {betamax}). Stopping.')
if beta2<betamin or beta2>betamax:
raise Exception(f'beta2 parameter {beta2} out of allowed range ({betamin} to {betamax}). Stopping.')
return [beta1, beta2, omega1o, omega2o]
raise Exception(f'Unrecognised warp scheme: {args.warp_scheme}')
def spectral_warp(exc,args):
exc_w = exc
if args.warp_scheme == 'beta' or args.warp_scheme == 'avgebeta':
exc_w = exc + args.warp_params[0]
if args.warp_scheme == 'alphabeta':
exc_w = exc*args.warp_params[0] + args.warp_params[1]
if args.warp_scheme == 'betabeta':
beta1, beta2, omega1, omega2 = args.warp_params
if (exc<omega1):
beta = beta1
elif (exc>omega2):
beta = beta2
elif (abs(exc-omega2)<abs(exc-omega1)):
beta = beta2
else:
beta = beta1
exc_w = exc + beta
return exc_w
def mask_spectrum(spectrum,lower,upper,swap=False):
spectrum_masked = spectrum.copy()
mask = (spectrum_masked[:,0]>lower)&(spectrum_masked[:,0]<upper)
if swap:
mask = np.logical_not(mask)
spectrum_masked[:,1] = spectrum_masked[:,1]*mask
return spectrum_masked
def spectral_value(wavelength,spectrum,args):
wav_in_eV=wavelength_eV_conv/wavelength
abs_value=0.0
p = 1/(args.broad*np.sqrt(2*np.pi))
for snap in spectrum:
for exc in snap:
abs_value=abs_value+p*exc[2]*np.exp(-0.5 *
((wav_in_eV-spectral_warp(exc[1],args))/args.broad)**2)
return abs_value
def total_source_val(wavelength,intensity,spectrum,args):
spectral_contribution=spectral_value(wavelength,spectrum,args)
path_length=10
exponential_term=np.exp(-spectral_contribution*path_length)
source_val=exponential_term*intensity
return source_val
[docs]def RGB_colour(spectrum,args):
"""
Finds the RGB colour corresponding to a given absorption spectrum and illumination.
spectrum: numpy array
Spectrum for which to find the colour
args: namespace or class
Full set of arguments for the spectra task. Relevant to this routine are:
``args.XYZresponse`` and ``args.illuminant`` which supply the Color Space XYZ response
spectra and the illuminant spectrum, respectively.
These must be on the same wavelength scale as the spectrum.
"""
xyz_funcs=np.genfromtxt(args.XYZresponse)
illu=np.genfromtxt(args.illuminant)
x_int=0.0
y_int=0.0
z_int=0.0
step_length=(illu[1,0]-illu[0,0])
h=step_length/2.0
counter = 0
# Numerical integration in a totally unpythonic way (to be fixed - must be a SciPy routine for this)
for x in illu:
wavelength=x[0]
temp_val=total_source_val(wavelength,x[1],spectrum,args)
if counter==0:
x_int = x_int+temp_val*xyz_funcs[counter,1]
y_int = y_int+temp_val*xyz_funcs[counter,2]
z_int = z_int+temp_val*xyz_funcs[counter,3]
else:
x_int = x_int+2.0*temp_val*xyz_funcs[counter,1]
y_int = y_int+2.0*temp_val*xyz_funcs[counter,2]
z_int = z_int+2.0*temp_val*xyz_funcs[counter,3]
counter=counter+1
counter=counter-1
wavelength=illu[counter,0]
temp_val=total_source_val(wavelength,illu[counter,1],spectrum,args)
x_int = (x_int-1.0*temp_val*xyz_funcs[counter,1])*h
y_int = (y_int-1.0*temp_val*xyz_funcs[counter,2])*h
z_int = (z_int-1.0*temp_val*xyz_funcs[counter,3])*h
# Conversion (linear transformation) from Color Space X, Y, Z values to R, G, B
r_int=0.41847*x_int-0.15866*y_int-0.082835*z_int
g_int=-0.091169*x_int+0.25243*y_int+0.015708*z_int
b_int=0.00092090*x_int-0.0025498*y_int+0.17860*z_int
# Black magic (need to check ref for why 384.0)
normalisation=384.0/(r_int+g_int+b_int)
if normalisation*r_int>255.0:
normalisation=255.0/r_int
if normalisation*g_int>255.0:
normalisation=255.0/g_int
if normalisation*b_int>255.0:
normalisation=255.0/b_int
rgb = np.array((r_int,g_int,b_int))*normalisation
print('R value:')
print(rgb[0])
print('G value:')
print(rgb[1])
print('B value:')
print(rgb[2])
return rgb
def plot_spectrum(broad_spectrum,rgb,fig,ax,label):
# Plot data
spec = ax.plot(broad_spectrum[:,0], broad_spectrum[:,1], 'r-',label=label)
pyplot.setp(spec,color=rgb*(1.0/256.0))
if label is not None:
ax.legend()
return spec
def plot_sticks(stick_spectrum,rgb,fig,ax,labels):
# Plot data
print(stick_spectrum[:,0])
print(stick_spectrum[:,1])
#spec = ax.stem(stick_spectrum[:,0], stick_spectrum[:,1])
spec = ax.stem([0,1,2],[12,51,1])
pyplot.setp(spec,color=rgb*(1.0/256.0))
return spec
# # Define input parser
# In[ ]:
# # Command-line driver
# In[ ]:
if __name__ == '__main__':
parser = make_parser()
args = parser.parse_args()
print('#',args)
main(args)