Skip to content

Commit

Permalink
Shared prior function in FitModel
Browse files Browse the repository at this point in the history
  • Loading branch information
tomasstolker committed Feb 3, 2024
1 parent 94c5889 commit b887861
Show file tree
Hide file tree
Showing 3 changed files with 133 additions and 105 deletions.
5 changes: 3 additions & 2 deletions species/fit/fit_evolution.py
Original file line number Diff line number Diff line change
@@ -1,6 +1,7 @@
"""
Module with functionalities for retrieving the age and bulk parameters
of one or multiple planets and/or brown dwarfs in a system.
Module with functionalities for retrieving the age and
bulk parameters of one or multiple planets and/or brown
dwarfs in a system.
"""

import configparser
Expand Down
227 changes: 128 additions & 99 deletions species/fit/fit_model.py
Original file line number Diff line number Diff line change
Expand Up @@ -1064,25 +1064,73 @@ def __init__(
print(f" - {phot_item} = {self.weights[phot_item]:.2e}")

@typechecked
def lnlike_func(
self, params, prior: Optional[Dict[str, Tuple[float, float]]]
def _prior_transform(
self, cube, bounds: Dict[str, Tuple[float, float]], cube_index: Dict[str, int]
):
"""
Function to transform the sampled unit cube into a
cube with sampled model parameters.
Parameters
----------
cube : LP_c_double, np.ndarray
Unit cube.
bounds : dict(str, tuple(float, float))
Dictionary with the prior boundaries.
cube_index : dict(str, int)
Dictionary with the indices for selecting the model
parameters in the ``cube``.
Returns
-------
np.ndarray
Cube with the sampled model parameters.
"""

n_modelpar = len(self.modelpar)

if isinstance(cube, np.ndarray):
# Create new array for UltraNest
param_out = cube.copy()
else:
# Convert from ctypes.c_double to np.ndarray
# Only required with MultiNest
# cube = np.ctypeslib.as_array(cube, shape=(n_modelpar,))

# Use the same object with MultiNest
param_out = cube

for item in cube_index:
if item in self.prior:
# Gaussian prior
param_out[cube_index[item]] = stats.norm.ppf(
param_out[cube_index[item]],
loc=self.prior[item][0],
scale=self.prior[item][1],
)

else:
# Uniform prior
param_out[cube_index[item]] = (
bounds[item][0]
+ (bounds[item][1] - bounds[item][0]) * param_out[cube_index[item]]
)

return param_out

@typechecked
def _lnlike_func(
self,
params,
) -> Union[np.float64, float]:
"""
Function for calculating the log-likelihood for the sampled
parameter cube.
Parameters
----------
params : np.ndarray, pymultinest.run.LP_c_double
params : LP_c_double, np.ndarray
Cube with sampled model parameters.
prior : dict(str, tuple(float, float))
Dictionary with Gaussian priors for one or multiple
parameters. The prior can be set for any of the atmosphere
or calibration parameters, e.g.
``prior={'teff': (1200., 100.)}``. Additionally, a prior
can be set for the mass, e.g. ``prior={'mass': (13., 3.)}``
for an expected mass of 13 Mjup with an uncertainty of
3 Mjup.
Returns
-------
Expand Down Expand Up @@ -1245,7 +1293,9 @@ def lnlike_func(
if "flux_scaling" in self.cube_index:
flux_scaling = params[self.cube_index["flux_scaling"]]
else:
flux_scaling = 10.**params[self.cube_index["log_flux_scaling"]]
flux_scaling = (
10.0 ** params[self.cube_index["log_flux_scaling"]]
)

else:
try:
Expand Down Expand Up @@ -1291,7 +1341,7 @@ def lnlike_func(

ln_like = 0.0

for key, value in prior.items():
for key, value in self.prior.items():
if key == "mass":
if "logg" in self.modelpar:
mass = logg_to_mass(
Expand Down Expand Up @@ -1784,13 +1834,13 @@ def run_multinest(
``MultiNest`` sampler. While ``PyMultiNest`` can be
installed with ``pip`` from the PyPI repository,
``MultiNest`` has to to be build manually. See the
``PyMultiNest`` documentation for details:
http://johannesbuchner.github.io/PyMultiNest/install.html.
Note that the library path of ``MultiNest`` should be set
to the environmental variable ``LD_LIBRARY_PATH`` on a
`PyMultiNest documentation <http://johannesbuchner.
github.io/PyMultiNest/install.html>`_. The library
path of ``MultiNest`` should be set to the
environmental variable ``LD_LIBRARY_PATH`` on a
Linux machine and ``DYLD_LIBRARY_PATH`` on a Mac.
Alternatively, the variable can be set before importing
the ``species`` package, for example:
Alternatively, the variable can be set before
importing the ``species`` package, for example:
.. code-block:: python
Expand All @@ -1809,12 +1859,11 @@ def run_multinest(
prior : dict(str, tuple(float, float)), None
Dictionary with Gaussian priors for one or multiple
parameters. The prior can be set for any of the
atmosphere or calibration parameters, for example
``prior={'teff': (1200., 100.)}``. Additionally, a
prior can be set for the mass, for example
``prior={'mass': (13., 3.)}`` for an expected mass
of 13 Mjup with an uncertainty of 3 Mjup. The
parameter is not used if set to ``None``.
atmosphere or calibration parameters, e.g.
``prior={'teff': (1200., 100.)}``. Additionally, a prior
can be set for the mass, e.g. ``prior={'mass': (13., 3.)}``
for an expected mass of 13 Mjup with an uncertainty of
3 Mjup. The parameter is not used if set to ``None``.
Returns
-------
Expand All @@ -1824,6 +1873,13 @@ def run_multinest(

print("Running nested sampling with MultiNest...")

# Set attributes

if prior is None:
self.prior = {}
else:
self.prior = prior

# Get the MPI rank of the process

try:
Expand All @@ -1841,19 +1897,19 @@ def run_multinest(

# Add parallax to dictionary with Gaussian priors

if prior is None:
prior = {}
if "parallax" in self.modelpar:
self.prior["parallax"] = (self.parallax[0], self.parallax[1])

@typechecked
def lnprior_multinest(cube, n_dim: int, n_param: int) -> None:
def _lnprior_multinest(cube, n_dim: int, n_param: int) -> None:
"""
Function to transform the unit cube into the parameter
cube. It is not clear how to pass additional arguments
to the function, therefore it is placed here.
Parameters
----------
cube : pymultinest.run.LP_c_double
cube : LP_c_double
Unit cube.
n_dim : int
Number of dimensions.
Expand All @@ -1866,48 +1922,38 @@ def lnprior_multinest(cube, n_dim: int, n_param: int) -> None:
None
"""

for item in self.cube_index:
if item == "parallax":
# Gaussian prior for the parallax
cube[self.cube_index[item]] = stats.norm.ppf(
cube[self.cube_index[item]],
loc=self.parallax[0],
scale=self.parallax[1],
)

else:
# Uniform priors for all parameters
cube[self.cube_index[item]] = (
self.bounds[item][0]
+ (self.bounds[item][1] - self.bounds[item][0])
* cube[self.cube_index[item]]
)
self._prior_transform(cube, self.bounds, self.cube_index)

@typechecked
def lnlike_multinest(params, n_dim: int, n_param: int) -> Union[float, np.float64]:
def _lnlike_multinest(
params, n_dim: int, n_param: int
) -> Union[float, np.float64]:
"""
Function for return the log-likelihood for the sampled parameter cube.
Function for return the log-likelihood for the
sampled parameter cube.
Parameters
----------
params : pymultinest.run.LP_c_double
params : LP_c_double
Cube with sampled model parameters.
n_dim : int
Number of dimensions. This parameter is mandatory but not used by the function.
Number of dimensions. This parameter is mandatory
but not used by the function.
n_param : int
Number of parameters. This parameter is mandatory but not used by the function.
Number of parameters. This parameter is mandatory
but not used by the function.
Returns
-------
float
Log-likelihood.
"""

return self.lnlike_func(params, prior=prior)
return self._lnlike_func(params)

pymultinest.run(
lnlike_multinest,
lnprior_multinest,
_lnlike_multinest,
_lnprior_multinest,
len(self.modelpar),
outputfiles_basename=output,
resume=False,
Expand Down Expand Up @@ -2027,19 +2073,23 @@ def run_ultranest(
tag : str
Database tag where the samples will be stored.
min_num_live_points : int
Minimum number of live points. The default of 400 is a reasonable number (see
https://johannesbuchner.github.io/UltraNest/issues.html). In principle, choosing a very
low number allows nested sampling to make very few iterations and go to the peak
quickly. However, the space will be poorly sampled, giving a large region and thus low
efficiency, and potentially not seeing interesting modes. Therefore, a value above 100
is typically useful.
Minimum number of live points. The default of 400 is a
`reasonable number <https://johannesbuchner.github.io/
UltraNest/issues.html>`_. In principle, choosing a very
low number allows nested sampling to make very few
iterations and go to the peak quickly. However, the space
will be poorly sampled, giving a large region and thus low
efficiency, and potentially not seeing interesting modes.
Therefore, a value above 100 is typically useful.
output : str
Path that is used for the output files from ``UltraNest``.
prior : dict(str, tuple(float, float)), None
Dictionary with Gaussian priors for one or multiple parameters. The prior can be set
for any of the atmosphere or calibration parameters, e.g.
``prior={'teff': (1200., 100.)}``. Additionally, a prior can be set for the mass, e.g.
``prior={'mass': (13., 3.)}`` for an expected mass of 13 Mjup with an uncertainty of
Dictionary with Gaussian priors for one or multiple
parameters. The prior can be set for any of the
atmosphere or calibration parameters, e.g.
``prior={'teff': (1200., 100.)}``. Additionally, a prior
can be set for the mass, e.g. ``prior={'mass': (13., 3.)}``
for an expected mass of 13 Mjup with an uncertainty of
3 Mjup. The parameter is not used if set to ``None``.
Returns
Expand All @@ -2050,14 +2100,12 @@ def run_ultranest(

print("Running nested sampling with UltraNest...")

# Print warning if n_live_points from PyMultiNest is used
# Set attributes

# if n_live_points is not None:
# warnings.warn('The \'n_live_points\' parameter has been deprecated because UltraNest '
# 'is used instead of PyMultiNest. UltraNest can be executed with the '
# '\'run_ultranest\' method of \'FitModel\' and uses the '
# '\'min_num_live_points\' parameter (see documentation for details).',
# DeprecationWarning)
if prior is None:
self.prior = {}
else:
self.prior = prior

# Get the MPI rank of the process

Expand All @@ -2076,11 +2124,11 @@ def run_ultranest(

# Add parallax to dictionary with Gaussian priors

if prior is None:
prior = {}
if "parallax" in self.modelpar:
self.prior["parallax"] = (self.parallax[0], self.parallax[1])

@typechecked
def lnprior_ultranest(cube: np.ndarray) -> np.ndarray:
def _lnprior_ultranest(cube: np.ndarray) -> np.ndarray:
"""
Function to transform the unit cube into the parameter
cube. It is not clear how to pass additional arguments
Expand All @@ -2089,37 +2137,18 @@ def lnprior_ultranest(cube: np.ndarray) -> np.ndarray:
Parameters
----------
cube : np.ndarray
Array with unit parameters.
Unit cube.
Returns
-------
np.ndarray
Array with sampled model parameters.
Cube with the sampled model parameters.
"""

params = cube.copy()

for item in self.cube_index:
if item == "parallax":
# Gaussian prior for the parallax
params[self.cube_index[item]] = stats.norm.ppf(
cube[self.cube_index[item]],
loc=self.parallax[0],
scale=self.parallax[1],
)

else:
# Uniform priors for all parameters
params[self.cube_index[item]] = (
self.bounds[item][0]
+ (self.bounds[item][1] - self.bounds[item][0])
* params[self.cube_index[item]]
)

return params
return self._prior_transform(cube, self.bounds, self.cube_index)

@typechecked
def lnlike_ultranest(params: np.ndarray) -> Union[float, np.float64]:
def _lnlike_ultranest(params: np.ndarray) -> Union[float, np.float64]:
"""
Function for returning the log-likelihood for the sampled parameter cube.
Expand All @@ -2134,7 +2163,7 @@ def lnlike_ultranest(params: np.ndarray) -> Union[float, np.float64]:
Log-likelihood.
"""

ln_like = self.lnlike_func(params, prior=prior)
ln_like = self._lnlike_func(params)

if not np.isfinite(ln_like):
# UltraNest can not handle np.inf in the likelihood
Expand All @@ -2144,8 +2173,8 @@ def lnlike_ultranest(params: np.ndarray) -> Union[float, np.float64]:

sampler = ultranest.ReactiveNestedSampler(
self.modelpar,
lnlike_ultranest,
transform=lnprior_ultranest,
_lnlike_ultranest,
transform=_lnprior_ultranest,
resume="subfolder",
log_dir=output,
)
Expand Down
Loading

0 comments on commit b887861

Please sign in to comment.