Further improvements of External Operators

[jorgensd]

After discussions with @a-latyshev we have decided that some nice things to add to this library would be

• Support of non-quadrature elements. This relates to work with @croci where we would like to call ufl.grad(external_operator). We have some work in progress on this (@a-latyshev will publish branch). We need a good test-case for this.
  • Allow basic UFL operation to be applied to external operators.
• Batch eval/coefficient packing. This is to deal with the fact that packing at quadrature spaces for high order quadrature is horrible for memory. First step would be to take extract integration domains from the current forms, and create new batched dolfinx.fem.Form objects.
  • Alternative would be to use the symbolic dolfinx.fem.create_form and dolfinx.fem.compile_form, but that is rather verbose.
• Better handling of codim meshes. This requires some upgrades to dolfinx.fem.Expression, which should take in entity_maps and ensure that packing is done correctly.
• Enable nested external operators. @jorgensd will provide a use case. Fixed in Support for nested external operators #40
• Support for mixed function spaces. Currently this doesn't fully work due to how extract blocks works. Need to rethink this.

[jorgensd]

Something along the lines of:

from mpi4py import MPI
import dolfinx.fem.petsc
import ufl
import numpy as np

mesh = dolfinx.mesh.create_unit_square(MPI.COMM_WORLD, 10, 10)
num_cells_local = mesh.topology.index_map(mesh.topology.dim).size_local
ct = dolfinx.mesh.meshtags(mesh, mesh.topology.dim, np.arange(num_cells_local, dtype=np.int32), np.ones(num_cells_local, dtype=np.int32))
V = dolfinx.fem.functionspace(mesh, ("Lagrange", 1))
u = dolfinx.fem.Function(V)
u.interpolate(lambda x: x[0] + x[1]**2)

v = ufl.TestFunction(V)

F = ufl.inner(u, v) * ufl.dx(subdomain_id=1, subdomain_data=ct)

F_c = dolfinx.fem.form(F)
F_s = dolfinx.fem.compile_form(mesh.comm, F)


i_g =  F_c._cpp_object.domains(dolfinx.fem.IntegralType.cell, 0)
batch_size = 10
F_batched = []
for i in range(batch_size):
    i_g_local = i_g[i*batch_size:(i+1)*batch_size]
    F_batch = dolfinx.fem.create_form(F_s, function_spaces=[V],
                                      msh=mesh, subdomains={dolfinx.fem.IntegralType.cell: [(1, i_g_local)]},
                                      coefficient_map={u: u}, constant_map={},
                                      entity_maps=[])
    F_batched.append(F_batch)


b = dolfinx.fem.petsc.assemble_vector(F_c)

b_batched = dolfinx.fem.petsc.create_vector(V)
for batch in F_batched:
    dolfinx.fem.petsc.assemble_vector(b_batched, batch)

np.testing.assert_allclose(b.array, b_batched)

Needs some further debugging.

[a-latyshev]

@jhale, what do you think about this batching strategy proposed by @jorgensd?

[jorgensd]

@jhale, what do you think about this batching strategy proposed by @jorgensd?

@jhale, I've made a refined version of the above code at: https://github.com/scientificcomputing/scifem/blob/dokken/batched/tests/test_batch_assembly.py

[jorgensd]

@jhale @a-latyshev: A more advance version of batching can be found at:
https://github.com/scientificcomputing/scifem/blob/dokken/batched/tests/test_batch_by_submesh.py
where one avoids allocating the full dofmap for the coefficient entirely by assembling over submesh patches. This means a dramatic reduction in memory usage for coefficients for large tensors.