tensorflow_probability/python/stats/quantiles.py
def find_bins(x,
              edges,
              extend_lower_interval=False,
              extend_upper_interval=False,
              dtype=None,
              name=None):
  """Bin values into discrete intervals.

  Given `edges = [c0, ..., cK]`, defining intervals
  `I0 = [c0, c1)`, `I1 = [c1, c2)`, ..., `I_{K-1} = [c_{K-1}, cK]`,
  This function returns `bins`, such that:
  `edges[bins[i]] <= x[i] < edges[bins[i] + 1]`.

  Args:
    x:  Numeric `N-D` `Tensor` with `N > 0`.
    edges:  `Tensor` of same `dtype` as `x`.  The first dimension indexes edges
      of intervals.  Must either be `1-D` or have
      `x.shape[1:] == edges.shape[1:]`.  If `rank(edges) > 1`, `edges[k]`
      designates a shape `edges.shape[1:]` `Tensor` of bin edges for the
      corresponding dimensions of `x`.
    extend_lower_interval:  Python `bool`.  If `True`, extend the lowest
      interval `I0` to `(-inf, c1]`.
    extend_upper_interval:  Python `bool`.  If `True`, extend the upper
      interval `I_{K-1}` to `[c_{K-1}, +inf)`.
    dtype: The output type (`int32` or `int64`). `Default value:` `x.dtype`.
      This effects the output values when `x` is below/above the intervals,
      which will be `-1/K+1` for `int` types and `NaN` for `float`s.
      At indices where `x` is `NaN`, the output values will be `0` for `int`
      types and `NaN` for floats.
    name:  A Python string name to prepend to created ops. Default: 'find_bins'

  Returns:
    bins: `Tensor` with same `shape` as `x` and `dtype`.
      Has whole number values.  `bins[i] = k` means the `x[i]` falls into the
      `kth` bin, ie, `edges[bins[i]] <= x[i] < edges[bins[i] + 1]`.

  Raises:
    ValueError:  If `edges.shape[0]` is determined to be less than 2.

  #### Examples

  Cut a `1-D` array

  ```python
  x = [0., 5., 6., 10., 20.]
  edges = [0., 5., 10.]
  tfp.stats.find_bins(x, edges)
  ==> [0., 0., 1., 1., np.nan]
  ```

  Cut `x` into its deciles

  ```python
  x = tf.random_uniform(shape=(100, 200))
  decile_edges = tfp.stats.quantiles(x, num_quantiles=10)
  bins = tfp.stats.find_bins(x, edges=decile_edges)
  bins.shape
  ==> (100, 200)
  tf.reduce_mean(bins == 0.)
  ==> approximately 0.1
  tf.reduce_mean(bins == 1.)
  ==> approximately 0.1
  ```

  """
  # TFP users may be surprised to see the "action" in the leftmost dim of
  # edges, rather than the rightmost (event) dim.  Why?
  # 1. Most likely you created edges by getting quantiles over samples, and
  #    quantile/percentile return these edges in the leftmost (sample) dim.
  # 2. Say you have event_shape = [5], then we expect the bin will be different
  #    for all 5 events, so the index of the bin should not be in the event dim.
  with tf.compat.v1.name_scope(
      name, default_name='find_bins', values=[x, edges]):
    in_type = dtype_util.common_dtype([x, edges],
                                      preferred_dtype=tf.float32)
    edges = tf.convert_to_tensor(value=edges, name='edges', dtype=in_type)
    x = tf.convert_to_tensor(value=x, name='x', dtype=in_type)

    if (tf.compat.dimension_value(edges.shape[0]) is not None and
        tf.compat.dimension_value(edges.shape[0]) < 2):
      raise ValueError(
          'First dimension of `edges` must have length > 1 to index 1 or '
          'more bin. Found: {}'.format(edges.shape))

    flattening_x = edges.shape.ndims == 1 and x.shape.ndims > 1

    if flattening_x:
      x_orig_shape = tf.shape(input=x)
      x = tf.reshape(x, [-1])

    if dtype is None:
      dtype = in_type
    dtype = tf.as_dtype(dtype)

    # Move first dims into the rightmost.
    x_permed = distribution_util.rotate_transpose(x, shift=-1)
    edges_permed = distribution_util.rotate_transpose(edges, shift=-1)

    # If...
    #   x_permed = [0, 1, 6., 10]
    #   edges = [0, 5, 10.]
    #   ==> almost_output = [0, 1, 2, 2]
    searchsorted_type = dtype if dtype in [tf.int32, tf.int64] else None
    almost_output_permed = tf.searchsorted(
        sorted_sequence=edges_permed,
        values=x_permed,
        side='right',
        out_type=searchsorted_type)
    # Move the rightmost dims back to the leftmost.
    almost_output = tf.cast(
        distribution_util.rotate_transpose(almost_output_permed, shift=1),
        dtype)

    # In above example, we want [0, 0, 1, 1], so correct this here.
    bins = tf.clip_by_value(almost_output - 1, tf.cast(0, dtype),
                            tf.cast(tf.shape(input=edges)[0] - 2, dtype))

    if not extend_lower_interval:
      low_fill = np.nan if dtype.is_floating else -1
      bins = tf.where(x < tf.expand_dims(edges[0], 0),
                      tf.fill(tf.shape(input=x), tf.cast(low_fill, dtype)),
                      bins)

    if not extend_upper_interval:
      up_fill = np.nan if dtype.is_floating else tf.shape(input=edges)[0] - 1
      bins = tf.where(x > tf.expand_dims(edges[-1], 0),
                      tf.fill(tf.shape(input=x), tf.cast(up_fill, dtype)), bins)

    if flattening_x:
      bins = tf.reshape(bins, x_orig_shape)

    return bins
def find_bins(x,
              edges,
              extend_lower_interval=False,
              extend_upper_interval=False,
              dtype=None,
              name=None):
  """Bin values into discrete intervals.

  Given `edges = [c0, ..., cK]`, defining intervals
  `I0 = [c0, c1)`, `I1 = [c1, c2)`, ..., `I_{K-1} = [c_{K-1}, cK]`,
  This function returns `bins`, such that:
  `edges[bins[i]] <= x[i] < edges[bins[i] + 1]`.

  Args:
    x:  Numeric `N-D` `Tensor` with `N > 0`.
    edges:  `Tensor` of same `dtype` as `x`.  The first dimension indexes edges
      of intervals.  Must either be `1-D` or have
      `x.shape[1:] == edges.shape[1:]`.  If `rank(edges) > 1`, `edges[k]`
      designates a shape `edges.shape[1:]` `Tensor` of bin edges for the
      corresponding dimensions of `x`.
    extend_lower_interval:  Python `bool`.  If `True`, extend the lowest
      interval `I0` to `(-inf, c1]`.
    extend_upper_interval:  Python `bool`.  If `True`, extend the upper
      interval `I_{K-1}` to `[c_{K-1}, +inf)`.
    dtype: The output type (`int32` or `int64`). `Default value:` `x.dtype`.
      This effects the output values when `x` is below/above the intervals,
      which will be `-1/K+1` for `int` types and `NaN` for `float`s.
      At indices where `x` is `NaN`, the output values will be `0` for `int`
      types and `NaN` for floats.
    name:  A Python string name to prepend to created ops. Default: 'find_bins'

  Returns:
    bins: `Tensor` with same `shape` as `x` and `dtype`.
      Has whole number values.  `bins[i] = k` means the `x[i]` falls into the
      `kth` bin, ie, `edges[bins[i]] <= x[i] < edges[bins[i] + 1]`.

  Raises:
    ValueError:  If `edges.shape[0]` is determined to be less than 2.

  #### Examples

  Cut a `1-D` array

  ```python
  x = [0., 5., 6., 10., 20.]
  edges = [0., 5., 10.]
  tfp.stats.find_bins(x, edges)
  ==> [0., 0., 1., 1., np.nan]
  ```

  Cut `x` into its deciles

  ```python
  x = tf.random_uniform(shape=(100, 200))
  decile_edges = tfp.stats.quantiles(x, num_quantiles=10)
  bins = tfp.stats.find_bins(x, edges=decile_edges)
  bins.shape
  ==> (100, 200)
  tf.reduce_mean(bins == 0.)
  ==> approximately 0.1
  tf.reduce_mean(bins == 1.)
  ==> approximately 0.1
  ```

  """
  # TFP users may be surprised to see the "action" in the leftmost dim of
  # edges, rather than the rightmost (event) dim.  Why?
  # 1. Most likely you created edges by getting quantiles over samples, and
  #    quantile/percentile return these edges in the leftmost (sample) dim.
  # 2. Say you have event_shape = [5], then we expect the bin will be different
  #    for all 5 events, so the index of the bin should not be in the event dim.
  with tf.compat.v1.name_scope(
      name, default_name='find_bins', values=[x, edges]):
    in_type = dtype_util.common_dtype([x, edges],
                                      preferred_dtype=tf.float32)
    edges = tf.convert_to_tensor(value=edges, name='edges', dtype=in_type)
    x = tf.convert_to_tensor(value=x, name='x', dtype=in_type)

    if (tf.compat.dimension_value(edges.shape[0]) is not None and
        tf.compat.dimension_value(edges.shape[0]) < 2):
      raise ValueError(
          'First dimension of `edges` must have length > 1 to index 1 or '
          'more bin. Found: {}'.format(edges.shape))

    flattening_x = edges.shape.ndims == 1 and x.shape.ndims > 1

    if flattening_x:
      x_orig_shape = tf.shape(input=x)
      x = tf.reshape(x, [-1])

    if dtype is None:
      dtype = in_type
    dtype = tf.as_dtype(dtype)

    # Move first dims into the rightmost.
    x_permed = distribution_util.rotate_transpose(x, shift=-1)
    edges_permed = distribution_util.rotate_transpose(edges, shift=-1)

    # If...
    #   x_permed = [0, 1, 6., 10]
    #   edges = [0, 5, 10.]
    #   ==> almost_output = [0, 1, 2, 2]
    searchsorted_type = dtype if dtype in [tf.int32, tf.int64] else None
    almost_output_permed = tf.searchsorted(
        sorted_sequence=edges_permed,
        values=x_permed,
        side='right',
        out_type=searchsorted_type)
    # Move the rightmost dims back to the leftmost.
    almost_output = tf.cast(
        distribution_util.rotate_transpose(almost_output_permed, shift=1),
        dtype)

    # In above example, we want [0, 0, 1, 1], so correct this here.
    bins = tf.clip_by_value(almost_output - 1, tf.cast(0, dtype),
                            tf.cast(tf.shape(input=edges)[0] - 2, dtype))

    if not extend_lower_interval:
      low_fill = np.nan if dtype.is_floating else -1
      bins = tf.where(x < tf.expand_dims(edges[0], 0),
                      tf.fill(tf.shape(input=x), tf.cast(low_fill, dtype)),
                      bins)

    if not extend_upper_interval:
      up_fill = np.nan if dtype.is_floating else tf.shape(input=edges)[0] - 1
      bins = tf.where(x > tf.expand_dims(edges[-1], 0),
                      tf.fill(tf.shape(input=x), tf.cast(up_fill, dtype)), bins)

    if flattening_x:
      bins = tf.reshape(bins, x_orig_shape)

    return bins
