<s> """<STR_LIT>""" <EOL> import os <EOL> import time <EOL> import json <EOL> import numpy as np <EOL> from scipy . interpolate import interp1d , splrep , splev , UnivariateSpline <EOL> from sqlalchemy import MetaData , create_engine <EOL> from sqlalchemy . orm import sessionmaker , mapper , clear_mappers <EOL> from sqlalchemy . pool import SingletonThreadPool <EOL> import sqlalchemy . dialects . sqlite <EOL> import larch <EOL> from larch_plugins . math import as_ndarray <EOL> def OLDas_ndarray ( obj ) : <EOL> """<STR_LIT>""" <EOL> if isinstance ( obj , ( float , int ) ) : <EOL> return np . array ( [ obj ] ) <EOL> return np . asarray ( obj ) <EOL> def make_engine ( dbname ) : <EOL> return create_engine ( '<STR_LIT>' % ( dbname ) , <EOL> poolclass = SingletonThreadPool ) <EOL> def isxrayDB ( dbname ) : <EOL> """<STR_LIT>""" <EOL> _tables = ( '<STR_LIT>' , '<STR_LIT>' , '<STR_LIT>' , <EOL> '<STR_LIT>' , <EOL> '<STR_LIT>' , '<STR_LIT>' , '<STR_LIT>' ) <EOL> result = False <EOL> try : <EOL> engine = make_engine ( dbname ) <EOL> meta = MetaData ( engine ) <EOL> meta . reflect ( ) <EOL> result = all ( [ t in meta . tables for t in _tables ] ) <EOL> except : <EOL> pass <EOL> return result <EOL> def json_encode ( val ) : <EOL> "<STR_LIT>" <EOL> if val is None or isinstance ( val , ( str , unicode ) ) : <EOL> return val <EOL> return json . dumps ( val ) <EOL> def elam_spline ( xin , yin , yspl_in , x ) : <EOL> """<STR_LIT>""" <EOL> x = as_ndarray ( x ) <EOL> x [ np . where ( x < min ( xin ) ) ] = min ( xin ) <EOL> x [ np . where ( x > max ( xin ) ) ] = max ( xin ) <EOL> lo , hi = np . array ( [ ( np . flatnonzero ( xin < e ) [ - <NUM_LIT:1> ] , <EOL> np . flatnonzero ( xin > e ) [ <NUM_LIT:0> ] ) <EOL> for e in x ] ) . transpose ( ) <EOL> diff = xin [ hi ] - xin [ lo ] <EOL> if any ( diff <= <NUM_LIT:0> ) : <EOL> raise ValueError ( '<STR_LIT>' ) <EOL> a = ( xin [ hi ] - x ) / diff <EOL> b = ( x - xin [ lo ] ) / diff <EOL> return ( a * yin [ lo ] + b * yin [ hi ] + <EOL> ( diff * diff / <NUM_LIT:6> ) * ( ( a * a - <NUM_LIT:1> ) * a * yspl_in [ lo ] + <EOL> ( b * b - <NUM_LIT:1> ) * b * yspl_in [ hi ] ) ) <EOL> class DBException ( Exception ) : <EOL> """<STR_LIT>""" <EOL> def __init__ ( self , msg ) : <EOL> Exception . __init__ ( self ) <EOL> self . msg = msg <EOL> def __str__ ( self ) : <EOL> return self . msg <EOL> class _BaseTable ( object ) : <EOL> "<STR_LIT>" <EOL> def __repr__ ( self ) : <EOL> el = getattr ( self , '<STR_LIT>' , '<STR_LIT>' ) <EOL> return "<STR_LIT>" % ( self . __class__ . __name__ , el ) <EOL> class CosterKronigTable ( _BaseTable ) : <EOL> ( id , element , initial_level , final_level , <EOL> transition_probability , total_transition_probability ) = [ None ] * <NUM_LIT:6> <EOL> class ElementsTable ( _BaseTable ) : <EOL> ( atomic_number , element , molar_mass , density ) = [ None ] * <NUM_LIT:4> <EOL> class PhotoAbsorptionTable ( _BaseTable ) : <EOL> ( id , element , log_energy , <EOL> log_photoabsorption , log_photoabsorption_spline ) = [ None ] * <NUM_LIT:5> <EOL> class ScatteringTable ( _BaseTable ) : <EOL> ( id , element , log_energy , <EOL> log_coherent_scatter , log_coherent_scatter_spline , <EOL> log_incoherent_scatter , log_incoherent_scatter_spline ) = [ None ] * <NUM_LIT:7> <EOL> class XrayLevelsTable ( _BaseTable ) : <EOL> ( id , element , iupac_symbol , <EOL> absorption_edge , fluorescence_yield , jump_ratio ) = [ None ] * <NUM_LIT:6> <EOL> def __repr__ ( self ) : <EOL> el = getattr ( self , '<STR_LIT>' , '<STR_LIT>' ) <EOL> edge = getattr ( self , '<STR_LIT>' , '<STR_LIT>' ) <EOL> return "<STR_LIT>" % ( self . __class__ . __name__ , el , edge ) <EOL> class XrayTransitionsTable ( _BaseTable ) : <EOL> ( id , element , iupac_symbol , siegbahn_symbol , initial_level , <EOL> final_level , emission_energy , intensity ) = [ None ] * <NUM_LIT:8> <EOL> def __repr__ ( self ) : <EOL> el = getattr ( self , '<STR_LIT>' , '<STR_LIT>' ) <EOL> line = getattr ( self , '<STR_LIT>' , '<STR_LIT>' ) <EOL> return "<STR_LIT>" % ( self . __class__ . __name__ , el , line ) <EOL> class WaasmaierTable ( _BaseTable ) : <EOL> ( id , atomic_number , element , ion , offset , scale , exponents ) = [ None ] * <NUM_LIT:7> <EOL> def __repr__ ( self ) : <EOL> el = getattr ( self , '<STR_LIT>' , '<STR_LIT>' ) <EOL> return "<STR_LIT>" % ( self . __class__ . __name__ , el ) <EOL> class KeskiRahkonenKrauseTable ( _BaseTable ) : <EOL> ( id , atomic_number , element , edge , width ) = [ None ] * <NUM_LIT:5> <EOL> def __repr__ ( self ) : <EOL> el = getattr ( self , '<STR_LIT>' , '<STR_LIT>' ) <EOL> edge = getattr ( self , '<STR_LIT>' , '<STR_LIT>' ) <EOL> return "<STR_LIT>" % ( self . __class__ . __name__ , el , edge ) <EOL> class ChantlerTable ( _BaseTable ) : <EOL> ( id , element , sigma_mu , mue_f2 , density , <EOL> corr_henke , corr_cl35 , corr_nucl , <EOL> energy , f1 , f2 , mu_photo , mu_incoh , mu_total ) = [ None ] * <NUM_LIT> <EOL> class xrayDB ( object ) : <EOL> "<STR_LIT>" <EOL> def __init__ ( self , dbname = '<STR_LIT>' , read_only = True ) : <EOL> "<STR_LIT>" <EOL> if not os . path . exists ( dbname ) : <EOL> parent , child = os . path . split ( __file__ ) <EOL> dbname = os . path . join ( parent , dbname ) <EOL> if not os . path . exists ( dbname ) : <EOL> raise IOError ( "<STR_LIT>" % dbname ) <EOL> if not isxrayDB ( dbname ) : <EOL> raise ValueError ( "<STR_LIT>" % dbname ) <EOL> self . dbname = dbname <EOL> self . engine = make_engine ( dbname ) <EOL> self . conn = self . engine . connect ( ) <EOL> kwargs = { } <EOL> if read_only : <EOL> kwargs = { '<STR_LIT>' : True , '<STR_LIT>' : False } <EOL> def readonly_flush ( * args , ** kwargs ) : <EOL> return <EOL> self . session = sessionmaker ( bind = self . engine , ** kwargs ) ( ) <EOL> self . session . flush = readonly_flush <EOL> else : <EOL> self . session = sessionmaker ( bind = self . engine , ** kwargs ) ( ) <EOL> self . metadata = MetaData ( self . engine ) <EOL> self . metadata . reflect ( ) <EOL> tables = self . tables = self . metadata . tables <EOL> clear_mappers ( ) <EOL> mapper ( ChantlerTable , tables [ '<STR_LIT>' ] ) <EOL> mapper ( WaasmaierTable , tables [ '<STR_LIT>' ] ) <EOL> mapper ( KeskiRahkonenKrauseTable , tables [ '<STR_LIT>' ] ) <EOL> mapper ( ElementsTable , tables [ '<STR_LIT>' ] ) <EOL> mapper ( XrayLevelsTable , tables [ '<STR_LIT>' ] ) <EOL> mapper ( XrayTransitionsTable , tables [ '<STR_LIT>' ] ) <EOL> mapper ( CosterKronigTable , tables [ '<STR_LIT>' ] ) <EOL> mapper ( PhotoAbsorptionTable , tables [ '<STR_LIT>' ] ) <EOL> mapper ( ScatteringTable , tables [ '<STR_LIT>' ] ) <EOL> def close ( self ) : <EOL> "<STR_LIT>" <EOL> self . session . flush ( ) <EOL> self . session . close ( ) <EOL> def query ( self , * args , ** kws ) : <EOL> "<STR_LIT>" <EOL> return self . session . query ( * args , ** kws ) <EOL> def f0_ions ( self , element = None ) : <EOL> """<STR_LIT>""" <EOL> rows = self . query ( WaasmaierTable ) <EOL> if element is not None : <EOL> if isinstance ( element , int ) : <EOL> rows = rows . filter ( WaasmaierTable . atomic_number == element ) <EOL> else : <EOL> rows = rows . filter ( WaasmaierTable . element == element . title ( ) ) <EOL> return [ str ( r . ion ) for r in rows . all ( ) ] <EOL> def f0 ( self , ion , q ) : <EOL> """<STR_LIT>""" <EOL> tab = WaasmaierTable <EOL> row = self . query ( tab ) <EOL> if isinstance ( ion , int ) : <EOL> row = row . filter ( tab . atomic_number == ion ) . all ( ) <EOL> else : <EOL> row = row . filter ( tab . ion == ion . title ( ) ) . all ( ) <EOL> if len ( row ) > <NUM_LIT:0> : <EOL> row = row [ <NUM_LIT:0> ] <EOL> if isinstance ( row , tab ) : <EOL> q = as_ndarray ( q ) <EOL> f0 = row . offset <EOL> for s , e in zip ( json . loads ( row . scale ) , json . loads ( row . exponents ) ) : <EOL> f0 += s * np . exp ( - e * q * q ) <EOL> return f0 <EOL> def _getChantler ( self , element , energy , column = '<STR_LIT>' , smoothing = <NUM_LIT:0> ) : <EOL> """<STR_LIT>""" <EOL> tab = ChantlerTable <EOL> row = self . query ( tab ) <EOL> if isinstance ( element , int ) : <EOL> row = row . filter ( tab . id == element ) . all ( ) <EOL> else : <EOL> row = row . filter ( tab . element == element . title ( ) ) . all ( ) <EOL> if len ( row ) > <NUM_LIT:0> : <EOL> row = row [ <NUM_LIT:0> ] <EOL> if isinstance ( row , tab ) : <EOL> energy = as_ndarray ( energy ) <EOL> emin , emax = min ( energy ) , max ( energy ) <EOL> te = np . array ( json . loads ( row . energy ) ) <EOL> nemin = max ( <NUM_LIT:0> , - <NUM_LIT:5> + max ( np . where ( te <= emin ) [ <NUM_LIT:0> ] ) ) <EOL> nemax = min ( len ( te ) , <NUM_LIT:6> + max ( np . where ( te <= emax ) [ <NUM_LIT:0> ] ) ) <EOL> region = np . arange ( nemin , nemax ) <EOL> te = te [ region ] <EOL> if column == '<STR_LIT>' : <EOL> column = '<STR_LIT>' <EOL> ty = np . array ( json . loads ( getattr ( row , column ) ) ) [ region ] <EOL> if column == '<STR_LIT>' : <EOL> out = UnivariateSpline ( te , ty , s = smoothing ) ( energy ) <EOL> else : <EOL> out = np . exp ( np . interp ( np . log ( energy ) , <EOL> np . log ( te ) , <EOL> np . log ( ty ) ) ) <EOL> if isinstance ( out , np . ndarray ) and len ( out ) == <NUM_LIT:1> : <EOL> return out [ <NUM_LIT:0> ] <EOL> return out <EOL> def chantler_energies ( self , element , emin = <NUM_LIT:0> , emax = <NUM_LIT> ) : <EOL> """<STR_LIT>""" <EOL> tab = ChantlerTable <EOL> row = self . query ( tab ) <EOL> if isinstance ( element , int ) : <EOL> row = row . filter ( tab . id == element ) . all ( ) <EOL> else : <EOL> row = row . filter ( tab . element == element . title ( ) ) . all ( ) <EOL> if len ( row ) > <NUM_LIT:0> : <EOL> row = row [ <NUM_LIT:0> ] <EOL> if not isinstance ( row , tab ) : <EOL> return None <EOL> te = np . array ( json . loads ( row . energy ) ) <EOL> tf1 = np . array ( json . loads ( row . f1 ) ) <EOL> tf2 = np . array ( json . loads ( row . f2 ) ) <EOL> if emin <= min ( te ) : <EOL> nemin = <NUM_LIT:0> <EOL> else : <EOL> nemin = max ( <NUM_LIT:0> , - <NUM_LIT:2> + max ( np . where ( te <= emin ) [ <NUM_LIT:0> ] ) ) <EOL> if emax > max ( te ) : <EOL> nemax = len ( te ) <EOL> else : <EOL> nemax = min ( len ( te ) , <NUM_LIT:2> + max ( np . where ( te <= emax ) [ <NUM_LIT:0> ] ) ) <EOL> region = np . arange ( nemin , nemax ) <EOL> return te [ region ] <EOL> def f1_chantler ( self , element , energy , ** kws ) : <EOL> """<STR_LIT>""" <EOL> return self . _getChantler ( element , energy , column = '<STR_LIT>' , ** kws ) <EOL> def f2_chantler ( self , element , energy , ** kws ) : <EOL> """<STR_LIT>""" <EOL> return self . _getChantler ( element , energy , column = '<STR_LIT>' , ** kws ) <EOL> def mu_chantler ( self , element , energy , incoh = False , photo = False ) : <EOL> """<STR_LIT>""" <EOL> col = '<STR_LIT>' <EOL> if photo : <EOL> col = '<STR_LIT>' <EOL> elif incoh : <EOL> col = '<STR_LIT>' <EOL> return self . _getChantler ( element , energy , column = col ) <EOL> def _getElementData ( self , element ) : <EOL> "<STR_LIT>" <EOL> tab = ElementsTable <EOL> row = self . query ( tab ) <EOL> if isinstance ( element , int ) : <EOL> row = row . filter ( tab . atomic_number == element ) . all ( ) <EOL> else : <EOL> row = row . filter ( tab . element == element . title ( ) ) . all ( ) <EOL> if len ( row ) > <NUM_LIT:0> : row = row [ <NUM_LIT:0> ] <EOL> return row <EOL> def zofsym ( self , element ) : <EOL> "<STR_LIT>" <EOL> return int ( self . _getElementData ( element ) . atomic_number ) <EOL> def symbol ( self , z ) : <EOL> "<STR_LIT>" <EOL> return self . _getElementData ( z ) . element <EOL> def molar_mass ( self , element ) : <EOL> "<STR_LIT>" <EOL> return self . _getElementData ( element ) . molar_mass <EOL> def density ( self , element ) : <EOL> "<STR_LIT>" <EOL> return self . _getElementData ( element ) . density <EOL> def xray_edges ( self , element ) : <EOL> """<STR_LIT>""" <EOL> if isinstance ( element , int ) : <EOL> element = self . symbol ( element ) <EOL> tab = XrayLevelsTable <EOL> out = { } <EOL> for r in self . query ( tab ) . filter ( tab . element == element . title ( ) ) . all ( ) : <EOL> out [ str ( r . iupac_symbol ) ] = ( r . absorption_edge , <EOL> r . fluorescence_yield , <EOL> r . jump_ratio ) <EOL> return out <EOL> def xray_edge ( self , element , edge ) : <EOL> """<STR_LIT>""" <EOL> edge = edge . title ( ) <EOL> edges = self . xray_edges ( element ) <EOL> if edge in edges : <EOL> return edges [ edge ] <EOL> def xray_lines ( self , element , initial_level = None , excitation_energy = None ) : <EOL> """<STR_LIT>""" <EOL> if isinstance ( element , int ) : <EOL> element = self . symbol ( element ) <EOL> tab = XrayTransitionsTable <EOL> row = self . query ( tab ) . filter ( tab . element == element . title ( ) ) <EOL> if excitation_energy is not None : <EOL> initial_level = [ ] <EOL> for ilevel , dat in self . xray_edges ( element ) . items ( ) : <EOL> if dat [ <NUM_LIT:0> ] < excitation_energy : <EOL> initial_level . append ( ilevel . title ( ) ) <EOL> if initial_level is not None : <EOL> if isinstance ( initial_level , ( list , tuple ) ) : <EOL> row = row . filter ( tab . initial_level . in_ ( initial_level ) ) <EOL> else : <EOL> row = row . filter ( tab . initial_level == initial_level . title ( ) ) <EOL> out = { } <EOL> for r in row . all ( ) : <EOL> out [ str ( r . siegbahn_symbol ) ] = ( r . emission_energy , r . intensity , <EOL> r . initial_level , r . final_level ) <EOL> return out <EOL> def CK_probability ( self , element , initial , final , total = True ) : <EOL> """<STR_LIT>""" <EOL> if isinstance ( element , int ) : <EOL> element = self . symbol ( element ) <EOL> tab = CosterKronigTable <EOL> row = self . query ( tab ) . filter ( <EOL> tab . element == element . title ( ) <EOL> ) . filter ( <EOL> tab . initial_level == initial . title ( ) <EOL> ) . filter ( <EOL> tab . final_level == final . title ( ) ) . all ( ) <EOL> if len ( row ) > <NUM_LIT:0> : <EOL> row = row [ <NUM_LIT:0> ] <EOL> if isinstance ( row , tab ) : <EOL> if total : <EOL> return row . total_transition_probability <EOL> else : <EOL> return row . transition_probability <EOL> def corehole_width ( self , element = None , edge = None ) : <EOL> """<STR_LIT>""" <EOL> tab = KeskiRahkonenKrauseTable <EOL> rows = self . query ( tab ) <EOL> has_elem = element is not None <EOL> has_edge = edge is not None <EOL> if has_elem : <EOL> if isinstance ( element , int ) : <EOL> rows = rows . filter ( tab . atomic_number == element ) <EOL> else : <EOL> rows = rows . filter ( tab . element == element . title ( ) ) <EOL> if has_edge : <EOL> rows = rows . filter ( tab . edge == edge . title ( ) ) <EOL> out = rows . all ( ) <EOL> if len ( out ) == <NUM_LIT:1> : <EOL> return ( out [ <NUM_LIT:0> ] . width ) <EOL> elif has_elem : <EOL> return [ ( r . edge , r . width ) for r in out ] <EOL> elif has_edge : <EOL> return [ ( r . atomic_number , r . width ) for r in out ] <EOL> else : <EOL> return [ ( r . atomic_number , r . edge , r . width ) for r in out ] <EOL> def Elam_CrossSection ( self , element , energies , kind = '<STR_LIT>' ) : <EOL> """<STR_LIT>""" <EOL> if isinstance ( element , int ) : <EOL> element = self . symbol ( element ) <EOL> energies = <NUM_LIT:1.0> * as_ndarray ( energies ) <EOL> tab = ScatteringTable <EOL> if kind == '<STR_LIT>' : <EOL> tab = PhotoAbsorptionTable <EOL> row = self . query ( tab ) . filter ( tab . element == element . title ( ) ) . all ( ) <EOL> if len ( row ) > <NUM_LIT:0> : <EOL> row = row [ <NUM_LIT:0> ] <EOL> if not isinstance ( row , tab ) : <EOL> return None <EOL> tab_lne = np . array ( json . loads ( row . log_energy ) ) <EOL> if kind . lower ( ) . startswith ( '<STR_LIT>' ) : <EOL> tab_val = np . array ( json . loads ( row . log_coherent_scatter ) ) <EOL> tab_spl = np . array ( json . loads ( row . log_coherent_scatter_spline ) ) <EOL> elif kind . lower ( ) . startswith ( '<STR_LIT>' ) : <EOL> tab_val = np . array ( json . loads ( row . log_incoherent_scatter ) ) <EOL> tab_spl = np . array ( json . loads ( row . log_incoherent_scatter_spline ) ) <EOL> else : <EOL> tab_val = np . array ( json . loads ( row . log_photoabsorption ) ) <EOL> tab_spl = np . array ( json . loads ( row . log_photoabsorption_spline ) ) <EOL> emin_tab = <NUM_LIT:10> * int ( <NUM_LIT> * np . exp ( tab_lne [ <NUM_LIT:0> ] ) ) <EOL> energies [ np . where ( energies < emin_tab ) ] = emin_tab <EOL> out = np . exp ( elam_spline ( tab_lne , tab_val , tab_spl , np . log ( energies ) ) ) <EOL> if len ( out ) == <NUM_LIT:1> : <EOL> return out [ <NUM_LIT:0> ] <EOL> return out <EOL> def mu_elam ( self , element , energies , kind = '<STR_LIT>' ) : <EOL> """<STR_LIT>""" <EOL> calc = self . Elam_CrossSection <EOL> xsec = calc ( element , energies , kind = '<STR_LIT>' ) <EOL> if kind . lower ( ) . startswith ( '<STR_LIT>' ) : <EOL> xsec += calc ( element , energies , kind = '<STR_LIT>' ) <EOL> xsec += calc ( element , energies , kind = '<STR_LIT>' ) <EOL> return xsec <EOL> def coherent_cross_section_elam ( self , element , energies ) : <EOL> """<STR_LIT>""" <EOL> return self . Elam_CrossSection ( element , energies , kind = '<STR_LIT>' ) <EOL> def incoherent_cross_section_elam ( self , element , energies ) : <EOL> """<STR_LIT>""" <EOL> return self . Elam_CrossSection ( element , energies , kind = '<STR_LIT>' ) </s>
