tensorflow_probability/python/distributions/von_mises_fisher.py
def _sample_3d(self, n, seed=None):
    """Specialized inversion sampler for 3D."""
    seed = seed_stream.SeedStream(seed, salt='von_mises_fisher_3d')
    u_shape = tf.concat([[n], self._batch_shape_tensor()], axis=0)
    z = tf.random.uniform(u_shape, seed=seed(), dtype=self.dtype)
    # TODO(bjp): Higher-order odd dim analytic CDFs are available in [1], could
    # be bisected for bounded sampling runtime (i.e. not rejection sampling).
    # [1]: Inversion sampler via: https://ieeexplore.ieee.org/document/7347705/
    # The inversion is: u = 1 + log(z + (1-z)*exp(-2*kappa)) / kappa
    # We must protect against both kappa and z being zero.
    safe_conc = tf.where(self.concentration > 0,
                         self.concentration,
                         tf.ones_like(self.concentration))
    safe_z = tf.where(z > 0, z, tf.ones_like(z))
    safe_u = 1 + tf.reduce_logsumexp(
        input_tensor=[
            tf.math.log(safe_z),
            tf.math.log1p(-safe_z) - 2 * safe_conc
        ],
        axis=0) / safe_conc
    # Limit of the above expression as kappa->0 is 2*z-1
    u = tf.where(self.concentration > tf.zeros_like(safe_u), safe_u,
                 2 * z - 1)
    # Limit of the expression as z->0 is -1.
    u = tf.where(tf.equal(z, 0), -tf.ones_like(u), u)
    if not self._allow_nan_stats:
      u = tf.debugging.check_numerics(u, 'u in _sample_3d')
    return u[..., tf.newaxis]
def _sample_3d(self, n, seed=None):
    seed = seed_stream.SeedStream(seed, salt='von_mises_fisher_3d')
    u_shape = tf.concat([[n], self._batch_shape_tensor()], axis=0)
    z = tf.random.uniform(u_shape, seed=seed(), dtype=self.dtype)
    # TODO(bjp): Higher-order odd dim analytic CDFs are available in [1], could
    # be bisected for bounded sampling runtime (i.e. not rejection sampling).
    # [1]: Inversion sampler via: https://ieeexplore.ieee.org/document/7347705/
    # The inversion is: u = 1 + log(z + (1-z)*exp(-2*kappa)) / kappa
    # We must protect against both kappa and z being zero.
    safe_conc = tf.where(self.concentration > 0,
                         self.concentration,
                         tf.ones_like(self.concentration))
    safe_z = tf.where(z > 0, z, tf.ones_like(z))
    safe_u = 1 + tf.reduce_logsumexp(
        input_tensor=[
            tf.math.log(safe_z),
            tf.math.log1p(-safe_z) - 2 * safe_conc
        ],
        axis=0) / safe_conc
    # Limit of the above expression as kappa->0 is 2*z-1
    u = tf.where(self.concentration > tf.zeros_like(safe_u), safe_u,
                 2 * z - 1)
    # Limit of the expression as z->0 is -1.
    u = tf.where(tf.equal(z, 0), -tf.ones_like(u), u)
    if not self._allow_nan_stats:
      u = tf.debugging.check_numerics(u, 'u in _sample_3d')
    return u[..., tf.newaxis]
def _sample_3d ( self , n , seed = None ) : seed = seed_stream . SeedStream ( seed , salt = 'von_mises_fisher_3d' ) u_shape = tf . concat ( [ [ n ] , self . _batch_shape_tensor ( ) ] , axis = 0 ) z = tf . random . uniform ( u_shape , seed = seed ( ) , dtype = self . dtype ) # TODO(bjp): Higher-order odd dim analytic CDFs are available in [1], could # be bisected for bounded sampling runtime (i.e. not rejection sampling). # [1]: Inversion sampler via: https://ieeexplore.ieee.org/document/7347705/ # The inversion is: u = 1 + log(z + (1-z)*exp(-2*kappa)) / kappa # We must protect against both kappa and z being zero. safe_conc = tf . where ( self . concentration > 0 , self . concentration , tf . ones_like ( self . concentration ) ) safe_z = tf . where ( z > 0 , z , tf . ones_like ( z ) ) safe_u = 1 + tf . reduce_logsumexp ( input_tensor = [ tf . math . log ( safe_z ) , tf . math . log1p ( - safe_z ) - 2 * safe_conc ] , axis = 0 ) / safe_conc # Limit of the above expression as kappa->0 is 2*z-1 u = tf . where ( self . concentration > tf . zeros_like ( safe_u ) , safe_u , 2 * z - 1 ) # Limit of the expression as z->0 is -1. u = tf . where ( tf . equal ( z , 0 ) , - tf . ones_like ( u ) , u ) if not self . _allow_nan_stats : u = tf . debugging . check_numerics ( u , 'u in _sample_3d' ) return u [ ... , tf . newaxis ]
Specialized inversion sampler for 3D.
Specialized inversion sampler for 3D .