<s> """<STR_LIT>""" <EOL> import numpy as np <EOL> from scipy import polyfit <EOL> from larch import ( Group , Parameter , ValidateLarchPlugin , <EOL> Make_CallArgs , Minimizer , isgroup , <EOL> parse_group_args ) <EOL> from larch_plugins . math import ( index_of , index_nearest , <EOL> remove_dups , remove_nans2 ) <EOL> from larch_plugins . xafs import set_xafsGroup <EOL> MODNAME = '<STR_LIT>' <EOL> MAX_NNORM = <NUM_LIT:5> <EOL> @ ValidateLarchPlugin <EOL> def find_e0 ( energy , mu = None , group = None , _larch = None ) : <EOL> """<STR_LIT>""" <EOL> energy , mu , group = parse_group_args ( energy , members = ( '<STR_LIT>' , '<STR_LIT>' ) , <EOL> defaults = ( mu , ) , group = group , <EOL> fcn_name = '<STR_LIT>' ) <EOL> energy = remove_dups ( energy ) <EOL> dmu = np . gradient ( mu ) / np . gradient ( energy ) <EOL> high_deriv_pts = np . where ( dmu > max ( dmu ) * <NUM_LIT> ) [ <NUM_LIT:0> ] <EOL> idmu_max , dmu_max = <NUM_LIT:0> , <NUM_LIT:0> <EOL> for i in high_deriv_pts : <EOL> if ( dmu [ i ] > dmu_max and <EOL> ( i + <NUM_LIT:1> in high_deriv_pts ) and <EOL> ( i - <NUM_LIT:1> in high_deriv_pts ) ) : <EOL> idmu_max , dmu_max = i , dmu [ i ] <EOL> e0 = energy [ idmu_max ] <EOL> if group is not None : <EOL> group = set_xafsGroup ( group , _larch = _larch ) <EOL> group . e0 = e0 <EOL> return e0 <EOL> def flat_resid ( pars ) : <EOL> c0 , c1 , c2 = pars . c0 . value , pars . c1 . value , pars . c2 . value <EOL> return ( pars . mu - ( c0 + pars . en * ( c1 + pars . en * c2 ) ) ) <EOL> def preedge ( energy , mu , e0 = None , step = None , <EOL> nnorm = <NUM_LIT:3> , nvict = <NUM_LIT:0> , pre1 = None , pre2 = - <NUM_LIT:50> , <EOL> norm1 = <NUM_LIT:100> , norm2 = None ) : <EOL> """<STR_LIT>""" <EOL> energy = remove_dups ( energy ) <EOL> if e0 is None or e0 < energy [ <NUM_LIT:0> ] or e0 > energy [ - <NUM_LIT:1> ] : <EOL> energy = remove_dups ( energy ) <EOL> dmu = np . gradient ( mu ) / np . gradient ( energy ) <EOL> high_deriv_pts = np . where ( dmu > max ( dmu ) * <NUM_LIT> ) [ <NUM_LIT:0> ] <EOL> idmu_max , dmu_max = <NUM_LIT:0> , <NUM_LIT:0> <EOL> for i in high_deriv_pts : <EOL> if ( dmu [ i ] > dmu_max and <EOL> ( i + <NUM_LIT:1> in high_deriv_pts ) and <EOL> ( i - <NUM_LIT:1> in high_deriv_pts ) ) : <EOL> idmu_max , dmu_max = i , dmu [ i ] <EOL> e0 = energy [ idmu_max ] <EOL> nnorm = max ( min ( nnorm , MAX_NNORM ) , <NUM_LIT:1> ) <EOL> ie0 = index_nearest ( energy , e0 ) <EOL> e0 = energy [ ie0 ] <EOL> if pre1 is None : pre1 = min ( energy ) - e0 <EOL> if norm2 is None : norm2 = max ( energy ) - e0 <EOL> if norm2 < <NUM_LIT:0> : norm2 = max ( energy ) - e0 - norm2 <EOL> pre1 = max ( pre1 , ( min ( energy ) - e0 ) ) <EOL> norm2 = min ( norm2 , ( max ( energy ) - e0 ) ) <EOL> if pre1 > pre2 : <EOL> pre1 , pre2 = pre2 , pre1 <EOL> if norm1 > norm2 : <EOL> norm1 , norm2 = norm2 , norm1 <EOL> p1 = index_of ( energy , pre1 + e0 ) <EOL> p2 = index_nearest ( energy , pre2 + e0 ) <EOL> if p2 - p1 < <NUM_LIT:2> : <EOL> p2 = min ( len ( energy ) , p1 + <NUM_LIT:2> ) <EOL> omu = mu * energy ** nvict <EOL> ex , mx = remove_nans2 ( energy [ p1 : p2 ] , omu [ p1 : p2 ] ) <EOL> precoefs = polyfit ( ex , mx , <NUM_LIT:1> ) <EOL> pre_edge = ( precoefs [ <NUM_LIT:0> ] * energy + precoefs [ <NUM_LIT:1> ] ) * energy ** ( - nvict ) <EOL> p1 = index_of ( energy , norm1 + e0 ) <EOL> p2 = index_nearest ( energy , norm2 + e0 ) <EOL> if p2 - p1 < <NUM_LIT:2> : <EOL> p2 = min ( len ( energy ) , p1 + <NUM_LIT:2> ) <EOL> coefs = polyfit ( energy [ p1 : p2 ] , omu [ p1 : p2 ] , nnorm ) <EOL> post_edge = <NUM_LIT:0> <EOL> norm_coefs = [ ] <EOL> for n , c in enumerate ( reversed ( list ( coefs ) ) ) : <EOL> post_edge += c * energy ** ( n - nvict ) <EOL> norm_coefs . append ( c ) <EOL> edge_step = step <EOL> if edge_step is None : <EOL> edge_step = post_edge [ ie0 ] - pre_edge [ ie0 ] <EOL> norm = ( mu - pre_edge ) / edge_step <EOL> out = { '<STR_LIT>' : e0 , '<STR_LIT>' : edge_step , '<STR_LIT>' : norm , <EOL> '<STR_LIT>' : pre_edge , '<STR_LIT>' : post_edge , <EOL> '<STR_LIT>' : norm_coefs , '<STR_LIT>' : nvict , <EOL> '<STR_LIT>' : nnorm , '<STR_LIT>' : norm1 , '<STR_LIT>' : norm2 , <EOL> '<STR_LIT>' : pre1 , '<STR_LIT>' : pre2 , '<STR_LIT>' : precoefs } <EOL> return out <EOL> @ ValidateLarchPlugin <EOL> @ Make_CallArgs ( [ "<STR_LIT>" , "<STR_LIT>" ] ) <EOL> def pre_edge ( energy , mu = None , group = None , e0 = None , step = None , <EOL> nnorm = <NUM_LIT:3> , nvict = <NUM_LIT:0> , pre1 = None , pre2 = - <NUM_LIT:50> , <EOL> norm1 = <NUM_LIT:100> , norm2 = None , make_flat = True , _larch = None ) : <EOL> """<STR_LIT>""" <EOL> energy , mu , group = parse_group_args ( energy , members = ( '<STR_LIT>' , '<STR_LIT>' ) , <EOL> defaults = ( mu , ) , group = group , <EOL> fcn_name = '<STR_LIT>' ) <EOL> pre_dat = preedge ( energy , mu , e0 = e0 , step = step , nnorm = nnorm , <EOL> nvict = nvict , pre1 = pre1 , pre2 = pre2 , norm1 = norm1 , <EOL> norm2 = norm2 ) <EOL> group = set_xafsGroup ( group , _larch = _larch ) <EOL> e0 = pre_dat [ '<STR_LIT>' ] <EOL> norm = pre_dat [ '<STR_LIT>' ] <EOL> norm1 = pre_dat [ '<STR_LIT>' ] <EOL> norm2 = pre_dat [ '<STR_LIT>' ] <EOL> flat = norm <EOL> ie0 = index_nearest ( energy , e0 ) <EOL> p1 = index_of ( energy , norm1 + e0 ) <EOL> p2 = index_nearest ( energy , norm2 + e0 ) <EOL> if p2 - p1 < <NUM_LIT:2> : <EOL> p2 = min ( len ( energy ) , p1 + <NUM_LIT:2> ) <EOL> if make_flat and p2 - p1 > <NUM_LIT:4> : <EOL> enx , mux = remove_nans2 ( energy [ p1 : p2 ] , norm [ p1 : p2 ] ) <EOL> fpars = Group ( c0 = Parameter ( <NUM_LIT:0> , vary = True ) , <EOL> c1 = Parameter ( <NUM_LIT:0> , vary = True ) , <EOL> c2 = Parameter ( <NUM_LIT:0> , vary = True ) , <EOL> en = enx , mu = mux ) <EOL> fit = Minimizer ( flat_resid , fpars , _larch = _larch , toler = <NUM_LIT> ) <EOL> try : <EOL> fit . leastsq ( ) <EOL> except ( TypeError , ValueError ) : <EOL> pass <EOL> fc0 , fc1 , fc2 = fpars . c0 . value , fpars . c1 . value , fpars . c2 . value <EOL> flat_diff = fc0 + energy * ( fc1 + energy * fc2 ) <EOL> flat = norm - flat_diff + flat_diff [ ie0 ] <EOL> flat [ : ie0 ] = norm [ : ie0 ] <EOL> group . e0 = e0 <EOL> group . norm = norm <EOL> group . flat = flat <EOL> group . dmude = np . gradient ( mu ) / np . gradient ( energy ) <EOL> group . edge_step = pre_dat [ '<STR_LIT>' ] <EOL> group . pre_edge = pre_dat [ '<STR_LIT>' ] <EOL> group . post_edge = pre_dat [ '<STR_LIT>' ] <EOL> group . pre_edge_details = Group ( ) <EOL> group . pre_edge_details . pre1 = pre_dat [ '<STR_LIT>' ] <EOL> group . pre_edge_details . pre2 = pre_dat [ '<STR_LIT>' ] <EOL> group . pre_edge_details . norm1 = pre_dat [ '<STR_LIT>' ] <EOL> group . pre_edge_details . norm2 = pre_dat [ '<STR_LIT>' ] <EOL> group . pre_edge_details . pre_slope = pre_dat [ '<STR_LIT>' ] [ <NUM_LIT:0> ] <EOL> group . pre_edge_details . pre_offset = pre_dat [ '<STR_LIT>' ] [ <NUM_LIT:1> ] <EOL> for i in range ( MAX_NNORM ) : <EOL> if hasattr ( group , '<STR_LIT>' % i ) : <EOL> delattr ( group , '<STR_LIT>' % i ) <EOL> for i , c in enumerate ( pre_dat [ '<STR_LIT>' ] ) : <EOL> setattr ( group . pre_edge_details , '<STR_LIT>' % i , c ) <EOL> return <EOL> def registerLarchPlugin ( ) : <EOL> return ( MODNAME , { '<STR_LIT>' : find_e0 , <EOL> '<STR_LIT>' : pre_edge } ) </s>
