API

regressor

class regressor.OptResult(res)

Bases: object

Class to make the output of a jaxopt optimisation appear like a scipy.optimize._optimize.OptimizeResult

Args:

res (ScipyMinimizeInfo):

The result of a jaxopt minimisation

class regressor.RoxyRegressor(fun, param_names, param_default, param_prior)

Bases: object

Regressor class which handles optimisation and MCMC for roxy. One can use this class to evaluate the function of interest and its derivative, optimise the parameters using and of the defined likelihoods and run an MCMC for these parameters.

Args:

fun (callable):

The function, f, to be considered by this regressor y = f(x, theta). The function must take two arguments, the first of which is the independent variable, the second of which are the parameters (as an array or list).

param_names (list):

The list of parameter names, in the order which they are supplied to fun

param_default (list):

The default valus of the parameters

param_prior (dict):

The prior range for each of the parameters. The prior is assumed to be uniform in this range. If either entry is None in the prior, then an infinite uniform prior is assumed.

compute_information_criterion(criterion, params_to_opt, xobs, yobs, errors, ngauss=1, infer_intrinsic=True, progress_bar=True, initial=None, nwarm=100, nsamp=100, method='mnr', gmm_prior='hierarchical', seed=1234, verbose=True, include_logdet=True)

Compute an information criterion for a given setup If an initial guess is not given, we first run a MCMC to get an initial guess for the maximum likelihood point, and we then an optimsier from this point to get a better estimate for this. Since we do not need good MCMC convergence for this, small values of nwarm and nsamp can be used.

Args:

criterion (str):

Which information criterion to use (supported: AIC and BIC)

params_to_opt (list):

The names of the parameters we wish to optimise

xobs (jnp.ndarray):

The observed x values

yobs (jnp.ndarray):

The observed y values

errors (jnp.ndarray):

[xerr, yerr], giving the error on the observed x and y values

ngauss (int, default = 1):

The number of Gaussians to use in the GMM prior. Only used if method=’gmm’

infer_intrinsic (bool, default=True):

Whether to infer the intrinsic scatter in the y direction

initial (jnp.ndarray, default=None):

The starting point for the optimised. If None, a MCMC is run.

progress_bar (bool, default=True):

Whether to display a progress bar for the MCMC

nwarm (int, default=100):

The number of warmup steps to use in the MCMC

nsamp (int, default=100):

The number of samples to obtain in the MCMC

method (str, default=’mnr’):

The name of the likelihood method to use (‘mnr’, ‘gmm’, ‘unif’ or ‘prof’). See roxy.likelihoods for more information.

gmm_prior (string, default=’hierarchical’):

If method=’gmm’, this decides what prior to put on the GMM componenents. If ‘uniform’, then the mean and widths have a uniform prior, and if ‘hierarchical’ mu and w^2 have a Normal and Inverse Gamma prior, respectively.

seed (int, default=1234):

The seed to use when initialising the sampler

verbose (bool, default=True):

Whether to print progress or not

include_logdet (bool, default=True):

For the method ‘prof’, whether to include the normalisation term in the likelihood proportional to log(det(S))

Returns:

negloglike (float):

The optimum negative log-likelihood value

metric (float):

The value of the information criterion

find_best_gmm(params_to_opt, xobs, yobs, xerr, yerr, max_ngauss, best_metric='BIC', infer_intrinsic=True, progress_bar=True, nwarm=100, nsamp=100, gmm_prior='hierarchical', seed=1234, verbose=True, include_logdet=True)

Find the number of Gaussians to use in a Gaussian Mixture Model hyper-prior on the true x values, accoridng to some metric.

Args:

params_to_opt (list):

The names of the parameters we wish to optimise

xobs (jnp.ndarray):

The observed x values

yobs (jnp.ndarray):

The observed y values

xerr (jnp.ndarray):

The error on the observed x values

yerr (jnp.ndarray):

The error on the observed y values

max_ngauss (int):

The maximum number of Gaussians to consider

best_metric (str):

Metric to use to compare fits (supported: AIC and BIC)

infer_intrinsic (bool, default=True):

Whether to infer the intrinsic scatter in the y direction

progress_bar (bool, default=True):

Whether to display a progress bar for the MCMC

nwarm (int, default=100):

The number of warmup steps to use in the MCMC

nsamp (int, default=100):

The number of samples to obtain in the MCMC

gmm_prior (string, default=’hierarchical’):

If method=’gmm’, this decides what prior to put on the GMM componenents. If ‘uniform’, then the mean and widths have a uniform prior, and if ‘hierarchical’ mu and w^2 have a Normal and Inverse Gamma prior, respectively.

seed (int, default=1234):

The seed to use when initialising the sampler

verbose (bool, default=True):

Whether to print progress or not

include_logdet (bool, default=True):

For the method ‘prof’, whether to include the normalisation term in the likelihood proportional to log(det(S))

Returns:

ngauss (int):

The best number of Gaussians to use according to the metric

get_param_index(params_to_opt, verbose=True)

If the function of interest if f(x, theta), find the index in theta for each of the parameters we wish to optimise

Args:

params_to_opt (list):

The names of the parameters we wish to optimise

verbose (bool, default=True):

Whether to print the names and values of parameters which are not fitted

Returns:

pidx (jnp.ndarray):

The indices of the parameters to optimise

gradient(x, theta)

If we are fitting the function f(x, theta), this is df/dx evaluated at (x, theta)

Args:

x (jnp.ndarray):

The x values

theta (jnp.ndarray):

The parameter values

Returns:

jnp.ndarray:

df/dx evaluated at (x, theta)

mcmc(params_to_opt, xobs, yobs, errors, nwarm, nsamp, y_is_detected=[], method='mnr', ngauss=1, infer_intrinsic=True, num_chains=1, progress_bar=True, covmat=False, gmm_prior='hierarchical', seed=1234, verbose=True, init=None, include_logdet=True, optimiser='l-bfgs-b')

