a27bd35443
Refactoring for 0.2 * Data module, data loader and dataset * Refactor LabelTensor * Refactor solvers Co-authored-by: dario-coscia <dariocos99@gmail.com>
521 lines
18 KiB
Python
521 lines
18 KiB
Python
""" Module for LabelTensor """
|
|
from copy import copy, deepcopy
|
|
import torch
|
|
from torch import Tensor
|
|
|
|
|
|
full_labels = True
|
|
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)
|
|
|
|
@property
|
|
def tensor(self):
|
|
return self.as_subclass(Tensor)
|
|
|
|
def __init__(self, x, labels, **kwargs):
|
|
"""
|
|
Construct a `LabelTensor` by passing a dict of the labels
|
|
|
|
:Example:
|
|
>>> from pina import LabelTensor
|
|
>>> tensor = LabelTensor(
|
|
>>> torch.rand((2000, 3)),
|
|
{1: {"name": "space"['a', 'b', 'c'])
|
|
|
|
"""
|
|
self.full = kwargs.get('full', full_labels)
|
|
if labels is not None:
|
|
self.labels = labels
|
|
else:
|
|
self._labels = {}
|
|
|
|
@property
|
|
def labels(self):
|
|
"""Property decorator for labels
|
|
|
|
:return: labels of self
|
|
:rtype: list
|
|
"""
|
|
if self.ndim - 1 in self._labels.keys():
|
|
return self._labels[self.ndim - 1]['dof']
|
|
|
|
@property
|
|
def full_labels(self):
|
|
"""Property decorator for labels
|
|
|
|
:return: labels of self
|
|
:rtype: list
|
|
"""
|
|
to_return_dict = {}
|
|
shape_tensor = self.shape
|
|
for i in range(len(shape_tensor)):
|
|
if i in self._labels.keys():
|
|
to_return_dict[i] = self._labels[i]
|
|
else:
|
|
to_return_dict[i] = {'dof': range(shape_tensor[i]), 'name': i}
|
|
return to_return_dict
|
|
|
|
@property
|
|
def stored_labels(self):
|
|
"""Property decorator for labels
|
|
|
|
:return: labels of self
|
|
:rtype: list
|
|
"""
|
|
return self._labels
|
|
|
|
@labels.setter
|
|
def labels(self, labels):
|
|
""""
|
|
Set properly the parameter _labels
|
|
|
|
:param labels: Labels to assign to the class variable _labels.
|
|
:type: labels: str | list(str) | dict
|
|
"""
|
|
if not hasattr(self, '_labels'):
|
|
self._labels = {}
|
|
if isinstance(labels, dict):
|
|
self._init_labels_from_dict(labels)
|
|
elif isinstance(labels, (list, range)):
|
|
self._init_labels_from_list(labels)
|
|
elif isinstance(labels, str):
|
|
labels = [labels]
|
|
self._init_labels_from_list(labels)
|
|
else:
|
|
raise ValueError("labels must be list, dict or string.")
|
|
|
|
def _init_labels_from_dict(self, labels):
|
|
"""
|
|
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
|
|
does not match with tensor shape
|
|
"""
|
|
tensor_shape = self.shape
|
|
|
|
# Set all labels if full_labels is True
|
|
if hasattr(self, 'full') and self.full:
|
|
labels = {
|
|
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]:
|
|
raise ValueError(
|
|
'Number of dof does not match tensor shape')
|
|
else:
|
|
raise ValueError('Illegal labels initialization')
|
|
# Assign labels values
|
|
self._labels[k] = v
|
|
|
|
def _init_labels_from_list(self, labels):
|
|
"""
|
|
Given a list of dof, this method update the internal label
|
|
representation
|
|
|
|
:param labels: The label(s) to update.
|
|
:type labels: list
|
|
"""
|
|
# Create a dict with labels
|
|
last_dim_labels = {
|
|
self.ndim - 1: {
|
|
'dof': labels,
|
|
'name': self.ndim - 1
|
|
}
|
|
}
|
|
self._init_labels_from_dict(last_dim_labels)
|
|
|
|
def extract(self, labels_to_extract):
|
|
"""
|
|
Extract the subset of the original tensor by returning all the columns
|
|
corresponding to the passed ``label_to_extract``.
|
|
|
|
:param labels_to_extract: The label(s) to extract.
|
|
:type labels_to_extract: str | list(str) | tuple(str)
|
|
: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
|
|
|
|
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
|
|
# means that we want to extract a values from the last dimension)
|
|
if isinstance(labels_to_extract, (tuple, list)):
|
|
if not ndim - 1 in stored_keys:
|
|
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
|
|
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')
|
|
|
|
dim_labels = labels[idx_dim]['dof']
|
|
labels_te = [labels_te] if isinstance(labels_te, (int, str)) else 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}})
|
|
|
|
tensor = super().__getitem__(extractor).as_subclass(LabelTensor)
|
|
tensor._labels = labels
|
|
return tensor
|
|
|
|
def __str__(self):
|
|
"""
|
|
returns a string with the representation of the class
|
|
"""
|
|
s = ''
|
|
for key, value in self._labels.items():
|
|
s += f"{key}: {value}\n"
|
|
s += '\n'
|
|
s += self.tensor.__str__()
|
|
return s
|
|
|
|
@staticmethod
|
|
def cat(tensors, dim=0):
|
|
"""
|
|
Stack a list of tensors. For example, given a tensor `a` of shape
|
|
`(n,m,dof)` and a tensor `b` of dimension `(n',m,dof)`
|
|
the resulting tensor is of shape `(n+n',m,dof)`
|
|
|
|
:param tensors: tensors to concatenate
|
|
:type tensors: list(LabelTensor)
|
|
:param dim: dimensions on which you want to perform the operation
|
|
(default 0)
|
|
:type dim: int
|
|
:rtype: LabelTensor
|
|
:raises ValueError: either number dof or dimensions names differ
|
|
"""
|
|
if len(tensors) == 0:
|
|
return []
|
|
if len(tensors) == 1 or isinstance(tensors, LabelTensor):
|
|
return tensors[0]
|
|
|
|
# Perform cat on tensors
|
|
new_tensor = torch.cat(tensors, dim=dim)
|
|
|
|
# --------- Start definition auxiliary function ------
|
|
# Compute and update labels
|
|
def create_labels_cat(tensors, dim, tensor_shape):
|
|
stored_labels = [tensor.stored_labels for tensor in tensors]
|
|
keys = stored_labels[0].keys()
|
|
|
|
if any(not all(stored_labels[i][k] == stored_labels[0][k] for i in
|
|
range(len(stored_labels))) for k in keys if k != dim):
|
|
raise RuntimeError('tensors must have the same shape and dof')
|
|
|
|
# Copy labels from the first tensor and update the 'dof' for dimension `dim`
|
|
labels = copy(stored_labels[0])
|
|
if dim in labels:
|
|
labels_list = [tensor[dim]['dof'] for tensor in stored_labels]
|
|
last_dim_dof = range(tensor_shape[dim]) if all(isinstance(label, range)
|
|
for label in labels_list) else sum(labels_list, [])
|
|
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
|
|
new_tensor._labels = labels
|
|
return new_tensor
|
|
|
|
@staticmethod
|
|
def stack(tensors):
|
|
new_tensor = torch.stack(tensors)
|
|
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):
|
|
lt = super().requires_grad_(mode)
|
|
lt._labels = self._labels
|
|
return lt
|
|
|
|
@property
|
|
def dtype(self):
|
|
return super().dtype
|
|
|
|
def to(self, *args, **kwargs):
|
|
"""
|
|
Performs Tensor dtype and/or device conversion. For more details, see
|
|
:meth:`torch.Tensor.to`.
|
|
"""
|
|
lt = super().to(*args, **kwargs)
|
|
lt._labels = self._labels
|
|
return lt
|
|
|
|
def clone(self, *args, **kwargs):
|
|
"""
|
|
Clone the LabelTensor. For more details, see
|
|
:meth:`torch.Tensor.clone`.
|
|
|
|
:return: A copy of the tensor.
|
|
:rtype: LabelTensor
|
|
"""
|
|
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
|
|
new_label_tensor = LabelTensor.cat([self, tensor],
|
|
dim=self.ndim - 1)
|
|
elif mode == 'cross':
|
|
# Crete tensor and call cat on last dimension
|
|
tensor1 = self
|
|
tensor2 = tensor
|
|
n1 = tensor1.shape[0]
|
|
n2 = tensor2.shape[0]
|
|
tensor1 = LabelTensor(tensor1.repeat(n2, 1), labels=tensor1.labels)
|
|
tensor2 = LabelTensor(tensor2.repeat_interleave(n1, dim=0),
|
|
labels=tensor2.labels)
|
|
new_label_tensor = LabelTensor.cat([tensor1, tensor2],
|
|
dim=self.ndim - 1)
|
|
else:
|
|
raise ValueError('mode must be either "std" or "cross"')
|
|
return new_label_tensor
|
|
|
|
@staticmethod
|
|
def vstack(label_tensors):
|
|
"""
|
|
Stack tensors vertically. For more details, see
|
|
:meth:`torch.vstack`.
|
|
|
|
:param list(LabelTensor) label_tensors: the tensors to stack. They need
|
|
to have equal labels.
|
|
:return: the stacked tensor
|
|
:rtype: LabelTensor
|
|
"""
|
|
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
|
|
:param old_labels: labels from which retrieve data
|
|
:param to_update_labels: labels to update
|
|
:param index: index of dof to retain
|
|
:param dim: label index
|
|
:return:
|
|
"""
|
|
old_dof = old_labels[to_update_dim]['dof']
|
|
if isinstance(index, slice):
|
|
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 -----
|
|
|
|
|
|
def __getitem__(self, index):
|
|
"""
|
|
TODO: Complete docstring
|
|
:param index:
|
|
:return:
|
|
"""
|
|
# 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
|
|
selected_lt = super().__getitem__(index)
|
|
stored_labels = self._labels
|
|
labels = copy(stored_labels)
|
|
|
|
# Put here because it is the most common case (int as index).
|
|
# Used by DataLoader -> put here for efficiency purpose
|
|
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
|
|
if index[0] == Ellipsis:
|
|
if len(self.shape) in labels:
|
|
self._update_single_label(stored_labels, labels, index, 0, 0)
|
|
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:
|
|
value for key, value in labels.items()}
|
|
continue
|
|
i += 1
|
|
selected_lt._labels = labels
|
|
return selected_lt
|
|
|
|
def sort_labels(self, dim=None):
|
|
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)
|
|
indexer = [slice(None)] * self.ndim
|
|
indexer[dim] = sorted_index
|
|
return self.__getitem__(tuple(indexer))
|
|
|
|
def __deepcopy__(self, memo):
|
|
cls = self.__class__
|
|
result = cls(deepcopy(self.tensor), deepcopy(self.stored_labels))
|
|
return result
|
|
|
|
def permute(self, *dims):
|
|
tensor = super().permute(*dims)
|
|
labels = self._labels
|
|
keys_list = list(*dims)
|
|
labels = {
|
|
keys_list.index(k): labels[k]
|
|
for k in labels.keys()
|
|
}
|
|
tensor._labels = labels
|
|
return tensor
|
|
|
|
def detach(self):
|
|
lt = super().detach()
|
|
lt._labels = self.stored_labels
|
|
return lt |