tensorflow_probability/python/mcmc/diagnostic.py
def _effective_sample_size_single_state(states, filter_beyond_lag,
                                        filter_threshold):
  """ESS computation for one single Tensor argument."""

  with tf.compat.v1.name_scope(
      'effective_sample_size_single_state',
      values=[states, filter_beyond_lag, filter_threshold]):

    states = tf.convert_to_tensor(value=states, name='states')
    dt = states.dtype

    # filter_beyond_lag == None ==> auto_corr is the full sequence.
    auto_corr = stats.auto_correlation(
        states, axis=0, max_lags=filter_beyond_lag)
    if filter_threshold is not None:
      filter_threshold = tf.convert_to_tensor(
          value=filter_threshold, dtype=dt, name='filter_threshold')
      # Get a binary mask to zero out values of auto_corr below the threshold.
      #   mask[i, ...] = 1 if auto_corr[j, ...] > threshold for all j <= i,
      #   mask[i, ...] = 0, otherwise.
      # So, along dimension zero, the mask will look like [1, 1, ..., 0, 0,...]
      # Building step by step,
      #   Assume auto_corr = [1, 0.5, 0.0, 0.3], and filter_threshold = 0.2.
      # Step 1:  mask = [False, False, True, False]
      mask = auto_corr < filter_threshold
      # Step 2:  mask = [0, 0, 1, 1]
      mask = tf.cast(mask, dtype=dt)
      # Step 3:  mask = [0, 0, 1, 2]
      mask = tf.cumsum(mask, axis=0)
      # Step 4:  mask = [1, 1, 0, 0]
      mask = tf.maximum(1. - mask, 0.)
      auto_corr *= mask

    # With R[k] := auto_corr[k, ...],
    # ESS = N / {1 + 2 * Sum_{k=1}^N (N - k) / N * R[k]}
    #     = N / {-1 + 2 * Sum_{k=0}^N (N - k) / N * R[k]} (since R[0] = 1)
    #     approx N / {-1 + 2 * Sum_{k=0}^M (N - k) / N * R[k]}
    # where M is the filter_beyond_lag truncation point chosen above.

    # Get the factor (N - k) / N, and give it shape [M, 1,...,1], having total
    # ndims the same as auto_corr
    n = _axis_size(states, axis=0)
    k = tf.range(0., _axis_size(auto_corr, axis=0))
    nk_factor = (n - k) / n
    if auto_corr.shape.ndims is not None:
      new_shape = [-1] + [1] * (auto_corr.shape.ndims - 1)
    else:
      new_shape = tf.concat(
          ([-1],
           tf.ones([tf.rank(auto_corr) - 1], dtype=tf.int32)),
          axis=0)
    nk_factor = tf.reshape(nk_factor, new_shape)

    return n / (-1 +
                2 * tf.reduce_sum(input_tensor=nk_factor * auto_corr, axis=0))