Run an MCMC using the NUTS sampler of numpyro for the parameters of the function given some data, under the assumption of an uncorrelated Gaussian likelihood, using the likelihood specififed by ‘method’.

Args:

params_to_opt (list):

The names of the parameters we wish to optimise

xobs (jnp.ndarray):

The observed x values

yobs (jnp.ndarray):

The observed y values

errors (jnp.ndarray):

If covmat=False, then this is [xerr, yerr], giving the error on the observed x and y values. Otherwise, this is the covariance matrix in the order (x, y)

nwarm (int):

The number of warmup steps to use in the MCMC

nsamp (int):

The number of samples to obtain in the MCMC

y_is_detected (jnp.ndarray):

A boolean array of the same length as xobs and yobs, giving whether each point is a detection (True) or an upper limit (False).

method (str, default=’mnr’):

The name of the likelihood method to use (‘mnr’, ‘gmm’, ‘unif’ or ‘prof’). See roxy.likelihoods for more information.

ngauss (int, default = 1):

The number of Gaussians to use in the GMM prior. Only used if method=’gmm’

infer_intrinsic (bool, default=True):

Whether to infer the intrinsic scatter in the y direction

num_chains (int, default=1):

The number of independent MCMC chains to run

progress_bar (bool, default=True):

Whether to display a progress bar for the MCMC

covmat (bool, default=False):

This determines whether the errors argument is [xerr, yerr] (False) or a covariance matrix (True).

gmm_prior (string, default=’hierarchical’):

If method=’gmm’, this decides what prior to put on the GMM componenents. If ‘uniform’, then the mean and widths have a uniform prior, and if ‘hierarchical’ mu and w^2 have a Normal and Inverse Gamma prior, respectively.

seed (int, default=1234):

The seed to use when initialising the sampler

verbose (bool, default=True):

Whether to print progress or not

init (dict, default=None):

A dictionary of values of initialise the MCMC at

include_logdet (bool, default=True):

For the method ‘prof’, whether to include the normalisation term in the likelihood proportional to log(det(S))

optimiser (str, default=’l-bfgs-b’):

The optimiser to use if a initial point is not specified. This must be a method supported by jaxopt.ScipyBoundedMinimize.

Returns:

samples (dict):

The MCMC samples, where the keys are the parameter names and values are ndarrays of the samples

mcmc2opt_index(labels, ngauss=1, method='mnr', gmm_prior='hierarchical', infer_intrinsic=True)

Find the indices which convert the samples produced by the MCMC to the order required for the optimiser

Args:

labels (list):

List of label names in the order produced by the MCMC

ngauss (int, default = 1):

The number of Gaussians to use in the GMM prior. Only used if method=’gmm’

method (str, default=’mnr’):

The name of the likelihood method to use (‘mnr’, ‘gmm’, ‘unif’ or ‘prof’). See roxy.likelihoods for more information.

gmm_prior (string, default=’hierarchical’):

If method=’gmm’, this decides what prior to put on the GMM componenents. If ‘uniform’, then the mean and widths have a uniform prior, and if ‘hierarchical’ mu and w^2 have a Normal and Inverse Gamma prior, respectively.

infer_intrinsic (bool, default=True):

Whether to infer the intrinsic scatter in the y direction

Returns:

param_idx (list):

The indices which convert the order of parameters from the MCMC to that expected by the optimiser

param_names (list):

The names of the parameters in the order expected by the optimiser

negloglike(theta, xobs, yobs, errors, y_is_detected=[], sig=0.0, mu_gauss=0.0, w_gauss=1.0, weights_gauss=1.0, method='mnr', covmat=False, test_prior=True, include_logdet=True)

Computes the negative log-likelihood under the assumption of an uncorrelated (correlated) Gaussian likelihood if covmat is False (True), using the likelihood specififed by ‘method’.

Args:

theta (jnp.ndarray):

The parameters of the function to use

xobs (jnp.ndarray):

The observed x values

yobs (jnp.ndarray):

The observed y values

errors (jnp.ndarray):

If covmat=False, then this is [xerr, yerr], giving the error on the observed x and y values. Otherwise, this is the covariance matrix in the order (x, y)

y_is_detected (jnp.ndarray):

A boolean array of the same length as xobs and yobs, giving whether each point is a detection (True) or an upper limit (False)

sig (float, default=0.):

The intrinsic scatter, which is added in quadrature with yerr

mu_gauss (float or jnp.ndarray, default=0.):

The mean of the Gaussian prior on the true x positions (only used if method=’mnr’ or ‘gmm’). If using ‘mnr’ and this is an array, only the first mean is used.

w_gauss (float or jnp.ndarray, default=1.):

The standard deviation of the Gaussian prior on the true x positions (only used if method=’mnr’).

weights_gauss (float or jnp.ndarray, default=1.):

The weights of the Gaussians in a GMM prior on the true x positions (only used if method=’gmm’).

method (str, default=’mnr’):

The name of the likelihood method to use (‘mnr’, ‘gmm’, ‘unif’ or ‘prof’). See roxy.likelihoods for more information

covmat (bool, default=False):

This determines whether the errors argument is [xerr, yerr] (False) or a covariance matrix (True).

test_prior (bool, default=True):

Whether to test sigma >= 0 and Gaussians weights >= 0

include_logdet (bool, default=True):

For the method ‘prof’, whether to include the normalisation term in the likelihood proportional to log(det(S))

optimise(params_to_opt, xobs, yobs, errors, y_is_detected=[], method='mnr', infer_intrinsic=True, initial=None, ngauss=1, covmat=False, gmm_prior='hierarchical', include_logdet=True, verbose=True, optimiser='l-bfgs-b', niter=10, nconv=3, tol=0.001)

Optimise the parameters of the function given some data, using multiple random restarts to check for convergence.

Args:

params_to_opt (list):

The names of the parameters we wish to optimise

xobs (jnp.ndarray):

The observed x values

yobs (jnp.ndarray):

