utils#

Utility functions for pioran.

Submodules#

Overview#

Function#
`EuclideanDistance`(xq, xp)	Compute the Euclidean distance between two arrays.
`progress_bar_factory`(num_samples, num_chains)	Factory that builds a progress bar decorator along
`save_sampling_results`(info, warmup, samples, log_prob, log_densitygrad, filename)	Save the results of the Monte Carlo runs in a ASDF file.
`SHO_power_spectrum`(f, A, f0)	Power spectrum of a stochastic harmonic oscillator.
`get_samples_psd`(psd_acvf, f, params_samples)	Just a wrapper for jax.vmap(get_psd_approx_samples,(None,None,0))(psd_acvf,f,params_samples)
`wrapper_psd_true_samples`(psd_acvf, f, params_samples)	Just a wrapper for jax.vmap(get_psd_true_samples,(None,None,0))(psd_acvf,f,params_samples)

Functions#

pioran.utils.EuclideanDistance(xq, xp)[source]#

Compute the Euclidean distance between two arrays.

(1)#\[D(\boldsymbol{x_q},\boldsymbol{x_p}) = \sqrt{(\boldsymbol{x_q} - \boldsymbol{x_p}^{\mathrm{T}})^2}\]

Parameters:

xq(n, 1) jax.Array: First array.
xp(m, 1) jax.Array: Second array.
Returns
——-
(n, m) :obj:`jax.Array`

pioran.utils.progress_bar_factory(num_samples: int, num_chains: int)[source]#

Factory that builds a progress bar decorator along with the set_tqdm_description and close_tqdm functions

progress bar obtained from numpyro source code pyro-ppl/numpyro and modified to work with scan using https://www.jeremiecoullon.com/2021/01/29/jax_progress_bar/

Parameters:

num_samples: :obj:`int`: The number of samples
num_chains: :obj:`int`: The number of chains

pioran.utils.save_sampling_results(info: dict, warmup: dict, samples: numpy.ndarray, log_prob: numpy.ndarray, log_densitygrad: numpy.ndarray, filename: str)[source]#

Save the results of the Monte Carlo runs in a ASDF file.

This file contains the following data: - info: a dictionary containing the information about the run, namely: - num_params: the number of parameters - num_samples: the number of samples - num_warmup: the number of warmup samples - num_chains: the number of chains - ESS: the effective sample size - Rhat-split: the split Rhat statistic - warmup: a numpy array containing the warmup samples - samples: a numpy array containing the samples - log_prob: a numpy array containing the log probabilities of the samples

Parameters:

info: :obj:`dict`: A dictionary containing the information about the run
warmup: :obj:`dict`: A numpy array containing the warmup samples
samples: :obj:`jax.Array`: A numpy array containing the samples
log_prob: :obj:`jax.Array`: A numpy array containing the log probabilities of the samples
log_densitygrad: :obj:`jax.Array`: A numpy array containing the log density gradients of the samples
filename: :obj:`str`: The name of the file to save the data to

pioran.utils.SHO_power_spectrum(f: jax.Array, A: float, f0: float) → jax.Array[source]#

Power spectrum of a stochastic harmonic oscillator.

(2)#\[\mathcal{P}(f) = \dfrac{A}{1 + (f/f_0)^4}.\]

with the amplitude \(A\), the position \(f_0\ge 0\).

Parameters:

fjax.Array: Frequency array.
Afloat: Amplitude.
f0float: Position.

Returns:

jax.Array

pioran.utils.get_samples_psd(psd_acvf: pioran.psdtoacv.PSDToACV, f: jax.Array, params_samples: jax.Array) → jax.Array[source]#

Just a wrapper for jax.vmap(get_psd_approx_samples,(None,None,0))(psd_acvf,f,params_samples)

Parameters:

psd_acvfPSDToACV: PSDToACV object.
fjax.Array: Frequency array.
params_samplesjax.Array: Parameters of the PSD model.

pioran.utils.wrapper_psd_true_samples(psd_acvf: pioran.psdtoacv.PSDToACV, f: jax.Array, params_samples: jax.Array) → jax.Array[source]#

Just a wrapper for jax.vmap(get_psd_true_samples,(None,None,0))(psd_acvf,f,params_samples)

Parameters:

psd_acvfPSDToACV: PSDToACV object.
fjax.Array: Frequency array.
params_samplesjax.Array: Parameters of the PSD model.