Package ComboCode :: Package cc :: Package ivs :: Package sigproc :: Module funclib
[hide private]
[frames] | no frames]

Source Code for Module ComboCode.cc.ivs.sigproc.funclib

  1  """ 
  2  Database with model functions. 
  3   
  4  To be used with the L{cc.ivs.sigproc.fit.minimizer} function or with the L{evaluate} 
  5  function in this module. 
  6   
  7  >>> p = plt.figure() 
  8  >>> x = np.linspace(-10,10,1000) 
  9  >>> p = plt.plot(x,evaluate('gauss',x,[5,1.,2.,0.5]),label='gauss') 
 10  >>> p = plt.plot(x,evaluate('voigt',x,[20.,1.,1.5,3.,0.5]),label='voigt') 
 11  >>> p = plt.plot(x,evaluate('lorentz',x,[5,1.,2.,0.5]),label='lorentz') 
 12  >>> leg = plt.legend(loc='best') 
 13  >>> leg.get_frame().set_alpha(0.5) 
 14   
 15  ]include figure]]ivs_sigproc_fit_funclib01.png] 
 16   
 17  >>> p = plt.figure() 
 18  >>> x = np.linspace(0,10,1000)[1:] 
 19  >>> p = plt.plot(x,evaluate('power_law',x,[2.,3.,1.5,0,0.5]),label='power_law') 
 20  >>> p = plt.plot(x,evaluate('power_law',x,[2.,3.,1.5,0,0.5])+evaluate('gauss',x,[1.,5.,0.5,0,0]),label='power_law + gauss') 
 21  >>> leg = plt.legend(loc='best') 
 22  >>> leg.get_frame().set_alpha(0.5) 
 23   
 24  ]include figure]]ivs_sigproc_fit_funclib02.png] 
 25   
 26  >>> p = plt.figure() 
 27  >>> x = np.linspace(0,10,1000) 
 28  >>> p = plt.plot(x,evaluate('sine',x,[1.,2.,0,0]),label='sine') 
 29  >>> p = plt.plot(x,evaluate('sine_linfreqshift',x,[1.,0.5,0,0,.5]),label='sine_linfreqshift') 
 30  >>> p = plt.plot(x,evaluate('sine_expfreqshift',x,[1.,0.5,0,0,1.2]),label='sine_expfreqshift') 
 31  >>> leg = plt.legend(loc='best') 
 32  >>> leg.get_frame().set_alpha(0.5) 
 33   
 34  ]include figure]]ivs_sigproc_fit_funclib03.png] 
 35   
 36  >>> p = plt.figure() 
 37  >>> p = plt.plot(x,evaluate('sine',x,[1.,2.,0,0]),label='sine') 
 38  >>> p = plt.plot(x,evaluate('sine_orbit',x,[1.,2.,0,0,0.1,10.,0.1]),label='sine_orbit') 
 39  >>> leg = plt.legend(loc='best') 
 40  >>> leg.get_frame().set_alpha(0.5) 
 41   
 42  ]include figure]]ivs_sigproc_fit_funclib03a.png] 
 43   
 44  >>> p = plt.figure() 
 45  >>> x_single = np.linspace(0,10,1000) 
 46  >>> x_double = np.vstack([x_single,x_single]) 
 47  >>> p = plt.plot(x_single,evaluate('kepler_orbit',x_single,[2.5,0.,0.5,0,3,1.]),label='kepler_orbit (single)') 
 48  >>> y_double = evaluate('kepler_orbit',x_double,[2.5,0.,0.5,0,3,2.,-4,2.],type='double') 
 49  >>> p = plt.plot(x_double[0],y_double[0],label='kepler_orbit (double 1)') 
 50  >>> p = plt.plot(x_double[1],y_double[1],label='kepler_orbit (double 2)') 
 51  >>> p = plt.plot(x,evaluate('box_transit',x,[2.,0.4,0.1,0.3,0.5]),label='box_transit') 
 52  >>> leg = plt.legend(loc='best') 
 53  >>> leg.get_frame().set_alpha(0.5) 
 54   
 55  ]include figure]]ivs_sigproc_fit_funclib04.png] 
 56   
 57  >>> p = plt.figure() 
 58  >>> x = np.linspace(-1,1,1000) 
 59  >>> gammas = [-0.25,0.1,0.25,0.5,1,2,4] 
 60  >>> y = np.array([evaluate('soft_parabola',x,[1.,0,1.,gamma]) for gamma in gammas]) 
 61  divide by zero encountered in power 
 62  >>> for iy,gamma in zip(y,gammas): p = plt.plot(x,iy,label="soft_parabola $\gamma$={:.2f}".format(gamma)) 
 63  >>> leg = plt.legend(loc='best') 
 64  >>> leg.get_frame().set_alpha(0.5) 
 65   
 66  ]include figure]]ivs_sigproc_fit_funclib05.png] 
 67   
 68  >>> p = plt.figure() 
 69  >>> x = np.logspace(-1,2,1000) 
 70  >>> blbo = evaluate('blackbody',x,[10000.,1.],wave_units='micron',flux_units='W/m3') 
 71  >>> raje = evaluate('rayleigh_jeans',x,[10000.,1.],wave_units='micron',flux_units='W/m3') 
 72  >>> wien = evaluate('wien',x,[10000.,1.],wave_units='micron',flux_units='W/m3') 
 73  >>> p = plt.subplot(221) 
 74  >>> p = plt.title(r'$\lambda$ vs $F_\lambda$') 
 75  >>> p = plt.loglog(x,blbo,label='Black Body') 
 76  >>> p = plt.loglog(x,raje,label='Rayleigh-Jeans') 
 77  >>> p = plt.loglog(x,wien,label='Wien') 
 78  >>> leg = plt.legend(loc='best') 
 79  >>> leg.get_frame().set_alpha(0.5) 
 80   
 81  >>> blbo = evaluate('blackbody',x,[10000.,1.],wave_units='micron',flux_units='Jy') 
 82  >>> raje = evaluate('rayleigh_jeans',x,[10000.,1.],wave_units='micron',flux_units='Jy') 
 83  >>> wien = evaluate('wien',x,[10000.,1.],wave_units='micron',flux_units='Jy') 
 84  >>> p = plt.subplot(223) 
 85  >>> p = plt.title(r"$\lambda$ vs $F_\\nu$") 
 86  >>> p = plt.loglog(x,blbo,label='Black Body') 
 87  >>> p = plt.loglog(x,raje,label='Rayleigh-Jeans') 
 88  >>> p = plt.loglog(x,wien,label='Wien') 
 89  >>> leg = plt.legend(loc='best') 
 90  >>> leg.get_frame().set_alpha(0.5) 
 91   
 92  >>> x = np.logspace(0.47,3.47,1000) 
 93  >>> blbo = evaluate('blackbody',x,[10000.,1.],wave_units='THz',flux_units='Jy') 
 94  >>> raje = evaluate('rayleigh_jeans',x,[10000.,1.],wave_units='THz',flux_units='Jy') 
 95  >>> wien = evaluate('wien',x,[10000.,1.],wave_units='THz',flux_units='Jy') 
 96  >>> p = plt.subplot(224) 
 97  >>> p = plt.title(r"$\\nu$ vs $F_\\nu$") 
 98  >>> p = plt.loglog(x,blbo,label='Black Body') 
 99  >>> p = plt.loglog(x,raje,label='Rayleigh-Jeans') 
100  >>> p = plt.loglog(x,wien,label='Wien') 
101  >>> leg = plt.legend(loc='best') 
102  >>> leg.get_frame().set_alpha(0.5) 
103   
104  >>> blbo = evaluate('blackbody',x,[10000.,1.],wave_units='THz',flux_units='W/m3') 
105  >>> raje = evaluate('rayleigh_jeans',x,[10000.,1.],wave_units='THz',flux_units='W/m3') 
106  >>> wien = evaluate('wien',x,[10000.,1.],wave_units='THz',flux_units='W/m3') 
107  >>> p = plt.subplot(222) 
108  >>> p = plt.title(r"$\\nu$ vs $F_\lambda$") 
109  >>> p = plt.loglog(x,blbo,label='Black Body') 
110  >>> p = plt.loglog(x,raje,label='Rayleigh-Jeans') 
111  >>> p = plt.loglog(x,wien,label='Wien') 
112  >>> leg = plt.legend(loc='best') 
113  >>> leg.get_frame().set_alpha(0.5) 
114   
115  ]include figure]]ivs_sigproc_fit_funclib06.png] 
116   
117  """ 
118  import numpy as np 
119  from numpy import pi,cos,sin,sqrt,tan,arctan 
120  from scipy.special import erf,jn 
121  from cc.ivs.sigproc.fit import Model, Function 
122  #import ivs.timeseries.keplerorbit as kepler 
123  from cc.ivs.sed import model as sed_model 
124  import logging 
125   
126  logger = logging.getLogger("SP.FUNCLIB") 
127   
128  #{ Function Library 
129 -def blackbody(wave_units='AA',flux_units='erg/s/cm2/AA',disc_integrated=True):
130 """ 131 Blackbody (T, scale). 132 133 @param wave_units: wavelength units 134 @type wave_units: string 135 @param flux_units: flux units 136 @type flux_units: string 137 @param disc_integrated: sets units equal to SED models 138 @type disc_integrated: bool 139 """ 140 pnames = ['T', 'scale'] 141 function = lambda p, x: p[1]*sed_model.blackbody(x, p[0], wave_units=wave_units,\ 142 flux_units=flux_units,\ 143 disc_integrated=disc_integrated) 144 function.__name__ = 'blackbody' 145 146 return Function(function=function, par_names=pnames)
147
148 -def rayleigh_jeans(wave_units='AA',flux_units='erg/s/cm2/AA',disc_integrated=True):
149 """ 150 Rayleigh-Jeans tail (T, scale). 151 152 @param wave_units: wavelength units 153 @type wave_units: string 154 @param flux_units: flux units 155 @type flux_units: string 156 @param disc_integrated: sets units equal to SED models 157 @type disc_integrated: bool 158 """ 159 pnames = ['T', 'scale'] 160 function = lambda p, x: p[1]*sed_model.rayleigh_jeans(x, p[0], wave_units=wave_units,\ 161 flux_units=flux_units,\ 162 disc_integrated=disc_integrated) 163 function.__name__ = 'rayleigh_jeans' 164 165 return Function(function=function, par_names=pnames)
166
167 -def wien(wave_units='AA',flux_units='erg/s/cm2/AA',disc_integrated=True):
168 """ 169 Wien approximation (T, scale). 170 171 @param wave_units: wavelength units 172 @type wave_units: string 173 @param flux_units: flux units 174 @type flux_units: string 175 @param disc_integrated: sets units equal to SED models 176 @type disc_integrated: bool 177 """ 178 pnames = ['T', 'scale'] 179 function = lambda p, x: p[1]*sed_model.wien(x, p[0], wave_units=wave_units,\ 180 flux_units=flux_units,\ 181 disc_integrated=disc_integrated) 182 function.__name__ = 'wien' 183 184 return Function(function=function, par_names=pnames)
185 186 # def kepler_orbit(type='single'): 187 # """ 188 # Kepler orbits ((p,t0,e,omega,K,v0) or (p,t0,e,omega,K1,v01,K2,v02)) 189 # 190 # A single kepler orbit 191 # parameters are: [p, t0, e, omega, k, v0] 192 # 193 # A double kepler orbit 194 # parameters are: [p, t0, e, omega, k_1, v0_1, k_2, v0_2] 195 # Warning: This function uses 2d input and output! 196 # """ 197 # if type == 'single': 198 # pnames = ['p','t0','e','omega','k','v0'] 199 # function = lambda p, x: kepler.radial_velocity(p, times=x, itermax=8) 200 # function.__name__ = 'kepler_orbit_single' 201 # 202 # return Function(function=function, par_names=pnames) 203 # 204 # elif type == 'double': 205 # pnames = ['p','t0','e','omega','k1','v01','k2','v02' ] 206 # function = lambda p, x: [kepler.radial_velocity([p[0],p[1],p[2],p[3],p[4],p[5]], times=x[0], itermax=8), 207 # kepler.radial_velocity([p[0],p[1],p[2],p[3],p[6],p[7]], times=x[1], itermax=8)] 208 # def residuals(syn, data, weights=None, errors=None, **kwargs): 209 # return np.hstack( [( data[0] - syn[0] ) * weights[0], ( data[1] - syn[1] ) * weights[1] ] ) 210 # function.__name__ = 'kepler_orbit_double' 211 # 212 # return Function(function=function, par_names=pnames, resfunc=residuals) 213
214 -def box_transit(t0=0.):
215 """ 216 Box transit model (cont,freq,ingress,egress,depth) 217 218 @param t0: reference time (defaults to 0) 219 @type t0: float 220 """ 221 pnames = 'cont','freq','ingress','egress','depth' 222 def function(p,x): 223 cont,freq,ingress,egress,depth = p 224 model = np.ones(len(x))*cont 225 phase = np.fmod((x - t0) * freq, 1.0) 226 phase = np.where(phase<0,phase+1,phase) 227 transit_place = (ingress<=phase) & (phase<=egress) 228 model = np.where(transit_place,model-depth,model) 229 return model
230 function.__name__ = 'box_transit' 231 return Function(function=function, par_names=pnames) 232 233
234 -def polynomial(d=1):
235 """ 236 Polynomial (a1,a0). 237 238 y(x) = ai*x**i + a(i-1)*x**(i-1) + ... + a1*x + a0 239 240 @param d: degree of the polynomial 241 @type d: int 242 """ 243 pnames = ['a{:d}'.format(i) for i in range(d,0,-1)]+['a0'] 244 function = lambda p, x: np.polyval(p,x) 245 function.__name__ = 'polynomial' 246 247 return Function(function=function, par_names=pnames)
248 249
250 -def soft_parabola():
251 """ 252 Soft parabola (ta,vlsr,vinf,gamma). 253 254 See Olofsson 1993ApJS...87...267O. 255 256 T_A(x) = T_A(0) * [ 1 - ((x- v_lsr)/v_inf)**2 ] ** (gamma/2) 257 """ 258 pnames = ['ta','vlsr','vinf','gamma'] 259 def function(p,x): 260 term = (x-p[1]) / p[2] 261 y = p[0] * (1- term**2)**(p[3]/2.) 262 if p[3]<=0: y[np.abs(term)>=1] = 0 263 y[np.isnan(y)] = 0 264 return y
265 function.__name__ = 'soft_parabola' 266 267 return Function(function=function, par_names=pnames) 268 269
270 -def gauss(use_jacobian=True):
271 """ 272 Gaussian (a,mu,sigma,c) 273 274 f(x) = a * exp( - (x - mu)**2 / (2 * sigma**2) ) + c 275 """ 276 pnames = ['a', 'mu', 'sigma', 'c'] 277 function = lambda p, x: p[0] * np.exp( -(x-p[1])**2 / (2.0*p[2]**2)) + p[3] 278 function.__name__ = 'gauss' 279 280 if not use_jacobian: 281 return Function(function=function, par_names=pnames) 282 else: 283 def jacobian(p, x): 284 ex = np.exp( -(x-p[1])**2 / (2.0*p[2]**2) ) 285 return np.array([-ex, -p[0] * (x-p[1]) * ex / p[2]**2, -p[0] * (x-p[1])**2 * ex / p[2]**3, [-1 for i in x] ]).T
286 return Function(function=function, par_names=pnames, jacobian=jacobian) 287
288 -def sine():
289 """ 290 Sine (ampl,freq,phase,const) 291 292 f(x) = ampl * sin(2pi*freq*x + 2pi*phase) + const 293 """ 294 pnames = ['ampl', 'freq', 'phase', 'const'] 295 function = lambda p, x: p[0] * sin(2*pi*(p[1]*x + p[2])) + p[3] 296 function.__name__ = 'sine' 297 298 return Function(function=function, par_names=pnames)
299 300
301 -def sine_linfreqshift(t0=0.):
302 """ 303 Sine with linear frequency shift (ampl,freq,phase,const,D). 304 305 Similar to C{sine}, but with extra parameter 'D', which is the linear 306 frequency shift parameter. 307 308 @param t0: reference time (defaults to 0) 309 @type t0: float 310 """ 311 pnames = ['ampl', 'freq', 'phase', 'const','D'] 312 def function(p,x): 313 freq = (p[1] + p[4]/2.*(x-t0))*(x-t0) 314 return p[0] * sin(2*pi*(freq + p[2])) + p[3]
315 function.__name__ = 'sine_linfreqshift' 316 return Function(function=function, par_names=pnames) 317
318 -def sine_expfreqshift(t0=0.):
319 """ 320 Sine with exponential frequency shift (ampl,freq,phase,const,K). 321 322 Similar to C{sine}, but with extra parameter 'K', which is the exponential 323 frequency shift parameter. 324 325 frequency(x) = freq / log(K) * (K**(x-t0)-1) 326 327 f(x) = ampl * sin( 2*pi * (frequency + phase)) 328 329 @param t0: reference time (defaults to 0) 330 @type t0: float 331 """ 332 pnames = ['ampl', 'freq', 'phase', 'const','K'] 333 def function(p,x): 334 freq = p[1] / np.log(p[4]) * (p[4]**(x-t0)-1.) 335 return p[0] * sin(2*pi*(freq + p[2])) + p[3]
336 function.__name__ = 'sine_expfreqshift' 337 return Function(function=function, par_names=pnames) 338 339
340 -def sine_orbit(t0=0.,nmax=10):
341 """ 342 Sine with a sinusoidal frequency shift (ampl,freq,phase,const,forb,asini,omega,(,ecc)) 343 344 Similar to C{sine}, but with extra parameter 'asini' and 'forb', which are 345 the orbital parameters. forb in cycles/day or something similar, asini in au. 346 347 For eccentric orbits, add longitude of periastron 'omega' (radians) and 348 'ecc' (eccentricity). 349 350 @param t0: reference time (defaults to 0) 351 @type t0: float 352 @param nmax: number of terms to include in series for eccentric orbit 353 @type nmax: int 354 """ 355 pnames = ['ampl', 'freq', 'phase', 'const','forb','asini','omega'] 356 def ane(n,e): return 2.*sqrt(1-e**2)/e/n*jn(n,n*e) 357 def bne(n,e): return 1./n*(jn(n-1,n*e)-jn(n+1,n*e)) 358 def function(p,x): 359 ampl,freq,phase,const,forb,asini,omega = p[:7] 360 ecc = None 361 if len(p)==8: 362 ecc = p[7] 363 cc = 173.144632674 # speed of light in AU/d 364 alpha = freq*asini/cc 365 if ecc is None: 366 frequency = freq*(x-t0) + alpha*(sin(2*pi*forb*x) - sin(2*pi*forb*t0)) 367 else: 368 ns = np.arange(1,nmax+1,1) 369 ans,bns = np.array([[ane(n,ecc),bne(n,ecc)] for n in ns]).T 370 ksins = sqrt(ans**2*cos(omega)**2+bns**2*sin(omega)**2) 371 thns = arctan(bns/ans*tan(omega)) 372 tau = -np.sum(bns*sin(omega)) 373 frequency = freq*(x-t0) + \ 374 alpha*(np.sum(np.array([ksins[i]*sin(2*pi*ns[i]*forb*(x-t0)+thns[i]) for i in range(nmax)]),axis=0)+tau) 375 return ampl * sin(2*pi*(frequency + phase)) + const
376 function.__name__ = 'sine_orbit' 377 return Function(function=function, par_names=pnames) 378 379 #def generic(func_name): 380 ##func = model.FUNCTION(function=getattr(evaluate,func_name), par_names) 381 #raise NotImplementedError 382
383 -def power_law():
384 """ 385 Power law (A,B,C,f0,const) 386 387 P(f) = A / (1+ B(f-f0))**C + const 388 """ 389 pnames = ['ampl','b','c','f0','const'] 390 function = lambda p, x: p[0] / (1 + (p[1]*(x-p[3]))**p[2]) + p[4] 391 function.__name__ = 'power_law' 392 return Function(function=function, par_names=pnames)
393 394
395 -def lorentz():
396 """ 397 Lorentz profile (ampl,mu,gamma,const) 398 399 P(f) = A / ( (x-mu)**2 + gamma**2) + const 400 """ 401 pnames = ['ampl','mu','gamma','const'] 402 function = lambda p,x: p[0] / ((x-p[1])**2 + p[2]**2) + p[3] 403 function.__name__ = 'lorentz' 404 return Function(function=function, par_names=pnames)
405
406 -def voigt():
407 """ 408 Voigt profile (ampl,mu,sigma,gamma,const) 409 410 z = (x + gamma*i) / (sigma*sqrt(2)) 411 V = A * Real[cerf(z)] / (sigma*sqrt(2*pi)) 412 """ 413 pnames = ['ampl','mu','sigma','gamma','const'] 414 def function(p,x): 415 x = x-p[1] 416 z = (x+1j*p[3])/(p[2]*sqrt(2)) 417 return p[0]*_complex_error_function(z).real/(p[2]*sqrt(2*pi))+p[4]
418 function.__name__ = 'voigt' 419 return Function(function=function, par_names=pnames) 420 421 #} 422 423 #{ Combination functions 424
425 -def multi_sine(n=10):
426 """ 427 Multiple sines. 428 429 @param n: number of sines 430 @type n: int 431 """ 432 return Model(functions=[sine() for i in range(n)])
433
434 -def multi_blackbody(n=3,**kwargs):
435 """ 436 Multiple black bodies. 437 438 @param n: number of blackbodies 439 @type n: int 440 """ 441 return Model(functions=[blackbody(**kwargs) for i in range(n)])
442 443 #} 444 445 #{ Internal Helper functions 446
447 -def _complex_error_function(x):
448 """ 449 Complex error function 450 """ 451 cef_value = np.exp(-x**2)*(1-erf(-1j*x)) 452 if sum(np.isnan(cef_value))>0: 453 logging.warning("Complex Error function: NAN encountered, results are biased") 454 noisnans = np.compress(1-np.isnan(cef_value),cef_value) 455 try: 456 last_value = noisnans[-1] 457 except: 458 last_value = 0 459 logging.warning("Complex Error function: all values are NAN, results are wrong") 460 cef_value = np.where(np.isnan(cef_value),last_value,cef_value) 461 462 return cef_value
463 464 #} 465 466
467 -def evaluate(funcname, domain, parameters, **kwargs):
468 """ 469 Evaluate a function on specified interval with specified parameters. 470 471 Extra keywords are passed to the funcname. 472 473 Example: 474 475 >>> x = np.linspace(-5,5,1000) 476 >>> y = evaluate('gauss',x,[1.,0.,2.,0.]) 477 478 @parameter funcname: name of the function 479 @type funcname: str 480 @parameter domain: domain to evaluate onto 481 @type domain: array 482 @parameter parameters: parameter of the function 483 @type parameters: array 484 """ 485 function = globals()[funcname](**kwargs) 486 function.setup_parameters(parameters) 487 return function.evaluate(domain)
488 489 if __name__=="__main__": 490 import doctest 491 from matplotlib import pyplot as plt 492 doctest.testmod() 493 plt.show() 494