Probability Density Function (PDF) Tree data structure

gistfile1.txt

The PdfTree is a hierarchical data structure that stores values and their associated probabilities, organized in a way that resembles a binary tree or mip-map pyramid. It allows for efficient random selection of values proportional to their probabilities (e.g., for Monte Carlo sampling in rendering or procedural generation). The tree is built as a "summed probability tree," where each level aggregates probabilities from the level below it, enabling fast traversal for sampling.

The tree is laid out linearly in probabilities, with each level concatenated. For example:
Mip 0: Original probabilities [p0, p1, p2, p3]
Mip 1: Pairwise sums [p0+p1, p2+p3]
Mip 2: Sum of mip 1 [p0+p1+p2+p3]

Mip Level Calculation:
Computes the total number of elements across all mip levels.

Each level halves the size of the previous one (rounded up with alignUp), forming a binary tree-like structure.
Example: If capacity = 4:
Mip 0: 4 elements
Mip 1: 2 elements (4/2)
Mip 2: 1 element (2/2)
Total = 7 elements, 3 mip levels.

Repopulate the tree with values and probabilities from an input array.
Step 1: Fill Values:
Copies the value field from elements into pTree->values.

Step 2: Clear and Fill Mip 0:
Resets the tree’s arrays.
Copies probabilities into mip0Probabilities (the base level), ensuring no negative probabilities with max(0.0f, ...).
Counts active elements (probability > 0).

Step 3: Build Higher Mip Levels:
Iteratively constructs each mip level by summing pairs of probabilities from the level below.

Example: For [1.0, 2.0, 3.0, 4.0]:
Mip 0: [1.0, 2.0, 3.0, 4.0]
Mip 1: [3.0, 7.0] (1+2, 3+4)
Mip 2: [10.0] (3+7)

Stores offsets and sizes for each level.

Setting an element’s probability to 0 and updates the tree.
Sets probabilities[index] (in mip 0) to 0.
Propagates the change upward:
For each mip level, recomputes the sum of the two child probabilities and updates the parent.

Example: If index = 1 in [1.0, 2.0, 3.0, 4.0]:
Mip 0: [1.0, 0.0, 3.0, 4.0]
Mip 1: [1.0, 7.0] (1+0, 3+4)
Mip 2: [8.0] (1+7)

Sampleing a value randomly based on its probability and deactivates it.
Generates a random float between 0 and the total probability (top of the tree).
Traverses the tree from top to bottom:
At each level, compares the random value to the left child’s probability.
If greater, moves to the right child and subtracts the left probability; otherwise, moves left.

Example: Total = 10.0, random = 4.5:
Mip 2: [10.0] (start)
Mip 1: [3.0, 7.0] (4.5 > 3.0, go right, 4.5 - 3.0 = 1.5)
Mip 0: [3.0, 4.0] (1.5 < 3.0, go left, pick index 0)

This structure is ideal for importance sampling in rendering (e.g., picking light sources or surface points proportional to their contribution) or procedural generation (e.g., spawning objects based on weighted probabilities). The hierarchical design ensures O(log n) sampling time, and deactivation allows for sampling without replacement.

pdf.c

#include <assert.h>
#include <stdint.h>
#include <stdlib.h>
#include <string.h>

#define PDF_MAX_CAP       1024
#define PDF_MAX_ELM_CNT   2048
#define PDF_MAX_MIP_LVL   11