The observed y values

errors (jnp.ndarray):

If covmat=False, then this is [xerr, yerr], giving the error on the observed x and y values. Otherwise, this is the covariance matrix in the order (x, y)

y_is_detected (jnp.ndarray):

A boolean array of the same length as xobs and yobs, giving whether each point is a detection (True) or an upper limit (False)

method (str, default=’mnr’):

The name of the likelihood method to use (‘mnr’, ‘gmm’, ‘unif’ or ‘prof’). See roxy.likelihoods for more information

infer_intrinsic (bool, default=True):

Whether to infer the intrinsic scatter in the y direction

initial (jnp.ndarray, default=None):

The starting point for the first optimisation. If None, a random value in the prior range is chosen. For subsequent iterations, None is used.

ngauss (int, default = 1):

The number of Gaussians to use in the GMM prior. Only used if method=’gmm’

covmat (bool, default=False):

This determines whether the errors argument is [xerr, yerr] (False) or a covariance matrix (True).

gmm_prior (string, default=’hierarchical’):

If method=’gmm’, this decides what prior to put on the GMM componenents. If ‘uniform’, then the mean and widths have a uniform prior, and if ‘hierarchical’ mu and w^2 have a Normal and Inverse Gamma prior, respectively.

include_logdet (bool, default=True):

For the method ‘prof’, whether to include the normalisation term in the likelihood proportional to log(det(S))

verbose (bool, default=True):

Whether to print progress or not

optimiser (str, default=’l-bfgs-b’):

The optimiser to use. This must be a method supported by jaxopt.ScipyBoundedMinimize.

niter (int, default=10):

The number of iterations (random restarts) to perform

nconv (int, default=3):

The number of times the best likelihood must be achieved (within tolerance) to declare convergence

tol (float, default=1e-3):

The tolerance for determining if two likelihoods are the same

Returns:

res (OptResult):

The result of the best optimisation

param_names (list):

List of parameter names in order of res.params

second_derivative(x, theta)

If we are fitting the function f(x, theta), this is d^2f/dx^2 evaluated at (x, theta)

Args:

x (jnp.ndarray):

The x values

theta (jnp.ndarray):

The parameter values

Returns:

jnp.ndarray:

d^2f/dx^2 evaluated at (x, theta)

value(x, theta)

If we are fitting the function f(x, theta), this is f(x, theta) evaluated at (x, theta)

Args:

x (jnp.ndarray):

The x values

theta (jnp.ndarray):

The parameter values

Returns:

jnp.ndarray:

f(x, theta) evaluated at (x, theta)

likelihoods

likelihoods.check_valid_covmat(D, tol=1e-08)

Check if a covariance matrix is valid (symmetric and positive semi-definite) - Symmetry (within tolerance) - Positive semi-definiteness (eigvals >= -tol)

Args:

D (jnp.ndarray):

The covariance matrix to check

tol (float, default=1e-8):

The tolerance for numerical checks

Returns:

is_valid (bool):

Whether the covariance matrix is valid

likelihoods.likelihood_warnings(method, infer_intrinsic, nx, errors, covmat)

Raise warnings if using asymptotically biased likelihood for the situation of interest.

Args:

method (str, default=’mnr’):

The name of the likelihood method to use (‘mnr’, ‘gmm’, ‘unif’ or ‘prof’).

infer_intrinsic (bool, default=True):

Whether to infer the intrinsic scatter in the y direction

nx (int):

The number of observed x values

errors (jnp.ndarray):

If covmat=False, then this is [xerr, yerr], giving the error on the observed x and y values. Otherwise, this is the covariance matrix in the order (x, y)

covmat (bool, default=False):

This determines whether the errors argument is [xerr, yerr] (False) or a covariance matrix (True).

likelihoods.negloglike_gmm(xobs, yobs, xerr, yerr, f, fprime, sig, all_mu_gauss, all_w_gauss, all_weights)

Computes the negative log-likelihood under the assumption of an uncorrelated Gaussian likelihood with a GMM prior on the true x positions.

Args:

xobs (jnp.ndarray):

The observed x values

yobs (jnp.ndarray):

The observed y values

xerr (jnp.ndarray):

The error on the observed x values

yerr (jnp.ndarray):

The error on the observed y values

f (jnp.ndarray):

If we are fitting the function f(x), this is f(x) evaluated at xobs

fprime (jnp.ndarray):

If we are fitting the function f(x), this is df/dx evaluated at xobs

sig (float):

The intrinsic scatter, which is added in quadrature with yerr

all_mu_gauss (jnp.ndarray):

The means of the Gaussians in the GMM prior on the true x positions

all_w_gauss (jnp.ndarray):

The standard deviations of the Gaussians in the GMM prior on the true x positions

all_weights (jnp.ndarray):

The weights of the Gaussians in the GMM prior on the true x positions

Returns:

neglog_p (float):

The negative log-likelihood

likelihoods.negloglike_mnr(xobs, yobs, xerr, yerr, f, fprime, sig, mu_gauss, w_gauss)

Computes the negative log-likelihood under the assumption of an uncorrelated Gaussian likelihood with a Gaussian prior on the true x positions.

Args:

xobs (jnp.ndarray):

The observed x values

yobs (jnp.ndarray):

The observed y values

xerr (jnp.ndarray):

The error on the observed x values

yerr (jnp.ndarray):

The error on the observed y values

f (jnp.ndarray):

If we are fitting the function f(x), this is f(x) evaluated at xobs

fprime (jnp.ndarray):

If we are fitting the function f(x), this is df/dx evaluated at xobs

sig (float):

The intrinsic scatter, which is added in quadrature with yerr

mu_gauss (float):

The mean of the Gaussian prior on the true x positions

w_gauss (float):

The standard deviation of the Gaussian prior on the true x positions

Returns:

neglog_p (float):

The negative log-likelihood

likelihoods.negloglike_mnr_mv(xobs, yobs, Sigma, f, G, sig, mu_gauss, w_gauss)

