folivo/PINA
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

Improve efficiency and refact LabelTensor, codacy correction and fix bug in PinaBatch

Browse Source
This commit is contained in:
FilippoOlivo
2024-10-23 15:04:28 +02:00
committed by Nicola Demo
parent ccc5f5a322
commit ea3d1924e7
13 changed files with 496 additions and 395 deletions
Download Patch File Download Diff File Expand all files Collapse all files

8
pina/__init__.py
View File

@@ -1,6 +1,7 @@
__all__ = [
"PINN", "Trainer", "LabelTensor", "Plotter", "Condition",
"SamplePointDataset", "PinaDataModule", "PinaDataLoader"
"Trainer", "LabelTensor", "Plotter", "Condition",
"SamplePointDataset", "PinaDataModule", "PinaDataLoader",
'TorchOptimizer', 'Graph'
]
from .meta import *
@@ -12,3 +13,6 @@ from .condition.condition import Condition
from .data import SamplePointDataset
from .data import PinaDataModule
from .data import PinaDataLoader
from .optim import TorchOptimizer
from .optim import TorchScheduler
from .graph import Graph

8
pina/collector.py
View File

@@ -1,6 +1,3 @@
from sympy.strategies.branch import condition
from . import LabelTensor
from .utils import check_consistency, merge_tensors
@@ -16,6 +13,8 @@ class Collector:
# }
# those variables are used for the dataloading
self._data_collections = {name: {} for name in self.problem.conditions}
self.conditions_name = {i: name for i, name in
enumerate(self.problem.conditions)}
# variables used to check that all conditions are sampled
self._is_conditions_ready = {
@@ -101,7 +100,8 @@ class Collector:
"""
Add input points to a sampled condition
:param new_points_dict: Dictonary of input points (condition_name: LabelTensor)
:param new_points_dict: Dictonary of input points (condition_name:
LabelTensor)
:raises RuntimeError: if at least one condition is not already sampled
"""
for k, v in new_points_dict.items():

166
pina/data/base_dataset.py
View File

@@ -1,10 +1,12 @@
"""
Basic data module implementation
"""
from torch.utils.data import Dataset
import torch
import logging
from torch.utils.data import Dataset
from ..label_tensor import LabelTensor
from ..graph import Graph
class BaseDataset(Dataset):
@@ -12,10 +14,9 @@ class BaseDataset(Dataset):
BaseDataset class, which handle initialization and data retrieval
:var condition_indices: List of indices
:var device: torch.device
:var condition_names: dict of condition index and corresponding name
"""
def __new__(cls, problem, device):
def __new__(cls, problem=None, device=torch.device('cpu')):
"""
Ensure correct definition of __slots__ before initialization
:param AbstractProblem problem: The formulation of the problem.
@@ -30,7 +31,7 @@ class BaseDataset(Dataset):
'Something is wrong, __slots__ must be defined in subclasses.')
return object.__new__(cls)
def __init__(self, problem, device):
def __init__(self, problem=None, device=torch.device('cpu')):
""""
Initialize the object based on __slots__
:param AbstractProblem problem: The formulation of the problem.
@@ -38,79 +39,118 @@ class BaseDataset(Dataset):
dataset will be loaded.
"""
super().__init__()
self.condition_names = {}
collector = problem.collector
self.empty = True
self.problem = problem
self.device = device
self.condition_indices = None
for slot in self.__slots__:
setattr(self, slot, [])
num_el_per_condition = []
idx = 0
for name, data in collector.data_collections.items():
self.num_el_per_condition = []
self.conditions_idx = []
if self.problem is not None:
self._init_from_problem(self.problem.collector.data_collections)
self.initialized = False
def _init_from_problem(self, collector_dict):
"""
TODO
"""
for name, data in collector_dict.items():
keys = list(data.keys())
current_cond_num_el = None
if sorted(self.__slots__) == sorted(keys):
for slot in self.__slots__:
slot_data = data[slot]
if isinstance(slot_data, (LabelTensor, torch.Tensor,
Graph)):
if current_cond_num_el is None:
current_cond_num_el = len(slot_data)
elif current_cond_num_el != len(slot_data):
raise ValueError('Different number of conditions')
current_list = getattr(self, slot)
current_list += [data[slot]] if not (
isinstance(data[slot], list)) else data[slot]
num_el_per_condition.append(current_cond_num_el)
self.condition_names[idx] = name
idx += 1
if num_el_per_condition:
if set(self.__slots__) == set(keys):
self._populate_init_list(data)
idx = [key for key, val in
self.problem.collector.conditions_name.items() if
val == name]
self.conditions_idx.append(idx)
self.initialize()
def add_points(self, data_dict, condition_idx, batching_dim=0):
"""
This method filled internal lists of data points
:param data_dict: dictionary containing data points
:param condition_idx: index of the condition to which the data points
belong to
:param batching_dim: dimension of the batching
:raises: ValueError if the dataset has already been initialized
"""
if not self.initialized:
self._populate_init_list(data_dict, batching_dim)
self.conditions_idx.append(condition_idx)
self.empty = False
else:
raise ValueError('Dataset already initialized')
def _populate_init_list(self, data_dict, batching_dim=0):
current_cond_num_el = None
for slot in data_dict.keys():
slot_data = data_dict[slot]
if batching_dim != 0:
if isinstance(slot_data, (LabelTensor, torch.Tensor)):
dims = len(slot_data.size())
slot_data = slot_data.permute(
[batching_dim] + [dim for dim in range(dims) if
dim != batching_dim])
if current_cond_num_el is None:
current_cond_num_el = len(slot_data)
elif current_cond_num_el != len(slot_data):
raise ValueError('Different dimension in same condition')
current_list = getattr(self, slot)
current_list += [slot_data] if not (
isinstance(slot_data, list)) else slot_data
self.num_el_per_condition.append(current_cond_num_el)
def initialize(self):
"""
Initialize the datasets tensors/LabelTensors/lists given the lists
already filled
"""
logging.debug(f'Initialize dataset {self.__class__.__name__}')
if self.num_el_per_condition:
self.condition_indices = torch.cat(
[
torch.tensor([i] * num_el_per_condition[i],
torch.tensor([i] * self.num_el_per_condition[i],
dtype=torch.uint8)
for i in range(len(num_el_per_condition))
for i in range(len(self.num_el_per_condition))
],
dim=0,
dim=0
)
for slot in self.__slots__:
current_attribute = getattr(self, slot)
if all(isinstance(a, LabelTensor) for a in current_attribute):
setattr(self, slot, LabelTensor.vstack(current_attribute))
else:
self.condition_indices = torch.tensor([], dtype=torch.uint8)
for slot in self.__slots__:
setattr(self, slot, torch.tensor([]))
self.device = device
self.initialized = True
def __len__(self):
"""
:return: Number of elements in the dataset
"""
return len(getattr(self, self.__slots__[0]))
def __getattribute__(self, item):
attribute = super().__getattribute__(item)
if isinstance(attribute,
LabelTensor) and attribute.dtype == torch.float32:
attribute = attribute.to(device=self.device).requires_grad_()
return attribute
def __getitem__(self, idx):
if isinstance(idx, str):
return getattr(self, idx).to(self.device)
if isinstance(idx, slice):
to_return_list = []
for i in self.__slots__:
to_return_list.append(getattr(self, i)[idx].to(self.device))
return to_return_list
"""
:param idx:
:return:
"""
if not isinstance(idx, (tuple, list, slice, int)):
raise IndexError("Invalid index")
tensors = []
for attribute in self.__slots__:
tensor = getattr(self, attribute)
if isinstance(attribute, (LabelTensor, torch.Tensor)):
tensors.append(tensor.__getitem__(idx))
elif isinstance(attribute, list):
if isinstance(idx, (list, tuple)):
tensor = [tensor[i] for i in idx]
tensors.append(tensor)
return tensors
if isinstance(idx, (tuple, list)):
if (len(idx) == 2 and isinstance(idx[0], str)
and isinstance(idx[1], (list, slice))):
tensor = getattr(self, idx[0])
return tensor[[idx[1]]].to(self.device)
if all(isinstance(x, int) for x in idx):
to_return_list = []
for i in self.__slots__:
to_return_list.append(
getattr(self, i)[[idx]].to(self.device))
return to_return_list
raise ValueError(f'Invalid index {idx}')
def apply_shuffle(self, indices):
for slot in self.__slots__:
if slot != 'equation':
attribute = getattr(self, slot)
if isinstance(attribute, (LabelTensor, torch.Tensor)):
setattr(self, 'slot', attribute[[indices]])
if isinstance(attribute, list):
setattr(self, 'slot', [attribute[i] for i in indices])

148
pina/data/data_module.py
View File

@@ -4,7 +4,8 @@ This module provide basic data management functionalities
import math
import torch
from lightning import LightningDataModule
import logging
from pytorch_lightning import LightningDataModule
from .sample_dataset import SamplePointDataset
from .supervised_dataset import SupervisedDataset
from .unsupervised_dataset import UnsupervisedDataset
@@ -22,8 +23,9 @@ class PinaDataModule(LightningDataModule):
problem,
device,
train_size=.7,
test_size=.2,
eval_size=.1,
test_size=.1,
val_size=.2,
predict_size=0.,
batch_size=None,
shuffle=True,
datasets=None):
@@ -37,37 +39,64 @@ class PinaDataModule(LightningDataModule):
:param batch_size: batch size used for training
:param datasets: list of datasets objects
"""
logging.debug('Start initialization of Pina DataModule')
logging.info('Start initialization of Pina DataModule')
super().__init__()
dataset_classes = [SupervisedDataset, UnsupervisedDataset,
SamplePointDataset]
self.problem = problem
self.device = device
self.dataset_classes = [SupervisedDataset, UnsupervisedDataset,
SamplePointDataset]
if datasets is None:
self.datasets = [DatasetClass(problem, device) for DatasetClass in
dataset_classes]
self.datasets = None
else:
self.datasets = datasets
self.split_length = []
self.split_names = []
self.loader_functions = {}
self.batch_size = batch_size
self.condition_names = problem.collector.conditions_name
if train_size > 0:
self.split_names.append('train')
self.split_length.append(train_size)
self.loader_functions['train_dataloader'] = lambda: PinaDataLoader(
self.splits['train'], self.batch_size, self.condition_names)
if test_size > 0:
self.split_length.append(test_size)
self.split_names.append('test')
if eval_size > 0:
self.split_length.append(eval_size)
self.split_names.append('eval')
self.batch_size = batch_size
self.condition_names = None
self.loader_functions['test_dataloader'] = lambda: PinaDataLoader(
self.splits['test'], self.batch_size, self.condition_names)
if val_size > 0:
self.split_length.append(val_size)
self.split_names.append('val')
self.loader_functions['val_dataloader'] = lambda: PinaDataLoader(
self.splits['val'], self.batch_size,
self.condition_names)
if predict_size > 0:
self.split_length.append(predict_size)
self.split_names.append('predict')
self.loader_functions[
'predict_dataloader'] = lambda: PinaDataLoader(
self.splits['predict'], self.batch_size,
self.condition_names)
self.splits = {k: {} for k in self.split_names}
self.shuffle = shuffle
for k, v in self.loader_functions.items():
setattr(self, k, v)
def prepare_data(self):
if self.datasets is None:
self._create_datasets()
def setup(self, stage=None):
"""
Perform the splitting of the dataset
"""
self.extract_conditions()
logging.debug('Start setup of Pina DataModule obj')
if self.datasets is None:
self._create_datasets()
if stage == 'fit' or stage is None:
for dataset in self.datasets:
if len(dataset) > 0:
@@ -82,53 +111,6 @@ class PinaDataModule(LightningDataModule):
else:
raise ValueError("stage must be either 'fit' or 'test'")
def extract_conditions(self):
"""
Extract conditions from dataset and update condition indices
"""
# Extract number of conditions
n_conditions = 0
for dataset in self.datasets:
if n_conditions != 0:
dataset.condition_names = {
key + n_conditions: value
for key, value in dataset.condition_names.items()
}
n_conditions += len(dataset.condition_names)
self.condition_names = {
key: value
for dataset in self.datasets
for key, value in dataset.condition_names.items()
}
def train_dataloader(self):
"""
Return the training dataloader for the dataset
:return: data loader
:rtype: PinaDataLoader
"""
return PinaDataLoader(self.splits['train'], self.batch_size,
self.condition_names)
def test_dataloader(self):
"""
Return the testing dataloader for the dataset
:return: data loader
:rtype: PinaDataLoader
"""
return PinaDataLoader(self.splits['test'], self.batch_size,
self.condition_names)
def eval_dataloader(self):
"""
Return the evaluation dataloader for the dataset
:return: data loader
:rtype: PinaDataLoader
"""
return PinaDataLoader(self.splits['eval'], self.batch_size,
self.condition_names)
@staticmethod
def dataset_split(dataset, lengths, seed=None, shuffle=True):
"""
@@ -141,30 +123,28 @@ class PinaDataModule(LightningDataModule):
:rtype: PinaSubset
"""
if sum(lengths) - 1 < 1e-3:
len_dataset = len(dataset)
lengths = [
int(math.floor(len(dataset) * length)) for length in lengths
int(math.floor(len_dataset * length)) for length in lengths
]
remainder = len(dataset) - sum(lengths)
for i in range(remainder):
lengths[i % len(lengths)] += 1
elif sum(lengths) - 1 >= 1e-3:
raise ValueError(f"Sum of lengths is {sum(lengths)} less than 1")
if sum(lengths) != len(dataset):
raise ValueError("Sum of lengths is not equal to dataset length")
if shuffle:
if seed is not None:
generator = torch.Generator()
generator.manual_seed(seed)
indices = torch.randperm(sum(lengths),
generator=generator).tolist()
generator=generator)
else:
indices = torch.arange(sum(lengths)).tolist()
else:
indices = torch.arange(0, sum(lengths), 1,
dtype=torch.uint8).tolist()
indices = torch.randperm(sum(lengths))
dataset.apply_shuffle(indices)
indices = torch.arange(0, sum(lengths), 1,
dtype=torch.uint8).tolist()
offsets = [
sum(lengths[:i]) if i > 0 else 0 for i in range(len(lengths))
]
@@ -172,3 +152,29 @@ class PinaDataModule(LightningDataModule):
PinaSubset(dataset, indices[offset:offset + length])
for offset, length in zip(offsets, lengths)
]
def _create_datasets(self):
"""
Create the dataset objects putting data
"""
logging.debug('Dataset creation in PinaDataModule obj')
collector = self.problem.collector
batching_dim = self.problem.batching_dimension
datasets_slots = [i.__slots__ for i in self.dataset_classes]
self.datasets = [dataset(device=self.device) for dataset in
self.dataset_classes]
logging.debug('Filling datasets in PinaDataModule obj')
for name, data in collector.data_collections.items():
keys = list(data.keys())
idx = [key for key, val in collector.conditions_name.items() if
val == name]
for i, slot in enumerate(datasets_slots):
if slot == keys:
self.datasets[i].add_points(data, idx[0], batching_dim)
continue
datasets = []
for dataset in self.datasets:
if not dataset.empty:
dataset.initialize()
datasets.append(dataset)
self.datasets = datasets

25
pina/data/pina_batch.py
View File

@@ -10,13 +10,15 @@ class Batch:
optimization.
"""
def __init__(self, dataset_dict, idx_dict):
def __init__(self, dataset_dict, idx_dict, require_grad=True):
self.attributes = []
for k, v in dataset_dict.items():
setattr(self, k, v)
self.attributes.append(k)
for k, v in idx_dict.items():
setattr(self, k + '_idx', v)
self.require_grad = require_grad
def __len__(self):
"""
@@ -31,9 +33,18 @@ class Batch:
length += len(getattr(self, dataset))
return length
def __getattribute__(self, item):
if item in super().__getattribute__('attributes'):
dataset = super().__getattribute__(item)
index = super().__getattribute__(item + '_idx')
return PinaSubset(
dataset.dataset,
dataset.indices[index])
else:
return super().__getattribute__(item)
def __getattr__(self, item):
if not item in dir(self):
raise AttributeError(f'Batch instance has no attribute {item}')
return PinaSubset(
getattr(self, item).dataset,
getattr(self, item).indices[self.coordinates_dict[item]])
if item == 'data' and len(self.attributes) == 1:
item = self.attributes[0]
return super().__getattribute__(item)
raise AttributeError(f"'Batch' object has no attribute '{item}'")

13
pina/data/pina_subset.py
View File

@@ -2,21 +2,22 @@
Module for PinaSubset class
"""
from pina import LabelTensor
from torch import Tensor
from torch import Tensor, float32
class PinaSubset:
"""
TODO
"""
__slots__ = ['dataset', 'indices']
__slots__ = ['dataset', 'indices', 'require_grad']
def __init__(self, dataset, indices):
def __init__(self, dataset, indices, require_grad=True):
"""
TODO
"""
self.dataset = dataset
self.indices = indices
self.require_grad = require_grad
def __len__(self):
"""
@@ -27,7 +28,9 @@ class PinaSubset:
def __getattr__(self, name):
tensor = self.dataset.__getattribute__(name)
if isinstance(tensor, (LabelTensor, Tensor)):
return tensor[self.indices]
tensor = tensor[[self.indices]].to(self.dataset.device)
return tensor.requires_grad_(
self.require_grad) if tensor.dtype == float32 else tensor
if isinstance(tensor, list):
return [tensor[i] for i in self.indices]
raise AttributeError("No attribute named {}".format(name))
raise AttributeError(f"No attribute named {name}")

21
pina/data/sample_dataset.py
View File

@@ -1,8 +1,9 @@
"""
Sample dataset module
"""
from copy import deepcopy
from .base_dataset import BaseDataset
from ..condition.input_equation_condition import InputPointsEquationCondition
from ..condition import InputPointsEquationCondition
class SamplePointDataset(BaseDataset):
@@ -12,3 +13,21 @@ class SamplePointDataset(BaseDataset):
"""
data_type = 'physics'
__slots__ = InputPointsEquationCondition.__slots__
def add_points(self, data_dict, condition_idx, batching_dim=0):
data_dict = deepcopy(data_dict)
data_dict.pop('equation')
super().add_points(data_dict, condition_idx)
def _init_from_problem(self, collector_dict, batching_dim=0):
for name, data in collector_dict.items():
keys = list(data.keys())
if set(self.__slots__) == set(keys):
data = deepcopy(data)
data.pop('equation')
self._populate_init_list(data)
idx = [key for key, val in
self.problem.collector.conditions_name.items() if
val == name]
self.conditions_idx.append(idx)
self.initialize()

