tensorflow_probability/python/math/interpolation.py
def _batch_interp_with_gather_nd(x, x_ref_min, x_ref_max, y_ref, nd, fill_value,
                                 batch_dims):
  """N-D interpolation that works with leading batch dims."""
  dtype = x.dtype

  # In this function,
  # x.shape = [A1, ..., An, D, nd], where n = batch_dims
  # and
  # y_ref.shape = [A1, ..., An, C1, C2,..., Cnd, B1,...,BM]
  # y_ref[A1, ..., An, i1,...,ind] is a shape [B1,...,BM] Tensor with the value
  # at index [i1,...,ind] in the interpolation table.
  #  and x_ref_max have shapes [A1, ..., An, nd].

  # ny[k] is number of y reference points in interp dim k.
  ny = tf.cast(tf.shape(input=y_ref)[batch_dims:batch_dims + nd], dtype)

  # Map [x_ref_min, x_ref_max] to [0, ny - 1].
  # This is the (fractional) index of x.
  # x_idx_unclipped[A1, ..., An, d, k] is the fractional index into dim k of
  # interpolation table for the dth x value.
  x_ref_min_expanded = tf.expand_dims(x_ref_min, axis=-2)
  x_ref_max_expanded = tf.expand_dims(x_ref_max, axis=-2)
  x_idx_unclipped = (ny - 1) * (x - x_ref_min_expanded) / (
      x_ref_max_expanded - x_ref_min_expanded)

  # Wherever x is NaN, x_idx_unclipped will be NaN as well.
  # Keep track of the nan indices here (so we can impute NaN later).
  # Also eliminate any NaN indices, since there is not NaN in 32bit.
  nan_idx = tf.math.is_nan(x_idx_unclipped)
  x_idx_unclipped = tf.where(nan_idx, tf.zeros_like(x_idx_unclipped),
                             x_idx_unclipped)

  # x_idx.shape = [A1, ..., An, D, nd]
  x_idx = tf.clip_by_value(x_idx_unclipped, tf.zeros((), dtype=dtype), ny - 1)

  # Get the index above and below x_idx.
  # Naively we could set idx_below = floor(x_idx), idx_above = ceil(x_idx),
  # however, this results in idx_below == idx_above whenever x is on a grid.
  # This in turn results in y_ref_below == y_ref_above, and then the gradient
  # at this point is zero.  So here we "jitter" one of idx_below, idx_above,
  # so that they are at different values.  This jittering does not affect the
  # interpolated value, but does make the gradient nonzero (unless of course
  # the y_ref values are the same).
  idx_below = tf.floor(x_idx)
  idx_above = tf.minimum(idx_below + 1, ny - 1)
  idx_below = tf.maximum(idx_above - 1, 0)

  # These are the values of y_ref corresponding to above/below indices.
  # idx_below_int32.shape = x.shape[:-1] + [nd]
  idx_below_int32 = tf.cast(idx_below, dtype=tf.int32)
  idx_above_int32 = tf.cast(idx_above, dtype=tf.int32)

  # idx_below_list is a length nd list of shape x.shape[:-1] int32 tensors.
  idx_below_list = tf.unstack(idx_below_int32, axis=-1)
  idx_above_list = tf.unstack(idx_above_int32, axis=-1)

  # Use t to get a convex combination of the below/above values.
  # t.shape = [A1, ..., An, D, nd]
  t = x_idx - idx_below

  # x, and tensors shaped like x, need to be added to, and selected with
  # (using tf.where) the output y.  This requires appending singletons.
  def _expand_x_fn(tensor):
    # Reshape tensor to tensor.shape + [1] * M.
    extended_shape = tf.concat([
        tf.shape(input=tensor),
        tf.ones_like(tf.shape(input=y_ref)[batch_dims + nd:])
    ],
                               axis=0)
    return tf.reshape(tensor, extended_shape)

  # Now, t.shape = [A1, ..., An, D, nd] + [1] * (rank(y_ref) - nd - batch_dims)
  t = _expand_x_fn(t)
  s = 1 - t

  # Re-insert NaN wherever x was NaN.
  nan_idx = _expand_x_fn(nan_idx)
  t = tf.where(nan_idx,
               tf.fill(tf.shape(input=t), tf.constant(np.nan, dtype)),
               t)

  terms = []
  # Our work above has located x's fractional index inside a cube of above/below
  # indices. The distance to the below indices is t, and to the above indices
  # is s.
  # Drawing lines from x to the cube walls, we get 2**nd smaller cubes. Each
  # term in the result is a product of a reference point, gathered from y_ref,
  # multiplied by a volume.  The volume is that of the cube opposite to the
  # reference point.  E.g. if the reference point is below x in every axis, the
  # volume is that of the cube with corner above x in every axis, s[0]*...*s[nd]
  # We could probably do this with one massive gather, but that would be very
  # unreadable and un-debuggable.  It also would create a large Tensor.
  for zero_ones_list in _binary_count(nd):
    gather_from_y_ref_idx = []
    opposite_volume_t_idx = []
    opposite_volume_s_idx = []
    for k, zero_or_one in enumerate(zero_ones_list):
      if zero_or_one == 0:
        # If the kth iterate has zero_or_one = 0,
        # Will gather from the "below" reference point along axis k.
        gather_from_y_ref_idx.append(idx_below_list[k])
        # Now append the index to gather for computing opposite_volume.
        # This could be done by initializing opposite_volume to 1, then here:
        #  opposite_volume *= tf.gather(s, indices=k, axis=tf.rank(x) - 1)
        # but that puts a gather in the "inner loop."  Better to append the
        # index and do one larger gather down below.
        opposite_volume_s_idx.append(k)
      else:
        gather_from_y_ref_idx.append(idx_above_list[k])
        # Append an index to gather, having the same effect as
        #   opposite_volume *= tf.gather(t, indices=k, axis=tf.rank(x) - 1)
        opposite_volume_t_idx.append(k)

    # Compute opposite_volume (volume of cube opposite the ref point):
    # Recall t.shape = s.shape = [D, nd] + [1, ..., 1]
    # Gather from t and s along the "nd" axis, which is rank(x) - 1.
    ov_axis = tf.rank(x) - 1
    opposite_volume = (
        tf.reduce_prod(
            input_tensor=tf.gather(
                t, indices=tf.cast(opposite_volume_t_idx, dtype=tf.int32),
                axis=ov_axis),
            axis=ov_axis) *
        tf.reduce_prod(
            input_tensor=tf.gather(
                s, indices=tf.cast(opposite_volume_s_idx, dtype=tf.int32),
                axis=ov_axis),
            axis=ov_axis)
    )  # pyformat: disable

    y_ref_pt = tf.gather_nd(
        y_ref,
        tf.stack(gather_from_y_ref_idx, axis=-1),
        batch_dims=batch_dims)

    terms.append(y_ref_pt * opposite_volume)

  y = tf.math.add_n(terms)

  if tf.debugging.is_numeric_tensor(fill_value):
    # Recall x_idx_unclipped.shape = [D, nd],
    # so here we check if it was out of bounds in any of the nd dims.
    # Thus, oob_idx.shape = [D].
    oob_idx = tf.reduce_any(
        input_tensor=(x_idx_unclipped < 0) | (x_idx_unclipped > ny - 1),
        axis=-1)

    # Now, y.shape = [D, B1,...,BM], so we'll have to broadcast oob_idx.

    oob_idx = _expand_x_fn(oob_idx)  # Shape [D, 1,...,1]
    oob_idx |= tf.fill(tf.shape(input=y), False)
    y = tf.where(oob_idx, tf.fill(tf.shape(input=y), fill_value), y)
  return y
