def compute_mu_derived_trace(self, solution):
from taurex.util.util import quantile_corner
from taurex.constants import AMU
sigma_frac = self._sigma_fraction
self._sigma_fraction = 1.0
mu_trace = []
weights = []
self.info('Computing derived mu......')
disableLogging()
for parameters, weight in self.sample_parameters(solution):
self.update_model(parameters)
self._model.initialize_profiles()
mu_trace.append(self._model.chemistry.muProfile[0] / AMU)
weights.append(weight)
enableLogging()
self.info('Done!')
self._sigma_fraction = sigma_frac
q_16, q_50, q_84 = \
quantile_corner(np.array(mu_trace), [0.16, 0.5, 0.84],
weights=np.array(weights))
mean = np.average(mu_trace, weights=weights, axis=0)
mu_derived = {
'value': q_50,
'sigma_m': q_50 - q_16,
'sigma_p': q_84 - q_50,
'trace': mu_trace,
'mean': mean
}
return mu_derived
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
def fit(self, output_size=OutputSize.heavy):
import time
"""
Performs fit.
"""
from tabulate import tabulate
self.compile_params()
fit_names = self.fit_names
fit_boundaries = self.fit_boundaries
fit_min = [x[0] for x in fit_boundaries]
fit_max = [x[1] for x in fit_boundaries]
fit_values = self.fit_values
self.info('')
self.info('-------------------------------------')
self.info('------Retrieval Parameters-----------')
self.info('-------------------------------------')
self.info('')
self.info('Dimensionality of fit: %s', len(fit_names))
self.info('')
output = tabulate(zip(fit_names, fit_values, fit_min, fit_max),
headers=['Param', 'Value', 'Bound-min', 'Bound-max'])
self.info('\n%s\n\n', output)
self.info('')
start_time = time.time()
disableLogging()
# Compute fit here
self.compute_fit()
enableLogging()
end_time = time.time()
self.info('Sampling time %s s', end_time - start_time)
solution = self.generate_solution(output_size=output_size)
self.info('')
self.info('-------------------------------------')
self.info('------Final results------------------')
self.info('-------------------------------------')
self.info('')
self.info('Dimensionality of fit: %s', len(fit_names))
self.info('')
for idx, optimized_map, optimized_median, values in self.get_solution(
):
self.info('\n%s', '---Solution {}------'.format(idx))
output = tabulate(zip(fit_names, optimized_map, optimized_median),
headers=['Param', 'MAP', 'Median'])
self.info('\n%s\n\n', output)
return solution
def fit(self, output_size=OutputSize.heavy):
import time
"""
Performs fit.
"""
from tabulate import tabulate
self.compile_params()
fit_names = self.fit_names
prior_type = [p.__class__.__name__ for p in self.fitting_priors]
args = [p.params() for p in self.fitting_priors]
fit_values = self.fit_values
self.info('')
self.info('-------------------------------------')
self.info('------Retrieval Parameters-----------')
self.info('-------------------------------------')
self.info('')
self.info('Dimensionality of fit: %s', len(fit_names))
self.info('')
output = tabulate(zip(fit_names, fit_values, prior_type, args),
headers=['Param', 'Value', 'Type', 'Args'])
self.info('\n%s\n\n', output)
self.info('')
start_time = time.time()
disableLogging()
# Compute fit here
self.compute_fit()
enableLogging()
end_time = time.time()
self.info('Sampling time %s s', end_time-start_time)
solution = self.generate_solution(output_size=output_size)
self.info('')
self.info('-------------------------------------')
self.info('------Final results------------------')
self.info('-------------------------------------')
self.info('')
self.info('Dimensionality of fit: %s', len(fit_names))
self.info('')
for idx, optimized_map, optimized_median, values in self.get_solution():
self.info('\n%s', '---Solution {}------'.format(idx))
output = tabulate(zip(fit_names, optimized_map, optimized_median), headers=[
'Param', 'MAP', 'Median'])
self.info('\n%s\n\n', output)
return solution
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 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
def generate_solution(self, output_size=OutputSize.heavy):
"""
Generates a dictionary with all solutions and other useful parameters
"""
from taurex.util.output import generate_profile_dict, \
store_contributions
solution_dict = {}
self.info('Post-processing - Generating spectra and profiles')
# Loop through each solution, grab optimized parameters and anything
# else we want to store
for solution, optimized_map, \
optimized_median, values in self.get_solution():
enableLogging()
self.info('Computing solution %s', solution)
sol_values = {}
# Include extra stuff we might want to store (provided by the child)
for k, v in values:
sol_values[k] = v
# Update the model with optimized map values
self.update_model(optimized_map)
opt_result = self._model.model(cutoff_grid=False) # Run the model
sol_values['Spectra'] = self._binner.generate_spectrum_output(
opt_result, output_size=output_size)
try:
sol_values['Spectra']['Contributions'] = store_contributions(
self._binner, self._model, output_size=output_size-3)
except Exception as e:
self.warning('Not bothering to store contributions since its broken')
self.warning('%s ', str(e))
# Update with the optimized median
self.update_model(optimized_median)
self._model.model(cutoff_grid=False)
# Store profiles here
sol_values['Profiles'] = self._model.generate_profiles()
profile_dict, spectrum_dict = self.generate_profiles(
solution, self._observed.wavenumberGrid)
for k, v in profile_dict.items():
if k in sol_values['Profiles']:
sol_values['Profiles'][k].update(v)
else:
sol_values['Profiles'][k] = v
for k, v in spectrum_dict.items():
if k in sol_values['Spectra']:
sol_values['Spectra'][k].update(v)
else:
sol_values['Spectra'][k] = v
solution_dict['solution{}'.format(solution)] = sol_values
if len(self.derived_names) > 0:
#solution_dict[f'solution{solution}']['derived_params'] = {}
# Compute derived
for solution, optimized_map, \
optimized_median, values in self.get_solution():
solution_dict[f'solution{solution}']['derived_params'] = {}
derived_dict = self.compute_derived_trace(solution)
if derived_dict is None:
continue
solution_dict[f'solution{solution}']['derived_params'].update(derived_dict)
enableLogging()
self.info('Post-processing - Complete')
return solution_dict