Basis Evaluation Struct

[dtoshniwal]

Hi @apalha @J15525 @dccabanas:

I started wrapping the evaluations of basis functions in a struct. The code is attached below. Since the earlier, explicit indexing was not obvious/was hard to remember, the getters here should help clarify what you are asking for. However, are these all the getters that we want? Is this enough functionality or would we like to access things in a different way?

"""
    EvaluatedBasis

A struct to store the evaluations of finite element basis functions and their derivatives.

The raw data is stored in a nested vector format, but should be accessed through the provided
getter methods for clarity.

# Fields
- `data::Vector{Vector{Vector{Matrix{Float64}}}}`: The core data store where:
    - `data[i][j][k]` corresponds to a matrix of evaluations where
     - `i`: derivative order `i-1`
     - `j`: `j`-th mixed derivative of order `i-1` (indexing specified by get_derivative_idx)
     - `k`: index for the component of a multicomponent function.
"""
struct EvaluatedBasis{manifold_dim}
    data::Vector{Vector{Vector{Matrix{Float64}}}}

    function EvaluatedBasis(manifold_dim::Int, data::Vector{Vector{Vector{Matrix{Float64}}}})
        return new{manifold_dim}(data)
    end
end

"""
    get_values(evals::EvaluatedBasis, component_id::Int=1)

Retrieves the evaluation matrix for the function values (0-th order derivatives).

# Arguments
- `evals::EvaluatedBasis`: The `EvaluatedBasis` object.
- `component_id::Int=1`: The component index for vector-valued elements. Defaults to 1 for scalar cases.

# Returns
- `Matrix{Float64}`: A matrix where `M[a, b]` is the value of the `b`-th basis function at the `a`-th evaluation point.
"""
function get_values(evals::EvaluatedBasis, component_id::Int=1)
    # Values are order 0, which corresponds to the first index `i=1`.
    # The derivative type index `j` is also 1.
    return @view evals.data[1][1][component_id]
end

"""
    get_partial_derivatives(evals::EvaluatedBasis, type::Vector{Int}, component_id::Int=1)

Retrieves the evaluation matrix for a specific derivative.

# Arguments
- `evals::EvaluatedBasis`: The `EvaluatedBasis` object.
- `type::Vector{Int}`: The type of mixed derivative. E.g.:
    -- `[1,0]` for ∂/∂x, `[0,1]` for ∂/∂y, `[1,1]` for ∂²/∂x∂y in 2D
    -- `[1,0,0]` for ∂/∂x, `[0,1,0]` for ∂/∂y, `[0,0,1]` for ∂/∂z,
       `[1,1,0]` for ∂²/∂x∂y, etc. in 3D.
- `component_id::Int=1`: (Optional) The component index. Defaults to 1.

# Returns
- `Matrix{Float64}`: A matrix where `M[a, b]` is the value of the specified derivative of the `b`-th basis function at the `a`-th evaluation point.
"""
function get_partial_derivatives(
    evals::EvaluatedBasis{manifold_dim},
    type::Vector{Int},
    component_id::Int=1
) where {manifold_dim}
    # The data index `i` is `order + 1`.
    order = sum(type)
    if order == 0
        return get_values(evals, component_id)
    end
    return @view evals.data[order+1][get_derivative_idx(type)][component_id]
end

[dtoshniwal]

PS: I added manifold_dim as a parameter for the struct because I was thinking about two checks for the method get_partial_derivatives:

• Checking that all entries of type are non-negative
• Checking that the length of type is equal to manifold_dim
  However, while useful, do you think it will be expensive to do these checks every time we have to access evaluations?

PPS: The second check could be avoided by setting type::NTuple{manifold_dim, Int}, but the get_derivative_idx method accepts Vector{Int} so that will require a conversion, a bit unnecessary.

[dtoshniwal]

Actually, instead of manifold_dim, maybe num_components is a more relevant parameter. Also, I think the function apply_extraction_op (see below) could be useful.

