Refactor value processing to fenics objects

[jhdark]

At the moment in the fenicsx branch sources, boundary conditions and fluxes (and define_temperature within HydorgenTransportProblem) have a function create_value_fenics or something similar in which we convert a user-given value toa fenics object. The type of object produced (e.g. fem.Constant or fem.function) is dependent on the type of value given. However this has resulted in a lot of code duplication and should be simplified. One potential solution could be to have some global functions within the helpers file, like so:

import festim as F
import numpy as np
from dolfinx import fem
import ufl

def as_fenics_constant(value, mesh):
    """Converts a value to a dolfinx.Constant

    Args:
        value (float, int or dolfinx.Constant): the value to convert
        mesh (dolfinx.mesh.Mesh): the mesh of the domiain

    Returns:
        dolfinx.fem.function.Constant: the converted value

    Raises:
        TypeError: if the value is not a float, an int or a dolfinx.Constant
    """
    if isinstance(value, (float, int)):
        return fem.Constant(mesh, float(value))
    elif isinstance(value, fem.Constant):
        return value
    else:
        raise TypeError(
            f"Value must be a float, an int or a dolfinx.Constant, not {type(value)}"
        )

def as_fenics_expr(
    value, mesh, function_space, temperature=None, t=None, species_dependent_value=None
):
    """Creates a dolfinx.fem.function.Expression object based on the input
    value given.
    If the value is a function of t, it is converted to a
    dolfinx.fem.function.Constant.
    Otherwise, it is converted to a dolfinx.fem.function.Expression and
    returned

    Args:
        value (float, int, fem.Constant or callable)
        mesh (dolfinx.mesh.Mesh) : the mesh
        function_space (dolfinx.fem.FunctionSpaceBase): the function space
        temperature (float, optional): the temperature
        t (dolfinx.fem.Constant, optional): the time
        species_dependent_value (dict, optional) : dictionary of species dependent arguments
            in value definition

    Returns:
        dolfinx.fem.function.Constant or dolfinx.fem.function.Expression:
            the value converted to a format suitable for fenics
    """
    x = ufl.SpatialCoordinate(mesh)

    arguments = value.__code__.co_varnames

    acceptable_args = "t,x,T"

    if "t" in arguments and "x" not in arguments and "T" not in arguments:
        # only t is an argument
        if not isinstance(value(t=float(t)), (float, int)):
            raise ValueError(
                f"value should return a float or an int, not {type(value(t=float(t)))} "
            )
        return as_fenics_constant(mesh=mesh, value=value(t=float(t))), None
    else:
        kwargs = {}
        if species_dependent_value is not None:
            for arg in species_dependent_value.keys():
                acceptable_args += f",{arg}"
                if not isinstance(species_dependent_value[arg], F.Species):
                    raise ValueError(
                        "F.Species objects must be passed to species_dependent_value"
                    )
                kwargs[arg] = species_dependent_value[arg].solution
        if "t" in arguments:
            kwargs["t"] = t
        if "x" in arguments:
            kwargs["x"] = x
        if "T" in arguments:
            kwargs["T"] = temperature

        for arg in arguments:
            if arg not in acceptable_args:
                raise ValueError(f"argument, {arg}, in value is not defined")

        return fem.Expression(
            value(**kwargs),
            function_space.element.interpolation_points(),
        )

def as_fenics_func_and_expr(
    value, mesh, function_space, temperature=None, t=None, species_dependent_value=None
):
    """Creates a dolfinx.fem.function.Function and a dolfinx.fem.function.Expression
    based on the input value given in the defined functionspace.

    Args:
        value (float, int, fem.Constant or callable)
        mesh (dolfinx.mesh.Mesh) : the mesh
        function_space (dolfinx.fem.FunctionSpaceBase): the function space
        temperature (dolfinx.fem.function.Constant, optional): the temperature
        t (dolfinx.fem.function.Constant, optional): the time
        species_dependent_value (dict, optional) : dictionary of species dependent arguments
            in value definition

    Returns:
        dolfinx.fem.function.Function: the value as a fenics function
        dolfinx.fem.function.Constant or dolfinx.fem.function.Expression:
            the value converted to a format suitable for fenics
    """
    value_fenics_func = fem.Function(function_space)

    value_fenics_expr = as_fenics_expr(
        value=value,
        mesh=mesh,
        function_space=function_space,
        temperature=temperature,
        t=t,
        species_dependent_value=species_dependent_value,
    )

    value_fenics_func.interpolate(value_fenics_expr)

    return value_fenics_func, value_fenics_expr

def create_fenics_objects_from_value(
    value, mesh, function_space, temperature=None, t=None, species_dependent_value=None
):
    """Processes a defined value into a fenics object depending the type of
    the defined value.
    If the value is a constant, it is converted to a fenics.Constant.
    If the value is a function of t, it is converted to a fenics.Constant.
    Otherwise, it is converted to a dolfinx.fem.function.Function and the
    associated expression of the function.

    Args:
        mesh (dolfinx.mesh.Mesh) : the mesh
        function_space (dolfinx.fem.FunctionSpaceBase): the function space
        temperature (float): the temperature
        t (dolfinx.fem.Constant): the time

    Returns:
        fem.Constant, fem.Function or ufl.core.Expr : the value as a fenics
            object
        None, dolfinx.fem.function.Constant or dolfinx.fem.function.Expression:
            the value associated expreesion for the fucntion, if needed
    """
    if isinstance(value, (int, float)):
        return as_fenics_constant(mesh=mesh, value=value), None

    elif isinstance(value, (fem.Constant, fem.Function, ufl.core.expr.Expr)):
        return value, None

    elif callable(value):
        return as_fenics_func_and_expr(
            value=value,
            mesh=mesh,
            function_space=function_space,
            temperature=temperature,
            t=t,
            species_dependent_value=species_dependent_value,
        )

Then, within each Problem class, we could process each of the sources, boundary conditions and fluxes like so:

for bc in self.boundary_conditions:
    my_bc.value_fenics, my_bc.bc_expr = create_fenics_objects_from_value(
        value=my_bc.value,
        mesh=self.mesh.mesh,
        function_space=bc.species.function_space,
        temperature=self.temperature_fenics,
        t=self.t,
        species_dependent_value=my_bc.species_dependent_value,
    )

However, we will still need the additional logic to know which function space to give, as in monospecies cases, it would be the subfunction space. In contrast, in multispecies cases, it would be the collapsed function space (as mixed elements are used). In the case of DirichletBCs, if a value is given as a function and the case is multispecies, both need to be given to obtain the dofs of the surface. I don't know if anyone knows how this could be integrated too. This may have to be a question for the fenics devs on their discourse.

Finally, the temperature is a optional argument in as_fenics_expr and as_fenics_func_and_expr so that we can also use these functions in HeatTransferProblem, however we could have an additional layer of checking by having all the base classes with the temperature_dependent property. Then within heat transfer relevant classes such as HeatSource we could have a check within the __init__:

if self.temperature_dependent:
            raise ValueError("Temperature source cannot be termperature dependent")