text: #!/usr/bin/env python
  # -*- coding: utf-8 -*-
  
  # This file is part of the Kramers-Kronig Calculator software package.
  #
  # Copyright (c) 2013 Benjamin Watts, Daniel J. Lauk
  #
  # The software is licensed under the terms of the zlib/libpng license.
  # For details see LICENSE.txt
  
  """This module implements the Kramers-Kronig transformation."""
  
  import logging, sys
  logger = logging.getLogger(__name__)
  if __name__ == '__main__':
  	logging.basicConfig(level=logging.DEBUG)
  	logging.StreamHandler(stream=sys.stdout)
  
  import math
  import numpy
  import os
  import kkcalc.data as data
  
  def calc_relativistic_correction(stoichiometry):
  	"""Calculate the relativistic correction to the Kramers-Kronig transform.
  
  	Parameters:
  	-----------
  	stoichiometry : array of integer/float pairs
  		Each pair in the list consists of an atomic number and the relative proportion of that element.
  
  	Returns
  	-------
  	This function returns a ``float`` holding the relativistic
  	corection to the Kramers-Kronig transform.
  
  	"""
  	correction = 0
  	for z, n in stoichiometry:
  		correction += (z - (z/82.5)**2.37) * n
  	return correction
  
  
  def KK_General_PP(Eval_Energy, Energy, imaginary_spectrum, orders, relativistic_correction):
  	"""Calculate Kramers-Kronig transform with "Piecewise Polynomial"
  	algorithm plus the Biggs and Lighthill extended data.
  
  	Parameters
  	----------
  	Eval_Energy : numpy vector of `float`
  		Set of photon energies describing points at which to evaluate the real spectrum
  	Energy : numpy vector of `float`
  		Set of photon energies describing intervals for which each row of `imaginary_spectrum` is valid
  	imaginary_spectrum : two-dimensional `numpy.array` of `float`
  		The array consists of columns of polynomial coefficients belonging to the power terms indicated by 'order'
  	orders : numpy vector of integers
  		The vector represents the polynomial indices corresponding to the columns of imaginary_spectrum
  	relativistic_correction : float
  		The relativistic correction to the Kramers-Kronig transform.
  		You can calculate the value using the `calc_relativistic_correction` function.
  
  	Returns
  	-------
  	This function returns the real part of the scattering factors evaluated at photon energies specified by Eval_Energy.
  
  	"""
  	logger = logging.getLogger(__name__)
  	logger.info("Calculate Kramers-Kronig transform using general piecewise-polynomial algorithm")
  	# Need to build x-E-n arrays
  	X = numpy.tile(Energy[:,numpy.newaxis,numpy.newaxis],(1,len(Eval_Energy),len(orders)))
  	E = numpy.tile(Eval_Energy[numpy.newaxis,:,numpy.newaxis],(len(Energy)-1,1,len(orders)))
  	C = numpy.tile(imaginary_spectrum[:,numpy.newaxis,:],(1,len(Eval_Energy),1))
  	N = numpy.tile(orders[numpy.newaxis,numpy.newaxis,:],(len(Energy)-1,len(Eval_Energy),1))
  	poles = numpy.equal(X,numpy.tile(Eval_Energy[numpy.newaxis,:,numpy.newaxis],(len(Energy),1,len(orders))))
  	
  	# all N, ln(x+E) and ln(x-E) terms and poles
  	Integral = numpy.sum(-C*(-E)**N*numpy.log(numpy.absolute((X[1:,:,:]+E)/(X[:-1,:,:]+E)))-C*E**N*(1-poles[1:,:,:])*numpy.log(numpy.absolute((X[1:,:,:]-E+poles[1:,:,:])/((1-poles[:-1,:,:])*X[:-1,:,:]+poles[:-1,:,:]*X[[0]+list(range(len(Energy)-2)),:,:]-E))),axis=(0,2))
  	
  	if numpy.any(orders<=-2): # N<=-2, ln(x) terms
  		i = [slice(None,None,None),slice(None,None,None),orders<=-2]
  		Integral += numpy.sum(C[i]*((-E[i])**N[i]+E[i]**N[i])*numpy.log(numpy.absolute((X[1:,:,orders<=-2])/(X[:-1,:,orders<=-2]))),axis=(0,2))
  	
  	if numpy.any(orders>=0): # N>=0,  x^k terms
  		for ni in numpy.where(orders>=0)[0]:
  			i = [slice(None,None,None),slice(None,None,None),ni]
  			n = orders[ni]
  			for k in range(n,0,-2):
  				Integral += numpy.sum(C[i]/float(-k)*2*E[i]**(n-k)*(X[1:,:,ni]**k-X[:-1,:,ni]**k),axis=0)
  	
  	if numpy.any(orders <=-3): # N<=-3, x^k terms
  		for ni in numpy.where(orders<=-3)[0]:
  			i = [slice(None,None,None),slice(None,None,None),ni]
  			n = orders[ni]
  			for k in range(n+2,0,2):
  				Integral += numpy.sum(C[i]/float(k)*((-1)**(n-k)+1)*E[i]**(n-k)*(X[1:,:,ni]**k-X[:-1,:,ni]**k),axis=0)
  	
  	logger.debug("Done!")
  	return Integral / math.pi + relativistic_correction
  
  def KK_PP(Eval_Energy, Energy, imaginary_spectrum, relativistic_correction):
  	"""Calculate Kramers-Kronig transform with "Piecewise Polynomial"
  	algorithm plus the Biggs and Lighthill extended data.
  
  	Parameters
  	----------
  	Eval_Energy : numpy vector of `float`
  		Set of photon energies describing points at which to evaluate the real spectrum
  	Energy : numpy vector of `float`
  		Set of photon energies describing intervals for which each row of `imaginary_spectrum` is valid
  	imaginary_spectrum : two-dimensional `numpy.array` of `float`
  		The array consists of five columns of polynomial coefficients: A_1, A_0, A_-1, A_-2, A_-3
  	relativistic_correction : float
  		The relativistic correction to the Kramers-Kronig transform.
  		You can calculate the value using the `calc_relativistic_correction` function.
  
  	Returns
  	-------
  	This function returns the real part of the scattering factors evaluated at photon energies specified by Eval_Energy.
  
  	"""
  	logger = logging.getLogger(__name__)
  	logger.info("Calculate Kramers-Kronig transform using (n from 1 to -3) piecewise-polynomial algorithm")
  	X1 = Energy[0:-1]
  	X2 = Energy[1:]
  	E = numpy.tile(Eval_Energy, (len(Energy)-1, 1)).T
  	Full_coeffs = imaginary_spectrum.T
  	Symb_1 = (( Full_coeffs[0, :]*E+Full_coeffs[1, :])*(X2-X1)+0.5*Full_coeffs[0, :]*(X2**2-X1**2)-(Full_coeffs[3, :]/E+Full_coeffs[4, :]*E**-2)*numpy.log(numpy.absolute(X2/X1))+Full_coeffs[4, :]/E*(X2**-1-X1**-1))
  	Symb_2 = ((-Full_coeffs[0, :]*E+Full_coeffs[1, :])*(X2-X1)+0.5*Full_coeffs[0, :]*(X2**2-X1**2)+(Full_coeffs[3, :]/E-Full_coeffs[4, :]*E**-2)*numpy.log(numpy.absolute(X2/X1))-Full_coeffs[4, :]/E*(X2**-1-X1**-1))+(Full_coeffs[0, :]*E**2-Full_coeffs[1, :]*E+Full_coeffs[2, :]-Full_coeffs[3, :]*E**-1+Full_coeffs[4, :]*E**-2)*numpy.log(numpy.absolute((X2+E)/(X1+E)))
  	Symb_3 = (1-1*((X2==E)|(X1==E)))*(Full_coeffs[0, :]*E**2+Full_coeffs[1, :]*E+Full_coeffs[2, :]+Full_coeffs[3, :]*E**-1+Full_coeffs[4, :]*E**-2)*numpy.log(numpy.absolute((X2-E+1*(X2==E))/(X1-E+1*(X1==E))))
  	Symb_B = numpy.sum(Symb_2 - Symb_1 - Symb_3, axis=1)  # Sum areas for approximate integral
  	# Patch singularities
  	hits = Energy[1:-1]==E[:,0:-1]
  	E_hits = numpy.append(numpy.insert(numpy.any(hits, axis=0),[0,0],False),[False,False])
  	Eval_hits = numpy.any(hits, axis=1)
  	X1 = Energy[E_hits[2:]]
  	XE = Energy[E_hits[1:-1]]
  	X2 = Energy[E_hits[:-2]]
  	C1 = Full_coeffs[:, E_hits[2:-1]]
  	C2 = Full_coeffs[:, E_hits[1:-2]]
  	Symb_singularities = numpy.zeros(len(Eval_Energy))
  	Symb_singularities[Eval_hits] = (C2[0, :]*XE**2+C2[1, :]*XE+C2[2, :]+C2[3, :]*XE**-1+C2[4, :]*XE**-2)*numpy.log(numpy.absolute((X2-XE)/(X1-XE)))
  	# Finish things off
  	KK_Re = (Symb_B-Symb_singularities) / (math.pi*Eval_Energy) + relativistic_correction
  	logger.debug("Done!")
  	return KK_Re
  
  def improve_accuracy(Full_E, Real_Spectrum, Imaginary_Spectrum, relativistic_correction, tolerance, recursion=50):
  	"""Calculate extra data points so that a linear interpolation is more accurate.
  	
  	Parameters
  	----------
  	Full_E : numpy vector of `float`
  		Set of photon energies describing intervals for which each row of `imaginary_spectrum` is valid
  	Real_Spectrum : numpy vector of `float`
  		The real part of the spectrum corresponding to magnitudes at photon energies in Full_E
  	Imaginary_Spectrum : two-dimensional `numpy.array` of `float`
  		The array consists of five columns of polynomial coefficients: A_1, A_0, A_-1, A_-2, A_-3
  	relativistic_correction : float
  		The relativistic correction to the Kramers-Kronig transform.
  		(You can calculate the value using the `calc_relativistic_correction` function.)
  	tolerance : float
  		Level of error in linear extrapolation of data values to be allowed.
  	recursion : integer
  		Number of times an energy interval can be halved before giving up.
  
  	Returns
  	-------
  	This function returns a numpy array with three columns respectively representing photon energy, the real spectrum and the imaginary spectrum.
  	"""
  	logger.debug("Improve data accuracy")
  	new_points = numpy.cumsum(numpy.ones((len(Full_E)-2,1),dtype=numpy.int8))+1
  	Im_values = data.coeffs_to_ASF(Full_E, numpy.vstack((Imaginary_Spectrum,Imaginary_Spectrum[-1])))
  	#plot_Im_values = Im_values
  	Re_values = Real_Spectrum
  	E_values = Full_E
  	temp_Im_spectrum = Imaginary_Spectrum[1:]
  	count = 0
  	improved = 1
  	total_improved_points = 0
  	while count<recursion and numpy.sum(improved)>0:
  		#get E_midpoints
  		midpoints = (E_values[new_points-1]+E_values[new_points])/2.
  		#evaluate at new points
  		Im_midpoints = data.coeffs_to_ASF(midpoints, temp_Im_spectrum)
  		Re_midpoints = KK_PP(midpoints, Full_E, Imaginary_Spectrum, relativistic_correction)
  		#evaluate error levels
  		Im_error = abs((Im_values[new_points-1]+Im_values[new_points])/2. - Im_midpoints)
  		Re_error = abs((Re_values[new_points-1]+Re_values[new_points])/2. - Re_midpoints)
  		improved = (Im_error>tolerance) | (Re_error>tolerance)
  		logger.debug(str(numpy.sum(improved))+" points (out of "+str(len(improved))+") can be improved in pass number "+str(count+1)+".")
  		total_improved_points += numpy.sum(improved)
  		#insert new points and values
  		Im_values = numpy.insert(Im_values,new_points[improved],Im_midpoints[improved])
  		Re_values = numpy.insert(Re_values,new_points[improved],Re_midpoints[improved])
  		E_values = numpy.insert(E_values,new_points[improved],midpoints[improved])
  		#prepare for next loop
  		temp_Im_spectrum =numpy.repeat(temp_Im_spectrum[improved],2,axis=0)
  		new_points = numpy.where(numpy.insert(numpy.zeros(Im_values.shape, dtype=numpy.bool),new_points[improved],True))[0]
  		new_points = numpy.vstack((new_points, new_points+1)).T.flatten()
  		count += 1
  	
  	#import matplotlib
  	#matplotlib.use('WXAgg')
  	#import pylab
  	#pylab.figure()
  	#pylab.plot(Full_E,plot_Im_values,'ok')
  	#pylab.plot(Full_E,Real_Spectrum,'og')
  	#pylab.plot(midpoints,Im_midpoints,'+b')
  	#pylab.plot(midpoints,Re_midpoints,'+r')
  	#pylab.plot(E_values,Im_values,'b-')
  	#pylab.plot(E_values,Re_values,'r-')
  	#pylab.plot(midpoints,Im_error,'b-')
  	#pylab.plot(midpoints,Re_error,'r-')
  	#pylab.xscale('log')
  	#pylab.show()
  	logger.info("Improved data accuracy by inserting "+str(total_improved_points)+" extra points.")
  	return numpy.vstack((E_values,Re_values,Im_values)).T
  	
  def kk_calculate_real(NearEdgeDataFile, ChemicalFormula, load_options=None, input_data_type=None, merge_points=None, add_background=False, fix_distortions=False, curve_tolerance=None, curve_recursion=50):
  	"""Do all data loading and processing and then calculate the kramers-Kronig transform.
  	Parameters
  	----------
  	NearEdgeDataFile : string
  		Path to file containg near-edge data
  	ChemicalFormula : string
  		A standard chemical formula string consisting of element symbols, numbers and parentheses.
  	merge_points : list or tuple pair of `float` values, or None
  		The photon energy values (low, high) at which the near-edge and scattering factor data values
  		are set equal so as to ensure continuity of the merged data set.
  
  	Returns
  	-------
  	This function returns a numpy array with columns consisting of the photon energy, the real and the imaginary parts of the scattering factors.
  	"""
  	Stoichiometry = data.ParseChemicalFormula(ChemicalFormula)
  	Relativistic_Correction = calc_relativistic_correction(Stoichiometry)
  	Full_E, Imaginary_Spectrum = data.calculate_asf(Stoichiometry)
  	if NearEdgeDataFile is not None:
  		NearEdge_Data = data.convert_data(data.load_data(NearEdgeDataFile, load_options),FromType=input_data_type,ToType='asf')
  		Full_E, Imaginary_Spectrum = data.merge_spectra(NearEdge_Data, Full_E, Imaginary_Spectrum, merge_points=merge_points, add_background=add_background, fix_distortions=fix_distortions)
  	Real_Spectrum = KK_PP(Full_E, Full_E, Imaginary_Spectrum, Relativistic_Correction)
  	if curve_tolerance is not None:
  		output_data = improve_accuracy(Full_E,Real_Spectrum,Imaginary_Spectrum, Relativistic_Correction, curve_tolerance, curve_recursion)
  	else:
  		Imaginary_Spectrum_Values = data.coeffs_to_ASF(Full_E, numpy.vstack((Imaginary_Spectrum,Imaginary_Spectrum[-1])))
  		output_data = numpy.vstack((Full_E,Real_Spectrum,Imaginary_Spectrum_Values)).T
  	return output_data
  
  if __name__ == '__main__':
  	#use argparse here to get command line arguments
  	#process arguments and pass to a pythonic function
  	
  	#I will abuse this section of code for initial testing
  	#Output = kk_calculate_real('../../data/Xy_norm_bgsub.txt', 'C10SH14', input_data_type='NEXAFS')
  	Output = kk_calculate_real('../../data/LaAlO3/LaAlO3_Exp.csv', 'LaAlO3', input_data_type='NEXAFS', fix_distortions=True, curve_tolerance=0.05)
  	#Output = kk_calculate_real('../../data/GaAs/As.xmu.csv', 'GaAs', input_data_type='NEXAFS', fix_distortions=True, curve_tolerance=0.05)
  	
  	Stoichiometry = data.ParseChemicalFormula('LaAlO3')
  	#Stoichiometry = data.ParseChemicalFormula('GaAs')
  	Relativistic_Correction = calc_relativistic_correction(Stoichiometry)
  	ASF_E, ASF_Data = data.calculate_asf(Stoichiometry)
  	ASF_Data3 = data.coeffs_to_linear(ASF_E, ASF_Data, 0.1)
  	ASF_Data2 = data.coeffs_to_ASF(ASF_E, numpy.vstack((ASF_Data,ASF_Data[-1])))
  	
  	#Test_E = (Output[1:,0]+Output[0:-1,0])*0.5
  	#Test_E = numpy.linspace(41257.87,41259.87,num=21)
  	#Real_Spectrum2 = KK_PP(Test_E, Output[:,0], Im, Relativistic_Correction)
  	
  	
  	#import matplotlib
  	#matplotlib.use('WXAgg')
  	#import pylab
  	
  	#pylab.figure()
  	#pylab.plot(Output[:,0],Output[:,1],'xg-',Output[:,0],Output[:,2],'xb-')
  	#pylab.plot(ASF_E,ASF_Data2,'+r')
  	#pylab.plot(ASF_E,ASF_Data22,'xr')
  	#pylab.plot(ASF_Data3[0],ASF_Data3[1],'r-')
  	#pylab.plot(Test_E,Real_Spectrum2,'*y')
  	#pylab.xscale('log')
  	#pylab.show()
  	pass
  