Vector laplacian behaves incorrectly on components with labels longer than one character

[Nemo-Malhomme]

Describe the bug
As implemented now, when applied to more than one component, the Laplacian operator will call scalar_laplace, which itself calls grad, for each requested component individually, giving grad the component label directly.
When a component label is itself an iterable (str) of length more than one, grad will interpret it as a list of components instead, and iterate over it as if it were several variables
This leads to the gradient being computed on the wrong variables, or on variables that do not exist, causing an error.

The following code reproduces the error:

import torch
from pina import LabelTensor
from pina.operators import laplacian

d = torch.rand((20, 2))
input_data = LabelTensor( d, ['x','y'] )
output_data = LabelTensor( 2*d, ['ab','cd'] )
laplacian(output_data, input_data, components=['ab','cd'])

Expected behavior
The Laplacians of 'ab' and 'cd' are computed.

Output
Attempts to compute the gradient of 'a', which causes an error:
ValueError: 'a' is not in list

Additional context
Using version 0.2.0 of pina-mathlab.

Suggestion
One possible fix would be to change line 233 of operator.py, currently:
result[:, idx] = scalar_laplace(output_, input_, c, d).flatten()
into:
result[:, idx] = scalar_laplace(output_, input_, [c], d).flatten()
This is the line that calls scalar_laplace when laplacian is called on several components. c is the variable that iterates over component labels, never containing more than one. By replacing c by [c], the grad function called by scalar_laplace now knows without ambiguity that it should be applied on a single variable, no matter the length of its label.

[GiovanniCanali]

Hi @Nemo-Malhomme, thank you for reporting the bug.
We will investigate on the issue and fix it as soon as possible!

[dario-coscia]

Hi @Nemo-Malhomme and @GiovanniCanali I tried to run the code but faced some issues, here I provide a clean snippet:

import torch
from pina import LabelTensor
from pina.operator import laplacian

d = torch.rand((20, 2))
input_data = LabelTensor(d, ['x','y']).requires_grad_(True)
output_data = LabelTensor(2*input_data**2, ['ab','cd'])
laplacian(output_data, input_data, components=['ab','cd'])

The error is the same as the one reported by @Nemo-Malhomme!

[GiovanniCanali]

Hi @Nemo-Malhomme,

We have fixed the bug you reported! However, as we are still running additional tests to further investigate the issue - which may take a few more days -, I want to share the corrected implementation of the Laplacian function with you. This way, your work with PINA won't be delayed.

Please, use the following method to compute the laplacian:

def laplacian(output_, input_, components=None, d=None, method="std"):
    """
    Compute the laplacian of the ``output_`` with respect to ``input``.

    This operator supports both vector-valued and scalar-valued functions with
    one or multiple input coordinates.

    :param LabelTensor output_: The output tensor on which the laplacian is
        computed.
    :param LabelTensor input_: The input tensor with respect to which the
        laplacian is computed.
    :param list[str] components: The names of the output variables for which to
        compute the laplacian. It must be a subset of the output labels.
        If ``None``, all output variables are considered. Default is ``None``.
    :param list[str] d: The names of the input variables with respect to which
        the laplacian is computed. It must be a subset of the input labels.
        If ``None``, all input variables are considered. Default is ``None``.
    :param str method: The method used to compute the Laplacian. Default is
        ``std``.
    :raises NotImplementedError: If ``std=divgrad``.
    :return: The computed laplacian tensor.
    :rtype: LabelTensor
    """

    def scalar_laplace(output_, input_, components, d):
        """
        Compute the laplacian of a scalar-valued ``output_``.

        :param LabelTensor output_: The output tensor on which the laplacian is
            computed. It must be a column tensor.
        :param LabelTensor input_: The input tensor with respect to which the
            laplacian is computed.
        :param list[str] components: The names of the output variables for which
            to compute the laplacian. It must be a subset of the output labels.
            If ``None``, all output variables are considered.
        :param list[str] d: The names of the input variables with respect to
            which the laplacian is computed. It must be a subset of the input
            labels. If ``None``, all input variables are considered.
        :return: The computed laplacian tensor.
        :rtype: LabelTensor
        """

        grad_output = grad(output_, input_, components=components, d=d)
        result = torch.zeros(output_.shape[0], 1, device=output_.device)

        for i, label in enumerate(grad_output.labels):
            gg = grad(grad_output, input_, d=d, components=[label])
            result[:, 0] += super(torch.Tensor, gg.T).__getitem__(i)

        return result

    if d is None:
        d = input_.labels

    if components is None:
        components = output_.labels

    if not isinstance(components, list):
        components = [components]

    if method == "divgrad":
        raise NotImplementedError("divgrad not implemented as method")

    if method == "std":

        result = torch.empty(
            input_.shape[0], len(components), device=output_.device
        )
        labels = [None] * len(components)
        for idx, c in enumerate(components):
            result[:, idx] = scalar_laplace(output_, input_, [c], d).flatten()
            labels[idx] = f"dd{c}"

        result = result.as_subclass(LabelTensor)
        result.labels = labels

    return result

[dario-coscia]

@Nemo-Malhomme We merged #524, which solved the problem. If you install PINA directly from GitHub, you can have the latest change; otherwise, remember that our PyPI updates once per month so it might be possible that the change is not visible today to you! Thank you again for reporting the issue :)

[Nemo-Malhomme]

@GiovanniCanali @dario-coscia Thank you both a lot for the help. This was a very efficient process. I will use the latest github version from now on.