Repository: mathLab/PINA

Issue:
Vector laplacian behaves incorrectly on components with labels longer than one character
**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.

Hints from issue comments:
@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:
```python
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:

```python
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.

Solution patch:
diff --git a/pina/operator.py b/pina/operator.py
index 68e2cb4..aad7160 100644
--- a/pina/operator.py
+++ b/pina/operator.py
@@ -27,12 +27,14 @@ def grad(output_, input_, components=None, d=None):
         computed.
     :param LabelTensor input_: The input tensor with respect to which the
         gradient is computed.
-    :param list[str] components: The names of the output variables for which to
+    :param components: The names of the output variables for which to
         compute the gradient. 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
+    :type components: str | list[str]
+    :param d: The names of the input variables with respect to which
         the gradient is computed. It must be a subset of the input labels.
         If ``None``, all input variables are considered. Default is ``None``.
+    :type d: str | list[str]
     :raises TypeError: If the input tensor is not a LabelTensor.
     :raises RuntimeError: If the output is a scalar field and the components
         are not equal to the output labels.
@@ -50,9 +52,10 @@ def grad(output_, input_, components=None, d=None):
             computed. It must be a column tensor.
         :param LabelTensor input_: The input tensor with respect to which the
             gradient is computed.
-        :param list[str] d: The names of the input variables with respect to
+        :param d: The names of the input variables with respect to
             which the gradient is computed. It must be a subset of the input
             labels. If ``None``, all input variables are considered.
+        :type d: str | list[str]
         :raises RuntimeError: If a vectorial function is passed.
         :raises RuntimeError: If missing derivative labels.
         :return: The computed gradient tensor.
@@ -89,6 +92,12 @@ def grad(output_, input_, components=None, d=None):
     if components is None:
         components = output_.labels
 
+    if not isinstance(components, list):
+        components = [components]
+
+    if not isinstance(d, list):
+        d = [d]
+
     if output_.shape[1] == 1:  # scalar output ################################
 
         if components != output_.labels:
@@ -120,12 +129,14 @@ def div(output_, input_, components=None, d=None):
         computed.
     :param LabelTensor input_: The input tensor with respect to which the
         divergence is computed.
-    :param list[str] components: The names of the output variables for which to
+    :param components: The names of the output variables for which to
         compute the divergence. 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
+    :type components: str | list[str]
+    :param d: The names of the input variables with respect to which
         the divergence is computed. It must be a subset of the input labels.
         If ``None``, all input variables are considered. Default is ``None``.
+    :type d: str | list[str]
     :raises TypeError: If the input tensor is not a LabelTensor.
     :raises ValueError: If the output is a scalar field.
     :raises ValueError: If the number of components is not equal to the number
@@ -142,6 +153,12 @@ def div(output_, input_, components=None, d=None):
     if components is None:
         components = output_.labels
 
+    if not isinstance(components, list):
+        components = [components]
+
+    if not isinstance(d, list):
+        d = [d]
+
     if output_.shape[1] < 2 or len(components) < 2:
         raise ValueError("div supported only for vector fields")
 
@@ -170,12 +187,14 @@ def laplacian(output_, input_, components=None, d=None, method="std"):
         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
+    :param 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
+    :type components: str | list[str]
+    :param 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``.
+    :type d: str | list[str]
     :param str method: The method used to compute the Laplacian. Default is
         ``std``.
     :raises NotImplementedError: If ``std=divgrad``.
@@ -191,12 +210,14 @@ def laplacian(output_, input_, components=None, d=None, method="std"):
             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
