Arty
hash_grid.h
1 #ifndef HASH_GRID_H
2 #define HASH_GRID_H
3 
4 #include <vector>
5 #include <algorithm>
6 #include <numeric>
7 #include <cstdint>
8 
9 #include "hash.h"
10 
11 class HashGrid {
12 public:
13  HashGrid() {}
14 
15  template <typename PositionFn>
16  void build(PositionFn positions, size_t num_photons, float radius) {
17  radius_sqr = radius * radius;
18  inv_size = 0.5f / radius;
19 
20  // Compute the global bounding box encompassing all the photons
21  bbox = BBox::empty();
22  #pragma omp parallel for reduction(bbox_extend: bbox)
23  for (size_t i = 0; i < num_photons; ++i)
24  bbox = extend(bbox, positions(i));
25 
26  // Slightly enlarge the bounding box of the photons to avoid numerical problems
27  auto extents = bbox.max - bbox.min;
28  bbox.max += extents * 0.001f;
29  bbox.min -= extents * 0.001f;
30 
31  photons.resize(num_photons);
32  cell_counts.resize(1 << (closest_log2(num_photons) + 1));
33  std::fill(cell_counts.begin(), cell_counts.end(), 0);
34 
35  // Count the number of photons per cell
36  #pragma omp parallel for
37  for (size_t i = 0; i < num_photons; i++) {
38  auto h = hash_photon(positions(i));
39  #pragma omp atomic
40  cell_counts[h]++;
41  }
42 
43  // Compute the insertion position for each cell
44  std::partial_sum(cell_counts.begin(), cell_counts.end(), cell_counts.begin());
45  assert(cell_counts.back() == photons.size());
46 
47  // Put the photons in their respective cells
48  #pragma omp parallel for
49  for (size_t i = 0; i < num_photons; i++) {
50  auto h = hash_photon(positions(i));
51  int old_count;
52  #pragma omp atomic capture
53  old_count = --cell_counts[h];
54  photons[old_count] = uint32_t(i);
55  }
56  }
57 
58  template <typename PositionFn, typename InsertFn>
59  void query(const float3& pos, PositionFn positions, InsertFn insert) const {
60  if (!is_inside(bbox, pos)) return;
61 
62  auto p = (pos - bbox.min) * inv_size;
63  int px1 = p.x;
64  int py1 = p.y;
65  int pz1 = p.z;
66  int px2 = px1 + (p.x - px1 > 0.5f ? 1 : -1);
67  int py2 = py1 + (p.y - py1 > 0.5f ? 1 : -1);
68  int pz2 = pz1 + (p.z - pz1 > 0.5f ? 1 : -1);
69 
70  for (int i = 0; i < 8; i++) {
71  auto range = cell_range(i & 1 ? px2 : px1,
72  i & 2 ? py2 : py1,
73  i & 4 ? pz2 : pz1);
74 
75  for (auto j = range.first; j < range.second; j++) {
76  int photon_id = photons[j];
77  auto photon_pos = positions(photon_id);
78  auto d = lensqr(pos - photon_pos);
79  if (d < radius_sqr)
80  insert(photon_id, d);
81  }
82  }
83  }
84 
85 private:
86  std::pair<size_t, size_t> cell_range(uint32_t x, uint32_t y, uint32_t z) const {
87  auto h = hash_cell(x, y, z);
88  return std::make_pair(cell_counts[h], h == cell_counts.size() - 1 ? photons.size() : cell_counts[h + 1]);
89  }
90 
91  uint32_t hash_cell(uint32_t x, uint32_t y, uint32_t z) const {
92  uint32_t h = bernstein_hash(bernstein_init(), x);
93  h = bernstein_hash(h, y);
94  h = bernstein_hash(h, z);
95  return h & (cell_counts.size() - 1);
96  }
97 
98  uint32_t hash_photon(const float3& pos) const {
99  auto p = (pos - bbox.min) * inv_size;
100  return hash_cell(p.x, p.y, p.z);
101  }
102 
103  std::vector<uint32_t> photons;
104  std::vector<uint32_t> cell_counts;
105  BBox bbox;
106  float inv_size;
107  float radius_sqr;
108 };
109 
110 #endif // HASH_GRID_H
Bounding box represented by its two extreme points.
Definition: bbox.h:9
Definition: float3.h:10
Definition: hash_grid.h:11