Source code for pyratbay.driver

# Copyright (c) 2021-2022 Patricio Cubillos
# Pyrat Bay is open-source software under the GNU GPL-2.0 license (see LICENSE)

__all__ = [
    'run',
]

import os

import numpy as np
import mc3

from . import tools as pt
from . import opacity as po
from . import constants as pc
from . import atmosphere as pa
from . import spectrum as ps
from . import io as io
from . import plots as pp
from .pyrat import Pyrat


[docs]@pt.ignore_system_exit def run(cfile, run_step='run', no_logfile=False): """ Pyrat Bay initialization driver. Parameters ---------- cfile: String A Pyrat Bay configuration file. run_step: String If 'dry': only read the configuration file into a Pyrat object If 'init': initialize a Pyrat object (but no spectra calculation). If 'run': run all the way (default). no_logfile: Bool If True, enforce not to write outputs to a log file (e.g., to prevent overwritting log of a previous run). """ pyrat = Pyrat(cfile, no_logfile=no_logfile) log = pyrat.log phy = pyrat.phy atm = pyrat.atm ret = pyrat.ret inputs = pyrat.inputs if run_step == 'dry': return pyrat # Call lineread package: if pyrat.runmode == 'tli': if pyrat.lt.tlifile is None: log.error('Undefined TLI file (tlifile).') po.make_tli( inputs.dblist, inputs.pflist, inputs.dbtype, pyrat.lt.tlifile[0], pyrat.spec.wllow, pyrat.spec.wlhigh, pyrat.spec.wlunits, log) return # Get gplanet from mplanet and rplanet if necessary: mplanet = phy.mplanet is not None gplanet = phy.gplanet is not None rplanet = phy.rplanet is not None # Check planetary surface gravity/mass/radius: if mplanet and rplanet and gplanet: gplanet = pc.G * phy.mplanet / phy.rplanet**2 if np.abs(gplanet-phy.gplanet)/phy.gplanet > 0.05: log.error( "All mplanet, rplanet, and gplanet were provided, but " f"values are inconsistent (>5%): g(M,R) = {gplanet:7.1f} " f"cm s-2 and gplanet = {phy.gplanet:7.1f} cm s-2.") elif not mplanet and rplanet and gplanet: phy.mplanet = phy.gplanet * phy.rplanet**2 / pc.G elif mplanet and not rplanet and gplanet: phy.rplanet = np.sqrt(pc.G * phy.mplanet / phy.gplanet) elif mplanet and rplanet and not gplanet: phy.gplanet = pc.G * phy.mplanet / phy.rplanet**2 require_atmospheric_model = ( pyrat.runmode in ['atmosphere', 'radeq'] or pt.isfile(atm.atmfile) != 1 or 'equil' in atm.molmodel) if require_atmospheric_model: # Compute pressure-temperature profile: read_atmosphere = ( pt.isfile(atm.atmfile) == 1 and (pyrat.runmode == 'radeq' or 'equil' in atm.molmodel) ) if pt.isfile(atm.ptfile) == 1: log.msg(f"\nReading pressure-temperature file: '{atm.ptfile}'.") units, _, pressure, temperature = io.read_atm(atm.ptfile)[0:4] pressure *= pt.u(units[0]) # pressure in barye elif read_atmosphere: units, inputs.species, pressure, temperature, _, _ = \ io.read_atm(atm.atmfile) pressure *= pt.u(units[0]) # pressure in barye else: check_pressure(pyrat) pressure = pa.pressure( atm.ptop, atm.pbottom, atm.nlayers, 'barye', log) check_temp(pyrat) temperature = pa.temperature( atm.tmodelname, pressure, atm.nlayers, log, atm.tpars) abundances = None species = None radius = None if 'equil' in atm.molmodel: atm.chemistry = 'tea' # Compute volume-mixing-ratio profiles: if atm.chemistry is not None or pyrat.runmode != 'atmosphere': check_atm(pyrat) chem_net = pa.chemistry( atm.chemistry, pressure, temperature, inputs.species, metallicity=atm.metallicity, e_abundances=atm.e_abundances, e_scale=atm.e_scale, e_ratio=atm.e_ratio, solar_file=inputs.solar, log=log, atmfile=atm.atmfile, punits=atm.punits, q_uniform=inputs.uniform, ) atm.chem_model = chem_net abundances = chem_net.vmr species = chem_net.species # Compute altitude profile: if abundances is not None and atm.rmodelname is not None: check_altitude(pyrat) # Mean molecular mass: mean_mass = pa.mean_weight(abundances, species) # Altitude profile: radius = pyrat.hydro( pressure, temperature, mean_mass, phy.gplanet, phy.mplanet, atm.refpressure, phy.rplanet, ) # Return atmospheric model if requested: if pyrat.runmode == 'atmosphere': header = '# Pyrat bay atmospheric model\n' # Guess radius units if not defined (by rplanet): if radius is not None and atm.runits is None: atm.runits = 'rjup' if phy.rplanet > 0.5*pc.rjup else 'rearth' if atm.atmfile is not None: io.write_atm( atm.atmfile, pressure, temperature, species, abundances, radius, atm.punits, atm.runits, header=header) log.msg(f"Output atmospheric file: '{atm.atmfile}'.") return pressure, temperature, abundances, species, radius # Check status of extinction-coefficient file if necessary: if pyrat.runmode != 'spectrum' and pt.isfile(pyrat.ex.extfile) == -1: log.error("Undefined extinction-coefficient file (extfile).") # Force to re-calculate extinction-coefficient file if requested: if pyrat.runmode == 'opacity' and pt.isfile(pyrat.ex.extfile) == 1: for extfile in pyrat.ex.extfile: os.remove(extfile) if pyrat.runmode == 'mcmc' and ret.mcmcfile is None: log.error('Undefined MCMC file (mcmcfile).') # Initialize pyrat object: pyrat.setup_spectrum() #if pyrat.inputs.resume: # Bypass writting all of the initialization log: # pyrat = Pyrat(args, log=None) # pyrat.log = log #else: # pyrat = Pyrat(args, log=log) # Stop and return if requested: if run_step == 'init' or pyrat.runmode == 'opacity': return pyrat # Compute spectrum and return pyrat object if requested: if pyrat.runmode == "spectrum": pyrat.run() return pyrat if pyrat.runmode == 'radeq': pyrat.radiative_equilibrium() return pyrat # Mute logging in pyrat object, but not in mc3: pyrat.log = mc3.utils.Log(verb=0, width=80) pyrat.spec.specfile = None # Avoid writing spectrum file during MCMC retmodel = False # Return only the band-integrated spectrum # Basename of the output files (no path, no extension): outfile = os.path.splitext(os.path.basename(ret.mcmcfile))[0] # Run MCMC: mc3_out = mc3.sample( data=pyrat.obs.data, uncert=pyrat.obs.uncert, func=pyrat.eval, indparams=[retmodel], params=ret.params, pmin=ret.pmin, pmax=ret.pmax, pstep=ret.pstep, prior=ret.prior, priorlow=ret.priorlow, priorup=ret.priorup, sampler=ret.sampler, nsamples=ret.nsamples, nchains=ret.nchains, burnin=ret.burnin, thinning=ret.thinning, grtest=True, grbreak=ret.grbreak, grnmin=ret.grnmin, log=log, ncpu=pyrat.ncpu, plots=True, showbp=True, pnames=ret.pnames, texnames=ret.texnames, resume=inputs.resume, savefile=ret.mcmcfile) if mc3_out is None: log.error("Error in MC3.") bestp = mc3_out['bestp'] ret.bestp = bestp posterior, zchain, zmask = mc3.utils.burn(mc3_out) ret.posterior = posterior # Best-fitting model: pyrat.spec.specfile = f"{outfile}_bestfit_spectrum.dat" ret.spec_best, ret.bestbandflux = pyrat.eval(bestp) header = "# MCMC best-fitting atmospheric model.\n\n" bestatm = f"{outfile}_bestfit_atmosphere.atm" io.write_atm( bestatm, atm.press, atm.temp, pyrat.mol.name, atm.q, radius=atm.radius, punits=atm.punits, runits='km', header=header) pyrat.plot_spectrum(spec='best', filename=f'{outfile}_bestfit_spectrum.png') # Temperature profiles: if atm.tmodelname is not None: ifree = ret.pstep[ret.itemp] > 0 itemp = np.arange(np.sum(ifree)) ret.temp_best = atm.tmodel(bestp[ret.itemp]) tpost = pa.temperature_posterior( posterior[:,itemp], atm.tmodel, ret.params[ret.itemp], ifree, atm.press) ret.temp_median = tpost[0] # 1sigma-low, 1sigma-high, 2sigma-low, 2sigma-high: ret.temp_post_boundaries = tpost[1:] pyrat.plot_temperature( filename=f'{outfile}_posterior_temperature_profile.png') is_emission = pyrat.od.rt_path in pc.emission_rt is_transmission = pyrat.od.rt_path in pc.transmission_rt if is_emission: cf = ps.contribution_function( pyrat.od.depth, atm.press, pyrat.od.B, ) bcf = ps.band_cf( cf, pyrat.obs.bandtrans, pyrat.spec.wn, pyrat.obs.bandidx) elif is_transmission: transmittance = ps.transmittance(pyrat.od.depth, pyrat.od.ideep) bcf = ps.band_cf( transmittance, pyrat.obs.bandtrans, pyrat.spec.wn, pyrat.obs.bandidx) path = 'transit' if is_transmission else 'emission' pp.contribution( bcf, 1.0/(pyrat.obs.bandwn*pc.um), path, atm.press, atm.radius, atm.rtop, filename=f"{outfile}_bestfit_cf.png") pyrat.log = log # Un-mute log.msg( "\nOutput MCMC posterior results, log, bestfit atmosphere, " "and spectrum:" f"\n'{outfile}.npz'" f"\n'{os.path.basename(inputs.logfile)}'" f"\n'{bestatm}'" f"\n'{pyrat.spec.specfile}'\n\n") return pyrat
def check_pressure(pyrat): """ Check the input arguments to calculate the pressure profile. """ if pyrat.atm.nlayers is None: pyrat.log.error("Undefined number of atmospheric layers (nlayers).") if pyrat.atm.ptop is None: pyrat.log.error("Undefined atmospheric top pressure (ptop)") if pyrat.atm.pbottom is None: pyrat.log.error("Undefined atmospheric bottom pressure (pbottom)") def check_temp(pyrat): """ Check the input arguments to calculate the temperature profile. """ if pyrat.atm.tmodelname is None: pyrat.log.error("Undefined temperature model (tmodel).") if pyrat.atm.tpars is None: pyrat.log.error("Undefined temperature-model parameters (tpars).") def check_atm(pyrat): """ Check the input arguments to calculate the atmospheric model. """ atm = pyrat.atm log = pyrat.log if atm.chemistry is None: log.error("Undefined chemistry model (chemistry).") if atm.atmfile is None: log.error("Undefined atmospheric file (atmfile).") if pyrat.inputs.species is None: log.error("Undefined atmospheric species list (species).") # Uniform-abundances profile: if atm.chemistry == 'uniform': if pyrat.inputs.uniform is None: log.error( "Undefined list of uniform volume mixing ratios " f"(uniform) for {atm.chemistry} chemistry model." ) nuniform = len(pyrat.inputs.uniform) nspecies = len(pyrat.inputs.species) if nuniform != nspecies: pyrat.log.error( f"Number of uniform abundances ({nuniform}) does " f"not match the number of species ({nspecies})." ) return pyrat.inputs.metallicity = pyrat.inputs.get_default( 'metallicity', 'Metallicity scaling factor (dex, relative to solar)', default=0.0) def check_altitude(pyrat): """Check input arguments to calculate altitude profile.""" phy = pyrat.phy log = pyrat.log rad_vars = ['mplanet', 'rplanet', 'gplanet'] missing = [rvar for rvar in rad_vars if getattr(phy,rvar) is None] if len(missing) == 2: err = f'either {missing[0]} or {missing[1]}' elif len(missing) == 3: err = 'at least two of mplanet, rplanet, or gplanet' if len(missing) > 0: log.error( 'Cannot compute hydrostatic-equilibrium radius profile. ' f'Must\ndefine {err}.' ) if pyrat.atm.refpressure is None: log.error( 'Cannot compute hydrostatic-equilibrium radius profile. ' 'Undefined reference pressure level (refpressure).' )