Fix returned array shape for accumulator

src/core/accumulators/MeanVarianceCalculator.hpp

@@ -40,6 +40,7 @@ class MeanVarianceCalculator : public AccumulatorBase {
      via the interface. */
   std::string get_internal_state() const;
   void set_internal_state(std::string const &);
+  std::vector<std::size_t> shape() const { return m_obs->shape(); }
 
 private:
   std::shared_ptr<Observables::Observable> m_obs;

src/core/accumulators/TimeSeries.hpp

@@ -42,6 +42,12 @@ class TimeSeries : public AccumulatorBase {
   void set_internal_state(std::string const &);
 
   const std::vector<std::vector<double>> &time_series() const { return m_data; }
+  std::vector<std::size_t> shape() const {
+    std::vector<std::size_t> shape{m_data.size()};
+    auto obs_shape = m_obs->shape();
+    shape.insert(shape.end(), obs_shape.begin(), obs_shape.end());
+    return shape;
+  }
   void clear() { m_data.clear(); }
 
 private:

src/python/espressomd/accumulators.py

@@ -44,11 +44,16 @@ class MeanVarianceCalculator(ScriptInterfaceHelper):
     _so_name = "Accumulators::MeanVarianceCalculator"
     _so_bind_methods = (
         "update",
-        "get_mean",
-        "get_variance"
+        "shape",
     )
     _so_creation_policy = "LOCAL"
 
+    def get_mean(self):
+        return np.array(self.call_method("get_mean")).reshape(self.shape())
+
+    def get_variance(self):
+        return np.array(self.call_method("get_variance")).reshape(self.shape())
+
 
 @script_interface_register
 class TimeSeries(ScriptInterfaceHelper):
@@ -75,11 +80,14 @@ class TimeSeries(ScriptInterfaceHelper):
     _so_name = "Accumulators::TimeSeries"
     _so_bind_methods = (
         "update",
-        "time_series",
+        "shape",
         "clear"
     )
     _so_creation_policy = "LOCAL"
 
+    def time_series(self):
+        return np.array(self.call_method("time_series")).reshape(self.shape())
+
 
 @script_interface_register
 class Correlator(ScriptInterfaceHelper):

src/script_interface/accumulators/MeanVarianceCalculator.hpp

@@ -69,6 +69,10 @@ class MeanVarianceCalculator : public AccumulatorBase {
 
     if (method == "get_variance")
       return mean_variance_calculator()->get_variance();
+    if (method == "shape") {
+      auto const shape = m_accumulator->shape();
+      return std::vector<int>{shape.begin(), shape.end()};
+    }
     return {};
   }
 

src/script_interface/accumulators/TimeSeries.hpp

@@ -63,6 +63,10 @@ class TimeSeries : public AccumulatorBase {
       m_accumulator->clear();
     }
 
+    if (method == "shape") {
+      auto const shape = m_accumulator->shape();
+      return std::vector<int>{shape.begin(), shape.end()};
+    }
     return {};
   }
 

testsuite/python/accumulator.py

@@ -29,6 +29,9 @@
 import espressomd.accumulators
 
 
+N_PART = 4
+
+
 class AccumulatorTest(ut.TestCase):
 
     """
@@ -41,24 +44,27 @@ def setUp(self):
         self.system = espressomd.System(box_l=[10.0] * 3)
         self.system.cell_system.skin = 0.4
         self.system.time_step = 0.01
-        self.system.part.add(id=0, pos=[0.0, 0.0, 0.0])
+        self.system.part.add(pos=np.zeros((N_PART, 3)))
         self.system.integrator.run(steps=0)
-        self.pos_obs = espressomd.observables.ParticlePositions(ids=(0,))
+        self.pos_obs = espressomd.observables.ParticlePositions(
+            ids=range(N_PART))
         self.pos_obs_acc = espressomd.accumulators.MeanVarianceCalculator(
             obs=self.pos_obs)
         self.system.auto_update_accumulators.add(self.pos_obs_acc)
-        self.positions = np.copy(self.system.box_l * np.random.rand(10, 3))
+        self.positions = np.copy(self.system.box_l) * \
+            np.random.random((10, N_PART, 3))
 
     def test_accumulator(self):
         """Check that accumulator results are the same as the respective numpy result.
 
         """
-        for i in range(self.positions.shape[0]):
-            self.system.part[0].pos = self.positions[i]
+        for pos in self.positions:
+            self.system.part[:].pos = pos
             self.system.integrator.run(1)
         self.assertEqual(self.pos_obs, self.pos_obs_acc.get_params()['obs'])
+        acc_mean = self.pos_obs_acc.get_mean()
         np.testing.assert_allclose(
-            self.pos_obs_acc.get_mean(),
+            acc_mean,
             np.mean(self.positions, axis=0), atol=1e-4)
         np.testing.assert_allclose(
             self.pos_obs_acc.get_variance(),

testsuite/python/brownian_dynamics.py

@@ -273,29 +273,33 @@ def test_08__noise_correlation(self):
 
         # test translational noise correlation
         vel = np.array(vel_series.time_series())
-        for i in range(6):
-            for j in range(i, 6):
-                corrcoef = np.dot(vel[:, i], vel[:, j]) / steps / kT
-                if i == j:
-                    self.assertAlmostEqual(corrcoef, 1.0, delta=0.04)
-                else:
-                    self.assertLessEqual(np.abs(corrcoef), 0.04)
+        for ind in range(2):
+            for i in range(3):
+                for j in range(i, 3):
+                    corrcoef = np.dot(
+                        vel[:, ind, i], vel[:, ind, j]) / steps / kT
+                    if i == j:
+                        self.assertAlmostEqual(corrcoef, 1.0, delta=0.04)
+                    else:
+                        self.assertLessEqual(np.abs(corrcoef), 0.04)
 
         # test rotational noise correlation
         if espressomd.has_features("ROTATION"):
             omega = np.array(omega_series.time_series())
-            for i in range(6):
-                for j in range(6):
-                    corrcoef = np.dot(omega[:, i], omega[:, j]) / steps / kT
-                    if i == j:
-                        self.assertAlmostEqual(corrcoef, 1.0, delta=0.04)
-                    else:
+            for ind in range(2):
+                for i in range(3):
+                    for j in range(3):
+                        corrcoef = np.dot(
+                            omega[:, ind, i], omega[:, ind, j]) / steps / kT
+                        if i == j:
+                            self.assertAlmostEqual(corrcoef, 1.0, delta=0.04)
+                        else:
+                            self.assertLessEqual(np.abs(corrcoef), 0.04)
+                        # translational and angular velocities should be
+                        # independent
+                        corrcoef = np.dot(
+                            vel[:, ind, i], omega[:, ind, j]) / steps / kT
                         self.assertLessEqual(np.abs(corrcoef), 0.04)
-            # translational and angular velocities should be independent
-            for i in range(6):
-                for j in range(6):
-                    corrcoef = np.dot(vel[:, i], omega[:, j]) / steps / kT
-                    self.assertLessEqual(np.abs(corrcoef), 0.04)
 
 
 if __name__ == "__main__":

testsuite/python/langevin_thermostat.py

@@ -522,29 +522,33 @@ def test_08__noise_correlation(self):
 
         # test translational noise correlation
         vel = np.array(vel_series.time_series())
-        for i in range(6):
-            for j in range(i, 6):
-                corrcoef = np.dot(vel[:, i], vel[:, j]) / steps / kT
-                if i == j:
-                    self.assertAlmostEqual(corrcoef, 1.0, delta=0.04)
-                else:
-                    self.assertLessEqual(np.abs(corrcoef), 0.04)
+        for ind in range(2):
+            for i in range(3):
+                for j in range(i, 3):
+                    corrcoef = np.dot(
+                        vel[:, ind, i], vel[:, ind, j]) / steps / kT
+                    if i == j:
+                        self.assertAlmostEqual(corrcoef, 1.0, delta=0.04)
+                    else:
+                        self.assertLessEqual(np.abs(corrcoef), 0.04)
 
         # test rotational noise correlation
         if espressomd.has_features("ROTATION"):
             omega = np.array(omega_series.time_series())
-            for i in range(6):
-                for j in range(6):
-                    corrcoef = np.dot(omega[:, i], omega[:, j]) / steps / kT
-                    if i == j:
-                        self.assertAlmostEqual(corrcoef, 1.0, delta=0.04)
-                    else:
+            for ind in range(2):
+                for i in range(3):
+                    for j in range(3):
+                        corrcoef = np.dot(
+                            omega[:, ind, i], omega[:, ind, j]) / steps / kT
+                        if i == j:
+                            self.assertAlmostEqual(corrcoef, 1.0, delta=0.04)
+                        else:
+                            self.assertLessEqual(np.abs(corrcoef), 0.04)
+                        # translational and angular velocities should be
+                        # independent
+                        corrcoef = np.dot(
+                            vel[:, ind, i], omega[:, ind, j]) / steps / kT
                         self.assertLessEqual(np.abs(corrcoef), 0.04)
-            # translational and angular velocities should be independent
-            for i in range(6):
-                for j in range(6):
-                    corrcoef = np.dot(vel[:, i], omega[:, j]) / steps / kT
-                    self.assertLessEqual(np.abs(corrcoef), 0.04)
 
 
 if __name__ == "__main__":

testsuite/python/test_checkpoint.py

@@ -270,12 +270,12 @@ def test_virtual_sites(self):
 
     def test_mean_variance_calculator(self):
         np.testing.assert_array_equal(
-            acc.get_mean(), np.array([1.0, 1.5, 2.0, 1.0, 1.0, 2.0]))
+            acc.get_mean(), np.array([[1.0, 1.5, 2.0], [1.0, 1.0, 2.0]]))
         np.testing.assert_array_equal(
-            acc.get_variance(), np.array([0., 0.5, 2., 0., 0., 0.]))
+            acc.get_variance(), np.array([[0., 0.5, 2.], [0., 0., 0.]]))
         np.testing.assert_array_equal(
             system.auto_update_accumulators[0].get_variance(),
-            np.array([0., 0.5, 2., 0., 0., 0.]))
+            np.array([[0., 0.5, 2.], [0., 0., 0.]]))
 
     @utx.skipIfMissingFeatures('P3M')
     @ut.skipIf('P3M.CPU' not in modes,

testsuite/python/time_series.py

@@ -58,7 +58,7 @@ def test_time_series(self):
 
         for result, expected in zip(time_series.time_series(), positions):
             np.testing.assert_array_equal(
-                np.array(result).reshape((N_PART, 3)), expected)
+                result, expected)
 
         time_series.clear()
         self.assertEqual(len(time_series.time_series()), 0)