sampling.py

import numpy as np
#np.random.seed(101111)
import matplotlib.pyplot as plt
plt.close('all')
from tedi import utils

def SamplingSun(initFile, finalFile, timespan, observationsNumber,
                planet = False, planetParams = [], saveFile = True):
    """ Parameters:
            initFile = cleaned SunAsAStar file
            finalFile = name of output Sun file
            timespan = timespan of the observations
            observationsNumber = Number of observations wated
            planet = True if you want to include a planet to the observations
            planateParams = [P, K, e, w, phi] of the planet to be added
            saveFile = True to save a .txt file with the results
        Returns:
            Time, RVfinal, RVerr, Bisobs, Biserr, FWHMobs, FWHMerr
    """
    #timespan need to be bigger than the number of observations
    assert timespan >= observationsNumber, \
    'The timespan need to be bigger than the number of observations!'
    
    #Cleaned Sun data
    time,rv,rverr,bis,biserr,fwhm,fwhmerr = np.loadtxt(initFile, skiprows=2, 
                                                       unpack=True, 
                                                       usecols=(0,1,2,3,4,5,6))
    #Dealing with the timespan
    initPosition = np.random.randint(0, time.size-timespan)
    initTime = int(time[initPosition])
    finalTime = initTime + timespan
    positions = np.where((time >= initTime) & (time <= finalTime))
    
    time = time[positions]
    rv = rv[positions]
    rverr = rverr[positions]
    bis = bis[positions]
    biserr = biserr[positions]
    fwhm = fwhm[positions]
    fwhmerr = fwhmerr[positions]
    
    # Sampling in diferent days
    pos = np.array([])
    timecheck = np.array([])
    i = 0
    while i < observationsNumber:
        randomValue = np.random.randint(0, rv.size)
        if (randomValue in pos) or (time[randomValue].astype(int) in timecheck):
            randomValue = np.random.randint(0, rv.size)
        else:
            pos = np.append(pos, randomValue)
            timecheck = np.append(timecheck, time[randomValue].astype(int)).astype(int)
            i += 1
    
    pos = np.sort(pos).astype(int)
    Time = time[pos]
    RVfinal = rv[pos]
    RVerr = rverr[pos]
    FWHMobs = fwhm[pos]
    FWHMerr = fwhmerr[pos]
    Bisobs = bis[pos]
    Biserr = biserr[pos]
    
    if planet:
        assert len(planetParams) == 5, \
        'Missing data on planetParams!'
        ### Adding one planet
        P, K, e, w, phi = planetParams
        planet = utils.keplerian(P = P, K = K, e = e, w = w, phi = phi, 
                                 t = Time)
        RVfinal = RVfinal + planet[1]
    
    if saveFile:
        results = np.stack((Time, RVfinal, RVerr, Bisobs, Biserr, FWHMobs, FWHMerr))
        results = results.T
        np.savetxt(finalFile, results, delimiter='\t',
                   header ="BJD\tRV\tRVerr\tBIS\tBISerr\tFWHM\tFWHMerr", 
                   comments='')
    return Time, RVfinal, RVerr, Bisobs, Biserr, FWHMobs, FWHMerr
    
    
    
#### Testing
initFile = "CleanedSun.txt"
finalFile = "SunDataToWork.txt"
data = SamplingSun(initFile, finalFile, timespan=100, observationsNumber=50,
                   planet=True, planetParams=[365,1,0,0,0])

plt.figure()
plt.title('RV data')
plt.plot(data[0], data[1], '.')
plt.xlabel('Time (BJD)')
plt.ylabel('RV (m/s)')