"""
    apply_extraction_op(
        evals::EvaluatedBasis{num_components},
        extraction_op::ExtractionOperator{num_components},
        element_id::Int,
    ) where {num_components}

Applies an extraction operator to the evaluated basis functions. It is assumed that the input `evals` contains evaluations of the component spaces for a finite element space. The extraction operator is applied to transform these component evaluations into the evaluations of the actual finite element basis functions.

# Arguments
- `evals::EvaluatedBasis{num_components}`: The evaluated basis functions of the
    component spaces.
- `extraction_op::ExtractionOperator{num_components}`: The extraction operator
    to be applied.
- `element_id::Int`: The identifier for the element on which the extraction
    operator is applied.

# Returns
- `EvaluatedBasis{num_components}`: A new `EvaluatedBasis` object containing
    the transformed evaluations after applying the extraction operator.
"""
function apply_extraction_operator(
    evals::EvaluatedBasis{num_components},
    extraction_op::ExtractionOperator{num_components},
    element_id::Int,
) where {num_components}
    num_basis = length(get_basis_indices(extraction_op, element_id))

    # allocate space for the transformed evaluations
    transformed_data = Vector{Vector{Vector{Matrix{Float64}}}}(undef, length(evals.data))
    for i in eachindex(evals.data)
        transformed_data[i] = Vector{Vector{Matrix{Float64}}}(undef, length(evals.data[i]))
        for j in eachindex(transformed_data[i])
            transformed_data[i][j] = [
                zeros(size(evals.data[i][j][k]), num_basis) for k in 1:num_components
            ]
        end
    end

    # transform the evaluations by applying the extraction operator
    for component_id in 1:num_components
        extraction_coefficients, J = get_extraction(extraction_op, element_id, component_id)
        for i in eachindex(transformed_data)
            for j in eachindex(transformed_data[i])
                LinearAlgebra.mul!(
                    view(transformed_data[i][j][component_id], :, J),
                    evals.data[i][j][component_id],
                    extraction_coefficients,
                )
            end
        end
    end

    return EvaluatedBasis{num_components}(transformed_data)
end

This would simplify the evaluate method in FiniteElementSpaces.jl:

function evaluate(
    space::AbstractFESpace{manifold_dim, num_components, num_patches},
    element_id::Int,
    xi::Points.AbstractPoints{manifold_dim},
    nderivatives::Int=0,
) where {manifold_dim, num_components, num_patches}
    basis_indices = get_basis_indices(space, element_id)
    component_wise_evaluations = get_component_space_evaluations(space, element_id, xi)
    space_evaluations = apply_extraction_operator(
        component_wise_evaluations,
        get_extraction_operator(space),
        element_id
    )

    return space_evaluations, basis_indices
end

[J15525]

I think it would be nice to add the FunctionSpace to the type parameter/field, and thus have it as an input as well. This will allow us to check to which space these evaluations belong (may not be so important now, but might be helpful in complicated cases as a safety/debugging check, and perhaps for some caching-related things?), and it will allow us to infer both the manifold_dim and num_components (and patches if need be) directly from the input types. I don't know now if that will help performance, but we will probably construct these structs on the fly in performance-sensitive code.

Also, is this specific for FunctionSpaces, or can we utilise similar functionality for geometry/analytical inputs as well?

[J15525]

PS: I added manifold_dim as a parameter for the struct because I was thinking about two checks for the method get_partial_derivatives:

• Checking that all entries of type are non-negative
• Checking that the length of type is equal to manifold_dim
  However, while useful, do you think it will be expensive to do these checks every time we have to access evaluations?

PPS: The second check could be avoided by setting type::NTuple{manifold_dim, Int}, but the get_derivative_idx method accepts Vector{Int} so that will require a conversion, a bit unnecessary.

I think it would be good to test this on the type. I think we can easily rewrite the get_derivative_idx functions to accept tuples, because I don't think it relies on vectors per se. The non-negativity check is already done in get_derivative_idx as well.