470
pina/label_tensor.py
View File

@@ -1,5 +1,5 @@
""" Module for LabelTensor """
from copy import deepcopy, copy
from copy import copy
import torch
from torch import Tensor
@@ -8,21 +8,29 @@ def issubset(a, b):
"""
Check if a is a subset of b.
"""
return set(a).issubset(set(b))
if isinstance(a, list) and isinstance(b, list):
return set(a).issubset(set(b))
elif isinstance(a, range) and isinstance(b, range):
return a.start <= b.start and a.stop >= b.stop
else:
return False
class LabelTensor(torch.Tensor):
"""Torch tensor with a label for any column."""
@staticmethod
def __new__(cls, x, labels, *args, **kwargs):
return super().__new__(cls, x, *args, **kwargs)
def __new__(cls, x, labels, full=True, *args, **kwargs):
if isinstance(x, LabelTensor):
return x
else:
return super().__new__(cls, x, *args, **kwargs)
@property
def tensor(self):
return self.as_subclass(Tensor)
def __init__(self, x, labels):
def __init__(self, x, labels, full=False):
"""
Construct a `LabelTensor` by passing a dict of the labels
@@ -34,8 +42,17 @@ class LabelTensor(torch.Tensor):
"""
self.dim_names = None
self.full = full
self.labels = labels
@classmethod
def __internal_init__(cls, x, labels, dim_names ,full=False, *args, **kwargs):
lt = cls.__new__(cls, x, labels, full, *args, **kwargs)
lt._labels = labels
lt.full = full
lt.dim_names = dim_names
return lt
@property
def labels(self):
"""Property decorator for labels
@@ -43,12 +60,29 @@ class LabelTensor(torch.Tensor):
:return: labels of self
:rtype: list
"""
return self._labels[self.tensor.ndim - 1]['dof']
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
"""
@@ -62,26 +96,77 @@ class LabelTensor(torch.Tensor):
:param labels: Labels to assign to the class variable _labels.
:type: labels: str | list(str) | dict
"""
if hasattr(self, 'labels') is False:
self.init_labels()
if not hasattr(self, '_labels'):
self._labels = {}
if isinstance(labels, dict):
self.update_labels_from_dict(labels)
self._init_labels_from_dict(labels)
elif isinstance(labels, list):
self.update_labels_from_list(labels)
self._init_labels_from_list(labels)
elif isinstance(labels, str):
labels = [labels]
self.update_labels_from_list(labels)
self._init_labels_from_list(labels)
else:
raise ValueError("labels must be list, dict or string.")
self.set_names()
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
if hasattr(self, 'full') and self.full:
labels = {i: labels[i] if i in labels else {'name': i} for i in
labels.keys()}
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) and list(v.keys()) == ['name']:
# Init from dict with only name key
v['dof'] = range(tensor_shape[k])
# Init from dict with both name and dof keys
elif isinstance(v, dict) and 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:
ValueError('Illegal labels initialization')
# Perform update
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 set_names(self):
labels = self.full_labels
labels = self.stored_labels
self.dim_names = {}
for dim in range(self.tensor.ndim):
for dim in labels.keys():
self.dim_names[labels[dim]['name']] = dim
def extract(self, label_to_extract):
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``.
@@ -91,78 +176,68 @@ class LabelTensor(torch.Tensor):
:raises TypeError: Labels are not ``str``.
:raises ValueError: Label to extract is not in the labels ``list``.
"""
if isinstance(label_to_extract, (str, int)):
label_to_extract = [label_to_extract]
if isinstance(label_to_extract, (tuple, list)):
return self._extract_from_list(label_to_extract)
if isinstance(label_to_extract, dict):
return self._extract_from_dict(label_to_extract)
raise ValueError('labels_to_extract must be str or list or dict')
# Convert str/int to string
if isinstance(labels_to_extract, (str, int)):
labels_to_extract = [labels_to_extract]
def _extract_from_list(self, labels_to_extract):
# Store locally all necessary obj/variables
ndim = self.tensor.ndim
labels = self.full_labels
tensor = self.tensor
last_dim_label = self.labels
# Store useful variables
labels = self.stored_labels
stored_keys = labels.keys()
dim_names = self.dim_names
ndim = len(super().shape)
# Verify if all the labels in labels_to_extract are in last dimension
if set(labels_to_extract).issubset(last_dim_label) is False:
raise ValueError(
'Cannot extract a dof which is not in the original LabelTensor')
# Extract index to extract
idx_to_extract = [last_dim_label.index(i) for i in labels_to_extract]
# Perform extraction
new_tensor = tensor[..., idx_to_extract]
# Manage labels
new_labels = copy(labels)
last_dim_new_label = {ndim - 1: {
'dof': list(labels_to_extract),
'name': labels[ndim - 1]['name']
}}
new_labels.update(last_dim_new_label)
return LabelTensor(new_tensor, new_labels)
def _extract_from_dict(self, labels_to_extract):
labels = self.full_labels
tensor = self.tensor
ndim = tensor.ndim
new_labels = deepcopy(labels)
new_tensor = tensor
for k, _ in labels_to_extract.items():
idx_dim = self.dim_names[k]
dim_labels = labels[idx_dim]['dof']
if isinstance(labels_to_extract[k], (int, str)):
labels_to_extract[k] = [labels_to_extract[k]]
if set(labels_to_extract[k]).issubset(dim_labels) is False:
# Convert tuple/list to dict
if isinstance(labels_to_extract, (tuple, list)):
if not ndim - 1 in stored_keys:
raise ValueError(
'Cannot extract a dof which is not in the original '
'LabelTensor')
idx_to_extract = [dim_labels.index(i) for i in labels_to_extract[k]]
indexer = [slice(None)] * idx_dim + [idx_to_extract] + [
slice(None)] * (ndim - idx_dim - 1)
new_tensor = new_tensor[indexer]
dim_new_label = {idx_dim: {
'dof': labels_to_extract[k],
'name': labels[idx_dim]['name']
}}
new_labels.update(dim_new_label)
return LabelTensor(new_tensor, new_labels)
"LabelTensor does not have labels in last dimension")
name = labels[max(stored_keys)]['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')
# Make copy of labels (avoid issue in consistency)
updated_labels = {k: copy(v) for k, v in labels.items()}
# Initialize list used to perform extraction
extractor = [slice(None) for _ in range(ndim)]
# Loop over labels_to_extract dict
for k, v in labels_to_extract.items():
# If label is not find raise value error
idx_dim = dim_names.get(k)
if idx_dim is None:
raise ValueError(
'Cannot extract label with is not in original labels')
dim_labels = labels[idx_dim]['dof']
v = [v] if isinstance(v, (int, str)) else v
if not isinstance(v, range):
extractor[idx_dim] = [dim_labels.index(i) for i in v] if len(
v) > 1 else slice(dim_labels.index(v[0]),
dim_labels.index(v[0]) + 1)
else:
extractor[idx_dim] = slice(v.start, v.stop)
updated_labels.update({idx_dim: {'dof': v, 'name': k}})
tensor = self.tensor
tensor = tensor[extractor]
return LabelTensor.__internal_init__(tensor, updated_labels, dim_names)
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 += super().__str__()
s += self.tensor.__str__()
return s
@staticmethod
@@ -174,55 +249,44 @@ class LabelTensor(torch.Tensor):
:param tensors: tensors to concatenate
:type tensors: list(LabelTensor)
:param dim: dimensions on which you want to perform the operation (default 0)
: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:
if len(tensors) == 1 or isinstance(tensors, LabelTensor):
return tensors[0]
new_labels_cat_dim = LabelTensor._check_validity_before_cat(tensors,
dim)
# Perform cat on tensors
new_tensor = torch.cat(tensors, dim=dim)
# Update labels
labels = tensors[0].full_labels
labels.pop(dim)
new_labels_cat_dim = new_labels_cat_dim if len(
set(new_labels_cat_dim)) == len(new_labels_cat_dim) \
else range(new_tensor.shape[dim])
labels[dim] = {'dof': new_labels_cat_dim,
'name': tensors[1].full_labels[dim]['name']}
return LabelTensor(new_tensor, labels)
labels = LabelTensor.__create_labels_cat(tensors,
dim)
return LabelTensor.__internal_init__(new_tensor, labels, tensors[0].dim_names)
@staticmethod
def _check_validity_before_cat(tensors, dim):
n_dims = tensors[0].ndim
new_labels_cat_dim = []
def __create_labels_cat(tensors, dim):
# Check if names and dof of the labels are the same in all dimensions
# except in dim
for i in range(n_dims):
name = tensors[0].full_labels[i]['name']
if i != dim:
dof = tensors[0].full_labels[i]['dof']
for tensor in tensors:
dof_to_check = tensor.full_labels[i]['dof']
name_to_check = tensor.full_labels[i]['name']
if dof != dof_to_check or name != name_to_check:
raise ValueError(
'dimensions must have the same dof and name')
else:
for tensor in tensors:
new_labels_cat_dim += tensor.full_labels[i]['dof']
name_to_check = tensor.full_labels[i]['name']
if name != name_to_check:
raise ValueError(
'Dimensions to concatenate must have the same name')
return new_labels_cat_dim
stored_labels = [tensor.stored_labels for tensor in tensors]
# check if:
# - labels dict have same keys
# - all labels are the same expect for dimension dim
if not all(all(stored_labels[i][k] == stored_labels[0][k]
for i in range(len(stored_labels)))
for k in stored_labels[0].keys() if k != dim):
raise RuntimeError('tensors must have the same shape and dof')
labels = {k: copy(v) for k, v in tensors[0].stored_labels.items()}
if dim in labels.keys():
last_dim_dof = [i for j in stored_labels for i in j[dim]['dof']]
labels[dim]['dof'] = last_dim_dof
return labels
def requires_grad_(self, mode=True):
lt = super().requires_grad_(mode)
@@ -251,52 +315,10 @@ class LabelTensor(torch.Tensor):
:return: A copy of the tensor.
:rtype: LabelTensor
"""
out = LabelTensor(super().clone(*args, **kwargs), self._labels)
labels = {k: copy(v) for k, v in self._labels.items()}
out = LabelTensor(super().clone(*args, **kwargs), labels)
return out
def init_labels(self):
self._labels = {
idx_: {
'dof': range(self.tensor.shape[idx_]),
'name': idx_
} for idx_ in range(self.tensor.ndim)
}
def update_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.tensor.shape
# Check dimensionality
for k, v in labels.items():
if len(v['dof']) != len(set(v['dof'])):
raise ValueError("dof must be unique")
if len(v['dof']) != tensor_shape[k]:
raise ValueError(
'Number of dof does not match with tensor dimension')
# Perform update
self._labels.update(labels)
def update_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.tensor.ndim - 1: {'dof': labels, 'name': self.tensor.ndim - 1}}
self.update_labels_from_dict(last_dim_labels)
@staticmethod
def summation(tensors):
if len(tensors) == 0:
@@ -304,25 +326,30 @@ class LabelTensor(torch.Tensor):
if len(tensors) == 1:
return tensors[0]
# Collect all labels
labels = tensors[0].full_labels
# Check labels of all the tensors in each dimension
for j in range(tensors[0].ndim):
for i in range(1, len(tensors)):
if labels[j] != tensors[i].full_labels[j]:
labels.pop(j)
break
# Sum tensors
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)
for tensor in tensors:
data += tensor.tensor
new_tensor = LabelTensor(data, labels)
return new_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.tensor.ndim - 1)
dim=self.ndim - 1)
elif mode == 'cross':
# Crete tensor and call cat on last dimension
tensor1 = self
@@ -333,7 +360,7 @@ class LabelTensor(torch.Tensor):
tensor2 = LabelTensor(tensor2.repeat_interleave(n1, dim=0),
labels=tensor2.labels)
new_label_tensor = LabelTensor.cat([tensor1, tensor2],
dim=self.tensor.ndim - 1)
dim=self.ndim - 1)
else:
raise ValueError('mode must be either "std" or "cross"')
return new_label_tensor
@@ -357,97 +384,76 @@ class LabelTensor(torch.Tensor):
:param index:
:return:
"""
if isinstance(index, str) or (isinstance(index, (tuple, list)) and all(
isinstance(a, str) for a in index)):
return self.extract(index)
selected_lt = super().__getitem__(index)
if isinstance(index, (int, slice)):
return self._getitem_int_slice(index, selected_lt)
index = [index]
if len(index) == self.tensor.ndim:
return self._getitem_full_dim_indexing(index, selected_lt)
if index[0] == Ellipsis:
index = [slice(None)] * (self.ndim - 1) + [index[1]]
if isinstance(index, torch.Tensor) or (
isinstance(index, (tuple, list)) and all(
isinstance(x, int) for x in index)):
return self._getitem_permutation(index, selected_lt)
raise ValueError('Not recognized index type')
def _getitem_int_slice(self, index, selected_lt):
"""
:param index:
:param selected_lt:
:return:
"""
if selected_lt.ndim == 1:
selected_lt = selected_lt.reshape(1, -1)
if hasattr(self, "labels"):
new_labels = deepcopy(self.full_labels)
to_update_dof = new_labels[0]['dof'][index]
to_update_dof = to_update_dof if isinstance(to_update_dof, (
tuple, list, range)) else [to_update_dof]
new_labels.update(
{0: {'dof': to_update_dof, 'name': new_labels[0]['name']}}
)
selected_lt.labels = new_labels
return selected_lt
def _getitem_full_dim_indexing(self, index, selected_lt):
new_labels = {}
old_labels = self.full_labels
if selected_lt.ndim == 1:
selected_lt = selected_lt.reshape(-1, 1)
new_labels = deepcopy(old_labels)
new_labels[1].update({'dof': old_labels[1]['dof'][index[1]],
'name': old_labels[1]['name']})
idx = 0
for j in range(selected_lt.ndim):
if not isinstance(index[j], int):
if hasattr(self, "labels"):
new_labels.update(
self._update_label_for_dim(old_labels, index[j], idx))
idx += 1
selected_lt.labels = new_labels
return selected_lt
def _getitem_permutation(self, index, selected_lt):
new_labels = deepcopy(self.full_labels)
new_labels.update(self._update_label_for_dim(self.full_labels, index,
0))
selected_lt.labels = self.labels
labels = {k: copy(v) for k, v in self.stored_labels.items()}
for j, idx in enumerate(index):
if isinstance(idx, int):
selected_lt = selected_lt.unsqueeze(j)
if j in labels.keys() and idx != slice(None):
self._update_single_label(labels, labels, idx, j)
selected_lt = LabelTensor.__internal_init__(selected_lt, labels,
self.dim_names)
return selected_lt
@staticmethod
def _update_label_for_dim(old_labels, index, dim):
def _update_single_label(old_labels, to_update_labels, index, dim):
"""
TODO
:param old_labels:
:param index:
:param dim:
: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[dim]['dof']
if not isinstance(index, (int, slice)) and len(index) == len(
old_dof) and isinstance(old_dof, range):
return
if isinstance(index, torch.Tensor):
index = index.nonzero()
index = index.nonzero(as_tuple=True)[
0] if index.dtype == torch.bool else index.tolist()
if isinstance(index, list):
return {dim: {'dof': [old_labels[dim]['dof'][i] for i in index],
'name': old_labels[dim]['name']}}
to_update_labels.update({dim: {
'dof': [old_dof[i] for i in index],
'name': old_labels[dim]['name']}})
else:
return {dim: {'dof': old_labels[dim]['dof'][index],
'name': old_labels[dim]['name']}}
to_update_labels.update({dim: {'dof': old_dof[index],
'name': old_labels[dim]['name']}})
def sort_labels(self, dim=None):
def argsort(lst):
def arg_sort(lst):
return sorted(range(len(lst)), key=lambda x: lst[x])
if dim is None:
dim = self.tensor.ndim - 1
labels = self.full_labels[dim]['dof']
sorted_index = argsort(labels)
indexer = [slice(None)] * self.tensor.ndim
dim = self.ndim - 1
labels = self.stored_labels[dim]['dof']
sorted_index = arg_sort(labels)
indexer = [slice(None)] * self.ndim
indexer[dim] = sorted_index
new_labels = deepcopy(self.full_labels)
new_labels[dim] = {'dof': sorted(labels),
'name': new_labels[dim]['name']}
return LabelTensor(self.tensor[indexer], new_labels)
return self.__getitem__(indexer)
def __deepcopy__(self, memo):
from copy import deepcopy
cls = self.__class__
result = cls(deepcopy(self.tensor), deepcopy(self.stored_labels))
return result
def permute(self, *dims):
tensor = super().permute(*dims)
stored_labels = self.stored_labels
keys_list = list(*dims)
labels = {keys_list.index(k): copy(stored_labels[k]) for k in
stored_labels.keys()}
return LabelTensor.__internal_init__(tensor, labels, self.dim_names)

