EROSITAwiki: attachment:eROSITA_LS.py of PeriodSearch

   1 ##in bash shell on ds machines: source /utils/anaconda-setup.sh
   2 import pyfits
   3 import numpy as np
   4 from astropy.modeling import models, fitting
   5 from scipy import signal
   6 from scipy import interpolate
   7 import gatspy
   8 
   9 #optionally plotting results
  10 import matplotlib.pylab as plt
  11 
  12 def eLS(lcfil,pe_min=0.4,pe_max=11.0,nper=1e5,with_cl="yes",nbsim=100):
  13      
  14      #create period vector for LS
  15      freq = np.linspace(1./pe_max,1./pe_min,nper) 
  16      f0=1./pe_max
  17      df=freq[1]-freq[0]
  18      per_vec=1./freq
  19 
  20      #duration eROSTA gap between 2 consecutive scans (4h)
  21      dur=4*3600.
  22 
  23 
  24      #read input file, set negative rates to zero and take only valid points
  25      hdu=pyfits.open(lcfil)
  26      time=hdu[1].data.field('TIME')
  27      rate= hdu[1].data.field('RATE')
  28      rate_err= hdu[1].data.field('RATE_ERR')
  29      rate[rate<=0]=0.
  30      rate_err[rate_err<=0]=0.
  31      #otherwise should be:
  32      indok=np.where(np.isnan(rate)==False)[0]
  33      time=time[indok]
  34      rate=rate[indok]
  35      rate_err=rate_err[indok]
  36 
  37      #apply Lomb-Scargle on input LC
  38      fmodel = gatspy.periodic.LombScargleFast(fit_period=True)
  39      fmodel.optimizer.period_range = (pe_min,pe_max)
  40      fmodel.fit(time,rate,rate_err)
  41      pgram = fmodel.score_frequency_grid(f0,df,nper)
  42 
  43      #temporary best period
  44      indmax=np.argmax(pgram)
  45      bestper=per_vec[indmax]
  46 
  47      #check if single scan or multiple scans
  48      diff=np.abs(time-np.roll(time,1))
  49      diff2=diff[diff>(0.9*dur)]
  50      if np.nanmax(diff)>0.9*dur:
  51 	gaps=1
  52 	print 'Multiple scans, fitting upper envelope of Periodogram...'
  53 	if np.min(diff2)>dur:
  54 	  dur=np.min(diff2)
  55      else:
  56 	gaps=0  
  57 	print 'Single scan'
  58 
  59 
  60      #define region around best peak
  61      if bestper>0.8*pe_max: 
  62 	maxi=pe_max
  63      else:
  64 	maxi=1.2*bestper 
  65      if bestper<0.2*pe_max: 
  66 	mini=pe_min
  67      else:
  68 	mini=0.8*bestper 
  69      indfit=np.where((per_vec>mini) & (per_vec<maxi))[0]
  70      pgramfit=pgram[indfit]
  71      per_vecfit=per_vec[indfit]
  72 
  73      plt.plot(per_vec,pgram)
  74 
  75      #fitting gaussian function around best period peak  
  76      gfit = fitting.LevMarLSQFitter()  
  77      #for single scan
  78      if gaps==0:
  79 	indmaxfit=np.argmax(pgramfit)
  80 	g_init = models.Gaussian1D(amplitude=np.max(pgramfit), mean=per_vecfit[indmaxfit], stddev=0.1)  
  81 	g_init.amplitude.fixed= True 
  82 	g = gfit(g_init, per_vecfit,pgramfit)
  83 	std2=np.abs(g.stddev.value)
  84 	bestper2=g.mean.value
  85 	plt.plot(per_vecfit,g(per_vecfit),color="green")
  86      #for multiple scans   
  87      else:
  88 	maxs = signal.argrelmax(pgramfit,order=5)[0]  
  89 	spl_max = interpolate.interp1d(maxs, pgramfit[maxs], kind='linear')
  90 	max_rng = np.arange(maxs.min(), maxs.max())
  91 	u_env = spl_max(max_rng) 
  92 	indmaxfit=np.argmax(u_env)
  93 	g_init = models.Gaussian1D(amplitude=np.max(u_env), mean=max_rng[indmaxfit], stddev=100.) 
  94 	g_init.amplitude.fixed= True 
  95 	g = gfit(g_init, max_rng,u_env)
  96 	std=np.abs(g.stddev.value)
  97 	xxx=per_vecfit[max_rng]
  98 	yyy=g(max_rng)
  99 	bestper2=xxx[np.argmax(yyy)]
 100 	plt.plot(xxx,u_env,color="red")
 101 	plt.plot(xxx,yyy,color="green")
 102 	std2=bestper2-xxx[int(np.argmax(yyy)+std)]
 103 
 104 
 105      if with_cl=="yes":
 106 
 107 	 #calculate confidence levels (taking time!!!)
 108 	 #use wild bootstrap methods, consisting in flux randomization
 109 	 D = np.zeros(nbsim)	
 110 	 for ll in range(0,nbsim):
 111 	   print "simulation number:",ll
 112 	   #flux randomization
 113 	   meanlc=np.mean(rate)
 114 	   resid=rate-meanlc
 115 	   rate2=meanlc+resid*np.random.normal(0, 1, len(rate))
 116 	   rate2[rate2<=0]=0
 117 	   #Lomb Scargle periodogram
 118 	   fmodel = gatspy.periodic.LombScargleFast(fit_period=True)
 119 	   fmodel.optimizer.period_range = (pe_min,pe_max) 
 120 	   fmodel.fit(time,rate2,rate_err)
 121 	   pgram = fmodel.score_frequency_grid(f0,df,nper)
 122 	   D[ll] = pgram.max() 
 123 	 #Confidence levels at 90%, 95% and 99%
 124 	 cl=  [90, 95,99]
 125 	 z1,z2,z3 = np.percentile(D, cl)
 126 	 Dsort=np.sort(D)
 127 	 zzz=np.array([z1,z2,z3])
 128 	 xlim = (pe_min,pe_max)
 129 	 plt.xlim([pe_min,pe_max])
 130 	 for zi, pi in zip(zzz, cl):
 131 	     plt.plot(xlim, (zi, zi), ':k', lw=1)
 132 	     plt.text(xlim[1] - 0.001, zi, "%d%%" % pi, ha='right', va='top')
 133 
 134      else:
 135          zzz=np.array([np.nan,np.nan,np.nan])
 136      plt.show() 
 137      
 138      return bestper2, std2, zzz