Lb nodes generator

src/core/grid_based_algorithms/lb_interface.cpp

@@ -1107,22 +1107,23 @@ void lb_lbfluid_load_checkpoint(const std::string &filename, int binary) {
   }
 }
 
-bool lb_lbnode_is_index_valid(const Utils::Vector3i &ind) {
-  auto within_bounds = [](const Utils::Vector3i &ind,
-                          const Utils::Vector3i &limits) {
-    return ind < limits && ind >= Utils::Vector3i{};
-  };
+Utils::Vector3i lb_lbfluid_get_shape() {
   if (lattice_switch == ActiveLB::GPU) {
 #ifdef CUDA
-    return within_bounds(ind, {static_cast<int>(lbpar_gpu.dim_x),
-                               static_cast<int>(lbpar_gpu.dim_y),
-                               static_cast<int>(lbpar_gpu.dim_z)});
+    return {static_cast<int>(lbpar_gpu.dim_x),
+            static_cast<int>(lbpar_gpu.dim_y),
+            static_cast<int>(lbpar_gpu.dim_z)};
 #endif
   }
   if (lattice_switch == ActiveLB::CPU) {
-    return within_bounds(ind, lblattice.global_grid);
+    return lblattice.global_grid;
   }
-  return false;
+  throw std::runtime_error("No LB active");
+}
+
+bool lb_lbnode_is_index_valid(const Utils::Vector3i &ind) {
+  auto const limit = lb_lbfluid_get_shape();
+  return ind < limit && ind >= Utils::Vector3i{};
 }
 
 double lb_lbnode_get_density(const Utils::Vector3i &ind) {

src/core/grid_based_algorithms/lb_interface.hpp

@@ -216,6 +216,11 @@ void lb_lbfluid_load_checkpoint(const std::string &filename, int binary);
  */
 bool lb_lbnode_is_index_valid(const Utils::Vector3i &ind);
 
+/**
+ * @brief returns the shape of the LB fluid lattice
+ */
+Utils::Vector3i lb_lbfluid_get_shape();
+
 void lb_lbfluid_on_lb_params_change(LBParam field);
 
 Utils::Vector3d lb_lbfluid_calc_fluid_momentum();

src/python/espressomd/lb.pxd

@@ -76,6 +76,7 @@ cdef extern from "grid_based_algorithms/lb_interface.hpp":
     ActiveLB lb_lbfluid_get_lattice_switch() except +
     Vector6d lb_lbfluid_get_stress() except +
     bool lb_lbnode_is_index_valid(const Vector3i & ind) except +
+    Vector3i lb_lbfluid_get_shape() except +
     const Vector3d lb_lbnode_get_velocity(const Vector3i & ind) except +
     void lb_lbnode_set_velocity(const Vector3i & ind, const Vector3d & u) except +
     double lb_lbnode_get_density(const Vector3i & ind) except +

src/python/espressomd/lb.pyx

@@ -20,6 +20,7 @@ from __future__ import print_function, absolute_import, division
 include "myconfig.pxi"
 import os
 import cython
+import itertools
 import numpy as np
 cimport numpy as np
 from libc cimport stdint
@@ -244,6 +245,28 @@ cdef class HydrodynamicInteraction(Actor):
 
     def _deactivate_method(self):
         lb_lbfluid_set_lattice_switch(NONE)
+
+    property shape:
+        def __get__(self):
+            cdef Vector3i shape = lb_lbfluid_get_shape()
+            return (shape[0], shape[1], shape[2])
+
+    property stress:
+        def __get__(self):
+            cdef Vector6d res
+            res = lb_lbfluid_get_stress() 
+            return array_locked((
+                res[0], res[1], res[2], res[3], res[4], res[5]))
+
+        def __set__(self, value):
+            raise NotImplementedError
+
+    def nodes(self):
+        """Provides a generator for iterating over all lb nodes"""
+
+        shape = self.shape
+        for i, j, k in itertools.product(range(shape[0]), range(shape[1]), range(shape[2])):
+            yield self[i, j, k]
 
 
 # LBFluid main class
@@ -261,17 +284,7 @@ cdef class LBFluid(HydrodynamicInteraction):
         self.validate_params()
         self._set_lattice_switch()
         self._set_params_in_es_core()
-
-    property stress:
-        def __get__(self):
-            cdef Vector6d res
-            res = lb_lbfluid_get_stress() 
-            return array_locked((
-                res[0], res[1], res[2], res[3], res[4], res[5]))
-
-        def __set__(self, value):
-            raise NotImplementedError
-
+
 IF CUDA:
     cdef class LBFluidGPU(HydrodynamicInteraction):
         """
@@ -330,6 +343,10 @@ cdef class LBFluidRoutines(object):
         self.node[2] = key[2]
         if not lb_lbnode_is_index_valid(self.node):
             raise ValueError("LB node index out of bounds")
+
+    property index:
+        def __get__(self):
+            return (self.node[0], self.node[1], self.node[2])
 
     property velocity:
         def __get__(self):

testsuite/python/lb.py

@@ -108,14 +108,6 @@ def test_mass_momentum_thermostat(self):
         all_temp_particle = []
         all_temp_fluid = []
 
-        # Cache the lb nodes
-        lb_nodes = []
-        n_nodes = int(self.params['box_l'] / self.params['agrid'])
-        for i in range(n_nodes):
-            for j in range(n_nodes):
-                for k in range(n_nodes):
-                    lb_nodes.append(self.lbf[i, j, k])
-
         # Integration
         for i in range(self.params['int_times']):
             self.system.integrator.run(self.params['int_steps'])
@@ -125,14 +117,15 @@ def test_mass_momentum_thermostat(self):
             fluid_temp = 0.0
 
             # Go over lb lattice
-            for lb_node in lb_nodes:
+            for lb_node in self.lbf.nodes():
                 dens = lb_node.density
                 fluid_mass += dens
                 fluid_temp += np.sum(lb_node.velocity**2) * dens
 
             # Normalize
-            fluid_mass /= len(lb_nodes)
-            fluid_temp *= self.system.volume() / (3. * len(lb_nodes)**2)
+            fluid_mass /= np.product(self.lbf.shape)
+            fluid_temp *= self.system.volume() / (
+                3. * np.product(self.lbf.shape)**2)
 
             # check mass conversation
             self.assertAlmostEqual(fluid_mass, self.params["dens"],
@@ -190,10 +183,8 @@ def test_stress_tensor(self):
         system.integrator.run(10)
         stress = np.zeros((3, 3))
         agrid = self.params["agrid"]
-        for i in range(int(system.box_l[0] / agrid)):
-            for j in range(int(system.box_l[1] / agrid)):
-                for k in range(int(system.box_l[2] / agrid)):
-                    stress += self.lbf[i, j, k].stress
+        for n in self.lbf.nodes():
+            stress += n.stress
 
         stress /= system.volume() / agrid**3
 
@@ -202,7 +193,6 @@ def test_stress_tensor(self):
         obs_stress = np.array([[obs_stress[0], obs_stress[1], obs_stress[3]],
                                [obs_stress[1], obs_stress[2], obs_stress[4]],
                                [obs_stress[3], obs_stress[4], obs_stress[5]]])
-        print(stress / obs_stress)
         np.testing.assert_allclose(stress, obs_stress, atol=1E-10)
 
     def test_lb_node_set_get(self):
@@ -215,6 +205,13 @@ def test_lb_node_set_get(self):
             tau=self.system.time_step,
             ext_force_density=[0, 0, 0])
         self.system.actors.add(self.lbf)
+
+        self.assertEqual(self.lbf.shape, 
+                         (
+                         int(self.system.box_l[0] / self.params["agrid"]),
+                         int(self.system.box_l[1] / self.params["agrid"]),
+                             int(self.system.box_l[2] / self.params["agrid"])))
+
         v_fluid = np.array([1.2, 4.3, 0.2])
         self.lbf[0, 0, 0].velocity = v_fluid
         np.testing.assert_allclose(
@@ -223,6 +220,8 @@ def test_lb_node_set_get(self):
         self.lbf[0, 0, 0].density = density
         self.assertAlmostEqual(self.lbf[0, 0, 0].density, density, delta=1e-4)
 
+        self.assertEqual(self.lbf[3, 2, 1].index, (3, 2, 1))
+
     def test_parameter_change_without_seed(self):
         self.system.actors.clear()
         self.lbf = self.lb_class(
@@ -324,9 +323,9 @@ def test_a_ext_force_density(self):
         # (force is applied only to the second half of the first integration step)
         fluid_velocity = np.array(ext_force_density) * self.system.time_step * (
             n_time_steps - 0.5) / self.params['dens']
-        for n in list(itertools.combinations(range(int(self.system.box_l[0] / self.params['agrid'])), 3)):
+        for n in self.lbf.nodes():
             np.testing.assert_allclose(
-                np.copy(self.lbf[n].velocity), fluid_velocity, atol=1E-6)
+                np.copy(n.velocity), fluid_velocity, atol=1E-6)
 
 
 class TestLBCPU(TestLB, ut.TestCase):