Computes the negative log-likelihood under the assumption of a correlated Gaussian likelihood (i.e. arbitrary covariance matrix) with a Gaussian prior on the true x positions.

Args:

xobs (jnp.ndarray):

The observed x values

yobs (jnp.ndarray):

The observed y values

Sigma (jnp.ndarray):

The covariance matrix giving the errors on the observed (x, y) values

f (jnp.ndarray):

If we are fitting the function f(x), this is f(x) evaluated at xobs

G (jnp.ndarray):

If we are fitting the function f(x), this is G_{ij} = df_i/dx_j evaluated at xobs

sig (float):

The intrinsic scatter, which is added in quadrature with yerr

mu_gauss (float):

The mean of the Gaussian prior on the true x positions

w_gauss (float):

The standard deviation of the Gaussian prior on the true x positions

Returns:

neglog_p (float):

The negative log-likelihood

likelihoods.negloglike_mnr_uplims(xobs, yobs, y_is_detected, xerr, yerr, f, fprime, sig, mu_gauss, w_gauss)

Computes the negative log-likelihood under the assumption of an uncorrelated Gaussian likelihood with a Gaussian prior on the true x positions,

where some of the y values are upper limits.

Args:

xobs (jnp.ndarray):

The observed x values

yobs (jnp.ndarray):

The observed y values

y_is_detected (jnp.ndarray):

A boolean array of the same length as xobs and yobs, giving whether each point is a detection (True) or an upper limit (False)

xerr (jnp.ndarray):

The error on the observed x values

yerr (jnp.ndarray):

The error on the observed y values

f (jnp.ndarray):

If we are fitting the function f(x), this is f(x) evaluated at xobs

fprime (jnp.ndarray):

If we are fitting the function f(x), this is df/dx evaluated at xobs

sig (float):

The intrinsic scatter, which is added in quadrature with yerr

mu_gauss (float):

The mean of the Gaussian prior on the true x positions

w_gauss (float):

The standard deviation of the Gaussian prior on the true x positions

likelihoods.negloglike_prof(xobs, yobs, xerr, yerr, f, fprime, sig, include_logdet=True)

Computes the negative log-likelihood under the assumption of an uncorrelated Gaussian likelihood, evaluated at the maximum likelihood values of xtrue (the profile likelihood)

Args:

xobs (jnp.ndarray):

The observed x values

yobs (jnp.ndarray):

The observed y values

xerr (jnp.ndarray):

The error on the observed x values

yerr (jnp.ndarray):

The error on the observed y values

f (jnp.ndarray):

If we are fitting the function f(x), this is f(x) evaluated at xobs

fprime (jnp.ndarray):

If we are fitting the function f(x), this is df/dx evaluated at xobs

sig (float):

The intrinsic scatter, which is added in quadrature with yerr

include_logdet (bool, default=True):

Whether to include the normalisation term in the likelihood proportional to log(det(S))

Returns:

neglog_p (float):

The negative log-likelihood

likelihoods.negloglike_prof_mv(xobs, yobs, Sigma, f, G, sig, include_logdet=True)

Computes the negative log-likelihood under the assumption of a correlated Gaussian likelihood (i.e. arbitrary covariance matrix), evaluated at the maximum likelihood values of xtrue (the profile likelihood)

Args:

xobs (jnp.ndarray):

The observed x values

yobs (jnp.ndarray):

The observed y values

Sigma (jnp.ndarray):

The covariance matrix giving the errors on the observed (x, y) values

f (jnp.ndarray):

If we are fitting the function f(x), this is f(x) evaluated at xobs

G (jnp.ndarray):

If we are fitting the function f(x), this is G_{ij} = df_i/dx_j evaluated at xobs

sig (float):

The intrinsic scatter, which is added in quadrature with yerr

include_logdet (bool, default=True):

Whether to include the normalisation term in the likelihood proportional to log(det(S))

Returns:

neglog_p (float):

The negative log-likelihood

likelihoods.negloglike_unif(xobs, yobs, xerr, yerr, f, fprime, sig)

Computes the negative log-likelihood under the assumption of an uncorrelated Gaussian likelihood, where we have marginalised over the true x values, assuming an infinite uniform prior on these.

Args:

xobs (jnp.ndarray):

The observed x values

yobs (jnp.ndarray):

The observed y values

xerr (jnp.ndarray):

The error on the observed x values

yerr (jnp.ndarray):

The error on the observed y values

f (jnp.ndarray):

If we are fitting the function f(x), this is f(x) evaluated at xobs

fprime (jnp.ndarray):

If we are fitting the function f(x), this is df/dx evaluated at xobs

sig (float):

The intrinsic scatter, which is added in quadrature with yerr

Returns:

neglog_p (float):

The negative log-likelihood

likelihoods.negloglike_unif_mv(xobs, yobs, Sigma, f, G, sig)

Computes the negative log-likelihood under the assumption of a correlated Gaussian likelihood (i.e. arbitrary covariance matrix), where we have marginalised over the true x values, assuming an infinite uniform prior on these.

Args:

xobs (jnp.ndarray):

The observed x values

yobs (jnp.ndarray):

The observed y values

Sigma (jnp.ndarray):

The covariance matrix giving the errors on the observed (x, y) values

f (jnp.ndarray):

If we are fitting the function f(x), this is f(x) evaluated at xobs

G (jnp.ndarray):

If we are fitting the function f(x), this is G_{ij} = df_i/dx_j evaluated at xobs

sig (float):

The intrinsic scatter, which is added in quadrature with yerr

Returns:

neglog_p (float):

The negative log-likelihood

mcmc

class mcmc.Likelihood_GMM(xobs, yobs, xerr, yerr, f, fprime, sig, all_mu_gauss, all_w_gauss, all_weights)

Bases: Distribution

Class to be used by numpyro to evaluate the log-likelihood under the assumption of an uncorrelated Gaussian likelihood with a Gaussian Mixture Model prior on the true x positions.

