src/lbmpy/db.py

+import json
+import six
+import inspect
+
+from pystencils.runhelper.db import PystencilsJsonEncoder
+from pystencils.simp import SimplificationStrategy
+from lbmpy import LBStencil, Method, CollisionSpace, SubgridScaleModel
+from lbmpy.creationfunctions import LBMConfig, LBMOptimisation
+from lbmpy.methods import CollisionSpaceInfo
+from lbmpy.forcemodels import AbstractForceModel
+from lbmpy.non_newtonian_models import CassonsParameters
+
+
+class LbmpyJsonEncoder(PystencilsJsonEncoder):
+
+    def default(self, obj):
+        if isinstance(obj, (LBMConfig, LBMOptimisation, CollisionSpaceInfo, CassonsParameters)):
+            return obj.__dict__
+        if isinstance(obj, (LBStencil, Method, CollisionSpace, SubgridScaleModel)):
+            return obj.name
+        if isinstance(obj, AbstractForceModel):
+            return obj.__class__.__name__
+        if isinstance(obj, SimplificationStrategy):
+            return obj.__str__()
+        if inspect.isclass(obj):
+            return obj.__name__
+        return PystencilsJsonEncoder.default(self, obj)
+
+
+class LbmpyJsonSerializer(object):
+
+    @classmethod
+    def serialize(cls, data):
+        if six.PY3:
+            if isinstance(data, bytes):
+                return json.dumps(data.decode('utf-8'), cls=LbmpyJsonEncoder, ensure_ascii=False).encode('utf-8')
+            else:
+                return json.dumps(data, cls=LbmpyJsonEncoder, ensure_ascii=False).encode('utf-8')
+        else:
+            return json.dumps(data, cls=LbmpyJsonEncoder, ensure_ascii=False).encode('utf-8')
+
+    @classmethod
+    def deserialize(cls, data):
+        if six.PY3:
+            return json.loads(data.decode('utf-8'))
+        else:
+            return json.loads(data.decode('utf-8'))

src/lbmpy/enums.py

+from enum import Enum, auto
+
+
+class Stencil(Enum):
+    """
+    The Stencil enumeration represents all possible lattice Boltzmann stencils that are available in lbmpy.
+    It should be passed to :class:`lbmpy.stencils.LBStenil`. This class then creates a stencils representation
+    containing the concrete neighbour directions as a tuple of tuples.
+
+    The number of spatial dimensions *d* and the number of discrete velocities *q* are stated in the DdQq notation
+    """
+    D2Q9 = auto()
+    """
+    A two dimensional stencil using 9 discrete velocities.
+    """
+    D2V17 = auto()
+    """
+    A two dimensional stencil using 17 discrete velocities. (long range stencil).
+    """
+    D2V37 = auto()
+    """
+    A two dimensional stencil using 37 discrete velocities. (long range stencil).
+    """
+    D3Q7 = auto()
+    """
+    A three dimensional stencil using 7 discrete velocities.
+    """
+    D3Q15 = auto()
+    """
+    A three dimensional stencil using 15 discrete velocities.
+    """
+    D3Q19 = auto()
+    """
+    A three dimensional stencil using 19 discrete velocities.
+    """
+    D3Q27 = auto()
+    """
+    A three dimensional stencil using 27 discrete velocities.
+    """
+
+
+class Method(Enum):
+    """
+    The Method enumeration represents all possible lattice Boltzmann collision operators that are available in lbmpy.
+    It should be passed to :class:`lbmpy.creationfunctions.LBMConfig`. The LBM configuration *dataclass* then derives
+    the respective collision equations when passed to the creations functions in the `lbmpy.creationfunctions`
+    module of lbmpy.
+
+    Note here, when using a specific enumeration to derive a particular LBM collision operator,
+    different parameters of the :class:`lbmpy.creationfunctions.LBMConfig` might become necessary.
+    For example, it does not make sense to define *relaxation_rates* for a single relaxation rate method, which
+    is essential for multiple relaxation rate methods. Important specific parameters are listed below to the enum value.
+    A specific creation function is stated for each case which explains these parameters in detail.
+    """
+    SRT = auto()
+    """
+    See :func:`lbmpy.methods.create_srt`, 
+    Single relaxation time method
+    """
+    TRT = auto()
+    """
+    See :func:`lbmpy.methods.create_trt`, 
+    Two relaxation time, the first relaxation rate is for even moments and determines the
+    viscosity (as in SRT). The second relaxation rate is used for relaxing odd moments and controls the
+    bulk viscosity. For details in the TRT collision operator see :cite:`TRT`
+    """
+    MRT_RAW = auto()
+    """
+    See :func:`lbmpy.methods.create_mrt_raw`, 
+    Non-orthogonal MRT where all relaxation rates can be specified independently, i.e. there are as many relaxation 
+    rates as stencil entries. Look at the generated method in Jupyter to see which moment<->relaxation rate mapping.
+    Originally defined in :cite:`raw_moments`
+    """
+    MRT = auto()
+    """
+    See :func:`lbmpy.methods.create_mrt_orthogonal`
+    Orthogonal multi relaxation time model, relaxation rates are used in this order for *shear modes*, *bulk modes*,
+    *third-order modes*, *fourth-order modes*, etc. Requires also a parameter *weighted* that should be `True` if the
+    moments should be orthogonal w.r.t. weighted scalar product using the lattice weights. If `False`, the normal
+    scalar product is used. For custom definition of the method, a *nested_moments* can be passed.
+    For example: [ [1, x, y], [x*y, x**2, y**2], ... ] that groups all moments together that should be relaxed
+    at the same rate. Literature values of this list can be obtained through 
+    :func:`lbmpy.methods.creationfunctions.mrt_orthogonal_modes_literature`.
+    WMRT collision operators are reported to be numerically more stable and more accurate, 
+    whilst also having a lower computational cos :cite:`FAKHARI201722`
+    """
+    CENTRAL_MOMENT = auto()
+    """
+    See :func:`lbmpy.methods.create_central_moment`
+    Creates moment based LB method where the collision takes place in the central moment space. By default, 
+    a raw-moment set is used where the bulk and the shear viscosity are separated. An original derivation can be 
+    found in :cite:`Geier2006`
+    """
+    TRT_KBC_N1 = auto()
+    """
+    See :func:`lbmpy.methods.create_trt_kbc`
+    Particular two-relaxation rate method. This is not the entropic method yet, only the relaxation pattern. 
+    To get the entropic method also *entropic* needs to be set to `True`. 
+    There are four KBC methods available in lbmpy. The naming is according to :cite:`karlin2015entropic`
+    """
+    TRT_KBC_N2 = auto()
+    """
+    See :func:`lbmpy.methods.create_trt_kbc`
+    Particular two-relaxation rate method. This is not the entropic method yet, only the relaxation pattern. 
+    To get the entropic method also *entropic* needs to be set to `True`. 
+    There are four KBC methods available in lbmpy. The naming is according to :cite:`karlin2015entropic`
+    """
+    TRT_KBC_N3 = auto()
+    """
+    See :func:`lbmpy.methods.create_trt_kbc`
+    Particular two-relaxation rate method. This is not the entropic method yet, only the relaxation pattern. 
+    To get the entropic method also *entropic* needs to be set to `True`. 
+    There are four KBC methods available in lbmpy. The naming is according to :cite:`karlin2015entropic`
+    """
+    TRT_KBC_N4 = auto()
+    """
+    See :func:`lbmpy.methods.create_trt_kbc`
+    Particular two-relaxation rate method. This is not the entropic method yet, only the relaxation pattern. 
+    To get the entropic method also *entropic* needs to be set to `True`. 
+    There are four KBC methods available in lbmpy. The naming is according to :cite:`karlin2015entropic`
+    """
+    CUMULANT = auto()
+    """
+    See :func:`lbmpy.methods.create_with_default_polynomial_cumulants`
+    Cumulant-based LB method which relaxes groups of polynomial cumulants chosen to optimize rotational invariance.
+    For details on the method see :cite:`geier2015`
+    """
+    MONOMIAL_CUMULANT = auto()
+    """
+    See :func:`lbmpy.methods.create_with_monomial_cumulants`
+    Cumulant-based LB method which relaxes monomial cumulants.
+    For details on the method see :cite:`geier2015` and :cite:`Coreixas2019`
+    """
+
+
+class CollisionSpace(Enum):
+    """
+    The CollisionSpace enumeration lists all possible spaces for collision to occur in.
+    """
+    POPULATIONS = auto()
+    """
+    Population space, meaning post-collision populations are obtained directly by relaxation of linear combinations of
+    pre-collision populations. Default for `lbmpy.enums.Method.SRT` and `lbmpy.enums.Method.TRT`.
+    Results in the creation of an instance of :class:`lbmpy.methods.momentbased.MomentBasedLbMethod`.
+    """
+    RAW_MOMENTS = auto()
+    """
+    Raw moment space, meaning relaxation is applied to a set of linearly independent, polynomial raw moments of the 
+    discrete population vector. Default for `lbmpy.enums.Method.MRT`.
+    Results in the creation of an instance of :class:`lbmpy.methods.momentbased.MomentBasedLbMethod`.
+    """
+    CENTRAL_MOMENTS = auto()
+    """
+    Central moment space, meaning relaxation is applied to a set of linearly independent, polynomial central moments 
+    of the discrete population vector. Default for `lbmpy.enums.Method.CENTRAL_MOMENT`.
+    Results in the creation of an instance of :class:`lbmpy.methods.momentbased.CentralMomentBasedLbMethod`.
+    """
+    CUMULANTS = auto()
+    """
+    Cumulant space, meaning relaxation is applied to a set of linearly independent, polynomial cumulants of the
+    discrete population vector. Default for `lbmpy.enums.Method.CUMULANT` and `lbmpy.enums.Method.MONOMIAL_CUMULANT`.
+    Results in the creation of an instance of :class:`lbmpy.methods.cumulantbased.CumulantBasedLbMethod`.
+    """
+
+    def compatible(self, method: Method):
+        """Determines if the given `lbmpy.enums.Method` is compatible with this collision space."""
+        compat_dict = {
+            CollisionSpace.POPULATIONS: {Method.SRT, Method.TRT, Method.MRT_RAW, Method.MRT,
+                                         Method.TRT_KBC_N1, Method.TRT_KBC_N2, Method.TRT_KBC_N3, Method.TRT_KBC_N4},
+            CollisionSpace.RAW_MOMENTS: {Method.SRT, Method.TRT, Method.MRT_RAW, Method.MRT},
+            CollisionSpace.CENTRAL_MOMENTS: {Method.CENTRAL_MOMENT},
+            CollisionSpace.CUMULANTS: {Method.MONOMIAL_CUMULANT, Method.CUMULANT}
+        }
+
+        return method in compat_dict[self]
+
+
+class ForceModel(Enum):
+    """
+    The ForceModel enumeration defines which force model is used to introduce forcing terms in the collision operator
+    of the lattice Boltzmann method. A short summary of the theory behind is shown in `lbmpy.forcemodels`.
+    More precise definitions are given in Chapter 6 and 10 of :cite:`lbm_book`
+    """
+    SIMPLE = auto()
+    """
+    See :class:`lbmpy.forcemodels.Simple`
+    """
+    LUO = auto()
+    """
+    See :class:`lbmpy.forcemodels.Luo`
+    """
+    GUO = auto()
+    """
+    See :class:`lbmpy.forcemodels.Guo`
+    """
+    BUICK = auto()
+    """
+    See :class:`lbmpy.forcemodels.Buick`
+    """
+    SILVA = auto()
+    """
+    See :class:`lbmpy.forcemodels.Buick`
+    """
+    EDM = auto()
+    """
+    See :class:`lbmpy.forcemodels.EDM`
+    """
+    KUPERSHTOKH = auto()
+    """
+    See :class:`lbmpy.forcemodels.EDM`
+    """
+    HE = auto()
+    """
+    See :class:`lbmpy.forcemodels.He`
+    """
+    SHANCHEN = auto()
+    """
+    See :class:`lbmpy.forcemodels.ShanChen`
+    """
+    CENTRALMOMENT = auto()
+    """
+    See :class:`lbmpy.forcemodels.CentralMoment`
+    """
+
+
+class SubgridScaleModel(Enum):
+    """
+    The SubgridScaleModel enumeration defines which subgrid-scale model (SGS) is used to perform
+    Large-Eddy-Simulations (LES).
+    """
+    SMAGORINSKY = auto()
+    """
+    See :func:`lbmpy.turbulence_models.add_smagorinsky_model`
+    """
+    QR = auto()
+    """
+    See :func:`lbmpy.turbulence_models.add_qr_model`
+    """

src/lbmpy/equilibrium/__init__.py

+r"""
+This module contains various classes encapsulating equilibrium distributions used in the lattice Boltzmann
+method. These include both the continuous and the discretized variants of the Maxwellian equilibrium of
+hydrodynamics. Furthermore, a lightweight wrapper class for custom discrete equilibria is provided.
+Custom equilibria may also be implemented by manually overriding the abstract base class 
+:class:`lbmpy.equilibrium.AbstractEquilibrium`.
+"""
+
+from .abstract_equilibrium import AbstractEquilibrium
+from .continuous_hydro_maxwellian import ContinuousHydrodynamicMaxwellian, default_background_distribution
+from .generic_discrete_equilibrium import GenericDiscreteEquilibrium, discrete_equilibrium_from_matching_moments
+from .discrete_hydro_maxwellian import DiscreteHydrodynamicMaxwellian
+
+__all__ = [
+    "AbstractEquilibrium",
+    "ContinuousHydrodynamicMaxwellian", "default_background_distribution",
+    "GenericDiscreteEquilibrium", "discrete_equilibrium_from_matching_moments",
+    "DiscreteHydrodynamicMaxwellian"
+]

src/lbmpy/equilibrium/abstract_equilibrium.py

+from abc import ABC, abstractmethod
+import sympy as sp
+from pystencils.cache import sharedmethodcache
+
+from lbmpy.moments import polynomial_to_exponent_representation
+
+
+class AbstractEquilibrium(ABC):
+    """
+    Abstract Base Class for description of equilibrium distribution functions used in lattice
+    Boltzmann methods.
+
+    **Equilibrium Representation:**
+
+    This class provides the common interface for describing equilibrium distribution functions,
+    which is then used by the various method classes in the derivation of collision equations.
+    An equilibrium distribution is defined by either its continuous equation (see :attr:`continuous_equation`)
+    or a set of discrete populations
+    (see :attr:`discrete_populations` and :class:`lbmpy.equilibrium.GenericDiscreteEquilibrium`).
+    The distribution function may be given either in its regular, absolute form; or only as its
+    deviation from the rest state, represented by the background distribution (see :attr:`background_distribution`).
+
+    **Computation of Statistical Modes:**
+
+    The major computational task of an equilbrium class is the computation of the distribution's
+    statistical modes. For discrete distributions, the subclass :class:`lbmpy.equilibrium.GenericDiscreteEquilibrium`
+    provides a generic implementation for their computation. For continuous distributions, computation
+    of raw moments, central moments, and cumulants is more complicated, but may also be simplified using special
+    tricks.
+
+    As the computation of statistical modes is a time-consuming process, the abstract base class provides caching
+    functionality to avoid recomputing quantities that are already known.
+
+    **Instructions to Override:**
+
+    If you wish to model a simple custom discrete distribution, just using the class
+    :class:`lbmpy.equilibrium.GenericDiscreteEquilibrium` might already be sufficient.
+    If, however, you need to implement more specific functionality, custom properties, 
+    a background distribution, etc., or if you wish to model a continuous distribution,
+    you will have to set up a custom subclass of :class:`AbstractEquilibrium`.
+
+    A subclass must implement all abstract properties according to their docstrings.
+    For computation of statistical modes, a large part of the infrastructure is already given in the abstract base
+    class. The public interface for computing e.g. raw moments reduces the computation of polynomial moments to their
+    contained monomials (for details on how moments are represented in *lbmpy*, see :mod:`lbmpy.moments`). The values
+    of both polynomial and monomial moments, once computed, will be cached per instance of the equilibrium class.
+    To take full advantage of the caching functionality, you will have to override only :func:`_monomial_raw_moment`
+    and its central moment and cumulant counterparts. These methods will be called only once for each monomial quantity
+    when it is required for the first time. Afterward, the cached value will be used.
+    """
+
+    def __init__(self, dim=3):
+        self._dim = dim
+
+    @property
+    def dim(self):
+        """This distribution's spatial dimensionality."""
+        return self._dim
+
+    #   -------------- Abstract Properties, to be overridden in subclass ----------------------------------------------
+
+    @property
+    @abstractmethod
+    def deviation_only(self):
+        """Whether or not this equilibrium distribution is represented only by its deviation
+        from the background distribution."""
+        raise NotImplementedError("'deviation_only' must be provided by subclass.")
+
+    @property
+    @abstractmethod
+    def continuous_equation(self):
+        """Returns the continuous equation defining this equilibrium distribution,
+        or `None` if no such equation is available."""
+        raise NotImplementedError("'continuous_equation' must be provided by subclass.")
+
+    @property
+    @abstractmethod
+    def discrete_populations(self):
+        """Returns the discrete populations of this equilibrium distribution as a tuple, 
+        or `None` if none are available."""
+        raise NotImplementedError("'discrete_populations' must be provided by subclass.")
+
+    @property
+    @abstractmethod
+    def background_distribution(self):
+        """Returns this equilibrium distribution's background distribution, which is
+        the distribution the discrete populations are centered around in the case of
+        zero-centered storage. If no background distribution is available, `None` must be 
+        returned."""
+        raise NotImplementedError("'background_distribution' must be provided by subclass.")
+
+    @property
+    @abstractmethod
+    def zeroth_order_moment_symbol(self):
+        """Returns a symbol referring to the zeroth-order moment of this distribution,
+        which is the area under it's curve."""
+        raise NotImplementedError("'zeroth_order_moment' must be provided by subclass.")
+
+    @property
+    @abstractmethod
+    def first_order_moment_symbols(self):
+        """Returns a vector of symbols referring to the first-order moment of this distribution,
+        which is its mean value."""
+        raise NotImplementedError("'first_order_moments' must be provided by subclass.")
+
+    #   -------------- Statistical Modes Interface --------------------------------------------------------------------
+
+    @sharedmethodcache("_moment_cache")
+    def moment(self, exponent_tuple_or_polynomial):
+        """Returns this equilibrium distribution's moment specified by ``exponent_tuple_or_polynomial``.
+
+        Args:
+            exponent_tuple_or_polynomial: Moment specification, see :mod:`lbmpy.moments`.
+        """
+        monomials = []
+        if isinstance(exponent_tuple_or_polynomial, tuple):
+            monomials = [(1, exponent_tuple_or_polynomial)]
+        else:
+            monomials = polynomial_to_exponent_representation(exponent_tuple_or_polynomial, dim=self._dim)
+
+        moment_value = sp.Integer(0)
+        for coeff, moment in monomials:
+            moment_value += coeff * self._cached_monomial_raw_moment(moment)
+
+        return moment_value.expand()
+
+    def moments(self, exponent_tuples_or_polynomials):
+        """Returns a tuple of this equilibrium distribution's moments specified by 'exponent_tuple_or_polynomial'.
+
+        Args:
+            exponent_tuples_or_polynomials: Sequence of moment specifications, see :mod:`lbmpy.moments`.
+        """
+        return tuple(self.moment(m) for m in exponent_tuples_or_polynomials)
+
+    @sharedmethodcache("_central_moment_cache")
+    def central_moment(self, exponent_tuple_or_polynomial, frame_of_reference):
+        """Returns this equilibrium distribution's central moment specified by
+        ``exponent_tuple_or_polynomial``, computed according to the given ``frame_of_reference``.
+
+        Args:
+            exponent_tuple_or_polynomial: Moment specification, see :mod:`lbmpy.moments`.
+            frame_of_reference: The frame of reference with respect to which the central moment should be computed.
+        """
+        monomials = []
+        if isinstance(exponent_tuple_or_polynomial, tuple):
+            monomials = [(1, exponent_tuple_or_polynomial)]
+        else:
+            monomials = polynomial_to_exponent_representation(exponent_tuple_or_polynomial, dim=self._dim)
+
+        moment_value = sp.Integer(0)
+        for coeff, moment in monomials:
+            moment_value += coeff * self._cached_monomial_central_moment(moment, frame_of_reference)
+
+        return moment_value.expand()
+
+    def central_moments(self, exponent_tuples_or_polynomials, frame_of_reference):
+        """Returns a list this equilibrium distribution's central moments specified by
+        ``exponent_tuples_or_polynomials``, computed according to the given ``frame_of_reference``.
+
+        Args:
+            exponent_tuples_or_polynomials: Sequence of moment specifications, see :mod:`lbmpy.moments`.
+            frame_of_reference: The frame of reference with respect to which the central moment should be computed.
+        """
+        return tuple(self.central_moment(m, frame_of_reference) for m in exponent_tuples_or_polynomials)
+
+    @sharedmethodcache("_cumulant_cache")
+    def cumulant(self, exponent_tuple_or_polynomial, rescale=True):
+        """Returns this equilibrium distribution's cumulant specified by ``exponent_tuple_or_polynomial``.
+
+        Args:
+            exponent_tuple_or_polynomial: Moment specification, see :mod:`lbmpy.moments`.
+            rescale: If ``True``, the cumulant value should be multiplied by the zeroth-order moment.
+        """
+        monomials = []
+        if isinstance(exponent_tuple_or_polynomial, tuple):
+            monomials = [(1, exponent_tuple_or_polynomial)]
+        else:
+            monomials = polynomial_to_exponent_representation(exponent_tuple_or_polynomial, dim=self._dim)
+
+        cumulant_value = sp.Integer(0)
+        for coeff, moment in monomials:
+            cumulant_value += coeff * self._cached_monomial_cumulant(moment, rescale=rescale)
+
+        return cumulant_value.expand()
+
+    def cumulants(self, exponent_tuples_or_polynomials, rescale=True):
+        """Returns a list of this equilibrium distribution's cumulants specified by ``exponent_tuples_or_polynomial``.
+
+        Args:
+            exponent_tuples_or_polynomials: Sequence of moment specifications, see :mod:`lbmpy.moments`.
+            rescale: If ``True``, the cumulant value should be multiplied by the zeroth-order moment.
+        """
+        return tuple(self.cumulant(m, rescale) for m in exponent_tuples_or_polynomials)
+
+    #   -------------- Monomial moment computation, to be overridden in subclass --------------------------------------
+
+    @abstractmethod
+    def _monomial_raw_moment(self, exponents):
+        """See :func:`lbmpy.equilibrium.AbstractEquilibrium.moment`."""
+        raise NotImplementedError("'_monomial_raw_moment' must be implemented by a subclass.")
+
+    @abstractmethod
+    def _monomial_central_moment(self, exponents, frame_of_reference):
+        """See :func:`lbmpy.equilibrium.AbstractEquilibrium.central_moment`."""
+        raise NotImplementedError("'_monomial_central_moment' must be implemented by a subclass.")
+
+    @abstractmethod
+    def _monomial_cumulant(self, exponents, rescale):
+        """See :func:`lbmpy.equilibrium.AbstractEquilibrium.cumulant`."""
+        raise NotImplementedError("'_monomial_cumulant' must be implemented by a subclass.")
+
+    #   -------------- Cached monomial moment computation methods -----------------------------------------------------
+
+    @sharedmethodcache("_moment_cache")
+    def _cached_monomial_raw_moment(self, exponents):
+        return self._monomial_raw_moment(exponents)
+
+    @sharedmethodcache("_central_moment_cache")
+    def _cached_monomial_central_moment(self, exponents, frame_of_reference):
+        return self._monomial_central_moment(exponents, frame_of_reference)
+
+    @sharedmethodcache("_cumulant_cache")
+    def _cached_monomial_cumulant(self, exponents, rescale):
+        return self._monomial_cumulant(exponents, rescale)
+
+    #   -------------- HTML Representation ----------------------------------------------------------------------------
+
+    def _repr_html_(self):
+        html = f"""
+        <table style="border:none; width: 100%">
+            <tr>
+                <th colspan="2" style="text-align: left">
+                    Instance of {self.__class__.__name__}
+                </th>
+            </tr>
+        """
+
+        cont_eq = self.continuous_equation
+        if cont_eq is not None:
+            html += f"""
+                <tr>
+                    <td> Continuous Equation: </td>
+                    <td style="text-align: center">
+                        ${sp.latex(self.continuous_equation)}$
+                    </td>
+                </tr>
+            """
+        else:
+            pdfs = self.discrete_populations
+            if pdfs is not None:
+                html += """
+                    <tr>
+                        <td colspan="2" style="text-align: right;"> Discrete Populations: </td>
+                    </tr>
+                    """
+                for f, eq in zip(sp.symbols(f"f_:{len(pdfs)}"), pdfs):
+                    html += f'<tr><td colspan="2" style="text-align: left;"> ${f} = {sp.latex(eq)}$ </td></tr>'
+
+        html += "</table>"
+
+        return html
+
+#   end class AbstractEquilibrium

