From 39aa2832b3d8e1a3e1dd2efb0bf9629d8190ed39 Mon Sep 17 00:00:00 2001
From: Adrian Kummerlaender
Date: Tue, 29 Oct 2019 20:41:33 +0100
Subject: Adapt to upstream changes

---
 simulation.py | 3 ++-
 1 file changed, 2 insertions(+), 1 deletion(-)

(limited to 'simulation.py')

diff --git a/simulation.py b/simulation.py
index e0ce3ba..7625609 100644
--- a/simulation.py
+++ b/simulation.py
@@ -41,6 +41,7 @@ class CellList:
 class Lattice:
     def __init__(self, geometry, kernel_src, descriptor, platform = 0, precision = 'single'):
         self.geometry = geometry
+        self.descriptor = descriptor
 
         self.float_type = {
             'single': (numpy.float32, 'float'),
@@ -90,7 +91,7 @@ class Lattice:
         self.queue.finish()
 
     def get_moments(self):
-        moments = numpy.ndarray(shape=(3, self.memory.volume), dtype=self.float_type[0])
+        moments = numpy.ndarray(shape=(self.memory.volume*(self.descriptor.d+1),1), dtype=self.float_type[0])
 
         if self.tick:
             self.program.collect_moments(
-- 
cgit v1.2.3