def generate_profiles(self, solution, binning):
from taurex import mpi
"""Generates sigma plots for profiles"""
sample_list = None
if mpi.get_rank() == 0:
sample_list = list(self.sample_parameters(solution))
sample_list = mpi.broadcast(sample_list)
self.debug('We all got %s', sample_list)
self.info('------------Variance generation step------------------')
self.info('We are sampling %s points for the profiles',
len(sample_list))
rank = mpi.get_rank()
size = mpi.nprocs()
enableLogging()
self.info(
'I will only iterate through partitioned %s '
'points (the rest is in parallel)',
len(sample_list) // size)
disableLogging()
count = 0
def sample_iter():
for parameters, weight in sample_list[rank::size]:
self.update_model(parameters)
enableLogging()
if rank == 0 and count % 10 == 0 and count > 0:
self.info('Progress {}%'.format(count * 100.0 /
(len(sample_list) / size)))
disableLogging()
yield weight
return self._model.compute_error(sample_iter,
wngrid=binning,
binner=self._binner)
def output_citations(model, instrument, optimizer):
from taurex.mpi import barrier, get_rank
barrier()
bib_tex = None
citation_string = None
if get_rank() == 0:
print('\n\n----------------------------------------------------------')
print('----------------------Bibiliography-----------------------')
print('----------------------------------------------------------')
print('If you use any of the results from this run please cite')
print('the following publications:')
citation_string = ''
all_citations = []
print('\n')
print('TauREx-Related')
print('--------------\n')
from taurex._citation import __citations__, taurex_citation
citation_string += __citations__
all_citations.extend(taurex_citation.citations())
print(__citations__)
print('Forward model')
print('-------------\n')
cite = model.nice_citation()
all_citations.extend(model.citations())
citation_string += cite
print(cite)
if optimizer is not None:
cite = optimizer.nice_citation()
all_citations.extend(optimizer.citations())
if len(cite) > 0:
citation_string += cite
print('Optimizer')
print('---------\n')
print(cite)
if instrument is not None:
cite = instrument.nice_citation()
all_citations.extend(instrument.citations())
if len(cite) > 0:
citation_string += cite
print('Instrument')
print('---------\n')
print(cite)
from taurex.core import to_bibtex
bib_tex = to_bibtex(all_citations)
barrier()
return bib_tex, citation_string
def open(self):
if get_rank() == 0:
self.fd = self._openFile(self.filename, self._append)
def __init__(self, stream=None):
from taurex.mpi import get_rank
super().__init__(stream=stream)
self._rank = get_rank()
def main():
import argparse
import datetime
import logging
from taurex.mpi import get_rank
from taurex.log import setLogLevel
from taurex.log.logger import root_logger
from taurex.parameter import ParameterParser
from taurex.output.hdf5 import HDF5Output
from taurex.util.output import generate_profile_dict, store_contributions
from .taurexdefs import OutputSize
from . import __version__ as version
import numpy as np
parser = argparse.ArgumentParser(description='TauREx {}'.format(version))
parser.add_argument("-i", "--input", dest='input_file', type=str,
required=True, help="Input par file to pass")
parser.add_argument("-R", "--retrieval", dest='retrieval', default=False,
help="When set, runs retrieval", action='store_true')
parser.add_argument("-p", "--plot", dest='plot', default=False,
help="Whether to plot after the run",
action='store_true')
parser.add_argument("-g", "--debug-log", dest='debug', default=False,
help="Debug log output", action='store_true')
parser.add_argument("-c", "--show-contrib", dest='contrib',
default=False, help="Show basic contributions",
action='store_true')
parser.add_argument("-C", "--full-contrib", dest='full_contrib',
default=False, help="Show ALL contributions",
action='store_true')
parser.add_argument("--light", dest='light', default=False,
help="Light outputs", action='store_true')
parser.add_argument("--lighter", dest='lighter', default=False,
help="Even Lighter outputs", action='store_true')
parser.add_argument("-o", "--output_file", dest='output_file', type=str)
parser.add_argument("-S", "--save-spectrum",
dest='save_spectrum', type=str)
args = parser.parse_args()
output_size = OutputSize.heavy
if args.light:
output_size = OutputSize.light
if args.lighter:
output_size = OutputSize.lighter
if args.debug:
setLogLevel(logging.DEBUG)
root_logger.info('TauREx %s', version)
root_logger.info('TauREx PROGRAM START AT %s', datetime.datetime.now())
# Parse the input file
pp = ParameterParser()
pp.read(args.input_file)
# Setup global parameters
pp.setup_globals()
# Generate a model from the input
model = pp.generate_appropriate_model()
# build the model
model.build()
# Get the spectrum
observation = pp.generate_observation()
binning = pp.generate_binning()
wngrid = None
if binning == 'observed' and observation is None:
logging.critical('Binning selected from Observation yet None provided')
quit()
if binning is None:
if observation is None or observation == 'self':
binning = model.defaultBinner()
wngrid = model.nativeWavenumberGrid
else:
binning = observation.create_binner()
wngrid = observation.wavenumberGrid
else:
if binning == 'native':
binning = model.defaultBinner()
wngrid = model.nativeWavenumberGrid
elif binning == 'observed':
binning = observation.create_binner()
wngrid = observation.wavenumberGrid
else:
binning, wngrid = binning
instrument = pp.generate_instrument(binner=binning)
num_obs = 1
if instrument is not None:
instrument, num_obs = instrument
if observation == 'self' and instrument is None:
logging.getLogger('taurex').critical(
'Instrument nust be specified when using self option')
raise ValueError('No instruemnt specified for self option')
inst_result = None
if instrument is not None:
inst_result = instrument.model_noise(
model, model_res=model.model(), num_observations=num_obs)
# Observation on self
if observation == 'self':
from taurex.data.spectrum import ArraySpectrum
from taurex.util.util import wnwidth_to_wlwidth
inst_wngrid, inst_spectrum, inst_noise, inst_width = inst_result
inst_wlgrid = 10000/inst_wngrid
inst_wlwidth = wnwidth_to_wlwidth(inst_wngrid, inst_width)
observation = ArraySpectrum(
np.vstack([inst_wlgrid, inst_spectrum,
inst_noise, inst_wlwidth]).T)
binning = observation.create_binner()
# Handle outputs
if args.output_file:
# Output taurex data
with HDF5Output(args.output_file) as o:
model.write(o)
optimizer = None
solution = None
if args.retrieval is True:
import time
if observation is None:
logging.critical('No spectrum is defined!!')
quit()
optimizer = pp.generate_optimizer()
optimizer.set_model(model)
optimizer.set_observed(observation)
fitting_parameters = pp.generate_fitting_parameters()
for key, value in fitting_parameters.items():
fit = value['fit']
bounds = value['bounds']
mode = value['mode']
factor = value['factor']
if fit:
logging.info('Fitting: {}'.format(key))
optimizer.enable_fit(key)
else:
optimizer.disable_fit(key)
if factor:
optimizer.set_factor_boundary(key, factor)
if bounds:
optimizer.set_boundary(key, bounds)
if mode:
optimizer.set_mode(key, mode.lower())
start_time = time.time()
solution = optimizer.fit(output_size=output_size)
end_time = time.time()
root_logger.info(
'Total Retrieval finish in %s seconds', end_time-start_time)
for _, optimized, _, _ in optimizer.get_solution():
optimizer.update_model(optimized)
break
result = model.model()
if args.save_spectrum is not None:
# with open(args.save_spectrum, 'w') as f:
from taurex.util.util import wnwidth_to_wlwidth, compute_bin_edges
save_wnwidth = compute_bin_edges(wngrid)[1]
save_wl = 10000/wngrid
save_wlwidth = wnwidth_to_wlwidth(wngrid, save_wnwidth)
save_model = binning.bin_model(result)[1]
save_error = np.zeros_like(save_wl)
if inst_result is not None:
inst_wngrid, inst_spectrum, inst_noise, inst_width = inst_result
save_model = inst_spectrum
save_wl = 10000/inst_wngrid
save_wlwidth = wnwidth_to_wlwidth(inst_wngrid, inst_width)
save_error = inst_noise
np.savetxt(args.save_spectrum,
np.vstack((save_wl, save_model, save_error,
save_wlwidth)).T)
if args.output_file:
# Output taurex data
with HDF5Output(args.output_file, append=True) as o:
out = o.create_group('Output')
if observation is not None:
obs = o.create_group('Observed')
observation.write(obs)
profiles = generate_profile_dict(model)
spectrum = \
binning.generate_spectrum_output(result,
output_size=output_size)
if inst_result is not None:
spectrum['instrument_wngrid'] = inst_result[0]
spectrum['instrument_wnwidth'] = inst_result[-1]
spectrum['instrument_wlgrid'] = 10000/inst_result[0]
spectrum['instrument_spectrum'] = inst_result[1]
spectrum['instrument_noise'] = inst_result[2]
spectrum['Contributions'] = \
store_contributions(binning, model, output_size=output_size-3)
if solution is not None:
out.store_dictionary(solution, group_name='Solutions')
priors = {}
priors['Profiles'] = profiles
priors['Spectra'] = spectrum
out.store_dictionary(priors, group_name='Priors')
else:
out.store_dictionary(profiles, group_name='Profiles')
out.store_dictionary(spectrum, group_name='Spectra')
if optimizer:
optimizer.write(o)
root_logger.info('TauREx PROGRAM END AT %s s', datetime.datetime.now())
if args.plot:
wlgrid = 10000/wngrid
if get_rank() == 0:
import matplotlib.pyplot as plt
fig = plt.figure()
ax = fig.add_subplot(1, 1, 1)
ax.set_xlabel(r'Wavelength $(\mu m)$')
ax.set_ylabel(r'$(R_p/R_s)^2$')
is_lightcurve = False
try:
from taurex.model.lightcurve.lightcurve import LightCurveModel
is_lightcurve = isinstance(model, LightCurveModel)
ax.set_xscale('linear')
except ImportError:
pass
if observation is not None:
if is_lightcurve:
ax.plot(observation.spectrum.flatten(),
label='observation')
else:
ax.errorbar(observation.wavelengthGrid,
observation.spectrum, observation.errorBar,
fmt='.',
label='observation')
if is_lightcurve:
ax.plot(result[1], label='forward model')
else:
if inst_result is not None:
from taurex.util.util import wnwidth_to_wlwidth
inst_wngrid, inst_spectrum, \
inst_noise, inst_width = inst_result
inst_wlgrid = 10000/inst_wngrid
inst_wlwidth = wnwidth_to_wlwidth(inst_wngrid, inst_width)
ax.errorbar(inst_wlgrid, inst_spectrum, inst_noise,
inst_wlwidth/2, '.', label='Instrument')
else:
ax.plot(wlgrid, binning.bin_model(
result)[1], label='forward model')
ax.set_xscale('log')
if args.contrib:
native_grid, contrib_result = model.model_contrib(
wngrid=wngrid)
for contrib_name, contrib in contrib_result.items():
flux, tau, extras = contrib
binned = binning.bindown(native_grid, flux)
if is_lightcurve:
ax.plot(binned[1], label=contrib_name)
else:
ax.plot(wlgrid, binned[1], label=contrib_name)
if args.full_contrib:
native_grid, contrib_result = model.model_full_contrib(
wngrid=wngrid)
for contrib_name, contrib in contrib_result.items():
for name, flux, tau, extras in contrib:
label = '{} - {}'.format(contrib_name, name)
binned = binning.bindown(native_grid, flux)
if is_lightcurve:
ax.plot(binned[1], label=label)
else:
ax.plot(wlgrid, binned[1], label=label)
plt.legend()
plt.show()
else:
logging.getLogger('taurex').warning(
'Number of processes > 1 so not plotting')
def __init__(self,
multi_nest_path=None,
observed=None,
model=None,
sampling_efficiency='parameter',
num_live_points=1500,
max_iterations=0,
search_multi_modes=True,
num_params_cluster=None,
maximum_modes=100,
constant_efficiency_mode=False,
evidence_tolerance=0.5,
mode_tolerance=-1e90,
importance_sampling=False,
resume=False,
multinest_prefix='1-',
verbose_output=True,
sigma_fraction=0.1):
super().__init__('Multinest', observed, model, sigma_fraction)
# sampling chains directory
self.nest_path = 'chains/'
self.nclust_par = -1
# sampling efficiency (parameter, ...)
self.sampling_eff = sampling_efficiency
# number of live points
self.n_live_points = int(num_live_points)
# maximum no. of iterations (0=inf)
self.max_iter = int(max_iterations)
# search for multiple modes
self.multimodes = search_multi_modes
# parameters on which to cluster, e.g. if nclust_par = 3, it will
# cluster on the first 3 parameters only.
# if ncluster_par = -1 it clusters on all parameters
self.nclust_par = num_params_cluster
# maximum number of modes
self.max_modes = int(maximum_modes)
# run in constant efficiency mode
self.const_eff = constant_efficiency_mode
# set log likelihood tolerance. If change is smaller, multinest will have converged
self.evidence_tolerance = evidence_tolerance
self.mode_tolerance = mode_tolerance
# importance nested sampling
self.imp_sampling = importance_sampling
if self.imp_sampling:
self.multimodes = False
self.multinest_prefix = multinest_prefix
self.dir_multinest = multi_nest_path
if get_rank() == 0:
if not os.path.exists(self.dir_multinest):
self.info('Directory %s does not exist, creating',
self.dir_multinest)
os.makedirs(self.dir_multinest)
barrier()
self.info('Found directory %s', self.dir_multinest)
self.resume = resume
self.verbose = verbose_output
temperature=temperature,
planet=planet,
star=star)
# Build it
model.build()
# Lets get the binner from the model
model_binner = model.defaultBinner()
# For fun lets make this work in mpi. This will have no
# effect if mpi is not installed
runs = list(range(30))
rank = get_rank()
size = nprocs()
# Now lets run the model several times and store
# the output to file
for n in runs[rank::size]:
# Run the forward model
model_result = model.model()
# Lets store it to HDF5 file, with light amount of data
small_size = OutputSize.light
with HDF5Output(f'myoutput_{n}_fm.h5') as o:
def compute_derived_trace(self, solution):
from taurex.util.util import quantile_corner
from taurex.constants import AMU
from taurex import mpi
enableLogging()
samples = self.get_samples(solution)
weights = self.get_weights(solution)
len_samples = len(samples)
rank = mpi.get_rank()
num_procs = mpi.nprocs()
count = 0
derived_param = {p: ([],[]) for p in self.derived_names}
weight_comb = []
if len(self.derived_names) == 0:
return
self.info('Computing derived parameters......')
disableLogging()
for idx in range(rank, len_samples, num_procs):
enableLogging()
if rank == 0 and count % 10 == 0 and count > 0:
self.info('Progress {}%'.format(
idx*100.0 / len_samples))
disableLogging()
parameters = samples[idx]
weight = weights[idx]
self.update_model(parameters)
self._model.initialize_profiles()
for p, v in zip(self.derived_names, self.derived_values):
derived_param[p][0].append(v)
derived_param[p][1].append(weight)
result_dict = {}
sorted_weights = weights.argsort()
for param, (trace, w) in derived_param.items():
all_trace = np.array(mpi.allreduce(trace, op='SUM')) # I cant remember why this works
all_weight = np.array(mpi.allreduce(w, op='SUM')) # I cant remember why this works
all_weight_sort = all_weight.argsort()
# Sort them into the right order
all_weight[sorted_weights] = all_weight[all_weight_sort]
all_trace[sorted_weights] = all_trace[all_weight_sort]
q_16, q_50, q_84 = \
quantile_corner(np.array(all_trace), [0.16, 0.5, 0.84],
weights=np.array(all_weight))
mean = np.average(all_trace, weights=all_weight, axis=0)
derived = {
'value': q_50,
'sigma_m': q_50-q_16,
'sigma_p': q_84-q_50,
'trace': all_trace,
'mean': mean
}
result_dict[f'{param}_derived'] = derived
return result_dict