src/lbmpy/equilibrium/continuous_hydro_maxwellian.py

+import sympy as sp
+
+from .abstract_equilibrium import AbstractEquilibrium
+
+from lbmpy.moments import contained_moments
+from lbmpy.maxwellian_equilibrium import continuous_maxwellian_equilibrium
+from lbmpy.continuous_distribution_measures import continuous_moment, continuous_cumulant
+
+from pystencils.sympyextensions import remove_higher_order_terms, simplify_by_equality
+
+
+def default_background_distribution(dim):
+    return ContinuousHydrodynamicMaxwellian(dim=dim, compressible=True, deviation_only=False,
+                                            rho=sp.Integer(1), delta_rho=0, u=(0,) * dim,
+                                            c_s_sq=sp.Rational(1, 3))
+
+
+class ContinuousHydrodynamicMaxwellian(AbstractEquilibrium):
+    r"""
+        The standard continuous Maxwellian equilibrium distribution for hydrodynamics.
+
+        This class represents the Maxwellian equilibrium distribution of hydrodynamics in its continuous form
+        in :math:`d` dimensions :cite:`lbm_book`:
+        
+        .. math::
+
+            \Psi \left( \rho, \mathbf{u}, \mathbf{\xi} \right)
+             = \rho \left( \frac{1}{2 \pi c_s^2} \right)^{d/2} 
+                    \exp \left( \frac{- (\mathbf{\xi} - \mathbf{u})^2 }{2 c_s^2} \right)
+
+        Beyond this classic, 'compressible' form of the equilibrium, an alternative form known as the
+        incompressible equilibrium of the LBM can be obtained by setting the flag ``compressible=False``.
+        The continuous incompressible equilibrium can be expressed as
+        (:cite:`HeIncompressible,GruszczynskiCascadedPhaseFieldModel`):
+
+        .. math::
+
+            \Psi^{\mathrm{incomp}} \left( \rho, \mathbf{u}, \mathbf{\xi} \right)
+                = \Psi \left( \rho_0, \mathbf{u}, \mathbf{\xi} \right) 
+                  + \Psi \left( \delta\rho, \mathbf{0}, \mathbf{\xi} \right)
+
+        Here, :math:`\rho_0` (typically :math:`\rho_0 = 1`) denotes the background density, and :math:`\delta\rho` is
+        the density deviation, such that the total fluid density amounts to :math:`\rho = \rho_0 + \delta\rho`.
+
+        To simplify computations when the zero-centered storage format is used for PDFs, both equilibrium variants can
+        also be expressed in a *deviation-only* or *delta-equilibrium* form, which is obtained by subtracting the
+        constant background distribution :math:`\Psi (\rho_0, \mathbf{0})`. The delta form expresses the equilibrium
+        distribution only by its deviation from the rest state:
+
+        .. math::
+
+            \delta\Psi \left( \rho, \mathbf{u}, \mathbf{\xi} \right)
+                &=   \Psi \left( \rho, \mathbf{u}, \mathbf{\xi} \right) 
+                    - \Psi \left( \rho_0, \mathbf{0}, \mathbf{\xi} \right) \\
+            \delta\Psi^{\mathrm{incomp}} \left( \rho, \mathbf{u}, \mathbf{\xi} \right)
+                &=   \Psi^{\mathrm{incomp}} \left( \rho, \mathbf{u}, \mathbf{\xi} \right) 
+                    - \Psi \left( \rho_0, \mathbf{0}, \mathbf{\xi} \right)
+
+        Parameters:
+            dim: Spatial dimensionality
+            compressible: If `False`, the incompressible equilibrium is created
+            deviation_only: If `True`, the delta-equilibrium is created
+            order: The discretization order in velocity to which computed statistical modes should be truncated
+            rho: Symbol or value for the density
+            rho_background: Symbol or value for the background density
+            delta_rho: Symbol or value for the density deviation
+            u: Sequence of symbols for the macroscopic velocity
+            v: Sequence of symbols for the particle velocity :math:`\xi`
+            c_s_sq: Symbol or value for the squared speed of sound
+    """
+
+    def __init__(self, dim=3, compressible=True, deviation_only=False,
+                 order=None,
+                 rho=sp.Symbol("rho"),
+                 rho_background=sp.Integer(1),
+                 delta_rho=sp.Symbol("delta_rho"),
+                 u=sp.symbols("u_:3"),
+                 v=sp.symbols("v_:3"),
+                 c_s_sq=sp.Symbol("c_s") ** 2):
+        super().__init__(dim=dim)
+
+        self._order = order
+        self._compressible = compressible
+        self._deviation_only = deviation_only
+        self._rho = rho
+        self._rho_background = rho_background
+        self._delta_rho = delta_rho
+        self._u = u[:dim]
+        self._v = v[:dim]
+
+        # trick to speed up sympy integration (otherwise it takes multiple minutes, or aborts):
+        # use a positive, real symbol to represent c_s_sq -> then replace this symbol afterwards with the real c_s_sq
+        # (see maxwellian_equilibrium.py)
+        self._c_s_sq = c_s_sq
+        self._c_s_sq_helper = sp.Symbol("csq_helper", positive=True, real=True)
+
+        def psi(rho, u):
+            return continuous_maxwellian_equilibrium(dim=self._dim,
+                                                     rho=rho,
+                                                     u=u,
+                                                     v=self._v,
+                                                     c_s_sq=self._c_s_sq_helper)
+
+        zeroth_moment_arg = self._rho if self._compressible else self._rho_background
+        self._base_equation = psi(zeroth_moment_arg, self._u)
+
+        self._corrections = []
+        
+        if not self._compressible:
+            zeroth_order_correction = psi(self._delta_rho, (sp.Integer(0), ) * self._dim)
+            self._corrections.append((sp.Integer(1), zeroth_order_correction))
+
+        if self._deviation_only:
+            rest_state = psi(self._rho_background, (sp.Integer(0), ) * self._dim)
+            self._corrections.append((sp.Integer(-1), rest_state))
+
+    @property
+    def order(self):
+        return self._order
+
+    @property
+    def deviation_only(self):
+        return self._deviation_only
+
+    @property
+    def compressible(self):
+        return self._compressible
+
+    @property
+    def density(self):
+        return self._rho
+
+    @property
+    def background_density(self):
+        return self._rho_background
+
+    @property
+    def density_deviation(self):
+        return self._delta_rho
+
+    @property
+    def velocity(self):
+        return self._u
+
+    @property
+    def continuous_equation(self):
+        eq = self._base_equation + sum(f * e for f, e in self._corrections)
+        eq = eq.subs(self._c_s_sq_helper, self._c_s_sq)
+        return eq
+
+    @property
+    def zeroth_order_moment_symbol(self):
+        return self._delta_rho if self._deviation_only else self._rho
+
+    @property
+    def first_order_moment_symbols(self):
+        return self._u
+
+    @property
+    def background_distribution(self):
+        return ContinuousHydrodynamicMaxwellian(dim=self.dim, compressible=True, deviation_only=False,
+                                                order=self._order, rho=self._rho_background,
+                                                rho_background=self._rho_background,
+                                                delta_rho=0, u=(0,) * self.dim, v=self._v,
+                                                c_s_sq=self._c_s_sq)
+
+    @property
+    def discrete_populations(self):
+        return None
+
+    def central_moment(self, exponent_tuple_or_polynomial, velocity=None):
+        if velocity is None:
+            velocity = self._u
+
+        return super().central_moment(exponent_tuple_or_polynomial, velocity)
+
+    def central_moments(self, exponent_tuples_or_polynomials, velocity=None):
+        if velocity is None:
+            velocity = self._u
+
+        return super().central_moments(exponent_tuples_or_polynomials, velocity)
+
+    def cumulant(self, exponent_tuple_or_polynomial, rescale=True):
+        if not self._compressible or self._deviation_only:
+            raise Exception("Cumulants can only be computed for the compressible, "
+                            "non-deviation maxwellian equilibrium!")
+
+        return super().cumulant(exponent_tuple_or_polynomial, rescale=rescale)
+
+    #   ------------------ Overridden Moment Computation ------------------------------------------
+
+    def _monomial_raw_moment(self, exponents):
+        moment_value = continuous_moment(self._base_equation, exponents, self._v)
+
+        for coeff, corr in self._corrections:
+            moment_value += coeff * continuous_moment(corr, exponents, self._v)
+
+        moment_value = self._correct_order_and_cssq(moment_value)
+        moment_value = self._simplify_moment(moment_value)
+        return moment_value
+
+    def _monomial_central_moment(self, cm_exponents, velocity):
+        #   Setting up the central moment-generating function using SymPy integration
+        #   will take unfeasibly long at times
+        #   So we compute the central moments by binomial expansion in raw moments
+
+        cm_order = sum(cm_exponents)
+        contained_raw_moments = contained_moments(cm_exponents, exclude_original=False)
+        moment_value = sp.Integer(0)
+
+        for rm_exponents in contained_raw_moments:
+            rm_order = sum(rm_exponents)
+            factor = (-1)**(cm_order - rm_order)
+            factor *= sp.Mul(*(u**(c - i) * sp.binomial(c, i)
+                               for u, c, i in zip(velocity, cm_exponents, rm_exponents)))
+            rm_value = self._cached_monomial_raw_moment(rm_exponents)
+            moment_value += factor * rm_value
+
+        moment_value = self._correct_order_and_cssq(moment_value)
+        moment_value = self._simplify_moment(moment_value)
+        return moment_value
+
+    def _monomial_cumulant(self, c_exponents, rescale):
+        #   this implementation works only for the compressible, non-deviation equilibrium
+        cumulant_value = continuous_cumulant(self._base_equation, c_exponents, self._v)
+        cumulant_value = self._correct_order_and_cssq(cumulant_value)
+        if rescale:
+            cumulant_value = self._rho * cumulant_value
+        return cumulant_value
+
+    def _correct_order_and_cssq(self, term):
+        term = term.subs(self._c_s_sq_helper, self._c_s_sq)
+        term = term.expand()
+        if self._order is not None:
+            return remove_higher_order_terms(term, order=self._order, symbols=self._u)
+        else:
+            return term
+
+    def _simplify_moment(self, moment_value):
+        if (self.deviation_only or not self.compressible) \
+                and isinstance(self.density, sp.Symbol) and isinstance(self.density_deviation, sp.Symbol):
+            moment_value = simplify_by_equality(moment_value, self.density,
+                                                self.density_deviation, self.background_density)
+        return moment_value
+
+    #   ------------------ Utility ----------------------------------------------------------------
+
+    def __repr__(self):
+        return f"ContinuousHydrodynamicMaxwellian({self.dim}D, " \
+               f"compressible={self.compressible}, deviation_only:{self.deviation_only}" \
+               f"order={self.order})"
+
+    def _repr_html_(self):
+        def stylized_bool(b):
+            return "✓" if b else "✗"
+
+        html = f"""
+        <table style="border:none; width: 100%">
+            <tr>
+                <th colspan="3" style="text-align: left">
+                    Continuous Hydrodynamic Maxwellian Equilibrium
+                </th>
+                <td rowspan="2" style="width: 50%; text-align: center">
+                    $f ({sp.latex(self._rho)}, {sp.latex(self._u)}, {sp.latex(self._v)}) 
+                        = {sp.latex(self.continuous_equation)}$
+                </td>
+            </tr>
+            <tr>
+                <td>Compressible: {stylized_bool(self._compressible)}</td>
+                <td>Deviation Only: {stylized_bool(self._deviation_only)}</td>
+                <td>Order: {"&#8734;" if self._order is None else self._order}</td>
+            </tr>
+        </table>
+        """
+
+        return html
+
+#   end class ContinuousHydrodynamicMaxwellian

src/lbmpy/equilibrium/discrete_hydro_maxwellian.py

+import sympy as sp
+
+from pystencils.sympyextensions import simplify_by_equality
+
+from lbmpy.maxwellian_equilibrium import discrete_maxwellian_equilibrium
+
+from .generic_discrete_equilibrium import GenericDiscreteEquilibrium
+
+
+class DiscreteHydrodynamicMaxwellian(GenericDiscreteEquilibrium):
+    r"""
+        The textbook discretization of the Maxwellian equilibrium distribution of hydrodynamics.
+
+        This class represents the default discretization of the Maxwellian in velocity space,
+        computed from the distribution's expansion in Hermite polynomials (cf. :cite:`lbm_book`).
+        In :math:`d` dimensions, its populations :math:`f_i` on a given stencil 
+        :math:`(\mathbf{c}_i)_{i=0,\dots,q-1}` are given by
+
+        .. math::
+
+            f_i (\rho, \mathbf{u}) 
+                = w_i \rho \left( 
+                    1 + \frac{\mathbf{c}_i \cdot \mathbf{u}}{c_s^2}
+                    + \frac{(\mathbf{c}_i \cdot \mathbf{u})^2}{2 c_s^4}
+                    - \frac{\mathbf{u} \cdot \mathbf{u}}{2 c_s^2}
+                  \right).
+
+        Here :math:`w_i` denote the Hermite integration weights, also called lattice weights.
+        The incompressible variant of this distribution :cite:`HeIncompressible` can be written as
+
+        .. math::
+            f_i^{\mathrm{incomp}} (\rho, \mathbf{u})
+                = w_i \rho + w_i \rho_0 \left(
+                    \frac{\mathbf{c}_i \cdot \mathbf{u}}{c_s^2}
+                    + \frac{(\mathbf{c}_i \cdot \mathbf{u})^2}{2 c_s^4}
+                    - \frac{\mathbf{u} \cdot \mathbf{u}}{2 c_s^2}
+                  \right).
+
+        Again, for usage with zero-centered PDF storage, both distributions may be expressed in a delta-form
+        by subtracting their values at the background rest state at :math:`\rho = \rho_0`,
+        :math:`\mathbf{u} = \mathbf{0}`, which are exactly the lattice weights:
+
+        .. math::
+            \delta f_i &= f_i - w_i \\
+            \delta f_i^{\mathrm{incomp}} &= f_i^{\mathrm{incomp}} - w_i \\
+
+        Parameters:
+            stencil: Discrete velocity set for the discretization, see :class:`lbmpy.stencils.LBStencil`
+            compressible: If `False`, the incompressible equilibrium is created
+            deviation_only: If `True`, the delta-equilibrium is created
+            order: The discretization order in velocity to which computed statistical modes should be truncated
+            rho: Symbol or value for the density
+            delta_rho: Symbol or value for the density deviation
+            u: Sequence of symbols for the macroscopic velocity
+            c_s_sq: Symbol or value for the squared speed of sound
+    """
+
+    def __init__(self, stencil, compressible=True, deviation_only=False,
+                 order=2,
+                 rho=sp.Symbol("rho"),
+                 delta_rho=sp.Symbol("delta_rho"),
+                 u=sp.symbols("u_:3"),
+                 c_s_sq=sp.Symbol("c_s") ** 2):
+        dim = stencil.D
+
+        if order is None:
+            order = 4
+
+        self._order = order
+        self._compressible = compressible
+        self._deviation_only = deviation_only
+        self._rho = rho
+        self._rho_background = sp.Integer(1)
+        self._delta_rho = delta_rho
+        self._u = u[:dim]
+        self._c_s_sq = c_s_sq
+
+        pdfs = discrete_maxwellian_equilibrium(stencil, rho=rho, u=u,
+                                               order=order, c_s_sq=c_s_sq,
+                                               compressible=compressible)
+        if deviation_only:
+            shift = discrete_maxwellian_equilibrium(stencil, rho=self._rho_background, u=(0,) * dim,
+                                                    order=0, c_s_sq=c_s_sq, compressible=False)
+            pdfs = tuple(simplify_by_equality(f - s, rho, delta_rho, self._rho_background) for f, s in zip(pdfs, shift))
+
+        zeroth_order_moment = delta_rho if deviation_only else rho
+        super().__init__(stencil, pdfs, zeroth_order_moment, u, deviation_only)
+
+    @property
+    def order(self):
+        return self._order
+
+    @property
+    def deviation_only(self):
+        return self._deviation_only
+
+    @property
+    def compressible(self):
+        return self._compressible
+
+    @property
+    def density(self):
+        return self._rho
+
+    @property
+    def background_density(self):
+        return self._rho_background
+
+    @property
+    def density_deviation(self):
+        return self._delta_rho
+
+    @property
+    def velocity(self):
+        return self._u
+
+    @property
+    def background_distribution(self):
+        """Returns the discrete Maxwellian background distribution, which amounts exactly to the
+        lattice weights."""
+        return DiscreteHydrodynamicMaxwellian(self._stencil, compressible=True, deviation_only=False,
+                                              order=self._order, rho=self._rho_background,
+                                              delta_rho=0, u=(0,) * self.dim, c_s_sq=self._c_s_sq)
+
+    def central_moment(self, exponent_tuple_or_polynomial, velocity=None):
+        if velocity is None:
+            velocity = self._u
+
+        return super().central_moment(exponent_tuple_or_polynomial, velocity)
+
+    def central_moments(self, exponent_tuples_or_polynomials, velocity=None):
+        if velocity is None:
+            velocity = self._u
+
+        return super().central_moments(exponent_tuples_or_polynomials, velocity)
+
+    def cumulant(self, exponent_tuple_or_polynomial, rescale=True):
+        if not self._compressible or self._deviation_only:
+            raise Exception("Cumulants can only be computed for the compressible, "
+                            "non-deviation maxwellian equilibrium!")
+
+        return super().cumulant(exponent_tuple_or_polynomial, rescale=rescale)
+
+    #   ------------------ Utility ----------------------------------------------------------------
+
+    def __repr__(self):
+        return f"DiscreteHydrodynamicMaxwellian({self.stencil}, " \
+               f"compressible={self.compressible}, deviation_only:{self.deviation_only}" \
+               f"order={self.order})"
+
+    def _repr_html_(self):
+        def stylized_bool(b):
+            return "✓" if b else "✗"
+
+        html = f"""
+        <div style="max-height: 150pt; overflow-y: auto;">
+        <table style="border:none; width: 100%">
+            <tr>
+                <th colspan="3" style="text-align: left">
+                    Discrete Hydrodynamic Maxwellian Equilibrium
+                </th>
+                <td>Compressible: {stylized_bool(self._compressible)}</td>
+                <td>Deviation Only: {stylized_bool(self._deviation_only)}</td>
+                <td>Order: {"&#8734;" if self._order is None else self._order}</td>
+            </tr>
+        """
+
+        pdfs = self.discrete_populations
+        for f, eq in zip(sp.symbols(f"f_:{len(pdfs)}"), pdfs):
+            html += f'<tr><td colspan="6" style="text-align: left;"> ${f} = {sp.latex(eq)}$ </td></tr>'
+        
+        html += "</table></div>"
+
+        return html
+
+# end class DiscreteHydrodynamicMaxwellian

