Skip to content
Snippets Groups Projects
Commit 4b3e0ae1 authored by Markus Holzer's avatar Markus Holzer Committed by Helen Schottenhamml
Browse files

Simplify equilibrium terms

parent 1c008ce4
1 merge request!134Simplify equilibrium terms
......@@ -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
......
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
......
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
......
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)
......
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]:
......
0% or .
You are about to add 0 people to the discussion. Proceed with caution.
Finish editing this message first!
Please register or to comment