<s> from glue . core . fitters import BaseFitter1D <EOL> from glue . core . simpleforms import IntOption <EOL> from glue . config import fit_plugin <EOL> import numpy as np <EOL> import emcee <EOL> def gaussian ( x , mean , amplitude , stddev ) : <EOL> return np . exp ( - ( x - mean ) ** <NUM_LIT:2> / ( <NUM_LIT:2> * stddev ** <NUM_LIT:2> ) ) * amplitude <EOL> def lnprob ( param , x , y , dy ) : <EOL> if param [ - <NUM_LIT:1> ] < <NUM_LIT:0> : <EOL> return - np . inf <EOL> yp = gaussian ( x , * param ) <EOL> diff = ( y - yp ) <EOL> if dy is not None : <EOL> diff /= dy <EOL> return - np . dot ( diff , diff ) <EOL> @ fit_plugin <EOL> class EmceeGaussianFitter ( BaseFitter1D ) : <EOL> label = "<STR_LIT>" <EOL> walkers = IntOption ( label = "<STR_LIT>" , min = <NUM_LIT:1> , max = <NUM_LIT:200> , default = <NUM_LIT:50> ) <EOL> burnin = IntOption ( label = "<STR_LIT>" , min = <NUM_LIT:1> , max = <NUM_LIT> , default = <NUM_LIT> ) <EOL> steps = IntOption ( label = "<STR_LIT>" , min = <NUM_LIT:10> , max = <NUM_LIT> , default = <NUM_LIT> ) <EOL> def fit ( self , x , y , dy , constraints , <EOL> walkers = <NUM_LIT:50> , burnin = <NUM_LIT> , steps = <NUM_LIT> ) : <EOL> ndim = <NUM_LIT:3> <EOL> amp = y . max ( ) <EOL> mean = ( x * y ) . sum ( ) / y . sum ( ) <EOL> var = ( ( x - mean ) ** <NUM_LIT:2> * y ) . sum ( ) / y . sum ( ) <EOL> p0 = np . array ( [ amp , mean , np . sqrt ( var ) ] ) . reshape ( <NUM_LIT:1> , - <NUM_LIT:1> ) <EOL> p0 = np . random . lognormal ( sigma = <NUM_LIT> , size = ( walkers , ndim ) ) * p0 <EOL> sampler = emcee . EnsembleSampler ( walkers , ndim , lnprob , <EOL> args = [ x , y , dy ] ) <EOL> pos , prob , state = sampler . run_mcmc ( p0 , burnin ) <EOL> sampler . reset ( ) <EOL> sampler . run_mcmc ( pos , steps ) <EOL> return sampler <EOL> def predict ( self , fit_result , x ) : <EOL> chain = fit_result . flatchain <EOL> params = np . mean ( chain , axis = <NUM_LIT:0> ) <EOL> return gaussian ( x , * params ) <EOL> def summarize ( self , fit_result , x , y , dy ) : <EOL> af = fit_result . acceptance_fraction . mean ( ) <EOL> chain = fit_result . flatchain <EOL> amp , mean , sigma = chain . mean ( axis = <NUM_LIT:0> ) <EOL> damp , dmean , dsigma = np . std ( chain , axis = <NUM_LIT:0> ) <EOL> walkers , steps , dim = fit_result . chain . shape <EOL> result = [ <EOL> "<STR_LIT>" % walkers , <EOL> "<STR_LIT>" % steps , <EOL> "<STR_LIT>" % af , <EOL> "<STR_LIT>" , <EOL> "<STR_LIT>" % ( amp , damp ) , <EOL> "<STR_LIT>" % ( mean , dmean ) , <EOL> "<STR_LIT>" % ( sigma , dsigma ) <EOL> ] <EOL> return '<STR_LIT:\n>' . join ( result ) <EOL> def plot ( self , fit_result , axes , x ) : <EOL> chain = fit_result . flatchain <EOL> result = [ ] <EOL> for i in range ( <NUM_LIT:100> ) : <EOL> row = np . random . randint ( <NUM_LIT:0> , chain . shape [ <NUM_LIT:0> ] ) <EOL> params = chain [ row ] <EOL> y = gaussian ( x , * params ) <EOL> result . extend ( axes . plot ( x , y , '<STR_LIT:k>' , alpha = <NUM_LIT> ) ) <EOL> result . extend ( <EOL> super ( EmceeGaussianFitter , self ) . plot ( fit_result , axes , x ) ) <EOL> return result </s>