def _effective_sample_size_single_state(states, filter_beyond_lag,
                                        filter_threshold):
  with tf.compat.v1.name_scope(
      'effective_sample_size_single_state',
      values=[states, filter_beyond_lag, filter_threshold]):

    states = tf.convert_to_tensor(value=states, name='states')
    dt = states.dtype

    # filter_beyond_lag == None ==> auto_corr is the full sequence.
    auto_corr = stats.auto_correlation(
        states, axis=0, max_lags=filter_beyond_lag)
    if filter_threshold is not None:
      filter_threshold = tf.convert_to_tensor(
          value=filter_threshold, dtype=dt, name='filter_threshold')
      # Get a binary mask to zero out values of auto_corr below the threshold.
      #   mask[i, ...] = 1 if auto_corr[j, ...] > threshold for all j <= i,
      #   mask[i, ...] = 0, otherwise.
      # So, along dimension zero, the mask will look like [1, 1, ..., 0, 0,...]
      # Building step by step,
      #   Assume auto_corr = [1, 0.5, 0.0, 0.3], and filter_threshold = 0.2.
      # Step 1:  mask = [False, False, True, False]
      mask = auto_corr < filter_threshold
      # Step 2:  mask = [0, 0, 1, 1]
      mask = tf.cast(mask, dtype=dt)
      # Step 3:  mask = [0, 0, 1, 2]
      mask = tf.cumsum(mask, axis=0)
      # Step 4:  mask = [1, 1, 0, 0]
      mask = tf.maximum(1. - mask, 0.)
      auto_corr *= mask

    # With R[k] := auto_corr[k, ...],
    # ESS = N / {1 + 2 * Sum_{k=1}^N (N - k) / N * R[k]}
    #     = N / {-1 + 2 * Sum_{k=0}^N (N - k) / N * R[k]} (since R[0] = 1)
    #     approx N / {-1 + 2 * Sum_{k=0}^M (N - k) / N * R[k]}
    # where M is the filter_beyond_lag truncation point chosen above.

    # Get the factor (N - k) / N, and give it shape [M, 1,...,1], having total
    # ndims the same as auto_corr
    n = _axis_size(states, axis=0)
    k = tf.range(0., _axis_size(auto_corr, axis=0))
    nk_factor = (n - k) / n
    if auto_corr.shape.ndims is not None:
      new_shape = [-1] + [1] * (auto_corr.shape.ndims - 1)
    else:
      new_shape = tf.concat(
          ([-1],
           tf.ones([tf.rank(auto_corr) - 1], dtype=tf.int32)),
          axis=0)
    nk_factor = tf.reshape(nk_factor, new_shape)

    return n / (-1 +
                2 * tf.reduce_sum(input_tensor=nk_factor * auto_corr, axis=0))
def _effective_sample_size_single_state ( states , filter_beyond_lag , filter_threshold ) : with tf . compat . v1 . name_scope ( 'effective_sample_size_single_state' , values = [ states , filter_beyond_lag , filter_threshold ] ) : states = tf . convert_to_tensor ( value = states , name = 'states' ) dt = states . dtype # filter_beyond_lag == None ==> auto_corr is the full sequence. auto_corr = stats . auto_correlation ( states , axis = 0 , max_lags = filter_beyond_lag ) if filter_threshold is not None : filter_threshold = tf . convert_to_tensor ( value = filter_threshold , dtype = dt , name = 'filter_threshold' ) # Get a binary mask to zero out values of auto_corr below the threshold. #   mask[i, ...] = 1 if auto_corr[j, ...] > threshold for all j <= i, #   mask[i, ...] = 0, otherwise. # So, along dimension zero, the mask will look like [1, 1, ..., 0, 0,...] # Building step by step, #   Assume auto_corr = [1, 0.5, 0.0, 0.3], and filter_threshold = 0.2. # Step 1:  mask = [False, False, True, False] mask = auto_corr < filter_threshold # Step 2:  mask = [0, 0, 1, 1] mask = tf . cast ( mask , dtype = dt ) # Step 3:  mask = [0, 0, 1, 2] mask = tf . cumsum ( mask , axis = 0 ) # Step 4:  mask = [1, 1, 0, 0] mask = tf . maximum ( 1. - mask , 0. ) auto_corr *= mask # With R[k] := auto_corr[k, ...], # ESS = N / {1 + 2 * Sum_{k=1}^N (N - k) / N * R[k]} #     = N / {-1 + 2 * Sum_{k=0}^N (N - k) / N * R[k]} (since R[0] = 1) #     approx N / {-1 + 2 * Sum_{k=0}^M (N - k) / N * R[k]} # where M is the filter_beyond_lag truncation point chosen above. # Get the factor (N - k) / N, and give it shape [M, 1,...,1], having total # ndims the same as auto_corr n = _axis_size ( states , axis = 0 ) k = tf . range ( 0. , _axis_size ( auto_corr , axis = 0 ) ) nk_factor = ( n - k ) / n if auto_corr . shape . ndims is not None : new_shape = [ - 1 ] + [ 1 ] * ( auto_corr . shape . ndims - 1 ) else : new_shape = tf . concat ( ( [ - 1 ] , tf . ones ( [ tf . rank ( auto_corr ) - 1 ] , dtype = tf . int32 ) ) , axis = 0 ) nk_factor = tf . reshape ( nk_factor , new_shape ) return n / ( - 1 + 2 * tf . reduce_sum ( input_tensor = nk_factor * auto_corr , axis = 0 ) )
ESS computation for one single Tensor argument.
ESS computation for one single Tensor argument .