def generateDirichletDomainCode(predefined, tensor, tempVars=True):
# predefined={}
keys = tensor.keys()
expressions = [tensor[i] for i in keys]
preamble, results = codegen.generateCode(predefined, expressions, tempVars=tempVars)
result = Variable('int', 'domainId')
return [assign(result, results[0])]
def generateCode(predefined, testFunctions, tensorMap, tempVars=True):
# build list of all expressions to compile
expressions = []
for phi in testFunctions:
if (phi, ) not in tensorMap:
continue
tensor = tensorMap[(phi, )]
keys = tensor.keys()
expressions += [tensor[i] for i in keys]
# compile all expressions at once
preamble, results = codegen.generateCode(predefined,
expressions,
tempVars=tempVars)
# extract generated code for expressions and build values
values = []
for phi in testFunctions:
value = tensors.fill(phi.ufl_shape, makeExpression(0))
if (phi, ) in tensorMap:
tensor = tensorMap[(phi, )]
keys = tensor.keys()
for i, r in zip(keys, results[:len(keys)]):
value = tensors.setItem(value, i, r)
results = results[len(keys):]
values += [
tensors.reformat(lambda row: InitializerList(*row), phi.ufl_shape,
value)
]
return preamble, values
def generateDirichletCode(predefined, tensor, tempVars=True):
keys = tensor.keys()
expressions = [tensor[i] for i in keys]
preamble, results = codegen.generateCode(predefined,
expressions,
tempVars=tempVars)
result = Variable('auto', 'result')
return preamble + [assign(result[i], r) for i, r in zip(keys, results)]
def generateLinearizedCode(predefined,
testFunctions,
trialFunctionMap,
tensorMap,
tempVars=True):
"""generate code for a bilinear form
Args:
predefined: list of predefined arguments or coefficients
testFunctions: list of arguments to interpret as test functions
trialFunctionMap: map of variable to list of arguments to interpret as trial functions
tensorMap: map of expression tensors of shape (testFunction x trialFunction)
tempVars: introduce temporary variables during code generation
"""
# build list of all expressions to compile
expressions = []
for var, trialFunctions in trialFunctionMap.items():
for phi in testFunctions:
for psi in trialFunctions:
if (phi, psi) not in tensorMap:
continue
tensor = tensorMap[(phi, psi)]
keys = tensor.keys()
expressions += [tensor[i] for i in keys]
# compile all expressions at once
preamble, results = codegen.generateCode(predefined,
expressions,
tempVars=tempVars)
# extract generated code for expressions and build values
values = {}
for var, trialFunctions in trialFunctionMap.items():
values[var] = []
for phi in testFunctions:
value = tensors.fill(phi.ufl_shape, None)
for idx in range(len(trialFunctions)):
psi = trialFunctions[idx]
if (phi, psi) in tensorMap:
tensor = tensorMap[(phi, psi)]
keys = tensor.keys()
for ij, r in zip(keys, results[:len(keys)]):
if isinstance(tensor[ij], Zero):
continue
i = ij[:len(phi.ufl_shape)]
j = ij[len(phi.ufl_shape):]
if isinstance(tensor[ij], IntValue) and int(
tensor[ij]) == 1:
r = var[idx][j]
else:
r = r * var[idx][j]
s = tensors.getItem(value, i)
s = r if s is None else s + r
value = tensors.setItem(value, i, s)
results = results[len(keys):]
value = tensors.apply(
lambda v: makeExpression(0)
if v is None else v, phi.ufl_shape, value)
values[var] += [
tensors.reformat(lambda row: InitializerList(*row),
phi.ufl_shape, value)
]
return preamble, values