def _batch_interp_with_gather_nd(x, x_ref_min, x_ref_max, y_ref, nd, fill_value,
                                 batch_dims):
  """N-D interpolation that works with leading batch dims."""
  dtype = x.dtype

  # In this function,
  # x.shape = [A1, ..., An, D, nd], where n = batch_dims
  # and
  # y_ref.shape = [A1, ..., An, C1, C2,..., Cnd, B1,...,BM]
  # y_ref[A1, ..., An, i1,...,ind] is a shape [B1,...,BM] Tensor with the value
  # at index [i1,...,ind] in the interpolation table.
  #  and x_ref_max have shapes [A1, ..., An, nd].

  # ny[k] is number of y reference points in interp dim k.
  ny = tf.cast(tf.shape(input=y_ref)[batch_dims:batch_dims + nd], dtype)

  # Map [x_ref_min, x_ref_max] to [0, ny - 1].
  # This is the (fractional) index of x.
  # x_idx_unclipped[A1, ..., An, d, k] is the fractional index into dim k of
  # interpolation table for the dth x value.
  x_ref_min_expanded = tf.expand_dims(x_ref_min, axis=-2)
  x_ref_max_expanded = tf.expand_dims(x_ref_max, axis=-2)
  x_idx_unclipped = (ny - 1) * (x - x_ref_min_expanded) / (
      x_ref_max_expanded - x_ref_min_expanded)

  # Wherever x is NaN, x_idx_unclipped will be NaN as well.
  # Keep track of the nan indices here (so we can impute NaN later).
  # Also eliminate any NaN indices, since there is not NaN in 32bit.
  nan_idx = tf.math.is_nan(x_idx_unclipped)
  x_idx_unclipped = tf.where(nan_idx, tf.zeros_like(x_idx_unclipped),
                             x_idx_unclipped)

  # x_idx.shape = [A1, ..., An, D, nd]
  x_idx = tf.clip_by_value(x_idx_unclipped, tf.zeros((), dtype=dtype), ny - 1)

  # Get the index above and below x_idx.
  # Naively we could set idx_below = floor(x_idx), idx_above = ceil(x_idx),
  # however, this results in idx_below == idx_above whenever x is on a grid.
  # This in turn results in y_ref_below == y_ref_above, and then the gradient
  # at this point is zero.  So here we "jitter" one of idx_below, idx_above,
  # so that they are at different values.  This jittering does not affect the
  # interpolated value, but does make the gradient nonzero (unless of course
  # the y_ref values are the same).
  idx_below = tf.floor(x_idx)
  idx_above = tf.minimum(idx_below + 1, ny - 1)
  idx_below = tf.maximum(idx_above - 1, 0)

  # These are the values of y_ref corresponding to above/below indices.
  # idx_below_int32.shape = x.shape[:-1] + [nd]
  idx_below_int32 = tf.cast(idx_below, dtype=tf.int32)
  idx_above_int32 = tf.cast(idx_above, dtype=tf.int32)

  # idx_below_list is a length nd list of shape x.shape[:-1] int32 tensors.
  idx_below_list = tf.unstack(idx_below_int32, axis=-1)
  idx_above_list = tf.unstack(idx_above_int32, axis=-1)

  # Use t to get a convex combination of the below/above values.
  # t.shape = [A1, ..., An, D, nd]
  t = x_idx - idx_below

  # x, and tensors shaped like x, need to be added to, and selected with
  # (using tf.where) the output y.  This requires appending singletons.
  def _expand_x_fn(tensor):
    # Reshape tensor to tensor.shape + [1] * M.
    extended_shape = tf.concat([
        tf.shape(input=tensor),
        tf.ones_like(tf.shape(input=y_ref)[batch_dims + nd:])
    ],
                               axis=0)
    return tf.reshape(tensor, extended_shape)

  # Now, t.shape = [A1, ..., An, D, nd] + [1] * (rank(y_ref) - nd - batch_dims)
  t = _expand_x_fn(t)
  s = 1 - t

  # Re-insert NaN wherever x was NaN.
  nan_idx = _expand_x_fn(nan_idx)
  t = tf.where(nan_idx,
               tf.fill(tf.shape(input=t), tf.constant(np.nan, dtype)),
               t)

  terms = []
  # Our work above has located x's fractional index inside a cube of above/below
  # indices. The distance to the below indices is t, and to the above indices
  # is s.
  # Drawing lines from x to the cube walls, we get 2**nd smaller cubes. Each
  # term in the result is a product of a reference point, gathered from y_ref,
  # multiplied by a volume.  The volume is that of the cube opposite to the
  # reference point.  E.g. if the reference point is below x in every axis, the
  # volume is that of the cube with corner above x in every axis, s[0]*...*s[nd]
  # We could probably do this with one massive gather, but that would be very
  # unreadable and un-debuggable.  It also would create a large Tensor.
  for zero_ones_list in _binary_count(nd):
    gather_from_y_ref_idx = []
    opposite_volume_t_idx = []
    opposite_volume_s_idx = []
    for k, zero_or_one in enumerate(zero_ones_list):
      if zero_or_one == 0:
        # If the kth iterate has zero_or_one = 0,
        # Will gather from the "below" reference point along axis k.
        gather_from_y_ref_idx.append(idx_below_list[k])
        # Now append the index to gather for computing opposite_volume.
        # This could be done by initializing opposite_volume to 1, then here:
        #  opposite_volume *= tf.gather(s, indices=k, axis=tf.rank(x) - 1)
        # but that puts a gather in the "inner loop."  Better to append the
        # index and do one larger gather down below.
        opposite_volume_s_idx.append(k)
      else:
        gather_from_y_ref_idx.append(idx_above_list[k])
        # Append an index to gather, having the same effect as
        #   opposite_volume *= tf.gather(t, indices=k, axis=tf.rank(x) - 1)
        opposite_volume_t_idx.append(k)

    # Compute opposite_volume (volume of cube opposite the ref point):
    # Recall t.shape = s.shape = [D, nd] + [1, ..., 1]
    # Gather from t and s along the "nd" axis, which is rank(x) - 1.
    ov_axis = tf.rank(x) - 1
    opposite_volume = (
        tf.reduce_prod(
            input_tensor=tf.gather(
                t, indices=tf.cast(opposite_volume_t_idx, dtype=tf.int32),
                axis=ov_axis),
            axis=ov_axis) *
        tf.reduce_prod(
            input_tensor=tf.gather(
                s, indices=tf.cast(opposite_volume_s_idx, dtype=tf.int32),
                axis=ov_axis),
            axis=ov_axis)
    )  # pyformat: disable

    y_ref_pt = tf.gather_nd(
        y_ref,
        tf.stack(gather_from_y_ref_idx, axis=-1),
        batch_dims=batch_dims)

    terms.append(y_ref_pt * opposite_volume)

  y = tf.math.add_n(terms)

  if tf.debugging.is_numeric_tensor(fill_value):
    # Recall x_idx_unclipped.shape = [D, nd],
    # so here we check if it was out of bounds in any of the nd dims.
    # Thus, oob_idx.shape = [D].
    oob_idx = tf.reduce_any(
        input_tensor=(x_idx_unclipped < 0) | (x_idx_unclipped > ny - 1),
        axis=-1)

    # Now, y.shape = [D, B1,...,BM], so we'll have to broadcast oob_idx.

    oob_idx = _expand_x_fn(oob_idx)  # Shape [D, 1,...,1]
    oob_idx |= tf.fill(tf.shape(input=y), False)
    y = tf.where(oob_idx, tf.fill(tf.shape(input=y), fill_value), y)
  return y
