Simplify LabelTensor class, fix #395, add docstrings, and resolve Python 3.8 compatibility issue in tests

pina/label_tensor.py

@@ -4,18 +4,13 @@ import torch
 from torch import Tensor
 
 
-full_labels = False
-MATH_FUNCTIONS = {torch.sin, torch.cos}
-
 class LabelTensor(torch.Tensor):
     """Torch tensor with a label for any column."""
 
     @staticmethod
     def __new__(cls, x, labels, *args, **kwargs):
-        full = kwargs.pop("full", full_labels)
 
         if isinstance(x, LabelTensor):
-            x.full = full
             return x
         return super().__new__(cls, x, *args, **kwargs)
 
@@ -34,7 +29,7 @@ class LabelTensor(torch.Tensor):
                     {1: {"name": "space"['a', 'b', 'c'])
 
         """
-        self.full = kwargs.get('full', full_labels)
+        super().__init__()
         if labels is not None:
             self.labels = labels
         else:
@@ -95,51 +90,49 @@ class LabelTensor(torch.Tensor):
         else:
             raise ValueError("labels must be list, dict or string.")
 
-    def _init_labels_from_dict(self, labels):
+    def _init_labels_from_dict(self, labels: dict):
         """
         Update the internal label representation according to the values
         passed as input.
 
         :param labels: The label(s) to update.
         :type labels: dict
-            :raises ValueError: dof list contain duplicates or number of dof
+        :raises ValueError: If the dof list contains duplicates or the number of
-            does not match with tensor shape
+                            dof does not match the tensor shape.
         """
         tensor_shape = self.shape
 
-        # Set all labels if full_labels is True
+        def validate_dof(dof_list, dim_size: int):
-        if hasattr(self, 'full') and self.full:
+            """Validate the 'dof' list for uniqueness and size."""
-            labels = {
+            if len(dof_list) != len(set(dof_list)):
-                i: labels[i] if i in labels else {
-                    'name': i, 'dof': range(tensor_shape[i])
-                }
-                for i in range(len(tensor_shape))
-            }
-
-        for k, v in labels.items():
-
-            # Init labels from str
-            if isinstance(v, str):
-                v = {'name': v, 'dof': range(tensor_shape[k])}
-
-            # Init labels from dict
-            elif isinstance(v, dict):
-                # Only name of the dimension if provided
-                if list(v.keys()) == ['name']:
-                    v['dof'] = range(tensor_shape[k])
-                # Both name and dof are provided
-                elif sorted(list(v.keys())) == ['dof', 'name']:
-                    dof_list = v['dof']
-                    dof_len = len(dof_list)
-                    if dof_len != len(set(dof_list)):
                 raise ValueError("dof must be unique")
-                    if dof_len != tensor_shape[k]:
+            if len(dof_list) != dim_size:
                 raise ValueError(
-                            'Number of dof does not match tensor shape')
+                    f"Number of dof ({len(dof_list)}) does not match "
+                    f"tensor shape ({dim_size})")
+
+        for dim, label in labels.items():
+            if isinstance(label, dict):
+                if 'name' not in label:
+                    label['name'] = dim
+                if 'dof' not in label:
+                    label['dof'] = range(tensor_shape[dim])
+                if 'dof' in label and 'name' in label:
+                    dof = label['dof']
+                    dof_list = dof if isinstance(dof, (list, range)) else [dof]
+                    if not isinstance(dof_list, (list, range)):
+                        raise ValueError(f"'dof' should be a list or range, not"
+                                         f" {type(dof_list)}")
+                    validate_dof(dof_list, tensor_shape[dim])
                 else:
-                raise ValueError('Illegal labels initialization')
+                    raise ValueError("Labels dictionary must contain either "
-            # Assign labels values
+                                     " both 'name' and 'dof' keys")
-            self._labels[k] = v
+            else:
+                raise ValueError(f"Invalid label format for {dim}: Expected "
+                                 f"list or dictionary, got {type(label)}")
+
+            # Assign validated label data to internal labels
+            self._labels[dim] = label
 
     def _init_labels_from_list(self, labels):
         """
@@ -168,61 +161,54 @@ class LabelTensor(torch.Tensor):
         :raises TypeError: Labels are not ``str``.
         :raises ValueError: Label to extract is not in the labels ``list``.
         """
-        # Convert str/int to string
-        def find_names(labels):
-            dim_names = {}
-            for dim in labels.keys():
-                dim_names[labels[dim]['name']] = dim
-            return dim_names
 
+        def get_label_indices(dim_labels, labels_te):
+            if isinstance(labels_te, (int, str)):
+                labels_te = [labels_te]
+            return [dim_labels.index(label) for label in labels_te] if len(
+                labels_te) > 1 else slice(dim_labels.index(labels_te[0]),
+                                          dim_labels.index(labels_te[0]) + 1)
+
+        # Ensure labels_to_extract is a list or dict
         if isinstance(labels_to_extract, (str, int)):
             labels_to_extract = [labels_to_extract]
 
-        # Store useful variables
         labels = copy(self._labels)
-        stored_keys = labels.keys()
-        dim_names = find_names(labels)
-        ndim = len(super().shape)
 
-        # Convert tuple/list to dict (having a list as input
+        # Get the dimension names and the respective dimension index
-        # means that we want to extract a values from the last dimension)
+        dim_names = {labels[dim]['name']: dim for dim in labels.keys()}
-        if isinstance(labels_to_extract, (tuple, list)):
+        ndim = super().ndim
-            if not ndim - 1 in stored_keys:
+        tensor = self.tensor.as_subclass(torch.Tensor)
-                raise ValueError(
-                    "LabelTensor does not have labels in last dimension")
-            name = labels[ndim-1]['name']
-            labels_to_extract = {name: list(labels_to_extract)}
 
-        # If labels_to_extract is not dict then rise error
+        # Convert list/tuple to a dict for the last dimension if applicable
+        if isinstance(labels_to_extract, (list, tuple)):
+            last_dim = ndim - 1
+            dim_name = labels[last_dim]['name']
+            labels_to_extract = {dim_name: list(labels_to_extract)}
+
+        # Validate the labels_to_extract type
         if not isinstance(labels_to_extract, dict):
-            raise ValueError('labels_to_extract must be str or list or dict')
-
-        # Initialize list used to perform extraction
-        extractor = [slice(None)]*ndim
-
-        # Loop over labels_to_extract dict
-        for dim_name, labels_te in labels_to_extract.items():
-
-            # If label is not find raise value error
-            idx_dim = dim_names.get(dim_name, None)
-            if idx_dim is None:
             raise ValueError(
-                    'Cannot extract label with is not in original labels')
+                "labels_to_extract must be a string, list, or dictionary.")
 
+        # Perform the extraction for each specified dimension
+        for dim_name, labels_te in labels_to_extract.items():
+            if dim_name not in dim_names:
+                raise ValueError(
+                    f"Cannot extract labels for dimension '{dim_name}' as it is"
+                    f" not present in the original labels.")
+
+            idx_dim = dim_names[dim_name]
             dim_labels = labels[idx_dim]['dof']
-            labels_te = [labels_te] if isinstance(labels_te, (int, str)) else labels_te
+            indices = get_label_indices(dim_labels, labels_te)
-            if not isinstance(labels_te, range):
-                #If is done to keep the dimension if there is only one extracted label
-                extractor[idx_dim] = [dim_labels.index(i) for i in labels_te] \
-                    if len(labels_te)>1 else slice(dim_labels.index(labels_te[0]), dim_labels.index(labels_te[0])+1)
-            else:
-                extractor[idx_dim] = slice(labels_te.start, labels_te.stop)
 
-            labels.update({idx_dim: {'dof': labels_te, 'name': dim_name}})
+            extractor = [slice(None)] * ndim
+            extractor[idx_dim] = indices
+            tensor = tensor[tuple(extractor)]
 
-        tensor = super().__getitem__(extractor).as_subclass(LabelTensor)
+            labels[idx_dim] = {'dof': labels_te, 'name': dim_name}
-        tensor._labels = labels
+
-        return tensor
+        return LabelTensor(tensor, labels)
 
     def __str__(self):
         """
@@ -243,62 +229,82 @@ class LabelTensor(torch.Tensor):
         the resulting tensor is of shape `(n+n',m,dof)`
 
         :param tensors: tensors to concatenate
-        :type tensors: list(LabelTensor)
+        :type tensors: list of LabelTensor
         :param dim: dimensions on which you want to perform the operation
-        (default 0)
+        (default is 0)
         :type dim: int
         :rtype: LabelTensor
         :raises ValueError: either number dof or dimensions names differ
         """
-        if len(tensors) == 0:
+        if not tensors:
-            return []
+            return []  # Handle empty list
-        if len(tensors) == 1 or isinstance(tensors, LabelTensor):
+        if len(tensors) == 1:
-            return tensors[0]
+            return tensors[0]  # Return single tensor as-is
 
-        # Perform cat on tensors
+            # Perform concatenation
-        new_tensor = torch.cat(tensors, dim=dim)
+        cat_tensor = torch.cat(tensors, dim=dim)
+        tensors_labels = [tensor.stored_labels for tensor in tensors]
 
-        # --------- Start definition auxiliary function ------
+        # Check label consistency across tensors, excluding the
-        # Compute and update labels
+        # concatenation dimension
-        def create_labels_cat(tensors, dim, tensor_shape):
+        for key in tensors_labels[0].keys():
-            stored_labels = [tensor.stored_labels for tensor in tensors]
+            if key != dim:
-            keys = stored_labels[0].keys()
+                if any(tensors_labels[i][key] != tensors_labels[0][key]
+                       for i in range(len(tensors_labels))):
+                    raise RuntimeError(
+                        f"Tensors must have the same labels along all "
+                        f"dimensions except {dim}.")
 
-            if any(not all(stored_labels[i][k] == stored_labels[0][k] for i in
+        # Copy and update the 'dof' for the concatenation dimension
-                        range(len(stored_labels))) for k in keys if k != dim):
+        cat_labels = {k: copy(v) for k, v in tensors_labels[0].items()}
-                raise RuntimeError('tensors must have the same shape and dof')
 
-            # Copy labels from the first tensor and update the 'dof' for dimension `dim`
+        # Update labels if the concatenation dimension has labels
-            labels = copy(stored_labels[0])
+        if dim in tensors[0].stored_labels:
-            if dim in labels:
+            if dim in cat_labels:
-                labels_list = [tensor[dim]['dof'] for tensor in stored_labels]
+                cat_dofs = [label[dim]['dof'] for label in
-                last_dim_dof = range(tensor_shape[dim]) if all(isinstance(label, range)
+                            tensors_labels]
-                                    for label in labels_list) else sum(labels_list, [])
+                cat_labels[dim]['dof'] = sum(cat_dofs, [])
-                labels[dim]['dof'] = last_dim_dof
-            return labels
-        # --------- End definition auxiliary function ------
-
-        # Update labels
-        if dim in tensors[0].stored_labels.keys():
-            new_tensor_shape = new_tensor.shape
-            labels = create_labels_cat(tensors, dim, new_tensor_shape)
         else:
-            labels = tensors[0].stored_labels
+            cat_labels = tensors[0].stored_labels
-        new_tensor._labels = labels
+
-        return new_tensor
+        # Assign updated labels to the concatenated tensor
+        cat_tensor._labels = cat_labels
+        return cat_tensor
 
     @staticmethod
     def stack(tensors):
+        """
+        Stacks a list of tensors along a new dimension.
+
+        :param tensors: A list of tensors to stack. All tensors must have the
+                        same shape.
+        :type tensors: list of LabelTensor
+        :return: A new tensor obtained by stacking the input tensors,
+                 with the updated labels.
+        :rtype: LabelTensor
+        """
+        # Perform stacking in torch
         new_tensor = torch.stack(tensors)
+
+        # Increase labels keys by 1
         labels = tensors[0]._labels
         labels = {key + 1: value for key, value in labels.items()}
-        if full_labels:
-            new_tensor.labels = labels
-        else:
         new_tensor._labels = labels
         return new_tensor
 
     def requires_grad_(self, mode=True):
+        """
+        Override the requires_grad_ method to update the labels in the new
+        tensor.
+
+        :param mode: A boolean value indicating whether the tensor should track
+                     gradients.If `True`, the tensor will track gradients; if `
+                     False`, it will not.
+        :type mode: bool, optional (default is `True`)
+        :return: The tensor itself with the updated `requires_grad` state and
+                 retained labels.
+        :rtype: LabelTensor
+        """
         lt = super().requires_grad_(mode)
         lt._labels = self._labels
         return lt
@@ -324,39 +330,10 @@ class LabelTensor(torch.Tensor):
         :return: A copy of the tensor.
         :rtype: LabelTensor
         """
-        out = LabelTensor(super().clone(*args, **kwargs), deepcopy(self._labels))
+        out = LabelTensor(super().clone(*args, **kwargs),
+                          deepcopy(self._labels))
         return out
 
-    @staticmethod
-    def summation(tensors):
-        if len(tensors) == 0:
-            raise ValueError('tensors list must not be empty')
-        if len(tensors) == 1:
-            return tensors[0]
-        # Collect all labels
-
-        # Check labels of all the tensors in each dimension
-        if not all(tensor.shape == tensors[0].shape for tensor in tensors) or \
-                not all(tensor.full_labels[i] == tensors[0].full_labels[i] for
-                        tensor in tensors for i in range(tensors[0].ndim - 1)):
-            raise RuntimeError('Tensors must have the same shape and labels')
-
-        last_dim_labels = []
-        data = torch.zeros(tensors[0].tensor.shape).to(tensors[0].device)
-        for tensor in tensors:
-            data += tensor.tensor
-            last_dim_labels.append(tensor.labels)
-
-        last_dim_labels = ['+'.join(items) for items in zip(*last_dim_labels)]
-        labels = {k: copy(v) for k, v in tensors[0].stored_labels.items()}
-        labels.update({
-            tensors[0].ndim - 1: {
-                'dof': last_dim_labels,
-                'name': tensors[0].name
-            }
-        })
-        return LabelTensor(data, labels)
-
     def append(self, tensor, mode='std'):
         if mode == 'std':
             # Call cat on last dimension
@@ -390,12 +367,11 @@ class LabelTensor(torch.Tensor):
         """
         return LabelTensor.cat(label_tensors, dim=0)
 
-        # ---------------------- Start auxiliary function definition -----
     # This method is used to update labels
     def _update_single_label(self, old_labels, to_update_labels, index, dim,
                              to_update_dim):
         """
-            TODO
+        Update the labels of the tensor by selecting only the labels
         :param old_labels: labels from which retrieve data
         :param to_update_labels: labels to update
         :param index: index of dof to retain
@@ -403,119 +379,198 @@ class LabelTensor(torch.Tensor):
         :return:
         """
         old_dof = old_labels[to_update_dim]['dof']
-            if isinstance(index, slice):
+        label_name = old_labels[dim]['name']
-                to_update_labels.update({
-                    dim: {
-                        'dof': old_dof[index],
-                        'name': old_labels[dim]['name']
-                    }
-                })
-                return
-            if isinstance(index, int):
-                index = [index]
-            if isinstance(index, (list, torch.Tensor)):
-                to_update_labels.update({
-                    dim: {
-                        'dof': [old_dof[i] for i in index] if isinstance(old_dof, list) else index,
-                        'name': old_labels[dim]['name']
-                    }
-                })
-                return
-            raise NotImplementedError(f'Getitem not implemented for '
-                                      f'{type(index)} values')
-        # ---------------------- End auxiliary function definition -----
 
+        if isinstance(index, slice):
+            # Handle slicing
+            to_update_labels[dim] = {'dof': old_dof[index], 'name': label_name}
+        elif isinstance(index, int):
+            # Handle single integer index
+            to_update_labels[dim] = {'dof': [old_dof[index]],
+                                     'name': label_name}
+        elif isinstance(index, (list, torch.Tensor)):
+            # Handle lists or tensors
+            indices = [index] if isinstance(index, (int, str)) else index
+            to_update_labels[dim] = {
+                'dof': [old_dof[i] for i in indices] if isinstance(old_dof,
+                                                                   list) else indices,
+                'name': label_name
+            }
+        else:
+            raise NotImplementedError(
+                f"Unsupported index type: {type(index)}. Expected slice, int, "
+                f"list, or torch.Tensor."
+            )
 
     def __getitem__(self, index):
+        """"
+        Override the __getitem__ method to handle the labels of the tensor.
+        Perform the __getitem__ operation on the tensor and update the labels.
+
+        :param index: The index used to access the item
+        :type index: Union[int, str, tuple, list]
+        :return: A tensor-like object with updated labels.
+        :rtype: LabelTensor
+        :raises KeyError: If an invalid label index is provided.
+        :raises IndexError: If an invalid index is accessed in the tensor.
         """
-        TODO: Complete docstring
+        # Handle string index
-        :param index:
+        if isinstance(index, str) or (isinstance(index, (tuple, list)) and all(
-        :return:
+                isinstance(i, str) for i in index)):
-        """
-        # Index are str --> call extract
-        if isinstance(index, str) or (isinstance(index, (tuple, list))
-                                      and all(
-                    isinstance(a, str) for a in index)):
             return self.extract(index)
 
-        # Store important variables
+        # Retrieve selected tensor and labels
-        selected_lt = super().__getitem__(index)
+        selected_tensor = super().__getitem__(index)
-        stored_labels = self._labels
+        original_labels = self._labels
-        labels = copy(stored_labels)
+        updated_labels = copy(original_labels)
 
-        # Put here because it is the most common case (int as index).
+        # Ensure the index is iterable
-        # Used by DataLoader -> put here for efficiency purpose
+        if not isinstance(index, tuple):
-        if isinstance(index, list):
-            if 0 in labels.keys():
-                self._update_single_label(stored_labels, labels, index,
-                                          0, 0)
-            selected_lt._labels = labels
-            return selected_lt
-
-        if isinstance(index, int):
-            labels.pop(0, None)
-            labels = {key - 1 if key > 0 else key: value for key, value in
-                      labels.items()}
-            selected_lt._labels = labels
-            return selected_lt
-
-        if not isinstance(index, (tuple, torch.Tensor)):
             index = [index]
 
-        # Ellipsis are used to perform operation on the last dimension
+        # Update labels based on the index
-        if index[0] == Ellipsis:
+        offset = 0
-            if len(self.shape) in labels:
+        for dim, idx in enumerate(index):
-                self._update_single_label(stored_labels, labels, index, 0, 0)
+            if dim in self.stored_labels.keys():
-            selected_lt._labels = labels
-            return selected_lt
-
-        i = 0
-        for j, idx in enumerate(index):
-            if j in self.stored_labels.keys():
-                if isinstance(idx, int) or (
-                        isinstance(idx, torch.Tensor) and idx.ndim == 0):
-                    selected_lt = selected_lt.unsqueeze(j)
-                if idx != slice(None):
-                    self._update_single_label(stored_labels, labels, idx, j, i)
-            else:
                 if isinstance(idx, int):
-                    labels = {key - 1 if key > j else key:
+                    selected_tensor = selected_tensor.unsqueeze(dim)
-                                  value for key, value in labels.items()}
+                if idx != slice(None):
+                    self._update_single_label(original_labels, updated_labels,
+                                              idx, dim, offset)
+            else:
+                # Adjust label keys if dimension is reduced (case of integer
+                # index on a non-labeled dimension)
+                if isinstance(idx, int):
+                    updated_labels = {
+                        key - 1 if key > dim else key: value
+                        for key, value in updated_labels.items()
+                    }
                     continue
-            i += 1
+            offset += 1
-        selected_lt._labels = labels
+
-        return selected_lt
+        # Update the selected tensor's labels
+        selected_tensor._labels = updated_labels
+        return selected_tensor
 
     def sort_labels(self, dim=None):
+        """
+        Sorts the labels along a specified dimension and returns a new tensor
+        with sorted labels.
+
+        :param dim: The dimension along which to sort the labels. If `None`,
+                    the last dimension (`ndim - 1`) is used.
+        :type dim: int, optional
+        :return: A new tensor with sorted labels along the specified dimension.
+        :rtype: LabelTensor
+        """
+
         def arg_sort(lst):
             return sorted(range(len(lst)), key=lambda x: lst[x])
+
         if dim is None:
             dim = self.ndim - 1
         if self.shape[dim] == 1:
             return self
         labels = self.stored_labels[dim]['dof']
         sorted_index = arg_sort(labels)
+        # Define an indexer to sort the tensor along the specified dimension
         indexer = [slice(None)] * self.ndim
+        # Assigned the sorted index to the specified dimension
         indexer[dim] = sorted_index
         return self.__getitem__(tuple(indexer))
 
     def __deepcopy__(self, memo):
+        """
+        Creates a deep copy of the object.
+
+        :param memo: LabelTensor object to be copied.
+        :type memo: LabelTensor
+        :return: A deep copy of the original LabelTensor object.
+        :rtype: LabelTensor
+        """
         cls = self.__class__
         result = cls(deepcopy(self.tensor), deepcopy(self.stored_labels))
         return result
 
     def permute(self, *dims):
+        """
+        Permutes the dimensions of the tensor and the associated labels
+        accordingly.
+
+        :param dims: The dimensions to permute the tensor to.
+        :type dims: tuple, list
+        :return: A new object with permuted dimensions and reordered labels.
+        :rtype: LabelTensor
+        """
+        # Call the base class permute method
         tensor = super().permute(*dims)
+
+        # Update lables
         labels = self._labels
         keys_list = list(*dims)
         labels = {
             keys_list.index(k): labels[k]
             for k in labels.keys()
         }
+
+        # Assign labels to the new tensor
         tensor._labels = labels
         return tensor
 
     def detach(self):
+        """
+        Detaches the tensor from the computation graph and retains the stored
+        labels.
+
+        :return: A new tensor detached from the computation graph.
+        :rtype: LabelTensor
+        """
         lt = super().detach()
+
+        # Copy the labels to the new tensor only if present
+        if hasattr(self, "_labels"):
             lt._labels = self.stored_labels
         return lt
+
+    @staticmethod
+    def summation(tensors):
+        """
+        Computes the summation of a list of tensors.
+
+        :param tensors: A list of tensors to sum. All tensors must have the same
+                        shape and labels.
+        :type tensors: list of LabelTensor
+        :return: A new `LabelTensor` containing the element-wise sum of the
+                 input tensors.
+        :rtype: LabelTensor
+        :raises ValueError: If the input `tensors` list is empty.
+        :raises RuntimeError: If the tensors have different shapes and/or
+                              mismatched labels.
+        """
+
+        if not tensors:
+            raise ValueError('The tensors list must not be empty.')
+
+        if len(tensors) == 1:
+            return tensors[0]
+
+        # Initialize result tensor and labels
+        data = torch.zeros_like(tensors[0].tensor).to(tensors[0].device)
+        last_dim_labels = []
+
+        # Accumulate tensors
+        for tensor in tensors:
+            data += tensor.tensor
+            last_dim_labels.append(tensor.labels)
+
+        # Construct last dimension labels
+        last_dim_labels = ['+'.join(items) for items in zip(*last_dim_labels)]
+
+        # Update the labels for the resulting tensor
+        labels = {k: copy(v) for k, v in tensors[0].stored_labels.items()}
+        labels[tensors[0].ndim - 1] = {
+            'dof': last_dim_labels,
+            'name': tensors[0].name
+        }
+
+        return LabelTensor(data, labels)

tests/test_label_tensor/test_label_tensor.py

@@ -7,7 +7,8 @@ data = torch.rand((20, 3))
 labels_column = {1: {"name": "space", "dof": ['x', 'y', 'z']}}
 labels_row = {0: {"name": "samples", "dof": range(20)}}
 labels_list = ['x', 'y', 'z']
-labels_all = labels_column | labels_row
+labels_all = labels_column.copy()
+labels_all.update(labels_row)
 
 
 @pytest.mark.parametrize("labels",