Implement slicing for LB

doc/sphinx/lb.rst

@@ -190,13 +190,25 @@ Appending three indices to the ``lb`` object returns an object that represents t
 
 All of these properties can be read and used in further calculations. Only the property ``population`` can be modified. The indices ``x,y,z`` are integers and enumerate the LB nodes in the three directions, starts with 0. To modify ``boundary``, refer to :ref:`Setting up boundary conditions`.
 
-Examples::
+Example::
 
     print(lb[0, 0, 0].velocity)
-
     lb[0, 0, 0].density = 1.2
 
-The first line prints the fluid velocity at node 0 0 0 to the screen. The second line sets this fluid node's density to the value ``1.2``.
+The first line prints the fluid velocity at node (0 0 0) to the screen.
+The second line sets this fluid node's density to the value ``1.2``.
+
+The nodes can be read and modified using slices. Example::
+
+    print(lb[0:4:2, 0:2, 0].velocity)
+    lb[0:4:2, 0:2, 0].density = [[[1.1], [1.2]], [[1.3], [1.4]]]
+
+The first line prints an array of shape (2, 2, 1, 3) with the velocities
+of nodes (0 0 0), (0 1 0), (2 0 0), (2 1 0). The second line updates
+these nodes with densities ranging from 1.1 to 1.4. You can set either
+a value that matches the length of the slice (which sets each node
+individually), or a single value that will be copied to every node
+(e.g. a scalar for density, or an array of length 3 for the velociy).
 
 .. _Removing total fluid momentum:
 

src/python/espressomd/lb.pyx

@@ -20,6 +20,7 @@ include "myconfig.pxi"
 import os
 import cython
 import itertools
+import functools
 import numpy as np
 cimport numpy as np
 from libc cimport stdint
@@ -82,12 +83,20 @@ cdef class HydrodynamicInteraction(Actor):
         return _construct, (self.__class__, self._params), None
 
     def __getitem__(self, key):
-        utils.check_type_or_throw_except(
-            key, 3, int, "The index of an lb fluid node consists of three integers, e.g. lbf[0,0,0]")
-        return LBFluidRoutines(key)
-
+        cdef Vector3i shape
+        if isinstance(key, (tuple, list, np.ndarray)):
+            if len(key) == 3:
+                if any(isinstance(typ, slice) for typ in key):
+                    shape = lb_lbfluid_get_shape()
+                    return LBSlice(key, (shape[0], shape[1], shape[2]))
+                else:
+                    return LBFluidRoutines(np.array(key))
+        else:
+            raise Exception(
+                "%s is not a valid key. Should be a point on the nodegrid e.g. lbf[0,0,0], or a slice" % key)
     # validate the given parameters on actor initialization
     ####################################################
+
     def validate_params(self):
         default_params = self.default_params()
 
@@ -444,9 +453,9 @@ IF CUDA:
             length = positions.shape[0]
             velocities = np.empty_like(positions)
             if three_point:
-                quadratic_velocity_interpolation( < double * >np.PyArray_GETPTR2(positions, 0, 0), < double * >np.PyArray_GETPTR2(velocities, 0, 0), length)
+                quadratic_velocity_interpolation(< double * >np.PyArray_GETPTR2(positions, 0, 0), < double * >np.PyArray_GETPTR2(velocities, 0, 0), length)
             else:
-                linear_velocity_interpolation( < double * >np.PyArray_GETPTR2(positions, 0, 0), < double * >np.PyArray_GETPTR2(velocities, 0, 0), length)
+                linear_velocity_interpolation(< double * >np.PyArray_GETPTR2(positions, 0, 0), < double * >np.PyArray_GETPTR2(velocities, 0, 0), length)
             return velocities * lb_lbfluid_get_lattice_speed()
 
 cdef class LBFluidRoutines:
@@ -542,3 +551,106 @@ cdef class LBFluidRoutines:
 
         def __set__(self, value):
             raise NotImplementedError
