1 files changed, 45 insertions, 15 deletions
diff --git a/lib/src/graph.cpp b/lib/src/graph.cpp
index 4043413..c6f0ba6 100644
--- a/lib/src/graph.cpp
+++ b/lib/src/graph.cpp
@@ -1,25 +1,55 @@
 
 #include <wolff/graph.hpp>
 
-graph_t::graph_t(D_t D, L_t L) {
-  nv = pow(L, D);
-  ne = D * nv;
+graph_t::graph_t(D_t D, L_t L, lattice_t lat) {
+  switch (lat) {
+    case SQUARE_LATTICE: {
+        nv = pow(L, D);
+        ne = D * nv;
 
-  v_adj.resize(nv);
-  coordinate.resize(nv);
+        v_adj.resize(nv);
+        coordinate.resize(nv);
 
-  for (std::vector<v_t> v_adj_i : v_adj) {
-    v_adj_i.reserve(2 * D);
-  }
+        for (std::vector<v_t> v_adj_i : v_adj) {
+          v_adj_i.reserve(2 * D);
+        }
 
-  for (v_t i = 0; i < nv; i++) {
-    coordinate[i].resize(D);
-    for (D_t j = 0; j < D; j++) {
-      coordinate[i][j] = (i / (v_t)pow(L, D - j - 1)) % L;
+        for (v_t i = 0; i < nv; i++) {
+          coordinate[i].resize(D);
+          for (D_t j = 0; j < D; j++) {
+            coordinate[i][j] = (i / (v_t)pow(L, D - j - 1)) % L;
+
+            v_adj[i].push_back(pow(L, j + 1) * (i / ((v_t)pow(L, j + 1))) + fmod(i + pow(L, j), pow(L, j + 1)));
+            v_adj[i].push_back(pow(L, j + 1) * (i / ((v_t)pow(L, j + 1))) + fmod(pow(L, j+1) + i - pow(L, j), pow(L, j + 1)));
+          }
+        }
+        break;
+      }
+    case DIAGONAL_LATTICE: {
+        nv = D * pow(L, D);
+        ne = D * nv;
+
+        v_adj.resize(nv);
+        coordinate.resize(nv);
+
+        for (std::vector<v_t> v_adj_i : v_adj) {
+          v_adj_i.reserve(4 * (D - 1));
+        }
+
+        for (D_t i = 0; i < D; i++) {
+          v_t sb = i * pow(L, D);
+
+          for (v_t j = 0; j < pow(L, D); j++) {
+            v_t vc = sb + j;
 
-      v_adj[i].push_back(pow(L, j + 1) * (i / ((v_t)pow(L, j + 1))) + fmod(i + pow(L, j), pow(L, j + 1)));
-      v_adj[i].push_back(pow(L, j + 1) * (i / ((v_t)pow(L, j + 1))) + fmod(pow(L, j+1) + i - pow(L, j), pow(L, j + 1)));
-    }
+            v_adj[vc].push_back(((i + 1) % D) * pow(L, D) + j);
+            v_adj[vc].push_back(((i + 1) % D) * pow(L, D) + pow(L, D - 1) * (j / (v_t)pow(L, D - 1)) + (j + 1 - 2 * (i % 2)) % L);
+            v_adj[vc].push_back(((i + 1) % D) * pow(L, D) + pow(L, D - 1) * ((L + (j/ (v_t)pow(L, D - 1)) - 1 + 2 * (i % 2)) % L) + (j - i) % L);
+            v_adj[vc].push_back(((i + 1) % D) * pow(L, D) + pow(L, D - 1) * ((L + (j/ (v_t)pow(L, D - 1)) - 1 + 2 * (i % 2)) % L) + (j + 1 - i) % L);
+          }
+        }
+        break;
+      }
   }
 }