I also think that you wouldn't have to special-case the 0-th partial derivative, as the correct key will be returned anyway.

[dccabanas]

I agree that we should make this more general. I don't see why the name is EvaluatedBasis since this framework works for FunctionSpaces, FormSpaces and FormFields, for instance. We should also make geometry evaluations behave like this for consistency.

And the type (name is kind of vague) should be, or at least accept tuples as well, I don't see why get_derivative_idx can't handle that.

For the type I don't really see why we would want the thing being evaluated as a parameter, I think manifold_dim and num_components (perhaps also num_patches) could be enough.

[dtoshniwal]

@J15525 @dccabanas
Thanks for the comments.

• Regarding NTuple vs Vector for get_derivative_idx: I know that the method get_derivative_idx will not need to be changed in any major way to make it work with NTuples. The point is that the input for this method comes from the method integer_sums. So we will have to modify that one to ensure that all outputs are now only NTuples instead of Vectors. This is also not a fundamental issue, in fact we already had this setup a while back. Not sure why we got rid of it, but it seems like we are turning back the clock here.
• Regarding the parameters, not sure about adding the function space (or geometry or ...) to the type. This seems like an object that can exist on its own without reference to how it was created. This will also mean that we will always need to update the type whenever we transform the object. For instance, whenever we apply an extraction operator, or when we compose the evaluated jacobian of a geometry with that of a mapping. And this update seems to have no effect on its functionality so it seems a bit unnecessary to me. Maybe manifold_dim, num_components are enough parameters to keep track of.
• The name EvaluatedBasis already works for Forms (including FormFields where there is just one basis function) and FunctionSpaces. But maybe it would be good to include the geometry as well. In that case, the object could be called Evaluations?

[dtoshniwal]

@J15525 @dccabanas Thanks for the comments.

* Regarding NTuple vs Vector for `get_derivative_idx`: I know that the method `get_derivative_idx` will not need to be changed in any major way to make it work with NTuples. The point is that the input for this method comes from the method `integer_sums`. So we will have to modify that one to ensure that all outputs are now only NTuples instead of Vectors. This is also not a fundamental issue, in fact we already had this setup a while back. Not sure why we got rid of it, but it seems like we are turning back the clock here.

* Regarding the parameters, not sure about adding the function space (or geometry or ...) to the type. This seems like an object that can exist on its own without reference to how it was created. This will also mean that we will always need to update the type whenever we transform the object. For instance, whenever we apply an extraction operator, or when we compose the evaluated jacobian of a geometry with that of a mapping. And this update seems to have no effect on its functionality so it seems a bit unnecessary to me. Maybe `manifold_dim, num_components` are enough parameters to keep track of.

* The name `EvaluatedBasis` already works for `Forms` (including `FormFields` where there is just one basis function) and `FunctionSpaces`. But maybe it would be good to include the geometry as well. In that case, the object could be called `Evaluations`?

To follow-up on the last comment: in this case, however, it will be difficult to write specialized methods (e.g., apply_extraction_operator) unless we keep track of the type of object being evaluated. The module in this case will be general enough that it will be included before FunctionSpaces, Geometry and Forms, but then type-specific specializations for Evaluations will be written in the appropriate modules.

[dtoshniwal]

An example of what it could look like with the type-of-object-being-evaluated parameter in there:

module Evaluations

"""
    Evaluations

A struct to store the evaluations of objects and their derivatives.

The raw data is stored in a nested vector format, but should be accessed through the provided
getter methods for clarity.

# Fields
- `data::Vector{Vector{Vector{Matrix{Float64}}}}`: The core data store where:
    - `data[i][j][k]` corresponds to a matrix of evaluations where
     - `i`: derivative order `i-1`
     - `j`: `j`-th mixed derivative of order `i-1` (indexing specified by get_derivative_idx)
     - `k`: index for the component of a multicomponent function.
"""
struct Evaluations{T}
    data::Vector{Vector{Vector{Matrix{Float64}}}}

    function Evaluations(T, data::Vector{Vector{Vector{Matrix{Float64}}}})
        return new{T}(data)
    end