+
+
+class LBSlice:
+
+    def __init__(self, key, shape):
+        self.x_indices, self.y_indices, self.z_indices = self.get_indices(
+            key, shape[0], shape[1], shape[2])
+
+    def get_indices(self, key, shape_x, shape_y, shape_z):
+        x_indices = np.atleast_1d(np.arange(shape_x)[key[0]])
+        y_indices = np.atleast_1d(np.arange(shape_y)[key[1]])
+        z_indices = np.atleast_1d(np.arange(shape_z)[key[2]])
+        return x_indices, y_indices, z_indices
+
+    def get_values(self, x_indices, y_indices, z_indices, prop_name):
+        shape_res = np.shape(
+            getattr(LBFluidRoutines(np.array([0, 0, 0])), prop_name))
+        res = np.zeros(
+            (x_indices.size,
+             y_indices.size,
+             z_indices.size,
+             *shape_res))
+        for i, x in enumerate(x_indices):
+            for j, y in enumerate(y_indices):
+                for k, z in enumerate(z_indices):
+                    res[i, j, k] = getattr(LBFluidRoutines(
+                        np.array([x, y, z])), prop_name)
+        if shape_res == (1,):
+            res = np.squeeze(res, axis=-1)
+        return array_locked(res)
+
+    def set_values(self, x_indices, y_indices, z_indices, prop_name, value):
+        for i, x in enumerate(x_indices):
+            for j, y in enumerate(y_indices):
+                for k, z in enumerate(z_indices):
+                    setattr(LBFluidRoutines(
+                        np.array([x, y, z])), prop_name, value[i, j, k])
+
+
+def _add_lb_slice_properties():
+    """
+    Automatically add all of LBFluidRoutines's properties to LBSlice.
+
+    """
+
+    def set_attribute(lb_slice, value, attribute):
+        """
+        Setter function that sets attribute on every member of lb_slice.
+        If values contains only one element, all members are set to it.
+
+        """
+
+        indices = [lb_slice.x_indices, lb_slice.y_indices, lb_slice.z_indices]
+        N = [len(x) for x in indices]
+
+        if N[0] * N[1] * N[2] == 0:
+            raise AttributeError("Cannot set properties of an empty LBSlice")
+
+        value = np.copy(value)
+        attribute_shape = lb_slice.get_values(
+            *np.zeros((3, 1), dtype=int), attribute).shape[3:]
+        target_shape = (*N, *attribute_shape)
+
+        # broadcast if only one element was provided
+        if value.shape == attribute_shape:
+            value = np.ones(target_shape) * value
+
+        if value.shape != target_shape:
+            raise ValueError(
+                f"Input-dimensions of {attribute} array {value.shape} does not match slice dimensions {target_shape}.")
+
+        lb_slice.set_values(*indices, attribute, value)
+
+    def get_attribute(lb_slice, attribute):
+        """
+        Getter function that copies attribute from every member of
+        lb_slice into an array (if possible).
+
+        """
+
+        indices = [lb_slice.x_indices, lb_slice.y_indices, lb_slice.z_indices]
+        N = [len(x) for x in indices]
+
+        if N[0] * N[1] * N[2] == 0:
+            return np.empty(0, dtype=type(None))
+
+        return lb_slice.get_values(*indices, attribute)
+
+    for attribute_name in dir(LBFluidRoutines):
+        if attribute_name in dir(LBSlice) or not isinstance(
+                getattr(LBFluidRoutines, attribute_name), type(LBFluidRoutines.density)):
+            continue
+
+        # synthesize a new property
+        new_property = property(
+            functools.partial(get_attribute, attribute=attribute_name),
+            functools.partial(set_attribute, attribute=attribute_name),
+            doc=getattr(LBFluidRoutines, attribute_name).__doc__ or f'{attribute_name} for a slice')
+        # attach the property to LBSlice
+        setattr(LBSlice, attribute_name, new_property)
+
+
+_add_lb_slice_properties()

src/python/espressomd/particle_data.pyx

