diff --git a/doc/sphinx/lb.rst b/doc/sphinx/lb.rst index b6e23176461..88c81d0fa04 100644 --- a/doc/sphinx/lb.rst +++ b/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: diff --git a/src/python/espressomd/lb.pyx b/src/python/espressomd/lb.pyx index 7576ba9c4ea..f0921ee9bef 100644 --- a/src/python/espressomd/lb.pyx +++ b/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() diff --git a/src/python/espressomd/particle_data.pyx b/src/python/espressomd/particle_data.pyx index f1b0612fcdb..f5d0559cead 100644 --- a/src/python/espressomd/particle_data.pyx +++ b/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) diff --git a/testsuite/python/CMakeLists.txt b/testsuite/python/CMakeLists.txt index a94ba79bc5d..abf4ea61c93 100644 --- a/testsuite/python/CMakeLists.txt +++ b/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) diff --git a/testsuite/python/lb_slice.py b/testsuite/python/lb_slice.py new file mode 100644 index 00000000000..f9241acc1ec --- /dev/null +++ b/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 . +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()