1 files changed, 59 insertions, 49 deletions
diff --git a/uniform.cpp b/uniform.cpp
index 23cb7f8..2c123ff 100644
--- a/uniform.cpp
+++ b/uniform.cpp
@@ -7,49 +7,50 @@
 using Real = long double;
 
 class AztecDiamond {
+
 public:
   typedef struct Edge {
     std::stack<Real> weights;
     Real probability = 0;
   } Edge;
 
-  using Face = std::array<std::reference_wrapper<Edge>, 4>;
+private:
+  std::tuple<Edge&, Edge&, Edge&, Edge&> face(unsigned i, unsigned j) {
+    unsigned x0 = n - i;
+    unsigned xx = 2 * (j % i);
+    unsigned yy = 2 * (j / i);
 
-  std::vector<Edge> edges;
-  std::vector<std::vector<Face>> lattices;
+    Edge& e1 = edges[2 * n * (x0 + yy) + x0 + xx];
+    Edge& e2 = edges[2 * n * (x0 + yy) + x0 + xx + 1];
+    Edge& e3 = edges[2 * n * (x0 + yy + 1) + x0 + xx];
+    Edge& e4 = edges[2 * n * (x0 + yy + 1) + x0 + xx + 1];
 
-  AztecDiamond(unsigned n) : edges(pow(2 * n, 2)), lattices(n) {
-    for (unsigned i = 1; i <= n; i++) {
-      lattices[n - i].reserve(pow(i, 2));
-      unsigned x0 = n - i;
-      for (unsigned j = 0; j < pow(i, 2); j++) {
-        unsigned x = 2 * (j % i);
-        unsigned y = 2 * (j / i);
-        lattices[n - i].push_back({
-              edges[2 * n * (x0 + y) + x0 + x],
-              edges[2 * n * (x0 + y) + x0 + x + 1],
-              edges[2 * n * (x0 + y + 1) + x0 + x],
-              edges[2 * n * (x0 + y + 1) + x0 + x + 1]
-              });
-      }
-    }
+    return {e1, e2, e3, e4};
   }
 
+public:
+  unsigned n;
+  std::vector<Edge> edges;
+
+  AztecDiamond(unsigned n) : edges(pow(2 * n, 2)), n(n) {}
+
   void computeWeights() {
-    for (std::vector<Face>& faces : lattices) {
+    for (unsigned i = n; i > 0; i--) {
 #pragma omp parallel for
-      for (Face& f : faces) {
-        Real w = f[0].get().weights.top();
-        Real x = f[1].get().weights.top();
-        Real y = f[2].get().weights.top();
-        Real z = f[3].get().weights.top();
+      for (unsigned j = 0; j < pow(i, 2); j++) {
+        auto [e1, e2, e3, e4] = face(i, j);
+
+        Real w = e1.weights.top();
+        Real x = e2.weights.top();
+        Real y = e3.weights.top();
+        Real z = e4.weights.top();
 
         Real cellFactor = w * z + x * y;
 
-        f[0].get().weights.push(z / cellFactor);
-        f[1].get().weights.push(y / cellFactor);
-        f[2].get().weights.push(x / cellFactor);
-        f[3].get().weights.push(w / cellFactor);
+        e1.weights.push(z / cellFactor);
+        e2.weights.push(y / cellFactor);
+        e3.weights.push(x / cellFactor);
+        e4.weights.push(w / cellFactor);
       }
     }
 
@@ -59,33 +60,42 @@ public:
     }
   }
 
-  void computeProbabilities() { // destroys *all* weights
-    for (auto it = lattices.rbegin(); it != lattices.rend(); it++) {
+  Real computeProbabilities() { // destroys *all* weights
+    Real partitionFunction = 1;
+
+    for (unsigned i = 1; i <= n; i++) {
 #pragma omp parallel for
-      for (Face& f : *it) {
-        Real p = f[0].get().probability;
-        Real q = f[1].get().probability;
-        Real r = f[2].get().probability;
-        Real s = f[3].get().probability;
+      for (unsigned j = 0; j < pow(i, 2); j++) {
+        auto [e1, e2, e3, e4] = face(i, j);
+
+        Real p = e1.probability;
+        Real q = e2.probability;
+        Real r = e3.probability;
+        Real s = e4.probability;
 
-        Real w = f[0].get().weights.top();
-        Real x = f[1].get().weights.top();
-        Real y = f[2].get().weights.top();
-        Real z = f[3].get().weights.top();
+        Real w = e1.weights.top();
+        Real x = e2.weights.top();
+        Real y = e3.weights.top();
+        Real z = e4.weights.top();
 
         Real cellFactor = w * z + x * y;
         Real deficit = 1 - p - q - r - s;
 
-        f[0].get().probability = s + deficit * w * z / cellFactor;
-        f[1].get().probability = r + deficit * x * y / cellFactor;
-        f[2].get().probability = q + deficit * x * y / cellFactor;
-        f[3].get().probability = p + deficit * w * z / cellFactor;
+        e1.probability = s + deficit * w * z / cellFactor;
+        e2.probability = r + deficit * x * y / cellFactor;
+        e3.probability = q + deficit * x * y / cellFactor;
+        e4.probability = p + deficit * w * z / cellFactor;
 
-        for (Edge& e : f) {
-          e.weights.pop();
-        }
+        e1.weights.pop();
+        e2.weights.pop();
+        e3.weights.pop();
+        e4.weights.pop();
+
+        partitionFunction *= cellFactor;
       }
     }
+
+    return partitionFunction;
   }
 };
 
@@ -138,9 +148,9 @@ int main(int argc, char* argv[]) {
   std::vector<Real> data_y(G.vertices.size());
 
   for (unsigned i = 0; i < G.edges.size(); i++) {
-    const Edge& e = G.edges[i];
-    const Vertex& vt = e.halfedges[0].getTail();
-    const Vertex& vh =  e.halfedges[0].getHead();
+    const Graph::Edge& e = G.edges[i];
+    const Graph::Vertex& vt = e.halfedges[0].getTail();
+    const Graph::Vertex& vh =  e.halfedges[0].getHead();
     data_x[vt.index] += avgProbabilities[i] * (vt.coordinate[0] - vh.coordinate[0]);
     data_y[vt.index] += avgProbabilities[i] * (vt.coordinate[1] - vh.coordinate[1]);
     data_x[vh.index] += avgProbabilities[i] * (vt.coordinate[0] - vh.coordinate[0]);