Commit 6dec1eba authored by Martin Bauer's avatar Martin Bauer
Browse files

pystencils: Assignment instead of sympy.Eq

- Previously sympy.Eq was used to represent assignments. However Eq
  represents equality not assignment. This means that sometimes sympy
  "simplified" an equation like a = a  to True,
-> replaced sp.Eq by pystencils.Assignment everywhere
- renamed EquationCollection to AssignmentCollection
parent 181e9b17
import sympy as sp
from lbmpy.simplificationfactory import createSimplificationStrategy
from pystencils import Field, Assignment
from pystencils.astnodes import SympyAssignment
from pystencils.sympyextensions import getSymmetricPart
from pystencils import Field
from lbmpy.boundaries.boundaryhandling import BoundaryOffsetInfo, LbmWeightInfo
from pystencils.data_types import createType
from lbmpy.simplificationfactory import createSimplificationStrategy
from lbmpy.boundaries.boundaryhandling import BoundaryOffsetInfo, LbmWeightInfo
class Boundary(object):
......@@ -64,7 +64,7 @@ class NoSlip(Boundary):
def __call__(self, pdfField, directionSymbol, lbMethod, **kwargs):
neighbor = BoundaryOffsetInfo.offsetFromDir(directionSymbol, lbMethod.dim)
inverseDir = BoundaryOffsetInfo.invDir(directionSymbol)
return [sp.Eq(pdfField[neighbor](inverseDir), pdfField(directionSymbol))]
return [Assignment(pdfField[neighbor](inverseDir), pdfField(directionSymbol))]
def __hash__(self):
# All boundaries of these class behave equal -> should also be equal (as long as name is equal)
......@@ -126,7 +126,7 @@ class UBB(Boundary):
if self._adaptVelocityToForce:
cqc = lbMethod.conservedQuantityComputation
shiftedVelEqs = cqc.equilibriumInputEquationsFromInitValues(velocity=velocity)
velocity = [eq.rhs for eq in shiftedVelEqs.extract(cqc.firstOrderMomentSymbols).mainEquations]
velocity = [eq.rhs for eq in shiftedVelEqs.extract(cqc.firstOrderMomentSymbols).mainAssignments]
c_s_sq = sp.Rational(1, 3)
weightOfDirection = LbmWeightInfo.weightOfDirection
......@@ -142,12 +142,12 @@ class UBB(Boundary):
densityEquations = cqc.outputEquationsFromPdfs(pdfFieldAccesses, {'density': densitySymbol})
densitySymbol = lbMethod.conservedQuantityComputation.definedSymbols()['density']
result = densityEquations.allEquations
result += [sp.Eq(pdfField[neighbor](inverseDir),
pdfField(direction) - velTerm * densitySymbol)]
result += [Assignment(pdfField[neighbor](inverseDir),
pdfField(direction) - velTerm * densitySymbol)]
return result
else:
return [sp.Eq(pdfField[neighbor](inverseDir),
pdfField(direction) - velTerm)]
return [Assignment(pdfField[neighbor](inverseDir),
pdfField(direction) - velTerm)]
class FixedDensity(Boundary):
......@@ -159,16 +159,16 @@ class FixedDensity(Boundary):
def __call__(self, pdfField, directionSymbol, lbMethod, **kwargs):
"""Boundary condition that fixes the density/pressure at the obstacle"""
def removeAsymmetricPartOfMainEquations(eqColl, dofs):
newMainEquations = [sp.Eq(e.lhs, getSymmetricPart(e.rhs, dofs)) for e in eqColl.mainEquations]
return eqColl.copy(newMainEquations)
def removeAsymmetricPartOfmainAssignments(eqColl, dofs):
newmainAssignments = [Assignment(e.lhs, getSymmetricPart(e.rhs, dofs)) for e in eqColl.mainAssignments]
return eqColl.copy(newmainAssignments)
neighbor = BoundaryOffsetInfo.offsetFromDir(directionSymbol, lbMethod.dim)
inverseDir = BoundaryOffsetInfo.invDir(directionSymbol)
cqc = lbMethod.conservedQuantityComputation
velocity = cqc.definedSymbols()['velocity']
symmetricEq = removeAsymmetricPartOfMainEquations(lbMethod.getEquilibrium(), dofs=velocity)
symmetricEq = removeAsymmetricPartOfmainAssignments(lbMethod.getEquilibrium(), dofs=velocity)
substitutions = {sym: pdfField(i) for i, sym in enumerate(lbMethod.preCollisionPdfSymbols)}
symmetricEq = symmetricEq.copyWithSubstitutionsApplied(substitutions)
......@@ -178,15 +178,15 @@ class FixedDensity(Boundary):
densitySymbol = cqc.definedSymbols()['density']
density = self._density
densityEq = cqc.equilibriumInputEquationsFromInitValues(density=density).insertSubexpressions().mainEquations[0]
densityEq = cqc.equilibriumInputEquationsFromInitValues(density=density).insertSubexpressions().mainAssignments[0]
assert densityEq.lhs == densitySymbol
transformedDensity = densityEq.rhs
conditions = [(eq_i.rhs, sp.Equality(directionSymbol, i))
for i, eq_i in enumerate(symmetricEq.mainEquations)] + [(0, True)]
for i, eq_i in enumerate(symmetricEq.mainAssignments)] + [(0, True)]
eq_component = sp.Piecewise(*conditions)
subExprs = [sp.Eq(eq.lhs, transformedDensity if eq.lhs == densitySymbol else eq.rhs)
subExprs = [Assignment(eq.lhs, transformedDensity if eq.lhs == densitySymbol else eq.rhs)
for eq in symmetricEq.subexpressions]
return subExprs + [SympyAssignment(pdfField[neighbor](inverseDir),
......@@ -197,8 +197,8 @@ class NeumannByCopy(Boundary):
def __call__(self, pdfField, directionSymbol, lbMethod, **kwargs):
neighbor = BoundaryOffsetInfo.offsetFromDir(directionSymbol, lbMethod.dim)
inverseDir = BoundaryOffsetInfo.invDir(directionSymbol)
return [sp.Eq(pdfField[neighbor](inverseDir), pdfField(inverseDir)),
sp.Eq(pdfField[neighbor](directionSymbol), pdfField(directionSymbol))]
return [Assignment(pdfField[neighbor](inverseDir), pdfField(inverseDir)),
Assignment(pdfField[neighbor](directionSymbol), pdfField(directionSymbol))]
def __hash__(self):
# All boundaries of these class behave equal -> should also be equal
......@@ -212,8 +212,8 @@ class StreamInZero(Boundary):
def __call__(self, pdfField, directionSymbol, lbMethod, **kwargs):
neighbor = BoundaryOffsetInfo.offsetFromDir(directionSymbol, lbMethod.dim)
inverseDir = BoundaryOffsetInfo.invDir(directionSymbol)
return [sp.Eq(pdfField[neighbor](inverseDir), 0),
sp.Eq(pdfField[neighbor](directionSymbol), 0)]
return [Assignment(pdfField[neighbor](inverseDir), 0),
Assignment(pdfField[neighbor](directionSymbol), 0)]
def __hash__(self):
# All boundaries of these class behave equal -> should also be equal
......
import numpy as np
import sympy as sp
from lbmpy.stencils import inverseDirection
from pystencils import TypedSymbol, createIndexedKernel
from pystencils import TypedSymbol, createIndexedKernel, Assignment
from pystencils.backends.cbackend import CustomCppCode
from pystencils.boundaries import BoundaryHandling
from pystencils.boundaries.boundaryhandling import BoundaryOffsetInfo
from lbmpy.stencils import inverseDirection
class LatticeBoltzmannBoundaryHandling(BoundaryHandling):
......@@ -99,6 +99,6 @@ def createLatticeBoltzmannBoundaryKernel(pdfField, indexField, lbMethod, boundar
elements = [BoundaryOffsetInfo(lbMethod.stencil), LbmWeightInfo(lbMethod)]
indexArrDtype = indexField.dtype.numpyDtype
dirSymbol = TypedSymbol("dir", indexArrDtype.fields['dir'][0])
elements += [sp.Eq(dirSymbol, indexField[0]('dir'))]
elements += [Assignment(dirSymbol, indexField[0]('dir'))]
elements += boundaryFunctor(pdfField=pdfField, directionSymbol=dirSymbol, lbMethod=lbMethod, indexField=indexField)
return createIndexedKernel(elements, [indexField], target=target, cpuOpenMP=openMP)
......@@ -77,7 +77,7 @@ class CeMoment(sp.Symbol):
class LbMethodEqMoments:
def __init__(self, lbMethod):
self._eq = tuple(e.rhs for e in lbMethod.getEquilibrium().mainEquations)
self._eq = tuple(e.rhs for e in lbMethod.getEquilibrium().mainAssignments)
self._momentCache = dict()
self._postCollisionMomentCache = dict()
self._stencil = lbMethod.stencil
......
......@@ -152,8 +152,12 @@ For example, to modify the AST one can run::
import sympy as sp
from copy import copy
from lbmpy.turbulence_models import addSmagorinskyModel
from pystencils.cache import diskcacheNoFallback
from pystencils.data_types import collateTypes
from pystencils.assignment_collection.assignment_collection import AssignmentCollection
from pystencils.field import getLayoutOfArray, Field
from pystencils import createKernel, Assignment
from lbmpy.turbulence_models import addSmagorinskyModel
from lbmpy.methods import createSRT, createTRT, createOrthogonalMRT, createKBCTypeTRT, \
createRawMRT, createThreeRelaxationRateMRT
from lbmpy.methods.entropic import addIterativeEntropyCondition, addEntropyCondition
......@@ -164,10 +168,6 @@ import lbmpy.forcemodels as forcemodels
from lbmpy.simplificationfactory import createSimplificationStrategy
from lbmpy.updatekernels import StreamPullTwoFieldsAccessor, PeriodicTwoFieldsAccessor, CollideOnlyInplaceAccessor, \
createLBMKernel, createStreamPullWithOutputKernel
from pystencils.data_types import collateTypes
from pystencils.equationcollection.equationcollection import EquationCollection
from pystencils.field import getLayoutOfArray, Field
from pystencils import createKernel
def createLatticeBoltzmannFunction(ast=None, optimizationParams={}, **kwargs):
......@@ -275,10 +275,10 @@ def createLatticeBoltzmannCollisionRule(lbMethod=None, optimizationParams={}, **
cqc = lbMethod.conservedQuantityComputation
if params['velocityInput'] is not None:
eqs = [sp.Eq(cqc.zerothOrderMomentSymbol, sum(lbMethod.preCollisionPdfSymbols))]
eqs = [Assignment(cqc.zerothOrderMomentSymbol, sum(lbMethod.preCollisionPdfSymbols))]
velocityField = params['velocityInput']
eqs += [sp.Eq(uSym, velocityField(i)) for i, uSym in enumerate(cqc.firstOrderMomentSymbols)]
eqs = EquationCollection(eqs, [])
eqs += [Assignment(uSym, velocityField(i)) for i, uSym in enumerate(cqc.firstOrderMomentSymbols)]
eqs = AssignmentCollection(eqs, [])
collisionRule = lbMethod.getCollisionRule(conservedQuantityEquations=eqs)
else:
collisionRule = lbMethod.getCollisionRule()
......
......@@ -29,7 +29,7 @@ def createLbmSplitGroups(lbmCollisionEqs):
if preCollisionSymbols.intersection(lbmCollisionEqs.getDependentSymbols([e.lhs]))}
otherWrittenFields = []
for eq in lbmCollisionEqs.mainEquations:
for eq in lbmCollisionEqs.mainAssignments:
if eq.lhs not in postCollisionSymbols and isinstance(eq.lhs, Field.Access):
otherWrittenFields.append(eq.lhs)
if eq.lhs not in nonCenterPostCollisionSymbols:
......@@ -44,7 +44,7 @@ def createLbmSplitGroups(lbmCollisionEqs):
directionGroups = defaultdict(list)
dim = len(stencil[0])
for direction, eq in zip(stencil, lbmCollisionEqs.mainEquations):
for direction, eq in zip(stencil, lbmCollisionEqs.mainAssignments):
if direction == tuple([0]*dim):
splitGroups[0].append(eq.lhs)
continue
......
import abc
import sympy as sp
from collections import namedtuple
from pystencils.equationcollection import EquationCollection
from pystencils.assignment_collection import AssignmentCollection
RelaxationInfo = namedtuple('RelaxationInfo', ['equilibriumValue', 'relaxationRate'])
class LbmCollisionRule(EquationCollection):
class LbmCollisionRule(AssignmentCollection):
def __init__(self, lbmMethod, *args, **kwargs):
super(LbmCollisionRule, self).__init__(*args, **kwargs)
self.method = lbmMethod
......
import abc
from collections import OrderedDict
import sympy as sp
from pystencils.equationcollection import EquationCollection
from pystencils.field import Field
from collections import OrderedDict
from pystencils.assignment_collection import AssignmentCollection
from pystencils.field import Field, Assignment
class AbstractConservedQuantityComputation(abc.ABCMeta('ABC', (object,), {})):
......@@ -23,7 +22,8 @@ class AbstractConservedQuantityComputation(abc.ABCMeta('ABC', (object,), {})):
"""
@abc.abstractproperty
@property
@abc.abstractmethod
def conservedQuantities(self):
"""
Dict, mapping names (symbol) to dimensionality (int)
......@@ -37,7 +37,8 @@ class AbstractConservedQuantityComputation(abc.ABCMeta('ABC', (object,), {})):
Returns a dict, mapping names of conserved quantities to their symbols
"""
@abc.abstractproperty
@property
@abc.abstractmethod
def defaultValues(self):
"""
Returns a dict of symbol to default value, where "default" means that
......@@ -129,84 +130,84 @@ class DensityVelocityComputation(AbstractConservedQuantityComputation):
def equilibriumInputEquationsFromPdfs(self, pdfs):
dim = len(self._stencil[0])
eqColl = getEquationsForZerothAndFirstOrderMoment(self._stencil, pdfs, self._symbolOrder0,
eq_coll = getEquationsForZerothAndFirstOrderMoment(self._stencil, pdfs, self._symbolOrder0,
self._symbolsOrder1[:dim])
if self._compressible:
eqColl = divideFirstOrderMomentsByRho(eqColl, dim)
eq_coll = divideFirstOrderMomentsByRho(eq_coll, dim)
eqColl = applyForceModelShift('equilibriumVelocityShift', dim, eqColl, self._forceModel, self._compressible)
return eqColl
eq_coll = applyForceModelShift('equilibriumVelocityShift', dim, eq_coll, self._forceModel, self._compressible)
return eq_coll
def equilibriumInputEquationsFromInitValues(self, density=1, velocity=(0, 0, 0)):
dim = len(self._stencil[0])
zerothOrderMoment = density
firstOrderMoments = velocity[:dim]
velOffset = [0] * dim
first_order_moments = velocity[:dim]
vel_offset = [0] * dim
if self._compressible:
if self._forceModel and hasattr(self._forceModel, 'macroscopicVelocityShift'):
velOffset = self._forceModel.macroscopicVelocityShift(zerothOrderMoment)
vel_offset = self._forceModel.macroscopicVelocityShift(zerothOrderMoment)
else:
if self._forceModel and hasattr(self._forceModel, 'macroscopicVelocityShift'):
velOffset = self._forceModel.macroscopicVelocityShift(sp.Rational(1, 1))
vel_offset = self._forceModel.macroscopicVelocityShift(sp.Rational(1, 1))
zerothOrderMoment -= sp.Rational(1, 1)
eqs = [sp.Eq(self._symbolOrder0, zerothOrderMoment)]
eqs = [Assignment(self._symbolOrder0, zerothOrderMoment)]
firstOrderMoments = [a - b for a, b in zip(firstOrderMoments, velOffset)]
eqs += [sp.Eq(l, r) for l, r in zip(self._symbolsOrder1, firstOrderMoments)]
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)]
return EquationCollection(eqs, [])
return AssignmentCollection(eqs, [])
def outputEquationsFromPdfs(self, pdfs, outputQuantityNamesToSymbols):
dim = len(self._stencil[0])
eqColl = getEquationsForZerothAndFirstOrderMoment(self._stencil, pdfs, self._symbolOrder0, self._symbolsOrder1)
ac = getEquationsForZerothAndFirstOrderMoment(self._stencil, pdfs, self._symbolOrder0, self._symbolsOrder1)
if self._compressible:
eqColl = divideFirstOrderMomentsByRho(eqColl, dim)
ac = divideFirstOrderMomentsByRho(ac, dim)
else:
eqColl = addDensityOffset(eqColl)
ac = addDensityOffset(ac)
eqColl = applyForceModelShift('macroscopicVelocityShift', dim, eqColl, self._forceModel, self._compressible)
ac = applyForceModelShift('macroscopicVelocityShift', dim, ac, self._forceModel, self._compressible)
mainEquations = []
eqs = OrderedDict([(eq.lhs, eq.rhs) for eq in eqColl.allEquations])
main_assignments = []
eqs = OrderedDict([(eq.lhs, eq.rhs) for eq in ac.allEquations])
if 'density' in outputQuantityNamesToSymbols:
densityOutputSymbol = outputQuantityNamesToSymbols['density']
if isinstance(densityOutputSymbol, Field):
densityOutputSymbol = densityOutputSymbol()
if densityOutputSymbol != self._symbolOrder0:
mainEquations.append(sp.Eq(densityOutputSymbol, self._symbolOrder0))
density_output_symbol = outputQuantityNamesToSymbols['density']
if isinstance(density_output_symbol, Field):
density_output_symbol = density_output_symbol()
if density_output_symbol != self._symbolOrder0:
main_assignments.append(Assignment(density_output_symbol, self._symbolOrder0))
else:
mainEquations.append(sp.Eq(self._symbolOrder0, eqs[self._symbolOrder0]))
main_assignments.append(Assignment(self._symbolOrder0, eqs[self._symbolOrder0]))
del eqs[self._symbolOrder0]
if 'velocity' in outputQuantityNamesToSymbols:
velOutputSymbols = outputQuantityNamesToSymbols['velocity']
if isinstance(velOutputSymbols, Field):
field = velOutputSymbols
velOutputSymbols = [field(i) for i in range(len(self._symbolsOrder1))]
if tuple(velOutputSymbols) != tuple(self._symbolsOrder1):
mainEquations += [sp.Eq(a, b) for a, b in zip(velOutputSymbols, self._symbolsOrder1)]
vel_output_symbols = outputQuantityNamesToSymbols['velocity']
if isinstance(vel_output_symbols, Field):
field = vel_output_symbols
vel_output_symbols = [field(i) for i in range(len(self._symbolsOrder1))]
if tuple(vel_output_symbols) != tuple(self._symbolsOrder1):
main_assignments += [Assignment(a, b) for a, b in zip(vel_output_symbols, self._symbolsOrder1)]
else:
for u_i in self._symbolsOrder1:
mainEquations.append(sp.Eq(u_i, eqs[u_i]))
main_assignments.append(Assignment(u_i, eqs[u_i]))
del eqs[u_i]
if 'momentumDensity' in outputQuantityNamesToSymbols:
# get zeroth and first moments again - force-shift them if necessary
# and add their values directly to the main equations assuming that subexpressions are already in
# main equation collection
# Is not optimal when velocity and momentumDensity are calculated together, but this is usually not the case
momentumDensityOutputSymbols = outputQuantityNamesToSymbols['momentumDensity']
momDensityEqColl = getEquationsForZerothAndFirstOrderMoment(self._stencil, pdfs,
momentum_density_output_symbols = outputQuantityNamesToSymbols['momentumDensity']
mom_density_eq_coll = getEquationsForZerothAndFirstOrderMoment(self._stencil, pdfs,
self._symbolOrder0, self._symbolsOrder1)
momDensityEqColl = applyForceModelShift('macroscopicVelocityShift', dim, momDensityEqColl,
mom_density_eq_coll = applyForceModelShift('macroscopicVelocityShift', dim, mom_density_eq_coll,
self._forceModel, self._compressible)
for sym, val in zip(momentumDensityOutputSymbols, momDensityEqColl.mainEquations[1:]):
mainEquations.append(sp.Eq(sym, val.rhs))
for sym, val in zip(momentum_density_output_symbols, mom_density_eq_coll.mainAssignments[1:]):
main_assignments.append(Assignment(sym, val.rhs))
eqColl = eqColl.copy(mainEquations, [sp.Eq(a, b) for a, b in eqs.items()])
return eqColl.newWithoutUnusedSubexpressions()
ac = ac.copy(main_assignments, [Assignment(a, b) for a, b in eqs.items()])
return ac.newWithoutUnusedSubexpressions()
def __repr__(self):
return "ConservedValueComputation for %s" % (", " .join(self.conservedQuantities.keys()),)
......@@ -231,7 +232,7 @@ def getEquationsForZerothAndFirstOrderMoment(stencil, symbolicPdfs, symbolicZero
:param symbolicZerothMoment: called :math:`\rho` above
:param symbolicFirstMoments: called :math:`u` above
"""
def filterOutPlusTerms(expr):
def filter_out_plus_terms(expr):
result = 0
for term in expr.args:
if not type(term) is sp.Mul:
......@@ -241,34 +242,34 @@ def getEquationsForZerothAndFirstOrderMoment(stencil, symbolicPdfs, symbolicZero
dim = len(stencil[0])
subexpressions = []
pdfSum = sum(symbolicPdfs)
pdf_sum = sum(symbolicPdfs)
u = [0] * dim
for f, offset in zip(symbolicPdfs, stencil):
for i in range(dim):
u[i] += f * int(offset[i])
plusTerms = [set(filterOutPlusTerms(u_i).args) for u_i in u]
plus_terms = [set(filter_out_plus_terms(u_i).args) for u_i in u]
for i in range(dim):
rhs = plusTerms[i]
rhs = plus_terms[i]
for j in range(i):
rhs -= plusTerms[j]
eq = sp.Eq(sp.Symbol("vel%dTerm" % (i,)), sum(rhs))
rhs -= plus_terms[j]
eq = Assignment(sp.Symbol("vel%dTerm" % (i,)), sum(rhs))
subexpressions.append(eq)
for subexpression in subexpressions:
pdfSum = pdfSum.subs(subexpression.rhs, subexpression.lhs)
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)
equations = []
equations += [sp.Eq(symbolicZerothMoment, pdfSum)]
equations += [sp.Eq(u_i_sym, u_i) for u_i_sym, u_i in zip(symbolicFirstMoments, u)]
equations += [Assignment(symbolicZerothMoment, pdf_sum)]
equations += [Assignment(u_i_sym, u_i) for u_i_sym, u_i in zip(symbolicFirstMoments, u)]
return EquationCollection(equations, subexpressions)
return AssignmentCollection(equations, subexpressions)
def divideFirstOrderMomentsByRho(equationCollection, dim):
def divideFirstOrderMomentsByRho(assignment_collection, dim):
"""
Assumes that the equations of the passed equation collection are the following
- rho = f_0 + f_1 + ...
......@@ -277,38 +278,39 @@ def divideFirstOrderMomentsByRho(equationCollection, dim):
Returns a new equation collection where the u terms (first order moments) are divided by rho.
The dim parameter specifies the number of first order moments. All subsequent equations are just copied over.
"""
oldEqs = equationCollection.mainEquations
oldEqs = assignment_collection.mainAssignments
rho = oldEqs[0].lhs
newFirstOrderMomentEq = [sp.Eq(eq.lhs, eq.rhs / rho) for eq in oldEqs[1:dim+1]]
newEqs = [oldEqs[0]] + newFirstOrderMomentEq + oldEqs[dim+1:]
return equationCollection.copy(newEqs)
new_first_order_moment_eq = [Assignment(eq.lhs, eq.rhs / rho) for eq in oldEqs[1:dim+1]]
new_eqs = [oldEqs[0]] + new_first_order_moment_eq + oldEqs[dim+1:]
return assignment_collection.copy(new_eqs)
def addDensityOffset(equationCollection, offset=sp.Rational(1, 1)):
def addDensityOffset(assignment_collection, offset=sp.Rational(1, 1)):
"""
Assumes that first equation is the density (zeroth moment). Changes the density equations by adding offset to it.
"""
oldEqs = equationCollection.mainEquations
newDensity = sp.Eq(oldEqs[0].lhs, oldEqs[0].rhs + offset)
return equationCollection.copy([newDensity] + oldEqs[1:])
oldEqs = assignment_collection.mainAssignments
newDensity = Assignment(oldEqs[0].lhs, oldEqs[0].rhs + offset)
return assignment_collection.copy([newDensity] + oldEqs[1:])
def applyForceModelShift(shiftMemberName, dim, equationCollection, forceModel, compressible, reverse=False):
def applyForceModelShift(shiftMemberName, dim, assignment_collection, forceModel, compressible, reverse=False):
"""
Modifies the first order moment equations in equationCollection according to the force model shift.
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 shiftMemberName. The equations 1: dim+1 of the passed
equation collection are assumed to be the velocity equations.
"""
if forceModel is not None and hasattr(forceModel, shiftMemberName):
oldEqs = equationCollection.mainEquations
density = oldEqs[0].lhs if compressible else sp.Rational(1, 1)
oldVelEqs = oldEqs[1:dim + 1]
shiftFunc = getattr(forceModel, shiftMemberName)
velOffsets = shiftFunc(density)
old_eqs = assignment_collection.mainAssignments
density = old_eqs[0].lhs if compressible else sp.Rational(1, 1)
old_vel_eqs = old_eqs[1:dim + 1]
shift_func = getattr(forceModel, shiftMemberName)
vel_offsets = shift_func(density)
if reverse:
velOffsets = [-v for v in velOffsets]
shiftedVelocityEqs = [sp.Eq(oldEq.lhs, oldEq.rhs + offset) for oldEq, offset in zip(oldVelEqs, velOffsets)]
newEqs = [oldEqs[0]] + shiftedVelocityEqs + oldEqs[dim + 1:]
return equationCollection.copy(newEqs)
vel_offsets = [-v for v in vel_offsets]
shifted_velocity_eqs = [Assignment(oldEq.lhs, oldEq.rhs + offset)
for oldEq, 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 equationCollection
return assignment_collection
import sympy as sp
from collections import OrderedDict
from pystencils import Assignment
from pystencils.sympyextensions import fastSubs, replaceAdditive
from lbmpy.methods.abstractlbmethod import AbstractLbMethod, LbmCollisionRule, RelaxationInfo
from lbmpy.methods.conservedquantitycomputation import AbstractConservedQuantityComputation
from lbmpy.moments import MOMENT_SYMBOLS, extractMonomials, momentMatrix, monomialToPolynomialTransformationMatrix
from lbmpy.cumulants import cumulantAsFunctionOfRawMoments, rawMomentAsFunctionOfCumulants
from pystencils.sympyextensions import fastSubs, replaceAdditive
class CumulantBasedLbMethod(AbstractLbMethod):
......@@ -100,7 +101,7 @@ class CumulantBasedLbMethod(AbstractLbMethod):
def getEquilibriumTerms(self):
equilibrium = self.getEquilibrium()
return sp.Matrix([eq.rhs for eq in equilibrium.mainEquations])
return sp.Matrix([eq.rhs for eq in equilibrium.mainAssignments])
def getCollisionRule(self, conservedQuantityEquations=None, momentSubexpressions=False,
preCollisionSubexpressions=True, postCollisionSubexpressions=False):
......@@ -110,14 +111,16 @@ class CumulantBasedLbMethod(AbstractLbMethod):
def _computeWeights(self):
replacements = self._conservedQuantityComputation.defaultValues
eqColl = self.getEquilibrium().copyWithSubstitutionsApplied(replacements).insertSubexpressions()
newMainEqs = [sp.Eq(e.lhs,
replaceAdditive(e.rhs, 1, sum(self.preCollisionPdfSymbols), requiredMatchReplacement=1.0))
for e in eqColl.mainEquations]
eq = self.getEquilibrium()
eqColl = eq.copyWithSubstitutionsApplied(replacements, substituteOnLhs=False).insertSubexpressions()
newMainEqs = [Assignment(e.lhs,
replaceAdditive(e.rhs, 1, sum(self.preCollisionPdfSymbols),
requiredMatchReplacement=1.0))
for e in eqColl.mainAssignments]
eqColl = eqColl.copy(newMainEqs)
weights = []
for eq in eqColl.mainEquations:
for eq in eqColl.mainAssignments:
value = eq.rhs.expand()
assert len(value.atoms(sp.Symbol)) == 0, "Failed to compute weights"
weights.append(value)
......@@ -133,9 +136,9 @@ class CumulantBasedLbMethod(AbstractLbMethod):
def substituteConservedQuantities(expressions, cqe):
cqe = cqe.insertSubexpressions()
substitutionDict = {eq.rhs: eq.lhs for eq in cqe.mainEquations}
density = cqe.mainEquations[0].lhs
substitutionDict.update({density * eq.rhs: density * eq.lhs for eq in cqe.mainEquations[1:]})
substitutionDict = {eq.rhs: eq.lhs for eq in cqe.mainAssignments}
density = cqe.mainAssignments[0].lhs
substitutionDict.update({density * eq.rhs: density * eq.lhs for eq in cqe.mainAssignments[1:]})
return [fastSubs(e, substitutionDict) for e in expressions]
f = self.preCollisionPdfSymbols
......@@ -161,7 +164,7 @@ class CumulantBasedLbMethod(AbstractLbMethod):
moments = substituteConservedQuantities(moments, conservedQuantityEquations)
if momentSubexpressions:
symbols = [tupleToSymbol(t, "m") for t in higherOrderIndices]
subexpressions += [sp.Eq(sym, moment) for sym, moment in zip(symbols, moments[numLowerOrderIndices:])]
subexpressions += [Assignment(sym, moment) for sym, moment in zip(symbols, moments[numLowerOrderIndices:])]
moments = moments[:numLowerOrderIndices] + symbols
# 3) Transform moments to monomial cumulants
......@@ -170,7 +173,7 @@ class CumulantBasedLbMethod(AbstractLbMethod):
if preCollisionSubexpressions:
symbols = [tupleToSymbol(t, "preC") for t in higherOrderIndices]
subexpressions += [sp.Eq(sym, c)
subexpressions += [Assignment(sym, c)
for sym, c in zip(symbols, monomialCumulants[numLowerOrderIndices:])]
monomialCumulants = monomialCumulants[:numLowerOrderIndices] + symbols
......@@ -183,7 +186,7 @@ class CumulantBasedLbMethod(AbstractLbMethod):
if postCollisionSubexpressions:
symbols = [tupleToSymbol(t, "postC") for t in higherOrderIndices]
subexpressions += [sp.Eq(sym, c)
subexpressions += [Assignment(sym, c)
for sym, c in zip(symbols, relaxedMonomialCumulants[numLowerOrderIndices:])]
relaxedMonomialCumulants = relaxedMonomialCumulants[:numLowerOrderIndices] + symbols
......@@ -191,20 +194,20 @@ class CumulantBasedLbMethod(AbstractLbMethod):
cumulantDict = {idx: value for idx, value in zip(indices, relaxedMonomialCumulants)}
collidedMoments = [rawMomentAsFunctionOfCumulants(idx, cumulantDict) for idx in indices]
result = momentTransformationMatrix.inv() * sp.Matrix(collidedMoments)
mainEquations = [sp.Eq(sym, val) for sym, val in zip(self.postCollisionPdfSymbols, result)]
mainAssignments = [Assignment(sym, val) for sym, val in zip(self.postCollisionPdfSymbols, result)]
# 6) Add forcing terms