@@ -2115,7 +2115,7 @@ def _add_particle_slice_properties():
 
     """
 
-    def seta(particle_slice, values, attribute):
+    def set_attribute(particle_slice, values, attribute):
         """
         Setter function that sets attribute on every member of particle_slice.
         If values contains only one element, all members are set to it. If it
@@ -2193,7 +2193,7 @@ def _add_particle_slice_properties():
 
                 return
 
-    def geta(particle_slice, attribute):
+    def get_attribute(particle_slice, attribute):
         """
         Getter function that copies attribute from every member of
         particle_slice into an array (if possible).
@@ -2231,8 +2231,10 @@ def _add_particle_slice_properties():
             continue
 
         # synthesize a new property
-        new_property = property(functools.partial(geta, attribute=attribute_name), functools.partial(
-            seta, attribute=attribute_name), doc=getattr(ParticleHandle, attribute_name).__doc__)
+        new_property = property(
+            functools.partial(get_attribute, attribute=attribute_name),
+            functools.partial(set_attribute, attribute=attribute_name),
+            doc=getattr(ParticleHandle, attribute_name).__doc__)
         # attach the property to ParticleSlice
         setattr(ParticleSlice, attribute_name, new_property)
 

testsuite/python/CMakeLists.txt

@@ -187,6 +187,7 @@ python_test(FILE actor.py MAX_NUM_PROC 1)
 python_test(FILE drude.py MAX_NUM_PROC 2)
 python_test(FILE thermalized_bond.py MAX_NUM_PROC 4)
 python_test(FILE thole.py MAX_NUM_PROC 4)
+python_test(FILE lb_slice.py MAX_NUM_PROC 1)
 python_test(FILE lb_switch.py MAX_NUM_PROC 1 LABELS gpu)
 python_test(FILE lb_boundary_velocity.py MAX_NUM_PROC 1)
 python_test(FILE lb_boundary_volume_force.py MAX_NUM_PROC 4)

testsuite/python/lb_slice.py

@@ -0,0 +1,122 @@
+# Copyright (C) 2010-2019 The ESPResSo project
+#
+# This file is part of ESPResSo.
+#
+# ESPResSo is free software: you can redistribute it and/or modify
+# it under the terms of the GNU General Public License as published by
+# the Free Software Foundation, either version 3 of the License, or
+# (at your option) any later version.
+#
+# ESPResSo is distributed in the hope that it will be useful,
+# but WITHOUT ANY WARRANTY; without even the implied warranty of
+# MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
+# GNU General Public License for more details.
+#
+# You should have received a copy of the GNU General Public License
+# along with this program.  If not, see <http://www.gnu.org/licenses/>.
+import espressomd.lb
+import unittest as ut
+import numpy as np
+
+
+class LBSliceTest(ut.TestCase):
+
+    """This simple test first writes random numbers and then reads them 
+    to same slices of LB nodes and compares if the results are the same, 
+    shape and value wise.
+    """
+
+    system = espressomd.System(box_l=[10.0, 10.0, 10.0])
+    system.time_step = .01
+    system.cell_system.skin = 0.1
+    np.random.seed(seed=42)
+
+    def test_slicing(self):
+        system = self.system
+
+        lb_fluid = espressomd.lb.LBFluid(
+            agrid=1.0, dens=1., visc=1., tau=0.01)
+        system.actors.add(lb_fluid)
+
+        # array locked
+        array = lb_fluid[1:-1:2, 5, 3:6:2].velocity
+        with self.assertRaisesRegex(ValueError, "ESPResSo array properties return non-writable arrays"):
+            array[0, 0, 0, 1] = 5.
+
+        # velocity on test slice [:-1, :-1, -1]
+        input_vel = np.random.rand(9, 9, 9, 3)
+        lb_fluid[:-1, :-1, :-1].velocity = input_vel
+        output_vel = lb_fluid[:-1, :-1, :-1].velocity
+        np.testing.assert_array_almost_equal(input_vel, np.copy(output_vel))
+
+        with self.assertRaisesRegex(ValueError, r"Input-dimensions of velocity array \(9, 9, 9, 2\) does not match slice dimensions \(9, 9, 9, 3\)"):
+            lb_fluid[:-1, :-1, :-1].velocity = input_vel[:, :, :, :2]
+
+        # velocity broadcast
+        lb_fluid[:, :, 0].velocity = [1, 2, 3]
+        np.testing.assert_array_almost_equal(
+            np.copy(lb_fluid[:, :, 0].velocity), 10 * [10 * [[[1, 2, 3]]]])
+
+        # density on test slice [1:-1:2, 5, 3:6:2]
+        input_dens = np.random.rand(4, 1, 2)
+        lb_fluid[1:-1:2, 5, 3:6:2].density = input_dens
+        output_dens = lb_fluid[1:-1:2, 5, 3:6:2].density
+        np.testing.assert_array_almost_equal(input_dens, np.copy(output_dens))
+
+        # density broadcast
+        lb_fluid[:, :, 0].density = 1.2
+        np.testing.assert_array_almost_equal(
+            np.copy(lb_fluid[:, :, 0].density), 1.2)
+
+        # population on test slice [:, :, :]
+        input_pop = np.random.rand(10, 10, 10, 19)
+        lb_fluid[:, :, :].population = input_pop
+        output_pop = lb_fluid[:, :, :].population
+        np.testing.assert_array_almost_equal(input_pop, np.copy(output_pop))
+
+        with self.assertRaisesRegex(ValueError, r"Input-dimensions of population array \(10, 10, 10, 5\) does not match slice dimensions \(10, 10, 10, 19\)"):
+            lb_fluid[:, :, :].population = input_pop[:, :, :, :5]
+
+        # pressure tensor on test slice [3, 6, 2:5]
+        output_pressure_shape = lb_fluid[3, 6, 2:5].pressure_tensor.shape
+        should_pressure_shape = (1, 1, 3, 3, 3)
+        np.testing.assert_array_almost_equal(
+            output_pressure_shape, should_pressure_shape)
+
+        with self.assertRaises(NotImplementedError):
+            lb_fluid[3, 6, 2:5].pressure_tensor = np.zeros(
+                should_pressure_shape)
+
+        # pressure tensor neq on test slice [3, 6, 2:10]
+        output_pressure_neq_shape = lb_fluid[3:5,
+                                             6:7, 2:10].pressure_tensor_neq.shape
+        should_pressure_neq_shape = (2, 1, 8, 3, 3)
+        np.testing.assert_array_almost_equal(
+            output_pressure_neq_shape, should_pressure_neq_shape)
+
+        with self.assertRaises(NotImplementedError):
+            lb_fluid[3:5, 6:7, 2:10].pressure_tensor_neq = np.zeros(
+                output_pressure_neq_shape)
+
+        # index on test slice [1, 1:5, 6:]
+        output_index_shape = lb_fluid[1, 1:5, 6:].index.shape
+        should_index_shape = (1, 4, 4, 3)
+        np.testing.assert_array_almost_equal(
+            output_index_shape, should_index_shape)
+
+        with self.assertRaisesRegex(AttributeError, "attribute 'index' of 'espressomd.lb.LBFluidRoutines' objects is not writable"):
+            lb_fluid[1, 1:5, 6:].index = np.zeros(output_index_shape)
+
+        # boundary on test slice [1:, 1:, 1:]
+        if espressomd.has_features('LB_BOUNDARIES'):
+            output_boundary_shape = lb_fluid[1:, 1:, 1:].boundary.shape
+            should_boundary_shape = (9, 9, 9)
+            np.testing.assert_array_almost_equal(
+                output_boundary_shape, should_boundary_shape)
+
+            with self.assertRaises(NotImplementedError):
+                lb_fluid[1:, 1:, 1:].boundary = np.zeros(should_boundary_shape)
+
+
+if __name__ == "__main__":
+    ut.main()