def _batch_interp_with_gather_nd ( x , x_ref_min , x_ref_max , y_ref , nd , fill_value , batch_dims ) : dtype = x . dtype # In this function, # x.shape = [A1, ..., An, D, nd], where n = batch_dims # and # y_ref.shape = [A1, ..., An, C1, C2,..., Cnd, B1,...,BM] # y_ref[A1, ..., An, i1,...,ind] is a shape [B1,...,BM] Tensor with the value # at index [i1,...,ind] in the interpolation table. #  and x_ref_max have shapes [A1, ..., An, nd]. # ny[k] is number of y reference points in interp dim k. ny = tf . cast ( tf . shape ( input = y_ref ) [ batch_dims : batch_dims + nd ] , dtype ) # Map [x_ref_min, x_ref_max] to [0, ny - 1]. # This is the (fractional) index of x. # x_idx_unclipped[A1, ..., An, d, k] is the fractional index into dim k of # interpolation table for the dth x value. x_ref_min_expanded = tf . expand_dims ( x_ref_min , axis = - 2 ) x_ref_max_expanded = tf . expand_dims ( x_ref_max , axis = - 2 ) x_idx_unclipped = ( ny - 1 ) * ( x - x_ref_min_expanded ) / ( x_ref_max_expanded - x_ref_min_expanded ) # Wherever x is NaN, x_idx_unclipped will be NaN as well. # Keep track of the nan indices here (so we can impute NaN later). # Also eliminate any NaN indices, since there is not NaN in 32bit. nan_idx = tf . math . is_nan ( x_idx_unclipped ) x_idx_unclipped = tf . where ( nan_idx , tf . zeros_like ( x_idx_unclipped ) , x_idx_unclipped ) # x_idx.shape = [A1, ..., An, D, nd] x_idx = tf . clip_by_value ( x_idx_unclipped , tf . zeros ( ( ) , dtype = dtype ) , ny - 1 ) # Get the index above and below x_idx. # Naively we could set idx_below = floor(x_idx), idx_above = ceil(x_idx), # however, this results in idx_below == idx_above whenever x is on a grid. # This in turn results in y_ref_below == y_ref_above, and then the gradient # at this point is zero.  So here we "jitter" one of idx_below, idx_above, # so that they are at different values.  This jittering does not affect the # interpolated value, but does make the gradient nonzero (unless of course # the y_ref values are the same). idx_below = tf . floor ( x_idx ) idx_above = tf . minimum ( idx_below + 1 , ny - 1 ) idx_below = tf . maximum ( idx_above - 1 , 0 ) # These are the values of y_ref corresponding to above/below indices. # idx_below_int32.shape = x.shape[:-1] + [nd] idx_below_int32 = tf . cast ( idx_below , dtype = tf . int32 ) idx_above_int32 = tf . cast ( idx_above , dtype = tf . int32 ) # idx_below_list is a length nd list of shape x.shape[:-1] int32 tensors. idx_below_list = tf . unstack ( idx_below_int32 , axis = - 1 ) idx_above_list = tf . unstack ( idx_above_int32 , axis = - 1 ) # Use t to get a convex combination of the below/above values. # t.shape = [A1, ..., An, D, nd] t = x_idx - idx_below # x, and tensors shaped like x, need to be added to, and selected with # (using tf.where) the output y.  This requires appending singletons. def _expand_x_fn ( tensor ) : # Reshape tensor to tensor.shape + [1] * M. extended_shape = tf . concat ( [ tf . shape ( input = tensor ) , tf . ones_like ( tf . shape ( input = y_ref ) [ batch_dims + nd : ] ) ] , axis = 0 ) return tf . reshape ( tensor , extended_shape ) # Now, t.shape = [A1, ..., An, D, nd] + [1] * (rank(y_ref) - nd - batch_dims) t = _expand_x_fn ( t ) s = 1 - t # Re-insert NaN wherever x was NaN. nan_idx = _expand_x_fn ( nan_idx ) t = tf . where ( nan_idx , tf . fill ( tf . shape ( input = t ) , tf . constant ( np . nan , dtype ) ) , t ) terms = [ ] # Our work above has located x's fractional index inside a cube of above/below # indices. The distance to the below indices is t, and to the above indices # is s. # Drawing lines from x to the cube walls, we get 2**nd smaller cubes. Each # term in the result is a product of a reference point, gathered from y_ref, # multiplied by a volume.  The volume is that of the cube opposite to the # reference point.  E.g. if the reference point is below x in every axis, the # volume is that of the cube with corner above x in every axis, s[0]*...*s[nd] # We could probably do this with one massive gather, but that would be very # unreadable and un-debuggable.  It also would create a large Tensor. for zero_ones_list in _binary_count ( nd ) : gather_from_y_ref_idx = [ ] opposite_volume_t_idx = [ ] opposite_volume_s_idx = [ ] for k , zero_or_one in enumerate ( zero_ones_list ) : if zero_or_one == 0 : # If the kth iterate has zero_or_one = 0, # Will gather from the "below" reference point along axis k. gather_from_y_ref_idx . append ( idx_below_list [ k ] ) # Now append the index to gather for computing opposite_volume. # This could be done by initializing opposite_volume to 1, then here: #  opposite_volume *= tf.gather(s, indices=k, axis=tf.rank(x) - 1) # but that puts a gather in the "inner loop."  Better to append the # index and do one larger gather down below. opposite_volume_s_idx . append ( k ) else : gather_from_y_ref_idx . append ( idx_above_list [ k ] ) # Append an index to gather, having the same effect as #   opposite_volume *= tf.gather(t, indices=k, axis=tf.rank(x) - 1) opposite_volume_t_idx . append ( k ) # Compute opposite_volume (volume of cube opposite the ref point): # Recall t.shape = s.shape = [D, nd] + [1, ..., 1] # Gather from t and s along the "nd" axis, which is rank(x) - 1. ov_axis = tf . rank ( x ) - 1 opposite_volume = ( tf . reduce_prod ( input_tensor = tf . gather ( t , indices = tf . cast ( opposite_volume_t_idx , dtype = tf . int32 ) , axis = ov_axis ) , axis = ov_axis ) * tf . reduce_prod ( input_tensor = tf . gather ( s , indices = tf . cast ( opposite_volume_s_idx , dtype = tf . int32 ) , axis = ov_axis ) , axis = ov_axis ) ) # pyformat: disable y_ref_pt = tf . gather_nd ( y_ref , tf . stack ( gather_from_y_ref_idx , axis = - 1 ) , batch_dims = batch_dims ) terms . append ( y_ref_pt * opposite_volume ) y = tf . math . add_n ( terms ) if tf . debugging . is_numeric_tensor ( fill_value ) : # Recall x_idx_unclipped.shape = [D, nd], # so here we check if it was out of bounds in any of the nd dims. # Thus, oob_idx.shape = [D]. oob_idx = tf . reduce_any ( input_tensor = ( x_idx_unclipped < 0 ) | ( x_idx_unclipped > ny - 1 ) , axis = - 1 ) # Now, y.shape = [D, B1,...,BM], so we'll have to broadcast oob_idx. oob_idx = _expand_x_fn ( oob_idx ) # Shape [D, 1,...,1] oob_idx |= tf . fill ( tf . shape ( input = y ) , False ) y = tf . where ( oob_idx , tf . fill ( tf . shape ( input = y ) , fill_value ) , y ) return y
N-D interpolation that works with leading batch dims.
N - D interpolation that works with leading batch dims .