src/lbmpy/equilibrium/generic_discrete_equilibrium.py

+import sympy as sp
+
+from .abstract_equilibrium import AbstractEquilibrium
+
+from lbmpy.moments import discrete_moment, moment_matrix
+from lbmpy.cumulants import discrete_cumulant
+
+
+def discrete_equilibrium_from_matching_moments(stencil, moment_constraints,
+                                               zeroth_order_moment_symbol,
+                                               first_order_moment_symbols,
+                                               deviation_only=False):
+    assert len(moment_constraints) == stencil.Q
+    moments = tuple(moment_constraints.keys())
+    mm = moment_matrix(moments, stencil)
+    try:
+        pdfs = mm.inv() * sp.Matrix(list(moment_constraints.values()))
+        pdfs = pdfs.expand()
+        return GenericDiscreteEquilibrium(stencil, pdfs, zeroth_order_moment_symbol,
+                                          first_order_moment_symbols, deviation_only=deviation_only)
+    except sp.matrices.inverse.NonInvertibleMatrixError as e:
+        raise ValueError("Could not construct equilibrium from given moment constraints.") from e
+
+
+class GenericDiscreteEquilibrium(AbstractEquilibrium):
+    """
+        Class for encapsulating arbitrary discrete equilibria, given by their equilibrium populations.
+
+        This class takes both a stencil and a sequence of populations modelling a discrete distribution function
+        and provides basic functionality for computing and caching that distribution's statistical modes.
+
+        Parameters:
+            stencil: Discrete velocity set, see :class:`lbmpy.stencils.LBStencil`.
+            equilibrium_pdfs: List of q populations, describing the particle distribution on the discrete velocity
+                              set given by the stencil.
+            zeroth_order_moment_symbol: Symbol corresponding to the distribution's zeroth-order moment, the area under
+                                        it's curve (see :attr:`zeroth_order_moment_symbol`).
+            first_order_moment_symbols: Sequence of symbols corresponding to the distribution's first-order moment, the
+                                        vector of its mean values (see :attr:`first_order_moment_symbols`).
+            deviation_only: Set to `True` if the given populations model only the deviation from a rest state, to be
+                            used in junction with the zero-centered storage format.
+    """
+
+    def __init__(self, stencil, equilibrium_pdfs,
+                 zeroth_order_moment_symbol,
+                 first_order_moment_symbols,
+                 deviation_only=False):
+        super().__init__(dim=stencil.D)
+
+        if len(equilibrium_pdfs) != stencil.Q:
+            raise ValueError(f"Wrong number of PDFs."
+                             f"On the {stencil} stencil, exactly {stencil.Q} populations must be passed!")
+
+        self._stencil = stencil
+        self._pdfs = tuple(equilibrium_pdfs)
+        self._zeroth_order_moment_symbol = zeroth_order_moment_symbol
+        self._first_order_moment_symbols = first_order_moment_symbols
+        self._deviation_only = deviation_only
+
+    @property
+    def stencil(self):
+        return self._stencil
+
+    @property
+    def deviation_only(self):
+        return self._deviation_only
+
+    @property
+    def continuous_equation(self):
+        """Always returns `None`."""
+        return None
+
+    @property
+    def discrete_populations(self):
+        return self._pdfs
+
+    @property
+    def background_distribution(self):
+        """Always returns `None`. To specify a background distribution, override this class."""
+        return None
+
+    @property
+    def zeroth_order_moment_symbol(self):
+        return self._zeroth_order_moment_symbol
+
+    @property
+    def first_order_moment_symbols(self):
+        return self._first_order_moment_symbols
+
+    #   Moment Computation
+
+    def _monomial_raw_moment(self, exponents):
+        return discrete_moment(self._pdfs, exponents, self._stencil)
+
+    def _monomial_central_moment(self, exponents, frame_of_reference):
+        return discrete_moment(self._pdfs, exponents, self._stencil, shift_velocity=frame_of_reference)
+
+    def _monomial_cumulant(self, exponents, rescale):
+        value = discrete_cumulant(self._pdfs, exponents, self._stencil)
+        if rescale:
+            value = self.zeroth_order_moment_symbol * value
+        return value

lbmpy/fieldaccess.py → src/lbmpy/fieldaccess.py

@@ -7,10 +7,15 @@ from pystencils import Field
 from pystencils.astnodes import LoopOverCoordinate
 from pystencils.stencil import inverse_direction
 
+from lbmpy.enums import Stencil
+from lbmpy.stencils import LBStencil
+
 __all__ = ['PdfFieldAccessor', 'CollideOnlyInplaceAccessor', 'StreamPullTwoFieldsAccessor',
            'AAEvenTimeStepAccessor', 'AAOddTimeStepAccessor',
            'PeriodicTwoFieldsAccessor', 'StreamPushTwoFieldsAccessor',
            'EsoTwistEvenTimeStepAccessor', 'EsoTwistOddTimeStepAccessor',
+           'EsoPullEvenTimeStepAccessor', 'EsoPullOddTimeStepAccessor',
+           'EsoPushEvenTimeStepAccessor', 'EsoPushOddTimeStepAccessor',
            'visualize_pdf_field_accessor', 'visualize_field_mapping']
 
 
@@ -51,11 +56,11 @@ class CollideOnlyInplaceAccessor(PdfFieldAccessor):
 
     @staticmethod
     def read(field, stencil):
-        return [field(i) for i in range(len(stencil))]
+        return [field(i) for i in range(stencil.Q)]
 
     @staticmethod
     def write(field, stencil):
-        return [field(i) for i in range(len(stencil))]
+        return [field(i) for i in range(stencil.Q)]
 
 
 class StreamPullTwoFieldsAccessor(PdfFieldAccessor):
@@ -67,7 +72,7 @@ class StreamPullTwoFieldsAccessor(PdfFieldAccessor):
 
     @staticmethod
     def write(field, stencil):
-        return [field(i) for i in range(len(stencil))]
+        return [field(i) for i in range(stencil.Q)]
 
 
 class StreamPushTwoFieldsAccessor(PdfFieldAccessor):
@@ -75,7 +80,7 @@ class StreamPushTwoFieldsAccessor(PdfFieldAccessor):
 
     @staticmethod
     def read(field, stencil):
-        return [field(i) for i in range(len(stencil))]
+        return [field(i) for i in range(stencil.Q)]
 
     @staticmethod
     def write(field, stencil):
@@ -125,7 +130,7 @@ class PeriodicTwoFieldsAccessor(PdfFieldAccessor):
 
     @staticmethod
     def write(field, stencil):
-        return [field(i) for i in range(len(stencil))]
+        return [field(i) for i in range(stencil.Q)]
 
 
 class AAEvenTimeStepAccessor(PdfFieldAccessor):
@@ -133,11 +138,11 @@ class AAEvenTimeStepAccessor(PdfFieldAccessor):
 
     @staticmethod
     def read(field, stencil):
-        return [field(i) for i in range(len(stencil))]
+        return [field(i) for i in range(stencil.Q)]
 
     @staticmethod
     def write(field, stencil):
-        return [field(stencil.index(inverse_direction(d))) for d in stencil]
+        return [field(stencil.inverse_index(d)) for d in stencil]
 
 
 class AAOddTimeStepAccessor(PdfFieldAccessor):
@@ -145,67 +150,179 @@ class AAOddTimeStepAccessor(PdfFieldAccessor):
 
     @staticmethod
     def read(field, stencil):
-        res = []
-        for i, d in enumerate(stencil):
-            inv_dir = inverse_direction(d)
-            field_access = field[inv_dir](stencil.index(inv_dir))
-            res.append(field_access)
-        return res
+        return [field[inverse_direction(d)](stencil.inverse_index(d)) for i, d in enumerate(stencil)]
 
     @staticmethod
     def write(field, stencil):
         return [field[d](i) for i, d in enumerate(stencil)]
 
 
+class EsoTwistEvenTimeStepAccessor(PdfFieldAccessor):
+    is_inplace = True
+
+    @staticmethod
+    def read(field, stencil):
+        return [field[tuple(max(-e, 0) for e in d)](i) for i, d in enumerate(stencil)]
+
+    @staticmethod
+    def write(field, stencil):
+        return [field[tuple(max(e, 0) for e in d)](stencil.inverse_index(d)) for d in stencil]
+
+
 class EsoTwistOddTimeStepAccessor(PdfFieldAccessor):
     is_inplace = True
 
     @staticmethod
     def read(field, stencil):
-        result = []
-        for direction in stencil:
-            inv_dir = inverse_direction(direction)
-            spatial_offset = tuple(max(e, 0) for e in inv_dir)
-            result.append(field[spatial_offset](stencil.index(inv_dir)))
+        return [field[tuple(max(e, 0) for e in inverse_direction(d))](stencil.inverse_index(d)) for d in stencil]
+
+    @staticmethod
+    def write(field, stencil):
+        return [field[tuple(max(e, 0) for e in d)](i) for i, d in enumerate(stencil)]
+
+
+class EsoPullEvenTimeStepAccessor(PdfFieldAccessor):
+    is_inplace = True
+
+    @staticmethod
+    def read(field, stencil):
+        lehmann_stencil = _get_lehmann_stencil(stencil)
+        center_cell = tuple([0] * stencil.D)
+        result = [field.center]
+        for i, d in enumerate(stencil):
+            if i == 0:
+                continue
+            if lehmann_stencil.index(d) % 2 == 0:
+                result.append(field[inverse_direction(d)](i))
+            else:
+                result.append(field[center_cell](i))
         return result
 
     @staticmethod
     def write(field, stencil):
-        result = []
-        for i, direction in enumerate(stencil):
-            spatial_offset = tuple(max(e, 0) for e in direction)
-            result.append(field[spatial_offset](i))
+        lehmann_stencil = _get_lehmann_stencil(stencil)
+        center_cell = tuple([0] * stencil.D)
+        result = [field.center]
+        for i, d in enumerate(stencil):
+            if i == 0:
+                continue
+            if lehmann_stencil.index(d) % 2 == 0:
+                result.append(field[center_cell](stencil.inverse_index(d)))
+            else:
+                result.append(field[d](stencil.inverse_index(d)))
+
         return result
 
 
-class EsoTwistEvenTimeStepAccessor(PdfFieldAccessor):
+class EsoPullOddTimeStepAccessor(PdfFieldAccessor):
     is_inplace = True
 
     @staticmethod
     def read(field, stencil):
-        result = []
-        for i, direction in enumerate(stencil):
-            spatial_offset = tuple(max(-e, 0) for e in direction)
-            result.append(field[spatial_offset](i))
+        lehmann_stencil = _get_lehmann_stencil(stencil)
+        center_cell = tuple([0] * stencil.D)
+        result = [field.center]
+        for i, d in enumerate(stencil):
+            if i == 0:
+                continue
+            if lehmann_stencil.index(d) % 2 == 0:
+                result.append(field[inverse_direction(d)](stencil.inverse_index(d)))
+            else:
+                result.append(field[center_cell](stencil.inverse_index(d)))
         return result
 
     @staticmethod
     def write(field, stencil):
-        result = []
-        for direction in stencil:
-            inv_dir = inverse_direction(direction)
-            spatial_offset = tuple(max(e, 0) for e in direction)
-            result.append(field[spatial_offset](stencil.index(inv_dir)))
+        lehmann_stencil = _get_lehmann_stencil(stencil)
+        center_cell = tuple([0] * stencil.D)
+        result = [field.center]
+        for i, d in enumerate(stencil):
+            if i == 0:
+                continue
+            if lehmann_stencil.index(d) % 2 == 0:
+                result.append(field[center_cell](i))
+            else:
+                result.append(field[d](i))
+
+        return result
+
+
+class EsoPushEvenTimeStepAccessor(PdfFieldAccessor):
+    is_inplace = True
+
+    @staticmethod
+    def read(field, stencil):
+        lehmann_stencil = _get_lehmann_stencil(stencil)
+        center_cell = tuple([0] * stencil.D)
+        result = [field.center]
+        for i, d in enumerate(stencil):
+            if i == 0:
+                continue
+            if lehmann_stencil.index(d) % 2 == 0:
+                result.append(field[center_cell](stencil.inverse_index(d)))
+            else:
+                result.append(field[inverse_direction(d)](stencil.inverse_index(d)))
+
+        return result
+
+    @staticmethod
+    def write(field, stencil):
+        lehmann_stencil = _get_lehmann_stencil(stencil)
+        center_cell = tuple([0] * stencil.D)
+        result = [field.center]
+        for i, d in enumerate(stencil):
+            if i == 0:
+                continue
+            if lehmann_stencil.index(d) % 2 == 0:
+                result.append(field[d](i))
+            else:
+                result.append(field[center_cell](i))
+
+        return result
+
+
+class EsoPushOddTimeStepAccessor(PdfFieldAccessor):
+    is_inplace = True
+
+    @staticmethod
+    def read(field, stencil):
+        lehmann_stencil = _get_lehmann_stencil(stencil)
+        center_cell = tuple([0] * stencil.D)
+        result = [field.center]
+        for i, d in enumerate(stencil):
+            inv_dir = inverse_direction(d)
+            if i == 0:
+                continue
+            if lehmann_stencil.index(d) % 2 == 0:
+                result.append(field[center_cell](i))
+            else:
+                result.append(field[inv_dir](i))
+
+        return result
+
+    @staticmethod
+    def write(field, stencil):
+        lehmann_stencil = _get_lehmann_stencil(stencil)
+        center_cell = tuple([0] * stencil.D)
+        result = [field.center]
+        for i, d in enumerate(stencil):
+            if i == 0:
+                continue
+            if lehmann_stencil.index(d) % 2 == 0:
+                result.append(field[d](stencil.inverse_index(d)))
+            else:
+                result.append(field[center_cell](stencil.inverse_index(d)))
+
         return result
 
 
 # -------------------------------------------- Visualization -----------------------------------------------------------
 
 
-def visualize_field_mapping(axes, stencil, field_mapping, color='b'):
+def visualize_field_mapping(axes, stencil, field_mapping, inverted=False, color='b'):
     from lbmpy.plot import LbGrid
     grid = LbGrid(3, 3)
-    grid.fill_with_default_arrows()
+    grid.fill_with_default_arrows(inverted=inverted)
     for field_access, direction in zip(field_mapping, stencil):
         field_position = stencil[field_access.index[0]]
         neighbor = field_access.offsets
@@ -214,14 +331,18 @@ def visualize_field_mapping(axes, stencil, field_mapping, color='b'):
     grid.draw(axes)
 
 
-def visualize_pdf_field_accessor(pdf_field_accessor, figure=None):
-    from lbmpy.stencils import get_stencil
+def visualize_pdf_field_accessor(pdf_field_accessor, title=True, read_plot_params=None, write_plot_params=None,
+                                 figure=None):
 
+    if write_plot_params is None:
+        write_plot_params = {}
+    if read_plot_params is None:
+        read_plot_params = {}
     if figure is None:
         import matplotlib.pyplot as plt
         figure = plt.gcf()
 
-    stencil = get_stencil('D2Q9')
+    stencil = LBStencil(Stencil.D2Q9)
 
     figure.patch.set_facecolor('white')
 
@@ -232,8 +353,36 @@ def visualize_pdf_field_accessor(pdf_field_accessor, figure=None):
     ax_left = figure.add_subplot(1, 2, 1)
     ax_right = figure.add_subplot(1, 2, 2)
 
-    visualize_field_mapping(ax_left, stencil, pre_collision_accesses, color='k')
-    visualize_field_mapping(ax_right, stencil, post_collision_accesses, color='r')
+    if 'color' not in read_plot_params:
+        read_plot_params['color'] = 'k'
+    if 'color' not in write_plot_params:
+        write_plot_params['color'] = 'r'
+
+    visualize_field_mapping(ax_left, stencil, pre_collision_accesses, **read_plot_params)
+    visualize_field_mapping(ax_right, stencil, post_collision_accesses, **write_plot_params)
+
+    if title:
+        ax_left.set_title("Read")
+        ax_right.set_title("Write")
+
+# -------------------------------------------- Helpers -----------------------------------------------------------
 
-    ax_left.set_title("Read")
-    ax_right.set_title("Write")
+
+def _get_lehmann_stencil(stencil):
+    """
+    EsoPull and EsoPush streaming is only simple to implement with a specific stencil ordering, that comes from
+    "High Performance Free Surface LBM on GPUs" by moritz lehmann
+
+    Args:
+        stencil: lattice Boltzmann stencil
+    """
+    if stencil.Q == 9:
+        return LBStencil(Stencil.D2Q9, ordering="lehmann")
+    elif stencil.Q == 15:
+        return LBStencil(Stencil.D3Q15, ordering="lehmann")
+    elif stencil.Q == 19:
+        return LBStencil(Stencil.D3Q19, ordering="lehmann")
+    elif stencil.Q == 27:
+        return LBStencil(Stencil.D3Q27, ordering="lehmann")
+    else:
+        ValueError("EsoPull or EsoPush is only available for D2Q9, D3Q15, D3Q19 and D3Q27 stencil")

src/lbmpy/flow_statistics.py

+import sympy as sp
+import pystencils as ps
+
+from pystencils.field import Field
+
+
+def welford_assignments(field, mean_field, sum_of_squares_field=None, sum_of_cubes_field=None):
+    r"""
+    Creates the assignments for the kernel creation of Welford's algorithm
+    (https://en.wikipedia.org/wiki/Algorithms_for_calculating_variance#Welford's_online_algorithm).
+    This algorithm is an online algorithm for the mean and variance calculation of sample data, here implemented for
+    the execution on scalar or vector fields, e.g., velocity.
+    The calculation of the variance / the third-order central moments is optional and only performed if a field for
+    the sum of squares / sum of cubes is given.
+
+    The mean value is directly updated in the mean vector field.
+    The variance / covariance must be retrieved in a post-processing step. Let :math `M_{2,n}` denote the value of the
+    sum of squares after the first :math `n` samples. According to Welford the biased sample variance
+    :math `\sigma_n^2` and the unbiased sample variance :math `s_n^2` are given by
+
+    .. math ::
+        \sigma_n^2 = \frac{M_{2,n}}{n}
+        s_n^2 = \frac{M_{2,n}}{n-1},
+
+    respectively.
+
+    Liekwise, to the 3rd-order central moment(s) of the (vector) field, the sum of cubes field must be divided
+    by :math `n` in a post-processing step.
+    """
+
+    if isinstance(field, Field):
+        dim = field.values_per_cell()
+        welford_field = field.center
+    elif isinstance(field, Field.Access):
+        dim = field.field.values_per_cell()
+        welford_field = field
+    else:
+        raise ValueError("Vector field has to be a pystencils Field or Field.Access")
+
+    if isinstance(mean_field, Field):
+        welford_mean_field = mean_field.center
+    elif isinstance(mean_field, Field.Access):
+        welford_mean_field = mean_field
+    else:
+        raise ValueError("Mean vector field has to be a pystencils Field or Field.Access")
+
+    if sum_of_squares_field is None:
+        # sum of products will not be calculated, i.e., the covariance matrix is not retrievable
+        welford_sum_of_squares_field = None
+    else:
+        if isinstance(sum_of_squares_field, Field):
+            welford_sum_of_squares_field = sum_of_squares_field.center
+        elif isinstance(sum_of_squares_field, Field.Access):
+            welford_sum_of_squares_field = sum_of_squares_field
+        else:
+            raise ValueError("Sum of squares field has to be a pystencils Field or Field.Access")
+
+        assert welford_sum_of_squares_field.field.values_per_cell() == dim ** 2, \
+            (f"Sum of squares field does not have the right layout. "
+             f"Index dimension: {welford_sum_of_squares_field.field.values_per_cell()}, expected: {dim ** 2}")
+
+    if sum_of_cubes_field is None:
+        # sum of cubes will not be calculated, i.e., the 3rd-order central moments are not retrievable
+        welford_sum_of_cubes_field = None
+    else:
+        if isinstance(sum_of_cubes_field, Field):
+            welford_sum_of_cubes_field = sum_of_cubes_field.center
+        elif isinstance(sum_of_cubes_field, Field.Access):
+            welford_sum_of_cubes_field = sum_of_cubes_field
+        else:
+            raise ValueError("Sum of cubes field has to be a pystencils Field or Field.Access")
+
+        assert welford_sum_of_cubes_field.field.values_per_cell() == dim ** 3, \
+            (f"Sum of cubes field does not have the right layout. "
+             f"Index dimension: {welford_sum_of_cubes_field.field.values_per_cell()}, expected: {dim ** 3}")
+
+    # for the calculation of the thrid-order moments, the variance must also be calculated
+    if welford_sum_of_cubes_field is not None:
+        assert welford_sum_of_squares_field is not None
+
+    # actual assignments
+
+    counter = sp.Symbol('counter')
+    delta = sp.symbols(f"delta_:{dim}")
+
+    main_assignments = list()
+
+    for i in range(dim):
+        main_assignments.append(ps.Assignment(delta[i], welford_field.at_index(i) - welford_mean_field.at_index(i)))
+        main_assignments.append(ps.Assignment(welford_mean_field.at_index(i),
+                                              welford_mean_field.at_index(i) + delta[i] / counter))
+
+    if sum_of_squares_field is not None:
+        delta2 = sp.symbols(f"delta2_:{dim}")
+        for i in range(dim):
+            main_assignments.append(
+                ps.Assignment(delta2[i], welford_field.at_index(i) - welford_mean_field.at_index(i)))
+        for i in range(dim):
+            for j in range(dim):
+                idx = i * dim + j
+                main_assignments.append(ps.Assignment(
+                    welford_sum_of_squares_field.at_index(idx),
+                    welford_sum_of_squares_field.at_index(idx) + delta[i] * delta2[j]))
+
+        if sum_of_cubes_field is not None:
+            for i in range(dim):
+                for j in range(dim):
+                    for k in range(dim):
+                        idx = (i * dim + j) * dim + k
+                        main_assignments.append(ps.Assignment(
+                            welford_sum_of_cubes_field.at_index(idx),
+                            welford_sum_of_cubes_field.at_index(idx)
+                            - delta[k] / counter * welford_sum_of_squares_field(i * dim + j)
+                            - delta[j] / counter * welford_sum_of_squares_field(k * dim + i)
+                            - delta[i] / counter * welford_sum_of_squares_field(j * dim + k)
+                            + delta2[i] * (2 * delta[j] - delta2[j]) * delta[k]
+                        ))
+
+    return main_assignments