7
pina/operators.py
View File

@@ -85,7 +85,8 @@ def grad(output_, input_, components=None, d=None):
raise RuntimeError
gradients = grad_scalar_output(output_, input_, d)
elif output_.shape[output_.ndim - 1] >= 2: # vector output ##############################
elif output_.shape[
output_.ndim - 1] >= 2: # vector output ##############################
tensor_to_cat = []
for i, c in enumerate(components):
c_output = output_.extract([c])
@@ -143,7 +144,6 @@ def div(output_, input_, components=None, d=None):
tensors_to_sum.append(grad_output.extract(c_fields))
labels[i] = c_fields
div_result = LabelTensor.summation(tensors_to_sum)
div_result.labels = ["+".join(labels)]
return div_result
@@ -249,7 +249,8 @@ def laplacian(output_, input_, components=None, d=None, method="std"):
result[:, idx] = grad(grad_output, input_, d=di).flatten()
to_append_tensors[idx] = grad(grad_output, input_, d=di)
labels[idx] = f"dd{ci[0]}dd{di[0]}"
result = LabelTensor.cat(tensors=to_append_tensors, dim=output_.tensor.ndim - 1)
result = LabelTensor.cat(tensors=to_append_tensors,
dim=output_.tensor.ndim - 1)
result.labels = labels
return result

9
pina/problem/abstract_problem.py
View File

@@ -32,11 +32,20 @@ class AbstractProblem(metaclass=ABCMeta):
# training all type self.collector.full, which returns true if all
# points are ready.
self.collector.store_fixed_data()
self._batching_dimension = 0
@property
def collector(self):
return self._collector
@property
def batching_dimension(self):
return self._batching_dimension
@batching_dimension.setter
def batching_dimension(self, value):
self._batching_dimension = value
# TODO this should be erase when dataloading will interface collector,
# kept only for back compatibility
@property

2
pina/solvers/solver.py
View File

@@ -94,7 +94,7 @@ class SolverInterface(pytorch_lightning.LightningModule, metaclass=ABCMeta):
pass
@abstractmethod
def training_step(self):
def training_step(self, batch, batch_idx):
pass
@abstractmethod

2
pina/trainer.py
View File

@@ -79,7 +79,7 @@ class Trainer(pytorch_lightning.Trainer):
data_module = PinaDataModule(problem=self.solver.problem, device=device,
train_size=self.train_size,
test_size=self.test_size,
eval_size=self.eval_size)
val_size=self.eval_size)
data_module.setup()
self._loader = data_module.train_dataloader()