Args:

xobs (jnp.ndarray):

The observed x values

yobs (jnp.ndarray):

The observed y values

xerr (jnp.ndarray):

The error on the observed x values

yerr (jnp.ndarray):

The error on the observed y values

f (jnp.ndarray):

If we are fitting the function f(x), this is f(x) evaluated at xobs

fprime (jnp.ndarray):

If we are fitting the function f(x), this is df/dx evaluated at xobs

sig (float):

The intrinsic scatter, which is added in quadrature with yerr

all_mu_gauss (jnp.ndarray):

The mean of the Gaussians in the GMM prior on the true x positions

all_w_gauss (jnp.ndarray):

The standard deviation of the Gaussians in the GMM prior on the true x positions

all_weights (jnp.ndarray):

The weights of the Gaussians in the GMM prior on the true x positions

log_prob(value)

Evaluates the log probability density for a batch of samples given by value.

Parameters:

value – A batch of samples from the distribution.

Returns:

an array with shape value.shape[:-self.event_shape]

Return type:

ArrayLike

sample(key, sample_shape=())

Returns a sample from the distribution having shape given by sample_shape + batch_shape + event_shape. Note that when sample_shape is non-empty, leading dimensions (of size sample_shape) of the returned sample will be filled with iid draws from the distribution instance.

Parameters:

• key (jax.random.PRNGKey) – the rng_key key to be used for the distribution.

• sample_shape (tuple) – the sample shape for the distribution.

Returns:

an array of shape sample_shape + batch_shape + event_shape

Return type:

numpy.ndarray

class mcmc.Likelihood_MNR(xobs, yobs, xerr, yerr, f, fprime, sig, mu_gauss, w_gauss)

Bases: Distribution

Class to be used by numpyro to evaluate the log-likelihood under the assumption of an uncorrelated Gaussian likelihood with a Gaussian prior on the true x positions.

Args:

xobs (jnp.ndarray):

The observed x values

yobs (jnp.ndarray):

The observed y values

xerr (jnp.ndarray):

The error on the observed x values

yerr (jnp.ndarray):

The error on the observed y values

f (jnp.ndarray):

If we are fitting the function f(x), this is f(x) evaluated at xobs

fprime (jnp.ndarray):

If we are fitting the function f(x), this is df/dx evaluated at xobs

sig (float):

The intrinsic scatter, which is added in quadrature with yerr

mu_gauss (float):

The mean of the Gaussian prior on the true x positions

w_gauss (float):

The standard deviation of the Gaussian prior on the true x positions

log_prob(value)

Evaluates the log probability density for a batch of samples given by value.

Parameters:

value – A batch of samples from the distribution.

Returns:

an array with shape value.shape[:-self.event_shape]

Return type:

ArrayLike

sample(key, sample_shape=())

Returns a sample from the distribution having shape given by sample_shape + batch_shape + event_shape. Note that when sample_shape is non-empty, leading dimensions (of size sample_shape) of the returned sample will be filled with iid draws from the distribution instance.

Parameters:

• key (jax.random.PRNGKey) – the rng_key key to be used for the distribution.

• sample_shape (tuple) – the sample shape for the distribution.

Returns:

an array of shape sample_shape + batch_shape + event_shape

Return type:

numpy.ndarray

class mcmc.Likelihood_MNR_MV(xobs, yobs, Sxx, Syy, Sxy, f, G, sig, mu_gauss, w_gauss)

Bases: Distribution

Class to be used by numpyro to evaluate the log-likelihood under the assumption of a correlated Gaussian likelihood (i.e. arbitrary covariance matrix), with a Gaussian prior on the true x positions.

Args:

xobs (jnp.ndarray):

The observed x values

yobs (jnp.ndarray):

The observed y values

Sxx (jnp.ndarray):

The xx component of the covariance matrix giving the errors on the observed (x, y) values

Syy (jnp.ndarray):

The yy component of the covariance matrix giving the errors on the observed (x, y) values

Sxy (jnp.ndarray):

The xy component of the covariance matrix giving the errors on the observed (x, y) values

f (jnp.ndarray):

If we are fitting the function f(x), this is f(x) evaluated at xobs

G (jnp.ndarray):

If we are fitting the function f(x), this is G_{ij} = df_i/dx_j evaluated at xobs

sig (float):

The intrinsic scatter, which is added in quadrature with yerr

mu_gauss (float):

The mean of the Gaussian prior on the true x positions

w_gauss (float):

The standard deviation of the Gaussian prior on the true x positions

log_prob(value)

Evaluates the log probability density for a batch of samples given by value.

Parameters:

value – A batch of samples from the distribution.

Returns:

an array with shape value.shape[:-self.event_shape]

Return type:

ArrayLike

sample(key, sample_shape=())

Returns a sample from the distribution having shape given by sample_shape + batch_shape + event_shape. Note that when sample_shape is non-empty, leading dimensions (of size sample_shape) of the returned sample will be filled with iid draws from the distribution instance.

Parameters:

• key (jax.random.PRNGKey) – the rng_key key to be used for the distribution.

• sample_shape (tuple) – the sample shape for the distribution.

Returns:

an array of shape sample_shape + batch_shape + event_shape

Return type:

numpy.ndarray

class mcmc.Likelihood_MNR_uplims(xobs, yobs, y_is_detected, xerr, yerr, f, fprime, sig, mu_gauss, w_gauss)

Bases: Distribution

Class to be used by numpyro to evaluate the log-likelihood under the assumption of an uncorrelated Gaussian likelihood with a Gaussian prior on the true x positions, where some of the points are upper limits (i.e. non-detections in y).

Args:

xobs (jnp.ndarray):

The observed x values

yobs (jnp.ndarray):

The observed y values

y_is_detected (jnp.ndarray):

A boolean array of the same length as xobs and yobs, giving whether each point is a detection (True) or an upper limit (False)

xerr (jnp.ndarray):

