Commit f5ad4934 authored by Martin Bauer's avatar Martin Bauer Committed by Martin Bauer
Browse files

new lbmpy: macroscopic quantities compatation

parent dbbddc93
......@@ -5,8 +5,6 @@ from pystencils.backends.cbackend import CustomCppCode
from pystencils.ast import Block, SympyAssignment, LoopOverCoordinate, KernelFunction
from pystencils.transformations import moveConstantsBeforeLoop, resolveFieldAccesses, typingFromSympyInspection, \
typeAllEquations
from pystencils.cpu import makePythonFunction as makePythonCpuFunction
from pystencils.gpucuda import makePythonFunction as makePythonGpuFunction
from lbmpy.boundaries.createindexlist import createBoundaryIndexList
INV_DIR_SYMBOL = TypedSymbol("invDir", "int")
......@@ -75,8 +73,10 @@ class BoundaryHandling:
ast = generateBoundaryHandling(self._symbolicPdfField, idxField, self._latticeModel, boundaryFunc)
if self._target == 'cpu':
from pystencils.cpu import makePythonFunction as makePythonCpuFunction
self._boundarySweeps.append(makePythonCpuFunction(ast, {'indexField': idxField}))
elif self._target == 'gpu':
from pystencils.gpucuda import makePythonFunction as makePythonGpuFunction
self._boundarySweeps.append(makePythonGpuFunction(ast, {'indexField': idxField}))
else:
assert False
......
......@@ -8,10 +8,6 @@ from lbmpy.methods.momentbased import createSRT, createTRT, createOrthogonalMRT
import lbmpy.forcemodels as forceModels
from lbmpy.simplificationfactory import createSimplificationStrategy
from lbmpy.updatekernels import createStreamPullKernel, createPdfArray
from pystencils.cpu.kernelcreation import createKernel as createCpuKernel, addOpenMP
from pystencils.gpucuda.kernelcreation import createCUDAKernel as createGpuKernel
from pystencils.cpu import makePythonFunction as makePythonCpuFunction
from pystencils.gpucuda import makePythonFunction as makePythonGpuFunction
def _getParams(params, optParams):
......@@ -56,6 +52,7 @@ def createLatticeBoltzmannFunction(optimizationParams={}, **kwargs):
ast = createLatticeBoltzmannKernel(**params, optimizationParams=optParams)
if params['target'] == 'cpu':
from pystencils.cpu import makePythonFunction as makePythonCpuFunction, addOpenMP
if 'openMP' in optParams:
if isinstance(optParams['openMP'], bool) and optParams['openMP']:
addOpenMP(ast)
......@@ -63,6 +60,7 @@ def createLatticeBoltzmannFunction(optimizationParams={}, **kwargs):
addOpenMP(ast, numThreads=optParams['openMP'])
return makePythonCpuFunction(ast)
elif params['target'] == 'gpu':
from pystencils.gpucuda import makePythonFunction as makePythonGpuFunction
return makePythonGpuFunction(ast)
else:
return ValueError("'target' has to be either 'cpu' or 'gpu'")
......@@ -74,13 +72,15 @@ def createLatticeBoltzmannKernel(optimizationParams={}, **kwargs):
updateRule = createLatticeBoltzmannUpdateRule(**params, optimizationParams=optimizationParams)
if params['target'] == 'cpu':
from pystencils.cpu import createKernel
if 'splitGroups' in updateRule.simplificationHints:
splitGroups = updateRule.simplificationHints['splitGroups']
else:
splitGroups = ()
return createCpuKernel(updateRule.allEquations, splitGroups=splitGroups)
return createKernel(updateRule.allEquations, splitGroups=splitGroups)
elif params['target'] == 'gpu':
return createGpuKernel(updateRule.allEquations)
from pystencils.gpucuda import createCUDAKernel
return createCUDAKernel(updateRule.allEquations)
else:
return ValueError("'target' has to be either 'cpu' or 'gpu'")
......
import sympy as sp
from lbmpy.simplificationfactory import createSimplificationStrategy
from pystencils.field import Field
import pystencils.cpu as cpu
import pystencils.gpucuda as gpu
def compileMacroscopicValuesGetter(lbMethod, outputQuantities, pdfArr=None, fieldLayout='numpy', target='cpu'):
"""
Create kernel to compute macroscopic value(s) from a pdf field (e.g. density or velocity)
:param lbMethod:
:param outputQuantities:
:param pdfArr:
:param lbMethod: instance of :class:`lbmpy.methods.AbstractLbMethod`
:param outputQuantities: sequence of quantities to compute e.g. ['density', 'velocity']
:param pdfArr: optional numpy array for pdf field - used to get optimal loop structure for kernel
:param fieldLayout: layout for output field, also used for pdf field if pdfArr is not given
:param target:
:param target: 'cpu' or 'gpu'
:return: a function to compute macroscopic values:
- pdfArray, can be omitted if pdf array was already passed while compiling
- densityOut, if not None, density is written to that array
- velocityOut, if not None, velocity is written to that array
- pdfArray
- keyword arguments from name of conserved quantity (as in outputQuantities) to numpy field
"""
if not (isinstance(outputQuantities, list) or isinstance(outputQuantities, tuple)):
outputQuantities = [outputQuantities]
cqc = lbMethod.conservedQuantityComputation
unknownQuantities = [oq for oq in outputQuantities if oq not in cqc.conservedQuantities]
if unknownQuantities:
......@@ -31,7 +33,7 @@ def compileMacroscopicValuesGetter(lbMethod, outputQuantities, pdfArr=None, fiel
outputMapping = {}
for outputQuantity in outputQuantities:
numberOfElements = cqc.conservedQuantityComputation[outputQuantity]
numberOfElements = cqc.conservedQuantities[outputQuantity]
assert numberOfElements >= 1
outputField = Field.createGeneric(outputQuantity, lbMethod.dim, layout=fieldLayout,
indexDimensions=0 if numberOfElements <= 1 else 1)
......@@ -42,90 +44,83 @@ def compileMacroscopicValuesGetter(lbMethod, outputQuantities, pdfArr=None, fiel
stencil = lbMethod.stencil
pdfSymbols = [pdfField(i) for i in range(len(stencil))]
eqs = cqc.outputEquationsFromPdfs(pdfSymbols, outputMapping)
eqs = cqc.outputEquationsFromPdfs(pdfSymbols, outputMapping).allEquations
if target == 'cpu':
import pystencils.cpu as cpu
kernel = cpu.makePythonFunction(cpu.createKernel(eqs))
elif target == 'gpu':
import pystencils.gpucuda as gpu
kernel = gpu.makePythonFunction(gpu.createCUDAKernel(eqs))
else:
raise ValueError("Unknown target '%s'. Possible targets are 'cpu' and 'gpu'" % (target,))
def getter(pdfs=None, **kwargs):
def getter(pdfs, **kwargs):
if pdfArr is not None:
assert pdfs.shape == pdfArr.shape and pdfs.strides == pdfArr.strides, \
"Pdf array not matching blueprint which was used to compile" + str(pdfs.shape) + str(pdfArr.shape)
if not set(kwargs.keys()).issubset(set(outputQuantities)):
raise ValueError("You have to specify the output field for each of the following quantities: %s"
% (str(outputQuantities),))
kernel(pdfs=pdfs, **kwargs)
if not set(outputQuantities) == set(kwargs.keys()):
raise ValueError("You have to specify the output field for each of the following quantities: %s"
% (str(outputQuantities),))
kernel(pdfs=pdfs, **kwargs)
return getter
def compileMacroscopicValuesGetter(lbMethod, pdfArr=None, macroscopicFieldLayout='numpy'):
def compileMacroscopicValuesSetter(lbMethod, quantitiesToSet, pdfArr=None, fieldLayout='numpy', target='cpu'):
"""
Creates a function that computes density and/or velocity and stores it into given arrays
:param lbMethod: lattice model (to get information about stencil, force velocity shift and compressibility)
:param pdfArr: array to set can already be specified here, or later when the returned function is called
:param macroscopicFieldLayout: layout specification for Field.createGeneric
:return: a function, which has three parameters:
- pdfArray, can be omitted if pdf array was already passed while compiling
- densityOut, if not None, density is written to that array
- velocityOut, if not None, velocity is written to that array
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
:param lbMethod: instance of :class:`lbmpy.methods.AbstractLbMethod`
:param quantitiesToSet: map from conserved quantity name to fixed value or numpy array
:param pdfArr: optional numpy array for pdf field - used to get optimal loop structure for kernel
"""
if pdfArr is None:
pdfField = Field.createGeneric('pdfs', lbMethod.dim, indexDimensions=1)
else:
if pdfArr is not None:
pdfField = Field.createFromNumpyArray('pdfs', pdfArr, indexDimensions=1)
else:
pdfField = Field.createGeneric('pdfs', lbMethod.dim, indexDimensions=1, layout=fieldLayout)
rhoField = Field.createGeneric('rho', lbMethod.dim, indexDimensions=0, layout=macroscopicFieldLayout)
velField = Field.createGeneric('vel', lbMethod.dim, indexDimensions=1, layout=macroscopicFieldLayout)
fixedKernelParameters = {}
lm = lbMethod
Q = len(lm.stencil)
symPdfs = [pdfField(i) for i in range(Q)]
subexpressions, rho, velSubexpressions, u = getDensityVelocityExpressions(lm.stencil, symPdfs, lm.compressible)
valueMap = {}
atLeastOneFieldInput = False
for quantityName, value in quantitiesToSet.items():
if hasattr(value, 'shape'):
fixedKernelParameters[quantityName] = value
value = Field.createFromNumpyArray(quantityName, value)
atLeastOneFieldInput = True
valueMap[quantityName] = value
uRhs = [u_i.rhs for u_i in u]
uLhs = [u_i.lhs for u_i in u]
if hasattr(lm.forceModel, "macroscopicVelocity"):
correctedVel = lm.forceModel.macroscopicVelocity(lm, uRhs, rho.lhs)
u = [sp.Eq(a, b) for a, b in zip(uLhs, correctedVel)]
cqc = lbMethod.conservedQuantityComputation
cqEquations = cqc.equilibriumInputEquationsFromInitValues(**valueMap)
eqsRhoKernel = subexpressions + [sp.Eq(rhoField(0), rho.rhs)]
eqsVelKernel = subexpressions + [rho] + velSubexpressions + [sp.Eq(velField(i), u[i].rhs) for i in range(lm.dim)]
eqsRhoAndVelKernel = eqsVelKernel + [sp.Eq(rhoField(0), rho.lhs)]
eq = lbMethod.getEquilibrium(conservedQuantityEquations=cqEquations)
if atLeastOneFieldInput:
simplification = createSimplificationStrategy(eq)
eq = simplification(eq)
else:
eq = eq.insertSubexpressions()
stencil = lbMethod.stencil
cqc = lbMethod.conservedQuantityComputation
pdfSymbols = [pdfField(i) for i in range(len(stencil))]
eqsDensity = cqc.outputEquationsFromPdfs(pdfSymbols, {'density': rhoField(0)})
eqsVelocity = cqc.outputEquationsFromPdfs(pdfSymbols, {'velocity': [velField(i) for i in range(lbMethod.dim)]})
substitutions = {sym: pdfField(i) for i, sym in enumerate(lbMethod.postCollisionPdfSymbols)}
eq = eq.copyWithSubstitutionsApplied(substitutions).allEquations
kernelRho = makePythonFunction(createKernel(eqsRhoKernel))
kernelVel = makePythonFunction(createKernel(eqsVelKernel))
kernelRhoAndVel = makePythonFunction(createKernel(eqsRhoAndVelKernel))
if target == 'cpu':
import pystencils.cpu as cpu
kernel = cpu.makePythonFunction(cpu.createKernel(eq), argumentDict=fixedKernelParameters)
elif target == 'gpu':
import pystencils.gpucuda as gpu
kernel = gpu.makePythonFunction(gpu.createCUDAKernel(eq), argumentDict=fixedKernelParameters)
else:
raise ValueError("Unknown target '%s'. Possible targets are 'cpu' and 'gpu'" % (target,))
def getter(pdfs=None, densityOut=None, velocityOut=None):
if pdfs is not None and pdfArr is not None:
def setter(pdfs):
if pdfArr is not None:
assert pdfs.shape == pdfArr.shape and pdfs.strides == pdfArr.strides, \
"Pdf array not matching blueprint which was used to compile" + str(pdfs.shape) + str(pdfArr.shape)
if pdfs is None:
assert pdfArr is not None, "Pdf array has to be passed either when compiling or when calling."
pdfs = pdfArr
assert not (densityOut is None and velocityOut is None), \
"Specify either densityOut or velocityOut parameter or both"
kernel(pdfs=pdfs)
if densityOut is not None and velocityOut is None:
kernelRho(pdfs=pdfs, rho=densityOut)
elif densityOut is None and velocityOut is not None:
kernelVel(pdfs=pdfs, vel=velocityOut)
else:
kernelRhoAndVel(pdfs=pdfs, rho=densityOut, vel=velocityOut)
return getter
return setter
def compileAdvancedVelocitySetter(collisionRule, velocityArray, pdfArr=None):
......@@ -157,74 +152,3 @@ def compileAdvancedVelocitySetter(collisionRule, velocityArray, pdfArr=None):
kernelAst = createKernel(newCollisionRule.equations)
return makePythonFunction(kernelAst, {'vel': velocityArray})
def compileMacroscopicValuesSetter(latticeModel, density=1, velocity=0, equilibriumOrder=2, pdfArr=None):
"""
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
:param latticeModel: lattice model (to get information about stencil, force velocity shift and compressibility)
:param density: density used for equilibrium. Can either be scalar (all cells are set to same density) or array
:param velocity: velocity for equilibrium. Can either be scalar (e.g. 0, sets all cells to (0,0,0) velocity)
or a tuple with D (2 or 3) entries to set same velocity in the complete domain, or an array
specifying the velocity for each cell
:param equilibriumOrder: approximation order of equilibrium
:param pdfArr: array to set can already be specified here, or later when the returned function is called
"""
if pdfArr is not None:
pdfField = Field.createFromNumpyArray('pdfs', pdfArr, indexDimensions=1)
else:
pdfField = Field.createGeneric('pdfs', latticeModel.dim, indexDimensions=1)
noOffset = tuple([0] * latticeModel.dim)
kernelArguments = {}
if hasattr(density, 'shape'):
densityValue = Field.createFromNumpyArray('rho', density, indexDimensions=0)[noOffset]
kernelArguments['rho'] = density
else:
densityValue = density
if hasattr(velocity, 'shape'):
assert velocity.shape[-1] == latticeModel.dim, "Wrong shape of velocity array"
velocityValue = [Field.createFromNumpyArray('vel', velocity, indexDimensions=1)[noOffset](i)
for i in range(latticeModel.dim)]
kernelArguments['vel'] = velocity
else:
if not hasattr(velocity, "__len__"):
velocity = [velocity] * latticeModel.dim
velocityValue = tuple(velocity)
# force shift
if latticeModel.forceModel and hasattr(latticeModel.forceModel, "macroscopicVelocity"):
# force model knows only about one direction - use sympy to solve the shift equations to get backward
fm = latticeModel.forceModel
unshiftedVel = [sp.Symbol("v_unshifted_%d" % (i,)) for i in range(latticeModel.dim)]
shiftedVel = fm.macroscopicVelocity(latticeModel, unshiftedVel, densityValue)
velShiftEqs = [sp.Eq(a, b) for a, b in zip(velocityValue, shiftedVel)]
solveRes = sp.solve(velShiftEqs, unshiftedVel)
assert len(solveRes) == latticeModel.dim
velocityValue = [solveRes[unshiftedVel_i] for unshiftedVel_i in unshiftedVel]
eq = standardDiscreteEquilibrium(latticeModel.stencil, densityValue, velocityValue, equilibriumOrder,
c_s_sq=sp.Rational(1, 3), compressible=latticeModel.compressible)
updateEquations = [sp.Eq(pdfField(i), eq[i]) for i in range(len(latticeModel.stencil))]
f = makePythonFunction(createKernel(updateEquations), kernelArguments)
def setter(pdfs=None, **kwargs):
if pdfs is not None and pdfArr is not None:
assert pdfs.shape == pdfArr.shape and pdfs.strides == pdfArr.strides, \
"Pdf array not matching blueprint which was used to compile"
if pdfs is None:
assert pdfArr is not None, "Pdf array has to be passed either when compiling or when calling."
pdfs = pdfArr
assert pdfs.shape[-1] == len(latticeModel.stencil), "Wrong shape of pdf array"
assert len(pdfs.shape) == latticeModel.dim + 1, "Wrong shape of pdf array"
if hasattr(density, 'shape'):
assert pdfs.shape[:-1] == density.shape, "Density array shape does not match pdf array shape"
if hasattr(velocity, 'shape'):
assert pdfs.shape[:-1] == velocity.shape[:-1], "Velocity array shape does not match pdf array shape"
f(pdfs=pdfs, **kwargs) # kwargs passed here, since force model might need additional fields
return setter
from lbmpy.methods.abstractlbmmethod import AbstractLbMethod
from lbmpy.methods.abstractlbmethod import AbstractLbMethod
from lbmpy.methods.momentbased import MomentBasedLbMethod, RelaxationInfo
from lbmpy.methods.momentbased import createSRT, createTRT, createTRTWithMagicNumber, createOrthogonalMRT, \
createWithContinuousMaxwellianEqMoments, createWithDiscreteMaxwellianEqMoments
......
......@@ -5,7 +5,7 @@ from collections import namedtuple, OrderedDict, defaultdict
from lbmpy.stencils import stencilsHaveSameEntries, getStencil
from lbmpy.maxwellian_equilibrium import getMomentsOfDiscreteMaxwellianEquilibrium, \
getMomentsOfContinuousMaxwellianEquilibrium
from lbmpy.methods.abstractlbmmethod import AbstractLbMethod, LbmCollisionRule
from lbmpy.methods.abstractlbmethod import AbstractLbMethod, LbmCollisionRule
from lbmpy.methods.conservedquantitycomputation import AbstractConservedQuantityComputation, DensityVelocityComputation
from lbmpy.moments import MOMENT_SYMBOLS, momentMatrix, isShearMoment, \
isEven, gramSchmidt, getOrder, getDefaultMomentSetForStencil
......@@ -143,9 +143,9 @@ class MomentBasedLbMethod(AbstractLbMethod):
raise NotImplementedError("Shear moments seem to be not relaxed separately - "
"Can not determine their relaxation rate automatically")
def getEquilibrium(self):
def getEquilibrium(self, conservedQuantityEquations=None):
D = sp.eye(len(self._relaxationRates))
return self._getCollisionRuleWithRelaxationMatrix(D)
return self._getCollisionRuleWithRelaxationMatrix(D, conservedQuantityEquations=conservedQuantityEquations)
def getCollisionRule(self):
D = sp.diag(*self._relaxationRates)
......@@ -161,7 +161,7 @@ class MomentBasedLbMethod(AbstractLbMethod):
def conservedQuantityComputation(self):
return self._conservedQuantityComputation
def _getCollisionRuleWithRelaxationMatrix(self, D, additionalSubexpressions=[]):
def _getCollisionRuleWithRelaxationMatrix(self, D, additionalSubexpressions=[], conservedQuantityEquations=None):
f = sp.Matrix(self.preCollisionPdfSymbols)
M = self._momentMatrix
m_eq = self._equilibriumMoments
......@@ -169,12 +169,14 @@ class MomentBasedLbMethod(AbstractLbMethod):
collisionRule = f + M.inv() * D * (m_eq - M * f)
collisionEqs = [sp.Eq(lhs, rhs) for lhs, rhs in zip(self.postCollisionPdfSymbols, collisionRule)]
eqValueEqs = self._conservedQuantityComputation.equilibriumInputEquationsFromPdfs(f)
simplificationHints = eqValueEqs.simplificationHints
if conservedQuantityEquations is None:
conservedQuantityEquations = self._conservedQuantityComputation.equilibriumInputEquationsFromPdfs(f)
simplificationHints = conservedQuantityEquations.simplificationHints
simplificationHints.update(self._conservedQuantityComputation.definedSymbols())
simplificationHints['relaxationRates'] = D.atoms(sp.Symbol)
allSubexpressions = additionalSubexpressions + eqValueEqs.subexpressions + eqValueEqs.mainEquations
allSubexpressions = additionalSubexpressions + conservedQuantityEquations.allEquations
return LbmCollisionRule(self, collisionEqs, allSubexpressions,
simplificationHints)
......
Markdown is supported
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