Method predict_density

internal/models.py:211–235

Helper function to output density.

(rng, means, covs, coord_freq_mask=None)

    dense_layer = functools.partial(nn.Dense, kernel_init=self.weight_init)

    def predict_density(rng, means, covs, coord_freq_mask=None):
      """Helper function to output density."""
      # Encode input positions
      inputs = mip.integrated_pos_enc(
          (means, covs), self.min_deg_point, self.max_deg_point)
      ## ---- add freq reg mask ----- ##
      if coord_freq_mask is not None:
        inputs = inputs * coord_freq_mask
      ## ---------------------------- ##
      # Evaluate network to output density
      x = inputs
      for i in range(self.net_depth):
        x = dense_layer(self.net_width)(x)
        x = self.net_activation(x)
        if i % self.skip_layer == 0 and i > 0:
          x = jnp.concatenate([x, inputs], axis=-1)
      raw_density = dense_layer(1)(x)[Ellipsis, 0]  # Hardcoded to a single channel.
      # Add noise to regularize the density predictions if needed.
      if (rng is not None) and (self.density_noise > 0):
        key, rng = random.split(rng)
        raw_density += self.density_noise * random.normal(
            key, raw_density.shape, dtype=raw_density.dtype)
      # Apply bias and activation to raw density
      density = self.density_activation(raw_density + self.density_bias)
      return density, x

    means, covs = samples
    ## ---- split freq reg mask to coordinates and viewdirs ----- ##