#
# Copyright (C) 2021-2023 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 itertools
import numpy as np
from . import utils
from .detail.walberla import VTKOutputBase, LatticeWalberla # pylint: disable=unused-import
from .script_interface import ScriptInterfaceHelper, script_interface_register, ScriptObjectList, array_variant
import espressomd.detail.walberla
import espressomd.shapes
[docs]
@script_interface_register
class EKFFT(ScriptInterfaceHelper):
"""
A FFT-based Poisson solver.
Intrinsically assumes periodic boundary conditions.
Parameters
----------
lattice : :obj:`espressomd.lb.LatticeWalberla <espressomd.detail.walberla.LatticeWalberla>`
Lattice object.
permittivity : :obj:`float`
permittivity of the fluid :math:`\\epsilon_0 \\epsilon_{\\mathrm{r}}`.
single_precision : :obj:`bool`, optional
Use single-precision floating-point arithmetic.
"""
_so_name = "walberla::EKFFT"
_so_features = ("WALBERLA_FFT",)
_so_creation_policy = "GLOBAL"
[docs]
@script_interface_register
class EKNone(ScriptInterfaceHelper):
"""
The default Poisson solver.
Imposes a null electrostatic potential everywhere.
Parameters
----------
lattice : :obj:`espressomd.lb.LatticeWalberla <espressomd.detail.walberla.LatticeWalberla>`
Lattice object.
single_precision : :obj:`bool`, optional
Use single-precision floating-point arithmetic.
"""
_so_name = "walberla::EKNone"
_so_features = ("WALBERLA",)
_so_creation_policy = "GLOBAL"
[docs]
@script_interface_register
class EKSpecies(ScriptInterfaceHelper,
espressomd.detail.walberla.LatticeModel):
"""
The advection-diffusion-reaction method for chemical species using waLBerla.
Parameters
----------
lattice : :obj:`espressomd.electrokinetics.LatticeWalberla <espressomd.detail.walberla.LatticeWalberla>`
Lattice object.
tau : :obj:`float`
EK time step, must be an integer multiple of the MD time step.
density : :obj:`float`
Species density.
diffusion : :obj:`float`
Species diffusion coefficient.
valency : :obj:`float`
Species valency.
advection : :obj:`bool`
Whether to enable advection.
friction_coupling : :obj:`bool`
Whether to enable friction coupling.
ext_efield : (3,) array_like of :obj:`float`, optional
External electrical field.
kT : :obj:`float`, optional
Thermal energy of the simulated heat bath (for thermalized species).
Set it to 0 for an unthermalized species.
single_precision : :obj:`bool`, optional
Use single-precision floating-point arithmetic.
thermalized : :obj:`bool`, optional
Whether the species is thermalized.
seed : :obj:`int`, optional
Seed for the random number generator.
Methods
-------
clear_density_boundaries()
Remove density boundary conditions.
clear_flux_boundaries()
Remove flux boundary conditions.
clear_boundaries()
Remove all boundary conditions.
save_checkpoint()
Write EK densities and boundary conditions to a file.
Parameters
----------
path : :obj:`str`
Destination file path.
binary : :obj:`bool`
Whether to write in binary or ASCII mode.
load_checkpoint()
Load EK densities and boundary conditions from a file.
Parameters
----------
path : :obj:`str`
File path to read from.
binary : :obj:`bool`
Whether to read in binary or ASCII mode.
add_vtk_writer()
Attach a VTK writer.
Parameters
----------
vtk : :class:`espressomd.electrokinetics.VTKOutput`
VTK writer.
remove_vtk_writer()
Detach a VTK writer.
Parameters
----------
vtk : :class:`espressomd.electrokinetics.VTKOutput`
VTK writer.
clear_vtk_writers()
Detach all VTK writers.
"""
_so_name = "walberla::EKSpecies"
_so_features = ("WALBERLA",)
_so_creation_policy = "GLOBAL"
_so_bind_methods = (
"clear_density_boundaries",
"clear_flux_boundaries",
"clear_boundaries",
"add_vtk_writer",
"remove_vtk_writer",
"clear_vtk_writers",
)
def __init__(self, *args, **kwargs):
if "sip" not in kwargs:
params = self.default_params()
params.update(kwargs)
super().__init__(*args, **params)
else:
super().__init__(**kwargs)
[docs]
def default_params(self):
return {"single_precision": False,
"kT": 0., "ext_efield": [0., 0., 0.],
"thermalized": False}
def __getitem__(self, key):
if isinstance(key, (tuple, list, np.ndarray)) and len(key) == 3:
if any(isinstance(item, slice) for item in key):
return EKSpeciesSlice(
parent_sip=self, slice_range=key, node_grid=self.shape)
else:
return EKSpeciesNode(parent_sip=self, index=np.array(key))
raise TypeError(
f"{key} is not a valid index. Should be a point on the "
"nodegrid e.g. ek[0,0,0], or a slice, e.g. ek[:,0,0]")
[docs]
def add_boundary_from_shape(self, shape, value, boundary_type):
"""
Set boundary conditions from a shape.
Parameters
----------
shape : :obj:`espressomd.shapes.Shape`
Shape to rasterize.
value : (O,) or (L, M, N, O) array_like of :obj:`float`, optional
Boundary numerical value. If a single value of shape ``(O,)``
is given, it will be broadcast to all nodes inside the shape,
otherwise ``L, M, N`` must be equal to the EK grid dimensions.
boundary_type : Union[:class:`~espressomd.electrokinetics.DensityBoundary`,
:class:`~espressomd.electrokinetics.FluxBoundary`] (optional)
Type of the boundary condition.
"""
if not issubclass(boundary_type, (FluxBoundary, DensityBoundary)):
raise TypeError(
"Parameter 'boundary_type' must be a subclass of FluxBoundary or DensityBoundary")
if not hasattr(value, "__iter__"):
value = (value, )
value = np.array(value, dtype=float)
utils.check_type_or_throw_except(
shape, 1, espressomd.shapes.Shape, "expected an espressomd.shapes.Shape")
if issubclass(boundary_type, FluxBoundary):
if np.shape(value) not in [(3,), tuple(self.shape) + (3,)]:
raise ValueError(
f"Cannot process flux value grid of shape {np.shape(value)}")
if issubclass(boundary_type, DensityBoundary):
if np.shape(value) not in [(1,), tuple(self.shape) + (1,)]:
raise ValueError(
f"Cannot process density value grid of shape {np.shape(value)}")
mask = self.get_shape_bitmask(shape=shape).astype(int)
if issubclass(boundary_type, FluxBoundary):
boundaries_update_method = "update_flux_boundary_from_shape"
else:
boundaries_update_method = "update_density_boundary_from_shape"
self.call_method(
boundaries_update_method,
raster=array_variant(mask.flatten()),
values=array_variant(value.flatten()))
[docs]
class FluxBoundary:
"""
Hold flux information for the flux boundary
condition at a single node.
"""
def __init__(self, flux):
utils.check_type_or_throw_except(
flux, 3, float, "FluxBoundary flux must be three floats")
self.flux = flux
[docs]
class DensityBoundary:
"""
Hold density information for the density boundary
condition at a single node.
"""
def __init__(self, density):
utils.check_type_or_throw_except(
density, 1, float, "DensityBoundary flux must be one float")
self.density = density
[docs]
@script_interface_register
class EKSpeciesNode(ScriptInterfaceHelper):
_so_name = "walberla::EKSpeciesNode"
_so_creation_policy = "GLOBAL"
[docs]
def required_keys(self):
return {"parent_sip", "index"}
[docs]
def validate_params(self, params):
utils.check_required_keys(self.required_keys(), params.keys())
utils.check_type_or_throw_except(
params["index"], 3, int, "The index of an EK species node consists of three integers.")
def __init__(self, *args, **kwargs):
if "sip" not in kwargs:
self.validate_params(kwargs)
super().__init__(*args, **kwargs)
utils.handle_errors("EKSpeciesNode instantiation failed")
else:
super().__init__(**kwargs)
def __reduce__(self):
raise NotImplementedError("Cannot serialize EK species node objects")
def __eq__(self, obj):
return isinstance(obj, EKSpeciesNode) and self.index == obj.index
def __hash__(self):
return hash(self.index)
@property
def index(self):
return tuple(self._index)
@index.setter
def index(self, value):
raise RuntimeError("Parameter 'index' is read-only.")
@property
def density(self):
return self.call_method("get_density")
@density.setter
def density(self, value):
self.call_method("set_density", value=value)
@property
def is_boundary(self):
return self.call_method("get_is_boundary")
@is_boundary.setter
def is_boundary(self, value):
raise RuntimeError("Property 'is_boundary' is read-only.")
@property
def density_boundary(self):
"""
Returns
-------
:class:`~espressomd.electrokinetics.DensityBoundary`
If the node is a boundary node
``None``
If the node is not a boundary node
"""
density = self.call_method("get_node_density_at_boundary")
if density is not None:
return DensityBoundary(density)
return None
@density_boundary.setter
def density_boundary(self, value):
"""
Parameters
----------
value : :class:`~espressomd.electrokinetics.DensityBoundary` or ``None``
If value is :class:`~espressomd.EkSpecies.DensityBoundary`,
set the node to be a boundary node with the specified density.
If value is ``None``, the node will become a domain node.
"""
if isinstance(value, DensityBoundary):
value = value.density
elif value is not None:
raise TypeError(
"Parameter 'value' must be an instance of DensityBoundary or None")
self.call_method("set_node_density_at_boundary", value=value)
@property
def flux_boundary(self):
"""
Returns
-------
:class:`~espressomd.electrokinetics.FluxBoundary`
If the node is a boundary node
``None``
If the node is not a boundary node
"""
flux = self.call_method("get_node_flux_at_boundary")
if flux is not None:
return FluxBoundary(flux)
return None
@flux_boundary.setter
def flux_boundary(self, value):
"""
Parameters
----------
value : :class:`~espressomd.electrokinetics.FluxBoundary` or ``None``
If value is :class:`~espressomd.EkSpecies.FluxBoundary`,
set the node to be a boundary node with the specified flux.
If value is ``None``, the node will become a domain node.
"""
if isinstance(value, FluxBoundary):
value = value.flux
elif value is not None:
raise TypeError(
"Parameter 'value' must be an instance of FluxBoundary or None")
self.call_method("set_node_flux_at_boundary", value=value)
[docs]
@script_interface_register
class EKSpeciesSlice(ScriptInterfaceHelper):
_so_name = "walberla::EKSpeciesSlice"
_so_creation_policy = "GLOBAL"
[docs]
def required_keys(self):
return {"parent_sip", "slice_range"}
[docs]
def validate_params(self, params):
utils.check_required_keys(self.required_keys(), params.keys())
def __init__(self, *args, **kwargs):
if "sip" in kwargs:
super().__init__(**kwargs)
else:
self.validate_params(kwargs)
slice_range = kwargs.pop("slice_range")
grid_size = kwargs["parent_sip"].shape
extra_kwargs = espressomd.detail.walberla.get_slice_bounding_box(
slice_range, grid_size)
node = EKSpeciesNode(index=np.array([0, 0, 0]), **kwargs)
super().__init__(*args, node_sip=node, **kwargs, **extra_kwargs)
utils.handle_errors("EKSpeciesSlice instantiation failed")
def __iter__(self):
lower, upper = self.call_method("get_slice_ranges")
indices = [list(range(lower[i], upper[i])) for i in range(3)]
lb_sip = self.call_method("get_ek_sip")
for index in itertools.product(*indices):
yield EKSpeciesNode(parent_sip=lb_sip, index=np.array(index))
def __reduce__(self):
raise NotImplementedError("Cannot serialize EK species slice objects")
def _getter(self, attr):
value_grid, shape = self.call_method(f"get_{attr}")
if attr == "flux_at_boundary":
value_grid = [
None if x is None else FluxBoundary(x) for x in value_grid]
elif attr == "density_at_boundary":
value_grid = [
None if x is None else DensityBoundary(x) for x in value_grid]
return utils.array_locked(np.reshape(value_grid, shape))
def _setter(self, attr, values):
dimensions = self.call_method("get_slice_size")
if 0 in dimensions:
raise AttributeError(
f"Cannot set properties of an empty '{self.__class__.__name__}' object")
values = np.copy(values)
value_shape = tuple(self.call_method("get_value_shape", name=attr))
target_shape = (*dimensions, *value_shape)
# broadcast if only one element was provided
if values.shape == value_shape or values.shape == () and value_shape == (1,):
values = np.full(target_shape, values)
def shape_squeeze(shape):
return tuple(x for x in shape if x != 1)
if shape_squeeze(values.shape) != shape_squeeze(target_shape):
target_shape = tuple([int(x) for x in target_shape])
raise ValueError(
f"Input-dimensions of '{attr}' array {values.shape} does not match slice dimensions {target_shape}")
self.call_method(f"set_{attr}", values=values.flatten())
@property
def density(self):
return self._getter("density",)
@density.setter
def density(self, value):
self._setter("density", value)
@property
def is_boundary(self):
return self._getter("is_boundary")
@is_boundary.setter
def is_boundary(self, value):
raise RuntimeError("Property 'is_boundary' is read-only.")
@property
def density_boundary(self):
"""
Returns
-------
(N, M, L) array_like of :class:`~espressomd.electrokinetics.DensityBoundary`
If the nodes are boundary nodes
(N, M, L) array_like of ``None``
If the nodes are not boundary nodes
"""
return self._getter("density_at_boundary")
@density_boundary.setter
def density_boundary(self, values):
"""
Parameters
----------
values : (N, M, L) array_like of :class:`~espressomd.electrokinetics.DensityBoundary` or obj:`None`
If values are :class:`~espressomd.electrokinetics.DensityBoundary`,
set the nodes to be boundary nodes with the specified density.
If values are obj:`None`, the nodes will become domain nodes.
"""
type_error_msg = "Parameter 'values' must be an array_like of DensityBoundary or None"
values = np.copy(values)
if values.dtype != np.dtype("O"):
raise TypeError(type_error_msg)
for index in np.ndindex(*values.shape):
if values[index] is not None:
if not isinstance(values[index], DensityBoundary):
raise TypeError(type_error_msg)
values[index] = np.array(values[index].density)
self._setter("density_at_boundary", values=values)
@property
def flux_boundary(self):
"""
Returns
-------
(N, M, L) array_like of :class:`~espressomd.electrokinetics.FluxBoundary`
If the nodes are boundary nodes
(N, M, L) array_like of `None``
If the nodes are not boundary nodes
"""
return self._getter("flux_at_boundary")
@flux_boundary.setter
def flux_boundary(self, values):
"""
Parameters
----------
values : (N, M, L) array_like of :class:`~espressomd.electrokinetics.FluxBoundary` or obj:`None`
If values are :class:`~espressomd.lb.FluxBoundary`,
set the nodes to be boundary nodes with the specified flux.
If values are obj:`None`, the nodes will become domain nodes.
"""
type_error_msg = "Parameter 'values' must be an array_like of FluxBoundary or None"
values = np.copy(values)
if values.dtype != np.dtype("O"):
raise TypeError(type_error_msg)
for index in np.ndindex(*values.shape):
if values[index] is not None:
if not isinstance(values[index], FluxBoundary):
raise TypeError(type_error_msg)
values[index] = np.array(values[index].flux)
self._setter("flux_at_boundary", values=values)
[docs]
@script_interface_register
class VTKOutput(VTKOutputBase):
"""
Create a VTK writer.
Files are written to ``<base_folder>/<identifier>/<prefix>_*.vtu``.
Summary is written to ``<base_folder>/<identifier>.pvd``.
Manual VTK callbacks can be called at any time to take a snapshot
of the current state of the EK species.
Automatic VTK callbacks can be disabled at any time and re-enabled later.
Please note that the internal VTK counter is no longer incremented when
an automatic callback is disabled, which means the number of EK steps
between two frames will not always be an integer multiple of ``delta_N``.
Parameters
----------
identifier : :obj:`str`
Name of the VTK writer.
observables : :obj:`list`, {'density',}
List of observables to write to the VTK files.
delta_N : :obj:`int`
Write frequency. If this value is 0 (default), the object is a
manual VTK callback that must be triggered manually. Otherwise,
it is an automatic callback that is added to the time loop and
writes every ``delta_N`` EK steps.
base_folder : :obj:`str` (optional), default is 'vtk_out'
Path to the output VTK folder.
prefix : :obj:`str` (optional), default is 'simulation_step'
Prefix for VTK files.
"""
_so_name = "walberla::EKVTKHandle"
_so_creation_policy = "GLOBAL"
_so_bind_methods = ("enable", "disable", "write")
[docs]
def required_keys(self):
return self.valid_keys() - self.default_params().keys()
def __repr__(self):
class_id = f"{self.__class__.__module__}.{self.__class__.__name__}"
if self.delta_N:
write_when = f"every {self.delta_N} EK steps"
if not self.enabled:
write_when += " (disabled)"
else:
write_when = "on demand"
return f"<{class_id}: write to '{self.vtk_uid}' {write_when}>"
[docs]
@script_interface_register
class EKReactant(ScriptInterfaceHelper):
"""
Reactant-object which specifies the contribution of a species to a reaction.
Parameters
----------
ekspecies : :obj:`~espressomd.electrokinetics.EKSpecies`
EK species to react
stoech_coeff: :obj:`float`
Stoechiometric coefficient of this species in the reaction.
Products have positive coefficients whereas educts have negative ones.
order: :obj:`float`
Partial-order of this species in the reaction.
"""
_so_name = "walberla::EKReactant"
_so_creation_policy = "GLOBAL"
[docs]
class EKBulkReaction(ScriptInterfaceHelper):
"""
Reaction type that is applied everywhere in the domain.
Parameters
----------
lattice : :obj:`espressomd.electrokinetics.LatticeWalberla <espressomd.detail.walberla.LatticeWalberla>`
Lattice object.
tau : :obj:`float`
EK time step, must be an integer multiple of the MD time step.
coefficient : :obj:`float`
Reaction rate constant of the reaction.
reactants: array_like of :obj:`~espressomd.electrokinetics.EKReactant`
Reactants that participate this reaction.
"""
_so_name = "walberla::EKBulkReaction"
_so_creation_policy = "GLOBAL"
[docs]
class EKIndexedReaction(ScriptInterfaceHelper):
"""
Reaction type that is applied only on specific cells in the domain.
Can be used to model surface-reactions.
Parameters
----------
lattice : :obj:`espressomd.electrokinetics.LatticeWalberla <espressomd.detail.walberla.LatticeWalberla>`
Lattice object.
tau : :obj:`float`
EK time step, must be an integer multiple of the MD time step.
coefficient : :obj:`float`
Reaction rate constant of the reaction.
reactants: array_like of :obj:`~espressomd.electrokinetics.EKReactant`
Reactants that participate this reaction.
"""
_so_name = "walberla::EKIndexedReaction"
_so_creation_policy = "GLOBAL"
[docs]
def add_node_to_index(self, node):
self.call_method("set_node_is_boundary", node=node, is_boundary=True)
[docs]
def remove_node_from_index(self, node):
self.call_method("set_node_is_boundary", node=node, is_boundary=False)
def __getitem__(self, key):
if isinstance(key, (tuple, list, np.ndarray)) and len(key) == 3:
if any(isinstance(typ, slice) for typ in key):
shape = self.shape
indices = [np.atleast_1d(np.arange(shape[i])[key[i]])
for i in range(3)]
dimensions = [ind.size for ind in indices]
value_grid = np.zeros((*dimensions,), dtype=bool)
indices = itertools.product(*map(enumerate, indices))
for (i, x), (j, y), (k, z) in indices:
value_grid[i, j, k] = self.call_method(
"get_node_is_boundary", node=(x, y, z))
return utils.array_locked(value_grid)
else:
return self.call_method("get_node_is_boundary", node=key)
raise TypeError(
f"{key} is not a valid index. Should be a point on the nodegrid or a slice")
def __setitem__(self, key, values):
if isinstance(key, (tuple, list, np.ndarray)) and len(key) == 3:
if any(isinstance(typ, slice) for typ in key):
shape = self.shape
indices = [np.atleast_1d(np.arange(shape[i])[key[i]])
for i in range(3)]
dimensions = tuple(ind.size for ind in indices)
values = np.copy(values)
# broadcast if only one element was provided
if values.shape == ():
values = np.full(dimensions, values)
if values.shape != dimensions:
raise ValueError(
f"Input-dimensions of array {values.shape} does not match slice dimensions {dimensions}.")
indices = itertools.product(*map(enumerate, indices))
for (i, x), (j, y), (k, z) in indices:
self.call_method("set_node_is_boundary", node=(
x, y, z), is_boundary=bool(values[i, j, k]))
else:
return self.call_method(
"set_node_is_boundary", node=key, is_boundary=values)
else:
raise TypeError(
f"{key} is not a valid index. Should be a point on the nodegrid or a slice")
[docs]
@script_interface_register
class EKReactions(ScriptObjectList):
"""
Container object holding all EK-reactions that are considered.
Methods
-------
clear()
Remove all reactions.
"""
_so_name = "walberla::EKReactions"
_so_creation_policy = "GLOBAL"
_so_bind_methods = ("clear",)
[docs]
def add(self, reaction):
"""
Add a reaction to the container.
Parameters
----------
reaction : :obj:`~espressomd.electrokinetics.EKBulkReaction` or :obj:`~espressomd.electrokinetics.EKIndexedReaction`
Reaction to be added.
"""
if not isinstance(reaction, (EKBulkReaction, EKIndexedReaction)):
raise TypeError("reaction object is not of correct type.")
self.call_method("add", object=reaction)
[docs]
def remove(self, reaction):
"""
Remove a reaction from the container.
Parameters
----------
reaction : :obj:`~espressomd.electrokinetics.EKBulkReaction` or :obj:`~espressomd.electrokinetics.EKIndexedReaction`
Reaction to be removed.
"""
self.call_method("remove", object=reaction)
[docs]
@script_interface_register
class EKContainer(ScriptObjectList):
"""
Container object holding the :obj:`~espressomd.electrokinetics.EKSpecies`.
Parameters
----------
tau : :obj:`float`
EK time step, must be an integer multiple of the MD time step.
solver : :obj:`~espressomd.electrokinetics.EKNone` or :obj:`~espressomd.electrokinetics.EKFFT`
Solver defining the treatment of the electrostatic Poisson-equation.
Methods
-------
clear()
Removes all species.
"""
_so_name = "walberla::EKContainer"
_so_creation_policy = "GLOBAL"
_so_features = ("WALBERLA",)
_so_bind_methods = ("clear",)
def __init__(self, *args, **kwargs):
if "sip" not in kwargs:
kwargs["reactions"] = EKReactions()
super().__init__(*args, **kwargs)
else:
super().__init__(**kwargs)
@property
def reactions(self):
"""
Returns
-------
Reactions-container of the current reactions (:obj:`~espressomd.electrokinetics.EKReactions`).
"""
return self._getter("reactions")
[docs]
def add(self, ekspecies):
"""
Add an :obj:`~espressomd.electrokinetics.EKSpecies` to the container.
Parameters
----------
ekspecies : :obj:`~espressomd.electrokinetics.EKSpecies`
Species to be added.
"""
self.call_method("add", object=ekspecies)
[docs]
def remove(self, ekspecies):
"""
Remove an :obj:`~espressomd.electrokinetics.EKSpecies` from the container.
Parameters
----------
ekspecies : :obj:`~espressomd.electrokinetics.EKSpecies`
Species to be removed.
"""
self.call_method("remove", object=ekspecies)