utils
===============

.. py:module:: jaxns.utils

.. rubric:: :code:`jaxns.utils`


.. rubric:: Module Contents

.. py:function:: resample(key, samples, log_weights, S = None, replace = True)

   Resample the weighted samples into uniformly weighted samples.

   :param key: PRNGKey
   :param samples: samples from nested sampled results
   :param log_weights: log-posterior weight
   :param S: number of samples to generate. Will use Kish's estimate of ESS if None.
   :param replace: whether to sample with replacement

   :returns: equally weighted samples


.. py:function:: marginalise_static_from_U(key, U_samples, model, log_weights, ESS, fun)

   Marginalises function over posterior samples, where ESS is static.

   :param key: PRNG key
   :param U_samples: array of U samples
   :param model: model
   :param log_weights: log weights from nested sampling
   :param ESS: static effective sample size
   :param fun: function to marginalise
   :type fun: :code:`callable(**kwargs)`

   :returns: expectation over resampled samples


.. py:function:: marginalise_dynamic_from_U(key, U_samples, model, log_weights, ESS, fun)

   Marginalises function over posterior samples, where ESS can be dynamic.

   :param key: PRNG key
   :param U_samples: array of U samples
   :param model: model
   :param log_weights: log weights from nested sampling
   :param ESS: dynamic effective sample size
   :param fun: function to marginalise
   :type fun: :code:`callable(**kwargs)`

   :returns: expectation of `func` over resampled samples.


.. py:function:: marginalise_static(key, samples, log_weights, ESS, fun)

   Marginalises function over posterior samples, where ESS is static.

   :param key: PRNG key
   :param samples: dict of batched array of nested sampling samples
   :type samples: dict
   :param log_weights: log weights from nested sampling
   :param ESS: static effective sample size
   :param fun: function to marginalise
   :type fun: :code:`callable(**kwargs)`

   :returns: expectation over resampled samples


.. py:function:: marginalise_dynamic(key, samples, log_weights, ESS, fun)

   Marginalises function over posterior samples, where ESS can be dynamic.

   :param key: PRNG key
   :param samples: dict of batched array of nested sampling samples
   :type samples: dict
   :param log_weights: log weights from nested sampling
   :param ESS: dynamic effective sample size
   :param fun: function to marginalise
   :type fun: :code:`callable(**kwargs)`

   :returns: expectation of `func` over resampled samples.


.. py:function:: maximum_a_posteriori_point(results)

   Get the MAP point of a nested sampling result.
   Does this by choosing the point with largest L(x) p(x).

   :param results: Nested sampler result
   :type results: NestedSamplerResult

   :returns: dict of samples at MAP-point.


.. py:function:: evaluate_map_estimate(results, fun)

   Marginalises function over posterior samples, where ESS is static.

   :param results: results from run
   :param fun: function to marginalise
   :type fun: :code:`callable(**kwargs)`

   :returns: estimate at MAP sample point


.. py:function:: summary(results, with_parametrised = False, f_obj = None)

   Gives a summary of the results of a nested sampling run.

   :param results: Nested sampler result
   :type results: NestedSamplerResults
   :param with_parametrised: whether to include parametrised samples
   :param f_obj: file-like object to write summary to. If None, prints to stdout.


.. py:function:: sample_evidence(key, num_live_points_per_sample, log_L_samples, S = 100)

   Sample the evidence distribution, but stochastically simulating the shrinkage distribution.

   Note: this produces approximate samples, since there is also an uncertainty in the placement of the contours during
   shrinkage. Incorporating this stochasticity into the simulation would require running an entire
   nested sampling many times.

   :param key: PRNGKey
   :param num_live_points_per_sample: the number of live points for each sample
   :param log_L_samples: the log-L of samples
   :param S: The number of samples to produce

   :returns: samples of log(Z)


.. py:function:: bruteforce_posterior_samples(model, S = 60)

   Compute the posterior with brute-force over a regular grid.

   :param model: model
   :param S: resolution of grid

   :returns: samples, and log-weight


.. py:function:: bruteforce_evidence(model, S = 60)

   Compute the evidence with brute-force over a regular grid.

   :param model: model
   :param S: resolution of grid

   :returns: log(Z)


.. py:function:: analytic_posterior_samples(model, S = 60)

   Compute the evidence with brute-force over a regular grid.

   :param model: model
   :param S: resolution of grid

   :returns: log(Z)


.. py:function:: save_pytree(pytree, save_file)

   Saves results of nested sampler in a json file.

   :param pytree: Nested sampler result
   :param save_file: filename


.. py:function:: save_results(results, save_file)

   Saves results of nested sampler in a json file.

   :param results: Nested sampler result
   :type results: NestedSamplerResults
   :param save_file: filename
   :type save_file: str


.. py:function:: load_pytree(save_file)

   Loads saved nested sampler results from a json file.

   :param save_file: filename
   :type save_file: str

   :returns: NestedSamplerResults


.. py:function:: load_results(save_file)

   Loads saved nested sampler results from a json file.

   :param save_file: filename
   :type save_file: str

   :returns: NestedSamplerResults