The error on the observed x values

yerr (jnp.ndarray):

The error on the observed y values

f (jnp.ndarray):

If we are fitting the function f(x), this is f(x) evaluated at xobs

fprime (jnp.ndarray):

If we are fitting the function f(x), this is df/dx evaluated at xobs

sig (float):

The intrinsic scatter, which is added in quadrature with yerr

mu_gauss (float):

The mean of the Gaussian prior on the true x positions

w_gauss (float):

The standard deviation of the Gaussian prior on the true x positions

log_prob(value)

Evaluates the log probability density for a batch of samples given by value.

Parameters:

value – A batch of samples from the distribution.

Returns:

an array with shape value.shape[:-self.event_shape]

Return type:

ArrayLike

sample(key, sample_shape=())

Returns a sample from the distribution having shape given by sample_shape + batch_shape + event_shape. Note that when sample_shape is non-empty, leading dimensions (of size sample_shape) of the returned sample will be filled with iid draws from the distribution instance.

Parameters:

• key (jax.random.PRNGKey) – the rng_key key to be used for the distribution.

• sample_shape (tuple) – the sample shape for the distribution.

Returns:

an array of shape sample_shape + batch_shape + event_shape

Return type:

numpy.ndarray

class mcmc.Likelihood_prof(xobs, yobs, xerr, yerr, f, fprime, sig, include_logdet=True)

Bases: Distribution

Class to be used by numpyro to evaluate the log-likelihood under the assumption of an uncorrelated Gaussian likelihood, evaluated at the maximum likelihood values of xtrue (the profile likelihood)

Args:

xobs (jnp.ndarray):

The observed x values

yobs (jnp.ndarray):

The observed y values

xerr (jnp.ndarray):

The error on the observed x values

yerr (jnp.ndarray):

The error on the observed y values

f (jnp.ndarray):

If we are fitting the function f(x), this is f(x) evaluated at xobs

fprime (jnp.ndarray):

If we are fitting the function f(x), this is df/dx evaluated at xobs

sig (float):

The intrinsic scatter, which is added in quadrature with yerr

include_logdet (bool, default=True):

Whether to include the normalisation term in the likelihood proportional to log(det(S))

log_prob(value)

Evaluates the log probability density for a batch of samples given by value.

Parameters:

value – A batch of samples from the distribution.

Returns:

an array with shape value.shape[:-self.event_shape]

Return type:

ArrayLike

sample(key, sample_shape=())

Returns a sample from the distribution having shape given by sample_shape + batch_shape + event_shape. Note that when sample_shape is non-empty, leading dimensions (of size sample_shape) of the returned sample will be filled with iid draws from the distribution instance.

Parameters:

• key (jax.random.PRNGKey) – the rng_key key to be used for the distribution.

• sample_shape (tuple) – the sample shape for the distribution.

Returns:

an array of shape sample_shape + batch_shape + event_shape

Return type:

numpy.ndarray

class mcmc.Likelihood_prof_MV(xobs, yobs, Sxx, Syy, Sxy, f, G, sig, include_logdet=True)

Bases: Distribution

Class to be used by numpyro to evaluate the log-likelihood under the assumption of a correlated Gaussian likelihood (i.e. arbitrary covariance matrix), evaluated at the maximum likelihood values of xtrue (the profile likelihood)

Args:

xobs (jnp.ndarray):

The observed x values

yobs (jnp.ndarray):

The observed y values

Sxx (jnp.ndarray):

The xx component of the covariance matrix giving the errors on the observed (x, y) values

Syy (jnp.ndarray):

The yy component of the covariance matrix giving the errors on the observed (x, y) values

Sxy (jnp.ndarray):

The xy component of the covariance matrix giving the errors on the observed (x, y) values

f (jnp.ndarray):

If we are fitting the function f(x), this is f(x) evaluated at xobs

G (jnp.ndarray):

If we are fitting the function f(x), this is G_{ij} = df_i/dx_j evaluated at xobs

sig (float):

The intrinsic scatter, which is added in quadrature with yerr

include_logdet (bool, default=True):

Whether to include the normalisation term in the likelihood proportional to log(det(S))

log_prob(value)

Evaluates the log probability density for a batch of samples given by value.

Parameters:

value – A batch of samples from the distribution.

Returns:

an array with shape value.shape[:-self.event_shape]

Return type:

ArrayLike

sample(key, sample_shape=())

Returns a sample from the distribution having shape given by sample_shape + batch_shape + event_shape. Note that when sample_shape is non-empty, leading dimensions (of size sample_shape) of the returned sample will be filled with iid draws from the distribution instance.

Parameters:

• key (jax.random.PRNGKey) – the rng_key key to be used for the distribution.

• sample_shape (tuple) – the sample shape for the distribution.

Returns:

an array of shape sample_shape + batch_shape + event_shape

Return type:

numpy.ndarray

class mcmc.Likelihood_unif(xobs, yobs, xerr, yerr, f, fprime, sig)

Bases: Distribution

Class to be used by numpyro to evaluate the log-likelihood under the assumption of an uncorrelated Gaussian likelihood, where we have marginalised over the true x values, assuming an infinite uniform prior on these.

Args:

xobs (jnp.ndarray):

The observed x values

yobs (jnp.ndarray):

The observed y values

xerr (jnp.ndarray):

The error on the observed x values

yerr (jnp.ndarray):

The error on the observed y values

f (jnp.ndarray):

If we are fitting the funciton f(x), this is f(x) evaluated at xobs

fprime (jnp.ndarray):

If we are fitting the funciton f(x), this is df/dx evaluated at xobs

sig (float):

The intrinsic scatter, which is added in quadrature with yerr

log_prob(value)

Evaluates the log probability density for a batch of samples given by value.

Parameters:

value – A batch of samples from the distribution.

Returns:

an array with shape value.shape[:-self.event_shape]

Return type:

ArrayLike

sample(key, sample_shape=())

