multi_ellipsoid_utils
===============================

.. py:module:: jaxns.samplers.multi_ellipsoid.multi_ellipsoid_utils

.. rubric:: :code:`jaxns.samplers.multi_ellipsoid.multi_ellipsoid_utils`


.. rubric:: Module Contents

.. py:class:: MultEllipsoidState

   Bases: :py:obj:`NamedTuple`


   .. py:attribute:: params
      :type:  EllipsoidParams


   .. py:attribute:: cluster_id
      :type:  jaxns.internals.types.IntArray


.. py:function:: sample_multi_ellipsoid(key, mu, radii, rotation, unit_cube_constraint = True)

   Sample from a set of intersecting ellipsoids.
   When unit_cube_constraint=True then reject points outside the closed unit-cube.

   :param key: PRNGKey
   :param mu: [K, D] centres of ellipses
   :param radii: [K, D] radii of ellipses
   :param rotation: [K,D,D] rotation matrices of ellipses

   :returns: ellipsoid selected, and a sample point i.i.d. sampled from union of ellipsoids, of shape [D]


.. py:function:: ellipsoid_clustering(key, points, log_VS, max_num_ellipsoids, method = 'em_gmm')

   Partition live_points into 2^depth ellipsoids in depth-first order.

   :param key: PRNGKey
   :param points: [N, D] points to partition
   :param log_VS: expected true volume of points
   :param max_num_ellipsoids: the maximum number of ellipsoids

   :returns: params of multi-ellipsoids and cluster id of points