def count_flops(n):
mesh = Mesh(VectorElement('CG', interval, 1))
tfs = FunctionSpace(mesh, TensorElement('DG', interval, 1, shape=(n, n)))
vfs = FunctionSpace(mesh, VectorElement('DG', interval, 1, dim=n))
ensemble_f = Coefficient(vfs)
ensemble2_f = Coefficient(vfs)
phi = TestFunction(tfs)
i, j = indices(2)
nc = 42 # magic number
L = ((IndexSum(
IndexSum(
Product(nc * phi[i, j], Product(ensemble_f[i], ensemble_f[i])),
MultiIndex((i, ))), MultiIndex((j, ))) * dx) +
(IndexSum(
IndexSum(
Product(nc * phi[i, j], Product(
ensemble2_f[j], ensemble2_f[j])), MultiIndex(
(i, ))), MultiIndex((j, ))) * dx) -
(IndexSum(
IndexSum(
2 * nc *
Product(phi[i, j], Product(ensemble_f[i], ensemble2_f[j])),
MultiIndex((i, ))), MultiIndex((j, ))) * dx))
kernel, = compile_form(L, parameters=dict(mode='spectral'))
return EstimateFlops().visit(kernel.ast)
def split_sum(input_, output):
if isinstance(input_, Sum):
for operand in input_.ufl_operands:
split_sum(operand, output)
elif isinstance(input_, IndexSum):
assert len(input_.ufl_operands) == 2
summand, indices = input_.ufl_operands
assert isinstance(indices, MultiIndex)
output_0 = list()
split_sum(summand, output_0)
output.extend([IndexSum(summand_0, indices) for summand_0 in output_0])
else:
output.append(input_)
def index_sum(self, o, Ap, i):
return IndexSum(Ap, i)