Problems with numpy arrays of functions over MixedElement

Hello,

I compute the solution of a system of time dependent PDEs as an np.ndarray of elements of of a FunctionSpace over a MixedElement (like arr_Y in the code below). This solution gets passed to a lot of other functions and I don’t want to pass the corresponding FucntionSpace with it all the time.

For “simple” function spaces this is not needed since I can get them from the solutions themselves via .function_space(), but somehow this does not work for the mixed case.

Below I have an example code showing this problem. I have marked three sections with i), ii) and iii) respectively for reference in my question (and hopefully your answers). For comparison I added the “simple” non mixed case in the code as well.

Q1: What should I do if I want to get the function_space in i) ?
Q2: Why does unpacking the numpy.ndarray still results in an numpy.ndarray in ii)?
Q3: How get I get the function space from an Indexed object in general (see iii)?

Of course these “problems” are not restricted to .function_space.

I would greatly appreciate any explanation of what is going on or any answers to any of my questions.

import numpy
from fenics import *

mesh = UnitSquareMesh(4, 4)

S1 = FiniteElement('CG', mesh.ufl_cell(), 1)
V = FunctionSpace(mesh, S1)

S1S1 = MixedElement([S1, S1])
Y = FunctionSpace(mesh, S1S1)

lst_V = [Function(V), Function(V)]
lst_Y = [Function(Y), Function(Y)]

arr_V = numpy.array(lst_V)
arr_Y = numpy.array(lst_Y)

# this works
assert hasattr(lst_V[0], 'function_space')
assert hasattr(lst_Y[0], 'function_space')

# this also works
assert hasattr(arr_V[0], 'function_space')

# i) but this not
try:
    assert hasattr(arr_Y[0], 'function_space')
except AssertionError:
    print('`arr_Y[0]` does not have attribute `function_space`!')

# ii) I do not see why but indexing `arr_Y` still gives an `ndarray`
assert isinstance(arr_Y[0], numpy.ndarray)
# So I can unpack as many times as I want ...
assert isinstance(arr_Y[:][:][:], numpy.ndarray)

# iii) but also this does not work
try:
    assert hasattr(arr_Y[0][0], 'function_space')
except AssertionError:
    print('`arr_Y[0][0]` does not have attribute `function_space`!')

Hi,

ad ii) observe that the arrays you create have different shapes (2-vector vs. 2x2 matrix). This is due to the fact that Function(Y), being a function in MixedFunctionSpace is iterable, (len(Function(Y)) is 2). Note that you’d get the same behavior with functions from VectorFunctionSpace.

ad i) it is not the most elegant solution but consider

x = numpy.zeros_like(arr_V)
x[0] = Function(Y)

ad iii) Operands of the Indexed node are the coefficient and the indices. Consider

(Function(Y)[0]).ufl_operands[0].function_space()

Hey,

thanks a lot, that really helped. I didn’t see that arr_y.shape was (2, 2).
Your suggestion of using .ufl_operands is also very good, but I it does not quite do what I expected.

Consider Y1, Y2 = Y.split().
For me it would make sense if
(Function(Y)[0]).ufl_operands[0].function_space() == Y1
and
(Function(Y)[1]).ufl_operands[0].function_space() == Y2
while in fact both times the result is the full space Y.
Is there also a way of extracting the correct subspace (without the information about the index)?
Something like this (except this does not work because of the reason above):

def print_function_space(func):
    if isinstance(func, ufl.indexed.Indexed):
        func = func.ufl_operands[0]
    print(func.function_space())

I think I found a solution. I misunderstood ufl_operands.
The index I seek can be found in ufl_operands[1].
With this the above function would read

def print_function_space(func):
    if isinstance(func, ufl.indexed.Indexed):
        whole_space = func.ufl_operands[0].function_space()
        idx = int(func.ufl_operands[1]indices()[0])
        this_space = whole_space.sub(idx)collapse()
    else:
        this_space = func.function_space()

    print(this_space)

Just for completeness I am posting a slightly more general version

import ufl


def function_space_of(f):
    try:
        return f.function_space()
    except AttributeError:
        pass
    
    if isinstance(f, ufl.indexed.Indexed):
        f, index = f.ufl_operands
        # Flat index for tensors
        if len(index) > 1:
            index = int(index[-1]) + sum(ufl.shape(f)[:-1])
        else:
            index = index[0]
        index = int(index)
        
        return function_space_of(f.split()[index])

    
def same_space(V, W):
    return all((V.ufl_element() == W.ufl_element(),))
 
    
# --------------------------------------------------------------------

if __name__ == '__main__':
    from dolfin import *

    mesh = UnitSquareMesh(2, 2)

    V0 = FunctionSpace(mesh, 'CG', 1)
    assert same_space(function_space_of(Function(V0)), V0)
    
    V1 = VectorFunctionSpace(mesh, 'CG', 1)
    assert same_space(function_space_of(Function(V1)[0]), V0)
    

    V2 = TensorFunctionSpace(mesh, 'CG', 1)
    assert same_space(function_space_of(Function(V2)[1, 0]), V0)

    V = FunctionSpace(mesh,
                      MixedElement([V0.ufl_element(),
                                    V1.ufl_element(),
                                    V2.ufl_element()]))
                                    
    assert same_space(function_space_of(Function(V)[0]), V0)
    assert same_space(function_space_of(Function(V)[1]), V1)
    assert same_space(function_space_of(Function(V)[2]), V2)

Very nice. Thank you for your help.
I will mark the thread as solved.