Function sample_intervals

utils/stepfun.py:251–294

Sample *intervals* (rather than points) from a step function. Args: rand: random number generator (or None for `linspace` sampling). t: [..., num_bins + 1], bin endpoint coordinates (must be sorted) w_logits: [..., num_bins], logits corresponding to bin weights num_samples: int, t

(rand,
                     t,
                     w_logits,
                     num_samples,
                     single_jitter=False,
                     domain=(-torch.inf, torch.inf))

def sample_intervals(rand,
                     t,
                     w_logits,
                     num_samples,
                     single_jitter=False,
                     domain=(-torch.inf, torch.inf)):
    """Sample *intervals* (rather than points) from a step function.

  Args:
    rand: random number generator (or None for `linspace` sampling).
    t: [..., num_bins + 1], bin endpoint coordinates (must be sorted)
    w_logits: [..., num_bins], logits corresponding to bin weights
    num_samples: int, the number of intervals to sample.
    single_jitter: bool, if True, jitter every sample along each ray by the same
      amount in the inverse CDF. Otherwise, jitter each sample independently.
    domain: (minval, maxval), the range of valid values for `t`.

  Returns:
    t_samples: [batch_size, num_samples].
  """
    if num_samples <= 1:
        raise ValueError(f'num_samples must be > 1, is {num_samples}.')

    # Sample a set of points from the step function.
    centers = sample(
        rand,
        t,
        w_logits,
        num_samples,
        single_jitter,
        deterministic_center=True)

    # The intervals we return will span the midpoints of each adjacent sample.
    mid = (centers[..., 1:] + centers[..., :-1]) / 2

    # Each first/last fencepost is the reflection of the first/last midpoint
    # around the first/last sampled center. We clamp to the limits of the input
    # domain, provided by the caller.
    minval, maxval = domain
    first = (2 * centers[..., :1] - mid[..., :1]).clamp_min(minval)
    last = (2 * centers[..., -1:] - mid[..., -1:]).clamp_max(maxval)

    t_samples = torch.cat([first, mid, last], dim=-1)
    return t_samples


def lossfun_distortion(t, w):