Only used mixed function spaces

[jhdark]

Proposed changes

In dolfinx v0.9.0, it became possible to have a mixed function space with a single element - something we can take advantage of, as it can simplify some logic within FESTIM. Currently, we only use a mixed function space if we have a multispecies case.

Types of changes

What types of changes does your code introduce to FESTIM?

• [ ] Bugfix (non-breaking change which fixes an issue)
• [ ] New feature (non-breaking change which adds functionality)
• [ ] Breaking change (fix or feature that would cause existing functionality to not work as expected)
• [x] Code refactoring
• [ ] Documentation Update (if none of the other choices apply)
• [ ] New tests

Checklist

• [x] Black formatted
• [ ] Unit tests pass locally with my changes
• [ ] I have added tests that prove my fix is effective or that my feature works
• [ ] I have added necessary documentation (if appropriate)

Further comments

If this is a relatively large or complex change, kick off the discussion by explaining why you chose the solution you did and what alternatives you considered, etc...

[codecov[bot]]

Codecov Report

:white_check_mark: All modified and coverable lines are covered by tests. :white_check_mark: Project coverage is 94.87%. Comparing base (5cdfd31) to head (226fbab). :warning: Report is 10 commits behind head on fenicsx.

Additional details and impacted files

@@             Coverage Diff             @@
##           fenicsx    #1042      +/-   ##
===========================================
- Coverage    94.93%   94.87%   -0.06%     
===========================================
  Files           44       44              
  Lines         3120     3085      -35     
===========================================
- Hits          2962     2927      -35     
  Misses         158      158              

:umbrella: View full report in Codecov by Sentry.
:loudspeaker: Have feedback on the report? Share it here.

:rocket: New features to boost your workflow:

• :snowflake: Test Analytics: Detect flaky tests, report on failures, and find test suite problems.

[jhdark]

@jorgensd Do you foresee any issue with using only mixed element function spaces used like this? Just simplifies some logic for us.

[jorgensd]

@jorgensd Do you foresee any issue with using only mixed element function spaces used like this? Just simplifies some logic for us.

I'll try to have a look this week. In general, if not every sub space is tightly coupled (i.e. every sub components in the matrix is non-zero) mixedfunctionspace is way better for memory handling). Here is an illustration:

from mpi4py import MPI

import dolfinx.fem.petsc
import basix.ufl
import ufl

N = 10
mesh = dolfinx.mesh.create_unit_cube(MPI.COMM_WORLD, N, N, N)
el = basix.ufl.element("P", mesh.basix_cell(), degree=1)


def mix_function_space(mesh, el, num_spaces: int):
    V = dolfinx.fem.functionspace(mesh, el)
    return ufl.MixedFunctionSpace(*[V.clone()] * num_spaces)


def mixed_element(mesh, el, num_spaces: int):
    return dolfinx.fem.functionspace(mesh, basix.ufl.mixed_element([el] * num_spaces))


def block_element(mesh, el, num_spaces: int):
    return dolfinx.fem.functionspace(
        mesh, basix.ufl.element("P", mesh.basix_cell(), degree=1, shape=(num_spaces,))
    )


num_spaces = [1, 2, 4, 8, 16, 32]

nz_allocated = {mix_function_space: [], mixed_element: [], block_element: []}

for ns in num_spaces:
    for const, eb in [
        (mix_function_space, True),
        (mixed_element, False),
        (block_element, False),
    ]:
        W = const(mesh, el, ns)
        u = ufl.TrialFunctions(W)
        v = ufl.TestFunctions(W)
        a = 0
        # for i in range(ns):
        #     for j in range(ns):
        #         if i <= j:
        #             a += ufl.inner(u[i], v[j]) * ufl.dx
        a = sum(ufl.inner(u[i], v[i]) * ufl.dx for i in range(ns))
        if eb:
            a = ufl.extract_blocks(a)
        A = dolfinx.fem.petsc.assemble_matrix(dolfinx.fem.form(a))
        A.assemble()
        info = A.getInfo()
        print(ns, A.getSizes(), info["nz_allocated"])

        A.destroy()
        nz_allocated[const].append((ns, info["nz_allocated"]))
nz_to_label = {
    mix_function_space: "MixedFunctionSpace",
    mixed_element: "MixedElement",
    block_element: "BlockElement",
}
nz_to_tick = {
    mix_function_space: "o",
    mixed_element: "s",
    block_element: "^",
}

import matplotlib.pyplot as plt

for const, nzs in nz_allocated.items():
    ns, nzs = zip(*nzs)
    plt.plot(ns, nzs, label=nz_to_label[const], marker=nz_to_tick[const])
plt.yscale("log")
plt.xlabel("Number of spaces")
plt.ylabel("Nonzeros allocated")
plt.legend()
plt.grid()
plt.savefig("nz_diag_allocated.png", dpi=400)

This above is the scenario with no coupling. If there is say lower-diagonal coupling, you get:

[RemDelaporteMathurin]

MixedFunctionSpace is always better no?

[jorgensd]

MixedFunctionSpace is always better no?

Worst case scenario is equivalent memory with a different memory layout

[jorgensd]

So I guess I misunderstood the title of the PR. This makes it such that you always use mixed-element, even in the single-species case. I think this is ok (especially after the long debate at: https://github.com/FEniCS/dolfinx/pull/3520

However, as I showed above, if the species are weakly coupled, we should rather consider each of these spaces as individual function spaces, and then having a main-mixedfunctionspace that brings them all together in a blocked matrix.

That is is a major rewrite (and would use ufl.extract_blocks to set up the blocked structure, rather than coding it by hand.