def initialPT2(date_dir, params, pressfile, mode, PTfunc, tepfile, tint=100.0): """ Compute a Temperature profile. Parameters: ----------- params: 1D Float ndarray Array of fitting parameters. pressfile: String File name of the pressure array. mode: String Chose the PT model: 'madhu' or 'line'. tepfile: String Filename of the planet's TEP file. tint: Float Internal planetary temperature. """ # Read pressures from file: pressure = pt.read_press_file(pressfile) # For extra PT arguments: PTargs = None # Read the TEP file: tep = rd.File(tepfile) # Stellar radius (in meters): rstar = float(tep.getvalue('Rs')[0]) * c.Rsun # Stellar temperature in K: tstar = float(tep.getvalue('Ts')[0]) # Semi-major axis (in meters): sma = float(tep.getvalue( 'a')[0]) * sc.au # Planetary radius (in meters): rplanet = float(tep.getvalue('Rp')[0]) * c.Rjup # Planetary mass (in kg): mplanet = float(tep.getvalue('Mp')[0]) * c.Mjup if mode == "line": # Planetary surface gravity (in cm s-2): gplanet = 100.0 * sc.G * mplanet / rplanet**2 # Additional PT arguments for Line case: PTargs = [rstar, tstar, tint, sma, gplanet] # Calculate temperature Temp = pt.PT_generator(pressure, params, PTfunc, PTargs) # Plot PT profile plt.figure(1) plt.semilogy(Temp, pressure, '-', color = 'r') plt.xlim(0.9*min(Temp), 1.1*max(Temp)) plt.ylim(max(pressure), min(pressure)) plt.title('Initial PT Line', fontsize=14) plt.xlabel('T (K)' , fontsize=14) plt.ylabel('logP (bar)', fontsize=14) # Save plot to current directory plt.savefig(date_dir + 'InitialPT.png') return Temp
def callTransit(atmfile, tepfile, MCfile, stepsize, molfit, solution, p0, tconfig, date_dir, burnin, abun_file, PTtype, PTfunc, T_int, T_int_type, filters, ctf=None, fs=15): """ Call Transit to produce best-fit outputs. Plot MCMC posterior PT plot. Parameters: ----------- atmfile: String Atmospheric file. tepfile: String Transiting extra-solar planet file. MCfile: String File with MCMC log and best-fitting results. stepsize: 1D float ndarray Specified step sizes for each parameter. molfit: 1D String ndarray List of molecule names to modify their abundances. solution: String Flag to indicate transit or eclipse geometry p0: Float Atmosphere's 'surface' pressure level. tconfig: String Transit configuration file. date_dir: String Directory where to store results. burnin: Integer abun_file: String Elemental abundances file. PTtype: String Pressure-temperature profile type ('line' for Line et al. 2013, 'madhu_noinv' or 'madhu_inv' for Madhusudhan & Seager 2009, 'iso' for isothermal) PTfunc: pointer to function Determines the method of evaluating the PT profile's temperature T_int: float. Internal temperature of the planet. T_int_type: string. Method to evaluate `T_int`. 'const' for constant value of `T_int`, 'thorngren' for Thorngren et al. (2019) filters: list, strings. Filter files associated with the eclipse/transit depths ctf: 2D array. Contribution or transmittance functions corresponding to `filters` fs: int Font size for plots. """ # make sure burnin is an integer burnin = int(burnin) # read atmfile molecules, pressure, temp, abundances = mat.readatm(atmfile) # get surface gravity grav, Rp = mat.get_g(tepfile) # get star data if needed if PTtype == 'line': R_star, T_star, sma, gstar = get_starData(tepfile) PTargs = [R_star, T_star, T_int, sma, grav * 1e2, T_int_type] else: PTargs = None # For non-Line profiles # Get best parameters bestP, uncer = read_MCMC_out(MCfile) allParams = bestP # get PTparams and abundances factors nparams = len(allParams) nmol = len(molfit) nradfit = int(solution == 'transit') nPTparams = nparams - nmol - nradfit PTparams = allParams[:nPTparams] # call PT profile generator to calculate temperature best_T = pt.PT_generator(pressure, PTparams, PTfunc, PTargs) # Plot best PT profile plt.figure(1) plt.clf() plt.semilogy(best_T, pressure, '-', color='r') plt.xlim(0.9 * min(best_T), 1.1 * max(best_T)) plt.ylim(max(pressure), min(pressure)) plt.title('Best PT', fontsize=fs) plt.xlabel('T (K)', fontsize=fs) plt.ylabel('logP (bar)', fontsize=fs) # Save plot to current directory plt.savefig(date_dir + 'Best_PT.png') # Update R0, if needed: if nradfit: Rp = allParams[nPTparams] # write best-fit atmospheric file write_atmfile(atmfile, abun_file, molfit, best_T, allParams[nPTparams + nradfit:], date_dir, p0, Rp, grav) # write new bestFit Transit config if solution == 'transit': bestFit_tconfig(tconfig, date_dir, allParams[nPTparams]) else: bestFit_tconfig(tconfig, date_dir) # Call Transit with the best-fit tconfig Transitdir = os.path.join(os.path.dirname(os.path.realpath(__file__)), "..", "modules", "transit", "") bf_tconfig = os.path.join(date_dir, 'bestFit_tconfig.cfg') Tcall = os.path.join(Transitdir, "transit", "transit") subprocess.call(["{:s} -c {:s}".format(Tcall, bf_tconfig)], shell=True, cwd=date_dir) # ========== plot MCMC PT profiles ========== # get MCMC data: MCMCdata = os.path.join(date_dir, "output.npy") data = np.load(MCMCdata) nchains, npars, niter = np.shape(data) # stuck chains: data_stack = data[0, :, burnin:] for c in np.arange(1, nchains): data_stack = np.hstack((data_stack, data[c, :, burnin:])) # create array of PT profiles PTprofiles = np.zeros((np.shape(data_stack)[1], len(pressure))) # current PT parameters for each chain, iteration curr_PTparams = PTparams # fill-in PT profiles array if ctf is None: print(" Plotting MCMC PT profile figure.") for k in np.arange(0, np.shape(data_stack)[1]): j = 0 for i in np.arange(len(PTparams)): if stepsize[i] != 0.0: curr_PTparams[i] = data_stack[j, k] j += 1 else: pass PTprofiles[k] = pt.PT_generator(pressure, curr_PTparams, PTfunc, PTargs) # get percentiles (for 1,2-sigma boundaries): low1 = np.percentile(PTprofiles, 15.87, axis=0) hi1 = np.percentile(PTprofiles, 84.13, axis=0) low2 = np.percentile(PTprofiles, 2.28, axis=0) hi2 = np.percentile(PTprofiles, 97.72, axis=0) median = np.median(PTprofiles, axis=0) # plot figure plt.figure(2, dpi=300) plt.clf() if ctf is not None: plt.subplots_adjust(wspace=0.15) ax1 = plt.subplot(121) else: ax1 = plt.subplot(111) ax1.fill_betweenx(pressure, low2, hi2, facecolor="#62B1FF", edgecolor="0.5") ax1.fill_betweenx(pressure, low1, hi1, facecolor="#1873CC", edgecolor="#1873CC") plt.semilogy(median, pressure, "-", lw=2, label='Median', color="k") plt.semilogy(best_T, pressure, "-", lw=2, label="Best fit", color="r") plt.ylim(pressure[0], pressure[-1]) plt.legend(loc="best") plt.xlabel("Temperature (K)", size=fs) plt.ylabel("Pressure (bar)", size=fs) if ctf is not None: nfilters = len(filters) # Add contribution or transmittance functions ax2 = plt.subplot(122, sharey=ax1) colormap = plt.cm.rainbow(np.linspace(0, 1, len(filters))) ax2.set_prop_cycle(plt.cycler('color', colormap)) # Plot with filter labels for i in np.arange(len(filters)): (head, tail) = os.path.split(filters[i]) lbl = tail[:-4] ax2.semilogy(ctf[i], pressure, '--', linewidth=1, label=lbl) # Hide y axis tick labels plt.setp(ax2.get_yticklabels(), visible=False) # Place legend off figure in case there are many filters lgd = ax2.legend(loc='center left', bbox_to_anchor=(1.05, 0.5), ncol=nfilters // 25 + int(nfilters % 25 != 0), prop={'size': 8}) if solution == 'eclipse': ax2.set_xlabel('Normalized Contribution\nFunctions', fontsize=fs) else: ax2.set_xlabel('Transmittance', fontsize=fs) # save figure if ctf is not None: savefile = date_dir + "MCMC_PTprofiles_cf.png" plt.savefig(savefile, bbox_extra_artists=(lgd, ), bbox_inches='tight') else: savefile = date_dir + "MCMC_PTprofiles.png" plt.savefig(savefile) plt.close()
def main(comm): """ This is a hacked version of MC3's func.py. This function directly call's the modeling function for the BART project. """ # Parse arguments: cparser = argparse.ArgumentParser( description=__doc__, add_help=False, formatter_class=argparse.RawDescriptionHelpFormatter) # Add config file option: cparser.add_argument("-c", "--config_file", help="Configuration file", metavar="FILE") # Remaining_argv contains all other command-line-arguments: args, remaining_argv = cparser.parse_known_args() # Get parameters from configuration file: cfile = args.config_file if cfile: config = ConfigParser.SafeConfigParser() config.optionxform = str config.read([cfile]) defaults = dict(config.items("MCMC")) else: defaults = {} parser = argparse.ArgumentParser(parents=[cparser]) parser.add_argument("--func", dest="func", type=mu.parray, action="store", default=None) parser.add_argument("--indparams", dest="indparams", type=mu.parray, action="store", default=[]) parser.add_argument("--params", dest="params", type=mu.parray, action="store", default=None, help="Model-fitting parameters [default: %(default)s]") parser.add_argument("--molfit", dest="molfit", type=mu.parray, action="store", default=None, help="Molecules fit [default: %(default)s]") parser.add_argument( "--Tmin", dest="Tmin", type=float, action="store", default=400.0, help="Lower Temperature boundary [default: %(default)s]") parser.add_argument( "--Tmax", dest="Tmax", type=float, action="store", default=3000.0, help="Higher Temperature boundary [default: %(default)s]") parser.add_argument("--quiet", action="store_true", help="Set verbosity level to minimum", dest="quiet") # Input-Converter Options: group = parser.add_argument_group("Input Converter Options") group.add_argument("--atmospheric_file", action="store", help="Atmospheric file [default: %(default)s]", dest="atmfile", type=str, default=None) group.add_argument("--PTtype", action="store", help="PT profile type.", dest="PTtype", type=str, default="none") group.add_argument("--tint", action="store", help="Internal temperature of the planet [default: " "%(default)s].", dest="tint", type=float, default=100.0) # transit Options: group = parser.add_argument_group("transit Options") group.add_argument( "--config", action="store", help="transit configuration file [default: %(default)s]", dest="config", type=str, default=None) # Output-Converter Options: group = parser.add_argument_group("Output Converter Options") group.add_argument("--filters", action="store", help="Waveband filter name [default: %(default)s]", dest="filters", type=mu.parray, default=None) group.add_argument("--tep_name", action="store", help="A TEP file [default: %(default)s]", dest="tep_name", type=str, default=None) group.add_argument("--kurucz_file", action="store", help="Stellar Kurucz file [default: %(default)s]", dest="kurucz", type=str, default=None) group.add_argument("--solution", action="store", help="Solution geometry [default: %(default)s]", dest="solution", type=str, default="None", choices=('transit', 'eclipse')) parser.set_defaults(**defaults) args2, unknown = parser.parse_known_args(remaining_argv) # Quiet all threads except rank 0: rank = comm.Get_rank() verb = rank == 0 # Get (Broadcast) the number of parameters and iterations from MPI: array1 = np.zeros(2, np.int) mu.comm_bcast(comm, array1) npars, niter = array1 # ::::::: Initialize the Input converter :::::::::::::::::::::::::: atmfile = args2.atmfile molfit = args2.molfit PTtype = args2.PTtype params = args2.params tepfile = args2.tep_name tint = args2.tint Tmin = args2.Tmin Tmax = args2.Tmax solution = args2.solution # Solution type # Dictionary of functions to calculate temperature for PTtype PTfunc = { 'iso': pt.PT_iso, 'line': pt.PT_line, 'madhu_noinv': pt.PT_NoInversion, 'madhu_inv': pt.PT_Inversion } # Extract necessary values from the TEP file: tep = rd.File(tepfile) # Stellar temperature in K: tstar = float(tep.getvalue('Ts')[0]) # Stellar radius (in meters): rstar = float(tep.getvalue('Rs')[0]) * c.Rsun # Semi-major axis (in meters): sma = float(tep.getvalue('a')[0]) * sc.au # Planetary radius (in meters): rplanet = float(tep.getvalue('Rp')[0]) * c.Rjup # Planetary mass (in kg): mplanet = float(tep.getvalue('Mp')[0]) * c.Mjup # Number of fitting parameters: nfree = len(params) # Total number of free parameters nmolfit = len(molfit) # Number of molecular free parameters nradfit = int(solution == 'transit') # 1 for transit, 0 for eclipse nPT = nfree - nmolfit - nradfit # Number of PT free parameters # Read atmospheric file to get data arrays: species, pressure, temp, abundances = mat.readatm(atmfile) # Reverse pressure order (for PT to work): pressure = pressure[::-1] nlayers = len(pressure) # Number of atmospheric layers nspecies = len(species) # Number of species in the atmosphere mu.msg(verb, "There are {:d} layers and {:d} species.".format(nlayers, nspecies)) # Find index for Hydrogen and Helium: species = np.asarray(species) iH2 = np.where(species == "H2")[0] iHe = np.where(species == "He")[0] # Get H2/He abundance ratio: ratio = (abundances[:, iH2] / abundances[:, iHe]).squeeze() # Find indices for the metals: imetals = np.where((species != "He") & (species != "H2") & \ (species != "H-") & (species != 'e-'))[0] # Index of molecular abundances being modified: imol = np.zeros(nmolfit, dtype='i') for i in np.arange(nmolfit): imol[i] = np.where(np.asarray(species) == molfit[i])[0] # Pressure-Temperature profile: if PTtype == "line": # Planetary surface gravity (in cm s-2): gplanet = 100.0 * sc.G * mplanet / rplanet**2 # Additional PT arguments: PTargs = [rstar, tstar, tint, sma, gplanet] else: PTargs = None # Allocate arrays for receiving and sending data to master: freepars = np.zeros(nfree, dtype='d') profiles = np.zeros((nspecies + 1, nlayers), dtype='d') # This are sub-sections of profiles, containing just the temperature and # the abundance profiles, respectively: tprofile = profiles[0, :] aprofiles = profiles[1:, :] # Store abundance profiles: for i in np.arange(nspecies): aprofiles[i] = abundances[:, i] # ::::::: Spawn transit code ::::::::::::::::::::::::::::::::::::: # # transit configuration file: transitcfile = args2.tconfig # Initialize the transit python module: transit_args = ["transit", "-c", transitcfile] trm.transit_init(len(transit_args), transit_args) # Get wavenumber array from transit: nwave = trm.get_no_samples() specwn = trm.get_waveno_arr(nwave) # ::::::: Output Converter ::::::::::::::::::::::::::::::::::::::: ffile = args2.filters # Filter files kurucz = args2.kurucz # Kurucz file # Log10(stellar gravity) gstar = float(tep.getvalue('loggstar')[0]) # Planet-to-star radius ratio: rprs = rplanet / rstar nfilters = len(ffile) # Number of filters: # FINDME: Separate filter/stellar interpolation? # Get stellar model: starfl, starwn, tmodel, gmodel = w.readkurucz(kurucz, tstar, gstar) # Read and resample the filters: nifilter = [] # Normalized interpolated filter istarfl = [] # interpolated stellar flux wnindices = [] # wavenumber indices used in interpolation for i in np.arange(nfilters): # Read filter: filtwaven, filttransm = w.readfilter(ffile[i]) # Check that filter boundaries lie within the spectrum wn range: if filtwaven[0] < specwn[0] or filtwaven[-1] > specwn[-1]: mu.exit(message="Wavenumber array ({:.2f} - {:.2f} cm-1) does not " "cover the filter[{:d}] wavenumber range ({:.2f} - {:.2f} " "cm-1).".format(specwn[0], specwn[-1], i, filtwaven[0], filtwaven[-1])) # Resample filter and stellar spectrum: nifilt, strfl, wnind = w.resample(specwn, filtwaven, filttransm, starwn, starfl) nifilter.append(nifilt) istarfl.append(strfl) wnindices.append(wnind) # Allocate arrays for receiving and sending data to master: spectrum = np.zeros(nwave, dtype='d') bandflux = np.zeros(nfilters, dtype='d') # Allocate array to receive parameters from MPI: params = np.zeros(npars, np.double) # :::::: Main MCMC Loop :::::::::::::::::::::::::::::::::::::::::: # :::::::::::::::::::::::::::::::::::::::::::::::::::::::::::::::::: while niter >= 0: niter -= 1 # Receive parameters from MCMC: mu.comm_scatter(comm, params) # Check for the MCMC-end flag: if params[0] == np.inf: break # Input converter calculate the profiles: try: tprofile[:] = pt.PT_generator(pressure, params[0:nPT], PTfunc[PTtype], PTargs)[::-1] except ValueError: mu.msg(verb, 'Input parameters give non-physical profile.') # FINDME: what to do here? # If the temperature goes out of bounds: if np.any(tprofile < Tmin) or np.any(tprofile > Tmax): mu.comm_gather(comm, -np.ones(nfilters), MPI.DOUBLE) continue # Scale abundance profiles: for i in np.arange(nmolfit): m = imol[i] # Use variable as the log10: aprofiles[m] = abundances[:, m] * 10.0**params[nPT + nradfit + i] # Update H2, He abundances so sum(abundances) = 1.0 in each layer: q = 1.0 - np.sum(aprofiles[imetals], axis=0) if np.any(q < 0.0): mu.comm_gather(comm, -np.ones(nfilters), MPI.DOUBLE) continue aprofiles[iH2] = ratio * q / (1.0 + ratio) aprofiles[iHe] = q / (1.0 + ratio) # Set the 'surface' level: if solution == "transit": trm.set_radius(params[nPT]) # Let transit calculate the model spectrum: spectrum = trm.run_transit(profiles.flatten(), nwave) # Calculate the band-integrated intensity per filter: for i in np.arange(nfilters): if solution == "eclipse": fluxrat = (spectrum[wnindices[i]] / istarfl[i]) * rprs * rprs bandflux[i] = w.bandintegrate(fluxrat, specwn, nifilter[i], wnindices[i]) elif solution == "transit": bandflux[i] = w.bandintegrate(spectrum[wnindices[i]], specwn, nifilter[i], wnindices[i]) # Send resutls back to MCMC: mu.comm_gather(comm, bandflux, MPI.DOUBLE) # :::::: End main Loop ::::::::::::::::::::::::::::::::::::::::::: # :::::::::::::::::::::::::::::::::::::::::::::::::::::::::::::::::: # Close communications and disconnect: mu.comm_disconnect(comm) trm.free_memory()