def __init__(self,
model,
log_dir=None,
stepsampling=False,
enable_plots=False,
**kwargs):
super(UltranestSampler, self).__init__(model)
import ultranest
log_likelihood_call, prior_call = setup_calls(model, copy_prior=True)
self._sampler = ultranest.ReactiveNestedSampler(list(
self.model.variable_params),
log_likelihood_call,
prior_call,
log_dir=log_dir,
resume=True)
if stepsampling:
import ultranest.stepsampler
self._sampler.stepsampler = ultranest.stepsampler.RegionBallSliceSampler(
nsteps=100,
adaptive_nsteps='move-distance',
region_filter=True)
self.enable_plots = enable_plots
self.nlive = 0
self.ndim = len(self.model.variable_params)
self.result = None
self.kwargs = kwargs # Keywords for the run method of ultranest
def run():
n_data = len(values)
samples = []
for i in range(n_data):
# draw normal random points
u = np.random.normal(size=400)
v = values[i] + np.where(u < 0, u * values_lo[i], u * values_hi[i])
samples.append(v)
data = np.array(samples)
Nobj, Nsamples = data.shape
minval = -80
maxval = +80
ndim = 8
viz_callback = None
bins = np.linspace(minval, maxval, ndim + 1)
binned_data = np.array([np.histogram(row, bins=bins)[0] for row in data])
param_names = ['bin%d' % (i + 1) for i in range(ndim)]
def likelihood(params):
"""Histogram model"""
return np.log(np.dot(binned_data, params) / Nsamples + 1e-300).sum()
def transform_dirichlet(quantiles):
"""Histogram distribution priors"""
# https://en.wikipedia.org/wiki/Dirichlet_distribution#Random_number_generation
# first inverse transform sample from Gamma(alpha=1,beta=1), which is Exponential(1)
gamma_quantiles = -np.log(quantiles)
# dirichlet variables
return gamma_quantiles / gamma_quantiles.sum()
stepsamplers = [
ultranest.stepsampler.RegionBallSliceSampler(
40, region_filter=False, adaptive_nsteps='move-distance'),
ultranest.stepsampler.RegionSliceSampler(40, region_filter=False),
ultranest.stepsampler.CubeSliceSampler(40, region_filter=True),
ultranest.stepsampler.RegionMHSampler(40, region_filter=False),
ultranest.stepsampler.RegionSequentialSliceSampler(
40, region_filter=True, adaptive_nsteps='move-distance'),
ultranest.stepsampler.SpeedVariableRegionSliceSampler(
step_matrix=[[1], [1, 2, 3], Ellipsis,
np.ones(len(param_names), dtype=bool)],
nsteps=40,
region_filter=False),
]
for stepsampler in stepsamplers:
print(stepsampler)
sampler = ultranest.ReactiveNestedSampler(param_names, likelihood,
transform_dirichlet)
sampler.stepsampler = stepsampler
sampler.run(frac_remain=0.5, viz_callback=viz_callback)
sampler.print_results()
def fit_nested(self):
freekeys = list(self.bounds.keys())
boundarray = np.array([self.bounds[k] for k in freekeys])
bounddiff = np.diff(boundarray, 1).reshape(-1)
def loglike(pars):
for i in range(len(pars)):
self.prior[freekeys[i]] = pars[i]
model, _, _ = moon_transit(self.time, self.prior)
return -0.5 * np.sum(((self.data - model) / self.dataerr)**2)
def prior_transform(upars):
# transform unit cube to prior volume
return (boundarray[:, 0] + bounddiff * upars)
if self.verbose:
self.results = ultranest.ReactiveNestedSampler(
freekeys, loglike, prior_transform).run(max_ncalls=5e5)
else:
self.results = ultranest.ReactiveNestedSampler(
freekeys, loglike,
prior_transform).run(max_ncalls=5e5,
show_status=self.verbose,
viz_callback=self.verbose)
self.errors = {}
self.quantiles = {}
self.parameters = copy.deepcopy(self.prior)
for i, key in enumerate(freekeys):
self.parameters[key] = self.results['maximum_likelihood']['point'][
i]
self.errors[key] = self.results['posterior']['stdev'][i]
self.quantiles[key] = [
self.results['posterior']['errlo'][i],
self.results['posterior']['errup'][i]
]
# self.results['maximum_likelihood']
self.create_fit_variables()
def test_sampling():
'''
We just need to initiate a sampler here, which we do by running un.ReactiveNestedSampler, since this automatically tests that the likelihood and prior transform etc works, thanks to the option num_test_samples.
'''
sampler = un.ReactiveNestedSampler(param_names,
lc_fit.log_likelihood,
lc_fit.prior_transform,
resume='overwrite',
log_dir=f'chains/',
num_test_samples=10,
vectorized=False,
draw_multiple=True)
def __init__(self, model, **kwargs):
super(UltranestSampler, self).__init__(model)
import ultranest
log_likelihood_call, prior_call = setup_calls(model, copy_prior=True)
self._sampler = ultranest.ReactiveNestedSampler(
list(self.model.variable_params),
log_likelihood_call,
prior_call)
self.nlive = 0
self.ndim = len(self.model.variable_params)
self.result = None
self.kwargs = kwargs # Keywords for the run method of ultranest
def run_ultranest(hierarchy: Hierarchy,
data: NeutrinoConstraint,
max_ncalls: int = 100_000):
hierarchy_str = 'nh' if hierarchy == Hierarchy.Normal else 'ih'
sum_str = str(data.sum_of_masses_one_sigma)[:5] # Max 5 sig fig
prefix = '' if data.sum_of_masses_offset == 0. else str(
data.sum_of_masses_offset) + '_'
run_name = prefix + 'likeli_' + hierarchy_str + '_' + sum_str
dir_name = "ultra_" + run_name
sampler = ultra.ReactiveNestedSampler(
PARAM_NAMES,
evaluate_log_likelihood_of_parameters,
prior_map, # num_live_points=400,
log_dir=dir_name,
resume='overwrite')
sampler.run(max_ncalls=max_ncalls)
def __init__(self,
model,
log_dir=None,
stepsampling=False,
enable_plots=False,
**kwargs):
super(UltranestSampler, self).__init__(model)
import ultranest
log_likelihood_call, prior_call = setup_calls(model, copy_prior=True)
# Check for cyclic boundaries
periodic = []
cyclic = self.model.prior_distribution.cyclic
for param in self.variable_params:
if param in cyclic:
logging.info('Param: %s will be cyclic', param)
periodic.append(True)
else:
periodic.append(False)
self._sampler = ultranest.ReactiveNestedSampler(
list(self.model.variable_params),
log_likelihood_call,
prior_call,
log_dir=log_dir,
wrapped_params=periodic,
resume=True)
if stepsampling:
import ultranest.stepsampler
self._sampler.stepsampler = ultranest.stepsampler.RegionBallSliceSampler(
nsteps=100,
adaptive_nsteps='move-distance',
region_filter=True)
self.enable_plots = enable_plots
self.nlive = 0
self.ndim = len(self.model.variable_params)
self.result = None
self.kwargs = kwargs # Keywords for the run method of ultranest
do_mpi, _, rank = use_mpi()
self.main = (not do_mpi) or (rank == 0)
def run(self):
"""
Executes the inference
"""
if self.prepare():
# Set up the inference
un = ultranest.ReactiveNestedSampler(list(self.pnames),
self.loglike,
self.prior,
log_dir=self.outputdir,
vectorized=True,
resume=self.resume)
# Run it
out = un.run(min_ess=self.min_ess,
max_iters=self.niter,
min_num_live_points=self.nlive,
frac_remain=self.frac_remain,
Lepsilon=self.Lepsilon,
dlogz=self.dlogz)
if self.verb:
un.print_results()
# Posterior and best parameters
self.bestp = np.array(un.results['maximum_likelihood']['point'])
self.outp = un.results['samples'].T
if self.kll is not None:
for i in range(self.outp.shape[-1]):
self.kll.update(self.model(self.outp[:, i], fullout=True))
# UltraNest Plotting
un.plot_corner()
un.plot_run()
un.plot_trace()
# Save posterior and bestfit params
if self.fsavefile is not None:
np.save(self.fsavefile, self.outp)
if self.fbestp is not None:
np.save(self.fbestp, self.bestp)
else:
if self.verb:
print("Sampler is not fully prepared to run. " + \
"Correct the above errors and try again.")
def run_ultranest_mcmc(
directory,
parameters,
freq,
flux,
err_flux,
model,
prior_transform,
resume="resume-similar",
run_num=1,
):
log = ultranest.utils.make_run_dir(directory, run_num=run_num)
# ultranest.utils.create_logger("ultranest", log_dir=directory)
sampler = ultranest.ReactiveNestedSampler(
param_names=parameters,
loglike=create_lnlike(freq, flux, err_flux, model),
transform=prior_transform,
log_dir=log["run_dir"],
resume=resume,
storage_backend="hdf5",
)
return sampler
def _fit(self,
model: AbstractPriorModel,
analysis,
log_likelihood_cap=None):
"""
Fit a model using Dynesty and the Analysis class which contains the data and returns the log likelihood from
instances of the model, which the `NonLinearSearch` seeks to maximize.
Parameters
----------
model : ModelMapper
The model which generates instances for different points in parameter space.
analysis : Analysis
Contains the data and the log likelihood function which fits an instance of the model to the data, returning
the log likelihood the `NonLinearSearch` maximizes.
Returns
-------
A result object comprising the Samples object that includes the maximum log likelihood instance and full
set of accepted ssamples of the fit.
"""
import ultranest
pool, pool_ids = self.make_pool()
fitness_function = self.fitness_function_from_model_and_analysis(
model=model,
analysis=analysis,
pool_ids=pool_ids,
log_likelihood_cap=log_likelihood_cap,
)
def prior_transform(cube):
return model.vector_from_unit_vector(unit_vector=cube)
sampler = ultranest.ReactiveNestedSampler(
param_names=model.parameter_names,
loglike=fitness_function.__call__,
transform=prior_transform,
log_dir=self.paths.samples_path,
**self.config_dict_search)
sampler.stepsampler = self.stepsampler
finished = False
while not finished:
try:
total_iterations = sampler.ncall
except AttributeError:
total_iterations = 0
if self.config_dict_run["max_ncalls"] is not None:
iterations = self.config_dict_run["max_ncalls"]
else:
iterations = total_iterations + self.iterations_per_update
if iterations > 0:
config_dict_run = self.config_dict_run
config_dict_run.pop("max_ncalls")
config_dict_run["dKL"] = config_dict_run.pop("dkl")
config_dict_run["Lepsilon"] = config_dict_run.pop("lepsilon")
config_dict_run["update_interval_ncall"] = iterations
sampler.run(max_ncalls=iterations, **config_dict_run)
self.paths.save_object("results", sampler.results)
self.perform_update(model=model,
analysis=analysis,
during_analysis=True)
iterations_after_run = sampler.ncall
if (total_iterations == iterations_after_run
or iterations_after_run
== self.config_dict_run["max_ncalls"]):
finished = True
def call(self, **kwargs):
"""
Runs the IMAGINE pipeline using the
`UltraNest <https://johannesbuchner.github.io/UltraNest/>`_
:py:class:`ReactiveNestedSampler <ultranest.integrator.ReactiveNestedSampler>`.
Any keyword argument provided is used to update the
`sampling_controllers`.
Returns
-------
results : dict
UltraNest sampling results in a dictionary containing the keys:
logZ (the log-evidence), logZerror (the error in log-evidence) and
samples (equal weighted posterior)
Notes
-----
See base class for other attributes/properties and methods
"""
log.debug('@ ultranest_pipeline::__call__')
default_init_params = {
'resume': True,
'num_test_samples': 2,
'num_bootstraps': 30,
'draw_multiple': True}
default_run_params = {
'dlogz': 0.5,
'dKL': 0.5,
'frac_remain': 0.01,
'Lepsilon': 0.001,
'min_ess': 500,
'max_iters': None,
'max_ncalls': None,
'max_num_improvement_loops': -1,
'min_num_live_points': 400,
'cluster_num_live_points': 40,
'update_interval_volume_fraction': 0.2}
# Keyword arguments can alter the sampling controllers
self.sampling_controllers = kwargs # Updates the dict
# Prepares initialization and run parameters from
# defaults and sampling controllers
init_params = { k : self.sampling_controllers.get(k, default)
for k, default in default_init_params.items()}
run_params = { k : self.sampling_controllers.get(k, default)
for k, default in default_run_params.items()}
# Updates the sampling controllers to reflect what is being used
self.sampling_controllers = init_params # Updates the dict
self.sampling_controllers = run_params # Updates the dict
# Ultranest files directory
ultranest_dir = path.join(self.chains_directory, 'ultranest')
# Creates directory, if needed
os.makedirs(ultranest_dir, exist_ok=True)
# Cleans up the chains directory if not resuming
if not init_params['resume']:
init_params['resume'] = 'overwrite'
# Removing manually as UltraNest's 'overwrite' option does not
# seem to be working correctly
self.clean_chains_directory()
# Creates directory, if needed
os.makedirs(ultranest_dir, exist_ok=True)
# Runs UltraNest
sampler = ultranest.ReactiveNestedSampler(
param_names=list(self.active_parameters),
loglike=self._likelihood_function,
transform=self.prior_transform,
log_dir=ultranest_dir,
vectorized=False,
wrapped_params=self.wrapped_parameters,
**init_params)
self.results = sampler.run(viz_callback=ultranest.viz.nicelogger,
**run_params)
self._samples_array = self.results['samples']
self._evidence = self.results['logz']
self._evidence_err = self.results['logzerr']
return self.results
def sample(self, quiet=False):
"""
sample using the UltraNest numerical integration method
:rtype:
:returns:
"""
if not self._is_setup:
print("You forgot to setup the sampler!")
return
loud = not quiet
self._update_free_parameters()
param_names = list(self._free_parameters.keys())
n_dim = len(param_names)
loglike, ultranest_prior = self._construct_unitcube_posterior(
return_copy=True)
# We need to check if the MCMC
# chains will have a place on
# the disk to write and if not,
# create one
chain_name = self._kwargs.pop("chain_name")
if chain_name is not None:
mcmc_chains_out_dir = ""
tmp = chain_name.split("/")
for s in tmp[:-1]:
mcmc_chains_out_dir += s + "/"
if using_mpi:
# if we are running in parallel and this is not the
# first engine, then we want to wait and let everything finish
if rank != 0:
# let these guys take a break
time.sleep(1)
else:
# create mcmc chains directory only on first engine
if not os.path.exists(mcmc_chains_out_dir):
os.makedirs(mcmc_chains_out_dir)
else:
if not os.path.exists(mcmc_chains_out_dir):
os.makedirs(mcmc_chains_out_dir)
# Multinest must be run parallel via an external method
# see the demo in the examples folder!!
if threeML_config["parallel"]["use-parallel"]:
raise RuntimeError(
"If you want to run ultranest in parallell you need to use an ad-hoc method"
)
else:
sampler = ultranest.ReactiveNestedSampler(
param_names,
loglike,
transform=ultranest_prior,
log_dir=chain_name,
vectorized=False,
wrapped_params=self._wrapped_params,
)
with use_astromodels_memoization(False):
sampler.run(show_status=loud, **self._kwargs)
process_fit = False
if using_mpi:
# if we are running in parallel and this is not the
# first engine, then we want to wait and let everything finish
if rank != 0:
# let these guys take a break
time.sleep(5)
# these engines do not need to read
process_fit = False
else:
# wait for a moment to allow it all to turn off
time.sleep(5)
process_fit = True
else:
process_fit = True
if process_fit:
results = sampler.results
self._sampler = sampler
ws = results["weighted_samples"]
weights = ws["weights"]
# Get the log. likelihood values from the chain
SQRTEPS = (float(np.finfo(np.float64).eps))**0.5
if abs(np.sum(weights) -
1.0) > SQRTEPS: # same tol as in np.random.choice.
raise ValueError("weights do not sum to 1")
rstate = np.random
N = len(weights)
# make N subdivisions, and choose positions with a consistent random offset
positions = (rstate.random() + np.arange(N)) / N
idx = np.zeros(N, dtype=np.int)
cumulative_sum = np.cumsum(weights)
i, j = 0, 0
while i < N:
if positions[i] < cumulative_sum[j]:
idx[i] = j
i += 1
else:
j += 1
self._log_like_values = ws["logl"][idx]
self._raw_samples = ws["points"][idx]
# now get the log probability
self._log_probability_values = self._log_like_values + np.array(
[self._log_prior(samples) for samples in self._raw_samples])
self._build_samples_dictionary()
self._marginal_likelihood = sampler.results["logz"] / np.log(10.0)
self._build_results()
# Display results
if loud:
self._results.display()
# now get the marginal likelihood
return self.samples
def fit_gaussian(
self,
tag: str,
min_num_live_points: float = 400,
bounds: Dict[str, Union[Tuple[float, float]]] = None,
output: str = "ultranest/",
plot_filename: Optional[str] = "line_fit.pdf",
show_status: bool = True,
double_gaussian: bool = False,
) -> None:
"""
Method for fitting a Gaussian profile to an emission line and
using ``UltraNest`` for sampling the posterior distributions
and estimating the evidence.
Parameters
----------
tag : str
Database tag where the posterior samples will be stored.
min_num_live_points : int
Minimum number of live points (see
https://johannesbuchner.github.io/UltraNest/issues.html).
bounds : dict(str, tuple(float, float)), None
The boundaries that are used for the uniform priors of the
3 Gaussian parameters (``gauss_amplitude``, ``gauss_mean``,
and ``gauss_sigma``). Conservative prior boundaries will
be estimated from the spectrum if the argument is set to
``None`` or if any of the required parameters is missing
in the ``bounds`` dictionary.
output : str
Path that is used for the output files from ``UltraNest``.
plot_filename : str
Filename for the plot with the best-fit line profile.
The plot is shown in an interface window if the
argument is set to ``None``.
show_status : bool
Print information about the convergence.
double_gaussian : bool
Set to ``True`` for fitting a double instead of a single
Gaussian. In that case, the ``bounds`` dictionary may also
contain ``'gauss_amplitude_2'``, ``'gauss_mean_2'``, and
``'gauss_sigma_2'`` (otherwise conservative parameter
boundaries are estimated from the data).
Returns
-------
NoneType
None
"""
high_spec_res = 1e5
@typechecked
def gaussian_function(
amplitude: float, mean: float, sigma: float, wavel: np.ndarray
):
return amplitude * np.exp(-0.5 * (wavel - mean) ** 2 / sigma ** 2)
# Model parameters
modelpar = ["gauss_amplitude", "gauss_mean", "gauss_sigma"]
if double_gaussian:
modelpar.append("gauss_amplitude_2")
modelpar.append("gauss_mean_2")
modelpar.append("gauss_sigma_2")
# Create a dictionary with the cube indices of the parameters
cube_index = {}
for i, item in enumerate(modelpar):
cube_index[item] = i
# Check if all prior boundaries are present
if bounds is None:
bounds = {}
if "gauss_amplitude" not in bounds:
bounds["gauss_amplitude"] = (0.0, 2.0 * np.amax(self.spectrum[:, 1]))
if "gauss_mean" not in bounds:
bounds["gauss_mean"] = (self.spectrum[0, 0], self.spectrum[-1, 0])
if "gauss_sigma" not in bounds:
bounds["gauss_sigma"] = (0.0, self.spectrum[-1, 0] - self.spectrum[0, 0])
if double_gaussian:
if "gauss_amplitude_2" not in bounds:
bounds["gauss_amplitude_2"] = (0.0, 2.0 * np.amax(self.spectrum[:, 1]))
if "gauss_mean_2" not in bounds:
bounds["gauss_mean_2"] = (self.spectrum[0, 0], self.spectrum[-1, 0])
if "gauss_sigma_2" not in bounds:
bounds["gauss_sigma_2"] = (
0.0,
self.spectrum[-1, 0] - self.spectrum[0, 0],
)
# Get the MPI rank of the process
try:
from mpi4py import MPI
mpi_rank = MPI.COMM_WORLD.Get_rank()
except ModuleNotFoundError:
mpi_rank = 0
# Create the output folder if required
if mpi_rank == 0 and not os.path.exists(output):
os.mkdir(output)
@typechecked
def lnprior_ultranest(cube: np.ndarray) -> np.ndarray:
"""
Function for transforming the unit cube
into the parameter cube.
Parameters
----------
cube : np.ndarray
Array with unit parameters.
Returns
-------
np.ndarray
Array with physical parameters.
"""
params = cube.copy()
for item in cube_index:
# Uniform priors for all parameters
params[cube_index[item]] = (
bounds[item][0]
+ (bounds[item][1] - bounds[item][0]) * params[cube_index[item]]
)
return params
@typechecked
def lnlike_ultranest(params: np.ndarray) -> np.float64:
"""
Function for calculating the log-likelihood for the
sampled parameter cube.
Parameters
----------
params : np.ndarray
Cube with physical parameters.
Returns
-------
float
Log-likelihood.
"""
data_flux = self.spectrum[:, 1]
data_var = self.spectrum[:, 2] ** 2
model_flux = gaussian_function(
params[cube_index["gauss_amplitude"]],
params[cube_index["gauss_mean"]],
params[cube_index["gauss_sigma"]],
self.spectrum[:, 0],
)
if double_gaussian:
model_flux += gaussian_function(
params[cube_index["gauss_amplitude_2"]],
params[cube_index["gauss_mean_2"]],
params[cube_index["gauss_sigma_2"]],
self.spectrum[:, 0],
)
chi_sq = -0.5 * (data_flux - model_flux) ** 2 / data_var
return np.nansum(chi_sq)
sampler = ultranest.ReactiveNestedSampler(
modelpar,
lnlike_ultranest,
transform=lnprior_ultranest,
resume="subfolder",
log_dir=output,
)
result = sampler.run(
show_status=show_status,
viz_callback=False,
min_num_live_points=min_num_live_points,
)
# Log-evidence
ln_z = result["logz"]
ln_z_error = result["logzerr"]
print(f"Log-evidence = {ln_z:.2f} +/- {ln_z_error:.2f}")
# Best-fit parameters
print("Best-fit parameters (mean +/- std):")
for i, item in enumerate(modelpar):
mean = np.mean(result["samples"][:, i])
std = np.std(result["samples"][:, i])
print(f" - {item} = {mean:.2e} +/- {std:.2e}")
# Maximum likelihood sample
print("Maximum likelihood sample:")
max_lnlike = result["maximum_likelihood"]["logl"]
print(f" - Log-likelihood = {max_lnlike:.2e}")
for i, item in enumerate(result["maximum_likelihood"]["point"]):
print(f" - {modelpar[i]} = {item:.2e}")
# Posterior samples
samples = result["samples"]
# Best-fit model parameters
model_param = {
"gauss_amplitude": np.median(samples[:, 0]),
"gauss_mean": np.median(samples[:, 1]),
"gauss_sigma": np.median(samples[:, 2]),
}
if double_gaussian:
model_param["gauss_amplitude_2"] = np.median(samples[:, 3])
model_param["gauss_mean_2"] = np.median(samples[:, 4])
model_param["gauss_sigma_2"] = np.median(samples[:, 5])
best_model = read_util.gaussian_spectrum(
self.wavel_range,
model_param,
spec_res=high_spec_res,
double_gaussian=double_gaussian,
)
# Interpolate high-resolution continuum
if self.continuum_check:
cont_interp = interp1d(
self.spectrum[:, 0], self.continuum_flux, bounds_error=False
)
cont_high_res = cont_interp(best_model.wavelength)
else:
cont_high_res = np.full(best_model.wavelength.shape[0], 0.0)
# Add FWHM velocity
modelpar.append("gauss_fwhm")
gauss_mean = samples[:, 1] # (um)
gauss_fwhm = 2.0 * np.sqrt(2.0 * np.log(2.0)) * samples[:, 2] # (um)
vel_fwhm = 1e-3 * constants.LIGHT * gauss_fwhm / gauss_mean # (km s-1)
vel_fwhm = vel_fwhm[..., np.newaxis]
samples = np.append(samples, vel_fwhm, axis=1)
# Add line flux and luminosity
print("Calculating line fluxes...", end="", flush=True)
modelpar.append("line_flux")
modelpar.append("line_luminosity")
line_flux = np.zeros(samples.shape[0])
line_lum = np.zeros(samples.shape[0])
if self.continuum_check:
modelpar.append("line_eq_width")
eq_width = np.zeros(samples.shape[0])
for i in range(samples.shape[0]):
model_param = {
"gauss_amplitude": samples[i, 0],
"gauss_mean": samples[i, 1],
"gauss_sigma": samples[i, 2],
}
if double_gaussian:
model_param["gauss_amplitude_2"] = samples[i, 3]
model_param["gauss_mean_2"] = samples[i, 4]
model_param["gauss_sigma_2"] = samples[i, 5]
model_box = read_util.gaussian_spectrum(
self.wavel_range,
model_param,
spec_res=high_spec_res,
double_gaussian=double_gaussian,
)
line_flux[i] = np.trapz(model_box.flux, model_box.wavelength) # (W m-2)
line_lum[i] = (
4.0 * np.pi * (1e3 * constants.PARSEC / self.parallax) ** 2 * line_flux[i]
) # (W)
line_lum[i] /= constants.L_SUN # (Lsun)
if self.continuum_check:
# Normalize the spectrum to the continuum
spec_norm = (model_box.flux + cont_high_res) / cont_high_res
# Check if the flux is NaN (due to interpolation errors at the spectrum edge)
indices = ~np.isnan(spec_norm)
eq_width[i] = np.trapz(
1.0 - spec_norm[indices], model_box.wavelength[indices]
) # (um)
eq_width[i] *= 1e4 # (A)
line_flux = line_flux[..., np.newaxis]
samples = np.append(samples, line_flux, axis=1)
line_lum = line_lum[..., np.newaxis]
samples = np.append(samples, line_lum, axis=1)
if self.continuum_check:
eq_width = eq_width[..., np.newaxis]
samples = np.append(samples, eq_width, axis=1)
if self.lambda_rest is not None:
# Radial velocity (km s-1)
if double_gaussian:
# Weighted (with Gaussian amplitudes) mean of the central wavelength
v_fit = (
samples[:, 0] * samples[:, 1] + samples[:, 3] * samples[:, 4]
) / (samples[:, 0] + samples[:, 3])
else:
v_fit = samples[:, 1]
v_rad = (
1e-3 * constants.LIGHT * (v_fit - self.lambda_rest) / self.lambda_rest
)
v_rad = v_rad[..., np.newaxis]
modelpar.append("line_vrad")
samples = np.append(samples, v_rad, axis=1)
print(" [DONE]")
# Log-likelihood
ln_prob = result["weighted_samples"]["logl"]
# Log-evidence
ln_z = result["logz"]
ln_z_error = result["logzerr"]
print(f"Log-evidence = {ln_z:.2f} +/- {ln_z_error:.2f}")
# Get the MPI rank of the process
try:
from mpi4py import MPI
mpi_rank = MPI.COMM_WORLD.Get_rank()
except ModuleNotFoundError:
mpi_rank = 0
# Add samples to the database
if mpi_rank == 0:
# Writing the samples to the database is only
# possible when using a single process
species_db = database.Database()
species_db.add_samples(
sampler="ultranest",
samples=samples,
ln_prob=ln_prob,
ln_evidence=(ln_z, ln_z_error),
mean_accept=None,
spectrum=("model", "gaussian"),
tag=tag,
modelpar=modelpar,
parallax=self.parallax,
spec_labels=None,
)
# Create plot
if plot_filename is None:
print("Plotting best-fit line profile...", end="", flush=True)
else:
print(f"Plotting best-fit line profile: {plot_filename}...", end="", flush=True)
mpl.rcParams["font.serif"] = ["Bitstream Vera Serif"]
mpl.rcParams["font.family"] = "serif"
plt.rc("axes", edgecolor="black", linewidth=2)
plt.rcParams["axes.axisbelow"] = False
plt.figure(1, figsize=(6, 6))
gs = mpl.gridspec.GridSpec(2, 1)
gs.update(wspace=0, hspace=0.1, left=0, right=1, bottom=0, top=1)
ax1 = plt.subplot(gs[0, 0])
ax2 = plt.subplot(gs[1, 0])
ax3 = ax1.twiny()
ax4 = ax2.twiny()
ax1.tick_params(
axis="both",
which="major",
colors="black",
labelcolor="black",
direction="in",
width=1,
length=5,
labelsize=12,
top=False,
bottom=True,
left=True,
right=True,
labelbottom=False,
)
ax1.tick_params(
axis="both",
which="minor",
colors="black",
labelcolor="black",
direction="in",
width=1,
length=3,
labelsize=12,
top=False,
bottom=True,
left=True,
right=True,
labelbottom=False,
)
ax2.tick_params(
axis="both",
which="major",
colors="black",
labelcolor="black",
direction="in",
width=1,
length=5,
labelsize=12,
top=False,
bottom=True,
left=True,
right=True,
)
ax2.tick_params(
axis="both",
which="minor",
colors="black",
labelcolor="black",
direction="in",
width=1,
length=3,
labelsize=12,
top=False,
bottom=True,
left=True,
right=True,
)
ax3.tick_params(
axis="both",
which="major",
colors="black",
labelcolor="black",
direction="in",
width=1,
length=5,
labelsize=12,
top=True,
bottom=False,
left=True,
right=True,
)
ax3.tick_params(
axis="both",
which="minor",
colors="black",
labelcolor="black",
direction="in",
width=1,
length=3,
labelsize=12,
top=True,
bottom=False,
left=True,
right=True,
)
ax4.tick_params(
axis="both",
which="major",
colors="black",
labelcolor="black",
direction="in",
width=1,
length=5,
labelsize=12,
top=True,
bottom=False,
left=True,
right=True,
labeltop=False,
)
ax4.tick_params(
axis="both",
which="minor",
colors="black",
labelcolor="black",
direction="in",
width=1,
length=3,
labelsize=12,
top=True,
bottom=False,
left=True,
right=True,
labeltop=False,
)
ax1.set_ylabel("Flux (W m$^{-2}$ µm$^{-1}$)", fontsize=16)
ax2.set_xlabel("Wavelength (µm)", fontsize=16)
ax2.set_ylabel("Flux (W m$^{-2}$ µm$^{-1}$)", fontsize=16)
ax3.set_xlabel("Velocity (km s$^{-1}$)", fontsize=16)
ax1.get_yaxis().set_label_coords(-0.1, 0.5)
ax2.get_xaxis().set_label_coords(0.5, -0.1)
ax2.get_yaxis().set_label_coords(-0.1, 0.5)
ax3.get_xaxis().set_label_coords(0.5, 1.12)
ax1.plot(
self.spectrum[:, 0],
self.spectrum[:, 1] + self.continuum_flux,
color="black",
label=self.spec_name,
)
ax1.plot(
best_model.wavelength,
best_model.flux + cont_high_res,
color="tab:blue",
label="Best-fit model (continuum + line)",
)
ax2.plot(
self.spectrum[:, 0],
self.spectrum[:, 1],
color="black",
label=self.spec_name,
)
ax2.plot(
best_model.wavelength,
best_model.flux,
color="tab:blue",
label="Best-fit line profile",
)
ax3.plot(
self.spec_vrad, self.spectrum[:, 1] + self.continuum_flux, ls="-", lw=0.0
)
ax4.plot(self.spec_vrad, self.spectrum[:, 1], ls="-", lw=0.0)
ax1.legend(loc="upper left", frameon=False, fontsize=12.0)
ax2.legend(loc="upper left", frameon=False, fontsize=12.0)
print(" [DONE]")
if plot_filename is None:
plt.show()
else:
plt.savefig(plot_filename, bbox_inches="tight")
plt.clf()
plt.close()
pmax = param.prior._defaults['pmax']
transforms.append(uniform_trans(pmin, pmax))
elif param.type.lower() == 'normal':
mu = param.prior._defaults['mu']
sigma = param.prior._defaults['sigma']
transforms.append(normal_trans(mu, sigma))
elif param.type.lower() == 'ace_swepam_parameter':
transforms.append(sw_trans())
def transform(quantile):
return np.array([t(q) for q, t in zip(quantile, transforms)])
sampler1 = ultranest.ReactiveNestedSampler(
pta.param_names,
pta.get_lnlikelihood,
transform,
log_dir=args.outdir,
resume=True,
)
ndim = len(pta.params)
sampler1.stepsampler = ultranest.stepsampler.RegionSliceSampler(nsteps=2 *
ndim)
sampler1.run(
dlogz=0.5 + 0.1 * ndim,
# update_interval_iter_fraction=0.4 if ndim > 20 else 0.2,
# max_num_improvement_loops=3,
min_num_live_points=400)
m = theta[0]
c = theta[1]
# normalisation
norm = -0.5 * M * LN2PI - M * LNSIGMA
# chi-squared (data, sigma and x are global variables defined early on in this notebook)
chisq = np.sum(((data - straight_line(x, m, c)) / sigma)**2)
return norm - 0.5 * chisq
tol = 0.5 # the stopping criterion
# set the ReactiveNestedSampler method
sampler = ultranest.ReactiveNestedSampler(["m", "c"], loglikelihood,
prior_transform)
tol = 0.5 # the stopping criterion
# run the nested sampling algorithm
result = sampler.run(dlogz=tol)
logZultranest = result['logz'] # value of logZ
logZerrultranest = result[
'logzerr'] # estimate of the statistcal uncertainty on logZ
# output marginal likelihood
print('Marginalised evidence is {} ± {}'.format(logZultranest,
logZerrultranest))
# get the posterior samples
return params
names = [r'$A$', r'$\alpha$', r'$\beta$', r'$x_0$', r'$e_0$']
def likelihood(params):
model = modified_sch_log0_list(xdata, *params)
like = -0.5 * (((model - pdf_xoff_list) / yerr_list)**2).sum()
#like = -np.sum(model-(pdf_xoff[indexx])*np.log(model))
return like
#######UNCOMMENT TO USE ULTRANEST#############
sampler_rseppi = ultranest.ReactiveNestedSampler(names, likelihood,
my_flat_priors)
print('running sampler...')
result = sampler_rseppi.run()
sampler_rseppi.print_results()
tafterfit = time.time()
tem = tafterfit - t0
print('ci ho messo ', tem, 's')
#from ultranest.plot import PredictionBand
#fit = PredictionBand([log1sig, logx])
#print(result)
popt_pdf_xoff_list = np.array(result['posterior']['mean'])
pvar_pdf_xoff_list = np.array(result['posterior']['stdev'])
print('Best fit parameters = ', popt_pdf_xoff)
def fit_source(fit_struct,
data_struct,
filter_struct,
model_struct,
Parallel=0):
# detection mask to differenciate
# the upper limits from detections
# necessary to feed the chi2 calculation
detection_mask = []
for i in range(len(data_struct)):
detection_mask.append(data_struct[i]['det_type'] == 'd')
# flatten the parameters structure for processing
ndim = len(ut.flatten_model_keyword(model_struct, 'param'))
flat_param = ut.flatten_model_keyword(model_struct, 'current')
flat_min = ut.flatten_model_keyword(model_struct, 'min')
flat_max = ut.flatten_model_keyword(model_struct, 'max')
coef_param = np.array([(flat_max - flat_min) / 2.]).flatten()
# # initiate the walkers "ball" - TODO update with emcee function
# if os.path.isfile(fit_struct['sampler_file']):
# print("file found at ",fit_struct['sampler_file'])
# #with open(fit_struct['sampler_file'],"rb") as f:
# # chain = pickle.load(f)
# #last_index = chain.chain[0].nonzero()[0][-1]
# #print last_index
# #pos = [chain.chain[i][last_index] for i in range(fit_struct['nwalkers'])]
# else:
pos = [
flat_param + 1e-4 * coef_param * np.random.randn(ndim)
for i in range(fit_struct['nwalkers'])
]
print("initialise walker positions")
last_index = 0
print()
print('HMC attempt for ' + fit_struct['source'])
# single processor
# if os.path.isfile(fit_struct['sampler_file']):
# with open(fit_struct['sampler_file'],"rb") as f:
# sampler = pickle.load(f)
# else:
if fit_struct['fit_method'] == 'emcee':
# create the backend for emcee
filename = fit_struct['sampler_file']
backend = emcee.backends.HDFBackend(filename)
# if restart:
# print('continuing last run')
# else:
# backend.reset(fit_struct['nwalkers'],ndim)
backend.reset(fit_struct['nwalkers'], ndim)
if Parallel == 0:
sampler = emcee.EnsembleSampler(fit_struct['nwalkers'],
ndim,
lnprob,
args=(fit_struct, data_struct,
filter_struct, model_struct,
detection_mask),
backend=backend)
else:
# multi-processing (pool created via pathos, allow to pickle the sampler)
tmp_pool = mp.ProcessingPool(Parallel)
sampler = emcee.EnsembleSampler(fit_struct['nwalkers'],
ndim,
lnprob,
args=(data_struct, filter_struct,
model_struct, detection_mask,
fit_struct['redshift']),
pool=tmp_pool,
backend=backend)
# progress bar (work for multiprocess or single process)
# for sample in sampler.sample(pos,iterations=fit_source['nsteps'],progress=True):
with tqdm(total=fit_struct['nsteps']) as pbar:
for i, result in enumerate(
sampler.sample(pos, iterations=fit_struct['nsteps'])):
pbar.update()
print('HMC done!')
#pbar = tqdm(total=fit_struct['nsteps'],initial=last_index)
#print fit_struct['nsteps']-last_index
# sampler.run_mcmc(None,fit_struct['nsteps'])
# for i,_ in enumerate(sampler.sample(None, fit_struct['nsteps'])):
# pbar.update()
# if i % 10 == 0:
# pass
# with open(fit_struct['sampler_file'], 'w') as output_savefile:
# pickle.dump(sampler, output_savefile, pickle.HIGHEST_PROTOCOL)
# print i,' sampler saved!'
#print('HMC done!')
# TODO: transformation of chains
# save the modified sampler (allows to save pools as well - pathos library allows to serialise pools)
#with open(fit_struct['sampler_file'], 'ab') as output_savefile:
# pickle.dump(sampler, output_savefile, pickle.HIGHEST_PROTOCOL)
# print 'sampler saved!'
return sampler
elif fit_struct['fit_method'] == 'ultranest':
param_names = list(ut.flatten_model_keyword(model_struct, 'param'))
# create the inputs functions from decorators
@decorator_like(fit_struct, data_struct, filter_struct, model_struct,
detection_mask)
def like_func(theta, *args):
return lnlike(theta, *args)
@decorator_prior(model_struct)
def prior_func(theta, *args):
return lnprior_un(theta, *args)
# sampler_un = ultranest.ReactiveNestedSampler(param_names, like_func, prior_func, resume=True, log_dir='outputs/tmp_ultra')
sampler_un = ultranest.ReactiveNestedSampler(param_names, like_func,
prior_func)
sampler_un.run()
print('Ultranest done!')
with open(fit_struct['sampler_file'], 'wb') as file:
dill.dump(sampler_un, file)
return sampler_un
else:
print('error, wrong fit_method provided')
return 0