struct pdf_tree_elm {
  int val;
  float prob; 
}
struct pdf_tree {
  int vals[PDF_MAX_CAP];
  int act_cnt;
  float prob[PDF_MAX_ELM_CNT];
  int mip_off[PDF_MAX_MIP_LVL];
  int mip_siz[PDF_MAX_MIP_LVL];
  int cap;
  int mip_cnt;
};
static inline int
align_up(int v, int a) {
    return (v + a - 1) & ~(a - 1);
}
// Initialize the PdfTree
static void
pdf_init(struct pdf_tree *pdf, int cap) {
  assert(cap <= MAX_CAPACITY);
  pdf->act_cnt = 0;
  pdf->cap = cap;

  // Calculate mip levels
  int mip_cnt = 1;
  int mip_elm_cnt = cap;
  int elm_cnt = mip_elm_cnt;
  while (mip_elm_cnt > 1) {
    mip_elm_cnt = align_up(mip_elm_cnt, 2) / 2;
    elm_cnt += mip_elm_cnt;
    mip_cnt++;
  }
  assert(elm_cnt <= MAX_TOTAL_ELEMENTS);
  assert(mip_cnt <= MAX_MIP_LEVELS);
  pdf->mip_cnt = mip_cnt;

  memset(pdf->vals, 0, sizeof(pdf->vals));
  memset(pdf->prob, 0, sizeof(pdf->prob));
  memset(pdf->mip_off, 0, sizeof(pdf->mip_off));
  memset(pdf->mip_siz, 0, sizeof(pdf->mip_siz));
}
static void
pdf_fill(struct pdf_tree* pdf, const struct pdf_tree_elm *elm, int cnt) {
  assert(cnt <= pdf->cap);
  for (int i = 0; i < cnt; i++) {
      pdf->vals[i] = elm[i].value;
  }
  // Reset prob and mip data
  memset(pdf->prob, 0, sizeof(pdf->prob));
  memset(pdf->mip_off, 0, sizeof(pdf->mip_off));
  memset(pdf->mip_siz, 0, sizeof(pdf->mip_siz));

  // fill mip 0 (base level)
  pdf->act_cnt = 0;
  for (int i = 0; i < cnt; i++) {
    float prob = elm[i].prob > 0.0f ? elm[i].prob : 0.0f;
    pdf->prob[i] = prob;
    if (prob > 0.0f) {
      pdf->act_cnt++;
    }
  }
  pdf->mip_off[0] = 0;
  pdf->mip_siz[0] = cnt;

  // Build higher mip levels
  int mip_lvl = 1;
  int mip_elm_cnt = cnt;

  while (mip_elm_cnt > 1) {
    mip_elm_cnt = align_up(mip_elm_cnt, 2) / 2;
    pdf->mip_off[mip_lvl] = pdf->mip_off[mip_lvl - 1] + pdf->mip_siz[mip_lvl - 1];
    pdf->mip_siz[mip_lvl] = mip_elm_cnt;

    for (int e = 0; e < mip_elm_cnt; e++) {
      int child_idx = 2 * e;
      int child_prob_idx = pdf->mip_off[mip_lvl - 1] + child_idx;
      float prob_sum = pdf->prob[child_prob_idx];
      if (child_idx + 1 < pdf->mip_siz[mip_lvl - 1]) {
        prob_sum += pdf->prob[child_prob_idx + 1];
      }
      pdf->prob[pdf->mip_off[mip_lvl] + e] = prob_sum;
    }
    mip_lvl++;
  }
}
static void
pdf_del(struct pdf_tree* pdf, int index) {
  assert(pdf->act_cnt > 0);
  pdf->prob[index] = 0.0f;

  // Update higher mip levels
  int mip_lvl = 1;
  int e = index / 2;
  while (mip_lvl < pdf->mip_cnt) {
    int child_idx = 2 * e;
    int child_prob_idx = pdf->mip_off[mip_lvl - 1] + child_idx;
    float probabilitySum = pdf->prob[child_prob_idx];
    if (child_idx + 1 < pdf->mip_siz[mip_lvl - 1]) {
        probabilitySum += pdf->prob[child_prob_idx + 1];
    }
    pdf->prob[pdf->mip_off[mip_lvl] + e] = probabilitySum;

    mip_lvl++;
    e /= 2;
  }
  pdf->act_cnt--;
}
static inline float
getUniformFloat(float min, float max) {
  return min + (float)rand() / RAND_MAX * (max - min);
}
static int
pdf_gen(struct pdf_tree *tree) {
  assert(pdf->mip_cnt > 0);
  assert(pdf->act_cnt > 0);
  assert(pdf->mip_siz[pdf->mip_cnt - 1] == 1);

  float tot_prob = pdf->prob[pdf->mip_off[pdf->mip_cnt - 1]];
  float randomValue = getUniformFloat(0.0f, tot_prob);

  int e = 0;
  int mip_lvl = pdf->mip_cnt - 1;
  float cur_prob = randomValue;
  while (mip_lvl > 0) {
    mip_lvl--;

    int mip_off = pdf->mip_off[mip_lvl];
    int mip_siz = pdf->mip_siz[mip_lvl];
    int nxt_mip_e = 2 * e;
    assert(nxt_mip_e < mip_siz);

    float mip_prob = pdf->prob[mip_off + nxt_mip_e];
    if (cur_prob > mip_prob && nxt_mip_e + 1 < mip_siz) {
      nxt_mip_e++;
      cur_prob -= mip_prob;
    }
    e = nxt_mip_e;
  }
  pdf_del(pdf, e);
  return pdf->vals[e];
}

// Example usage
int main() {
  struct pdf_tree tree;
  pdf_init(&tree, 4);

  Element elm[] = {
    {0, 1.0f},
    {1, 2.0f},
    {2, 3.0f},
    {3, 4.0f}
  };
  pdf_fill(&tree, elm, 4);

  srand(1234);  // Seed RNG for reproducibility
  for (int i = 0; i < 4; i++) {
    int value = pdf_gen(&tree);
    printf("Chosen value: %u\n", value);
  }
  return 0;
}