refactor

1 files changed, 87 insertions, 0 deletions

diff --git a/octree.hpp b/octree.hpp
new file mode 100644
index 0000000..1be6e0a
--- /dev/null
+++ b/octree.hpp
@@ -0,0 +1,87 @@
+
+#pragma once
+
+#include <array>
+#include <set>
+#include <unordered_map>
+
+#include "spin.hpp"
+#include "vector.hpp"
+
+namespace std {
+template <typename T, size_t N> struct hash<array<T, N>> {
+  typedef array<T, N> argument_type;
+  typedef size_t result_type;
+
+  result_type operator()(const argument_type& a) const {
+    hash<T> hasher;
+    result_type h = 0;
+    for (result_type i = 0; i < N; ++i) {
+      h = h * 31 + hasher(a[i]);
+    }
+    return h;
+  }
+};
+} // namespace std
+
+template <class T, int D, class S> class Octree {
+private:
+  unsigned N;
+  T L;
+  std::unordered_map<std::array<signed, D>, std::set<Spin<T, D, S>*>> data;
+
+public:
+  Octree(T L, unsigned N) {
+    L = L;
+    N = N;
+  }
+
+  std::array<signed, D> ind(Vector<T, D> x) const {
+    std::array<signed, D> ind;
+
+    for (unsigned i = 0; i < D; i++) {
+      ind[i] = (signed)std::floor(N * x(i) / L);
+    }
+
+    return ind;
+  }
+
+  void insert(Spin<T, D, S>* s) { data[ind(s->x)].insert(s); };
+
+  void remove(Spin<T, D, S>* s) {
+    data[ind(s->x)].erase(s);
+    if (data[ind(s->x)].empty()) {
+      data.erase(ind(s->x));
+    }
+  };
+
+  std::set<Spin<T, D, S>*> neighbors(const Vector<T, D>& x) const {
+    std::array<signed, D> i0 = ind(x);
+    std::set<Spin<T, D, S>*> ns;
+
+    nearest_neighbors_of(i0, D + 1, ns);
+
+    return ns;
+  };
+
+  void nearest_neighbors_of(std::array<signed, D> i0, unsigned depth,
+                            std::set<Spin<T, D, S>*>& ns) const {
+    if (depth == 0) {
+      auto it = data.find(i0);
+      if (it != data.end()) {
+        ns.insert(it->second.begin(), it->second.end());
+      }
+    } else {
+      for (signed j : {-1, 0, 1}) {
+        std::array<signed, D> i1 = i0;
+        if (N < 2) {
+          i1[depth - 1] += j;
+        } else {
+          i1[depth - 1] = (N + i1[depth - 1] + j) % N;
+        }
+        nearest_neighbors_of(i1, depth - 1, ns);
+      }
+    }
+  };
+};
+