end

"""
    get_values(evals::Evaluations, component_id::Int=1)

Retrieves the evaluation matrix for the function values (0-th order derivatives).

# Arguments
- `evals::Evaluations`: The `Evaluations` object.
- `component_id::Int=1`: The component index for vector-valued elements. Defaults to 1 for scalar cases.

# Returns
- `Matrix{Float64}`: A matrix where `M[a, b]` is the value of the `b`-th basis function at the `a`-th evaluation point.
"""
function get_value(evals::Evaluations, component_id::Int=1)
    # Values are order 0, which corresponds to the first index `i=1`.
    # The derivative type index `j` is also 1.
    return @view evals.data[1][1][component_id]
end

"""
    get_partial_derivative(evals::Evaluations, type::Vector{Int}, component_id::Int=1)

Retrieves the evaluation matrix for a specific derivative.

# Arguments
- `evals::Evaluations`: The `Evaluations` object.
- `type::Vector{Int}`: The type of mixed derivative. E.g.:
    -- `[1,0]` for ∂/∂x, `[0,1]` for ∂/∂y, `[1,1]` for ∂²/∂x∂y in 2D
    -- `[1,0,0]` for ∂/∂x, `[0,1,0]` for ∂/∂y, `[0,0,1]` for ∂/∂z,
       `[1,1,0]` for ∂²/∂x∂y, etc. in 3D.
- `component_id::Int=1`: (Optional) The component index. Defaults to 1.

# Returns
- `Matrix{Float64}`: A matrix where `M[a, b]` is the value of the specified derivative of the `b`-th basis function at the `a`-th evaluation point.
"""
function get_partial_derivative(
    evals::Evaluations,
    order::Vector{Int},
    component_id::Int=1
)
    # The data index `i` is `order + 1`.
    return @view evals.data[sum(order)+1][get_derivative_idx(order)][component_id]
end

"""
    apply_extraction_op(
        evals::Evaluations{num_components},
        extraction_op::ExtractionOperator{num_components},
        element_id::Int,
    ) where {num_components}

Applies an extraction operator to the evaluated basis functions. It is assumed that the input `evals` contains evaluations of the component spaces for a finite element space. The extraction operator is applied to transform these component evaluations into the evaluations of the actual finite element basis functions.

# Arguments
- `evals::Evaluations{num_components}`: The evaluated basis functions of the
    component spaces.
- `extraction_op::ExtractionOperator{num_components}`: The extraction operator
    to be applied.
- `element_id::Int`: The identifier for the element on which the extraction
    operator is applied.

# Returns
- `Evaluations{num_components}`: A new `Evaluations` object containing
    the transformed evaluations after applying the extraction operator.
"""
function local_to_global_extraction(
    local_evals::Evaluations{T1},
    space::T2,
    element_id::Int,
) where {
    manifold_dim, num_components,
    T1 <: AbstractFESpace{manifold_dim, num_components},
    T2 <: AbstractFESpace{manifold_dim, num_components}
}
    num_basis = length(get_basis_indices(space, element_id))

    # allocate space for the transformed evaluations
    transformed_data = Vector{Vector{Vector{Matrix{Float64}}}}(undef, length(local_evals.data))
    for i in eachindex(local_evals.data)
        transformed_data[i] = Vector{Vector{Matrix{Float64}}}(undef, length(local_evals.data[i]))
        for j in eachindex(transformed_data[i])
            transformed_data[i][j] = [
                zeros(size(local_evals.data[i][j][k]), num_basis) for k in 1:num_components
            ]
        end
    end

    # transform the evaluations by applying the extraction operator
    for component_id in 1:num_components
        extraction_coefficients, J = get_extraction(space, element_id, component_id)
        for i in eachindex(transformed_data)
            for j in eachindex(transformed_data[i])
                LinearAlgebra.mul!(
                    view(transformed_data[i][j][component_id], :, J),
                    evals.data[i][j][component_id],
                    extraction_coefficients,
                )
            end
        end
    end

    return Evaluations(T2, transformed_data)