Bin values into discrete intervals.

  Given `edges = [c0, ..., cK]`, defining intervals
  `I0 = [c0, c1)`, `I1 = [c1, c2)`, ..., `I_{K-1} = [c_{K-1}, cK]`,
  This function returns `bins`, such that:
  `edges[bins[i]] <= x[i] < edges[bins[i] + 1]`.

  Args:
    x:  Numeric `N-D` `Tensor` with `N > 0`.
    edges:  `Tensor` of same `dtype` as `x`.  The first dimension indexes edges
      of intervals.  Must either be `1-D` or have
      `x.shape[1:] == edges.shape[1:]`.  If `rank(edges) > 1`, `edges[k]`
      designates a shape `edges.shape[1:]` `Tensor` of bin edges for the
      corresponding dimensions of `x`.
    extend_lower_interval:  Python `bool`.  If `True`, extend the lowest
      interval `I0` to `(-inf, c1]`.
    extend_upper_interval:  Python `bool`.  If `True`, extend the upper
      interval `I_{K-1}` to `[c_{K-1}, +inf)`.
    dtype: The output type (`int32` or `int64`). `Default value:` `x.dtype`.
      This effects the output values when `x` is below/above the intervals,
      which will be `-1/K+1` for `int` types and `NaN` for `float`s.
      At indices where `x` is `NaN`, the output values will be `0` for `int`
      types and `NaN` for floats.
    name:  A Python string name to prepend to created ops. Default: 'find_bins'

  Returns:
    bins: `Tensor` with same `shape` as `x` and `dtype`.
      Has whole number values.  `bins[i] = k` means the `x[i]` falls into the
      `kth` bin, ie, `edges[bins[i]] <= x[i] < edges[bins[i] + 1]`.

  Raises:
    ValueError:  If `edges.shape[0]` is determined to be less than 2.

  #### Examples

  Cut a `1-D` array

  ```python
  x = [0., 5., 6., 10., 20.]
  edges = [0., 5., 10.]
  tfp.stats.find_bins(x, edges)
  ==> [0., 0., 1., 1., np.nan]
  ```

  Cut `x` into its deciles

  ```python
  x = tf.random_uniform(shape=(100, 200))
  decile_edges = tfp.stats.quantiles(x, num_quantiles=10)
  bins = tfp.stats.find_bins(x, edges=decile_edges)
  bins.shape
  ==> (100, 200)
  tf.reduce_mean(bins == 0.)
  ==> approximately 0.1
  tf.reduce_mean(bins == 1.)
  ==> approximately 0.1
  ```