diff --git a/pina/__init__.py b/pina/__init__.py index 0fe93752d..3bc28ae6c 100644 --- a/pina/__init__.py +++ b/pina/__init__.py @@ -1,11 +1,6 @@ __all__ = [ - "PINN", - "Trainer", - "LabelTensor", - "Plotter", - "Condition", - "SamplePointDataset", - "SamplePointLoader", + "Trainer", "LabelTensor", "Plotter", "Condition", "SamplePointDataset", + "PinaDataModule", "PinaDataLoader", 'TorchOptimizer', 'Graph' ] from .meta import * @@ -15,4 +10,8 @@ from .plotter import Plotter from .condition.condition import Condition from .data import SamplePointDataset -from .data import SamplePointLoader +from .data import PinaDataModule +from .data import PinaDataLoader +from .optim import TorchOptimizer +from .optim import TorchScheduler +from .graph import Graph diff --git a/pina/collector.py b/pina/collector.py new file mode 100644 index 000000000..3219b2b6a --- /dev/null +++ b/pina/collector.py @@ -0,0 +1,114 @@ +from .utils import check_consistency, merge_tensors + + +class Collector: + + def __init__(self, problem): + # creating a hook between collector and problem + self.problem = problem + + # this variable is used to store the data in the form: + # {'[condition_name]' : + # {'input_points' : Tensor, + # '[equation/output_points/conditional_variables]': Tensor} + # } + # 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 = { + name: False + for name in self.problem.conditions + } + self.full = False + + @property + def full(self): + return all(self._is_conditions_ready.values()) + + @full.setter + def full(self, value): + check_consistency(value, bool) + self._full = value + + @property + def data_collections(self): + return self._data_collections + + @property + def problem(self): + return self._problem + + @problem.setter + def problem(self, value): + self._problem = value + + def store_fixed_data(self): + # loop over all conditions + for condition_name, condition in self.problem.conditions.items(): + # if the condition is not ready and domain is not attribute + # of condition, we get and store the data + if (not self._is_conditions_ready[condition_name]) and (not hasattr( + condition, "domain")): + # get data + keys = condition.__slots__ + values = [getattr(condition, name) for name in keys] + self.data_collections[condition_name] = dict(zip(keys, values)) + # condition now is ready + self._is_conditions_ready[condition_name] = True + + def store_sample_domains(self, n, mode, variables, sample_locations): + # loop over all locations + for loc in sample_locations: + # get condition + condition = self.problem.conditions[loc] + keys = ["input_points", "equation"] + # if the condition is not ready, we get and store the data + if (not self._is_conditions_ready[loc]): + # if it is the first time we sample + if not self.data_collections[loc]: + already_sampled = [] + # if we have sampled the condition but not all variables + else: + already_sampled = [ + self.data_collections[loc]['input_points'] + ] + # if the condition is ready but we want to sample again + else: + self._is_conditions_ready[loc] = False + already_sampled = [] + + # get the samples + samples = [ + condition.domain.sample(n=n, mode=mode, variables=variables) + ] + already_sampled + pts = merge_tensors(samples) + if (set(pts.labels).issubset(sorted(self.problem.input_variables))): + pts = pts.sort_labels() + if sorted(pts.labels) == sorted(self.problem.input_variables): + self._is_conditions_ready[loc] = True + values = [pts, condition.equation] + self.data_collections[loc] = dict(zip(keys, values)) + else: + raise RuntimeError( + 'Try to sample variables which are not in problem defined ' + 'in the problem') + + def add_points(self, new_points_dict): + """ + Add input points to a sampled condition + + :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(): + if not self._is_conditions_ready[k]: + raise RuntimeError( + 'Cannot add points on a non sampled condition') + self.data_collections[k]['input_points'] = self.data_collections[k][ + 'input_points'].vstack(v) diff --git a/pina/condition/__init__.py b/pina/condition/__init__.py index ff329c1bc..4c89b75d1 100644 --- a/pina/condition/__init__.py +++ b/pina/condition/__init__.py @@ -1,10 +1,12 @@ __all__ = [ 'Condition', 'ConditionInterface', - 'DomainOutputCondition', - 'DomainEquationCondition' + 'DomainEquationCondition', + 'InputPointsEquationCondition', + 'InputOutputPointsCondition', ] from .condition_interface import ConditionInterface -from .domain_output_condition import DomainOutputCondition -from .domain_equation_condition import DomainEquationCondition \ No newline at end of file +from .domain_equation_condition import DomainEquationCondition +from .input_equation_condition import InputPointsEquationCondition +from .input_output_condition import InputOutputPointsCondition \ No newline at end of file diff --git a/pina/condition/condition.py b/pina/condition/condition.py index d815838bb..3a62143e4 100644 --- a/pina/condition/condition.py +++ b/pina/condition/condition.py @@ -1,27 +1,22 @@ """ Condition module. """ -from ..label_tensor import LabelTensor -from ..domain import DomainInterface -from ..equation.equation import Equation - -from . import DomainOutputCondition, DomainEquationCondition - - -def dummy(a): - """Dummy function for testing purposes.""" - return None +from .domain_equation_condition import DomainEquationCondition +from .input_equation_condition import InputPointsEquationCondition +from .input_output_condition import InputOutputPointsCondition +from .data_condition import DataConditionInterface class Condition: """ The class ``Condition`` is used to represent the constraints (physical equations, boundary conditions, etc.) that should be satisfied in the - problem at hand. Condition objects are used to formulate the PINA :obj:`pina.problem.abstract_problem.AbstractProblem` object. - Conditions can be specified in three ways: + problem at hand. Condition objects are used to formulate the + PINA :obj:`pina.problem.abstract_problem.AbstractProblem` object. + Conditions can be specified in four ways: 1. By specifying the input and output points of the condition; in such a case, the model is trained to produce the output points given the input - points. + points. Those points can either be torch.Tensor, LabelTensors, Graph 2. By specifying the location and the equation of the condition; in such a case, the model is trained to minimize the equation residual by @@ -29,79 +24,42 @@ class Condition: 3. By specifying the input points and the equation of the condition; in such a case, the model is trained to minimize the equation residual by - evaluating it at the passed input points. + evaluating it at the passed input points. The input points must be + a LabelTensor. + + 4. By specifying only the data matrix; in such a case the model is + trained with an unsupervised costum loss and uses the data in training. + Additionaly conditioning variables can be passed, whenever the model + has extra conditioning variable it depends on. Example:: - >>> example_domain = Span({'x': [0, 1], 'y': [0, 1]}) - >>> def example_dirichlet(input_, output_): - >>> value = 0.0 - >>> return output_.extract(['u']) - value - >>> example_input_pts = LabelTensor( - >>> torch.tensor([[0, 0, 0]]), ['x', 'y', 'z']) - >>> example_output_pts = LabelTensor(torch.tensor([[1, 2]]), ['a', 'b']) - >>> - >>> Condition( - >>> input_points=example_input_pts, - >>> output_points=example_output_pts) - >>> Condition( - >>> location=example_domain, - >>> equation=example_dirichlet) - >>> Condition( - >>> input_points=example_input_pts, - >>> equation=example_dirichlet) + >>> TODO """ - # def _dictvalue_isinstance(self, dict_, key_, class_): - # """Check if the value of a dictionary corresponding to `key` is an instance of `class_`.""" - # if key_ not in dict_.keys(): - # return True - - # return isinstance(dict_[key_], class_) - - # def __init__(self, *args, **kwargs): - # """ - # Constructor for the `Condition` class. - # """ - # self.data_weight = kwargs.pop("data_weight", 1.0) - - # if len(args) != 0: - # raise ValueError( - # f"Condition takes only the following keyword arguments: {Condition.__slots__}." - # ) + __slots__ = list( + set(InputOutputPointsCondition.__slots__ + + InputPointsEquationCondition.__slots__ + + DomainEquationCondition.__slots__ + + DataConditionInterface.__slots__)) def __new__(cls, *args, **kwargs): - if sorted(kwargs.keys()) == sorted(["input_points", "output_points"]): - return DomainOutputCondition( - domain=kwargs["input_points"], - output_points=kwargs["output_points"] - ) - elif sorted(kwargs.keys()) == sorted(["domain", "output_points"]): - return DomainOutputCondition(**kwargs) - elif sorted(kwargs.keys()) == sorted(["domain", "equation"]): + if len(args) != 0: + raise ValueError("Condition takes only the following keyword " + f"arguments: {Condition.__slots__}.") + + sorted_keys = sorted(kwargs.keys()) + if sorted_keys == sorted(InputOutputPointsCondition.__slots__): + return InputOutputPointsCondition(**kwargs) + elif sorted_keys == sorted(InputPointsEquationCondition.__slots__): + return InputPointsEquationCondition(**kwargs) + elif sorted_keys == sorted(DomainEquationCondition.__slots__): return DomainEquationCondition(**kwargs) + elif sorted_keys == sorted(DataConditionInterface.__slots__): + return DataConditionInterface(**kwargs) + elif sorted_keys == DataConditionInterface.__slots__[0]: + return DataConditionInterface(**kwargs) else: raise ValueError(f"Invalid keyword arguments {kwargs.keys()}.") - # TODO: remove, not used anymore - ''' - if ( - sorted(kwargs.keys()) != sorted(["input_points", "output_points"]) - and sorted(kwargs.keys()) != sorted(["location", "equation"]) - and sorted(kwargs.keys()) != sorted(["input_points", "equation"]) - ): - raise ValueError(f"Invalid keyword arguments {kwargs.keys()}.") - # TODO: remove, not used anymore - if not self._dictvalue_isinstance(kwargs, "input_points", LabelTensor): - raise TypeError("`input_points` must be a torch.Tensor.") - if not self._dictvalue_isinstance(kwargs, "output_points", LabelTensor): - raise TypeError("`output_points` must be a torch.Tensor.") - if not self._dictvalue_isinstance(kwargs, "location", Location): - raise TypeError("`location` must be a Location.") - if not self._dictvalue_isinstance(kwargs, "equation", Equation): - raise TypeError("`equation` must be a Equation.") - - for key, value in kwargs.items(): - setattr(self, key, value) - ''' \ No newline at end of file diff --git a/pina/condition/condition_interface.py b/pina/condition/condition_interface.py index 0626a6d83..f2fe5db97 100644 --- a/pina/condition/condition_interface.py +++ b/pina/condition/condition_interface.py @@ -1,21 +1,33 @@ - -from abc import ABCMeta, abstractmethod +from abc import ABCMeta class ConditionInterface(metaclass=ABCMeta): - def __init__(self) -> None: + condition_types = ['physics', 'supervised', 'unsupervised'] + + def __init__(self, *args, **kwargs): + self._condition_type = None self._problem = None - @abstractmethod - def residual(self, model): - """ - Compute the residual of the condition. + @property + def problem(self): + return self._problem + + @problem.setter + def problem(self, value): + self._problem = value - :param model: The model to evaluate the condition. - :return: The residual of the condition. - """ - pass + @property + def condition_type(self): + return self._condition_type - def set_problem(self, problem): - self._problem = problem + @condition_type.setter + def condition_type(self, values): + if not isinstance(values, (list, tuple)): + values = [values] + for value in values: + if value not in ConditionInterface.condition_types: + raise ValueError( + 'Unavailable type of condition, expected one of' + f' {ConditionInterface.condition_types}.') + self._condition_type = values diff --git a/pina/condition/data_condition.py b/pina/condition/data_condition.py new file mode 100644 index 000000000..c6777231c --- /dev/null +++ b/pina/condition/data_condition.py @@ -0,0 +1,33 @@ +import torch + +from . import ConditionInterface +from ..label_tensor import LabelTensor +from ..graph import Graph +from ..utils import check_consistency + + +class DataConditionInterface(ConditionInterface): + """ + Condition for data. This condition must be used every + time a Unsupervised Loss is needed in the Solver. The conditionalvariable + can be passed as extra-input when the model learns a conditional + distribution + """ + + __slots__ = ["input_points", "conditional_variables"] + + def __init__(self, input_points, conditional_variables=None): + """ + TODO + """ + super().__init__() + self.input_points = input_points + self.conditional_variables = conditional_variables + self._condition_type = 'unsupervised' + + def __setattr__(self, key, value): + if (key == 'input_points') or (key == 'conditional_variables'): + check_consistency(value, (LabelTensor, Graph, torch.Tensor)) + DataConditionInterface.__dict__[key].__set__(self, value) + elif key in ('_problem', '_condition_type'): + super().__setattr__(key, value) diff --git a/pina/condition/domain_equation_condition.py b/pina/condition/domain_equation_condition.py index 15df3f85f..58dca70be 100644 --- a/pina/condition/domain_equation_condition.py +++ b/pina/condition/domain_equation_condition.py @@ -1,34 +1,34 @@ +import torch + from .condition_interface import ConditionInterface +from ..utils import check_consistency +from ..domain import DomainInterface +from ..equation.equation_interface import EquationInterface + class DomainEquationCondition(ConditionInterface): """ - Condition for input/output data. + Condition for domain/equation data. This condition must be used every + time a Physics Informed Loss is needed in the Solver. """ __slots__ = ["domain", "equation"] def __init__(self, domain, equation): """ - Constructor for the `InputOutputCondition` class. + TODO """ super().__init__() self.domain = domain self.equation = equation + self._condition_type = 'physics' - def residual(self, model): - """ - Compute the residual of the condition. - """ - self.batch_residual(model, self.domain, self.equation) - - @staticmethod - def batch_residual(model, input_pts, equation): - """ - Compute the residual of the condition for a single batch. Input and - output points are provided as arguments. - - :param torch.nn.Module model: The model to evaluate the condition. - :param torch.Tensor input_pts: The input points. - :param torch.Tensor equation: The output points. - """ - return equation.residual(input_pts, model(input_pts)) \ No newline at end of file + def __setattr__(self, key, value): + if key == 'domain': + check_consistency(value, (DomainInterface)) + DomainEquationCondition.__dict__[key].__set__(self, value) + elif key == 'equation': + check_consistency(value, (EquationInterface)) + DomainEquationCondition.__dict__[key].__set__(self, value) + elif key in ('_problem', '_condition_type'): + super().__setattr__(key, value) diff --git a/pina/condition/domain_output_condition.py b/pina/condition/domain_output_condition.py deleted file mode 100644 index f847720b5..000000000 --- a/pina/condition/domain_output_condition.py +++ /dev/null @@ -1,44 +0,0 @@ - -from . import ConditionInterface - -class DomainOutputCondition(ConditionInterface): - """ - Condition for input/output data. - """ - - __slots__ = ["domain", "output_points"] - - def __init__(self, domain, output_points): - """ - Constructor for the `InputOutputCondition` class. - """ - super().__init__() - print(self) - self.domain = domain - self.output_points = output_points - - @property - def input_points(self): - """ - Get the input points of the condition. - """ - return self._problem.domains[self.domain] - - def residual(self, model): - """ - Compute the residual of the condition. - """ - return self.batch_residual(model, self.domain, self.output_points) - - @staticmethod - def batch_residual(model, input_points, output_points): - """ - Compute the residual of the condition for a single batch. Input and - output points are provided as arguments. - - :param torch.nn.Module model: The model to evaluate the condition. - :param torch.Tensor input_points: The input points. - :param torch.Tensor output_points: The output points. - """ - - return output_points - model(input_points) \ No newline at end of file diff --git a/pina/condition/input_equation_condition.py b/pina/condition/input_equation_condition.py index 288022c00..bf05130c0 100644 --- a/pina/condition/input_equation_condition.py +++ b/pina/condition/input_equation_condition.py @@ -1,23 +1,37 @@ +import torch -from . import ConditionInterface +from .condition_interface import ConditionInterface +from ..label_tensor import LabelTensor +from ..graph import Graph +from ..utils import check_consistency +from ..equation.equation_interface import EquationInterface -class InputEquationCondition(ConditionInterface): + +class InputPointsEquationCondition(ConditionInterface): """ - Condition for input/output data. + Condition for input_points/equation data. This condition must be used every + time a Physics Informed Loss is needed in the Solver. """ - __slots__ = ["input_points", "output_points"] + __slots__ = ["input_points", "equation"] - def __init__(self, input_points, output_points): + def __init__(self, input_points, equation): """ - Constructor for the `InputOutputCondition` class. + TODO """ super().__init__() self.input_points = input_points - self.output_points = output_points + self.equation = equation + self._condition_type = 'physics' - def residual(self, model): - """ - Compute the residual of the condition. - """ - return self.output_points - model(self.input_points) \ No newline at end of file + def __setattr__(self, key, value): + if key == 'input_points': + check_consistency( + value, (LabelTensor) + ) # for now only labeltensors, we need labels for the operators! + InputPointsEquationCondition.__dict__[key].__set__(self, value) + elif key == 'equation': + check_consistency(value, (EquationInterface)) + InputPointsEquationCondition.__dict__[key].__set__(self, value) + elif key in ('_problem', '_condition_type'): + super().__setattr__(key, value) diff --git a/pina/condition/input_output_condition.py b/pina/condition/input_output_condition.py new file mode 100644 index 000000000..08ed21d96 --- /dev/null +++ b/pina/condition/input_output_condition.py @@ -0,0 +1,31 @@ +import torch + +from .condition_interface import ConditionInterface +from ..label_tensor import LabelTensor +from ..graph import Graph +from ..utils import check_consistency + + +class InputOutputPointsCondition(ConditionInterface): + """ + Condition for domain/equation data. This condition must be used every + time a Physics Informed or a Supervised Loss is needed in the Solver. + """ + + __slots__ = ["input_points", "output_points"] + + def __init__(self, input_points, output_points): + """ + TODO + """ + super().__init__() + self.input_points = input_points + self.output_points = output_points + self._condition_type = ['supervised', 'physics'] + + def __setattr__(self, key, value): + if (key == 'input_points') or (key == 'output_points'): + check_consistency(value, (LabelTensor, Graph, torch.Tensor)) + InputOutputPointsCondition.__dict__[key].__set__(self, value) + elif key in ('_problem', '_condition_type'): + super().__setattr__(key, value) diff --git a/pina/data/__init__.py b/pina/data/__init__.py index fba19b92c..2b3a126a7 100644 --- a/pina/data/__init__.py +++ b/pina/data/__init__.py @@ -1,7 +1,15 @@ +""" +Import data classes +""" __all__ = [ + 'PinaDataLoader', 'SupervisedDataset', 'SamplePointDataset', + 'UnsupervisedDataset', 'Batch', 'PinaDataModule', 'BaseDataset' ] -from .pina_dataloader import SamplePointLoader -from .data_dataset import DataPointDataset +from .pina_dataloader import PinaDataLoader +from .supervised_dataset import SupervisedDataset from .sample_dataset import SamplePointDataset -from .pina_batch import Batch \ No newline at end of file +from .unsupervised_dataset import UnsupervisedDataset +from .pina_batch import Batch +from .data_module import PinaDataModule +from .base_dataset import BaseDataset diff --git a/pina/data/base_dataset.py b/pina/data/base_dataset.py new file mode 100644 index 000000000..2c28ba30b --- /dev/null +++ b/pina/data/base_dataset.py @@ -0,0 +1,157 @@ +""" +Basic data module implementation +""" +import torch +import logging + +from torch.utils.data import Dataset + +from ..label_tensor import LabelTensor + + +class BaseDataset(Dataset): + """ + BaseDataset class, which handle initialization and data retrieval + :var condition_indices: List of indices + :var device: torch.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. + :param torch.device device: The device on which the + dataset will be loaded. + """ + if cls is BaseDataset: + raise TypeError( + 'BaseDataset cannot be instantiated directly. Use a subclass.') + if not hasattr(cls, '__slots__'): + raise TypeError( + 'Something is wrong, __slots__ must be defined in subclasses.') + return object.__new__(cls) + + def __init__(self, problem=None, device=torch.device('cpu')): + """" + Initialize the object based on __slots__ + :param AbstractProblem problem: The formulation of the problem. + :param torch.device device: The device on which the + dataset will be loaded. + """ + super().__init__() + self.empty = True + self.problem = problem + self.device = device + self.condition_indices = None + for slot in self.__slots__: + setattr(self, slot, []) + 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()) + 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] * self.num_el_per_condition[i], + dtype=torch.uint8) + for i in range(len(self.num_el_per_condition)) + ], + 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)) + self.initialized = True + + def __len__(self): + """ + :return: Number of elements in the dataset + """ + return len(getattr(self, self.__slots__[0])) + + def __getitem__(self, idx): + """ + :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 + + 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]) diff --git a/pina/data/data_dataset.py b/pina/data/data_dataset.py deleted file mode 100644 index 9dff2d7ed..000000000 --- a/pina/data/data_dataset.py +++ /dev/null @@ -1,41 +0,0 @@ -from torch.utils.data import Dataset -import torch -from ..label_tensor import LabelTensor - - -class DataPointDataset(Dataset): - - def __init__(self, problem, device) -> None: - super().__init__() - input_list = [] - output_list = [] - self.condition_names = [] - - for name, condition in problem.conditions.items(): - if hasattr(condition, "output_points"): - input_list.append(problem.conditions[name].input_points) - output_list.append(problem.conditions[name].output_points) - self.condition_names.append(name) - - self.input_pts = LabelTensor.stack(input_list) - self.output_pts = LabelTensor.stack(output_list) - - if self.input_pts != []: - self.condition_indeces = torch.cat( - [ - torch.tensor([i] * len(input_list[i])) - for i in range(len(self.condition_names)) - ], - dim=0, - ) - else: # if there are no data points - self.condition_indeces = torch.tensor([]) - self.input_pts = torch.tensor([]) - self.output_pts = torch.tensor([]) - - self.input_pts = self.input_pts.to(device) - self.output_pts = self.output_pts.to(device) - self.condition_indeces = self.condition_indeces.to(device) - - def __len__(self): - return self.input_pts.shape[0] \ No newline at end of file diff --git a/pina/data/data_module.py b/pina/data/data_module.py new file mode 100644 index 000000000..bd117b54b --- /dev/null +++ b/pina/data/data_module.py @@ -0,0 +1,179 @@ +""" +This module provide basic data management functionalities +""" + +import math +import torch +import logging +from pytorch_lightning import LightningDataModule +from .sample_dataset import SamplePointDataset +from .supervised_dataset import SupervisedDataset +from .unsupervised_dataset import UnsupervisedDataset +from .pina_dataloader import PinaDataLoader +from .pina_subset import PinaSubset + + +class PinaDataModule(LightningDataModule): + """ + This class extend LightningDataModule, allowing proper creation and + management of different types of Datasets defined in PINA + """ + + def __init__(self, + problem, + device, + train_size=.7, + test_size=.1, + val_size=.2, + predict_size=0., + batch_size=None, + shuffle=True, + datasets=None): + """ + Initialize the object, creating dataset based on input problem + :param AbstractProblem problem: PINA problem + :param device: Device used for training and testing + :param train_size: number/percentage of elements in train split + :param test_size: number/percentage of elements in test split + :param eval_size: number/percentage of elements in evaluation split + :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__() + self.problem = problem + self.device = device + self.dataset_classes = [ + SupervisedDataset, UnsupervisedDataset, SamplePointDataset + ] + if datasets is None: + 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') + 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 + """ + 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: + splits = self.dataset_split(dataset, + self.split_length, + shuffle=self.shuffle) + for i in range(len(self.split_length)): + self.splits[self.split_names[i]][ + dataset.data_type] = splits[i] + elif stage == 'test': + raise NotImplementedError("Testing pipeline not implemented yet") + else: + raise ValueError("stage must be either 'fit' or 'test'") + + @staticmethod + def dataset_split(dataset, lengths, seed=None, shuffle=True): + """ + Perform the splitting of the dataset + :param dataset: dataset object we wanted to split + :param lengths: lengths of elements in dataset + :param seed: random seed + :param shuffle: shuffle dataset + :return: split dataset + :rtype: PinaSubset + """ + if sum(lengths) - 1 < 1e-3: + len_dataset = len(dataset) + 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 shuffle: + if seed is not None: + generator = torch.Generator() + generator.manual_seed(seed) + indices = torch.randperm(sum(lengths), generator=generator) + else: + 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)) + ] + return [ + 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 diff --git a/pina/data/pina_batch.py b/pina/data/pina_batch.py index cb1296ede..c5d1b61dc 100644 --- a/pina/data/pina_batch.py +++ b/pina/data/pina_batch.py @@ -1,36 +1,47 @@ +""" +Batch management module +""" +from .pina_subset import PinaSubset class Batch: """ - This class is used to create a dataset of sample points. + Implementation of the Batch class used during training to perform SGD + optimization. """ - def __init__(self, type_, idx, *args, **kwargs) -> None: - """ - """ - if type_ == "sample": - - if len(args) != 2: - raise RuntimeError - - input = args[0] - conditions = args[1] - - self.input = input[idx] - self.condition = conditions[idx] + 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) - elif type_ == "data": + for k, v in idx_dict.items(): + setattr(self, k + '_idx', v) + self.require_grad = require_grad - if len(args) != 3: - raise RuntimeError - - input = args[0] - output = args[1] - conditions = args[2] - - self.input = input[idx] - self.output = output[idx] - self.condition = conditions[idx] - - else: - raise ValueError("Invalid number of arguments.") \ No newline at end of file + def __len__(self): + """ + Returns the number of elements in the batch + :return: number of elements in the batch + :rtype: int + """ + length = 0 + for dataset in dir(self): + attribute = getattr(self, dataset) + if isinstance(attribute, list): + 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]) + return super().__getattribute__(item) + + def __getattr__(self, 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}'") diff --git a/pina/data/pina_dataloader.py b/pina/data/pina_dataloader.py index 2c8967c50..e2d3fb76e 100644 --- a/pina/data/pina_dataloader.py +++ b/pina/data/pina_dataloader.py @@ -1,11 +1,11 @@ -import torch - -from .sample_dataset import SamplePointDataset -from .data_dataset import DataPointDataset +""" +This module is used to create an iterable object used during training +""" +import math from .pina_batch import Batch -class SamplePointLoader: +class PinaDataLoader: """ This class is used to create a dataloader to use during the training. @@ -14,198 +14,55 @@ class SamplePointLoader: :vartype condition_names: list[str] """ - def __init__( - self, sample_dataset, data_dataset, batch_size=None, shuffle=True - ) -> None: - """ - Constructor. - - :param SamplePointDataset sample_pts: The sample points dataset. - :param int batch_size: The batch size. If ``None``, the batch size is - set to the number of sample points. Default is ``None``. - :param bool shuffle: If ``True``, the sample points are shuffled. - Default is ``True``. - """ - if not isinstance(sample_dataset, SamplePointDataset): - raise TypeError( - f"Expected SamplePointDataset, got {type(sample_dataset)}" - ) - if not isinstance(data_dataset, DataPointDataset): - raise TypeError( - f"Expected DataPointDataset, got {type(data_dataset)}" - ) - - self.n_data_conditions = len(data_dataset.condition_names) - self.n_phys_conditions = len(sample_dataset.condition_names) - data_dataset.condition_indeces += self.n_phys_conditions - - self._prepare_sample_dataset(sample_dataset, batch_size, shuffle) - self._prepare_data_dataset(data_dataset, batch_size, shuffle) - - self.condition_names = ( - sample_dataset.condition_names + data_dataset.condition_names - ) - - self.batch_list = [] - for i in range(len(self.batch_sample_pts)): - self.batch_list.append(("sample", i)) - - for i in range(len(self.batch_input_pts)): - self.batch_list.append(("data", i)) - - if shuffle: - self.random_idx = torch.randperm(len(self.batch_list)) - else: - self.random_idx = torch.arange(len(self.batch_list)) - - self._prepare_batches() - - def _prepare_data_dataset(self, dataset, batch_size, shuffle): - """ - Prepare the dataset for data points. - - :param SamplePointDataset dataset: The dataset. - :param int batch_size: The batch size. - :param bool shuffle: If ``True``, the sample points are shuffled. - """ - self.sample_dataset = dataset - - if len(dataset) == 0: - self.batch_data_conditions = [] - self.batch_input_pts = [] - self.batch_output_pts = [] - return - - if batch_size is None: - batch_size = len(dataset) - batch_num = len(dataset) // batch_size - if len(dataset) % batch_size != 0: - batch_num += 1 - - output_labels = dataset.output_pts.labels - input_labels = dataset.input_pts.labels - self.tensor_conditions = dataset.condition_indeces - - if shuffle: - idx = torch.randperm(dataset.input_pts.shape[0]) - self.input_pts = dataset.input_pts[idx] - self.output_pts = dataset.output_pts[idx] - self.tensor_conditions = dataset.condition_indeces[idx] - - self.batch_input_pts = torch.tensor_split(dataset.input_pts, batch_num) - self.batch_output_pts = torch.tensor_split( - dataset.output_pts, batch_num - ) - #print(input_labels) - for i in range(len(self.batch_input_pts)): - self.batch_input_pts[i].labels = input_labels - self.batch_output_pts[i].labels = output_labels - - self.batch_data_conditions = torch.tensor_split( - self.tensor_conditions, batch_num - ) - - def _prepare_sample_dataset(self, dataset, batch_size, shuffle): + def __init__(self, dataset_dict, batch_size, condition_names) -> None: """ - Prepare the dataset for sample points. - - :param DataPointDataset dataset: The dataset. - :param int batch_size: The batch size. - :param bool shuffle: If ``True``, the sample points are shuffled. + Initialize local variables + :param dataset_dict: Dictionary of datasets + :type dataset_dict: dict + :param batch_size: Size of the batch + :type batch_size: int + :param condition_names: Names of the conditions + :type condition_names: list[str] """ + self.condition_names = condition_names + self.dataset_dict = dataset_dict + self._init_batches(batch_size) - self.sample_dataset = dataset - if len(dataset) == 0: - self.batch_sample_conditions = [] - self.batch_sample_pts = [] - return - - if batch_size is None: - batch_size = len(dataset) - - batch_num = len(dataset) // batch_size - if len(dataset) % batch_size != 0: - batch_num += 1 - - self.tensor_pts = dataset.pts - self.tensor_conditions = dataset.condition_indeces - - # if shuffle: - # idx = torch.randperm(self.tensor_pts.shape[0]) - # self.tensor_pts = self.tensor_pts[idx] - # self.tensor_conditions = self.tensor_conditions[idx] - - self.batch_sample_pts = torch.tensor_split(self.tensor_pts, batch_num) - for i in range(len(self.batch_sample_pts)): - self.batch_sample_pts[i].labels = dataset.pts.labels - - self.batch_sample_conditions = torch.tensor_split( - self.tensor_conditions, batch_num - ) - - def _prepare_batches(self): + def _init_batches(self, batch_size=None): """ - Prepare the batches. + Create batches according to the batch_size provided in input. """ self.batches = [] - for i in range(len(self.batch_list)): - type_, idx_ = self.batch_list[i] - - if type_ == "sample": - batch = Batch( - "sample", idx_, - self.batch_sample_pts, - self.batch_sample_conditions) + n_elements = sum(len(v) for v in self.dataset_dict.values()) + if batch_size is None: + batch_size = n_elements + indexes_dict = {} + n_batches = int(math.ceil(n_elements / batch_size)) + for k, v in self.dataset_dict.items(): + if n_batches != 1: + indexes_dict[k] = math.floor(len(v) / (n_batches - 1)) else: - batch = Batch( - "data", idx_, - self.batch_input_pts, - self.batch_output_pts, - self.batch_data_conditions) - - self.batches.append(batch) + indexes_dict[k] = len(v) + for i in range(n_batches): + temp_dict = {} + for k, v in indexes_dict.items(): + if i != n_batches - 1: + temp_dict[k] = slice(i * v, (i + 1) * v) + else: + temp_dict[k] = slice(i * v, len(self.dataset_dict[k])) + self.batches.append( + Batch(idx_dict=temp_dict, dataset_dict=self.dataset_dict)) def __iter__(self): """ - Return an iterator over the points. Any element of the iterator is a - dictionary with the following keys: - - ``pts``: The input sample points. It is a LabelTensor with the - shape ``(batch_size, input_dimension)``. - - ``output``: The output sample points. This key is present only - if data conditions are present. It is a LabelTensor with the - shape ``(batch_size, output_dimension)``. - - ``condition``: The integer condition indeces. It is a tensor - with the shape ``(batch_size, )`` of type ``torch.int64`` and - indicates for any ``pts`` the corresponding problem condition. - - :return: An iterator over the points. - :rtype: iter + Makes dataloader object iterable """ - # for i in self.random_idx: - for i in self.random_idx: - yield self.batches[i] - - # for i in range(len(self.batch_list)): - # type_, idx_ = self.batch_list[i] - - # if type_ == "sample": - # d = { - # "pts": self.batch_sample_pts[idx_].requires_grad_(True), - # "condition": self.batch_sample_conditions[idx_], - # } - # else: - # d = { - # "pts": self.batch_input_pts[idx_].requires_grad_(True), - # "output": self.batch_output_pts[idx_], - # "condition": self.batch_data_conditions[idx_], - # } - # yield d + yield from self.batches def __len__(self): """ Return the number of batches. - :return: The number of batches. :rtype: int """ - return len(self.batch_list) + return len(self.batches) diff --git a/pina/data/pina_subset.py b/pina/data/pina_subset.py new file mode 100644 index 000000000..275541e97 --- /dev/null +++ b/pina/data/pina_subset.py @@ -0,0 +1,36 @@ +""" +Module for PinaSubset class +""" +from pina import LabelTensor +from torch import Tensor, float32 + + +class PinaSubset: + """ + TODO + """ + __slots__ = ['dataset', 'indices', 'require_grad'] + + def __init__(self, dataset, indices, require_grad=True): + """ + TODO + """ + self.dataset = dataset + self.indices = indices + self.require_grad = require_grad + + def __len__(self): + """ + TODO + """ + return len(self.indices) + + def __getattr__(self, name): + tensor = self.dataset.__getattribute__(name) + if isinstance(tensor, (LabelTensor, Tensor)): + 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(f"No attribute named {name}") diff --git a/pina/data/sample_dataset.py b/pina/data/sample_dataset.py index 84af2920f..bc3bca335 100644 --- a/pina/data/sample_dataset.py +++ b/pina/data/sample_dataset.py @@ -1,43 +1,35 @@ -from torch.utils.data import Dataset -import torch +""" +Sample dataset module +""" +from copy import deepcopy +from .base_dataset import BaseDataset +from ..condition import InputPointsEquationCondition -from ..label_tensor import LabelTensor - -class SamplePointDataset(Dataset): +class SamplePointDataset(BaseDataset): """ - This class is used to create a dataset of sample points. + This class extends the BaseDataset to handle physical datasets + composed of only input points. """ - - def __init__(self, problem, device) -> None: - """ - :param dict input_pts: The input points. - """ - super().__init__() - pts_list = [] - self.condition_names = [] - - for name, condition in problem.conditions.items(): - if not hasattr(condition, "output_points"): - pts_list.append(problem.input_pts[name]) - self.condition_names.append(name) - - self.pts = LabelTensor.stack(pts_list) - - if self.pts != []: - self.condition_indeces = torch.cat( - [ - torch.tensor([i] * len(pts_list[i])) - for i in range(len(self.condition_names)) - ], - dim=0, - ) - else: # if there are no sample points - self.condition_indeces = torch.tensor([]) - self.pts = torch.tensor([]) - - self.pts = self.pts.to(device) - self.condition_indeces = self.condition_indeces.to(device) - - def __len__(self): - return self.pts.shape[0] \ No newline at end of file + 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): + 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() diff --git a/pina/data/supervised_dataset.py b/pina/data/supervised_dataset.py new file mode 100644 index 000000000..be601050a --- /dev/null +++ b/pina/data/supervised_dataset.py @@ -0,0 +1,13 @@ +""" +Supervised dataset module +""" +from .base_dataset import BaseDataset + + +class SupervisedDataset(BaseDataset): + """ + This class extends the BaseDataset to handle datasets that consist of + input-output pairs. + """ + data_type = 'supervised' + __slots__ = ['input_points', 'output_points'] diff --git a/pina/data/unsupervised_dataset.py b/pina/data/unsupervised_dataset.py new file mode 100644 index 000000000..18cf296f5 --- /dev/null +++ b/pina/data/unsupervised_dataset.py @@ -0,0 +1,14 @@ +""" +Unsupervised dataset module +""" +from .base_dataset import BaseDataset + + +class UnsupervisedDataset(BaseDataset): + """ + This class extend BaseDataset class to handle + unsupervised dataset,composed of input points + and, optionally, conditional variables + """ + data_type = 'unsupervised' + __slots__ = ['input_points', 'conditional_variables'] diff --git a/pina/domain/cartesian.py b/pina/domain/cartesian.py index 51270549a..5fe99d644 100644 --- a/pina/domain/cartesian.py +++ b/pina/domain/cartesian.py @@ -30,6 +30,10 @@ def __init__(self, cartesian_dict): else: raise TypeError + @property + def sample_modes(self): + return ["random", "grid", "lh", "chebyshev", "latin"] + @property def variables(self): """Spatial variables. @@ -164,9 +168,8 @@ def _1d_sampler(n, mode, variables): for variable in variables: if variable in self.fixed_.keys(): value = self.fixed_[variable] - pts_variable = torch.tensor([[value]]).repeat( - result.shape[0], 1 - ) + pts_variable = torch.tensor([[value] + ]).repeat(result.shape[0], 1) pts_variable = pts_variable.as_subclass(LabelTensor) pts_variable.labels = [variable] @@ -199,9 +202,8 @@ def _Nd_sampler(n, mode, variables): for variable in variables: if variable in self.fixed_.keys(): value = self.fixed_[variable] - pts_variable = torch.tensor([[value]]).repeat( - result.shape[0], 1 - ) + pts_variable = torch.tensor([[value] + ]).repeat(result.shape[0], 1) pts_variable = pts_variable.as_subclass(LabelTensor) pts_variable.labels = [variable] diff --git a/pina/domain/difference_domain.py b/pina/domain/difference_domain.py index d2ba414f0..4015a3860 100644 --- a/pina/domain/difference_domain.py +++ b/pina/domain/difference_domain.py @@ -77,7 +77,7 @@ def sample(self, n, mode="random", variables="all"): 5 """ - if mode != "random": + if mode not in self.sample_modes: raise NotImplementedError( f"{mode} is not a valid mode for sampling." ) diff --git a/pina/domain/domain_interface.py b/pina/domain/domain_interface.py index 5906d2851..916bf3e90 100644 --- a/pina/domain/domain_interface.py +++ b/pina/domain/domain_interface.py @@ -9,6 +9,37 @@ class DomainInterface(metaclass=ABCMeta): Any geometry entity should inherit from this class. """ + available_sampling_modes = ["random", "grid", "lh", "chebyshev", "latin"] + + @property + @abstractmethod + def sample_modes(self): + """ + Abstract method returing available samples modes for the Domain. + """ + pass + + @property + @abstractmethod + def variables(self): + """ + Abstract method returing Domain variables. + """ + pass + + @sample_modes.setter + def sample_modes(self, values): + """ + TODO + """ + if not isinstance(values, (list, tuple)): + values = [values] + for value in values: + if value not in DomainInterface.available_sampling_modes: + raise TypeError(f"mode {value} not valid. Expected at least " + "one in " + f"{DomainInterface.available_sampling_modes}.") + @abstractmethod def sample(self): """ diff --git a/pina/domain/ellipsoid.py b/pina/domain/ellipsoid.py index dcbc55543..18e28d4b7 100644 --- a/pina/domain/ellipsoid.py +++ b/pina/domain/ellipsoid.py @@ -55,7 +55,6 @@ def __init__(self, ellipsoid_dict, sample_surface=False): # perform operation only for not fixed variables (if any) if self.range_: - # convert dict vals to torch [dim, 2] matrix list_dict_vals = list(self.range_.values()) tmp = torch.tensor(list_dict_vals, dtype=torch.float) @@ -71,6 +70,10 @@ def __init__(self, ellipsoid_dict, sample_surface=False): self._centers = dict(zip(self.range_.keys(), centers.tolist())) self._axis = dict(zip(self.range_.keys(), ellipsoid_axis.tolist())) + @property + def sample_modes(self): + return ["random"] + @property def variables(self): """Spatial variables. @@ -281,7 +284,7 @@ def _single_points_sample(n, variables): if variables == "all": variables = list(self.range_.keys()) + list(self.fixed_.keys()) - if mode in ["random"]: + if mode in self.sample_modes: return _Nd_sampler(n, mode, variables) else: raise NotImplementedError(f"mode={mode} is not implemented.") diff --git a/pina/domain/exclusion_domain.py b/pina/domain/exclusion_domain.py index ed63db314..a05b1543e 100644 --- a/pina/domain/exclusion_domain.py +++ b/pina/domain/exclusion_domain.py @@ -76,7 +76,7 @@ def sample(self, n, mode="random", variables="all"): 5 """ - if mode != "random": + if mode not in self.sample_modes: raise NotImplementedError( f"{mode} is not a valid mode for sampling." ) diff --git a/pina/domain/intersection_domain.py b/pina/domain/intersection_domain.py index b40d36950..bb0499b50 100644 --- a/pina/domain/intersection_domain.py +++ b/pina/domain/intersection_domain.py @@ -78,7 +78,7 @@ def sample(self, n, mode="random", variables="all"): 5 """ - if mode != "random": + if mode not in self.sample_modes: raise NotImplementedError( f"{mode} is not a valid mode for sampling." ) diff --git a/pina/domain/operation_interface.py b/pina/domain/operation_interface.py index acd4cf44b..0300f5248 100644 --- a/pina/domain/operation_interface.py +++ b/pina/domain/operation_interface.py @@ -24,6 +24,10 @@ def __init__(self, geometries): # assign geometries self._geometries = geometries + @property + def sample_modes(self): + return ["random"] + @property def geometries(self): """ diff --git a/pina/domain/simplex.py b/pina/domain/simplex.py index 8d26422ae..931f861a7 100644 --- a/pina/domain/simplex.py +++ b/pina/domain/simplex.py @@ -74,6 +74,10 @@ def __init__(self, simplex_matrix, sample_surface=False): # build cartesian_bound self._cartesian_bound = self._build_cartesian(self._vertices_matrix) + @property + def sample_modes(self): + return ["random"] + @property def variables(self): return self._vertices_matrix.labels @@ -88,13 +92,12 @@ def _build_cartesian(self, vertices): """ span_dict = {} - for i, coord in enumerate(self.variables): - sorted_vertices = sorted(vertices, key=lambda vertex: vertex[i]) + sorted_vertices = torch.sort(vertices[coord].tensor.squeeze()) # respective coord bounded by the lowest and highest values span_dict[coord] = [ - float(sorted_vertices[0][i]), - float(sorted_vertices[-1][i]), + float(sorted_vertices.values[0]), + float(sorted_vertices.values[-1]), ] return CartesianDomain(span_dict) @@ -141,7 +144,7 @@ def is_inside(self, point, check_border=False): return all(torch.gt(lambdas, 0.0)) and all(torch.lt(lambdas, 1.0)) return all(torch.ge(lambdas, 0)) and ( - any(torch.eq(lambdas, 0)) or any(torch.eq(lambdas, 1)) + any(torch.eq(lambdas, 0)) or any(torch.eq(lambdas, 1)) ) def _sample_interior_randomly(self, n, variables): @@ -231,7 +234,7 @@ def sample(self, n, mode="random", variables="all"): in ``variables``. """ - if mode in ["random"]: + if mode in self.sample_modes: if self._sample_surface: sample_pts = self._sample_boundary_randomly(n) else: diff --git a/pina/domain/union_domain.py b/pina/domain/union_domain.py index da2ead90d..0af8e1bd1 100644 --- a/pina/domain/union_domain.py +++ b/pina/domain/union_domain.py @@ -33,6 +33,19 @@ def __init__(self, geometries): """ super().__init__(geometries) + @property + def sample_modes(self): + self.sample_modes = list( + set([geom.sample_modes for geom in self.geometries]) + ) + + @property + def variables(self): + variables = [] + for geom in self.geometries: + variables += geom.variables + return list(set(variables)) + def is_inside(self, point, check_border=False): """ Check if a point is inside the ``Union`` domain. diff --git a/pina/label_tensor.py b/pina/label_tensor.py index 7646dd8a3..719975c51 100644 --- a/pina/label_tensor.py +++ b/pina/label_tensor.py @@ -1,13 +1,18 @@ """ Module for LabelTensor """ - +from copy import copy, deepcopy import torch from torch import Tensor + 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)) + if isinstance(a, range) and isinstance(b, range): + return a.start <= b.start and a.stop >= b.stop + return False class LabelTensor(torch.Tensor): @@ -15,16 +20,16 @@ class LabelTensor(torch.Tensor): @staticmethod def __new__(cls, x, labels, *args, **kwargs): - return super().__new__(cls, x, *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 __len__(self) -> int: - return super().__len__() - - def __init__(self, x, labels): + def __init__(self, x, labels, **kwargs): """ Construct a `LabelTensor` by passing a dict of the labels @@ -35,18 +40,145 @@ def __init__(self, x, labels): {1: {"name": "space"['a', 'b', 'c']) """ - self.labels = None + self.dim_names = None + self.full = kwargs.get('full', True) + self.labels = labels + + @classmethod + def __internal_init__(cls, + x, + labels, + dim_names, + *args, + **kwargs): + lt = cls.__new__(cls, x, labels, *args, **kwargs) + lt._labels = labels + lt.full = kwargs.get('full', True) + lt.dim_names = dim_names + return lt + + @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.update_labels(labels) + self._init_labels_from_dict(labels) elif isinstance(labels, list): - self.init_labels_from_list(labels) + self._init_labels_from_list(labels) elif isinstance(labels, str): labels = [labels] - self.init_labels_from_list(labels) + self._init_labels_from_list(labels) else: - raise ValueError(f"labels must be list, dict or string.") + 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: + raise 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 - def extract(self, label_to_extract): + :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.stored_labels + self.dim_names = {} + for dim in labels.keys(): + self.dim_names[labels[dim]['name']] = dim + + 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``. @@ -56,106 +188,123 @@ def extract(self, label_to_extract): :raises TypeError: Labels are not ``str``. :raises ValueError: Label to extract is not in the labels ``list``. """ - from copy import deepcopy - if isinstance(label_to_extract, (str, int)): - label_to_extract = [label_to_extract] - if isinstance(label_to_extract, (tuple, list)): - last_dim_label = self.labels[self.tensor.ndim - 1]['dof'] - if set(label_to_extract).issubset(last_dim_label) is False: - raise ValueError('Cannot extract a dof which is not in the original LabelTensor') - idx_to_extract = [last_dim_label.index(i) for i in label_to_extract] - new_tensor = deepcopy(self.tensor) - new_tensor = new_tensor[..., idx_to_extract] - new_labels = deepcopy(self.labels) - last_dim_new_label = {self.tensor.ndim - 1: { - 'dof': label_to_extract, - 'name': self.labels[self.tensor.ndim - 1]['name'] - }} - new_labels.update(last_dim_new_label) - elif isinstance(label_to_extract, dict): - new_labels = (deepcopy(self.labels)) - new_tensor = deepcopy(self.tensor) - for k, v in label_to_extract.items(): - idx_dim = None - for kl, vl in self.labels.items(): - if vl['name'] == k: - idx_dim = kl - break - dim_labels = self.labels[idx_dim]['dof'] - if isinstance(label_to_extract[k], (int, str)): - label_to_extract[k] = [label_to_extract[k]] - if set(label_to_extract[k]).issubset(dim_labels) is False: - raise ValueError('Cannot extract a dof which is not in the original labeltensor') - idx_to_extract = [dim_labels.index(i) for i in label_to_extract[k]] - indexer = [slice(None)] * idx_dim + [idx_to_extract] + [slice(None)] * (self.tensor.ndim - idx_dim - 1) - new_tensor = new_tensor[indexer] - dim_new_label = {idx_dim: { - 'dof': label_to_extract[k], - 'name': self.labels[idx_dim]['name'] - }} - new_labels.update(dim_new_label) - else: + # Convert str/int to string + if isinstance(labels_to_extract, (str, int)): + labels_to_extract = [labels_to_extract] + + # Store useful variables + labels = self.stored_labels + stored_keys = labels.keys() + dim_names = self.dim_names + ndim = len(super().shape) + + # Convert tuple/list to dict + 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[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') - return LabelTensor(new_tensor, new_labels) + + # 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(): + for key, value in self._labels.items(): s += f"{key}: {value}\n" s += '\n' - s += super().__str__() + 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)` + 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) + :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] - n_dims = tensors[0].ndim - new_labels_cat_dim = [] - for i in range(n_dims): - name = tensors[0].labels[i]['name'] - if i != dim: - dof = tensors[0].labels[i]['dof'] - for tensor in tensors: - dof_to_check = tensor.labels[i]['dof'] - name_to_check = tensor.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.labels[i]['dof'] - name_to_check = tensor.labels[i]['name'] - if name != name_to_check: - raise ValueError('dimensions must have the same dof and name') + # Perform cat on tensors new_tensor = torch.cat(tensors, dim=dim) - labels = tensors[0].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].labels[dim]['name']} - return LabelTensor(new_tensor, labels) + + # Update labels + labels = LabelTensor.__create_labels_cat(tensors, dim) + + return LabelTensor.__internal_init__(new_tensor, labels, + tensors[0].dim_names) + + @staticmethod + def __create_labels_cat(tensors, dim): + # Check if names and dof of the labels are the same in all dimensions + # except in 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) - lt.labels = self.labels + lt.labels = self._labels return lt @property @@ -180,38 +329,161 @@ def clone(self, *args, **kwargs): :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 update_labels(self, labels): + @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) + 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): """ - Update the internal label representation according to the values passed as input. + Stack tensors vertically. For more details, see + :meth:`torch.vstack`. - :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 - """ - self.labels = { - idx_: { - 'dof': range(self.tensor.shape[idx_]), - 'name': idx_ - } for idx_ in range(self.tensor.ndim) - } - tensor_shape = self.tensor.shape - 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') - self.labels.update(labels) + :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) - def init_labels_from_list(self, labels): + def __getitem__(self, index): """ - Given a list of dof, this method update the internal label representation + TODO: Complete docstring + :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) - :param labels: The label(s) to update. - :type labels: list + selected_lt = super().__getitem__(index) + + if isinstance(index, (int, slice)): + index = [index] + + if index[0] == Ellipsis: + index = [slice(None)] * (self.ndim - 1) + [index[1]] + + if hasattr(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_single_label(old_labels, to_update_labels, index, dim): """ - last_dim_labels = {self.tensor.ndim - 1: {'dof': labels, 'name': self.tensor.ndim - 1}} - self.update_labels(last_dim_labels) \ No newline at end of file + 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[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( + as_tuple=True + )[0] if index.dtype == torch.bool else index.tolist() + if isinstance(index, list): + to_update_labels.update({ + dim: { + 'dof': [old_dof[i] for i in index], + 'name': old_labels[dim]['name'] + } + }) + else: + to_update_labels.update( + {dim: { + 'dof': old_dof[index], + 'name': old_labels[dim]['name'] + }}) + + 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 + labels = self.stored_labels[dim]['dof'] + sorted_index = arg_sort(labels) + indexer = [slice(None)] * self.ndim + indexer[dim] = sorted_index + return self.__getitem__(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) + 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) diff --git a/pina/model/avno.py b/pina/model/avno.py index 2ac3b3f7e..e27ce1f1e 100644 --- a/pina/model/avno.py +++ b/pina/model/avno.py @@ -1,7 +1,7 @@ """Module Averaging Neural Operator.""" import torch -from torch import nn, concatenate +from torch import nn, cat from .layers import AVNOBlock from .base_no import KernelNeuralOperator from pina.utils import check_consistency @@ -110,9 +110,9 @@ def forward(self, x): """ points_tmp = x.extract(self.coordinates_indices) new_batch = x.extract(self.field_indices) - new_batch = concatenate((new_batch, points_tmp), dim=-1) + new_batch = cat((new_batch, points_tmp), dim=-1) new_batch = self._lifting_operator(new_batch) new_batch = self._integral_kernels(new_batch) - new_batch = concatenate((new_batch, points_tmp), dim=-1) + new_batch = cat((new_batch, points_tmp), dim=-1) new_batch = self._projection_operator(new_batch) return new_batch diff --git a/pina/model/lno.py b/pina/model/lno.py index 077a6b929..b09f1aec1 100644 --- a/pina/model/lno.py +++ b/pina/model/lno.py @@ -1,7 +1,7 @@ """Module LowRank Neural Operator.""" import torch -from torch import nn, concatenate +from torch import nn, cat from pina.utils import check_consistency @@ -145,4 +145,4 @@ def forward(self, x): for module in self._integral_kernels: x = module(x, coords) # projecting - return self._projection_operator(concatenate((x, coords), dim=-1)) + return self._projection_operator(cat((x, coords), dim=-1)) diff --git a/pina/operators.py b/pina/operators.py index 17b45d814..0b306dfb9 100644 --- a/pina/operators.py +++ b/pina/operators.py @@ -5,9 +5,7 @@ to which computing the operator, the name of the output variables to calculate the operator for (in case of multidimensional functions), and the variables name on which the operator is calculated. """ - import torch - from pina.label_tensor import LabelTensor @@ -58,9 +56,9 @@ def grad_scalar_output(output_, input_, d): gradients = torch.autograd.grad( output_, input_, - grad_outputs=torch.ones( - output_.size(), dtype=output_.dtype, device=output_.device - ), + grad_outputs=torch.ones(output_.size(), + dtype=output_.dtype, + device=output_.device), create_graph=True, retain_graph=True, allow_unused=True, @@ -87,16 +85,13 @@ def grad_scalar_output(output_, input_, d): raise RuntimeError gradients = grad_scalar_output(output_, input_, d) - elif output_.shape[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]) - if i == 0: - gradients = grad_scalar_output(c_output, input_, d) - else: - gradients = gradients.append( - grad_scalar_output(c_output, input_, d) - ) + tensor_to_cat.append(grad_scalar_output(c_output, input_, d)) + gradients = LabelTensor.cat(tensor_to_cat, dim=output_.tensor.ndim - 1) else: raise NotImplementedError @@ -142,16 +137,14 @@ def div(output_, input_, components=None, d=None): raise ValueError grad_output = grad(output_, input_, components, d) - div = torch.zeros(input_.shape[0], 1, device=output_.device) labels = [None] * len(components) + tensors_to_sum = [] for i, (c, d) in enumerate(zip(components, d)): c_fields = f"d{c}d{d}" - div[:, 0] += grad_output.extract(c_fields).sum(axis=1) + tensors_to_sum.append(grad_output.extract(c_fields)) labels[i] = c_fields - - div = div.as_subclass(LabelTensor) - div.labels = ["+".join(labels)] - return div + div_result = LabelTensor.summation(tensors_to_sum) + return div_result def laplacian(output_, input_, components=None, d=None, method="std"): @@ -194,19 +187,19 @@ def laplacian(output_, input_, components=None, d=None, method="std"): if len(components) == 1: grad_output = grad(output_, input_, components=components, d=d) - result = torch.zeros(output_.shape[0], 1, device=output_.device) + to_append_tensors = [] 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 - ) # TODO improve + to_append_tensors.append(gg.extract([gg.labels[i]])) labels = [f"dd{components[0]}"] - + result = LabelTensor.summation(tensors=to_append_tensors) + result.labels = labels else: - result = torch.empty( - input_.shape[0], len(components), device=output_.device - ) + result = torch.empty(input_.shape[0], + len(components), + device=output_.device) labels = [None] * len(components) + to_append_tensors = [None] * len(components) for idx, (ci, di) in enumerate(zip(components, d)): if not isinstance(ci, list): @@ -216,10 +209,11 @@ def laplacian(output_, input_, components=None, d=None, method="std"): grad_output = grad(output_, input_, components=ci, d=di) result[:, idx] = grad(grad_output, input_, d=di).flatten() - labels[idx] = f"dd{ci}dd{di}" - - result = result.as_subclass(LabelTensor) - result.labels = labels + 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.labels = labels return result @@ -249,11 +243,8 @@ def advection(output_, input_, velocity_field, components=None, d=None): if components is None: components = output_.labels - tmp = ( - grad(output_, input_, components, d) - .reshape(-1, len(components), len(d)) - .transpose(0, 1) - ) + tmp = (grad(output_, input_, components, d).reshape(-1, len(components), + len(d)).transpose(0, 1)) tmp *= output_.extract(velocity_field) return tmp.sum(dim=2).T diff --git a/pina/optim/torch_optimizer.py b/pina/optim/torch_optimizer.py index 239819a4f..54818d5a5 100644 --- a/pina/optim/torch_optimizer.py +++ b/pina/optim/torch_optimizer.py @@ -5,6 +5,7 @@ from ..utils import check_consistency from .optimizer_interface import Optimizer + class TorchOptimizer(Optimizer): def __init__(self, optimizer_class, **kwargs): @@ -12,8 +13,8 @@ def __init__(self, optimizer_class, **kwargs): self.optimizer_class = optimizer_class self.kwargs = kwargs + self.optimizer_instance = None def hook(self, parameters): - self.optimizer_instance = self.optimizer_class( - parameters, **self.kwargs - ) \ No newline at end of file + self.optimizer_instance = self.optimizer_class(parameters, + **self.kwargs) diff --git a/pina/optim/torch_scheduler.py b/pina/optim/torch_scheduler.py index 50e1d91f7..9aa187d0c 100644 --- a/pina/optim/torch_scheduler.py +++ b/pina/optim/torch_scheduler.py @@ -5,13 +5,13 @@ from torch.optim.lr_scheduler import LRScheduler # torch >= 2.0 except ImportError: from torch.optim.lr_scheduler import ( - _LRScheduler as LRScheduler, - ) # torch < 2.0 + _LRScheduler as LRScheduler, ) # torch < 2.0 from ..utils import check_consistency from .optimizer_interface import Optimizer from .scheduler_interface import Scheduler + class TorchScheduler(Scheduler): def __init__(self, scheduler_class, **kwargs): @@ -23,5 +23,4 @@ def __init__(self, scheduler_class, **kwargs): def hook(self, optimizer): check_consistency(optimizer, Optimizer) self.scheduler_instance = self.scheduler_class( - optimizer.optimizer_instance, **self.kwargs - ) \ No newline at end of file + optimizer.optimizer_instance, **self.kwargs) diff --git a/pina/problem/abstract_problem.py b/pina/problem/abstract_problem.py index bf0a2dccd..6897fbb74 100644 --- a/pina/problem/abstract_problem.py +++ b/pina/problem/abstract_problem.py @@ -1,11 +1,11 @@ """ Module for AbstractProblem class """ from abc import ABCMeta, abstractmethod -from ..utils import merge_tensors, check_consistency +from ..utils import check_consistency +from ..domain import DomainInterface +from ..condition.domain_equation_condition import DomainEquationCondition +from ..collector import Collector from copy import deepcopy -import torch - -from .. import LabelTensor class AbstractProblem(metaclass=ABCMeta): @@ -20,26 +20,42 @@ class AbstractProblem(metaclass=ABCMeta): def __init__(self): - self._discretized_domains = {} - - for name, domain in self.domains.items(): - if isinstance(domain, (torch.Tensor, LabelTensor)): - self._discretized_domains[name] = domain + # create collector to manage problem data + self._collector = Collector(self) + # create hook conditions <-> problems for condition_name in self.conditions: - self.conditions[condition_name].set_problem(self) + self.conditions[condition_name].problem = self - # # variable storing all points - self.input_pts = {} + # store in collector all the available fixed points + # note that some points could not be stored at this stage (e.g. when + # sampling locations). To check that all data points are ready for + # training all type self.collector.full, which returns true if all + # points are ready. + self.collector.store_fixed_data() + self._batching_dimension = 0 - # # varible to check if sampling is done. If no location - # # element is presented in Condition this variable is set to true - # self._have_sampled_points = {} - for condition_name in self.conditions: - self._discretized_domains[condition_name] = False + @property + def collector(self): + return self._collector - # # put in self.input_pts all the points that we don't need to sample - self._span_condition_points() + @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 + def input_pts(self): + to_return = {} + for k, v in self.collector.data_collections.items(): + if 'input_points' in v.keys(): + to_return[k] = v['input_points'] + return to_return def __deepcopy__(self, memo): """ @@ -85,19 +101,6 @@ def input_variables(self): def input_variables(self, variables): raise RuntimeError - @property - @abstractmethod - def domains(self): - """ - The domain(s) where the conditions of the AbstractProblem are valid. - If more than one domain type is passed, a list of Location is - retured. - - :return: the domain(s) of ``self`` - :rtype: list[Location] - """ - pass - @property @abstractmethod def output_variables(self): @@ -114,35 +117,11 @@ def conditions(self): """ return self._conditions - - - def _span_condition_points(self): - """ - Simple function to get the condition points - """ - for condition_name in self.conditions: - condition = self.conditions[condition_name] - if hasattr(condition, "input_points"): - samples = condition.input_points - self.input_pts[condition_name] = samples - self._discretized_domains[condition_name] = True - if hasattr(self, "unknown_parameter_domain"): - # initialize the unknown parameters of the inverse problem given - # the domain the user gives - self.unknown_parameters = {} - for i, var in enumerate(self.unknown_variables): - range_var = self.unknown_parameter_domain.range_[var] - tensor_var = ( - torch.rand(1, requires_grad=True) * range_var[1] - + range_var[0] - ) - self.unknown_parameters[var] = torch.nn.Parameter( - tensor_var - ) - - def discretise_domain( - self, n, mode="random", variables="all", domains="all" - ): + def discretise_domain(self, + n, + mode="random", + variables="all", + locations="all"): """ Generate a set of points to span the `Location` of all the conditions of the problem. @@ -172,103 +151,44 @@ def discretise_domain( ``CartesianDomain``. """ - # check consistecy n + # check consistecy n, mode, variables, locations check_consistency(n, int) - - # check consistency mode check_consistency(mode, str) - if mode not in ["random", "grid", "lh", "chebyshev", "latin"]: + check_consistency(variables, str) + check_consistency(locations, str) + + # check correct sampling mode + if mode not in DomainInterface.available_sampling_modes: raise TypeError(f"mode {mode} not valid.") - # check consistency variables + # check correct variables if variables == "all": variables = self.input_variables - else: - check_consistency(variables, str) - - if sorted(variables) != sorted(self.input_variables): - TypeError( - f"Wrong variables for sampling. Variables ", - f"should be in {self.input_variables}.", - ) - - # check consistency location - if domains == "all": - domains = [condition for condition in self.conditions] - else: - check_consistency(domains, str) - print(domains) - if sorted(domains) != sorted(self.conditions): - TypeError( - f"Wrong locations for sampling. Location ", - f"should be in {self.conditions}.", - ) - - # sampling - for d in domains: - condition = self.conditions[d] - - # we try to check if we have already sampled - try: - already_sampled = [self.input_pts[d]] - # if we have not sampled, a key error is thrown - except KeyError: - already_sampled = [] - - # if we have already sampled fully the condition - # but we want to sample again we set already_sampled - # to an empty list since we need to sample again, and - # self._have_sampled_points to False. - if self._discretized_domains[d]: - already_sampled = [] - self._discretized_domains[d] = False - print(condition.domain) - print(d) - # build samples - samples = [ - self.domains[d].sample(n=n, mode=mode, variables=variables) - ] + already_sampled - pts = merge_tensors(samples) - self.input_pts[d] = pts - - # the condition is sampled if input_pts contains all labels - if sorted(self.input_pts[d].labels) == sorted( - self.input_variables - ): - self._have_sampled_points[d] = True - - def add_points(self, new_points): - """ - Adding points to the already sampled points. - - :param dict new_points: a dictionary with key the location to add the points - and values the torch.Tensor points. - """ - - if sorted(new_points.keys()) != sorted(self.conditions): - TypeError( - f"Wrong locations for new points. Location ", - f"should be in {self.conditions}.", - ) - - for location in new_points.keys(): - # extract old and new points - old_pts = self.input_pts[location] - new_pts = new_points[location] - - # if they don't have the same variables error - if sorted(old_pts.labels) != sorted(new_pts.labels): + for variable in variables: + if variable not in self.input_variables: TypeError( - f"Not matching variables for old and new points " - f"in condition {location}." + f"Wrong variables for sampling. Variables ", + f"should be in {self.input_variables}.", ) - if old_pts.labels != new_pts.labels: - new_pts = torch.hstack( - [new_pts.extract([i]) for i in old_pts.labels] + + # check correct location + if locations == "all": + locations = [ + name for name in self.conditions.keys() + if isinstance(self.conditions[name], DomainEquationCondition) + ] + else: + if not isinstance(locations, (list)): + locations = [locations] + for loc in locations: + if not isinstance(self.conditions[loc], DomainEquationCondition): + raise TypeError( + f"Wrong locations passed, locations for sampling " + f"should be in {[loc for loc in locations if isinstance(self.conditions[loc], DomainEquationCondition)]}.", ) - new_pts.labels = old_pts.labels - # merging - merged_pts = torch.vstack([old_pts, new_points[location]]) - merged_pts.labels = old_pts.labels - self.input_pts[location] = merged_pts \ No newline at end of file + # store data + self.collector.store_sample_domains(n, mode, variables, locations) + + def add_points(self, new_points_dict): + self.collector.add_points(new_points_dict) diff --git a/pina/problem/inverse_problem.py b/pina/problem/inverse_problem.py index 5a83566ae..51cbd3ca6 100644 --- a/pina/problem/inverse_problem.py +++ b/pina/problem/inverse_problem.py @@ -45,6 +45,20 @@ class InverseProblem(AbstractProblem): >>> 'data': Condition(CartesianDomain({'x': [0, 1]}), Equation(solution_data)) """ + def __init__(self): + super().__init__() + # storing unknown_parameters for optimization + self.unknown_parameters = {} + for i, var in enumerate(self.unknown_variables): + range_var = self.unknown_parameter_domain.range_[var] + tensor_var = ( + torch.rand(1, requires_grad=True) * range_var[1] + + range_var[0] + ) + self.unknown_parameters[var] = torch.nn.Parameter( + tensor_var + ) + @abstractmethod def unknown_parameter_domain(self): """ diff --git a/pina/solvers/pinns/basepinn.py b/pina/solvers/pinns/basepinn.py index 9e841d65d..543f823f9 100644 --- a/pina/solvers/pinns/basepinn.py +++ b/pina/solvers/pinns/basepinn.py @@ -27,13 +27,13 @@ class PINNInterface(SolverInterface, metaclass=ABCMeta): """ def __init__( - self, - models, - problem, - optimizers, - optimizers_kwargs, - extra_features, - loss, + self, + models, + problem, + optimizers, + optimizers_kwargs, + extra_features, + loss, ): """ :param models: Multiple torch neural network models instances. @@ -177,7 +177,7 @@ def compute_residual(self, samples, equation): try: residual = equation.residual(samples, self.forward(samples)) except ( - TypeError + TypeError ): # this occurs when the function has three inputs, i.e. inverse problem residual = equation.residual( samples, self.forward(samples), self._params diff --git a/pina/solvers/solver.py b/pina/solvers/solver.py index a27e93641..e00bc8d59 100644 --- a/pina/solvers/solver.py +++ b/pina/solvers/solver.py @@ -10,168 +10,6 @@ import sys -# class SolverInterface(pytorch_lightning.LightningModule, metaclass=ABCMeta): -# """ -# Solver base class. This class inherits is a wrapper of -# LightningModule class, inheriting all the -# LightningModule methods. -# """ - -# def __init__( -# self, -# models, -# problem, -# optimizers, -# optimizers_kwargs, -# extra_features=None, -# ): -# """ -# :param models: A torch neural network model instance. -# :type models: torch.nn.Module -# :param problem: A problem definition instance. -# :type problem: AbstractProblem -# :param list(torch.optim.Optimizer) optimizer: A list of neural network optimizers to -# use. -# :param list(dict) optimizer_kwargs: A list of optimizer constructor keyword args. -# :param list(torch.nn.Module) extra_features: The additional input -# features to use as augmented input. If ``None`` no extra features -# are passed. If it is a list of :class:`torch.nn.Module`, the extra feature -# list is passed to all models. If it is a list of extra features' lists, -# each single list of extra feature is passed to a model. -# """ -# super().__init__() - -# # check consistency of the inputs -# check_consistency(models, torch.nn.Module) -# check_consistency(problem, AbstractProblem) -# check_consistency(optimizers, torch.optim.Optimizer, subclass=True) -# check_consistency(optimizers_kwargs, dict) - -# # put everything in a list if only one input -# if not isinstance(models, list): -# models = [models] -# if not isinstance(optimizers, list): -# optimizers = [optimizers] -# optimizers_kwargs = [optimizers_kwargs] - -# # number of models and optimizers -# len_model = len(models) -# len_optimizer = len(optimizers) -# len_optimizer_kwargs = len(optimizers_kwargs) - -# # check length consistency optimizers -# if len_model != len_optimizer: -# raise ValueError( -# "You must define one optimizer for each model." -# f"Got {len_model} models, and {len_optimizer}" -# " optimizers." -# ) - -# # check length consistency optimizers kwargs -# if len_optimizer_kwargs != len_optimizer: -# raise ValueError( -# "You must define one dictionary of keyword" -# " arguments for each optimizers." -# f"Got {len_optimizer} optimizers, and" -# f" {len_optimizer_kwargs} dicitionaries" -# ) - -# # extra features handling -# if (extra_features is None) or (len(extra_features) == 0): -# extra_features = [None] * len_model -# else: -# # if we only have a list of extra features -# if not isinstance(extra_features[0], (tuple, list)): -# extra_features = [extra_features] * len_model -# else: # if we have a list of list extra features -# if len(extra_features) != len_model: -# raise ValueError( -# "You passed a list of extrafeatures list with len" -# f"different of models len. Expected {len_model} " -# f"got {len(extra_features)}. If you want to use " -# "the same list of extra features for all models, " -# "just pass a list of extrafeatures and not a list " -# "of list of extra features." -# ) - -# # assigning model and optimizers -# self._pina_models = [] -# self._pina_optimizers = [] - -# for idx in range(len_model): -# model_ = Network( -# model=models[idx], -# input_variables=problem.input_variables, -# output_variables=problem.output_variables, -# extra_features=extra_features[idx], -# ) -# optim_ = optimizers[idx]( -# model_.parameters(), **optimizers_kwargs[idx] -# ) -# self._pina_models.append(model_) -# self._pina_optimizers.append(optim_) - -# # assigning problem -# self._pina_problem = problem - -# @abstractmethod -# def forward(self, *args, **kwargs): -# pass - -# @abstractmethod -# def training_step(self): -# pass - -# @abstractmethod -# def configure_optimizers(self): -# pass - -# @property -# def models(self): -# """ -# The torch model.""" -# return self._pina_models - -# @property -# def optimizers(self): -# """ -# The torch model.""" -# return self._pina_optimizers - -# @property -# def problem(self): -# """ -# The problem formulation.""" -# return self._pina_problem - -# def on_train_start(self): -# """ -# On training epoch start this function is call to do global checks for -# the different solvers. -# """ - -# # 1. Check the verison for dataloader -# dataloader = self.trainer.train_dataloader -# if sys.version_info < (3, 8): -# dataloader = dataloader.loaders -# self._dataloader = dataloader - -# return super().on_train_start() - - # @model.setter - # def model(self, new_model): - # """ - # Set the torch.""" - # check_consistency(new_model, nn.Module, 'torch model') - # self._model= new_model - - # @problem.setter - # def problem(self, problem): - # """ - # Set the problem formulation.""" - # check_consistency(problem, AbstractProblem, 'pina problem') - # self._problem = problem - class SolverInterface(pytorch_lightning.LightningModule, metaclass=ABCMeta): """ Solver base class. This class inherits is a wrapper of @@ -179,13 +17,13 @@ class SolverInterface(pytorch_lightning.LightningModule, metaclass=ABCMeta): LightningModule methods. """ - def __init__( - self, - model, - problem, - optimizer, - scheduler, - ): + def __init__(self, + models, + problem, + optimizers, + schedulers, + extra_features, + use_lt=True): """ :param model: A torch neural network model instance. :type model: torch.nn.Module @@ -197,43 +35,66 @@ def __init__( super().__init__() # check consistency of the inputs - check_consistency(model, torch.nn.Module) check_consistency(problem, AbstractProblem) - check_consistency(optimizer, Optimizer) - check_consistency(scheduler, Scheduler) - - # put everything in a list if only one input - if not isinstance(model, list): - model = [model] - if not isinstance(scheduler, list): - scheduler = [scheduler] - if not isinstance(optimizer, list): - optimizer = [optimizer] - - # number of models and optimizers - len_model = len(model) - len_optimizer = len(optimizer) + self._check_solver_consistency(problem) + + # Check consistency of models argument and encapsulate in list + if not isinstance(models, list): + check_consistency(models, torch.nn.Module) + # put everything in a list if only one input + models = [models] + else: + for idx in range(len(models)): + # Check consistency + check_consistency(models[idx], torch.nn.Module) + len_model = len(models) + + # If use_lt is true add extract operation in input + if use_lt is True: + for idx, model in enumerate(models): + models[idx] = Network( + model=model, + input_variables=problem.input_variables, + output_variables=problem.output_variables, + extra_features=extra_features, + ) + + # Check scheduler consistency + encapsulation + if not isinstance(schedulers, list): + check_consistency(schedulers, Scheduler) + schedulers = [schedulers] + else: + for scheduler in schedulers: + check_consistency(scheduler, Scheduler) + + # Check optimizer consistency + encapsulation + if not isinstance(optimizers, list): + check_consistency(optimizers, Optimizer) + optimizers = [optimizers] + else: + for optimizer in optimizers: + check_consistency(optimizer, Optimizer) + len_optimizer = len(optimizers) # check length consistency optimizers if len_model != len_optimizer: - raise ValueError( - "You must define one optimizer for each model." - f"Got {len_model} models, and {len_optimizer}" - " optimizers." - ) + raise ValueError("You must define one optimizer for each model." + f"Got {len_model} models, and {len_optimizer}" + " optimizers.") # extra features handling + + self._pina_models = models + self._pina_optimizers = optimizers + self._pina_schedulers = schedulers self._pina_problem = problem - self._pina_model = model - self._pina_optimizer = optimizer - self._pina_scheduler = scheduler @abstractmethod def forward(self, *args, **kwargs): pass @abstractmethod - def training_step(self): + def training_step(self, batch, batch_idx): pass @abstractmethod @@ -244,13 +105,13 @@ def configure_optimizers(self): def models(self): """ The torch model.""" - return self._pina_model + return self._pina_models @property def optimizers(self): """ The torch model.""" - return self._pina_optimizer + return self._pina_optimizers @property def problem(self): @@ -272,16 +133,13 @@ def on_train_start(self): return super().on_train_start() - # @model.setter - # def model(self, new_model): - # """ - # Set the torch.""" - # check_consistency(new_model, nn.Module, 'torch model') - # self._model= new_model - - # @problem.setter - # def problem(self, problem): - # """ - # Set the problem formulation.""" - # check_consistency(problem, AbstractProblem, 'pina problem') - # self._problem = problem + def _check_solver_consistency(self, problem): + """ + TODO + """ + for _, condition in problem.conditions.items(): + if not set(self.accepted_condition_types).issubset( + condition.condition_type): + raise ValueError( + f'{self.__name__} support only dose not support condition ' + f'{condition.condition_type}') diff --git a/pina/solvers/supervised.py b/pina/solvers/supervised.py index c44d5a1e2..62fc99149 100644 --- a/pina/solvers/supervised.py +++ b/pina/solvers/supervised.py @@ -2,9 +2,7 @@ import torch from torch.nn.modules.loss import _Loss - - -from ..optim import Optimizer, Scheduler, TorchOptimizer, TorchScheduler +from ..optim import TorchOptimizer, TorchScheduler from .solver import SolverInterface from ..label_tensor import LabelTensor from ..utils import check_consistency @@ -39,15 +37,16 @@ class SupervisedSolver(SolverInterface): we are seeking to approximate multiple (discretised) functions given multiple (discretised) input functions. """ - - def __init__( - self, - problem, - model, - loss=None, - optimizer=None, - scheduler=None, - ): + accepted_condition_types = ['supervised'] + __name__ = 'SupervisedSolver' + + def __init__(self, + problem, + model, + loss=None, + optimizer=None, + scheduler=None, + extra_features=None): """ :param AbstractProblem problem: The formualation of the problem. :param torch.nn.Module model: The neural network model to use. @@ -57,11 +56,8 @@ def __init__( features to use as augmented input. :param torch.optim.Optimizer optimizer: The neural network optimizer to use; default is :class:`torch.optim.Adam`. - :param dict optimizer_kwargs: Optimizer constructor keyword args. - :param float lr: The learning rate; default is 0.001. :param torch.optim.LRScheduler scheduler: Learning rate scheduler. - :param dict scheduler_kwargs: LR scheduler constructor keyword args. """ if loss is None: loss = torch.nn.MSELoss() @@ -70,22 +66,20 @@ def __init__( optimizer = TorchOptimizer(torch.optim.Adam, lr=0.001) if scheduler is None: - scheduler = TorchScheduler( - torch.optim.lr_scheduler.ConstantLR) + scheduler = TorchScheduler(torch.optim.lr_scheduler.ConstantLR) - super().__init__( - model=model, - problem=problem, - optimizer=optimizer, - scheduler=scheduler, - ) + super().__init__(models=model, + problem=problem, + optimizers=optimizer, + schedulers=scheduler, + extra_features=extra_features) # check consistency check_consistency(loss, (LossInterface, _Loss), subclass=False) self._loss = loss - self._model = self._pina_model[0] - self._optimizer = self._pina_optimizer[0] - self._scheduler = self._pina_scheduler[0] + self._model = self._pina_models[0] + self._optimizer = self._pina_optimizers[0] + self._scheduler = self._pina_schedulers[0] def forward(self, x): """Forward pass implementation for the solver. @@ -97,12 +91,7 @@ def forward(self, x): output = self._model(x) - output.labels = { - 1: { - "name": "output", - "dof": self.problem.output_variables - } - } + output.labels = self.problem.output_variables return output def configure_optimizers(self): @@ -113,10 +102,8 @@ def configure_optimizers(self): """ self._optimizer.hook(self._model.parameters()) self._scheduler.hook(self._optimizer) - return ( - [self._optimizer.optimizer_instance], - [self._scheduler.scheduler_instance] - ) + return ([self._optimizer.optimizer_instance], + [self._scheduler.scheduler_instance]) def training_step(self, batch, batch_idx): """Solver training step. @@ -128,34 +115,28 @@ def training_step(self, batch, batch_idx): :return: The sum of the loss functions. :rtype: LabelTensor """ - - condition_idx = batch.condition + condition_idx = batch.supervised.condition_indices for condition_id in range(condition_idx.min(), condition_idx.max() + 1): condition_name = self._dataloader.condition_names[condition_id] condition = self.problem.conditions[condition_name] - pts = batch.input - out = batch.output - + pts = batch.supervised.input_points + out = batch.supervised.output_points if condition_name not in self.problem.conditions: raise RuntimeError("Something wrong happened.") # for data driven mode if not hasattr(condition, "output_points"): raise NotImplementedError( - f"{type(self).__name__} works only in data-driven mode." - ) - + f"{type(self).__name__} works only in data-driven mode.") output_pts = out[condition_idx == condition_id] input_pts = pts[condition_idx == condition_id] input_pts.labels = pts.labels output_pts.labels = out.labels - loss = ( - self.loss_data(input_pts=input_pts, output_pts=output_pts) - ) + loss = self.loss_data(input_pts=input_pts, output_pts=output_pts) loss = loss.as_subclass(torch.Tensor) self.log("mean_loss", float(loss), prog_bar=True, logger=True) @@ -167,8 +148,8 @@ def loss_data(self, input_pts, output_pts): the network output against the true solution. This function should not be override if not intentionally. - :param LabelTensor input_tensor: The input to the neural networks. - :param LabelTensor output_tensor: The true solution to compare the + :param LabelTensor input_pts: The input to the neural networks. + :param LabelTensor output_pts: The true solution to compare the network solution. :return: The residual loss averaged on the input coordinates :rtype: torch.Tensor @@ -181,7 +162,7 @@ def scheduler(self): Scheduler for training. """ return self._scheduler - + @property def optimizer(self): """ diff --git a/pina/trainer.py b/pina/trainer.py index 758bbaaf0..58c66f67c 100644 --- a/pina/trainer.py +++ b/pina/trainer.py @@ -3,13 +3,19 @@ import torch import pytorch_lightning from .utils import check_consistency -from .data import SamplePointDataset, SamplePointLoader, DataPointDataset +from .data import PinaDataModule from .solvers.solver import SolverInterface class Trainer(pytorch_lightning.Trainer): - def __init__(self, solver, batch_size=None, **kwargs): + def __init__(self, + solver, + batch_size=None, + train_size=.7, + test_size=.2, + eval_size=.1, + **kwargs): """ PINA Trainer class for costumizing every aspect of training via flags. @@ -31,35 +37,23 @@ def __init__(self, solver, batch_size=None, **kwargs): check_consistency(solver, SolverInterface) if batch_size is not None: check_consistency(batch_size, int) - - self._model = solver + self.train_size = train_size + self.test_size = test_size + self.eval_size = eval_size + self.solver = solver self.batch_size = batch_size - self._create_loader() self._move_to_device() - # create dataloader - # if solver.problem.have_sampled_points is False: - # raise RuntimeError( - # f"Input points in {solver.problem.not_sampled_points} " - # "training are None. Please " - # "sample points in your problem by calling " - # "discretise_domain function before train " - # "in the provided locations." - # ) - - # self._create_or_update_loader() - def _move_to_device(self): device = self._accelerator_connector._parallel_devices[0] - + # move parameters to device - pb = self._model.problem + pb = self.solver.problem if hasattr(pb, "unknown_parameters"): for key in pb.unknown_parameters: pb.unknown_parameters[key] = torch.nn.Parameter( - pb.unknown_parameters[key].data.to(device) - ) + pb.unknown_parameters[key].data.to(device)) def _create_loader(self): """ @@ -67,29 +61,45 @@ def _create_loader(self): during training, there is no need to define to touch the trainer dataloader, just call the method. """ + if not self.solver.problem.collector.full: + error_message = '\n'.join([ + f"""{" " * 13} ---> Condition {key} {"sampled" if value else + "not sampled"}""" for key, value in + self._solver.problem.collector._is_conditions_ready.items() + ]) + raise RuntimeError('Cannot create Trainer if not all conditions ' + 'are sampled. The Trainer got the following:\n' + f'{error_message}') devices = self._accelerator_connector._parallel_devices if len(devices) > 1: raise RuntimeError("Parallel training is not supported yet.") device = devices[0] - dataset_phys = SamplePointDataset(self._model.problem, device) - dataset_data = DataPointDataset(self._model.problem, device) - self._loader = SamplePointLoader( - dataset_phys, dataset_data, batch_size=self.batch_size, shuffle=True - ) + + data_module = PinaDataModule(problem=self.solver.problem, + device=device, + train_size=self.train_size, + test_size=self.test_size, + val_size=self.eval_size) + data_module.setup() + self._loader = data_module.train_dataloader() def train(self, **kwargs): """ Train the solver method. """ - return super().fit( - self._model, train_dataloaders=self._loader, **kwargs - ) + return super().fit(self.solver, + train_dataloaders=self._loader, + **kwargs) @property def solver(self): """ Returning trainer solver. """ - return self._model + return self._solver + + @solver.setter + def solver(self, solver): + self._solver = solver diff --git a/pina/utils.py b/pina/utils.py index 282dd5332..84d3e7419 100644 --- a/pina/utils.py +++ b/pina/utils.py @@ -40,9 +40,9 @@ def check_consistency(object, object_instance, subclass=False): raise ValueError(f"{type(obj).__name__} must be {object_instance}.") -def number_parameters( - model, aggregate=True, only_trainable=True -): # TODO: check +def number_parameters(model, + aggregate=True, + only_trainable=True): # TODO: check """ Return the number of parameters of a given `model`. @@ -80,9 +80,8 @@ def merge_two_tensors(tensor1, tensor2): n2 = tensor2.shape[0] tensor1 = LabelTensor(tensor1.repeat(n2, 1), labels=tensor1.labels) - tensor2 = LabelTensor( - tensor2.repeat_interleave(n1, dim=0), labels=tensor2.labels - ) + tensor2 = LabelTensor(tensor2.repeat_interleave(n1, dim=0), + labels=tensor2.labels) return tensor1.append(tensor2) diff --git a/tests/test_condition.py b/tests/test_condition.py index 5f1c6236b..9165c3fa1 100644 --- a/tests/test_condition.py +++ b/tests/test_condition.py @@ -18,27 +18,30 @@ def test_init_inputoutput(): Condition(input_points=example_input_pts, output_points=example_output_pts) with pytest.raises(ValueError): Condition(example_input_pts, example_output_pts) - with pytest.raises(TypeError): + with pytest.raises(ValueError): Condition(input_points=3., output_points='example') - with pytest.raises(TypeError): + with pytest.raises(ValueError): Condition(input_points=example_domain, output_points=example_domain) -def test_init_locfunc(): - Condition(location=example_domain, equation=FixedValue(0.0)) +test_init_inputoutput() + + +def test_init_domainfunc(): + Condition(domain=example_domain, equation=FixedValue(0.0)) with pytest.raises(ValueError): Condition(example_domain, FixedValue(0.0)) - with pytest.raises(TypeError): - Condition(location=3., equation='example') - with pytest.raises(TypeError): - Condition(location=example_input_pts, equation=example_output_pts) + with pytest.raises(ValueError): + Condition(domain=3., equation='example') + with pytest.raises(ValueError): + Condition(domain=example_input_pts, equation=example_output_pts) def test_init_inputfunc(): Condition(input_points=example_input_pts, equation=FixedValue(0.0)) with pytest.raises(ValueError): Condition(example_domain, FixedValue(0.0)) - with pytest.raises(TypeError): + with pytest.raises(ValueError): Condition(input_points=3., equation='example') - with pytest.raises(TypeError): + with pytest.raises(ValueError): Condition(input_points=example_domain, equation=example_output_pts) diff --git a/tests/test_dataset.py b/tests/test_dataset.py index 40f219228..87fd9a15b 100644 --- a/tests/test_dataset.py +++ b/tests/test_dataset.py @@ -1,122 +1,227 @@ +import math import torch -import pytest - -from pina.data.dataset import SamplePointDataset, SamplePointLoader, DataPointDataset +from pina.data import SamplePointDataset, SupervisedDataset, PinaDataModule, \ + UnsupervisedDataset +from pina.data import PinaDataLoader from pina import LabelTensor, Condition from pina.equation import Equation from pina.domain import CartesianDomain -from pina.problem import SpatialProblem -from pina.model import FeedForward +from pina.problem import SpatialProblem, AbstractProblem from pina.operators import laplacian from pina.equation.equation_factory import FixedValue +from pina.graph import Graph def laplace_equation(input_, output_): - force_term = (torch.sin(input_.extract(['x'])*torch.pi) * - torch.sin(input_.extract(['y'])*torch.pi)) + force_term = (torch.sin(input_.extract(['x']) * torch.pi) * + torch.sin(input_.extract(['y']) * torch.pi)) delta_u = laplacian(output_.extract(['u']), input_) return delta_u - force_term + my_laplace = Equation(laplace_equation) in_ = LabelTensor(torch.tensor([[0., 1.]]), ['x', 'y']) out_ = LabelTensor(torch.tensor([[0.]]), ['u']) in2_ = LabelTensor(torch.rand(60, 2), ['x', 'y']) out2_ = LabelTensor(torch.rand(60, 1), ['u']) + class Poisson(SpatialProblem): output_variables = ['u'] spatial_domain = CartesianDomain({'x': [0, 1], 'y': [0, 1]}) conditions = { - 'gamma1': Condition( - location=CartesianDomain({'x': [0, 1], 'y': 1}), - equation=FixedValue(0.0)), - 'gamma2': Condition( - location=CartesianDomain({'x': [0, 1], 'y': 0}), - equation=FixedValue(0.0)), - 'gamma3': Condition( - location=CartesianDomain({'x': 1, 'y': [0, 1]}), - equation=FixedValue(0.0)), - 'gamma4': Condition( - location=CartesianDomain({'x': 0, 'y': [0, 1]}), - equation=FixedValue(0.0)), - 'D': Condition( - input_points=LabelTensor(torch.rand(size=(100, 2)), ['x', 'y']), - equation=my_laplace), - 'data': Condition( - input_points=in_, - output_points=out_), - 'data2': Condition( - input_points=in2_, - output_points=out2_) + 'gamma1': + Condition(domain=CartesianDomain({ + 'x': [0, 1], + 'y': 1 + }), + equation=FixedValue(0.0)), + 'gamma2': + Condition(domain=CartesianDomain({ + 'x': [0, 1], + 'y': 0 + }), + equation=FixedValue(0.0)), + 'gamma3': + Condition(domain=CartesianDomain({ + 'x': 1, + 'y': [0, 1] + }), + equation=FixedValue(0.0)), + 'gamma4': + Condition(domain=CartesianDomain({ + 'x': 0, + 'y': [0, 1] + }), + equation=FixedValue(0.0)), + 'D': + Condition(input_points=LabelTensor(torch.rand(size=(100, 2)), + ['x', 'y']), + equation=my_laplace), + 'data': + Condition(input_points=in_, output_points=out_), + 'data2': + Condition(input_points=in2_, output_points=out2_), + 'unsupervised': + Condition( + input_points=LabelTensor(torch.rand(size=(45, 2)), ['x', 'y']), + conditional_variables=LabelTensor(torch.ones(size=(45, 1)), + ['alpha']), + ), + 'unsupervised2': + Condition( + input_points=LabelTensor(torch.rand(size=(90, 2)), ['x', 'y']), + conditional_variables=LabelTensor(torch.ones(size=(90, 1)), + ['alpha']), + ) } + boundaries = ['gamma1', 'gamma2', 'gamma3', 'gamma4'] poisson = Poisson() poisson.discretise_domain(10, 'grid', locations=boundaries) + def test_sample(): sample_dataset = SamplePointDataset(poisson, device='cpu') assert len(sample_dataset) == 140 - assert sample_dataset.pts.shape == (140, 2) - assert sample_dataset.pts.labels == ['x', 'y'] - assert sample_dataset.condition_indeces.dtype == torch.int64 - assert sample_dataset.condition_indeces.max() == torch.tensor(4) - assert sample_dataset.condition_indeces.min() == torch.tensor(0) + assert sample_dataset.input_points.shape == (140, 2) + assert sample_dataset.input_points.labels == ['x', 'y'] + assert sample_dataset.condition_indices.dtype == torch.uint8 + assert sample_dataset.condition_indices.max() == torch.tensor(4) + assert sample_dataset.condition_indices.min() == torch.tensor(0) + def test_data(): - dataset = DataPointDataset(poisson, device='cpu') + dataset = SupervisedDataset(poisson, device='cpu') assert len(dataset) == 61 - assert dataset.input_pts.shape == (61, 2) - assert dataset.input_pts.labels == ['x', 'y'] - assert dataset.output_pts.shape == (61, 1 ) - assert dataset.output_pts.labels == ['u'] - assert dataset.condition_indeces.dtype == torch.int64 - assert dataset.condition_indeces.max() == torch.tensor(1) - assert dataset.condition_indeces.min() == torch.tensor(0) + assert dataset['input_points'].shape == (61, 2) + assert dataset.input_points.shape == (61, 2) + assert dataset['input_points'].labels == ['x', 'y'] + assert dataset.input_points.labels == ['x', 'y'] + assert dataset.input_points[3:].shape == (58, 2) + assert dataset.output_points[:3].labels == ['u'] + assert dataset.output_points.shape == (61, 1) + assert dataset.output_points.labels == ['u'] + assert dataset.condition_indices.dtype == torch.uint8 + assert dataset.condition_indices.max() == torch.tensor(1) + assert dataset.condition_indices.min() == torch.tensor(0) + + +def test_unsupervised(): + dataset = UnsupervisedDataset(poisson, device='cpu') + assert len(dataset) == 135 + assert dataset.input_points.shape == (135, 2) + assert dataset.input_points.labels == ['x', 'y'] + assert dataset.input_points[3:].shape == (132, 2) + + assert dataset.conditional_variables.shape == (135, 1) + assert dataset.conditional_variables.labels == ['alpha'] + assert dataset.condition_indices.dtype == torch.uint8 + assert dataset.condition_indices.max() == torch.tensor(1) + assert dataset.condition_indices.min() == torch.tensor(0) + + +def test_data_module(): + data_module = PinaDataModule(poisson, device='cpu') + data_module.setup() + loader = data_module.train_dataloader() + assert isinstance(loader, PinaDataLoader) + assert isinstance(loader, PinaDataLoader) + + data_module = PinaDataModule(poisson, + device='cpu', + batch_size=10, + shuffle=False) + data_module.setup() + loader = data_module.train_dataloader() + assert len(loader) == 24 + for i in loader: + assert len(i) <= 10 + len_ref = sum( + [math.ceil(len(dataset) * 0.7) for dataset in data_module.datasets]) + len_real = sum( + [len(dataset) for dataset in data_module.splits['train'].values()]) + assert len_ref == len_real + + supervised_dataset = SupervisedDataset(poisson, device='cpu') + data_module = PinaDataModule(poisson, + device='cpu', + batch_size=10, + shuffle=False, + datasets=[supervised_dataset]) + data_module.setup() + loader = data_module.train_dataloader() + for batch in loader: + assert len(batch) <= 10 + + physics_dataset = SamplePointDataset(poisson, device='cpu') + data_module = PinaDataModule(poisson, + device='cpu', + batch_size=10, + shuffle=False, + datasets=[physics_dataset]) + data_module.setup() + loader = data_module.train_dataloader() + for batch in loader: + assert len(batch) <= 10 + + unsupervised_dataset = UnsupervisedDataset(poisson, device='cpu') + data_module = PinaDataModule(poisson, + device='cpu', + batch_size=10, + shuffle=False, + datasets=[unsupervised_dataset]) + data_module.setup() + loader = data_module.train_dataloader() + for batch in loader: + assert len(batch) <= 10 + def test_loader(): - sample_dataset = SamplePointDataset(poisson, device='cpu') - data_dataset = DataPointDataset(poisson, device='cpu') - loader = SamplePointLoader(sample_dataset, data_dataset, batch_size=10) + data_module = PinaDataModule(poisson, device='cpu', batch_size=10) + data_module.setup() + loader = data_module.train_dataloader() + assert isinstance(loader, PinaDataLoader) + assert len(loader) == 24 + for i in loader: + assert len(i) <= 10 + assert i.supervised.input_points.labels == ['x', 'y'] + assert i.physics.input_points.labels == ['x', 'y'] + assert i.unsupervised.input_points.labels == ['x', 'y'] + assert i.supervised.input_points.requires_grad == True + assert i.physics.input_points.requires_grad == True + assert i.unsupervised.input_points.requires_grad == True + + +coordinates = LabelTensor(torch.rand((100, 100, 2)), labels=['x', 'y']) +data = LabelTensor(torch.rand((100, 100, 3)), labels=['ux', 'uy', 'p']) + + +class GraphProblem(AbstractProblem): + output = LabelTensor(torch.rand((100, 3)), labels=['ux', 'uy', 'p']) + input = [ + Graph.build('radius', + nodes_coordinates=coordinates[i, :, :], + nodes_data=data[i, :, :], + radius=0.2) for i in range(100) + ] + output_variables = ['u'] - for batch in loader: - assert len(batch) in [2, 3] - assert batch['pts'].shape[0] <= 10 - assert batch['pts'].requires_grad == True - assert batch['pts'].labels == ['x', 'y'] - - loader2 = SamplePointLoader(sample_dataset, data_dataset, batch_size=None) - assert len(list(loader2)) == 2 - -def test_loader2(): - poisson2 = Poisson() - del poisson.conditions['data2'] - del poisson2.conditions['data'] - poisson2.discretise_domain(10, 'grid', locations=boundaries) - sample_dataset = SamplePointDataset(poisson, device='cpu') - data_dataset = DataPointDataset(poisson, device='cpu') - loader = SamplePointLoader(sample_dataset, data_dataset, batch_size=10) + conditions = { + 'graph_data': Condition(input_points=input, output_points=output) + } - for batch in loader: - assert len(batch) == 2 # only phys condtions - assert batch['pts'].shape[0] <= 10 - assert batch['pts'].requires_grad == True - assert batch['pts'].labels == ['x', 'y'] - -def test_loader3(): - poisson2 = Poisson() - del poisson.conditions['gamma1'] - del poisson.conditions['gamma2'] - del poisson.conditions['gamma3'] - del poisson.conditions['gamma4'] - del poisson.conditions['D'] - sample_dataset = SamplePointDataset(poisson, device='cpu') - data_dataset = DataPointDataset(poisson, device='cpu') - loader = SamplePointLoader(sample_dataset, data_dataset, batch_size=10) - for batch in loader: - assert len(batch) == 2 # only phys condtions - assert batch['pts'].shape[0] <= 10 - assert batch['pts'].requires_grad == True - assert batch['pts'].labels == ['x', 'y'] +graph_problem = GraphProblem() + + +def test_loader_graph(): + data_module = PinaDataModule(graph_problem, device='cpu', batch_size=10) + data_module.setup() + loader = data_module.train_dataloader() + for i in loader: + assert len(i) <= 10 + assert isinstance(i.supervised.input_points, list) + assert all(isinstance(x, Graph) for x in i.supervised.input_points) diff --git a/tests/test_geometry/test_cartesian.py b/tests/test_geometry/test_cartesian.py index 65026c332..fc30757b6 100644 --- a/tests/test_geometry/test_cartesian.py +++ b/tests/test_geometry/test_cartesian.py @@ -3,6 +3,7 @@ from pina import LabelTensor from pina.domain import CartesianDomain + def test_constructor(): CartesianDomain({'x': [0, 1], 'y': [0, 1]}) diff --git a/tests/test_geometry/test_simplex.py b/tests/test_geometry/test_simplex.py index 7fc34ce2d..25224aae3 100644 --- a/tests/test_geometry/test_simplex.py +++ b/tests/test_geometry/test_simplex.py @@ -40,7 +40,6 @@ def test_constructor(): LabelTensor(torch.tensor([[-.5, .5]]), labels=["x", "y"]), ]) - def test_sample(): # sampling inside simplex = SimplexDomain([ diff --git a/tests/test_label_tensor.py b/tests/test_label_tensor.py deleted file mode 100644 index f87d3abb1..000000000 --- a/tests/test_label_tensor.py +++ /dev/null @@ -1,148 +0,0 @@ -import torch -import pytest - -from pina.label_tensor import LabelTensor -#import pina - -data = torch.rand((20, 3)) -labels_column = { - 1: { - "name": "space", - "dof": ['x', 'y', 'z'] - } -} -labels_row = { - 0: { - "name": "samples", - "dof": range(20) - } -} -labels_all = labels_column | labels_row - -@pytest.mark.parametrize("labels", [labels_column, labels_row, labels_all]) -def test_constructor(labels): - LabelTensor(data, labels) - -def test_wrong_constructor(): - with pytest.raises(ValueError): - LabelTensor(data, ['a', 'b']) - -@pytest.mark.parametrize("labels", [labels_column, labels_all]) -@pytest.mark.parametrize("labels_te", ['z', ['z'], {'space': ['z']}]) -def test_extract_column(labels, labels_te): - tensor = LabelTensor(data, labels) - new = tensor.extract(labels_te) - assert new.ndim == tensor.ndim - assert new.shape[1] == 1 - assert new.shape[0] == 20 - assert torch.all(torch.isclose(data[:, 2].reshape(-1, 1), new)) - -@pytest.mark.parametrize("labels", [labels_row, labels_all]) -@pytest.mark.parametrize("labels_te", [{'samples': [2]}]) -def test_extract_row(labels, labels_te): - tensor = LabelTensor(data, labels) - new = tensor.extract(labels_te) - assert new.ndim == tensor.ndim - assert new.shape[1] == 3 - assert new.shape[0] == 1 - assert torch.all(torch.isclose(data[2].reshape(1, -1), new)) - -@pytest.mark.parametrize("labels_te", [ - {'samples': [2], 'space': ['z']}, - {'space': 'z', 'samples': 2} -]) -def test_extract_2D(labels_te): - labels = labels_all - tensor = LabelTensor(data, labels) - new = tensor.extract(labels_te) - assert new.ndim == tensor.ndim - assert new.shape[1] == 1 - assert new.shape[0] == 1 - assert torch.all(torch.isclose(data[2,2].reshape(1, 1), new)) - -def test_extract_3D(): - labels = labels_all - data = torch.rand(20, 3, 4) - labels = { - 1: { - "name": "space", - "dof": ['x', 'y', 'z'] - }, - 2: { - "name": "time", - "dof": range(4) - }, - } - labels_te = { - 'space': ['x', 'z'], - 'time': range(1, 4) - } - - tensor = LabelTensor(data, labels) - new = tensor.extract(labels_te) - assert new.ndim == tensor.ndim - assert new.shape[0] == 20 - assert new.shape[1] == 2 - assert new.shape[2] == 3 - assert torch.all(torch.isclose( - data[:, 0::2, 1:4].reshape(20, 2, 3), - new - )) - -def test_concatenation_3D(): - data_1 = torch.rand(20, 3, 4) - labels_1 = ['x', 'y', 'z', 'w'] - lt1 = LabelTensor(data_1, labels_1) - data_2 = torch.rand(50, 3, 4) - labels_2 = ['x', 'y', 'z', 'w'] - lt2 = LabelTensor(data_2, labels_2) - lt_cat = LabelTensor.cat([lt1, lt2]) - assert lt_cat.shape == (70, 3, 4) - assert lt_cat.labels[0]['dof'] == range(70) - assert lt_cat.labels[1]['dof'] == range(3) - assert lt_cat.labels[2]['dof'] == ['x', 'y', 'z', 'w'] - - data_1 = torch.rand(20, 3, 4) - labels_1 = ['x', 'y', 'z', 'w'] - lt1 = LabelTensor(data_1, labels_1) - data_2 = torch.rand(20, 2, 4) - labels_2 = ['x', 'y', 'z', 'w'] - lt2 = LabelTensor(data_2, labels_2) - lt_cat = LabelTensor.cat([lt1, lt2], dim=1) - assert lt_cat.shape == (20, 5, 4) - assert lt_cat.labels[0]['dof'] == range(20) - assert lt_cat.labels[1]['dof'] == range(5) - assert lt_cat.labels[2]['dof'] == ['x', 'y', 'z', 'w'] - - data_1 = torch.rand(20, 3, 2) - labels_1 = ['x', 'y'] - lt1 = LabelTensor(data_1, labels_1) - data_2 = torch.rand(20, 3, 3) - labels_2 = ['z', 'w', 'a'] - lt2 = LabelTensor(data_2, labels_2) - lt_cat = LabelTensor.cat([lt1, lt2], dim=2) - assert lt_cat.shape == (20, 3, 5) - assert lt_cat.labels[2]['dof'] == ['x', 'y', 'z', 'w', 'a'] - assert lt_cat.labels[0]['dof'] == range(20) - assert lt_cat.labels[1]['dof'] == range(3) - - data_1 = torch.rand(20, 2, 4) - labels_1 = ['x', 'y', 'z', 'w'] - lt1 = LabelTensor(data_1, labels_1) - data_2 = torch.rand(20, 3, 4) - labels_2 = ['x', 'y', 'z', 'w'] - lt2 = LabelTensor(data_2, labels_2) - with pytest.raises(ValueError): - LabelTensor.cat([lt1, lt2], dim=2) - - data_1 = torch.rand(20, 3, 2) - labels_1 = ['x', 'y'] - lt1 = LabelTensor(data_1, labels_1) - data_2 = torch.rand(20, 3, 3) - labels_2 = ['x', 'w', 'a'] - lt2 = LabelTensor(data_2, labels_2) - lt_cat = LabelTensor.cat([lt1, lt2], dim=2) - assert lt_cat.shape == (20, 3, 5) - assert lt_cat.labels[2]['dof'] == range(5) - assert lt_cat.labels[0]['dof'] == range(20) - assert lt_cat.labels[1]['dof'] == range(3) diff --git a/tests/test_label_tensor/test_label_tensor.py b/tests/test_label_tensor/test_label_tensor.py new file mode 100644 index 000000000..61e479951 --- /dev/null +++ b/tests/test_label_tensor/test_label_tensor.py @@ -0,0 +1,282 @@ +import torch +import pytest + +from pina.label_tensor import LabelTensor + +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 + + +@pytest.mark.parametrize("labels", + [labels_column, labels_row, labels_all, labels_list]) +def test_constructor(labels): + print(LabelTensor(data, labels)) + + +def test_wrong_constructor(): + with pytest.raises(ValueError): + LabelTensor(data, ['a', 'b']) + + +@pytest.mark.parametrize("labels", [labels_column, labels_all]) +@pytest.mark.parametrize("labels_te", ['z', ['z'], {'space': ['z']}]) +def test_extract_column(labels, labels_te): + tensor = LabelTensor(data, labels) + new = tensor.extract(labels_te) + assert new.ndim == tensor.ndim + assert new.shape[1] == 1 + assert new.shape[0] == 20 + assert torch.all(torch.isclose(data[:, 2].reshape(-1, 1), new)) + + +@pytest.mark.parametrize("labels", [labels_row, labels_all]) +@pytest.mark.parametrize("labels_te", [{'samples': [2]}]) +def test_extract_row(labels, labels_te): + tensor = LabelTensor(data, labels) + new = tensor.extract(labels_te) + assert new.ndim == tensor.ndim + assert new.shape[1] == 3 + assert new.shape[0] == 1 + assert torch.all(torch.isclose(data[2].reshape(1, -1), new)) + + +@pytest.mark.parametrize("labels_te", [{ + 'samples': [2], + 'space': ['z'] +}, { + 'space': 'z', + 'samples': 2 +}]) +def test_extract_2D(labels_te): + labels = labels_all + tensor = LabelTensor(data, labels) + new = tensor.extract(labels_te) + assert new.ndim == tensor.ndim + assert new.shape[1] == 1 + assert new.shape[0] == 1 + assert torch.all(torch.isclose(data[2, 2].reshape(1, 1), new)) + + +def test_extract_3D(): + data = torch.rand(20, 3, 4) + labels = { + 1: { + "name": "space", + "dof": ['x', 'y', 'z'] + }, + 2: { + "name": "time", + "dof": range(4) + }, + } + labels_te = {'space': ['x', 'z'], 'time': range(1, 4)} + + tensor = LabelTensor(data, labels) + new = tensor.extract(labels_te) + tensor2 = LabelTensor(data, labels) + assert new.ndim == tensor.ndim + assert new.shape[0] == 20 + assert new.shape[1] == 2 + assert new.shape[2] == 3 + assert torch.all(torch.isclose(data[:, 0::2, 1:4].reshape(20, 2, 3), new)) + assert tensor2.ndim == tensor.ndim + assert tensor2.shape == tensor.shape + assert tensor.full_labels == tensor2.full_labels + assert new.shape != tensor.shape + + +def test_concatenation_3D(): + data_1 = torch.rand(20, 3, 4) + labels_1 = ['x', 'y', 'z', 'w'] + lt1 = LabelTensor(data_1, labels_1) + data_2 = torch.rand(50, 3, 4) + labels_2 = ['x', 'y', 'z', 'w'] + lt2 = LabelTensor(data_2, labels_2) + lt_cat = LabelTensor.cat([lt1, lt2]) + assert lt_cat.shape == (70, 3, 4) + assert lt_cat.full_labels[0]['dof'] == range(70) + assert lt_cat.full_labels[1]['dof'] == range(3) + assert lt_cat.full_labels[2]['dof'] == ['x', 'y', 'z', 'w'] + + data_1 = torch.rand(20, 3, 4) + labels_1 = ['x', 'y', 'z', 'w'] + lt1 = LabelTensor(data_1, labels_1) + data_2 = torch.rand(20, 2, 4) + labels_2 = ['x', 'y', 'z', 'w'] + lt2 = LabelTensor(data_2, labels_2) + lt_cat = LabelTensor.cat([lt1, lt2], dim=1) + assert lt_cat.shape == (20, 5, 4) + assert lt_cat.full_labels[0]['dof'] == range(20) + assert lt_cat.full_labels[1]['dof'] == range(5) + assert lt_cat.full_labels[2]['dof'] == ['x', 'y', 'z', 'w'] + + data_1 = torch.rand(20, 3, 2) + labels_1 = ['x', 'y'] + lt1 = LabelTensor(data_1, labels_1) + data_2 = torch.rand(20, 3, 3) + labels_2 = ['z', 'w', 'a'] + lt2 = LabelTensor(data_2, labels_2) + lt_cat = LabelTensor.cat([lt1, lt2], dim=2) + assert lt_cat.shape == (20, 3, 5) + assert lt_cat.full_labels[2]['dof'] == ['x', 'y', 'z', 'w', 'a'] + assert lt_cat.full_labels[0]['dof'] == range(20) + assert lt_cat.full_labels[1]['dof'] == range(3) + + data_1 = torch.rand(20, 2, 4) + labels_1 = ['x', 'y', 'z', 'w'] + lt1 = LabelTensor(data_1, labels_1) + data_2 = torch.rand(20, 3, 4) + labels_2 = ['x', 'y', 'z', 'w'] + lt2 = LabelTensor(data_2, labels_2) + with pytest.raises(RuntimeError): + LabelTensor.cat([lt1, lt2], dim=2) + data_1 = torch.rand(20, 3, 2) + labels_1 = ['x', 'y'] + lt1 = LabelTensor(data_1, labels_1) + data_2 = torch.rand(20, 3, 3) + labels_2 = ['z', 'w', 'a'] + lt2 = LabelTensor(data_2, labels_2) + lt_cat = LabelTensor.cat([lt1, lt2], dim=2) + assert lt_cat.shape == (20, 3, 5) + assert lt_cat.full_labels[2]['dof'] == ['x', 'y', 'z', 'w', 'a'] + assert lt_cat.full_labels[0]['dof'] == range(20) + assert lt_cat.full_labels[1]['dof'] == range(3) + + +def test_summation(): + lt1 = LabelTensor(torch.ones(20, 3), labels_all) + lt2 = LabelTensor(torch.ones(30, 3), ['x', 'y', 'z']) + with pytest.raises(RuntimeError): + LabelTensor.summation([lt1, lt2]) + lt1 = LabelTensor(torch.ones(20, 3), labels_all) + lt2 = LabelTensor(torch.ones(20, 3), labels_all) + lt_sum = LabelTensor.summation([lt1, lt2]) + assert lt_sum.ndim == lt_sum.ndim + assert lt_sum.shape[0] == 20 + assert lt_sum.shape[1] == 3 + assert lt_sum.full_labels[0] == labels_all[0] + assert lt_sum.labels == ['x+x', 'y+y', 'z+z'] + assert torch.eq(lt_sum.tensor, torch.ones(20, 3) * 2).all() + lt1 = LabelTensor(torch.ones(20, 3), labels_all) + lt2 = LabelTensor(torch.ones(20, 3), labels_all) + lt3 = LabelTensor(torch.zeros(20, 3), labels_all) + lt_sum = LabelTensor.summation([lt1, lt2, lt3]) + assert lt_sum.ndim == lt_sum.ndim + assert lt_sum.shape[0] == 20 + assert lt_sum.shape[1] == 3 + assert lt_sum.full_labels[0] == labels_all[0] + assert lt_sum.labels == ['x+x+x', 'y+y+y', 'z+z+z'] + assert torch.eq(lt_sum.tensor, torch.ones(20, 3) * 2).all() + + +def test_append_3D(): + data_1 = torch.rand(20, 3, 2) + labels_1 = ['x', 'y'] + lt1 = LabelTensor(data_1, labels_1) + data_2 = torch.rand(20, 3, 2) + labels_2 = ['z', 'w'] + lt2 = LabelTensor(data_2, labels_2) + lt1 = lt1.append(lt2) + assert lt1.shape == (20, 3, 4) + assert lt1.full_labels[0]['dof'] == range(20) + assert lt1.full_labels[1]['dof'] == range(3) + assert lt1.full_labels[2]['dof'] == ['x', 'y', 'z', 'w'] + + +def test_append_2D(): + data_1 = torch.rand(20, 2) + labels_1 = ['x', 'y'] + lt1 = LabelTensor(data_1, labels_1) + data_2 = torch.rand(20, 2) + labels_2 = ['z', 'w'] + lt2 = LabelTensor(data_2, labels_2) + lt1 = lt1.append(lt2, mode='cross') + assert lt1.shape == (400, 4) + assert lt1.full_labels[0]['dof'] == range(400) + assert lt1.full_labels[1]['dof'] == ['x', 'y', 'z', 'w'] + + +def test_vstack_3D(): + data_1 = torch.rand(20, 3, 2) + labels_1 = { + 1: { + 'dof': ['a', 'b', 'c'], + 'name': 'first' + }, + 2: { + 'dof': ['x', 'y'], + 'name': 'second' + } + } + lt1 = LabelTensor(data_1, labels_1) + data_2 = torch.rand(20, 3, 2) + labels_1 = { + 1: { + 'dof': ['a', 'b', 'c'], + 'name': 'first' + }, + 2: { + 'dof': ['x', 'y'], + 'name': 'second' + } + } + lt2 = LabelTensor(data_2, labels_1) + lt_stacked = LabelTensor.vstack([lt1, lt2]) + assert lt_stacked.shape == (40, 3, 2) + assert lt_stacked.full_labels[0]['dof'] == range(40) + assert lt_stacked.full_labels[1]['dof'] == ['a', 'b', 'c'] + assert lt_stacked.full_labels[2]['dof'] == ['x', 'y'] + assert lt_stacked.full_labels[1]['name'] == 'first' + assert lt_stacked.full_labels[2]['name'] == 'second' + + +def test_vstack_2D(): + data_1 = torch.rand(20, 2) + labels_1 = {1: {'dof': ['x', 'y'], 'name': 'second'}} + lt1 = LabelTensor(data_1, labels_1) + data_2 = torch.rand(20, 2) + labels_1 = {1: {'dof': ['x', 'y'], 'name': 'second'}} + lt2 = LabelTensor(data_2, labels_1) + lt_stacked = LabelTensor.vstack([lt1, lt2]) + assert lt_stacked.shape == (40, 2) + assert lt_stacked.full_labels[0]['dof'] == range(40) + assert lt_stacked.full_labels[1]['dof'] == ['x', 'y'] + assert lt_stacked.full_labels[0]['name'] == 0 + assert lt_stacked.full_labels[1]['name'] == 'second' + + +def test_sorting(): + data = torch.ones(20, 5) + data[:, 0] = data[:, 0] * 4 + data[:, 1] = data[:, 1] * 2 + data[:, 2] = data[:, 2] + data[:, 3] = data[:, 3] * 5 + data[:, 4] = data[:, 4] * 3 + labels = ['d', 'b', 'a', 'e', 'c'] + lt_data = LabelTensor(data, labels) + lt_sorted = LabelTensor.sort_labels(lt_data) + assert lt_sorted.shape == (20, 5) + assert lt_sorted.labels == ['a', 'b', 'c', 'd', 'e'] + assert torch.eq(lt_sorted.tensor[:, 0], torch.ones(20) * 1).all() + assert torch.eq(lt_sorted.tensor[:, 1], torch.ones(20) * 2).all() + assert torch.eq(lt_sorted.tensor[:, 2], torch.ones(20) * 3).all() + assert torch.eq(lt_sorted.tensor[:, 3], torch.ones(20) * 4).all() + assert torch.eq(lt_sorted.tensor[:, 4], torch.ones(20) * 5).all() + + data = torch.ones(20, 4, 5) + data[:, 0, :] = data[:, 0] * 4 + data[:, 1, :] = data[:, 1] * 2 + data[:, 2, :] = data[:, 2] + data[:, 3, :] = data[:, 3] * 3 + labels = {1: {'dof': ['d', 'b', 'a', 'c'], 'name': 1}} + lt_data = LabelTensor(data, labels) + lt_sorted = LabelTensor.sort_labels(lt_data, dim=1) + assert lt_sorted.shape == (20, 4, 5) + assert lt_sorted.full_labels[1]['dof'] == ['a', 'b', 'c', 'd'] + assert torch.eq(lt_sorted.tensor[:, 0, :], torch.ones(20, 5) * 1).all() + assert torch.eq(lt_sorted.tensor[:, 1, :], torch.ones(20, 5) * 2).all() + assert torch.eq(lt_sorted.tensor[:, 2, :], torch.ones(20, 5) * 3).all() + assert torch.eq(lt_sorted.tensor[:, 3, :], torch.ones(20, 5) * 4).all() diff --git a/tests/test_label_tensor/test_label_tensor_01.py b/tests/test_label_tensor/test_label_tensor_01.py new file mode 100644 index 000000000..57aafb8c9 --- /dev/null +++ b/tests/test_label_tensor/test_label_tensor_01.py @@ -0,0 +1,118 @@ +import torch +import pytest + +from pina import LabelTensor + +data = torch.rand((20, 3)) +labels = ['a', 'b', 'c'] + + +def test_constructor(): + LabelTensor(data, labels) + + +def test_wrong_constructor(): + with pytest.raises(ValueError): + LabelTensor(data, ['a', 'b']) + + +def test_labels(): + tensor = LabelTensor(data, labels) + assert isinstance(tensor, torch.Tensor) + assert tensor.labels == labels + with pytest.raises(ValueError): + tensor.labels = labels[:-1] + + +def test_extract(): + label_to_extract = ['a', 'c'] + tensor = LabelTensor(data, labels) + new = tensor.extract(label_to_extract) + assert new.labels == label_to_extract + assert new.shape[1] == len(label_to_extract) + assert torch.all(torch.isclose(data[:, 0::2], new)) + + +def test_extract_onelabel(): + label_to_extract = ['a'] + tensor = LabelTensor(data, labels) + new = tensor.extract(label_to_extract) + assert new.ndim == 2 + assert new.labels == label_to_extract + assert new.shape[1] == len(label_to_extract) + assert torch.all(torch.isclose(data[:, 0].reshape(-1, 1), new)) + + +def test_wrong_extract(): + label_to_extract = ['a', 'cc'] + tensor = LabelTensor(data, labels) + with pytest.raises(ValueError): + tensor.extract(label_to_extract) + + +def test_extract_order(): + label_to_extract = ['c', 'a'] + tensor = LabelTensor(data, labels) + new = tensor.extract(label_to_extract) + expected = torch.cat((data[:, 2].reshape(-1, 1), data[:, 0].reshape(-1, 1)), + dim=1) + assert new.labels == label_to_extract + assert new.shape[1] == len(label_to_extract) + assert torch.all(torch.isclose(expected, new)) + + +def test_merge(): + tensor = LabelTensor(data, labels) + tensor_a = tensor.extract('a') + tensor_b = tensor.extract('b') + tensor_c = tensor.extract('c') + + tensor_bc = tensor_b.append(tensor_c) + assert torch.allclose(tensor_bc, tensor.extract(['b', 'c'])) + + +def test_merge2(): + tensor = LabelTensor(data, labels) + tensor_b = tensor.extract('b') + tensor_c = tensor.extract('c') + + tensor_bc = tensor_b.append(tensor_c) + assert torch.allclose(tensor_bc, tensor.extract(['b', 'c'])) + + +def test_getitem(): + tensor = LabelTensor(data, labels) + tensor_view = tensor['a'] + assert tensor_view.labels == ['a'] + assert torch.allclose(tensor_view.flatten(), data[:, 0]) + + tensor_view = tensor['a', 'c'] + assert tensor_view.labels == ['a', 'c'] + assert torch.allclose(tensor_view, data[:, 0::2]) + + +def test_getitem2(): + tensor = LabelTensor(data, labels) + tensor_view = tensor[:5] + assert tensor_view.labels == labels + assert torch.allclose(tensor_view, data[:5]) + + idx = torch.randperm(tensor.shape[0]) + tensor_view = tensor[idx] + assert tensor_view.labels == labels + + +def test_slice(): + tensor = LabelTensor(data, labels) + tensor_view = tensor[:5, :2] + assert tensor_view.labels == labels[:2] + assert torch.allclose(tensor_view, data[:5, :2]) + + tensor_view2 = tensor[3] + + assert tensor_view2.labels == labels + assert torch.allclose(tensor_view2, data[3]) + + tensor_view3 = tensor[:, 2] + assert tensor_view3.labels == labels[2] + assert torch.allclose(tensor_view3, data[:, 2].reshape(-1, 1)) diff --git a/tests/test_operators.py b/tests/test_operators.py index aa9ea9d8b..e446c8f37 100644 --- a/tests/test_operators.py +++ b/tests/test_operators.py @@ -10,15 +10,14 @@ def func_vec(x): def func_scalar(x): - print('X') x_ = x.extract(['x']) y_ = x.extract(['y']) mu_ = x.extract(['mu']) return x_**2 + y_**2 + mu_**3 -data = torch.rand((20, 3), requires_grad=True) -inp = LabelTensor(data, ['x', 'y', 'mu']) +data = torch.rand((20, 3)) +inp = LabelTensor(data, ['x', 'y', 'mu']).requires_grad_(True) labels = ['a', 'b', 'c'] tensor_v = LabelTensor(func_vec(inp), labels) tensor_s = LabelTensor(func_scalar(inp).reshape(-1, 1), labels[0]) @@ -31,7 +30,7 @@ def test_grad_scalar_output(): f'd{tensor_s.labels[0]}d{i}' for i in inp.labels ] grad_tensor_s = grad(tensor_s, inp, d=['x', 'y']) - assert grad_tensor_s.shape == (inp.shape[0], 2) + assert grad_tensor_s.shape == (20, 2) assert grad_tensor_s.labels == [ f'd{tensor_s.labels[0]}d{i}' for i in ['x', 'y'] ] diff --git a/tests/test_optimizer.py b/tests/test_optimizer.py index 489bbdc05..bdc87cac5 100644 --- a/tests/test_optimizer.py +++ b/tests/test_optimizer.py @@ -1,20 +1,18 @@ - import torch import pytest from pina import TorchOptimizer opt_list = [ - torch.optim.Adam, - torch.optim.AdamW, - torch.optim.SGD, - torch.optim.RMSprop + torch.optim.Adam, torch.optim.AdamW, torch.optim.SGD, torch.optim.RMSprop ] + @pytest.mark.parametrize("optimizer_class", opt_list) def test_constructor(optimizer_class): TorchOptimizer(optimizer_class, lr=1e-3) + @pytest.mark.parametrize("optimizer_class", opt_list) def test_hook(optimizer_class): opt = TorchOptimizer(optimizer_class, lr=1e-3) - opt.hook(torch.nn.Linear(10, 10).parameters()) \ No newline at end of file + opt.hook(torch.nn.Linear(10, 10).parameters()) diff --git a/tests/test_problem.py b/tests/test_problem.py index 3e1e1ee23..f1aafcd98 100644 --- a/tests/test_problem.py +++ b/tests/test_problem.py @@ -27,31 +27,31 @@ class Poisson(SpatialProblem): conditions = { 'gamma1': - Condition(location=CartesianDomain({ + Condition(domain=CartesianDomain({ 'x': [0, 1], 'y': 1 }), equation=FixedValue(0.0)), 'gamma2': - Condition(location=CartesianDomain({ + Condition(domain=CartesianDomain({ 'x': [0, 1], 'y': 0 }), equation=FixedValue(0.0)), 'gamma3': - Condition(location=CartesianDomain({ + Condition(domain=CartesianDomain({ 'x': 1, 'y': [0, 1] }), equation=FixedValue(0.0)), 'gamma4': - Condition(location=CartesianDomain({ + Condition(domain=CartesianDomain({ 'x': 0, 'y': [0, 1] }), equation=FixedValue(0.0)), 'D': - Condition(location=CartesianDomain({ + Condition(domain=CartesianDomain({ 'x': [0, 1], 'y': [0, 1] }), @@ -67,12 +67,9 @@ def poisson_sol(self, pts): truth_solution = poisson_sol -# make the problem -poisson_problem = Poisson() - - def test_discretise_domain(): n = 10 + poisson_problem = Poisson() boundaries = ['gamma1', 'gamma2', 'gamma3', 'gamma4'] poisson_problem.discretise_domain(n, 'grid', locations=boundaries) for b in boundaries: @@ -92,31 +89,68 @@ def test_discretise_domain(): poisson_problem.discretise_domain(n, 'lh', locations=['D']) assert poisson_problem.input_pts['D'].shape[0] == n + poisson_problem.discretise_domain(n) + def test_sampling_few_variables(): n = 10 + poisson_problem = Poisson() poisson_problem.discretise_domain(n, 'grid', locations=['D'], variables=['x']) assert poisson_problem.input_pts['D'].shape[1] == 1 - assert poisson_problem._have_sampled_points['D'] is False + assert poisson_problem.collector._is_conditions_ready['D'] is False -# def test_sampling_all_args(): -# n = 10 -# poisson_problem.discretise_domain(n, 'grid', locations=['D']) +def test_variables_correct_order_sampling(): + n = 10 + poisson_problem = Poisson() + poisson_problem.discretise_domain(n, + 'grid', + locations=['D'], + variables=['x']) + poisson_problem.discretise_domain(n, + 'grid', + locations=['D'], + variables=['y']) + assert poisson_problem.input_pts['D'].labels == sorted( + poisson_problem.input_variables) -# def test_sampling_all_kwargs(): -# n = 10 -# poisson_problem.discretise_domain(n=n, mode='latin', locations=['D']) + poisson_problem.discretise_domain(n, 'grid', locations=['D']) + assert poisson_problem.input_pts['D'].labels == sorted( + poisson_problem.input_variables) + + poisson_problem.discretise_domain(n, + 'grid', + locations=['D'], + variables=['y']) + poisson_problem.discretise_domain(n, + 'grid', + locations=['D'], + variables=['x']) + assert poisson_problem.input_pts['D'].labels == sorted( + poisson_problem.input_variables) -# def test_sampling_dict(): -# n = 10 -# poisson_problem.discretise_domain( -# {'variables': ['x', 'y'], 'mode': 'grid', 'n': n}, locations=['D']) -# def test_sampling_mixed_args_kwargs(): -# n = 10 -# with pytest.raises(ValueError): -# poisson_problem.discretise_domain(n, mode='latin', locations=['D']) +def test_add_points(): + poisson_problem = Poisson() + poisson_problem.discretise_domain(0, + 'random', + locations=['D'], + variables=['x', 'y']) + new_pts = LabelTensor(torch.tensor([[0.5, -0.5]]), labels=['x', 'y']) + poisson_problem.add_points({'D': new_pts}) + assert torch.isclose(poisson_problem.input_pts['D'].extract('x'), + new_pts.extract('x')) + assert torch.isclose(poisson_problem.input_pts['D'].extract('y'), + new_pts.extract('y')) + + +def test_collector(): + poisson_problem = Poisson() + collector = poisson_problem.collector + assert collector.full is False + assert collector._is_conditions_ready['data'] is True + poisson_problem.discretise_domain(10) + assert collector.full is True diff --git a/tests/test_solvers/test_supervised_solver.py b/tests/test_solvers/test_supervised_solver.py index 912480bb8..8ceadcd93 100644 --- a/tests/test_solvers/test_supervised_solver.py +++ b/tests/test_solvers/test_supervised_solver.py @@ -1,51 +1,28 @@ import torch - -from pina.problem import AbstractProblem +import pytest +from pina.problem import AbstractProblem, SpatialProblem from pina import Condition, LabelTensor from pina.solvers import SupervisedSolver -from pina.trainer import Trainer from pina.model import FeedForward -from pina.loss import LpLoss -from pina.solvers import GraphSupervisedSolver +from pina.equation.equation import Equation +from pina.equation.equation_factory import FixedValue +from pina.operators import laplacian +from pina.domain import CartesianDomain +from pina.trainer import Trainer + +in_ = LabelTensor(torch.tensor([[0., 1.]]), ['u_0', 'u_1']) +out_ = LabelTensor(torch.tensor([[0.]]), ['u']) + class NeuralOperatorProblem(AbstractProblem): input_variables = ['u_0', 'u_1'] output_variables = ['u'] - domains = { - 'pts': LabelTensor( - torch.rand(100, 2), - labels={1: {'name': 'space', 'dof': ['u_0', 'u_1']}} - ) - } - conditions = { - 'data' : Condition( - domain='pts', - output_points=LabelTensor( - torch.rand(100, 1), - labels={1: {'name': 'output', 'dof': ['u']}} - ) - ) - } -class NeuralOperatorProblemGraph(AbstractProblem): - input_variables = ['x', 'y', 'u_0', 'u_1'] - output_variables = ['u'] - domains = { - 'pts': LabelTensor( - torch.rand(100, 4), - labels={1: {'name': 'space', 'dof': ['x', 'y', 'u_0', 'u_1']}} - ) - } conditions = { - 'data' : Condition( - domain='pts', - output_points=LabelTensor( - torch.rand(100, 1), - labels={1: {'name': 'output', 'dof': ['u']}} - ) - ) + 'data': Condition(input_points=in_, output_points=out_), } + class myFeature(torch.nn.Module): """ Feature: sin(x) @@ -61,117 +38,106 @@ def forward(self, x): problem = NeuralOperatorProblem() -problem_graph = NeuralOperatorProblemGraph() -# make the problem + extra feats extra_feats = [myFeature()] -model = FeedForward(len(problem.input_variables), - len(problem.output_variables)) +model = FeedForward(len(problem.input_variables), len(problem.output_variables)) model_extra_feats = FeedForward( - len(problem.input_variables) + 1, - len(problem.output_variables)) + len(problem.input_variables) + 1, len(problem.output_variables)) def test_constructor(): SupervisedSolver(problem=problem, model=model) -# def test_constructor_extra_feats(): -# SupervisedSolver(problem=problem, model=model_extra_feats, extra_features=extra_feats) - -''' -class AutoSolver(SupervisedSolver): - - def forward(self, input): - from pina.graph import Graph - print(Graph) - print(input) - if not isinstance(input, Graph): - input = Graph.build('radius', nodes_coordinates=input, nodes_data=torch.rand(input.shape), radius=0.2) - print(input) - print(input.data.edge_index) - print(input.data) - g = self._model(input.data, edge_index=input.data.edge_index) - g.labels = {1: {'name': 'output', 'dof': ['u']}} - return g - du_dt_new = LabelTensor(self.model(graph).reshape(-1,1), labels = ['du']) - - return du_dt_new -''' - -class GraphModel(torch.nn.Module): - def __init__(self, in_channels, out_channels): - from torch_geometric.nn import GCNConv, NNConv - super().__init__() - self.conv1 = GCNConv(in_channels, 16) - self.conv2 = GCNConv(16, out_channels) - - def forward(self, data, edge_index): - print(data) - x = data.x - print(x) - x = self.conv1(x, edge_index) - x = x.relu() - x = self.conv2(x, edge_index) - return x - -def test_graph(): - solver = GraphSupervisedSolver(problem=problem_graph, model=GraphModel(2, 1), loss=LpLoss(), - nodes_coordinates=['x', 'y'], nodes_data=['u_0', 'u_1']) - trainer = Trainer(solver=solver, max_epochs=30, accelerator='cpu', batch_size=20) - trainer.train() +test_constructor() + + +def laplace_equation(input_, output_): + force_term = (torch.sin(input_.extract(['x']) * torch.pi) * + torch.sin(input_.extract(['y']) * torch.pi)) + delta_u = laplacian(output_.extract(['u']), input_) + return delta_u - force_term + + +my_laplace = Equation(laplace_equation) + + +class Poisson(SpatialProblem): + output_variables = ['u'] + spatial_domain = CartesianDomain({'x': [0, 1], 'y': [0, 1]}) + + conditions = { + 'gamma1': + Condition(domain=CartesianDomain({ + 'x': [0, 1], + 'y': 1 + }), + equation=FixedValue(0.0)), + 'gamma2': + Condition(domain=CartesianDomain({ + 'x': [0, 1], + 'y': 0 + }), + equation=FixedValue(0.0)), + 'gamma3': + Condition(domain=CartesianDomain({ + 'x': 1, + 'y': [0, 1] + }), + equation=FixedValue(0.0)), + 'gamma4': + Condition(domain=CartesianDomain({ + 'x': 0, + 'y': [0, 1] + }), + equation=FixedValue(0.0)), + 'D': + Condition(domain=CartesianDomain({ + 'x': [0, 1], + 'y': [0, 1] + }), + equation=my_laplace), + 'data': + Condition(input_points=in_, output_points=out_) + } + + def poisson_sol(self, pts): + return -(torch.sin(pts.extract(['x']) * torch.pi) * + torch.sin(pts.extract(['y']) * torch.pi)) / (2 * torch.pi ** 2) + + truth_solution = poisson_sol + + +def test_wrong_constructor(): + poisson_problem = Poisson() + with pytest.raises(ValueError): + SupervisedSolver(problem=poisson_problem, model=model) def test_train_cpu(): - solver = SupervisedSolver(problem = problem, model=model, loss=LpLoss()) - trainer = Trainer(solver=solver, max_epochs=300, accelerator='cpu', batch_size=20) + solver = SupervisedSolver(problem=problem, model=model) + trainer = Trainer(solver=solver, + max_epochs=200, + accelerator='gpu', + batch_size=5, + train_size=1, + test_size=0., + eval_size=0.) trainer.train() +test_train_cpu() -# def test_train_restore(): -# tmpdir = "tests/tmp_restore" -# solver = SupervisedSolver(problem=problem, -# model=model, -# extra_features=None, -# loss=LpLoss()) -# trainer = Trainer(solver=solver, -# max_epochs=5, -# accelerator='cpu', -# default_root_dir=tmpdir) -# trainer.train() -# ntrainer = Trainer(solver=solver, max_epochs=15, accelerator='cpu') -# t = ntrainer.train( -# ckpt_path=f'{tmpdir}/lightning_logs/version_0/checkpoints/epoch=4-step=5.ckpt') -# import shutil -# shutil.rmtree(tmpdir) - - -# def test_train_load(): -# tmpdir = "tests/tmp_load" -# solver = SupervisedSolver(problem=problem, -# model=model, -# extra_features=None, -# loss=LpLoss()) -# trainer = Trainer(solver=solver, -# max_epochs=15, -# accelerator='cpu', -# default_root_dir=tmpdir) -# trainer.train() -# new_solver = SupervisedSolver.load_from_checkpoint( -# f'{tmpdir}/lightning_logs/version_0/checkpoints/epoch=14-step=15.ckpt', -# problem = problem, model=model) -# test_pts = LabelTensor(torch.rand(20, 2), problem.input_variables) -# assert new_solver.forward(test_pts).shape == (20, 1) -# assert new_solver.forward(test_pts).shape == solver.forward(test_pts).shape -# torch.testing.assert_close( -# new_solver.forward(test_pts), -# solver.forward(test_pts)) -# import shutil -# shutil.rmtree(tmpdir) - -# def test_train_extra_feats_cpu(): -# pinn = SupervisedSolver(problem=problem, -# model=model_extra_feats, -# extra_features=extra_feats) -# trainer = Trainer(solver=pinn, max_epochs=5, accelerator='cpu') -# trainer.train() -test_graph() \ No newline at end of file +def test_extra_features_constructor(): + SupervisedSolver(problem=problem, + model=model_extra_feats, + extra_features=extra_feats) + + +def test_extra_features_train_cpu(): + solver = SupervisedSolver(problem=problem, + model=model_extra_feats, + extra_features=extra_feats) + trainer = Trainer(solver=solver, + max_epochs=200, + accelerator='gpu', + batch_size=5) + trainer.train()