Returns a sample from the distribution having shape given by sample_shape + batch_shape + event_shape. Note that when sample_shape is non-empty, leading dimensions (of size sample_shape) of the returned sample will be filled with iid draws from the distribution instance.

Parameters:

• key (jax.random.PRNGKey) – the rng_key key to be used for the distribution.

• sample_shape (tuple) – the sample shape for the distribution.

Returns:

an array of shape sample_shape + batch_shape + event_shape

Return type:

numpy.ndarray

class mcmc.Likelihood_unif_MV(xobs, yobs, Sxx, Syy, Sxy, f, G, sig)

Bases: Distribution

Class to be used by numpyro to evaluate the log-likelihood under the assumption of a correlated Gaussian likelihood (i.e. arbitrary covariance matrix), where we have marginalised over the true x values, assuming an infinite uniform prior on these.

Args:

xobs (jnp.ndarray):

The observed x values

yobs (jnp.ndarray):

The observed y values

Sxx (jnp.ndarray):

The xx component of the covariance matrix giving the errors on the observed (x, y) values

Syy (jnp.ndarray):

The yy component of the covariance matrix giving the errors on the observed (x, y) values

Sxy (jnp.ndarray):

The xy component of the covariance matrix giving the errors on the observed (x, y) values

f (jnp.ndarray):

If we are fitting the function f(x), this is f(x) evaluated at xobs

G (jnp.ndarray):

If we are fitting the function f(x), this is G_{ij} = df_i/dx_j evaluated at xobs

sig (float):

The intrinsic scatter, which is added in quadrature with yerr

log_prob(value)

Evaluates the log probability density for a batch of samples given by value.

Parameters:

value – A batch of samples from the distribution.

Returns:

an array with shape value.shape[:-self.event_shape]

Return type:

ArrayLike

sample(key, sample_shape=())

Returns a sample from the distribution having shape given by sample_shape + batch_shape + event_shape. Note that when sample_shape is non-empty, leading dimensions (of size sample_shape) of the returned sample will be filled with iid draws from the distribution instance.

Parameters:

• key (jax.random.PRNGKey) – the rng_key key to be used for the distribution.

• sample_shape (tuple) – the sample shape for the distribution.

Returns:

an array of shape sample_shape + batch_shape + event_shape

Return type:

numpy.ndarray

class mcmc.OrderedNormal(loc=0.0, scale=1.0, *, validate_args=None)

Bases: Distribution

arg_constraints: dict[str, Any] = {'loc': Real(), 'scale': Positive(lower_bound=0.0)}

cdf(value)

The cumulative distribution function of this distribution.

Parameters:

value – samples from this distribution.

Returns:

output of the cumulative distribution function evaluated at value.

icdf(q)

The inverse cumulative distribution function of this distribution.

Parameters:

q – quantile values, should belong to [0, 1].

Returns:

the samples whose cdf values equals to q.

log_cdf(value)

log_prob(value)

Evaluates the log probability density for a batch of samples given by value.

Parameters:

value – A batch of samples from the distribution.

Returns:

an array with shape value.shape[:-self.event_shape]

Return type:

ArrayLike

property mean

Mean of the distribution.

reparametrized_params: list[str] = ['loc', 'scale']

sample(key, sample_shape=())

Returns a sample from the distribution having shape given by sample_shape + batch_shape + event_shape. Note that when sample_shape is non-empty, leading dimensions (of size sample_shape) of the returned sample will be filled with iid draws from the distribution instance.

Parameters:

• key (jax.random.PRNGKey) – the rng_key key to be used for the distribution.

• sample_shape (tuple) – the sample shape for the distribution.

Returns:

an array of shape sample_shape + batch_shape + event_shape

Return type:

numpy.ndarray

support = OrderedVector()

property variance

Variance of the distribution.

mcmc.compute_bias(samples, truths, verbose=True)

Computes the bias between MCMC samples and the true values of the parameters. For the intrinsic scatter, a truncated normal distribution is first fitted to the parameters since this parameter can only take positive values.

Args:

samples (dict):

The MCMC samples, where the keys are the parameter names and values are ndarrays of the samples

truths (dict):

The true values of the parameters. The keys should be a subset of the keys of samples

verbose (bool, default=True):

Whether to print biases or not

Reurns:

: biases (dict): The biases in each of the parameters (units=number of sigmas)

mcmc.samples_to_array(samples)

Converts a dictionary of samples returned by numpro to a list of names and an array of samples.

Args:

samples (dict):

The MCMC samples, where the keys are the parameter names and values are ndarrays of the samples

Returns:

labels (np.ndarray):

The names of the sampled variables

all_samples (np.ndarray):

The sampled values for these variables. Shape = (number of samples, number of parameters).

plotting

plotting.posterior_predictive_plot(reg, samples, xobs, yobs, xerr, yerr, y_is_detected=[], savename=None, show=True, xlabel='$x$', ylabel='$y$', errorbar_kwargs={'alpha': 1, 'capsize': 1, 'color': 'k', 'elinewidth': 0.5, 'fmt': '.', 'markersize': 1, 'zorder': 10}, fgivenx_kwargs={}, xscale='linear', yscale='linear', xlim=None, ylim=None)

Make the posterior predictive plot showing the 1, 2 and 3 sigma predictions of the function given the inferred parameters and plot the observed points on the same plot.

Args:

reg (roxy.regressor.RoxyRegressor):

The regressor object used for the inference

samples (dict):

The MCMC samples, where the keys are the parameter names and values are ndarrays of the samples

xobs (jnp.ndarray):

The observed x values

yobs (jnp.ndarray):

The observed y values

xerr (jnp.ndarray):

The error on the observed x values

yerr (jnp.ndarray):

The error on the observed y values

y_is_detected (array-like, default=[]):

Boolean array of the same length as yobs, where True indicates a detected point and False indicates an upper limit.

savename (str, default=None):

If not None, save the figure to the file given by this argument.

show (bool, default=True):

If True, display the figure with plt.show()