end

end

[dtoshniwal]

@dccabanas @apalha @J15525
Yet another alternative would be to only have a basic evaluation storage struct, no specialized methods (like local_to_global_extraction) for it.

• The reason is that there are several different types of transformations that are applied to evaluations in our code (e.g., pullback to canonical, scalings of derivatives, extraction, etc.) and maybe it doesn't make sense to have specialized methods for these things because these are transformations that are not always type-specific, and a lot of times it wouldn't make sense to create a copy of the evaluations just because a simple scaling was applied.
• Instead, every module modifies and views evaluations in its own way through setters and getters. Something like this:

module Evaluations

"""
    Evaluations

A struct to store the evaluations of objects and their derivatives.

The raw data is stored in a nested vector format, but should be accessed through the provided
getter methods for clarity.

# Fields
- `data::Vector{Vector{Vector{Matrix{Float64}}}}`: The core data store where:
    - `data[i][j][k]` corresponds to a matrix of evaluations where
     - `i`: derivative order `i-1`
     - `j`: `j`-th mixed derivative of order `i-1` (indexing specified by get_derivative_idx)
     - `k`: index for the component of a multicomponent function.
"""
struct Evaluations
    data::Vector{Vector{Vector{Matrix{Float64}}}}
    size::NTuple{5, Int}  # (num_derivative_orders, num_derivative_types, num_components, num_eval_points, num_basis_functions)

    function Evaluations(size::NTuple{5, Int})
        data = Vector{Vector{Vector{Matrix{Float64}}}}(undef, size[1])
        for i in 1:size[1]
            data[i] = Vector{Vector{Matrix{Float64}}}(undef, size[2])
            for j in 1:size[2]
                data[i][j] = [zeros(size[4], size[5]) for _ in 1:size[3]]
            end
        end
        return new(data, size)
    end
end

Base.size(evals::Evaluations) = evals.size
Base.similar(evals::Evaluations) = Evaluations(Base.similar(evals.data), evals.size)

"""
    set_value!(
        evals::Evaluations,
        value::Matrix{Float64},
        component_id::Int=1
    )
"""
function set_value!(
    evals::Evaluations,
    value::Matrix{Float64},
    component_id::Int=1
)
    evals.data[1][1][component_id] .= value
    return nothing
end

"""
    set_partial_derivative!(
        evals::Evaluations,
        value::Matrix{Float64},
        order::Vector{Int},
        component_id::Int=1
    )

Sets the evaluation matrix for a specific derivative.

# Arguments
- `evals::Evaluations`: The `Evaluations` object.
- `value::Matrix{Float64}`: The matrix to set, where `M[a
, b]` is the value of the specified derivative of the `b`-th basis function at the `a`-th evaluation point.
- `order::Vector{Int}`: The type of mixed derivative. E.g.:
    -- `[1,0]` for ∂/∂x, `[0,1]` for ∂/∂y, `[1,1]` for ∂²/∂x∂y in 2D
    -- `[1,0,0]` for ∂/∂x, `[0,1,0]` for ∂/∂y, `[0,0,1]` for ∂/∂z,
       `[1,1,0]` for ∂²/∂x∂y, etc. in 3D.
- `component_id::Int=1`: (Optional) The component index. Defaults to 1.

# Returns
- `Nothing`: The function modifies the `evals` object in place.
"""
function set_partial_derivative!(
    evals::Evaluations,
    value::Matrix{Float64},
    order::Vector{Int},
    component_id::Int=1
)
    evals.data[sum(order)+1][get_derivative_idx(order)][component_id] .= value
    return nothing
end