lbmpy/fluctuatinglb.py → src/lbmpy/fluctuatinglb.py

@@ -6,7 +6,8 @@ correction term is added to the collision rule
 import numpy as np
 import sympy as sp
 
-from lbmpy.moments import MOMENT_SYMBOLS
+from lbmpy.moments import MOMENT_SYMBOLS, is_shear_moment, get_order
+from lbmpy.equilibrium import ContinuousHydrodynamicMaxwellian
 from pystencils import Assignment, TypedSymbol
 from pystencils.rng import PhiloxFourFloats, random_symbol
 from pystencils.simp.assignment_collection import SymbolGen
@@ -18,14 +19,17 @@ def add_fluctuations_to_collision_rule(collision_rule, temperature=None, amplitu
     """"""
     if not (temperature and not amplitudes) or (temperature and amplitudes):
         raise ValueError("Fluctuating LBM: Pass either 'temperature' or 'amplitudes'.")
-
+    
     method = collision_rule.method
     if not amplitudes:
         amplitudes = fluctuation_amplitude_from_temperature(method, temperature, c_s_sq)
     if block_offsets == 'walberla':
         block_offsets = tuple(TypedSymbol("block_offset_{}".format(i), np.uint32) for i in range(3))
 
-    rng_symbol_gen = random_symbol(collision_rule.subexpressions, seed,
+    if not method.is_weighted_orthogonal:
+        raise ValueError("Fluctuations can only be added to weighted-orthogonal methods")
+
+    rng_symbol_gen = random_symbol(collision_rule.subexpressions, seed=seed,
                                    rng_node=rng_node, dim=method.dim, offsets=block_offsets)
     correction = fluctuation_correction(method, rng_symbol_gen, amplitudes)
 
@@ -38,9 +42,7 @@ def fluctuation_amplitude_from_temperature(method, temperature, c_s_sq=sp.Symbol
     """Produces amplitude equations according to (2.60) and (3.54) in Schiller08"""
     normalization_factors = sp.matrix_multiply_elementwise(method.moment_matrix, method.moment_matrix) * \
         sp.Matrix(method.weights)
-    density = method.zeroth_order_equilibrium_moment_symbol
-    if method.conserved_quantity_computation.zero_centered_pdfs:
-        density += 1
+    density = method._cqc.density_symbol
     mu = temperature * density / c_s_sq
     return [sp.sqrt(mu * norm * (1 - (1 - rr) ** 2))
             for norm, rr in zip(normalization_factors, method.relaxation_rates)]
@@ -57,3 +59,22 @@ def fluctuation_correction(method, rng_generator, amplitudes=SymbolGen("phi")):
 
     # corrections are applied in real space hence we need to convert to real space here
     return method.moment_matrix.inv() * amplitude_matrix * random_matrix
+
+
+class ThermalizedEquilibrium(ContinuousHydrodynamicMaxwellian):
+    """TODO: Remove Again! 
+        
+    This class is currently only used in the tutorial notebook `demo_thermalized_lbm.ipynb`
+    and has been added only temporarily, until the thermalized LBM is updated to our new
+    equilibrium framework."""
+    def __init__(self, random_number_symbols, **kwargs):
+        super().__init__(**kwargs)
+        self.random_number_symbols = random_number_symbols
+
+    def moment(self, exponent_tuple_or_polynomial):
+        value = super().moment(exponent_tuple_or_polynomial)
+        if is_shear_moment(exponent_tuple_or_polynomial, dim=self.dim):
+            value += self.random_number_symbols[0] * 0.001
+        elif get_order(exponent_tuple_or_polynomial) > 2:
+            value += self.random_number_symbols[1] * 0.001
+        return value

src/lbmpy/forcemodels.py

+r"""
+.. module:: forcemodels
+    :synopsis: Collection of forcing terms for hydrodynamic LBM simulations
+
+
+Get started:
+------------
+
+This module offers different models to introduce a body force in the lattice Boltzmann scheme.
+If you don't know which model to choose, use :class:`lbmpy.forcemodels.Guo`.
+
+
+Detailed information:
+---------------------
+
+Force models add a term :math:`C_F` to the collision equation:
+
+.. math ::
+
+    f(\mathbf{x} + c_q \Delta t, t + \Delta t) - f(\mathbf{x},t) = \Omega(f, f^{(\mathrm{eq})})
+                                                            + \underbrace{F_q}_{\mbox{forcing term}}
+
+The form of this term depends on the concrete force model: the first moment of this forcing term is equal
+to the acceleration :math:`\mathbf{a}` for all force models. Here :math:`\mathbf{F}` is the D dimensional force vector,
+which defines the force for each spatial dircetion.
+
+.. math ::
+
+    \sum_q \mathbf{c}_q \mathbf{F} = \mathbf{a}
+
+
+The second order moment is different for the forcing models - if it is zero the model is suited for
+incompressible flows. For weakly compressible collision operators a force model with a corrected second order moment
+should be chosen.
+
+.. math ::
+
+    \sum_q c_{qi} c_{qj} f_q &= F_i u_j + F_j u_i  &\qquad \mbox{for Guo, Luo models}
+    
+    \sum_q c_{qi} c_{qj} f_q &= 0  &\qquad \mbox{for Simple, Buick}
+    
+Models with zero second order moment have:
+
+.. math ::
+    
+    F_q = \frac{w_q}{c_s^2} c_{qi} \; a_i
+
+Models with nonzero second moment have:
+
+.. math ::
+    
+    F_q = \frac{w_q}{c_s^2} c_{qi} \; a_i + \frac{w_q}{c_s^4} (c_{qi} c_{qj} - c_s^2 \delta_{ij} ) u_j \, a_i
+
+
+For all force models the computation of the macroscopic velocity has to be adapted (shifted) by adding a term
+:math:`S_{macro}` that we call "macroscopic velocity shift"
+    
+    .. math ::
+    
+        \mathbf{u} &= \sum_q \mathbf{c}_q f_q + S_{\mathrm{macro}}
+        
+        S_{\mathrm{macro}} &= \frac{\Delta t}{2 \cdot \rho} \sum_q F_q
+
+
+Some models also shift the velocity entering the equilibrium distribution.
+
+Comparison
+----------
+ 
+Force models can be distinguished by 2 options:
+
+**Option 1**:
+    :math:`C_F = 1` and equilibrium velocity is not shifted, or :math:`C_F=(1 - \frac{\omega}{2})`
+    and equilibrum is shifted.
+ 
+**Option 2**: 
+    second velocity moment is zero or :math:`F_i u_j + F_j u_i`
+
+
+=====================  ====================  =================
+ Option2 \\ Option1    no equilibrium shift  equilibrium shift
+=====================  ====================  =================
+second moment zero      :class:`Simple`      :class:`Buick`
+second moment nonzero   :class:`Luo`         :class:`Guo`         
+=====================  ====================  =================
+  
+"""
+from warnings import warn
+import abc
+import sympy as sp
+
+from pystencils.sympyextensions import scalar_product
+
+from lbmpy.maxwellian_equilibrium import (
+    discrete_maxwellian_equilibrium, get_equilibrium_values_of_maxwell_boltzmann_function
+)
+from lbmpy.moments import (
+    MOMENT_SYMBOLS, exponent_tuple_sort_key, exponents_to_polynomial_representations,
+    extract_monomials, moment_sort_key, moment_matrix,
+    monomial_to_polynomial_transformation_matrix, set_up_shift_matrix)
+
+FORCE_SYMBOLS = sp.symbols("F_x, F_y, F_z")
+
+
+class AbstractForceModel(abc.ABC):
+    r"""
+    Abstract base class for all force models. All force models have to implement the __call__, which should return a
+    q dimensional vector added to the PDFs in the population space. If an MRT method is used, it is also possible
+    to apply the force directly in the moment space. This is done by additionally providing the function
+    moment_space_forcing. The MRT method will check if it is available and apply the force directly in the moment
+    space. For MRT methods in the central moment space the central_moment_space_forcing function can be provided,
+    which shifts the force vector to the central moment space. Applying the force in the collision space has the
+    advantage of saving FLOPs. Furthermore, it is sometimes easier to apply the correct force vector in the
+    collision space because often, the relaxation matrix has to be taken into account.
+
+    Args:
+        force: force vector of size dim which contains the force for each spatial dimension.
+    """
+
+    def __init__(self, force):
+        self._force = force
+
+        # All force models work internaly with a pure symbolic list of the forcing vector.
+        # Each entry of the original force vector which is not a symbol is mapped to a symbol and a subs dict is
+        # created. The subs dict should be used inside the LB method for the creation of the collision rule.
+        self._symbolic_force = [x if isinstance(x, sp.Symbol) else y for x, y in zip(force, FORCE_SYMBOLS)]
+        self._subs_dict_force = {x: y for (x, y) in zip(self._symbolic_force, force) if x != y}
+
+        # The following booleans should indicate if a force model is has the function moment_space_forcing which
+        # transfers the forcing terms to the moment space or central_moment_space_forcing which transfers them to the
+        # central moment space
+        self.has_moment_space_forcing = hasattr(self, "moment_space_forcing")
+        self.has_central_moment_space_forcing = hasattr(self, "central_moment_space_forcing")
+        self.has_symmetric_central_moment_forcing = hasattr(self, "symmetric_central_moment_forcing")
+
+    def __call__(self, lb_method):
+        r"""
+        This function returns a vector of size q which is added to the PDFs in the PDF space. It has to be implemented
+        by all forcing models and returns a sympy Matrix containing the q dimensional force vector.
+
+        Args:
+            lb_method: LB method, see lbmpy.creationfunctions.create_lb_method
+        """
+        raise NotImplementedError("Force model class has to overwrite __call__")
+
+    def macroscopic_velocity_shift(self, density):
+        r"""
+        macroscopic velocity shift by :math:`\frac{\Delta t}{2 \cdot \rho}`
+
+        Args:
+            density: Density symbol which is needed for the shift
+        """
+        return default_velocity_shift(density, self.symbolic_force_vector)
+
+    def macroscopic_momentum_density_shift(self, *_):
+        r"""
+        macroscopic momentum density shift by :math:`\frac{\Delta t}{2}`
+        """
+        return default_momentum_density_shift(self.symbolic_force_vector)
+
+    def equilibrium_velocity_shift(self, density):
+        r"""
+        Some models also shift the velocity entering the equilibrium distribution. By default the shift is zero
+
+        Args:
+            density: Density symbol which is needed for the shift
+        """
+        return [0] * len(self.symbolic_force_vector)
+
+    @property
+    def symbolic_force_vector(self):
+        return self._symbolic_force
+
+    @property
+    def subs_dict_force(self):
+        return self._subs_dict_force
+
+
+class Simple(AbstractForceModel):
+    r"""
+    A simple force model which introduces the following additional force term in the
+    collision process :math:`\frac{w_q}{c_s^2} \mathbf{c_{q}} \cdot \mathbf{F}` (often: force = rho * acceleration
+    where rho is the zeroth moment to be consistent with the above definition)
+    Should only be used with constant forces!
+    Shifts the macroscopic velocity by :math:`\frac{\mathbf{F}}{2}`, but does not change the equilibrium velocity.
+    """
+
+    def __call__(self, lb_method):
+        force = self.symbolic_force_vector
+        assert len(force) == lb_method.dim, "Force vectore must match with the dimensions of the lb method"
+        cs_sq = sp.Rational(1, 3)  # squared speed of sound
+
+        result = [(w_i / cs_sq) * scalar_product(force, direction)
+                  for direction, w_i in zip(lb_method.stencil, lb_method.weights)]
+
+        return sp.Matrix(result)
+
+    def moment_space_forcing(self, lb_method):
+        return (lb_method.moment_matrix * self(lb_method)).expand()
+
+    def central_moment_space_forcing(self, lb_method):
+        moments = (lb_method.moment_matrix * sp.Matrix(self(lb_method))).expand()
+        return lb_method.shift_matrix * moments
+
+    def symmetric_central_moment_forcing(self, lb_method, central_moments):
+        u = lb_method.first_order_equilibrium_moment_symbols
+        cm_matrix = moment_matrix(central_moments, lb_method.stencil, shift_velocity=u)
+        before = sp.Matrix([0] * lb_method.stencil.Q)
+        after = cm_matrix @ sp.Matrix(self(lb_method))
+        return before, after
+
+
+class CentralMoment(AbstractForceModel):
+    r"""
+    A force model that only shifts the macroscopic and equilibrium velocities and does not introduce a forcing term to
+    the collision process. Forcing is then applied through relaxation of the first central moments in the shifted
+    frame of reference (cf. https://doi.org/10.1016/j.camwa.2015.05.001).
+    """
+    def __call__(self, lb_method):
+        raise ValueError("This force model can only be combined with the Cumulant collision model")
+
+    def symmetric_central_moment_forcing(self, lb_method, *_):
+        before = sp.Matrix([0, ] * lb_method.stencil.Q)
+        after = sp.Matrix([0, ] * lb_method.stencil.Q)
+        for i, sf in enumerate(self.symbolic_force_vector):
+            before[i + 1] = sp.Rational(1, 2) * sf
+            after[i + 1] = sp.Rational(1, 2) * sf
+
+        return before, after
+
+    def equilibrium_velocity_shift(self, density):
+        return default_velocity_shift(density, self.symbolic_force_vector)
+
+
+class Luo(AbstractForceModel):
+    r"""Force model by Luo :cite:`luo1993lattice`.
+
+    Shifts the macroscopic velocity by :math:`\frac{\mathbf{F}}{2}`, but does not change the equilibrium velocity.
+    """
+
+    def __call__(self, lb_method):
+        u = sp.Matrix(lb_method.first_order_equilibrium_moment_symbols)
+        force = sp.Matrix(self.symbolic_force_vector)
+
+        cs_sq = sp.Rational(1, 3)  # squared speed of sound
+
+        result = []
+        for direction, w_i in zip(lb_method.stencil, lb_method.weights):
+            direction = sp.Matrix(direction)
+            first_summand = (direction - u) / cs_sq
+            second_summand = (direction * direction.dot(u)) / cs_sq ** 2
+
+            fq = w_i * force.dot(first_summand + second_summand)
+
+            result.append(fq.simplify())
+        return sp.Matrix(result)
+
+    def moment_space_forcing(self, lb_method):
+        return (lb_method.moment_matrix * self(lb_method)).expand()
+
+    def central_moment_space_forcing(self, lb_method):
+        moments = lb_method.moment_matrix * self(lb_method)
+        return (lb_method.shift_matrix * moments).expand()
+
+    def symmetric_central_moment_forcing(self, lb_method, central_moments):
+        u = lb_method.first_order_equilibrium_moment_symbols
+        cm_matrix = moment_matrix(central_moments, lb_method.stencil, shift_velocity=u)
+        before = sp.Matrix([0] * lb_method.stencil.Q)
+        after = (cm_matrix @ sp.Matrix(self(lb_method))).expand()
+        return before, after
+
+
+class Guo(AbstractForceModel):
+    r"""
+    Force model by Guo  :cite:`guo2002discrete`, generalized to MRT,
+    which makes it equivalent to :cite:`schiller2008thermal`, equation 4.67
+    Adapts the calculation of the macroscopic velocity as well as the equilibrium velocity
+    (both shifted by :math:`\frac{\mathbf{F}}{2}`)!
+    """
+
+    def __call__(self, lb_method):
+        if len(set(lb_method.relaxation_rates)) == 1:
+            #   It's an SRT method!
+            rr = lb_method.symbolic_relaxation_matrix[0]
+            force_terms = Luo(self.symbolic_force_vector)(lb_method)
+            correction_factor = (1 - rr / 2)
+            result = correction_factor * force_terms
+        else:
+            force_terms = self.moment_space_forcing(lb_method)
+            result = (lb_method.moment_matrix.inv() * force_terms).expand()
+        return result
+
+    def moment_space_forcing(self, lb_method):
+        luo = Luo(self.symbolic_force_vector)
+        q = len(lb_method.stencil)
+        correction_factor = sp.eye(q) - sp.Rational(1, 2) * lb_method.symbolic_relaxation_matrix
+        moments = correction_factor * (lb_method.moment_matrix * sp.Matrix(luo(lb_method))).expand()
+        return moments
+
+    def central_moment_space_forcing(self, lb_method):
+        luo = Luo(self.symbolic_force_vector)
+        q = len(lb_method.stencil)
+        correction_factor = sp.eye(q) - sp.Rational(1, 2) * lb_method.symbolic_relaxation_matrix
+        moments = (lb_method.moment_matrix * sp.Matrix(luo(lb_method)))
+        central_moments = correction_factor * (lb_method.shift_matrix * moments).expand()
+
+        return central_moments
+
+    def symmetric_central_moment_forcing(self, lb_method, central_moments):
+        luo = Luo(self.symbolic_force_vector)
+        _, force_cms = luo.symmetric_central_moment_forcing(lb_method, central_moments)
+        force_cms = sp.Rational(1, 2) * force_cms
+        return force_cms, force_cms
+
+    def equilibrium_velocity_shift(self, density):
+        return default_velocity_shift(density, self.symbolic_force_vector)
+
+
+class He(AbstractForceModel):
+    r"""
+    Force model by He  :cite:`HeForce`
+    Adapts the calculation of the macroscopic velocity as well as the equilibrium velocity
+    (both shifted by :math:`\frac{\mathbf{F}}{2}`)!
+
+    Force moments are derived from the continuous maxwellian equilibrium. From the
+    moment integrals of the continuous force term
+
+    .. math::
+
+        F (\mathbf{c})
+        = \frac{1}{\rho c_s^2}
+          \mathbf{F} \cdot ( \mathbf{c} - \mathbf{u} )
+          f^{\mathrm{eq}} (\mathbf{c})
+
+    the following analytical expresson for the monomial raw moments of the force is found:
+
+    .. math::
+
+        m_{\alpha\beta\gamma}^{F, \mathrm{He}}
+            = \frac{1}{\rho c_s^2} \left(
+                F_x m^{\mathrm{eq}}_{\alpha+1,\beta,\gamma}
+                + F_y m^{\mathrm{eq}}_{\alpha,\beta+1,\gamma}
+                + F_z m^{\mathrm{eq}}_{\alpha,\beta,\gamma+1}
+                - m^{eq}_{\alpha\beta\gamma} ( \mathbf{F} \cdot \mathbf{u} )
+            \right)
+    """
+
+    def __init__(self, force):
+        super(He, self).__init__(force)
+
+    def forcing_terms(self, lb_method):
+        u = sp.Matrix(lb_method.first_order_equilibrium_moment_symbols)
+        force = sp.Matrix(self.symbolic_force_vector)
+
+        cs_sq = sp.Rational(1, 3)  # squared speed of sound
+        # eq. 6.31 in the LB book by Krüger et al. shows that the equilibrium terms are devided by rho.
+        # This is done here with subs({rho: 1}) to be consistent with compressible and incompressible force models
+        cqc = lb_method.conserved_quantity_computation
+        eq_terms = discrete_maxwellian_equilibrium(lb_method.stencil, rho=sp.Integer(1),
+                                                   u=cqc.velocity_symbols, c_s_sq=sp.Rational(1, 3))
+
+        result = []
+        for direction, eq in zip(lb_method.stencil, eq_terms):
+            direction = sp.Matrix(direction)
+            eu_dot_f = (direction - u).dot(force)
+            result.append(eq * eu_dot_f / cs_sq)
+
+        return sp.Matrix(result)
+
+    def continuous_force_raw_moments(self, lb_method, moments=None):
+        rho = lb_method.zeroth_order_equilibrium_moment_symbol
+        u = lb_method.first_order_equilibrium_moment_symbols
+        dim = lb_method.dim
+        c_s_sq = sp.Rational(1, 3)
+        force = sp.Matrix(self.symbolic_force_vector)
+
+        moment_polynomials = lb_method.moments if moments is None else moments
+        moment_exponents = sorted(extract_monomials(moment_polynomials), key=exponent_tuple_sort_key)
+        moment_monomials = exponents_to_polynomial_representations(moment_exponents)
+        extended_monomials = set()
+        for m in moment_monomials:
+            extended_monomials |= {m} | {m * x for x in MOMENT_SYMBOLS[:dim]}
+
+        extended_monomials = sorted(extended_monomials, key=moment_sort_key)
+        moment_eq_values = get_equilibrium_values_of_maxwell_boltzmann_function(extended_monomials, dim, rho=rho,
+                                                                                u=u, c_s_sq=c_s_sq)
+        moment_to_eq_dict = {m: v for m, v in zip(extended_monomials, moment_eq_values)}
+
+        monomial_force_moments = []
+        for moment in moment_monomials:
+            m_base = moment_to_eq_dict[moment]
+            m_shifted = sp.Matrix([moment_to_eq_dict[moment * x] for x in MOMENT_SYMBOLS[:dim]])
+            m_force = (c_s_sq * rho)**(-1) * (force.dot(m_shifted) - m_base * force.dot(u))
+            monomial_force_moments.append(m_force.expand())
+
+        mono_to_poly_matrix = monomial_to_polynomial_transformation_matrix(moment_exponents, moment_polynomials)
+        polynomial_force_moments = mono_to_poly_matrix * sp.Matrix(monomial_force_moments)
+        return polynomial_force_moments
+
+    def continuous_force_central_moments(self, lb_method, moments=None):
+        if moments is None:
+            moments = lb_method.moments
+        raw_moments = self.continuous_force_raw_moments(lb_method, moments=moments)
+        u = lb_method.first_order_equilibrium_moment_symbols
+        shift_matrix = set_up_shift_matrix(moments, lb_method.stencil, velocity_symbols=u)
+        return (shift_matrix * raw_moments).expand()
+
+    def __call__(self, lb_method):
+        if len(set(lb_method.relaxation_rates)) == 1:
+            #   It's an SRT method!
+            rr = lb_method.symbolic_relaxation_matrix[0]
+            force_terms = self.forcing_terms(lb_method)
+            correction_factor = (1 - rr / 2)
+            result = correction_factor * force_terms
+        else:
+            force_terms = self.moment_space_forcing(lb_method)
+            result = (lb_method.moment_matrix.inv() * force_terms).expand()
+        return result
+
+    def moment_space_forcing(self, lb_method):
+        correction_factor = sp.eye(len(lb_method.stencil)) - sp.Rational(1, 2) * lb_method.symbolic_relaxation_matrix
+        moments = self.continuous_force_raw_moments(lb_method)
+        moments = (correction_factor * moments).expand()
+        return moments
+
+    def central_moment_space_forcing(self, lb_method):
+        correction_factor = sp.eye(len(lb_method.stencil)) - sp.Rational(1, 2) * lb_method.symbolic_relaxation_matrix
+        central_moments = self.continuous_force_central_moments(lb_method)
+        central_moments = (correction_factor * central_moments).expand()
+        return central_moments
+
+    def symmetric_central_moment_forcing(self, lb_method, central_moments):
+        central_moments = exponents_to_polynomial_representations(central_moments)
+        force_cms = sp.Rational(1, 2) * self.continuous_force_central_moments(lb_method, moments=central_moments)
+        return force_cms, force_cms
+
+    def equilibrium_velocity_shift(self, density):
+        return default_velocity_shift(density, self.symbolic_force_vector)
+
+
+class Schiller(Guo):
+    r"""
+    Force model by Schiller  :cite:`schiller2008thermal`, equation 4.67
+    Equivalent to the generalized Guo model.
+    """
+
+    def __init__(self, force):
+        warn("The Schiller force model is deprecated, please use the Guo model, which is equivalent",
+             DeprecationWarning)
+        super(Schiller, self).__init__(force)
+
+
+class Buick(AbstractForceModel):
+    r"""
+    This force model :cite:`buick2000gravity` has a force term with zero second moment. 
+    It is suited for incompressible lattice models. However it should be used with care because such a LB body form
+    model is only consistent when applied to the solution of steady - state hydrodynamic problems. More information
+    on an analysis of the Buick force model can be found in :cite:`silva2010` and in :cite:`silva2020`
+    """
+
+    def __call__(self, lb_method, **kwargs):
+        if len(set(lb_method.relaxation_rates)) == 1:
+            #   It's an SRT method!
+            rr = lb_method.symbolic_relaxation_matrix[0]
+            force_terms = Simple(self.symbolic_force_vector)(lb_method)
+            correction_factor = (1 - rr / 2)
+            result = correction_factor * force_terms
+        else:
+            force_terms = self.moment_space_forcing(lb_method)
+            result = (lb_method.moment_matrix.inv() * force_terms).expand()
+        return result
+
+    def moment_space_forcing(self, lb_method):
+        simple = Simple(self.symbolic_force_vector)
+        q = len(lb_method.stencil)
+        correction_factor = sp.eye(q) - sp.Rational(1, 2) * lb_method.symbolic_relaxation_matrix
+        moments = correction_factor * (lb_method.moment_matrix * sp.Matrix(simple(lb_method)))
+        return moments.expand()
+
+    def central_moment_space_forcing(self, lb_method):
+        simple = Simple(self.symbolic_force_vector)
+        q = len(lb_method.stencil)
+        correction_factor = sp.eye(q) - sp.Rational(1, 2) * lb_method.symbolic_relaxation_matrix
+        moments = (lb_method.moment_matrix * sp.Matrix(simple(lb_method)))
+        central_moments = correction_factor * (lb_method.shift_matrix * moments)
+
+        return central_moments.expand()
+
+    def equilibrium_velocity_shift(self, density):
+        return default_velocity_shift(density, self.symbolic_force_vector)
+
+
+class EDM(AbstractForceModel):
+    r"""Exact differencing force model as shown in :cite:`lbm_book` in eq. 6.32"""
+
+    def __call__(self, lb_method):
+        cqc = lb_method.conserved_quantity_computation
+        reference_density = cqc.density_symbol if cqc.compressible else 1
+        rho = cqc.density_symbol
+        delta_rho = cqc.density_deviation_symbol
+        rho_0 = cqc.background_density
+        u = cqc.velocity_symbols
+
+        equilibrium_terms = lb_method.get_equilibrium_terms()
+        equilibrium_terms = equilibrium_terms.subs({delta_rho: rho - rho_0})
+
+        shifted_u = (u_i + f_i / reference_density for u_i, f_i in zip(u, self._force))
+        shifted_eq = equilibrium_terms.subs({u_i: su_i for u_i, su_i in zip(u, shifted_u)})
+        return shifted_eq - equilibrium_terms
+
+    def moment_space_forcing(self, lb_method):
+        moments = lb_method.moment_matrix * self(lb_method)
+        return moments.expand()
+
+    def central_moment_space_forcing(self, lb_method):
+        moments = lb_method.moment_matrix * self(lb_method)
+        central_moments = lb_method.shift_matrix * moments.expand()
+        return central_moments.expand()
+
+
+class ShanChen(AbstractForceModel):
+    r"""Shan and Chen force model. The implementation is done according to :cite:`silva2020`.
+        For reference compare table 1 which is the Shan and Chen model for an SRT collision operator. These terms are
+        transfered to the moment space and then all representations for the different collision operators are derived
+        from that.
+    """
+
+    def forcing_terms(self, lb_method):
+        q = len(lb_method.stencil)
+        cqc = lb_method.conserved_quantity_computation
+        rho = cqc.density_symbol if cqc.compressible else 1
+        u = cqc.velocity_symbols
+
+        F = sp.Matrix(self.symbolic_force_vector)
+        uf = sp.Matrix(u).dot(F)
+        F2 = sp.Matrix(F).dot(sp.Matrix(F))
+        Fq = sp.zeros(q, 1)
+        uq = sp.zeros(q, 1)
+        for i, cq in enumerate(lb_method.stencil):
+            Fq[i] = sp.Matrix(cq).dot(sp.Matrix(F))
+            uq[i] = sp.Matrix(cq).dot(u)
+
+        linear_term = sp.zeros(q, 1)
+        non_linear_term = sp.zeros(q, 1)
+        for i, w_i in enumerate(lb_method.weights):
+            linear_term[i] = w_i * (Fq[i] + 3 * uq[i] * Fq[i] - uf)
+            non_linear_term[i] = ((w_i / (2 * rho)) * (3 * Fq[i] ** 2 - F2))
+
+        return linear_term, non_linear_term
+
+    def __call__(self, lb_method):
+        force_terms = self.moment_space_forcing(lb_method)
+        result = lb_method.moment_matrix.inv() * force_terms
+        return result.expand()
+
+    def moment_space_forcing(self, lb_method):
+        linear_term, non_linear_term = self.forcing_terms(lb_method)
+        q = len(lb_method.stencil)
+
+        rel = lb_method.symbolic_relaxation_matrix
+        cs_sq = sp.Rational(1, 3)
+
+        correction_factor = 1 / cs_sq * (sp.eye(q) - sp.Rational(1, 2) * rel)
+        M = lb_method.moment_matrix
+        moments = correction_factor * (M * linear_term) + correction_factor ** 2 * (M * non_linear_term)
+        return moments.expand()
+
+    def central_moment_space_forcing(self, lb_method):
+        linear_term, non_linear_term = self.forcing_terms(lb_method)
+        q = len(lb_method.stencil)
+
+        rel = lb_method.symbolic_relaxation_matrix
+        cs_sq = sp.Rational(1, 3)
+
+        correction_factor = 1 / cs_sq * (sp.eye(q) - sp.Rational(1, 2) * rel)
+        M = lb_method.moment_matrix
+        N = lb_method.shift_matrix
+        moments_linear_term = (M * linear_term)
+        moments_non_linear_term = (M * non_linear_term)
+        central_moments_linear_term = correction_factor * (N * moments_linear_term)
+        central_moments_non_linear_term = correction_factor ** 2 * (N * moments_non_linear_term)
+        central_moments = central_moments_linear_term + central_moments_non_linear_term
+        return central_moments.expand()
+
+    def equilibrium_velocity_shift(self, density):
+        return default_velocity_shift(density, self.symbolic_force_vector)
+
+
+# --------------------------------  Helper functions  ------------------------------------------------------------------
+
+
+def default_velocity_shift(density, force):
+    return [f_i / (2 * density) for f_i in force]
+
+
+def default_momentum_density_shift(force):
+    return [f_i / 2 for f_i in force]

lbmpy/geometry.py → src/lbmpy/geometry.py

@@ -231,17 +231,17 @@ def add_black_and_white_image(boundary_handling, image_file, target_slice=None,
     # binarize
     zoomed_image[zoomed_image <= 254] = 0
     zoomed_image[zoomed_image > 254] = 1
-    zoomed_image = np.logical_not(zoomed_image.astype(np.bool))
+    zoomed_image = np.logical_not(zoomed_image.astype(bool))
 
     # resize necessary if aspect ratio should be constant
     if zoomed_image.shape != target_size:
-        resized_image = np.zeros(target_size, dtype=np.bool)
+        resized_image = np.zeros(target_size, dtype=bool)
         mid = [(ts - s) // 2 for ts, s in zip(target_size, zoomed_image.shape)]
         resized_image[mid[0]:zoomed_image.shape[0] + mid[0], mid[1]:zoomed_image.shape[1] + mid[1]] = zoomed_image
         zoomed_image = resized_image
 
     def callback(*coordinates):
-        result = np.zeros_like(coordinates[0], dtype=np.bool)
+        result = np.zeros_like(coordinates[0], dtype=bool)
         mask_start = [int(coordinates[i][(0,) * dim] - 0.5) for i in range(dim)]
         mask_end = [int(coordinates[i][(-1,) * dim] + 1 - 0.5) for i in range(dim)]
 

lbmpy/lbstep.py → src/lbmpy/lbstep.py

 from types import MappingProxyType
+from dataclasses import replace
 
 import numpy as np
 
 from lbmpy.boundaries.boundaryhandling import LatticeBoltzmannBoundaryHandling
-from lbmpy.creationfunctions import (
-    create_lb_function, switch_to_symbolic_relaxation_rates_for_omega_adapting_methods,
-    update_with_default_parameters)
+from lbmpy.creationfunctions import (create_lb_function, update_with_default_parameters)
+from lbmpy.enums import Stencil
 from lbmpy.macroscopic_value_kernels import (
     create_advanced_velocity_setter_collision_rule, pdf_initialization_assignments)
 from lbmpy.simplificationfactory import create_simplification_strategy
-from lbmpy.stencils import get_stencil
-from pystencils import create_data_handling, create_kernel, make_slice
+from lbmpy.stencils import LBStencil
+from pystencils import create_data_handling, create_kernel, make_slice, Target, Backend
 from pystencils.slicing import SlicedGetter
 from pystencils.timeloop import TimeLoop
 
@@ -22,32 +22,50 @@ class LatticeBoltzmannStep:
                  velocity_data_name=None, density_data_name=None, density_data_index=None,
                  compute_velocity_in_every_step=False, compute_density_in_every_step=False,
                  velocity_input_array_name=None, time_step_order='stream_collide', flag_interface=None,
-                 alignment_if_vectorized=64, fixed_loop_sizes=True, fixed_relaxation_rates=True, **method_parameters):
+                 alignment_if_vectorized=64, fixed_loop_sizes=True,
+                 timeloop_creation_function=TimeLoop,
+                 lbm_config=None, lbm_optimisation=None, config=None, **method_parameters):
+
+        if optimization is None:
+            optimization = {}
+        self._timeloop_creation_function = timeloop_creation_function
 
         # --- Parameter normalization  ---
         if data_handling is not None:
             if domain_size is not None:
                 raise ValueError("When passing a data_handling, the domain_size parameter can not be specified")
 
+        if config is not None:
+            target = config.target
+        else:
+            target = optimization.get('target', Target.CPU)
+
         if data_handling is None:
             if domain_size is None:
                 raise ValueError("Specify either domain_size or data_handling")
-            data_handling = create_data_handling(domain_size, default_ghost_layers=1,
-                                                 periodicity=periodicity, parallel=False)
+            data_handling = create_data_handling(domain_size,
+                                                 default_ghost_layers=1,
+                                                 periodicity=periodicity,
+                                                 default_target=target,
+                                                 parallel=False)
+
+        if lbm_config:
+            method_parameters['stencil'] = lbm_config.stencil
 
         if 'stencil' not in method_parameters:
-            method_parameters['stencil'] = 'D2Q9' if data_handling.dim == 2 else 'D3Q27'
+            method_parameters['stencil'] = LBStencil(Stencil.D2Q9) \
+                if data_handling.dim == 2 else LBStencil(Stencil.D3Q27)
 
-        method_parameters, optimization = update_with_default_parameters(method_parameters, optimization)
-        field_dtype = np.float64 if optimization['double_precision'] else np.float32
+        lbm_config, lbm_optimisation, config = update_with_default_parameters(method_parameters, optimization,
+                                                                              lbm_config, lbm_optimisation, config)
 
-        del method_parameters['kernel_type']
+        # the parallel datahandling understands only numpy datatypes. Strings lead to an errors
+        field_dtype = config.data_type.default_factory().numpy_dtype
 
         if lbm_kernel:
-            q = len(lbm_kernel.method.stencil)
+            q = lbm_kernel.method.stencil.Q
         else:
-            q = len(get_stencil(method_parameters['stencil']))
-        target = optimization['target']
+            q = lbm_config.stencil.Q
 
         self.name = name
         self._data_handling = data_handling
@@ -56,13 +74,12 @@ class LatticeBoltzmannStep:
         self.velocity_data_name = name + "_velocity" if velocity_data_name is None else velocity_data_name
         self.density_data_name = name + "_density" if density_data_name is None else density_data_name
         self.density_data_index = density_data_index
-        self._optimization = optimization
 
-        self._gpu = target == 'gpu'
-        layout = optimization['field_layout']
+        self._gpu = target == Target.GPU
+        layout = lbm_optimisation.field_layout
 
         alignment = False
-        if optimization['target'] == 'cpu' and optimization['vectorization']:
+        if config.backend == Backend.C and config.cpu_vectorize_info:
             alignment = alignment_if_vectorized
 
         self._data_handling.add_array(self._pdf_arr_name, values_per_cell=q, gpu=self._gpu, layout=layout,
@@ -80,49 +97,54 @@ class LatticeBoltzmannStep:
                                           layout=layout, latex_name='ρ', dtype=field_dtype, alignment=alignment)
 
         if compute_velocity_in_every_step:
-            method_parameters['output']['velocity'] = self._data_handling.fields[self.velocity_data_name]
+            lbm_config.output['velocity'] = self._data_handling.fields[self.velocity_data_name]
         if compute_density_in_every_step:
             density_field = self._data_handling.fields[self.density_data_name]
             if self.density_data_index is not None:
                 density_field = density_field(density_data_index)
-            method_parameters['output']['density'] = density_field
+            lbm_config.output['density'] = density_field
         if velocity_input_array_name is not None:
-            method_parameters['velocity_input'] = self._data_handling.fields[velocity_input_array_name]
-        if method_parameters['omega_output_field'] and isinstance(method_parameters['omega_output_field'], str):
-            method_parameters['omega_output_field'] = data_handling.add_array(method_parameters['omega_output_field'],
-                                                                              dtype=field_dtype, alignment=alignment)
+            lbm_config = replace(lbm_config, velocity_input=self._data_handling.fields[velocity_input_array_name])
+        if isinstance(lbm_config.omega_output_field, str):
+            lbm_config = replace(lbm_config, omega_output_field=data_handling.add_array(lbm_config.omega_output_field,
+                                                                                        dtype=field_dtype,
+                                                                                        alignment=alignment,
+                                                                                        values_per_cell=1))
 
         self.kernel_params = kernel_params.copy()
 
         # --- Kernel creation ---
         if lbm_kernel is None:
-            switch_to_symbolic_relaxation_rates_for_omega_adapting_methods(method_parameters, self.kernel_params,
-                                                                           force=not fixed_relaxation_rates)
+
             if fixed_loop_sizes:
-                optimization['symbolic_field'] = data_handling.fields[self._pdf_arr_name]
-            method_parameters['field_name'] = self._pdf_arr_name
-            method_parameters['temporary_field_name'] = self._tmp_arr_name
+                lbm_optimisation = replace(lbm_optimisation, symbolic_field=data_handling.fields[self._pdf_arr_name])
+            lbm_config = replace(lbm_config, field_name=self._pdf_arr_name)
+            lbm_config = replace(lbm_config, temporary_field_name=self._tmp_arr_name)
+
             if time_step_order == 'stream_collide':
-                self._lbmKernels = [create_lb_function(optimization=optimization,
-                                                       **method_parameters)]
+                self._lbmKernels = [create_lb_function(lbm_config=lbm_config,
+                                                       lbm_optimisation=lbm_optimisation,
+                                                       config=config)]
             elif time_step_order == 'collide_stream':
-                self._lbmKernels = [create_lb_function(optimization=optimization,
-                                                       kernel_type='collide_only',
-                                                       **method_parameters),
-                                    create_lb_function(optimization=optimization,
-                                                       kernel_type='stream_pull_only',
-                                                       ** method_parameters)]
+                self._lbmKernels = [create_lb_function(lbm_config=lbm_config,
+                                                       lbm_optimisation=lbm_optimisation,
+                                                       config=config,
+                                                       kernel_type='collide_only'),
+                                    create_lb_function(lbm_config=lbm_config,
+                                                       lbm_optimisation=lbm_optimisation,
+                                                       config=config,
+                                                       kernel_type='stream_pull_only')]
 
         else:
             assert self._data_handling.dim == lbm_kernel.method.dim, \
-                "Error: %dD Kernel for %d dimensional domain" % (lbm_kernel.method.dim, self._data_handling.dim)
+                f"Error: {lbm_kernel.method.dim}D Kernel for {self._data_handling.dim} dimensional domain"
             self._lbmKernels = [lbm_kernel]
 
         self.method = self._lbmKernels[0].method
         self.ast = self._lbmKernels[0].ast
 
         # -- Boundary Handling  & Synchronization ---
-        stencil_name = method_parameters['stencil']
+        stencil_name = lbm_config.stencil.name
         self._sync_src = data_handling.synchronization_function([self._pdf_arr_name], stencil_name, target,
                                                                 stencil_restricted=True)
         self._sync_tmp = data_handling.synchronization_function([self._tmp_arr_name], stencil_name, target,
@@ -131,7 +153,11 @@ class LatticeBoltzmannStep:
         self._boundary_handling = LatticeBoltzmannBoundaryHandling(self.method, self._data_handling, self._pdf_arr_name,
                                                                    name=name + "_boundary_handling",
                                                                    flag_interface=flag_interface,
-                                                                   target=target, openmp=optimization['openmp'])
+                                                                   target=target, openmp=config.cpu_openmp)
+
+        self._lbm_config = lbm_config
+        self._lbm_optimisation = lbm_optimisation
+        self._config = config
 
         # -- Macroscopic Value Kernels
         self._getterKernel, self._setterKernel = self._compile_macroscopic_setter_and_getter()
@@ -184,6 +210,21 @@ class LatticeBoltzmannStep:
     def pdf_array_name(self):
         return self._pdf_arr_name
 
+    @property
+    def lbm_config(self):
+        """LBM configuration of the scenario"""
+        return self._lbm_config
+
+    @property
+    def lbm_optimisation(self):
+        """LBM optimisation parameters"""
+        return self._lbm_optimisation
+
+    @property
+    def config(self):
+        """Configutation of pystencils parameters"""
+        return self.config
+
     def _get_slice(self, data_name, slice_obj, masked):
         if slice_obj is None:
             slice_obj = make_slice[:, :] if self.dim == 2 else make_slice[:, :, 0.5]
@@ -243,7 +284,7 @@ class LatticeBoltzmannStep:
     def get_time_loop(self):
         self.pre_run()  # make sure GPU arrays are allocated
 
-        fixed_loop = TimeLoop(steps=2)
+        fixed_loop = self._timeloop_creation_function(steps=2)
         fixed_loop.add_pre_run_function(self.pre_run)
         fixed_loop.add_post_run_function(self.post_run)
         fixed_loop.add_single_step_function(self.time_step)
@@ -328,7 +369,7 @@ class LatticeBoltzmannStep:
             tuple (residuum, steps_run) if successful or raises ValueError if not converged
         """
         dh = self.data_handling
-        gpu = self._optimization['target'] == 'gpu'
+        gpu = self._gpu
 
         def on_first_call():
             self._velocity_init_vel_backup = 'velocity_init_vel_backup'
@@ -337,11 +378,11 @@ class LatticeBoltzmannStep:
 
             collision_rule = create_advanced_velocity_setter_collision_rule(self.method, vel_backup_field,
                                                                             velocity_relaxation_rate)
-            optimization = self._optimization.copy()
-            optimization['symbolic_field'] = dh.fields[self._pdf_arr_name]
+            self._lbm_optimisation.symbolic_field = dh.fields[self._pdf_arr_name]
 
             kernel = create_lb_function(collision_rule=collision_rule, field_name=self._pdf_arr_name,
-                                        temporary_field_name=self._tmp_arr_name, optimization=optimization)
+                                        temporary_field_name=self._tmp_arr_name,
+                                        lbm_optimisation=self._lbm_optimisation)
             self._velocity_init_kernel = kernel
 
         def make_velocity_backup():
@@ -377,7 +418,7 @@ class LatticeBoltzmannStep:
             self._data_handling.all_to_cpu()
             self._data_handling.run_kernel(self._getterKernel, **self.kernel_params)
             global_residuum = compute_residuum()
-            print("Initialization iteration {}, residuum {}".format(steps_run, global_residuum))
+            print(f"Initialization iteration {steps_run}, residuum {global_residuum}")
             if np.isnan(global_residuum) or global_residuum < convergence_threshold:
                 break
 
@@ -385,8 +426,8 @@ class LatticeBoltzmannStep:
         converged = global_residuum < convergence_threshold
         if not converged:
             restore_velocity_backup()
-            raise ValueError("Iterative initialization did not converge after %d steps.\n"
-                             "Current residuum is %s" % (steps_run, global_residuum))
+            raise ValueError(f"Iterative initialization did not converge after {steps_run} steps.\n"
+                             f"Current residuum is {global_residuum}")
 
         return global_residuum, steps_run
 
@@ -400,10 +441,10 @@ class LatticeBoltzmannStep:
 
         getter_eqs = cqc.output_equations_from_pdfs(pdf_field.center_vector,
                                                     {'density': rho_field, 'velocity': vel_field})
-        getter_kernel = create_kernel(getter_eqs, target='cpu', cpu_openmp=self._optimization['openmp']).compile()
+        getter_kernel = create_kernel(getter_eqs, target=Target.CPU, cpu_openmp=self._config.cpu_openmp).compile()
 
         setter_eqs = pdf_initialization_assignments(lb_method, rho_field,
                                                     vel_field.center_vector, pdf_field.center_vector)
         setter_eqs = create_simplification_strategy(lb_method)(setter_eqs)
-        setter_kernel = create_kernel(setter_eqs, target='cpu', cpu_openmp=self._optimization['openmp']).compile()
+        setter_kernel = create_kernel(setter_eqs, target=Target.CPU, cpu_openmp=self._config.cpu_openmp).compile()
         return getter_kernel, setter_kernel

lbmpy/macroscopic_value_kernels.py → src/lbmpy/macroscopic_value_kernels.py

 import functools
+import sympy as sp
 from copy import deepcopy
-
 from lbmpy.simplificationfactory import create_simplification_strategy
+from pystencils import create_kernel, CreateKernelConfig, Assignment
 from pystencils.field import Field, get_layout_of_array
+from pystencils.enums import Target
+
+from lbmpy.advanced_streaming.utility import get_accessor, Timestep
+from lbmpy.relaxationrates import get_shear_relaxation_rate
+from lbmpy.utils import second_order_moment_tensor
+
+
+def get_field_accesses(lb_method, pdfs, streaming_pattern, previous_timestep, pre_collision_pdfs):
+    if isinstance(pdfs, Field):
+        accessor = get_accessor(streaming_pattern, previous_timestep)
+        if pre_collision_pdfs:
+            field_accesses = accessor.read(pdfs, lb_method.stencil)
+        else:
+            field_accesses = accessor.write(pdfs, lb_method.stencil)
+    elif streaming_pattern == 'pull' and not pre_collision_pdfs:
+        field_accesses = pdfs
+    else:
+        raise ValueError("Invalid value of pdfs: A PDF field reference is required to derive "
+                         + f"initialization assignments for streaming pattern {streaming_pattern}.")
+    return field_accesses
 
 
-def pdf_initialization_assignments(lb_method, density, velocity, pdfs):
+def pdf_initialization_assignments(lb_method, density, velocity, pdfs,
+                                   streaming_pattern='pull', previous_timestep=Timestep.BOTH,
+                                   set_pre_collision_pdfs=False):
     """Assignments to initialize the pdf field with equilibrium"""
+    field_accesses = get_field_accesses(lb_method, pdfs, streaming_pattern, previous_timestep, set_pre_collision_pdfs)
+
+    if isinstance(density, Field):
+        density = density.center
+    if isinstance(velocity, Field):
+        velocity = velocity.center_vector
+
     cqc = lb_method.conserved_quantity_computation
-    inp_eqs = cqc.equilibrium_input_equations_from_init_values(density, velocity)
+    inp_eqs = cqc.equilibrium_input_equations_from_init_values(density, velocity, force_substitution=False)
     setter_eqs = lb_method.get_equilibrium(conserved_quantity_equations=inp_eqs)
-    setter_eqs = setter_eqs.new_with_substitutions({sym: pdfs[i]
+    setter_eqs = setter_eqs.new_with_substitutions({sym: field_accesses[i]
                                                     for i, sym in enumerate(lb_method.post_collision_pdf_symbols)})
     return setter_eqs
 
 
-def macroscopic_values_getter(lb_method, density, velocity, pdfs):
+def macroscopic_values_getter(lb_method, density, velocity, pdfs,
+                              streaming_pattern='pull', previous_timestep=Timestep.BOTH,
+                              use_pre_collision_pdfs=False):
+
+    field_accesses = get_field_accesses(lb_method, pdfs, streaming_pattern, previous_timestep, use_pre_collision_pdfs)
+
     cqc = lb_method.conserved_quantity_computation
     assert not (velocity is None and density is None)
     output_spec = {}
@@ -23,13 +58,37 @@ def macroscopic_values_getter(lb_method, density, velocity, pdfs):
         output_spec['velocity'] = velocity
     if density is not None:
         output_spec['density'] = density
-    return cqc.output_equations_from_pdfs(pdfs, output_spec)
+    return cqc.output_equations_from_pdfs(field_accesses, output_spec)
 
 
 macroscopic_values_setter = pdf_initialization_assignments
 
 
-def compile_macroscopic_values_getter(lb_method, output_quantities, pdf_arr=None, field_layout='numpy', target='cpu'):
+def strain_rate_tensor_getter(lb_method, strain_rate_tensor, pdfs, streaming_pattern='pull',
+                              previous_timestep=Timestep.BOTH, use_pre_collision_pdfs=False):
+
+    field_accesses = get_field_accesses(lb_method, pdfs, streaming_pattern, previous_timestep, use_pre_collision_pdfs)
+
+    if isinstance(strain_rate_tensor, Field):
+        strain_rate_tensor = strain_rate_tensor.center_vector
+
+    omega_s = get_shear_relaxation_rate(lb_method)
+    equilibrium = lb_method.equilibrium_distribution
+    rho = equilibrium.density if equilibrium.compressible else equilibrium.background_density
+
+    f_neq = sp.Matrix([field_accesses[i] for i in range(lb_method.stencil.Q)]) - lb_method.get_equilibrium_terms()
+    pi = second_order_moment_tensor(f_neq, lb_method.stencil)
+    strain_rate_tensor_equ = - 1.5 * (omega_s / rho) * pi
+
+    assignments = [Assignment(strain_rate_tensor[i * lb_method.stencil.D + j], strain_rate_tensor_equ[i, j])
+                   for i in range(lb_method.stencil.D) for j in range(lb_method.stencil.D)]
+    return assignments
+
+
+def compile_macroscopic_values_getter(lb_method, output_quantities, pdf_arr=None,
+                                      ghost_layers=1, iteration_slice=None,
+                                      field_layout='numpy', target=Target.CPU,
+                                      streaming_pattern='pull', previous_timestep=Timestep.BOTH):
     """
     Create kernel to compute macroscopic value(s) from a pdf field (e.g. density or velocity)
 
@@ -37,8 +96,13 @@ def compile_macroscopic_values_getter(lb_method, output_quantities, pdf_arr=None
         lb_method: instance of :class:`lbmpy.methods.AbstractLbMethod`
         output_quantities: sequence of quantities to compute e.g. ['density', 'velocity']
         pdf_arr: optional numpy array for pdf field - used to get optimal loop structure for kernel
+        ghost_layers: a sequence of pairs for each coordinate with lower and upper nr of ghost layers
+                      that should be excluded from the iteration. If None, the number of ghost layers 
+                      is determined automatically and assumed to be equal for all dimensions.        
+        iteration_slice: if not None, iteration is done only over this slice of the field
         field_layout: layout for output field, also used for pdf field if pdf_arr is not given
-        target: 'cpu' or 'gpu'
+        target: `Target.CPU` or `Target.GPU`
+        previous_step_accessor: The accessor used by the streaming pattern of the previous timestep
 
     Returns:
         a function to compute macroscopic values:
@@ -83,17 +147,13 @@ def compile_macroscopic_values_getter(lb_method, output_quantities, pdf_arr=None
             output_mapping[output_quantity] = output_mapping[output_quantity][0]
 
     stencil = lb_method.stencil
-    pdf_symbols = [pdf_field(i) for i in range(len(stencil))]
+    previous_step_accessor = get_accessor(streaming_pattern, previous_timestep)
+    pdf_symbols = previous_step_accessor.write(pdf_field, stencil)
+
     eqs = cqc.output_equations_from_pdfs(pdf_symbols, output_mapping).all_assignments
 
-    if target == 'cpu':
-        import pystencils.cpu as cpu
-        kernel = cpu.make_python_function(cpu.create_kernel(eqs))
-    elif target == 'gpu':
-        import pystencils.gpucuda as gpu
-        kernel = gpu.make_python_function(gpu.create_cuda_kernel(eqs))
-    else:
-        raise ValueError("Unknown target '%s'. Possible targets are 'cpu' and 'gpu'" % (target,))
+    config = CreateKernelConfig(target=target, ghost_layers=ghost_layers, iteration_slice=iteration_slice)
+    kernel = create_kernel(eqs, config=config).compile()
 
     def getter(pdfs, **kwargs):
         if pdf_arr is not None:
@@ -107,7 +167,10 @@ def compile_macroscopic_values_getter(lb_method, output_quantities, pdf_arr=None
     return getter
 
 
-def compile_macroscopic_values_setter(lb_method, quantities_to_set, pdf_arr=None, field_layout='numpy', target='cpu'):
+def compile_macroscopic_values_setter(lb_method, quantities_to_set, pdf_arr=None,
+                                      ghost_layers=1, iteration_slice=None,
+                                      field_layout='numpy', target=Target.CPU,
+                                      streaming_pattern='pull', previous_timestep=Timestep.BOTH):
     """
     Creates a function that sets a pdf field to specified macroscopic quantities
     The returned function can be called with the pdf field to set as single argument
@@ -116,8 +179,13 @@ def compile_macroscopic_values_setter(lb_method, quantities_to_set, pdf_arr=None
         lb_method: instance of :class:`lbmpy.methods.AbstractLbMethod`
         quantities_to_set: map from conserved quantity name to fixed value or numpy array
         pdf_arr: optional numpy array for pdf field - used to get optimal loop structure for kernel
+        ghost_layers: a sequence of pairs for each coordinate with lower and upper nr of ghost layers
+                      that should be excluded from the iteration. If None, the number of ghost layers 
+                      is determined automatically and assumed to be equal for all dimensions.        
+        iteration_slice: if not None, iteration is done only over this slice of the field
         field_layout: layout of the pdf field if pdf_arr was not given
-        target: 'cpu' or 'gpu'
+        target: `Target.CPU` or `Target.GPU`
+        previous_step_accessor: The accessor used by the streaming pattern of the previous timestep
 
     Returns:
         function taking pdf array as single argument and which sets the field to the given values
@@ -146,7 +214,7 @@ def compile_macroscopic_values_setter(lb_method, quantities_to_set, pdf_arr=None
 
         value_map[quantity_name] = value
 
-    cq_equations = cqc.equilibrium_input_equations_from_init_values(**value_map)
+    cq_equations = cqc.equilibrium_input_equations_from_init_values(**value_map, force_substitution=False)
 
     eq = lb_method.get_equilibrium(conserved_quantity_equations=cq_equations)
     if at_least_one_field_input:
@@ -155,19 +223,15 @@ def compile_macroscopic_values_setter(lb_method, quantities_to_set, pdf_arr=None
     else:
         eq = eq.new_without_subexpressions()
 
-    substitutions = {sym: pdf_field(i) for i, sym in enumerate(lb_method.post_collision_pdf_symbols)}
+    previous_step_accessor = get_accessor(streaming_pattern, previous_timestep)
+    write_accesses = previous_step_accessor.write(pdf_field, lb_method.stencil)
+
+    substitutions = {sym: write_accesses[i] for i, sym in enumerate(lb_method.post_collision_pdf_symbols)}
     eq = eq.new_with_substitutions(substitutions).all_assignments
 
-    if target == 'cpu':
-        import pystencils.cpu as cpu
-        kernel = cpu.make_python_function(cpu.create_kernel(eq))
-        kernel = functools.partial(kernel, **fixed_kernel_parameters)
-    elif target == 'gpu':
-        import pystencils.gpucuda as gpu
-        kernel = gpu.make_python_function(gpu.create_cuda_kernel(eq))
-        kernel = functools.partial(kernel, **fixed_kernel_parameters)
-    else:
-        raise ValueError("Unknown target '%s'. Possible targets are 'cpu' and 'gpu'" % (target,))
+    config = CreateKernelConfig(target=target)
+    kernel = create_kernel(eq, config=config).compile()
+    kernel = functools.partial(kernel, **fixed_kernel_parameters)
 
     def setter(pdfs, **kwargs):
         if pdf_arr is not None:
@@ -193,15 +257,17 @@ def create_advanced_velocity_setter_collision_rule(method, velocity_field: Field
         LB collision rule
     """
     cqc = method.conserved_quantity_computation
-    density_symbol = cqc.defined_symbols(order=0)[1]
-    velocity_symbols = cqc.defined_symbols(order=1)[1]
+    density_symbol = cqc.density_symbol
+    velocity_symbols = cqc.velocity_symbols
 
     # density is computed from pdfs
-    eq_input_from_pdfs = cqc.equilibrium_input_equations_from_pdfs(method.pre_collision_pdf_symbols)
+    eq_input_from_pdfs = cqc.equilibrium_input_equations_from_pdfs(
+        method.pre_collision_pdf_symbols, force_substitution=False)
     eq_input_from_pdfs = eq_input_from_pdfs.new_filtered([density_symbol])
     # velocity is read from input field
     vel_symbols = [velocity_field(i) for i in range(method.dim)]
-    eq_input_from_field = cqc.equilibrium_input_equations_from_init_values(velocity=vel_symbols)
+    eq_input_from_field = cqc.equilibrium_input_equations_from_init_values(
+        velocity=vel_symbols, force_substitution=False)
     eq_input_from_field = eq_input_from_field.new_filtered(velocity_symbols)
     # then both are merged together
     eq_input = eq_input_from_pdfs.new_merged(eq_input_from_field)

lbmpy/max_domain_size_info.py → src/lbmpy/max_domain_size_info.py

@@ -27,16 +27,16 @@ import warnings
 
 import numpy as np
 
-# Optional packages cpuinfo, pycuda and psutil for hardware queries
+# Optional packages cpuinfo, cupy and psutil for hardware queries
 try:
     from cpuinfo import get_cpu_info
 except ImportError:
     get_cpu_info = None
 
 try:
-    from pycuda.autoinit import device
+    import cupy
 except ImportError:
-    device = None
+    cupy = None
 
 try:
     from psutil import virtual_memory
@@ -107,13 +107,17 @@ def memory_sizes_of_current_machine():
 
     if get_cpu_info:
         cpu_info = get_cpu_info()
-        result.update({'L1': cpu_info['l1_data_cache_size'],
-                       'L2': cpu_info['l2_cache_size'],
-                       'L3': cpu_info['l3_cache_size']})
-
-    if device:
-        size = device.total_memory() / (1024 * 1024)
-        result['GPU'] = "{0:.0f} MB".format(size)
+        if 'l1_data_cache_size' in cpu_info:
+            result['L1'] = cpu_info['l1_data_cache_size']
+        result['L2'] = cpu_info['l2_cache_size']
+        if 'l3_cache_size' in cpu_info:
+            result['L3'] = cpu_info['l3_cache_size']
+
+    if cupy:
+        for i in range(cupy.cuda.runtime.getDeviceCount()):
+            device = cupy.cuda.Device(i)
+            size = device.mem_info[1] / (1024 * 1024)
+            result[f'GPU{i}'] = "{0:.0f} MB".format(size)
 
     if virtual_memory:
         mem = virtual_memory()
@@ -122,7 +126,7 @@ def memory_sizes_of_current_machine():
 
     if not result:
         warnings.warn("Couldn't query for any local memory size."
-                      "Install py-cpuinfo to get cache sizes, psutil for RAM size and pycuda for GPU memory size.")
+                      "Install py-cpuinfo to get cache sizes, psutil for RAM size and cupy for GPU memory size.")
 
     return result
 

lbmpy/maxwellian_equilibrium.py → src/lbmpy/maxwellian_equilibrium.py

 """
-This module contains the continuous Maxwell-Boltzmann equilibrium and its discrete polynomial approximation, often
-used to formulate lattice-Boltzmann methods for hydrodynamics.
+This module contains a functional interface to the continuous Maxwell-Boltzmann equilibrium and its discrete 
+polynomial approximation, often used to formulate lattice-Boltzmann methods for hydrodynamics.
 Additionally functions are provided to compute moments and cumulants of these distributions.
+
+The functionality of this module has mostly been replaced by the :mod:`lbmpy.equilibrium` module.
+In particular, the continuous and discrete Maxwellians are now represented by 
+:class:`lbmpy.equilibrium.ContinuousHydrodynamicMaxwellian` and
+:class:`lbmpy.equilibrium.DiscreteHydrodynamicMaxwellian`.
 """
 
 import warnings
 
 import sympy as sp
-from sympy import Rational as R
+from sympy.core.numbers import Zero
 
 from pystencils.cache import disk_cache
+from pystencils.sympyextensions import remove_higher_order_terms
 
+from lbmpy.moments import MOMENT_SYMBOLS
+from lbmpy.continuous_distribution_measures import continuous_moment, continuous_central_moment, continuous_cumulant
 
-def get_weights(stencil, c_s_sq):
-    q = len(stencil)
 
+def get_weights(stencil, c_s_sq=sp.Rational(1, 3)):
     if c_s_sq != sp.Rational(1, 3) and c_s_sq != sp.Symbol("c_s") ** 2:
         warnings.warn("Weights of discrete equilibrium are only valid if c_s^2 = 1/3")
 
     def weight_for_direction(direction):
         abs_sum = sum([abs(d) for d in direction])
-        return get_weights.weights[q][abs_sum]
+        return get_weights.weights[stencil.Q][abs_sum]
     return [weight_for_direction(d) for d in stencil]
 
 
 get_weights.weights = {
     9: {
-        0: R(4, 9),
-        1: R(1, 9),
-        2: R(1, 36),
+        0: sp.Rational(4, 9),
+        1: sp.Rational(1, 9),
+        2: sp.Rational(1, 36),
+    },
+    7: {
+        0: Zero(),
+        1: sp.Rational(1, 6),
     },
     15: {
-        0: R(2, 9),
-        1: R(1, 9),
-        3: R(1, 72),
+        0: sp.Rational(2, 9),
+        1: sp.Rational(1, 9),
+        3: sp.Rational(1, 72),
     },
     19: {
-        0: R(1, 3),
-        1: R(1, 18),
-        2: R(1, 36),
+        0: sp.Rational(1, 3),
+        1: sp.Rational(1, 18),
+        2: sp.Rational(1, 36),
     },
     27: {
-        0: R(8, 27),
-        1: R(2, 27),
-        2: R(1, 54),
-        3: R(1, 216),
+        0: sp.Rational(8, 27),
+        1: sp.Rational(2, 27),
+        2: sp.Rational(1, 54),
+        3: sp.Rational(1, 216),
     }
 }
 
 
 @disk_cache
-def discrete_maxwellian_equilibrium(stencil, rho=sp.Symbol("rho"), u=tuple(sp.symbols("u_0 u_1 u_2")), order=2,
+def discrete_maxwellian_equilibrium(stencil, rho=sp.Symbol("rho"), u=sp.symbols("u_:3"), order=2,
                                     c_s_sq=sp.Symbol("c_s") ** 2, compressible=True):
     """
     Returns the common discrete LBM equilibrium as a list of sympy expressions
@@ -64,65 +75,74 @@ def discrete_maxwellian_equilibrium(stencil, rho=sp.Symbol("rho"), u=tuple(sp.sy
         compressible: compressibility
     """
     weights = get_weights(stencil, c_s_sq)
-    assert len(stencil) == len(weights)
+    assert stencil.Q == len(weights)
+    u = u[:stencil.D]
 
-    dim = len(stencil[0])
-    u = u[:dim]
+    res = [discrete_equilibrium(e_q, u, rho, w_q, order, c_s_sq, compressible) for w_q, e_q in zip(weights, stencil)]
+    return tuple(res)
 
+
+def discrete_equilibrium(v=sp.symbols("v_:3"), u=sp.symbols("u_:3"), rho=sp.Symbol("rho"), weight=sp.Symbol("w"),
+                         order=2, c_s_sq=sp.Symbol("c_s") ** 2, compressible=True):
+    """
+    Returns the common discrete LBM equilibrium depending on the mesoscopic velocity and the directional lattice weight
+
+    Args:
+        v: symbols for mesoscopic velocity
+        u: symbols for macroscopic velocity
+        rho: sympy symbol for the density
+        weight: symbol for stencil weights
+        order: highest order of velocity terms (for hydrodynamics order 2 is sufficient)
+        c_s_sq: square of speed of sound
+        compressible: compressibility
+    """
     rho_outside = rho if compressible else sp.Rational(1, 1)
     rho_inside = rho if not compressible else sp.Rational(1, 1)
 
-    res = []
-    for w_q, e_q in zip(weights, stencil):
-        e_times_u = 0
-        for c_q_alpha, u_alpha in zip(e_q, u):
-            e_times_u += c_q_alpha * u_alpha
-
-        fq = rho_inside + e_times_u / c_s_sq
+    e_times_u = 0
+    for c_q_alpha, u_alpha in zip(v, u):
+        e_times_u += c_q_alpha * u_alpha
 
-        if order <= 1:
-            res.append(fq * rho_outside * w_q)
-            continue
+    fq = rho_inside + e_times_u / c_s_sq
 
-        u_times_u = 0
-        for u_alpha in u:
-            u_times_u += u_alpha * u_alpha
-        fq += sp.Rational(1, 2) / c_s_sq**2 * e_times_u ** 2 - sp.Rational(1, 2) / c_s_sq * u_times_u
+    if order <= 1:
+        return fq * rho_outside * weight
 
-        if order <= 2:
-            res.append(fq * rho_outside * w_q)
-            continue
+    u_times_u = 0
+    for u_alpha in u:
+        u_times_u += u_alpha * u_alpha
+    fq += sp.Rational(1, 2) / c_s_sq**2 * e_times_u ** 2 - sp.Rational(1, 2) / c_s_sq * u_times_u
 
-        fq += sp.Rational(1, 6) / c_s_sq**3 * e_times_u**3 - sp.Rational(1, 2) / c_s_sq**2 * u_times_u * e_times_u
+    if order <= 2:
+        return fq * rho_outside * weight
 
-        res.append(sp.expand(fq * rho_outside * w_q))
+    fq += sp.Rational(1, 6) / c_s_sq**3 * e_times_u**3 - sp.Rational(1, 2) / c_s_sq**2 * u_times_u * e_times_u
 
-    return tuple(res)
+    return sp.expand(fq * rho_outside * weight)
 
 
 @disk_cache
-def generate_equilibrium_by_matching_moments(stencil, moments, rho=sp.Symbol("rho"), u=tuple(sp.symbols("u_0 u_1 u_2")),
+def generate_equilibrium_by_matching_moments(stencil, moments, rho=sp.Symbol("rho"), u=sp.symbols("u_:3"),
                                              c_s_sq=sp.Symbol("c_s") ** 2, order=None):
     """
     Computes discrete equilibrium, by setting the discrete moments to values taken from the continuous Maxwellian.
     The number of moments has to match the number of directions in the stencil. For documentation of other parameters
-    see :func:`get_moments_of_continuous_maxwellian_equilibrium`
+    see :func:`get_equilibrium_values_of_maxwell_boltzmann_function`
     """
     from lbmpy.moments import moment_matrix
-    dim = len(stencil[0])
-    Q = len(stencil)
-    assert len(moments) == Q, "Moment count(%d) does not match stencil size(%d)" % (len(moments), Q)
-    continuous_moments_vector = get_moments_of_continuous_maxwellian_equilibrium(moments, dim, rho, u, c_s_sq, order)
+    assert len(moments) == stencil.Q, f"Moment count({len(moments)}) does not match stencil size({stencil.Q})"
+    continuous_moments_vector = get_equilibrium_values_of_maxwell_boltzmann_function(moments, stencil.D, rho, u, c_s_sq,
+                                                                                     order, space="moment")
     continuous_moments_vector = sp.Matrix(continuous_moments_vector)
     M = moment_matrix(moments, stencil)
-    assert M.rank() == Q, "Rank of moment matrix (%d) does not match stencil size (%d)" % (M.rank(), Q)
+    assert M.rank() == stencil.Q, f"Rank of moment matrix ({M.rank()}) does not match stencil size ({stencil.Q})"
     return M.inv() * continuous_moments_vector
 
 
 @disk_cache
 def continuous_maxwellian_equilibrium(dim=3, rho=sp.Symbol("rho"),
-                                      u=tuple(sp.symbols("u_0 u_1 u_2")),
-                                      v=tuple(sp.symbols("v_0 v_1 v_2")),
+                                      u=sp.symbols("u_:3"),
+                                      v=sp.symbols("v_:3"),
                                       c_s_sq=sp.Symbol("c_s") ** 2):
     """
     Returns sympy expression of Maxwell Boltzmann distribution
@@ -141,15 +161,14 @@ def continuous_maxwellian_equilibrium(dim=3, rho=sp.Symbol("rho"),
     return rho / (2 * sp.pi * c_s_sq) ** (sp.Rational(dim, 2)) * sp.exp(- vel_term / (2 * c_s_sq))
 
 
-# -------------------------------- Equilibrium moments/cumulants  ------------------------------------------------------
-
-
+# -------------------------------- Equilibrium moments  ----------------------------------------------------------------
 @disk_cache
-def get_moments_of_continuous_maxwellian_equilibrium(moments, dim, rho=sp.Symbol("rho"),
-                                                     u=tuple(sp.symbols("u_0 u_1 u_2")),
-                                                     c_s_sq=sp.Symbol("c_s") ** 2, order=None):
+def get_equilibrium_values_of_maxwell_boltzmann_function(moments, dim, rho=sp.Symbol("rho"),
+                                                         u=sp.symbols("u_:3"),
+                                                         c_s_sq=sp.Symbol("c_s") ** 2, order=None,
+                                                         space="moment"):
     """
-    Computes moments of the continuous Maxwell Boltzmann equilibrium distribution
+    Computes equilibrium values from the continuous Maxwell Boltzmann equilibrium.
 
     Args:
         moments: moments to compute, either in polynomial or exponent-tuple form
@@ -159,19 +178,27 @@ def get_moments_of_continuous_maxwellian_equilibrium(moments, dim, rho=sp.Symbol
         c_s_sq: symbol for speed of sound squared, defaults to symbol c_s**2
         order: if this parameter is not None, terms that have a higher polynomial order in the macroscopic velocity
                are removed
+        space: function space of the equilibrium values. Either moment, central moment or cumulant space are supported.
 
-    >>> get_moments_of_continuous_maxwellian_equilibrium( ( (0,0,0), (1,0,0), (0,1,0), (0,0,1), (2,0,0) ), dim=3 )
+    >>> get_equilibrium_values_of_maxwell_boltzmann_function( ( (0,0,0), (1,0,0), (0,1,0), (0,0,1), (2,0,0) ), dim=3 )
     [rho, rho*u_0, rho*u_1, rho*u_2, rho*(c_s**2 + u_0**2)]
     """
-    from pystencils.sympyextensions import remove_higher_order_terms
-    from lbmpy.moments import MOMENT_SYMBOLS
-    from lbmpy.continuous_distribution_measures import continuous_moment
-
     # trick to speed up sympy integration (otherwise it takes multiple minutes, or aborts):
     # use a positive, real symbol to represent c_s_sq -> then replace this symbol afterwards with the real c_s_sq
     c_s_sq_helper = sp.Symbol("csq_helper", positive=True, real=True)
     mb = continuous_maxwellian_equilibrium(dim, rho, u, MOMENT_SYMBOLS[:dim], c_s_sq_helper)
-    result = [continuous_moment(mb, moment, MOMENT_SYMBOLS[:dim]).subs(c_s_sq_helper, c_s_sq) for moment in moments]
+    if space == "moment":
+        result = [continuous_moment(mb, moment, MOMENT_SYMBOLS[:dim]).subs(c_s_sq_helper, c_s_sq)
+                  for moment in moments]
+    elif space == "central moment":
+        result = [continuous_central_moment(mb, moment, MOMENT_SYMBOLS[:dim], velocity=u).subs(c_s_sq_helper, c_s_sq)
+                  for moment in moments]
+    elif space == "cumulant":
+        result = [continuous_cumulant(mb, moment, MOMENT_SYMBOLS[:dim]).subs(c_s_sq_helper, c_s_sq)
+                  for moment in moments]
+    else:
+        raise ValueError("Only moment, central moment or cumulant space are supported")
+
     if order is not None:
         result = [remove_higher_order_terms(r, order=order, symbols=u) for r in result]
 
@@ -180,7 +207,7 @@ def get_moments_of_continuous_maxwellian_equilibrium(moments, dim, rho=sp.Symbol
 
 @disk_cache
 def get_moments_of_discrete_maxwellian_equilibrium(stencil, moments,
-                                                   rho=sp.Symbol("rho"), u=tuple(sp.symbols("u_0 u_1 u_2")),
+                                                   rho=sp.Symbol("rho"), u=sp.symbols("u_:3"),
                                                    c_s_sq=sp.Symbol("c_s") ** 2, order=None, compressible=True):
     """Compute moments of discrete maxwellian equilibrium.
 
@@ -201,7 +228,7 @@ def get_moments_of_discrete_maxwellian_equilibrium(stencil, moments,
 
 
 def compressible_to_incompressible_moment_value(term, rho, u):
-    """Compressible to incompressible equilibrium moments
+    """**[DEPRECATED]** Compressible to incompressible equilibrium moments
 
     Transforms so-called compressible equilibrium moments (as obtained from the continuous Maxwellian) by
     removing the density factor in all monomials where velocity components are multiplied to the density.
@@ -219,6 +246,8 @@ def compressible_to_incompressible_moment_value(term, rho, u):
     Returns:
         incompressible equilibrium value
     """
+    warnings.warn("Usage of `compressible_to_incompressible_moment_value` is deprecated,"
+                  " and the method will be removed soon.")
     term = sp.sympify(term)
     term = term.expand()
     if term.func != sp.Add:
@@ -235,32 +264,12 @@ def compressible_to_incompressible_moment_value(term, rho, u):
             res += t
     return res
 
-
-# -------------------------------- Equilibrium moments -----------------------------------------------------------------
-
-
-def get_cumulants_of_continuous_maxwellian_equilibrium(cumulants, dim, rho=sp.Symbol("rho"),
-                                                       u=tuple(sp.symbols("u_0 u_1 u_2")), c_s_sq=sp.Symbol("c_s") ** 2,
-                                                       order=None):
-    from lbmpy.moments import MOMENT_SYMBOLS
-    from lbmpy.continuous_distribution_measures import continuous_cumulant
-    from pystencils.sympyextensions import remove_higher_order_terms
-
-    # trick to speed up sympy integration (otherwise it takes multiple minutes, or aborts):
-    # use a positive, real symbol to represent c_s_sq -> then replace this symbol afterwards with the real c_s_sq
-    c_s_sq_helper = sp.Symbol("csq_helper", positive=True, real=True)
-    mb = continuous_maxwellian_equilibrium(dim, rho, u, MOMENT_SYMBOLS[:dim], c_s_sq_helper)
-    result = [continuous_cumulant(mb, cumulant, MOMENT_SYMBOLS[:dim]).subs(c_s_sq_helper, c_s_sq)
-              for cumulant in cumulants]
-    if order is not None:
-        result = [remove_higher_order_terms(r, order=order, symbols=u) for r in result]
-
-    return result
+# -------------------------------- Equilibrium cumulants ---------------------------------------------------------------
 
 
 @disk_cache
 def get_cumulants_of_discrete_maxwellian_equilibrium(stencil, cumulants,
-                                                     rho=sp.Symbol("rho"), u=tuple(sp.symbols("u_0 u_1 u_2")),
+                                                     rho=sp.Symbol("rho"), u=sp.symbols("u_:3"),
                                                      c_s_sq=sp.Symbol("c_s") ** 2, order=None, compressible=True):
     from lbmpy.cumulants import discrete_cumulant
     if order is None:

src/lbmpy/methods/__init__.py

+from .creationfunctions import (
+    CollisionSpaceInfo,
+    create_mrt_orthogonal, create_mrt_raw, create_central_moment, create_srt, create_trt, create_trt_kbc,
+    create_trt_with_magic_number, create_with_continuous_maxwellian_equilibrium,
+    create_with_discrete_maxwellian_equilibrium, create_from_equilibrium,
+    create_cumulant, create_with_default_polynomial_cumulants, create_with_monomial_cumulants)
+
+from .default_moment_sets import mrt_orthogonal_modes_literature, cascaded_moment_sets_literature
+
+from .abstractlbmethod import LbmCollisionRule, AbstractLbMethod, RelaxationInfo
+from .conservedquantitycomputation import AbstractConservedQuantityComputation, DensityVelocityComputation
+
+
+__all__ = ['CollisionSpaceInfo', 'RelaxationInfo', 
+           'AbstractLbMethod', 'LbmCollisionRule',
+           'AbstractConservedQuantityComputation', 'DensityVelocityComputation',
+           'create_srt', 'create_trt', 'create_trt_with_magic_number', 'create_trt_kbc',
+           'create_mrt_orthogonal', 'create_mrt_raw', 'create_central_moment',
+           'create_with_continuous_maxwellian_equilibrium', 'create_with_discrete_maxwellian_equilibrium',
+           'create_from_equilibrium',
+           'mrt_orthogonal_modes_literature', 'cascaded_moment_sets_literature',
+           'create_cumulant', 'create_with_default_polynomial_cumulants',
+           'create_with_monomial_cumulants']

lbmpy/methods/abstractlbmethod.py → src/lbmpy/methods/abstractlbmethod.py

@@ -2,13 +2,18 @@ import abc
 from collections import namedtuple
 
 import sympy as sp
+from sympy.core.numbers import Zero
 
-from pystencils import AssignmentCollection
+from pystencils import Assignment, AssignmentCollection
+from lbmpy.stencils import LBStencil
 
 RelaxationInfo = namedtuple('RelaxationInfo', ['equilibrium_value', 'relaxation_rate'])
 
 
 class LbmCollisionRule(AssignmentCollection):
+    """
+    A pystencils AssignmentCollection that additionally holds an `AbstractLbMethod`
+    """
     def __init__(self, lb_method, *args, **kwargs):
         super(LbmCollisionRule, self).__init__(*args, **kwargs)
         self.method = lb_method
@@ -17,34 +22,66 @@ class LbmCollisionRule(AssignmentCollection):
 class AbstractLbMethod(abc.ABC):
     """Abstract base class for all LBM methods."""
 
-    def __init__(self, stencil):
+    def __init__(self, stencil: LBStencil):
         self._stencil = stencil
 
     @property
     def stencil(self):
-        """Discrete set of velocities, represented as nested tuple"""
+        """Discrete set of velocities, represented by :class:`lbmpy.stencils.LBStencil`"""
         return self._stencil
 
     @property
     def dim(self):
-        return len(self.stencil[0])
+        """The method's spatial dimensionality"""
+        return self._stencil.D
 
     @property
     def pre_collision_pdf_symbols(self):
         """Tuple of symbols representing the pdf values before collision"""
-        return sp.symbols("f_:%d" % (len(self.stencil),))
+        return sp.symbols(f"f_:{self._stencil.Q}")
 
     @property
     def post_collision_pdf_symbols(self):
         """Tuple of symbols representing the pdf values after collision"""
-        return sp.symbols("d_:%d" % (len(self.stencil),))
+        return sp.symbols(f"d_:{self._stencil.Q}")
+
+    @property
+    @abc.abstractmethod
+    def relaxation_rates(self):
+        """Tuple containing the relaxation rates of each moment"""
+
+    @property
+    def relaxation_matrix(self):
+        """Returns a qxq diagonal matrix which contains the relaxation rate for each moment on the diagonal"""
+        d = sp.zeros(len(self.relaxation_rates))
+        for i, w in enumerate(self.relaxation_rates):
+            # note that 0.0 is converted to sp.Zero here. It is not possible to prevent this.
+            d[i, i] = w
+        return d
+
+    @property
+    def symbolic_relaxation_matrix(self):
+        """Returns a qxq diagonal matrix which contains the relaxation rate for each moment on the diagonal.
+           In contrast to the normal relaxation matrix all numeric values are replaced by sympy symbols"""
+        _, d = self._generate_symbolic_relaxation_matrix()
+        return d
+
+    @property
+    def subs_dict_relaxation_rate(self):
+        """returns a dictionary which maps the replaced numerical relaxation rates to its original numerical value"""
+        result = dict()
+        for i in range(self._stencil.Q):
+            result[self.symbolic_relaxation_matrix[i, i]] = self.relaxation_matrix[i, i]
+        return result
 
     # ------------------------- Abstract Methods & Properties ----------------------------------------------------------
 
+    @property
     @abc.abstractmethod
     def conserved_quantity_computation(self):
         """Returns an instance of class :class:`lbmpy.methods.AbstractConservedQuantityComputation`"""
 
+    @property
     @abc.abstractmethod
     def weights(self):
         """Returns a sequence of weights, one for each lattice direction"""
@@ -59,3 +96,36 @@ class AbstractLbMethod(abc.ABC):
     def get_collision_rule(self):
         """Returns an LbmCollisionRule i.e. an equation collection with a reference to the method.
          This collision rule defines the collision operator."""
+
+    # -------------------------------- Helper Functions ----------------------------------------------------------------
+
+    def _generate_symbolic_relaxation_matrix(self, relaxation_rates=None, relaxation_rates_modifier=None):
+        """
+        This function replaces the numbers in the relaxation matrix with symbols in this case, and returns also
+        the subexpressions, that assign the number to the newly introduced symbol
+        """
+        rr = relaxation_rates if relaxation_rates is not None else self.relaxation_rates
+        subexpressions = {}
+        symbolic_relaxation_rates = list()
+        for relaxation_rate in rr:
+            relaxation_rate = sp.sympify(relaxation_rate)
+            if isinstance(relaxation_rate, sp.Symbol):
+                symbolic_relaxation_rate = relaxation_rate
+            else:
+                if isinstance(relaxation_rate, Zero):
+                    relaxation_rate = 0.0
+                if relaxation_rate in subexpressions:
+                    symbolic_relaxation_rate = subexpressions[relaxation_rate]
+                else:
+                    symbolic_relaxation_rate = sp.Symbol(f"rr_{len(subexpressions)}")
+                    subexpressions[relaxation_rate] = symbolic_relaxation_rate
+            symbolic_relaxation_rates.append(symbolic_relaxation_rate)
+
+        substitutions = [Assignment(e[1], e[0]) for e in subexpressions.items()]
+        if relaxation_rates_modifier is not None:
+            symbolic_relaxation_rates = [r * relaxation_rates_modifier for r in symbolic_relaxation_rates]
+        else:
+            for srr in symbolic_relaxation_rates:
+                assert isinstance(srr, sp.Symbol)
+
+        return substitutions, sp.diag(*symbolic_relaxation_rates)

lbmpy/methods/conservedquantitycomputation.py → src/lbmpy/methods/conservedquantitycomputation.py

 import abc
-from collections import OrderedDict
 
+from warnings import warn
 import sympy as sp
 
 from pystencils import Assignment, AssignmentCollection, Field
@@ -48,7 +48,7 @@ class AbstractConservedQuantityComputation(abc.ABC):
         """
 
     @abc.abstractmethod
-    def equilibrium_input_equations_from_pdfs(self, pdfs):
+    def equilibrium_input_equations_from_pdfs(self, pdfs, **kwargs):
         """
         Returns an equation collection that defines all necessary quantities to compute the equilibrium as functions
         of the pdfs.
@@ -59,7 +59,7 @@ class AbstractConservedQuantityComputation(abc.ABC):
         """
 
     @abc.abstractmethod
-    def output_equations_from_pdfs(self, pdfs, output_quantity_names_to_symbols):
+    def output_equations_from_pdfs(self, pdfs, output_quantity_names_to_symbols, **kwargs):
         """
         Returns an equation collection that defines conserved quantities for output. These conserved quantities might
         be slightly different that the ones used as input for the equilibrium e.g. due to a force model.
@@ -82,119 +82,259 @@ class AbstractConservedQuantityComputation(abc.ABC):
 
 
 class DensityVelocityComputation(AbstractConservedQuantityComputation):
-    def __init__(self, stencil, compressible, force_model=None,
-                 zeroth_order_moment_symbol=sp.Symbol("rho"),
-                 first_order_moment_symbols=sp.symbols("u_:3")):
-        dim = len(stencil[0])
+    r"""
+    This class emits equations for in- and output of the conserved quantities of regular
+    hydrodynamic lattice Boltzmann methods, which are density and velocity. The following symbolic
+    quantities are manages by this class:
+
+    - Density :math:`\rho`, background density :math:`\rho_0` (typically set to :math:`1`) and the
+      density deviation :math:`\delta \rho`. They are connected through :math:`\rho = \rho_0 + \delta\rho`.
+    - Velocity :math:`\mathbf{u} = (u_0, u_1, u_2)`.
+
+    Furthermore, this class provides output functionality for the second-order moment tensor :math:`p`.
+
+    Parameters:
+        stencil: see :class:`lbmpy.stencils.LBStencil`
+        compressible: `True` indicates the usage of a compressible equilibrium 
+                        (see :class:`lbmpy.equilibrium.ContinuousHydrodynamicMaxwellian`),
+                        and sets the reference density to :math:`\rho`.
+                        `False` indicates an incompressible equilibrium, using :math:`\rho` as
+                        reference density.
+        zero_centered: Indicates whether or not PDFs are stored in regular or zero-centered format.
+        force_model: Employed force model. See :mod:`lbmpy.forcemodels`.
+    """
+    
+    def __init__(self, stencil, compressible, zero_centered, force_model=None,
+                 background_density=sp.Integer(1),
+                 density_symbol=sp.Symbol("rho"),
+                 density_deviation_symbol=sp.Symbol("delta_rho"),
+                 velocity_symbols=sp.symbols("u_:3"),
+                 c_s_sq=sp.Symbol("c_s")**2,
+                 second_order_moment_symbols=sp.symbols("p_:9"),
+                 zeroth_order_moment_symbol=None,
+                 first_order_moment_symbols=None):
+        if zeroth_order_moment_symbol is not None:
+            warn("Usage of parameter 'zeroth_order_moment_symbol' is deprecated."
+                 " Use 'density_symbol' or 'density_deviation_symbol' instead.")
+            density_symbol = zeroth_order_moment_symbol
+
+        if first_order_moment_symbols is not None:
+            warn("Usage of parameter 'first_order_moment_symbols' is deprecated."
+                 " Use 'velocity_symbols' instead.")
+            velocity_symbols = first_order_moment_symbols
+
+        dim = stencil.D
         self._stencil = stencil
         self._compressible = compressible
-        self._forceModel = force_model
-        self._symbolOrder0 = zeroth_order_moment_symbol
-        self._symbolsOrder1 = first_order_moment_symbols[:dim]
+        self._zero_centered = zero_centered
+        self._force_model = force_model
+        self._background_density = background_density
+        self._density_symbol = density_symbol
+        self._density_deviation_symbol = density_deviation_symbol
+        self._velocity_symbols = velocity_symbols[:dim]
+        self._second_order_moment_symbols = second_order_moment_symbols[:(dim * dim)]
+        self._c_s_sq = c_s_sq
 
     @property
     def conserved_quantities(self):
         return {'density': 1,
-                'velocity': len(self._stencil[0])}
+                'density_deviation': 1,
+                'velocity': self._stencil.D}
 
     @property
     def compressible(self):
+        """Indicates whether a compressible or incompressible equilibrium is used."""
         return self._compressible
 
     def defined_symbols(self, order='all'):
         if order == 'all':
-            return {'density': self._symbolOrder0,
-                    'velocity': self._symbolsOrder1}
+            return {'density': self._density_symbol,
+                    'density_deviation': self._density_deviation_symbol,
+                    'velocity': self._velocity_symbols}
         elif order == 0:
-            return 'density', self._symbolOrder0
+            return {'density': self._density_symbol,
+                    'density_deviation': self._density_deviation_symbol, }
         elif order == 1:
-            return 'velocity', self._symbolsOrder1
+            return {'velocity': self._velocity_symbols}
         else:
-            return None
+            return dict()
 
     @property
     def zero_centered_pdfs(self):
-        return not self._compressible
+        """Whether regular or zero-centered storage is employed."""
+        return self._zero_centered
 
     @property
     def zeroth_order_moment_symbol(self):
-        return self._symbolOrder0
+        """Symbol corresponding to the zeroth-order moment of the stored PDFs, 
+        i.e. `density_deviation_symbol` if zero-centered, else `density_symbol`."""
+        warn("Usage of 'zeroth_order_moment_symbol' is deprecated."
+             "Use 'density_symbol' or 'density_deviation_symbol' instead.")
+        # to be consistent with previous behaviour, this method returns `density_symbol` for non-zero-centered methods,
+        # and `density_deviation_symbol` for zero-centered methods
+        return self._density_deviation_symbol if self.zero_centered_pdfs else self._density_symbol
 
     @property
     def first_order_moment_symbols(self):
-        return self._symbolsOrder1
+        """Symbol corresponding to the first-order moment vector of the stored PDFs, 
+        divided by the reference density."""
+        warn("Usage of 'first_order_moment_symbols' is deprecated. Use 'velocity_symbols' instead.")
+        return self._velocity_symbols
+
+    @property
+    def density_symbol(self):
+        """Symbol for the density."""
+        return self._density_symbol
+
+    @property
+    def density_deviation_symbol(self):
+        """Symbol for the density deviation."""
+        return self._density_deviation_symbol
+
+    @property
+    def background_density(self):
+        """Symbol or value of the background density."""
+        return self._background_density
+
+    @property
+    def velocity_symbols(self):
+        """Symbols for the velocity."""
+        return self._velocity_symbols
+
+    @property
+    def force_model(self):
+        return self._force_model
+
+    def set_force_model(self, force_model):
+        self._force_model = force_model
 
     @property
     def default_values(self):
-        result = {self._symbolOrder0: 1}
-        for s in self._symbolsOrder1:
+        result = {self._density_symbol: 1, self._density_deviation_symbol: 0}
+        for s in self._velocity_symbols:
             result[s] = 0
         return result
 
-    def equilibrium_input_equations_from_pdfs(self, pdfs):
-        dim = len(self._stencil[0])
-        eq_coll = get_equations_for_zeroth_and_first_order_moment(self._stencil, pdfs, self._symbolOrder0,
-                                                                  self._symbolsOrder1[:dim])
-        if self._compressible:
-            eq_coll = divide_first_order_moments_by_rho(eq_coll, dim)
-
-        eq_coll = apply_force_model_shift('equilibrium_velocity_shift', dim, eq_coll,
-                                          self._forceModel, self._compressible)
-        return eq_coll
-
-    def equilibrium_input_equations_from_init_values(self, density=1, velocity=(0, 0, 0)):
-        dim = len(self._stencil[0])
-        zeroth_order_moment = density
-        first_order_moments = velocity[:dim]
-        vel_offset = [0] * dim
+    def equilibrium_input_equations_from_pdfs(self, pdfs, force_substitution=True):
+        #   Compute moments from stored PDFs.
+        #   If storage is zero_centered, this yields the density_deviation, else density as zeroth moment.
+        zeroth_moment_symbol = self._density_deviation_symbol if self.zero_centered_pdfs else self._density_symbol
+        reference_density = self._density_symbol if self.compressible else self._background_density
+        ac = get_equations_for_zeroth_and_first_order_moment(self._stencil, pdfs, zeroth_moment_symbol,
+                                                             self._velocity_symbols)
+
+        main_assignments_dict = ac.main_assignments_dict
 
-        if self._compressible:
-            if self._forceModel and hasattr(self._forceModel, 'macroscopic_velocity_shift'):
-                vel_offset = self._forceModel.macroscopic_velocity_shift(zeroth_order_moment)
+        #   If zero_centered, rho must still be computed, else delta_rho
+        if self.zero_centered_pdfs:
+            main_assignments_dict[self._density_symbol] = self._background_density + self._density_deviation_symbol
         else:
-            if self._forceModel and hasattr(self._forceModel, 'macroscopic_velocity_shift'):
-                vel_offset = self._forceModel.macroscopic_velocity_shift(sp.Rational(1, 1))
-            zeroth_order_moment -= sp.Rational(1, 1)
-        eqs = [Assignment(self._symbolOrder0, zeroth_order_moment)]
+            main_assignments_dict[self._density_deviation_symbol] = self._density_symbol - self._background_density
 
-        first_order_moments = [a - b for a, b in zip(first_order_moments, vel_offset)]
-        eqs += [Assignment(l, r) for l, r in zip(self._symbolsOrder1, first_order_moments)]
+        #   If compressible, velocities must be divided by rho
+        for u in self._velocity_symbols:
+            main_assignments_dict[u] /= reference_density
 
-        return AssignmentCollection(eqs, [])
+        #   If force model is present, apply force model shift
+        if self._force_model is not None:
+            velocity_shifts = self._force_model.equilibrium_velocity_shift(reference_density)
+            for u, s in zip(self._velocity_symbols, velocity_shifts):
+                main_assignments_dict[u] += s
 
-    def output_equations_from_pdfs(self, pdfs, output_quantity_names_to_symbols):
-        dim = len(self._stencil[0])
+            if force_substitution:
+                ac = add_symbolic_force_substitutions(ac, self._force_model._subs_dict_force)
 
-        ac = get_equations_for_zeroth_and_first_order_moment(self._stencil, pdfs,
-                                                             self._symbolOrder0, self._symbolsOrder1)
+        ac.set_main_assignments_from_dict(main_assignments_dict)
+        ac.topological_sort()
+        return ac
+
+    def equilibrium_input_equations_from_init_values(self, density=1, velocity=(0, 0, 0), force_substitution=True):
+        dim = self._stencil.D
+        density_deviation = self._density_symbol - self._background_density
+        velocity = velocity[:dim]
+        vel_offset = [0] * dim
+        reference_density = self._density_symbol if self.compressible else self._background_density
+
+        if self._force_model is not None:
+            vel_offset = self._force_model.macroscopic_velocity_shift(reference_density)
+
+        eqs = [Assignment(self._density_symbol, density),
+               Assignment(self._density_deviation_symbol, density_deviation)]
+
+        velocity_eqs = [u - o for u, o in zip(velocity, vel_offset)]
+        eqs += [Assignment(l, r) for l, r in zip(self._velocity_symbols, velocity_eqs)]
+
+        result = AssignmentCollection(eqs, [])
+
+        if self._force_model is not None and force_substitution:
+            result = add_symbolic_force_substitutions(result, self._force_model._subs_dict_force)
+
+        return result
 
-        if self._compressible:
-            ac = divide_first_order_moments_by_rho(ac, dim)
+    def output_equations_from_pdfs(self, pdfs, output_quantity_names_to_symbols, force_substitution=True):
+        dim = self._stencil.D
+
+        zeroth_moment_symbol = self._density_deviation_symbol if self.zero_centered_pdfs else self._density_symbol
+        first_moment_symbols = sp.symbols(f'momdensity_:{dim}')
+        reference_density = self._density_symbol if self.compressible else self._background_density
+        ac = get_equations_for_zeroth_and_first_order_moment(self._stencil, pdfs,
+                                                             zeroth_moment_symbol, first_moment_symbols,
+                                                             self._second_order_moment_symbols)
+
+        eqs = {**ac.main_assignments_dict, **ac.subexpressions_dict}
+
+        #   If zero_centered, rho must still be computed, else delta_rho
+        if self.zero_centered_pdfs:
+            eqs[self._density_symbol] = self._background_density + self._density_deviation_symbol
+            dim = self._stencil.D
+            for j in range(dim):
+                #   Diagonal entries are shifted by c_s^2
+                eqs[self._second_order_moment_symbols[dim * j + j]] += self._c_s_sq
         else:
-            ac = add_density_offset(ac)
+            eqs[self._density_deviation_symbol] = self._density_symbol - self._background_density
+
+        #   If compressible, velocities must be divided by rho
+        for m, u in zip(first_moment_symbols, self._velocity_symbols):
+            eqs[u] = m / reference_density
+
+        #   If force model is present, apply force model shift
+        mom_density_shifts = [0] * dim
+        if self._force_model is not None:
+            velocity_shifts = self._force_model.macroscopic_velocity_shift(reference_density)
+            for u, s in zip(self._velocity_symbols, velocity_shifts):
+                eqs[u] += s
 
-        ac = apply_force_model_shift('macroscopic_velocity_shift', dim, ac, self._forceModel, self._compressible)
+            mom_density_shifts = self._force_model.macroscopic_momentum_density_shift()
 
         main_assignments = []
-        eqs = OrderedDict([(eq.lhs, eq.rhs) for eq in ac.all_assignments])
 
         if 'density' in output_quantity_names_to_symbols:
             density_output_symbol = output_quantity_names_to_symbols['density']
             if isinstance(density_output_symbol, Field):
                 density_output_symbol = density_output_symbol.center
-            if density_output_symbol != self._symbolOrder0:
-                main_assignments.append(Assignment(density_output_symbol, self._symbolOrder0))
+            if density_output_symbol != self._density_symbol:
+                main_assignments.append(Assignment(density_output_symbol, self._density_symbol))
             else:
-                main_assignments.append(Assignment(self._symbolOrder0, eqs[self._symbolOrder0]))
-                del eqs[self._symbolOrder0]
+                main_assignments.append(Assignment(self._density_symbol, eqs[self._density_symbol]))
+                del eqs[self._density_symbol]
+        if 'density_deviation' in output_quantity_names_to_symbols:
+            density_deviation_output_symbol = output_quantity_names_to_symbols['density_deviation']
+            if isinstance(density_deviation_output_symbol, Field):
+                density_deviation_output_symbol = density_deviation_output_symbol.center
+            if density_deviation_output_symbol != self._density_deviation_symbol:
+                main_assignments.append(Assignment(density_deviation_output_symbol, self._density_deviation_symbol))
+            else:
+                main_assignments.append(Assignment(self._density_deviation_symbol,
+                                        eqs[self._density_deviation_symbol]))
+                del eqs[self._density_deviation_symbol]
         if 'velocity' in output_quantity_names_to_symbols:
             vel_output_symbols = output_quantity_names_to_symbols['velocity']
             if isinstance(vel_output_symbols, Field):
                 vel_output_symbols = vel_output_symbols.center_vector
-            if tuple(vel_output_symbols) != tuple(self._symbolsOrder1):
-                main_assignments += [Assignment(a, b) for a, b in zip(vel_output_symbols, self._symbolsOrder1)]
+            if tuple(vel_output_symbols) != tuple(self._velocity_symbols):
+                main_assignments += [Assignment(a, b) for a, b in zip(vel_output_symbols, self._velocity_symbols)]
             else:
-                for u_i in self._symbolsOrder1:
+                for u_i in self._velocity_symbols:
                     main_assignments.append(Assignment(u_i, eqs[u_i]))
                     del eqs[u_i]
         if 'momentum_density' in output_quantity_names_to_symbols:
@@ -204,13 +344,8 @@ class DensityVelocityComputation(AbstractConservedQuantityComputation):
             # Is not optimal when velocity and momentum_density are calculated together,
             # but this is usually not the case
             momentum_density_output_symbols = output_quantity_names_to_symbols['momentum_density']
-            mom_density_eq_coll = get_equations_for_zeroth_and_first_order_moment(self._stencil, pdfs,
-                                                                                  self._symbolOrder0,
-                                                                                  self._symbolsOrder1)
-            mom_density_eq_coll = apply_force_model_shift('macroscopic_velocity_shift', dim, mom_density_eq_coll,
-                                                          self._forceModel, self._compressible)
-            for sym, val in zip(momentum_density_output_symbols, mom_density_eq_coll.main_assignments[1:]):
-                main_assignments.append(Assignment(sym, val.rhs))
+            for sym, m, s in zip(momentum_density_output_symbols, first_moment_symbols, mom_density_shifts):
+                main_assignments.append(Assignment(sym, m + s))
         if 'moment0' in output_quantity_names_to_symbols:
             moment0_output_symbol = output_quantity_names_to_symbols['moment0']
             if isinstance(moment0_output_symbol, Field):
@@ -222,19 +357,31 @@ class DensityVelocityComputation(AbstractConservedQuantityComputation):
                 moment1_output_symbol = moment1_output_symbol.center_vector
             main_assignments.extend([Assignment(lhs, sum(d[i] * pdf for d, pdf in zip(self._stencil, pdfs)))
                                      for i, lhs in enumerate(moment1_output_symbol)])
+        if 'moment2' in output_quantity_names_to_symbols:
+            moment2_output_symbol = output_quantity_names_to_symbols['moment2']
+            if isinstance(moment2_output_symbol, Field):
+                moment2_output_symbol = moment2_output_symbol.center_vector
+            for i, p in enumerate(moment2_output_symbol):
+                main_assignments.append(Assignment(p, eqs[self._second_order_moment_symbols[i]]))
+                del eqs[self._second_order_moment_symbols[i]]
+
+        ac = AssignmentCollection(main_assignments, eqs)
+        if self._force_model is not None and force_substitution:
+            ac = add_symbolic_force_substitutions(ac, self._force_model._subs_dict_force)
+
+        ac.topological_sort()
 
-        ac = ac.copy(main_assignments, [Assignment(a, b) for a, b in eqs.items()])
         return ac.new_without_unused_subexpressions()
 
     def __repr__(self):
-        return "ConservedValueComputation for %s" % (", " .join(self.conserved_quantities.keys()),)
+        return "ConservedValueComputation for %s" % (", ".join(self.conserved_quantities.keys()),)
 
 
 # -----------------------------------------  Helper functions ----------------------------------------------------------
 
 
-def get_equations_for_zeroth_and_first_order_moment(stencil, symbolic_pdfs,
-                                                    symbolic_zeroth_moment, symbolic_first_moments):
+def get_equations_for_zeroth_and_first_order_moment(stencil, symbolic_pdfs, symbolic_zeroth_moment,
+                                                    symbolic_first_moments, symbolic_second_moments=None):
     r"""
     Returns an equation system that computes the zeroth and first order moments with the least amount of operations
 
@@ -244,13 +391,16 @@ def get_equations_for_zeroth_and_first_order_moment(stencil, symbolic_pdfs,
 
         \rho = \sum_{d \in S} f_d
         u_j = \sum_{d \in S} f_d u_jd
+        p_j = \sum_{d \in S} {d \in S} f_d u_jd
 
     Args:
         stencil: called :math:`S` above
         symbolic_pdfs: called :math:`f` above
         symbolic_zeroth_moment:  called :math:`\rho` above
         symbolic_first_moments: called :math:`u` above
+        symbolic_second_moments: called :math:`p` above
     """
+
     def filter_out_plus_terms(expr):
         result = 0
         for term in expr.args:
@@ -258,32 +408,46 @@ def get_equations_for_zeroth_and_first_order_moment(stencil, symbolic_pdfs,
                 result += term
         return result
 
-    dim = len(stencil[0])
+    dim = stencil.D
 
-    subexpressions = []
+    subexpressions = list()
     pdf_sum = sum(symbolic_pdfs)
     u = [0] * dim
     for f, offset in zip(symbolic_pdfs, stencil):
         for i in range(dim):
             u[i] += f * int(offset[i])
 
+    p = [0] * dim * dim
+    for f, offset in zip(symbolic_pdfs, stencil):
+        for i in range(dim):
+            for j in range(dim):
+                p[dim * i + j] += f * int(offset[i]) * int(offset[j])
+
     plus_terms = [set(filter_out_plus_terms(u_i).args) for u_i in u]
+
+    velo_terms = sp.symbols(f"vel:{dim}Term")
     for i in range(dim):
         rhs = plus_terms[i]
         for j in range(i):
             rhs -= plus_terms[j]
-        eq = Assignment(sp.Symbol("vel%dTerm" % (i,)), sum(rhs))
+        eq = Assignment(velo_terms[i], sum(rhs))
         subexpressions.append(eq)
+        if len(rhs) == 0:  # if one of the substitutions is not found the simplification can not be applied
+            subexpressions = []
+            break
 
     for subexpression in subexpressions:
         pdf_sum = pdf_sum.subs(subexpression.rhs, subexpression.lhs)
 
-    for i in range(dim):
-        u[i] = u[i].subs(subexpressions[i].rhs, subexpressions[i].lhs)
+    if len(subexpressions) > 0:
+        for i in range(dim):
+            u[i] = u[i].subs(subexpressions[i].rhs, subexpressions[i].lhs)
 
     equations = []
     equations += [Assignment(symbolic_zeroth_moment, pdf_sum)]
     equations += [Assignment(u_i_sym, u_i) for u_i_sym, u_i in zip(symbolic_first_moments, u)]
+    if symbolic_second_moments:
+        equations += [Assignment(symbolic_second_moments[i], p[i]) for i in range(dim ** 2)]
 
     return AssignmentCollection(equations, subexpressions)
 
@@ -313,23 +477,45 @@ def add_density_offset(assignment_collection, offset=sp.Rational(1, 1)):
     return assignment_collection.copy([new_density] + old_eqs[1:])
 
 
-def apply_force_model_shift(shift_member_name, dim, assignment_collection, force_model, compressible, reverse=False):
+def apply_force_model_shift(shift_func, dim, assignment_collection, compressible, reverse=False):
     """
     Modifies the first order moment equations in assignment collection according to the force model shift.
     It is applied if force model has a method named shift_member_name. The equations 1: dim+1 of the passed
     equation collection are assumed to be the velocity equations.
+
+    Args:
+        shift_func: shift function which is applied. See lbmpy.forcemodels.AbstractForceModel for details
+        dim: number of spatial dimensions
+        assignment_collection: assignment collection containing the conserved quantity computation
+        compressible: True if a compressible LB method is used. Otherwise the Helmholtz decomposition was applied
+                      for rho
+        reverse: If True the sign of the shift is flipped
+    """
+    old_eqs = assignment_collection.main_assignments
+    density = old_eqs[0].lhs if compressible else sp.Rational(1, 1)
+    old_vel_eqs = old_eqs[1:dim + 1]
+    vel_offsets = shift_func(density)
+    if reverse:
+        vel_offsets = [-v for v in vel_offsets]
+    shifted_velocity_eqs = [Assignment(old_eq.lhs, old_eq.rhs + offset)
+                            for old_eq, offset in zip(old_vel_eqs, vel_offsets)]
+    new_eqs = [old_eqs[0]] + shifted_velocity_eqs + old_eqs[dim + 1:]
+    return assignment_collection.copy(new_eqs)
+
+
+def add_symbolic_force_substitutions(assignment_collection, subs_dict):
+    """
+    Every force model holds a symbolic representation of the forcing terms internally. This function adds the
+    equations for the D-dimensional force vector to the symbolic replacements
+
+    Args:
+        assignment_collection: assignment collection which will be modified
+        subs_dict: substitution dict which can be obtained from the force model
     """
-    if force_model is not None and hasattr(force_model, shift_member_name):
-        old_eqs = assignment_collection.main_assignments
-        density = old_eqs[0].lhs if compressible else sp.Rational(1, 1)
-        old_vel_eqs = old_eqs[1:dim + 1]
-        shift_func = getattr(force_model, shift_member_name)
-        vel_offsets = shift_func(density)
-        if reverse:
-            vel_offsets = [-v for v in vel_offsets]
-        shifted_velocity_eqs = [Assignment(old_eq.lhs, old_eq.rhs + offset)
-                                for old_eq, offset in zip(old_vel_eqs, vel_offsets)]
-        new_eqs = [old_eqs[0]] + shifted_velocity_eqs + old_eqs[dim + 1:]
-        return assignment_collection.copy(new_eqs)
-    else:
-        return assignment_collection
+    old_eqs = assignment_collection.subexpressions
+    subs_equations = []
+    for key, value in zip(subs_dict.keys(), subs_dict.values()):
+        subs_equations.append(Assignment(key, value))
+
+    new_eqs = subs_equations + old_eqs
+    return assignment_collection.copy(main_assignments=None, subexpressions=new_eqs)