xlabel (str, default=’$x$’):

The label to use for the x axis

ylabel (str, default=’$x$’):

The label to use for the y axis

errorbar_kwargs (dict):

Dictionary of kwargs to pass to plt.errorbar

fgivenx_kwargs (dict):

Dictionary of kwargs to pass to fgivenx.plot_contours

xscale (str, default=’linear’):

Scale to use for x axis (‘linear’ or ‘log’)

yscale (str, default=’linear’):

Scale to use for y axis (‘linear’ or ‘log’)

xlim (tuple, default=None):

If not None, set the x limits to this value

ylim (tuple, default=None):

If not None, set the y limits to this value

Returns:

fig (matplotlib.figure.Figure):

The figure containing the posterior predictive plot

plotting.trace_plot(samples, to_plot='all', truths=None, savename=None, show=True)

Plot the trace of the parameter values as a function of MCMC step

Args:

samples (dict):

The MCMC samples, where the keys are the parameter names and values are ndarrays of the samples

to_plot (list, default=’all’):

If ‘all’, then use all parameters. If a list, then only use the parameters given in that list

truths (dict, default=None):

If not None, use this to specify the true values of the parameters to plot.

savename (str, default=None):

If not None, save the figure to the file given by this argument.

show (bool, default=True):

If True, display the figure with plt.show()

plotting.triangle_plot(samples, labels=None, to_plot='all', module='corner', truths=None, param_prior=None, savename=None, show=True)

Plot the 1D and 2D posterior distributions of the parameters in a triangle plot.

Args:

samples (dict):

The MCMC samples, where the keys are the parameter names and values are ndarrays of the samples

labels (dict, default=None):

Dictionary of parameter labels ot use in the plot. If None, then use the names given as keys in samples.

to_plot (list, default=’all’):

If ‘all’, then use all parameters. If a list, then only use the parameters given in that list

module (str, default=’corner’):

Which module to use to make the triangle plot (‘corner’ or ‘getdist’ currently available)

truths (dict, default=None):

If not None, use this to specify the true values of the parameters to plot.

param_prior (dict, default=None):

If not None and using ‘getdist’, use this to specify the range of the varibales to prevent undesirable smoothing effects.

savename (str, default=None):

If not None, save the figure to the file given by this argument.

show (bool, default=True):

If True, display the figure with plt.show()

causality

causality.assess_causality(fun, fun_inv, xobs, yobs, errors, param_names, param_default, param_prior, method='mnr', criterion='hsic', ngauss=1, covmat=False, gmm_prior='hierarchical', savename=None, show=True)

Due to the asymmetry between x and y, this function assesses whether one should fit y(x) or x(y) with the intrinsic scatter in the dependent variable. The code runs both forward fits, i.e. y(x) and x(y), and the reverse fits, where one uses the parameters obtained from the forward fit and uses the inverse function. The correlation coefficient of the residuals (normalised by the square root of the sum of the squares of the vertical errors and the intrindic scatter) is then found, and the method which minimsies this recommended as the best option. Plots of the fits and the residuals are produced, since these can be more informative for some functions. The recommended setup is indicated by a star in the plots.

Args:

fun (callable):

The function, f, to be considered, y = f(x, theta). The function must take two arguments, the first of which is the independent variable, the second of which are the parameters (as an array or list).

fun_inv (callable):

If fun is y = f(x, theta), then this callable gives x = f^{-1}(y, theta) for the same theta. The function must take two arguments, the first corresponding to y and the second is theta.

xobs (jnp.ndarray):

The observed x values

yobs (jnp.ndarray):

The observed y values

errors (jnp.ndarray):

If covmat=False, then this is [xerr, yerr], giving the error on the observed x and y values. Otherwise, this is the covariance matrix in the order (x, y)

param_names (list):

The list of parameter names, in the order which they are supplied to fun

param_default (list):

The default valus of the parameters

param_prior (dict):

The prior range for each of the parameters. The prior is assumed to be uniform in this range. If either entry is None in the prior, then an infinite uniform prior is assumed.

method (str, default=’mnr’):

The name of the likelihood method to use (‘mnr’, ‘gmm’, ‘unif’ or ‘prof’). See roxy.likelihoods for more information

criterion (str, default=’hsic’):

The method used to determine the weakest correlation of the residuals. One of ‘hsic’, ‘spearman’ or ‘pearson’.

ngauss (int, default = 1):

The number of Gaussians to use in the GMM prior. Only used if method=’gmm’

covmat (bool, default=False):

This determines whether the errors argument is [xerr, yerr] (False) or a covariance matrix (True).

gmm_prior (string, default=’hierarchical’):

If method=’gmm’, this decides what prior to put on the GMM componenents. If ‘uniform’, then the mean and widths have a uniform prior, and if ‘hierarchical’ mu and w^2 have a Normal and Inverse Gamma prior, respectively.

savename (str, default=None):

If not None, save the figure to the file given by this argument.

show (bool, default=True):

If True, display the figure with plt.show()

causality.compute_hsic(x, y, alph=0.05)

Python implementation of Hilbert Schmidt Independence Criterion using a Gamma approximation. This code is largely taken from https://github.com/amber0309/HSIC/tree/master (Shoubo (shoubo.sub AT gmail.com) 09/11/2016), with the new addition of the computation of the significance level.

Gretton, A., Fukumizu, K., Teo, C. H., Song, L., Scholkopf, B., & Smola, A. J. (2007). A kernel statistical test of independence. In Advances in neural information processing systems (pp. 585-592).

Args:

x (np.ndarray):

Numpy vector of dependent variable. The row gives the sample and the column is the dimension.

y (np.ndarray):

Numpy vector of independent variable. The row gives the sample and the column is the dimension.

alph (float):

Worst significance level to consider. If the correlation is stronger than this, then we don’t compute the correlation coefficient.

Returns:

test_stat (float):

The test statistic

best_alph (float):

The significance of this result (if calculated)