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)
• Code refactoring
• Documentation Update (if none of the other choices apply)
• New tests

Checklist

• 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]

Codecov Report

✅ All modified and coverable lines are covered by tests.
✅ Project coverage is 94.87%. Comparing base (5cdfd31) to head (226fbab).
⚠️ 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              

☔ View full report in Codecov by Sentry.
📢 Have feedback on the report? Share it here.

🚀 New features to boost your workflow:

• ❄️ 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]

Will there be strong coupling across all these species?

If not, one should probably have a "global" MixedFunctionSpace where each sub-space is its own function space, and this class would keep track of the local->global indexing.

[jhdark]

There will be some strong coupling between different combinations of species from the reaction framework.

[jorgensd]

That there is coupling between them is on for either approach, I’ll try to rephrase:

Given M species, if we make an MxM matrix where M_ij =1 means that there is a coupling term, M_ij=0 means that there is not, how many non-zeros do we usually get?

[jorgensd]

Not just relevant for this PR, but currently, by assigning data in this way, without having a getter/setter structure, users can overwrite spe.sub_function_space in their own code, and invalidate the logic in FESTIM.
In general, I think it is a good idea to do something like:

from mpi4py import MPI
import dolfinx


class Species:
    _sub_function_space: None | dolfinx.fem.function.FunctionSpace

    def __init__(self):
        self._sub_function_space = None

    @property
    def sub_function_space(self) -> dolfinx.fem.FunctionSpace:
        if self._sub_function_space is None:
            raise ValueError("sub_function_space has not been set yet")
        else:
            return self._sub_function_space

    def set_sub_function_space(self, function_space: dolfinx.fem.FunctionSpace) -> None:
        if self._sub_function_space is not None:
            raise ValueError("sub_function_space has already been set")
        self._sub_function_space = function_space

to ensure that data is not set twice (mistakenly).
If you want these to be adaptive and be possible to set multiple times, I would advice having some warnings on it.

[jhdark]

Yeah, this would be better. I think we often overwrite this or assign it ourselves in a lot of the tests to simplify them, though, so it may make them more complicated

[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: FEniCS/dolfinx#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.