diff --git a/lbmpy/methods/abstractlbmethod.py b/lbmpy/methods/abstractlbmethod.py
index 3954fcfc913fb1e29253d7e5f52cef40e329a5ad..03747ac9a0808f545982490e35cb34dee1852394 100644
--- a/lbmpy/methods/abstractlbmethod.py
+++ b/lbmpy/methods/abstractlbmethod.py
@@ -5,6 +5,7 @@ import sympy as sp
from sympy.core.numbers import Zero
from pystencils import Assignment, AssignmentCollection
+from lbmpy.stencils import LBStencil
RelaxationInfo = namedtuple('RelaxationInfo', ['equilibrium_value', 'relaxation_rate'])
@@ -21,7 +22,7 @@ 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
@@ -32,17 +33,17 @@ class AbstractLbMethod(abc.ABC):
@property
def dim(self):
"""The method's spatial dimensionality"""
- return len(self.stencil[0])
+ return self._stencil.D
@property
def pre_collision_pdf_symbols(self):
"""Tuple of symbols representing the pdf values before collision"""
- return sp.symbols(f"f_:{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(f"d_:{len(self.stencil)}")
+ return sp.symbols(f"d_:{self._stencil.Q}")
@property
@abc.abstractmethod
@@ -69,7 +70,7 @@ class AbstractLbMethod(abc.ABC):
def subs_dict_relxation_rate(self):
"""returns a dictonary which maps the replaced numerical relaxation rates to its original numerical value"""
result = dict()
- for i in range(len(self.stencil)):
+ for i in range(self._stencil.Q):
result[self.symbolic_relaxation_matrix[i, i]] = self.relaxation_matrix[i, i]
return result
diff --git a/lbmpy/methods/cumulantbased/cumulantbasedmethod.py b/lbmpy/methods/cumulantbased/cumulantbasedmethod.py
index 94e5699252e4b70b08879be7fd127e60a7332e3e..c0bd07008daf2afb4e0a6867c54544512602dd0a 100644
--- a/lbmpy/methods/cumulantbased/cumulantbasedmethod.py
+++ b/lbmpy/methods/cumulantbased/cumulantbasedmethod.py
@@ -1,9 +1,11 @@
import sympy as sp
from collections import OrderedDict
+from typing import Set
from warnings import filterwarnings
from pystencils import Assignment, AssignmentCollection
from pystencils.sympyextensions import is_constant
+from pystencils.simp import apply_to_all_assignments
from lbmpy.methods.abstractlbmethod import AbstractLbMethod, LbmCollisionRule, RelaxationInfo
from lbmpy.methods.conservedquantitycomputation import AbstractConservedQuantityComputation
@@ -210,10 +212,11 @@ class CumulantBasedLbMethod(AbstractLbMethod):
assert len(weights) == len(self.stencil)
self._weights = weights
- def get_equilibrium(self, conserved_quantity_equations=None, subexpressions=False, pre_simplification=False,
- keep_cqc_subexpressions=True, include_force_terms=False):
+ def get_equilibrium(self, conserved_quantity_equations: AssignmentCollection = None, subexpressions: bool = False,
+ pre_simplification: bool = False, keep_cqc_subexpressions: bool = True,
+ include_force_terms: bool = False) -> LbmCollisionRule:
"""Returns equation collection, to compute equilibrium values.
- The equations have the post collision symbols as left hand sides and are
+ The equations have the post collision symbols as left-hand sides and are
functions of the conserved quantities
Args:
@@ -224,39 +227,46 @@ class CumulantBasedLbMethod(AbstractLbMethod):
keep_cqc_subexpressions: if equilibrium is returned without subexpressions keep_cqc_subexpressions
determines if also subexpressions to calculate conserved quantities should be
plugged into the main assignments
+ include_force_terms: if set to True the equilibrium is shifted by forcing terms coming from the force model
+ of the method.
"""
- r_info_dict = {c: RelaxationInfo(info.equilibrium_value, sp.Integer(1))
- for c, info in self.relaxation_info_dict.items()}
- ac = self._centered_cumulant_collision_rule(
- r_info_dict, conserved_quantity_equations, pre_simplification,
- include_force_terms=include_force_terms, symbolic_relaxation_rates=False)
+ r_info_dict = OrderedDict({c: RelaxationInfo(info.equilibrium_value, sp.Integer(1))
+ for c, info in self.relaxation_info_dict.items()})
+ ac = self._centered_cumulant_collision_rule(cumulant_to_relaxation_info_dict=r_info_dict,
+ conserved_quantity_equations=conserved_quantity_equations,
+ pre_simplification=pre_simplification,
+ include_force_terms=include_force_terms,
+ symbolic_relaxation_rates=False)
- # from .cumulant_simplifications import insert_logs
- # ac = insert_logs(ac)
+ expand_all_assignments = apply_to_all_assignments(sp.expand)
if not subexpressions:
if keep_cqc_subexpressions:
bs = self._bound_symbols_cqc(conserved_quantity_equations)
- return ac.new_without_subexpressions(subexpressions_to_keep=bs)
+ ac = expand_all_assignments(ac.new_without_subexpressions(subexpressions_to_keep=bs))
+ return ac.new_without_unused_subexpressions()
else:
- return ac.new_without_subexpressions()
+ ac = expand_all_assignments(ac.new_without_subexpressions())
+ return ac.new_without_unused_subexpressions()
else:
- return ac
+ return ac.new_without_unused_subexpressions()
- def get_equilibrium_terms(self):
+ def get_equilibrium_terms(self) -> sp.Matrix:
equilibrium = self.get_equilibrium()
return sp.Matrix([eq.rhs for eq in equilibrium.main_assignments])
- def get_collision_rule(self, conserved_quantity_equations=None, pre_simplification=False):
+ def get_collision_rule(self, conserved_quantity_equations: AssignmentCollection = None,
+ pre_simplification: bool = False) -> AssignmentCollection:
"""Returns an LbmCollisionRule i.e. an equation collection with a reference to the method.
This collision rule defines the collision operator."""
- return self._centered_cumulant_collision_rule(
- self.relaxation_info_dict, conserved_quantity_equations, pre_simplification, True,
- symbolic_relaxation_rates=True)
+ return self._centered_cumulant_collision_rule(cumulant_to_relaxation_info_dict=self.relaxation_info_dict,
+ conserved_quantity_equations=conserved_quantity_equations,
+ pre_simplification=pre_simplification,
+ include_force_terms=True, symbolic_relaxation_rates=True)
# ------------------------------- Internals --------------------------------------------
- def _bound_symbols_cqc(self, conserved_quantity_equations=None):
+ def _bound_symbols_cqc(self, conserved_quantity_equations: AssignmentCollection = None) -> Set[sp.Symbol]:
f = self.pre_collision_pdf_symbols
cqe = conserved_quantity_equations
@@ -283,11 +293,11 @@ class CumulantBasedLbMethod(AbstractLbMethod):
return [w for w in weights]
- def _centered_cumulant_collision_rule(self, cumulant_to_relaxation_info_dict,
- conserved_quantity_equations=None,
- pre_simplification=False,
- include_force_terms=False,
- symbolic_relaxation_rates=False):
+ def _centered_cumulant_collision_rule(self, cumulant_to_relaxation_info_dict: OrderedDict,
+ conserved_quantity_equations: AssignmentCollection = None,
+ pre_simplification: bool = False,
+ include_force_terms: bool = False,
+ symbolic_relaxation_rates: bool = False) -> LbmCollisionRule:
# Filter out JobLib warnings. They are not usefull for use:
# https://github.com/joblib/joblib/issues/683
diff --git a/lbmpy/methods/momentbased/centralmomentbasedmethod.py b/lbmpy/methods/momentbased/centralmomentbasedmethod.py
index 7cdacd2f521b166d49e81c1c26e8976114a33f5c..e56701d3f475cd5072bcb6339fdd4d5403779c74 100644
--- a/lbmpy/methods/momentbased/centralmomentbasedmethod.py
+++ b/lbmpy/methods/momentbased/centralmomentbasedmethod.py
@@ -1,8 +1,10 @@
import sympy as sp
from collections import OrderedDict
+from typing import Set
from pystencils import Assignment, AssignmentCollection
from pystencils.sympyextensions import is_constant
+from pystencils.simp import apply_to_all_assignments
from lbmpy.methods.abstractlbmethod import AbstractLbMethod, LbmCollisionRule, RelaxationInfo
from lbmpy.methods.conservedquantitycomputation import AbstractConservedQuantityComputation
@@ -152,15 +154,15 @@ class CentralMomentBasedLbMethod(AbstractLbMethod):
self._force_model = force_model
@property
- def moment_matrix(self):
+ def moment_matrix(self) -> sp.Matrix:
return moment_matrix(self.moments, self.stencil)
@property
- def shift_matrix(self):
+ def shift_matrix(self) -> sp.Matrix:
return set_up_shift_matrix(self.moments, self.stencil)
@property
- def relaxation_matrix(self):
+ def relaxation_matrix(self) -> sp.Matrix:
d = sp.zeros(len(self.relaxation_rates))
for i in range(0, len(self.relaxation_rates)):
d[i, i] = self.relaxation_rates[i]
@@ -218,10 +220,11 @@ class CentralMomentBasedLbMethod(AbstractLbMethod):
# ----------------------- Overridden Abstract Members --------------------------
- def get_equilibrium(self, conserved_quantity_equations=None, subexpressions=False, pre_simplification=False,
- keep_cqc_subexpressions=True, include_force_terms=False):
+ def get_equilibrium(self, conserved_quantity_equations: AssignmentCollection = None, subexpressions: bool = False,
+ pre_simplification: bool = False, keep_cqc_subexpressions: bool = True,
+ include_force_terms: bool = False) -> LbmCollisionRule:
"""Returns equation collection, to compute equilibrium values.
- The equations have the post collision symbols as left hand sides and are
+ The equations have the post collision symbols as left-hand sides and are
functions of the conserved quantities
Args:
@@ -232,34 +235,48 @@ class CentralMomentBasedLbMethod(AbstractLbMethod):
keep_cqc_subexpressions: if equilibrium is returned without subexpressions keep_cqc_subexpressions
determines if also subexpressions to calculate conserved quantities should be
plugged into the main assignments
+ include_force_terms: if set to True the equilibrium is shifted by forcing terms coming from the force model
+ of the method.
"""
- r_info_dict = {c: RelaxationInfo(info.equilibrium_value, sp.Integer(1))
- for c, info in self.relaxation_info_dict.items()}
- ac = self._central_moment_collision_rule(r_info_dict, conserved_quantity_equations, pre_simplification,
+ _, d = self._generate_symbolic_relaxation_matrix()
+ rr_sub_expressions = set([Assignment(d[i, i], sp.Integer(1)) for i in range(len(self.relaxation_rates))])
+ r_info_dict = OrderedDict({c: RelaxationInfo(info.equilibrium_value, sp.Integer(1))
+ for c, info in self.relaxation_info_dict.items()})
+ ac = self._central_moment_collision_rule(moment_to_relaxation_info_dict=r_info_dict,
+ conserved_quantity_equations=conserved_quantity_equations,
+ pre_simplification=pre_simplification,
include_force_terms=include_force_terms,
symbolic_relaxation_rates=False)
+ ac.subexpressions = list(rr_sub_expressions) + ac.subexpressions
+ expand_all_assignments = apply_to_all_assignments(sp.expand)
+
if not subexpressions:
if keep_cqc_subexpressions:
bs = self._bound_symbols_cqc(conserved_quantity_equations)
- return ac.new_without_subexpressions(subexpressions_to_keep=bs)
+ ac = expand_all_assignments(ac.new_without_subexpressions(subexpressions_to_keep=bs))
+ return ac.new_without_unused_subexpressions()
else:
- return ac.new_without_subexpressions()
+ ac = expand_all_assignments(ac.new_without_subexpressions())
+ return ac.new_without_unused_subexpressions()
else:
- return ac
+ return ac.new_without_unused_subexpressions()
- def get_equilibrium_terms(self):
+ def get_equilibrium_terms(self) -> sp.Matrix:
equilibrium = self.get_equilibrium()
return sp.Matrix([eq.rhs for eq in equilibrium.main_assignments])
- def get_collision_rule(self, conserved_quantity_equations=None, pre_simplification=False):
+ def get_collision_rule(self, conserved_quantity_equations: AssignmentCollection = None,
+ pre_simplification: bool = False) -> LbmCollisionRule:
"""Returns an LbmCollisionRule i.e. an equation collection with a reference to the method.
This collision rule defines the collision operator."""
- return self._central_moment_collision_rule(self.relaxation_info_dict, conserved_quantity_equations,
- pre_simplification, True, symbolic_relaxation_rates=True)
+ return self._central_moment_collision_rule(moment_to_relaxation_info_dict=self.relaxation_info_dict,
+ conserved_quantity_equations=conserved_quantity_equations,
+ pre_simplification=pre_simplification,
+ include_force_terms=True, symbolic_relaxation_rates=True)
# ------------------------------- Internals --------------------------------------------
- def _bound_symbols_cqc(self, conserved_quantity_equations=None):
+ def _bound_symbols_cqc(self, conserved_quantity_equations: AssignmentCollection = None) -> Set[sp.Symbol]:
f = self.pre_collision_pdf_symbols
cqe = conserved_quantity_equations
@@ -285,11 +302,11 @@ class CentralMomentBasedLbMethod(AbstractLbMethod):
return [w for w in weights]
- def _central_moment_collision_rule(self, moment_to_relaxation_info_dict,
- conserved_quantity_equations=None,
- pre_simplification=False,
- include_force_terms=False,
- symbolic_relaxation_rates=False):
+ def _central_moment_collision_rule(self, moment_to_relaxation_info_dict: OrderedDict,
+ conserved_quantity_equations: AssignmentCollection = None,
+ pre_simplification: bool = False,
+ include_force_terms: bool = False,
+ symbolic_relaxation_rates: bool = False) -> LbmCollisionRule:
stencil = self.stencil
f = self.pre_collision_pdf_symbols
density = self.zeroth_order_equilibrium_moment_symbol
diff --git a/lbmpy/methods/momentbased/momentbasedmethod.py b/lbmpy/methods/momentbased/momentbasedmethod.py
index 6dedcc3d7123c976dd50eab52aed3cacf8b18053..2ce9d51580b6c9d8ea37344fe2d5e606feaae87f 100644
--- a/lbmpy/methods/momentbased/momentbasedmethod.py
+++ b/lbmpy/methods/momentbased/momentbasedmethod.py
@@ -1,4 +1,5 @@
from collections import OrderedDict
+from typing import Iterable, Set
import sympy as sp
import numpy as np
@@ -127,31 +128,57 @@ class MomentBasedLbMethod(AbstractLbMethod):
assert len(weights) == len(self.stencil)
self._weights = weights
- def get_equilibrium(self, conserved_quantity_equations=None, include_force_terms=False,
- pre_simplification=False, subexpressions=False, keep_cqc_subexpressions=True):
- relaxation_matrix = sp.eye(len(self.relaxation_rates))
- ac = self._collision_rule_with_relaxation_matrix(relaxation_matrix,
- conserved_quantity_equations=conserved_quantity_equations,
+ def get_equilibrium(self, conserved_quantity_equations: AssignmentCollection = None,
+ include_force_terms: bool = False, pre_simplification: bool = False,
+ subexpressions: bool = False, keep_cqc_subexpressions: bool = True) -> LbmCollisionRule:
+ """Returns equation collection, to compute equilibrium values.
+ The equations have the post collision symbols as left-hand sides and are
+ functions of the conserved quantities
+
+ Args:
+ conserved_quantity_equations: equations to compute conserved quantities.
+ include_force_terms: if set to True the equilibrium is shifted by forcing terms coming from the force model
+ of the method.
+ pre_simplification: with or without pre-simplifications for the calculation of the collision
+ subexpressions: if set to false all subexpressions of the equilibrium assignments are plugged
+ into the main assignments
+ keep_cqc_subexpressions: if equilibrium is returned without subexpressions keep_cqc_subexpressions
+ determines if also subexpressions to calculate conserved quantities should be
+ plugged into the main assignments
+
+ """
+ _, d = self._generate_symbolic_relaxation_matrix()
+ rr_sub_expressions = set([Assignment(d[i, i], sp.Integer(1)) for i in range(len(self.relaxation_rates))])
+
+ ac = self._collision_rule_with_relaxation_matrix(d=d,
+ additional_subexpressions=rr_sub_expressions,
include_force_terms=include_force_terms,
+ conserved_quantity_equations=conserved_quantity_equations,
pre_simplification=pre_simplification)
+
if not subexpressions:
if keep_cqc_subexpressions:
bs = self._bound_symbols_cqc(conserved_quantity_equations)
- return ac.new_without_subexpressions(subexpressions_to_keep=bs)
+ ac = ac.new_without_subexpressions(subexpressions_to_keep=bs)
+ return ac.new_without_unused_subexpressions()
else:
- return ac.new_without_subexpressions()
+ ac = ac.new_without_subexpressions()
+ return ac.new_without_unused_subexpressions()
else:
- return ac
+ return ac.new_without_unused_subexpressions()
def get_equilibrium_terms(self):
"""Returns this method's equilibrium populations as a vector."""
equilibrium = self.get_equilibrium()
return sp.Matrix([eq.rhs for eq in equilibrium.main_assignments])
- def get_collision_rule(self, conserved_quantity_equations=None, pre_simplification=True):
- relaxation_rate_sub_expressions, d = self._generate_symbolic_relaxation_matrix()
- ac = self._collision_rule_with_relaxation_matrix(d, relaxation_rate_sub_expressions,
- True, conserved_quantity_equations,
+ def get_collision_rule(self, conserved_quantity_equations: AssignmentCollection = None,
+ pre_simplification: bool = True) -> LbmCollisionRule:
+ rr_sub_expressions, d = self._generate_symbolic_relaxation_matrix()
+ ac = self._collision_rule_with_relaxation_matrix(d=d,
+ additional_subexpressions=rr_sub_expressions,
+ include_force_terms=True,
+ conserved_quantity_equations=conserved_quantity_equations,
pre_simplification=pre_simplification)
return ac
@@ -179,27 +206,27 @@ class MomentBasedLbMethod(AbstractLbMethod):
self._cqc.set_force_model(force_model)
@property
- def collision_matrix(self):
+ def collision_matrix(self) -> sp.Matrix:
pdfs_to_moments = self.moment_matrix
d = self.relaxation_matrix
return pdfs_to_moments.inv() * d * pdfs_to_moments
@property
- def inverse_collision_matrix(self):
+ def inverse_collision_matrix(self) -> sp.Matrix:
pdfs_to_moments = self.moment_matrix
inverse_relaxation_matrix = self.relaxation_matrix.inv()
return pdfs_to_moments.inv() * inverse_relaxation_matrix * pdfs_to_moments
@property
- def moment_matrix(self):
+ def moment_matrix(self) -> sp.Matrix:
return moment_matrix(self.moments, self.stencil)
@property
- def is_orthogonal(self):
+ def is_orthogonal(self) -> bool:
return (self.moment_matrix * self.moment_matrix.T).is_diagonal()
@property
- def is_weighted_orthogonal(self):
+ def is_weighted_orthogonal(self) -> bool:
weights_matrix = sp.Matrix([self.weights] * len(self.weights))
moment_matrix_times_weights = sp.Matrix(np.multiply(self.moment_matrix, weights_matrix))
@@ -273,7 +300,7 @@ class MomentBasedLbMethod(AbstractLbMethod):
return [w for w in weights]
- def _bound_symbols_cqc(self, conserved_quantity_equations=None):
+ def _bound_symbols_cqc(self, conserved_quantity_equations: AssignmentCollection = None) -> Set[sp.Symbol]:
f = self.pre_collision_pdf_symbols
cqe = conserved_quantity_equations
@@ -282,8 +309,13 @@ class MomentBasedLbMethod(AbstractLbMethod):
return cqe.bound_symbols
- def _collision_rule_with_relaxation_matrix(self, d, additional_subexpressions=(), include_force_terms=True,
- conserved_quantity_equations=None, pre_simplification=False):
+ def _collision_rule_with_relaxation_matrix(self, d: sp.Matrix,
+ additional_subexpressions: Iterable[Assignment] = None,
+ include_force_terms: bool = True,
+ conserved_quantity_equations: AssignmentCollection = None,
+ pre_simplification: bool = False) -> LbmCollisionRule:
+ if additional_subexpressions is None:
+ additional_subexpressions = list()
f = sp.Matrix(self.pre_collision_pdf_symbols)
moment_polynomials = list(self.moments)
diff --git a/lbmpy_tests/test_cumulant_equivalences.py b/lbmpy_tests/test_cumulant_equivalences.py
index 865d9fd470d8107180a10c4bc903eefcc5a78094..11d1360a89aecc9fa4ba13a0808cb67518b9e1f3 100644
--- a/lbmpy_tests/test_cumulant_equivalences.py
+++ b/lbmpy_tests/test_cumulant_equivalences.py
@@ -1,25 +1,16 @@
import pytest
import sympy as sp
-from dataclasses import replace
-from lbmpy.enums import Method, ForceModel, Stencil
-from lbmpy.moments import (
- extract_monomials, get_default_moment_set_for_stencil, non_aliased_polynomial_raw_moments,
- exponent_tuple_sort_key)
+from lbmpy.enums import Stencil
from lbmpy.stencils import LBStencil
from lbmpy.methods.creationfunctions import create_with_monomial_cumulants
from lbmpy.maxwellian_equilibrium import get_weights
-from lbmpy.moment_transforms import (
- PdfsToMomentsByMatrixTransform, PdfsToMomentsByChimeraTransform
-)
-
@pytest.mark.parametrize('stencil', [Stencil.D2Q9, Stencil.D3Q19])
def test_zero_centering_equilibrium_equivalence(stencil):
stencil = LBStencil(stencil)
- transform_class = PdfsToMomentsByChimeraTransform
omega = sp.Symbol('omega')
r_rates = (omega,) * stencil.Q
@@ -28,8 +19,8 @@ def test_zero_centering_equilibrium_equivalence(stencil):
rho = sp.Symbol("rho")
rho_background = sp.Integer(1)
delta_rho = sp.Symbol("delta_rho")
-
- subs = { delta_rho : rho - rho_background }
+
+ subs = {delta_rho: rho - rho_background}
eqs = []
for zero_centered in [False, True]: