Implement slicing for LB (espressomd#4195)

Fixes espressomd#4143

Description of changes:
- Modified __getitem__ to detect slices in input keys
- Introduced LBSlice class that calculates node indices from slice input and calls single node getter from LBFLuidRoutines class for every node in the slice.

The output is a multidimensional array containing the requested quantity and is of dimension len(slice_x)*len(slice_y)*len(slice_z)*len(quantity), where len(quantity) is the dimensions of the quantity, for example 3 in case of velocity. The output array is arranged in the way numpy cycles through slices. Therefore, when setting attributes one has to set the dimensions of the input array exactly the same way numpy cycles through the sliced nodes.
An error is issued if the dimensions of the slices and the input does not match.

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()