diff --git a/MANIFEST.in b/MANIFEST.in new file mode 100644 index 0000000..51930f2 --- /dev/null +++ b/MANIFEST.in @@ -0,0 +1 @@ +recursive-include src/mors/data * diff --git a/README.md b/README.md index 344d0e2..3417f2c 100644 --- a/README.md +++ b/README.md @@ -62,8 +62,7 @@ and add the export command to the bottom of the file. After that, run the following command in a python environment or at the beginning of your python script ``` -from mors.data import DownloadEvolutionTracks -DownloadEvolutionTracks() +mors.DownloadEvolutionTracks() ``` This will download and extract package stellar evolution tracks data. diff --git a/pyproject.toml b/pyproject.toml index 5a405f0..f71e1d9 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -27,8 +27,7 @@ dependencies = [ 'matplotlib', 'numpy', 'osfclient', - 'scipy', - 'tqdm' + 'scipy' ] [project.urls] diff --git a/src/mors/__init__.py b/src/mors/__init__.py index ccd98be..7147762 100644 --- a/src/mors/__init__.py +++ b/src/mors/__init__.py @@ -25,6 +25,12 @@ # Spectral synthesis from .synthesis import * +# Data download +from .data import DownloadEvolutionTracks + +# Baraffe tracks +from .baraffe import * + # Spectrum stuff from .spectrum import * diff --git a/src/mors/baraffe.py b/src/mors/baraffe.py new file mode 100644 index 0000000..ae053a9 --- /dev/null +++ b/src/mors/baraffe.py @@ -0,0 +1,241 @@ +"""Module for loading and interpolating Baraffe tracks data. +Original tracks data can be found on the website +http://perso.ens-lyon.fr/isabelle.baraffe/BHAC15dir/BHAC15_tracks+structure""" +import numpy as np +import os +from scipy.interpolate import PchipInterpolator + +import mors.constants as const + +#Short cut to Baraffe tracks mass and temporal range +MassGrid = [0.010, 0.015, 0.020, 0.030, 0.040, 0.050, + 0.060, 0.070, 0.072, 0.075, 0.080, 0.090, + 0.100, 0.110, 0.130, 0.150, 0.170, 0.200, + 0.300, 0.400, 0.500, 0.600, 0.700, 0.800, + 0.900, 1.000, 1.100, 1.200, 1.300, 1.400 ] +tminGrid = [ 5.695210, 5.693523, 5.689884, 5.695099, 5.694350, 5.694123, + 5.694808, 5.692913, 5.694756, 5.694352, 5.694252, 5.689450, + 5.703957, 5.709729, 5.702708, 5.701659, 5.693846, 5.693823, + 5.690044, 5.689995, 5.694573, 5.690793, 5.691618, 5.693354, + 5.692710, 5.693063, 5.694315, 5.693263, 5.694626, 5.692977 ] +tmaxGrid = [ 7.612341, 8.050550, 8.089757, 8.514580, 8.851534, 9.178493, + 9.499851,10.000078, 9.999692,10.000130,10.000343, 9.999745, + 9.999914, 9.999209, 9.999428,10.000104, 9.999246, 9.999735, + 10.000004, 9.999193, 9.999099, 9.999787, 9.999811, 9.998991, + 10.000065, 9.920501, 9.789029, 9.651666, 9.538982, 9.425490 ] + +fwl_data_dir = os.getenv('FWL_DATA') + +class BaraffeTrack: + """Class to hold interpolated baraffe tracks data for a given star mass + + Attributes + ---------- + mstar : float + Star mass + track : dict + Dictionary containing track data + """ + + def __init__(self, Mstar): + + self.mstar = Mstar + + #If Mstar is out of the mass range, code breaks + if(Mstar < MassGrid[0]): + raise Exception("Stellar mass is too low for the Baraffe tracks") + elif(Mstar > MassGrid[29]): + raise Exception("Stellar mass is too high for the Baraffe tracks") + + #If Mstar matches with values in the mass grid + elif(Mstar in MassGrid): + track = BaraffeLoadTrack(Mstar) + + #If Mstar is between two points in the mass grid, need mass interpolation + else: + for i in range(29): + if not Mstar > MassGrid[i+1]: + #search for common temporal range + tmin = max(tminGrid[i],tminGrid[i+1]) + tmax = min(tmaxGrid[i],tmaxGrid[i+1]) + + #load neighbouring tracks with time interpolation on the same time grid + track = BaraffeLoadTrack(MassGrid[i ], tmin=tmin, tmax=tmax) + trackp = BaraffeLoadTrack(MassGrid[i+1], tmin=tmin, tmax=tmax) + + #perform linear mass interpolation for each array + mass_ratio=(Mstar-MassGrid[i])/(MassGrid[i+1]-MassGrid[i]) + track['Teff' ]=(trackp['Teff' ]-track['Teff' ])*mass_ratio + track['Teff' ] + track['Lstar']=(trackp['Lstar']-track['Lstar'])*mass_ratio + track['Lstar'] + track['Rstar']=(trackp['Rstar']-track['Rstar'])*mass_ratio + track['Rstar'] + + break + + self.track = track + self.tmin = track['t'][0] + self.tmax = track['t'][-1] + + return + + def BaraffeLuminosity(self, tstar): + """Calculates the star luminosity at a given time. + + Parameters + ---------- + tstar : float + Star's age + + Returns + ---------- + Lstar : float + Luminosity Flux in units of solar luminosity + """ + + # Get time and check that it is in range + if (tstar < self.tmin): + print("Star age too low! Clipping to %.1g Myr" % int(self.tmin*1.e-6)) + tstar = self.tmin + if (tstar > self.tmax): + print("Star age too high! Clipping to %.1g Myr" % int(self.tmax*1.e-6)) + tstar = self.tmax + + # Find closest row in track + iclose = (np.abs(self.track['t'] - tstar)).argmin() + + # Get data from track + Lstar = self.track['Lstar'][iclose] + + return Lstar + + def BaraffeSolarConstant(self, tstar, mean_distance): + """Calculates the bolometric flux of the star at a given time. + Flux is scaled to the star-planet distance. + + Parameters + ---------- + tstar : float + Star's age + mean_distance : float + Star-planet distance + + Returns + ---------- + inst : float + Flux at planet's orbital separation (solar constant) in W/m^2 + """ + + Lstar = self.BaraffeLuminosity(tstar) + Lstar *= const.LbolSun_SI + mean_distance *= const.AU_SI + + inst = Lstar / ( 4. * np.pi * mean_distance * mean_distance ) + + return inst + + def BaraffeStellarRadius(self, tstar): + """Calculates the star's radius at a time t. + + Parameters + ---------- + tstar : float + Star's age + + Returns + ---------- + Rstar : float + Radius of star in units of solar radius + """ + + # Get time and check that it is in range + if (tstar < self.tmin): + print("Star age too low! Clipping to %.1g Myr" % int(self.tmin*1.e-6)) + tstar = self.tmin + if (tstar > self.tmax): + print("Star age too high! Clipping to %.1g Myr" % int(self.tmax*1.e-6)) + tstar = self.tmax + + # Find closest row in track + iclose = (np.abs(self.track['t'] - tstar)).argmin() + + # Get data from track + Rstar = self.track['Rstar'][iclose] + + return Rstar + +def BaraffeLoadTrack(Mstar, pre_interp = True, tmin = None, tmax = None): + """Load a baraffe track into memory and optionally interpolate it into a fine time-grid + + Parameters + ---------- + Mstar : float + Star mass (in unit of solar mass) + It assumes the value has been checked and matches with the mass grid. + pre_interp : bool + Pre-interpolate the tracks onto a higher resolution time-grid + tmin : float + Minimum value of the temporal grid + tmax : float + Maximum value of the temporal grid + + Returns + ---------- + track : dict + Dictionary containing track data + """ + + # Load data + formatted_mass = f"{Mstar:.3f}" + file = (fwl_data_dir + + "/stellar_evolution_tracks/Baraffe/BHAC15-M" + + str(formatted_mass).replace('.', 'p') + + ".txt") + if not os.path.isfile(file): + raise Exception( + "Cannot find Baraffe track file %s. " + "Did you set the FWL_DATA environment variable? " + "Did you run the DownloadEvolutionTracks function to get access to the Baraffe track data?" + % file + ) + + data = np.loadtxt(file,skiprows=1).T + + # Parse data + t = data[1] + Teff = data[2] + Lstar = data[3] + Rstar = data[5] + + # Pre-interpolate in time if required + if pre_interp: + # Params for interpolation + grid_count = 5e4 + if tmin==None: tmin = t[0] + if tmax==None: tmax = t[-1] + + # Do interpolation + #log.info("Interpolating stellar track onto a grid of size %d" % grid_count) + interp_Teff = PchipInterpolator(t,Teff) + interp_Lstar = PchipInterpolator(t,Lstar) + interp_Rstar = PchipInterpolator(t,Rstar) + + new_t = np.logspace(tmin, tmax, int(grid_count)) + new_t = np.log10(new_t) + new_Teff = interp_Teff(new_t) + new_Lstar = interp_Lstar(new_t) + new_Rstar = interp_Rstar(new_t) + + track = { + 't': 10.0**new_t, # yr + 'Teff': new_Teff, # K + 'Lstar': 10.0**new_Lstar, # L_sun + 'Rstar': new_Rstar # R_sun + } + else: + track = { + 't': 10.0**t, # yr + 'Teff': Teff, # K + 'Lstar': 10.0**Lstar, # L_sun + 'Rstar': Rstar # R_sun + } + + return track diff --git a/src/mors/constants.py b/src/mors/constants.py index 756fbcd..07cf9b4 100644 --- a/src/mors/constants.py +++ b/src/mors/constants.py @@ -54,6 +54,7 @@ # Other solar quantities LbolSun = 3.828e33 # erg s^-1 +LbolSun_SI = 3.828e26 # W AgeSun = 4567.0 # Myr OmegaSun = 2.67e-6 # rad s^-1 - our assumed solar rotation ProtSun = 2.0*3.142 / ( OmegaSun ) / 86400.0 # days - corresponding solar rotation period diff --git a/src/mors/data.py b/src/mors/data.py index 9aee64e..2cbd151 100644 --- a/src/mors/data.py +++ b/src/mors/data.py @@ -1,12 +1,37 @@ import os import subprocess from osfclient.api import OSF -from tqdm import tqdm #project ID of the stellar evolution tracks folder in the OSF project_id = '9u3fb' -def DownloadEvolutionTracks(): +def download_folder(storage, folder_name, local_path): + '''' + Download a specific folder in the OSF repository + + Inputs : + - storage : OSF storage name + - folder_name : folder name to be downloaded + - local_path : local repository where data are saved + ''' + for file in storage.files: + if file.path.startswith(folder_name): + local_file_path = local_path + file.path + #Create local directory if needed + os.makedirs(os.path.dirname(local_file_path), exist_ok=True) + #Download the file + with open(local_file_path, 'wb') as local_file: + file.write_to(local_file) + return + +def DownloadEvolutionTracks(fname=""): + ''' + Download evolution track data + + Inputs : + - fname (optional) : folder name, "/Spada" or "/Baraffe" + if not provided download both + ''' #Check if data environment variable is set up fwl_data_dir = os.getenv('FWL_DATA') @@ -23,33 +48,31 @@ def DownloadEvolutionTracks(): project = osf.project(project_id) storage = project.storage('osfstorage') - with tqdm(unit='files') as pbar: - for store in project.storages: - - #Loop over all the remote files in the OSF project - for file_ in store.files: - - #Set up local path for remote file - path = file_.path - if path.startswith('/'): - path = path[1:] - path = os.path.join(data_dir, path) - - #Create local directory if needed - directory, _ = os.path.split(path) - os.makedirs(directory, exist_ok=True) - - #Download and write remote file to local path - with open(path, "wb") as f: - file_.write_to(f) - - pbar.update() + #If no folder name specified download both Spada and Baraffe + #If local directory exists, assumes the data are already there + unzip_spada = False + if not fname: + if not os.path.exists(data_dir+"/Spada"): + download_folder(storage,"/Spada",data_dir) + unzip_spada = True + if not os.path.exists(data_dir+"/Baraffe"): + download_folder(storage,"/Baraffe",data_dir) + elif fname == "/Spada": + if not os.path.exists(data_dir+"/Spada"): + download_folder(storage,"/Spada",data_dir) + unzip_spada = True + elif fname == "/Baraffe": + if not os.path.exists(data_dir+"/Baraffe"): + download_folder(storage,"/Baraffe",data_dir) + else: + print("Unrecognised folder name in DownloadEvolutionTrackst st") - #Unzip Spada evolution tracks - wrk_dir = os.getcwd() - os.chdir(data_dir + '/Spada') - subprocess.call( ['tar','xvfz', 'fs255_grid.tar.gz'] ) - subprocess.call( ['rm','-f', 'fs255_grid.tar.gz'] ) - os.chdir(wrk_dir) + if unzip_spada: + #Unzip Spada evolution tracks + wrk_dir = os.getcwd() + os.chdir(data_dir + '/Spada') + subprocess.call( ['tar','xvfz', 'fs255_grid.tar.gz'] ) + subprocess.call( ['rm','-f', 'fs255_grid.tar.gz'] ) + os.chdir(wrk_dir) return