+        :param 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
+        :type components: str | list[str]
+        :param 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.
+        :type d: str | list[str]
         :return: The computed laplacian tensor.
         :rtype: LabelTensor
         """
@@ -216,22 +237,24 @@ def laplacian(output_, input_, components=None, d=None, method="std"):
     if components is None:
         components = output_.labels
 
+    if not isinstance(components, list):
+        components = [components]
+
+    if not isinstance(d, list):
+        d = [d]
+
     if method == "divgrad":
         raise NotImplementedError("divgrad not implemented as method")
 
     if method == "std":
-        if len(components) == 1:
-            result = scalar_laplace(output_, input_, components, d)
-            labels = [f"dd{components[0]}"]
-
-        else:
-            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 = 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
@@ -251,12 +274,14 @@ def advection(output_, input_, velocity_field, components=None, d=None):
         is computed.
     :param str velocity_field: The name of the output variable used as velocity
         field. It must be chosen among the output labels.
-    :param list[str] components: The names of the output variables for which
+    :param components: The names of the output variables for which
         to compute the advection. 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
+    :type components: str | list[str]
+    :param d: The names of the input variables with respect to which
         the advection is computed. It must be a subset of the input labels.
         If ``None``, all input variables are considered. Default is ``None``.
+    :type d: str | list[str]
     :return: The computed advection tensor.
     :rtype: LabelTensor
     """
@@ -266,6 +291,12 @@ def advection(output_, input_, velocity_field, components=None, d=None):
     if components is None:
         components = output_.labels
 
+    if not isinstance(components, list):
+        components = [components]
+
+    if not isinstance(d, list):
+        d = [d]
+
     tmp = (
         grad(output_, input_, components, d)
         .reshape(-1, len(components), len(d))


Test patch:
diff --git a/tests/test_operator.py b/tests/test_operator.py
index e274fda..83a41a0 100644
--- a/tests/test_operator.py
+++ b/tests/test_operator.py
@@ -164,3 +164,42 @@ def test_laplacian_vector_output2():
 
     assert torch.allclose(lap_f.extract("ddu"), lap_u)
     assert torch.allclose(lap_f.extract("ddv"), lap_v)
+
+
+def test_label_format():
+    # Testing the format of `components` or `d` in case of single str of length
+    # greater than 1; e.g.: "aaa".
+    # This test is conducted only for gradient and laplacian, since div is not
+    # implemented for single components.
+    inp.labels = ["xx", "yy", "zz"]
+    tensor_v = LabelTensor(func_vector(inp), ["aa", "bbb", "c"])
+    comp = tensor_v.labels[0]
+    single_d = inp.labels[0]
+
+    # Single component as string + list of d
+    grad_tensor_v = grad(tensor_v, inp, components=comp, d=None)
+    assert grad_tensor_v.labels == [f"d{comp}d{i}" for i in inp.labels]
+
+    lap_tensor_v = laplacian(tensor_v, inp, components=comp, d=None)
+    assert lap_tensor_v.labels == [f"dd{comp}"]
+
+    # Single component as list + list of d
+    grad_tensor_v = grad(tensor_v, inp, components=[comp], d=None)
+    assert grad_tensor_v.labels == [f"d{comp}d{i}" for i in inp.labels]
+
+    lap_tensor_v = laplacian(tensor_v, inp, components=[comp], d=None)
+    assert lap_tensor_v.labels == [f"dd{comp}"]
+
+    # List of components + single d as string
+    grad_tensor_v = grad(tensor_v, inp, components=None, d=single_d)
+    assert grad_tensor_v.labels == [f"d{i}d{single_d}" for i in tensor_v.labels]
+
+    lap_tensor_v = laplacian(tensor_v, inp, components=None, d=single_d)
+    assert lap_tensor_v.labels == [f"dd{i}" for i in tensor_v.labels]
+
+    # List of components + single d as list
+    grad_tensor_v = grad(tensor_v, inp, components=None, d=[single_d])
+    assert grad_tensor_v.labels == [f"d{i}d{single_d}" for i in tensor_v.labels]
+
+    lap_tensor_v = laplacian(tensor_v, inp, components=None, d=[single_d])
+    assert lap_tensor_v.labels == [f"dd{i}" for i in tensor_v.labels]