"""
    get_value(evals::Evaluations, component_id::Int=1)

Retrieves the evaluation matrix for the function values (0-th order derivatives).

# Arguments
- `evals::Evaluations`: The `Evaluations` object.
- `component_id::Int=1`: The component index for vector-valued elements. Defaults to 1 for scalar cases.

# Returns
- `Matrix{Float64}`: A matrix where `M[a, b]` is the value of the `b`-th basis function at the `a`-th evaluation point.
"""
function get_value(evals::Evaluations, component_id::Int=1)
    # Values are order 0, which corresponds to the first index `i=1`.
    # The derivative type index `j` is also 1.
    return @view evals.data[1][1][component_id]
end

"""
    get_partial_derivative(evals::Evaluations, type::Vector{Int}, component_id::Int=1)

Retrieves the evaluation matrix for a specific derivative.

# Arguments
- `evals::Evaluations`: The `Evaluations` object.
- `type::Vector{Int}`: The type of mixed derivative. E.g.:
    -- `[1,0]` for ∂/∂x, `[0,1]` for ∂/∂y, `[1,1]` for ∂²/∂x∂y in 2D
    -- `[1,0,0]` for ∂/∂x, `[0,1,0]` for ∂/∂y, `[0,0,1]` for ∂/∂z,
       `[1,1,0]` for ∂²/∂x∂y, etc. in 3D.
- `component_id::Int=1`: (Optional) The component index. Defaults to 1.

# Returns
- `Matrix{Float64}`: A matrix where `M[a, b]` is the value of the specified derivative of the `b`-th basis function at the `a`-th evaluation point.
"""
function get_partial_derivative(
    evals::Evaluations,
    order::Vector{Int},
    component_id::Int=1
)
    # The data index `i` is `order + 1`.
    return @view evals.data[sum(order)+1][get_derivative_idx(order)][component_id]
end

end

[dtoshniwal]

PS: In the above case, we would omit the paramter {T}.

[dtoshniwal]

@apalha @J15525 @dccabanas

I have flipped the order of derivatives and components in the evaluation struct, and have accordingly modified the canonical spaces (except Lagrange). This is going to be a big change (within "src", ~35 files affected in multiple places), so it would be nice if you all can take a look and see if the interface for setting/getting, and if the struct itself makes sense.

[J15525]

@apalha @J15525 @dccabanas

I have flipped the order of derivatives and components in the evaluation struct, and have accordingly modified the canonical spaces (except Lagrange). This is going to be a big change (within "src", ~35 files affected in multiple places), so it would be nice if you all can take a look and see if the interface for setting/getting, and if the struct itself makes sense.

You now have

Memoization.@memoize function evaluate(
    polynomials::Bernstein, xi::Points.CartesianPoints{1}, nderivatives::Int=0
)
    # store the values and derivatives here
    neval = length(xi)
    # degree
    p = get_polynomial_degree(polynomials)

    # allocate space for evaluations
    ders = Evaluations.BasisEvaluations(1, (1, nderivatives + 1, neval, p + 1))
    # loop over the evaluation points and evaluate all derivatives at each point
    for i in 0:nderivatives
        Evaluations.set_partial_derivative!(
            ders,
            _evaluate(polynomials, xi, i),
            1,
            [i]
        )
    end

    return ders
end

as new evaluate, and you mentioned in MantisFEM/MantisDev.jl#87 that it is slower and allocates more. What happens if you instead use this:

Memoization.@memoize function evaluate(
    polynomials::Bernstein, xi::Points.CartesianPoints{1}, nderivatives::Int=0
)
    # store the values and derivatives here
    neval = Points.get_num_points(xi)
    # degree
    p = get_polynomial_degree(polynomials)

    # allocate space for evaluations
    ders = Evaluations.BasisEvaluations(1, (1, nderivatives + 1, neval, p + 1))
    for cart_ind in CartesianIndices((nderivatives + 1, neval, p + 1))
        (derivative, point, basis) = Tuple(cart_ind)
        Evaluations.set_partial_derivative!(
            ders,
            _dbpoly(p, basis, derivative, xi[point][1]),
            1,
            [derivative],
            point,
            basis,
        )
    end

    return ders
end

I did not run this so I hope I interpreted everything correctly.