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

Improve conditions and refactor dataset classes (#475)

Browse Source 
* Reimplement conditions

* Refactor datasets and implement LabelBatch

---------

Co-authored-by: Dario Coscia <dariocos99@gmail.com>
This commit is contained in:
Filippo Olivo
2025-03-07 11:24:09 +01:00
committed by Nicola Demo
parent bdad144461
commit a0cbf1c44a
40 changed files with 943 additions and 550 deletions
Download Patch File Download Diff File Expand all files Collapse all files

37
pina/condition/__init__.py
View File

@@ -1,12 +1,41 @@
"""
Module for conditions.
"""
__all__ = [
"Condition",
"ConditionInterface",
"DomainEquationCondition",
"InputPointsEquationCondition",
"InputOutputPointsCondition",
"InputTargetCondition",
"TensorInputTensorTargetCondition",
"TensorInputGraphTargetCondition",
"GraphInputTensorTargetCondition",
"GraphInputGraphTargetCondition",
"InputEquationCondition",
"InputTensorEquationCondition",
"InputGraphEquationCondition",
"DataCondition",
"GraphDataCondition",
"TensorDataCondition",
]
from .condition_interface import ConditionInterface
from .condition import Condition
from .domain_equation_condition import DomainEquationCondition
from .input_equation_condition import InputPointsEquationCondition
from .input_output_condition import InputOutputPointsCondition
from .input_target_condition import (
InputTargetCondition,
TensorInputTensorTargetCondition,
TensorInputGraphTargetCondition,
GraphInputTensorTargetCondition,
GraphInputGraphTargetCondition,
)
from .input_equation_condition import (
InputEquationCondition,
InputTensorEquationCondition,
InputGraphEquationCondition,
)
from .data_condition import (
DataCondition,
GraphDataCondition,
TensorDataCondition,
)

88
pina/condition/condition.py
View File

@@ -1,10 +1,10 @@
"""Condition module."""
from .domain_equation_condition import DomainEquationCondition
from .input_equation_condition import InputPointsEquationCondition
from .input_output_condition import InputOutputPointsCondition
from .data_condition import DataConditionInterface
import warnings
from .data_condition import DataCondition
from .domain_equation_condition import DomainEquationCondition
from .input_equation_condition import InputEquationCondition
from .input_target_condition import InputTargetCondition
from ..utils import custom_warning_format
# Set the custom format for warnings
@@ -12,6 +12,21 @@ warnings.formatwarning = custom_warning_format
warnings.filterwarnings("always", category=DeprecationWarning)
def warning_function(new, old):
"""Handle the deprecation warning.
:param new: Object to use instead of the old one.
:type new: str
:param old: Object to deprecate.
:type old: str
"""
warnings.warn(
f"'{old}' is deprecated and will be removed "
f"in future versions. Please use '{new}' instead.",
DeprecationWarning,
)
class Condition:
"""
The class ``Condition`` is used to represent the constraints (physical
@@ -40,16 +55,32 @@ class Condition:
Example::
>>> TODO
>>> from pina import Condition
>>> condition = Condition(
... input=input,
... target=target
... )
>>> condition = Condition(
... domain=location,
... equation=equation
... )
>>> condition = Condition(
... input=input,
... equation=equation
... )
>>> condition = Condition(
... input=data,
... conditional_variables=conditional_variables
... )
"""
__slots__ = list(
set(
InputOutputPointsCondition.__slots__
+ InputPointsEquationCondition.__slots__
InputTargetCondition.__slots__
+ InputEquationCondition.__slots__
+ DomainEquationCondition.__slots__
+ DataConditionInterface.__slots__
+ DataCondition.__slots__
)
)
@@ -62,25 +93,30 @@ class Condition:
)
# back-compatibility 0.1
if "location" in kwargs.keys():
keys = list(kwargs.keys())
if "location" in keys:
kwargs["domain"] = kwargs.pop("location")
warnings.warn(
f"'location' is deprecated and will be removed "
f"in future versions. Please use 'domain' instead.",
DeprecationWarning,
)
warning_function(new="domain", old="location")
if "input_points" in keys:
kwargs["input"] = kwargs.pop("input_points")
warning_function(new="input", old="input_points")
if "output_points" in keys:
kwargs["target"] = kwargs.pop("output_points")
warning_function(new="target", old="output_points")
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__):
if sorted_keys == sorted(InputTargetCondition.__slots__):
return InputTargetCondition(**kwargs)
if sorted_keys == sorted(InputEquationCondition.__slots__):
return InputEquationCondition(**kwargs)
if 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()}.")
if (
sorted_keys == sorted(DataCondition.__slots__)
or sorted_keys[0] == DataCondition.__slots__[0]
):
return DataCondition(**kwargs)
raise ValueError(f"Invalid keyword arguments {kwargs.keys()}.")

82
pina/condition/condition_interface.py
View File

@@ -1,34 +1,84 @@
"""
Module that defines the ConditionInterface class.
"""
from abc import ABCMeta
from torch_geometric.data import Data
from ..label_tensor import LabelTensor
from ..graph import Graph
class ConditionInterface(metaclass=ABCMeta):
"""
Abstract class which defines a common interface for all the conditions.
"""
condition_types = ["physics", "supervised", "unsupervised"]
def __init__(self, *args, **kwargs):
self._condition_type = None
def __init__(self):
self._problem = None
@property
def problem(self):
"""
Return the problem to which the condition is associated.
:return: Problem to which the condition is associated
:rtype: pina.problem.AbstractProblem
"""
return self._problem
@problem.setter
def problem(self, value):
self._problem = value
@property
def condition_type(self):
return self._condition_type
@staticmethod
def _check_graph_list_consistency(data_list):
@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:
# If the data is a Graph or Data object, return (do not need to check
# anything)
if isinstance(data_list, (Graph, Data)):
return
# check all elements in the list are of the same type
if not all(isinstance(i, (Graph, Data)) for i in data_list):
raise ValueError(
"Invalid input types. "
"Please provide either Data or Graph objects."
)
data = data_list[0]
# Store the keys of the first element in the list
keys = sorted(list(data.keys()))
# Store the type of each tensor inside first element Data/Graph object
data_types = {name: tensor.__class__ for name, tensor in data.items()}
# Store the labels of each LabelTensor inside first element Data/Graph
# object
labels = {
name: tensor.labels
for name, tensor in data.items()
if isinstance(tensor, LabelTensor)
}
# Iterate over the list of Data/Graph objects
for data in data_list[1:]:
# Check if the keys of the current element are the same as the first
# element
if sorted(list(data.keys())) != keys:
raise ValueError(
"Unavailable type of condition, expected one of"
f" {ConditionInterface.condition_types}."
"All elements in the list must have the same keys."
)
self._condition_type = values
for name, tensor in data.items():
# Check if the type of each tensor inside the current element
# is the same as the first element
if tensor.__class__ is not data_types[name]:
raise ValueError(
f"Data {name} must be a {data_types[name]}, got "
f"{tensor.__class__}"
)
# If the tensor is a LabelTensor, check if the labels are the
# same as the first element
if isinstance(tensor, LabelTensor):
if tensor.labels != labels[name]:
raise ValueError(
"LabelTensor must have the same labels"
)

76
pina/condition/data_condition.py
View File

@@ -1,12 +1,15 @@
import torch
"""
DataCondition class
"""
from . import ConditionInterface
import torch
from torch_geometric.data import Data
from .condition_interface import ConditionInterface
from ..label_tensor import LabelTensor
from ..graph import Graph
from ..utils import check_consistency
class DataConditionInterface(ConditionInterface):
class DataCondition(ConditionInterface):
"""
Condition for data. This condition must be used every
time a Unsupervised Loss is needed in the Solver. The conditionalvariable
@@ -14,19 +17,64 @@ class DataConditionInterface(ConditionInterface):
distribution
"""
__slots__ = ["input_points", "conditional_variables"]
__slots__ = ["input", "conditional_variables"]
_avail_input_cls = (torch.Tensor, LabelTensor, Data, Graph, list, tuple)
_avail_conditional_variables_cls = (torch.Tensor, LabelTensor)
def __init__(self, input_points, conditional_variables=None):
def __new__(cls, input, conditional_variables=None):
"""
TODO : add docstring
Instanciate the correct subclass of DataCondition by checking the type
of the input data (input and conditional_variables).
:param input: torch.Tensor or Graph/Data object containing the input
data
:type input: torch.Tensor or Graph or Data
:param conditional_variables: torch.Tensor or LabelTensor containing
the conditional variables
:type conditional_variables: torch.Tensor or LabelTensor
:return: DataCondition subclass
:rtype: TensorDataCondition or GraphDataCondition
"""
if cls != DataCondition:
return super().__new__(cls)
if isinstance(input, (torch.Tensor, LabelTensor)):
subclass = TensorDataCondition
return subclass.__new__(subclass, input, conditional_variables)
if isinstance(input, (Graph, Data, list, tuple)):
cls._check_graph_list_consistency(input)
subclass = GraphDataCondition
return subclass.__new__(subclass, input, conditional_variables)
raise ValueError(
"Invalid input types. "
"Please provide either Data or Graph objects."
)
def __init__(self, input, conditional_variables=None):
"""
Initialize the DataCondition, storing the input and conditional
variables (if any).
:param input: torch.Tensor or Graph/Data object containing the input
data
:type input: torch.Tensor or Graph or Data
:param conditional_variables: torch.Tensor or LabelTensor containing
the conditional variables
:type conditional_variables: torch.Tensor or LabelTensor
"""
super().__init__()
self.input_points = input_points
self.input = input
self.conditional_variables = conditional_variables
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)
class TensorDataCondition(DataCondition):
"""
DataCondition for torch.Tensor input data
"""
class GraphDataCondition(DataCondition):
"""
DataCondition for Graph/Data input data
"""

12
pina/condition/domain_equation_condition.py
View File

@@ -1,4 +1,6 @@
import torch
"""
DomainEquationCondition class definition.
"""
from .condition_interface import ConditionInterface
from ..utils import check_consistency
@@ -16,7 +18,11 @@ class DomainEquationCondition(ConditionInterface):
def __init__(self, domain, equation):
"""
TODO : add docstring
Initialize the DomainEquationCondition, storing the domain and equation.
:param DomainInterface domain: Domain object containing the domain data
:param EquationInterface equation: Equation object containing the
equation data
"""
super().__init__()
self.domain = domain
@@ -29,5 +35,5 @@ class DomainEquationCondition(ConditionInterface):
elif key == "equation":
check_consistency(value, (EquationInterface))
DomainEquationCondition.__dict__[key].__set__(self, value)
elif key in ("_problem", "_condition_type"):
elif key in ("_problem"):
super().__setattr__(key, value)

103
pina/condition/input_equation_condition.py
View File

@@ -1,5 +1,8 @@
import torch
"""
Module to define InputEquationCondition class and its subclasses.
"""
from torch_geometric.data import Data
from .condition_interface import ConditionInterface
from ..label_tensor import LabelTensor
from ..graph import Graph
@@ -7,30 +10,100 @@ from ..utils import check_consistency
from ..equation.equation_interface import EquationInterface
class InputPointsEquationCondition(ConditionInterface):
class InputEquationCondition(ConditionInterface):
"""
Condition for input_points/equation data. This condition must be used every
Condition for input/equation data. This condition must be used every
time a Physics Informed Loss is needed in the Solver.
"""
__slots__ = ["input_points", "equation"]
__slots__ = ["input", "equation"]
_avail_input_cls = (LabelTensor, Graph, list, tuple)
_avail_equation_cls = EquationInterface
def __init__(self, input_points, equation):
def __new__(cls, input, equation):
"""
TODO : add docstring
Instanciate the correct subclass of InputEquationCondition by checking
the type of the input data (only `input`).
:param input: torch.Tensor or Graph/Data object containing the input
:type input: torch.Tensor or Graph or Data
:param EquationInterface equation: Equation object containing the
equation function
:return: InputEquationCondition subclass
:rtype: InputTensorEquationCondition or InputGraphEquationCondition
"""
# If the class is already a subclass, return the instance
if cls != InputEquationCondition:
return super().__new__(cls)
# Instanciate the correct subclass
if isinstance(input, (Graph, Data, list, tuple)):
subclass = InputGraphEquationCondition
cls._check_graph_list_consistency(input)
subclass._check_label_tensor(input)
return subclass.__new__(subclass, input, equation)
if isinstance(input, LabelTensor):
subclass = InputTensorEquationCondition
return subclass.__new__(subclass, input, equation)
# If the input is not a LabelTensor or a Graph object raise an error
raise ValueError(
"The input data object must be a LabelTensor or a Graph object."
)
def __init__(self, input, equation):
"""
Initialize the InputEquationCondition by storing the input and equation.
:param input: torch.Tensor or Graph/Data object containing the input
:type input: torch.Tensor or Graph or Data
:param EquationInterface equation: Equation object containing the
equation function
"""
super().__init__()
self.input_points = input_points
self.input = input
self.equation = equation
def __setattr__(self, key, value):
if key == "input_points":
check_consistency(
value, (LabelTensor)
) # for now only labeltensors, we need labels for the operator!
InputPointsEquationCondition.__dict__[key].__set__(self, value)
if key == "input":
check_consistency(value, self._avail_input_cls)
InputEquationCondition.__dict__[key].__set__(self, value)
elif key == "equation":
check_consistency(value, (EquationInterface))
InputPointsEquationCondition.__dict__[key].__set__(self, value)
elif key in ("_problem", "_condition_type"):
check_consistency(value, self._avail_equation_cls)
InputEquationCondition.__dict__[key].__set__(self, value)
elif key in ("_problem"):
super().__setattr__(key, value)
class InputTensorEquationCondition(InputEquationCondition):
"""
InputEquationCondition subclass for LabelTensor input data.
"""
class InputGraphEquationCondition(InputEquationCondition):
"""
InputEquationCondition subclass for Graph input data.
"""
@staticmethod
def _check_label_tensor(input):
"""
Check if the input is a LabelTensor.
:param input: input data
:type input: torch.Tensor or Graph or Data
"""
# Store the fist element of the list/tuple if input is a list/tuple
# it is anougth to check the first element because all elements must
# have the same type and structure (already checked)
data = input[0] if isinstance(input, (list, tuple)) else input
# Check if the input data contains at least one LabelTensor
for v in data.values():
if isinstance(v, LabelTensor):
return
raise ValueError(
"The input data object must contain at least one LabelTensor."
)

34
pina/condition/input_output_condition.py
View File

@@ -1,34 +0,0 @@
import torch
import torch_geometric
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 : add docstring
"""
super().__init__()
self.input_points = input_points
self.output_points = output_points
def __setattr__(self, key, value):
if (key == "input_points") or (key == "output_points"):
check_consistency(
value,
(LabelTensor, Graph, torch.Tensor, torch_geometric.data.Data),
)
InputOutputPointsCondition.__dict__[key].__set__(self, value)
elif key in ("_problem", "_condition_type"):
super().__setattr__(key, value)

121
pina/condition/input_target_condition.py Normal file
View File

@@ -0,0 +1,121 @@
"""
This module contains condition classes for supervised learning tasks.
"""
import torch
from torch_geometric.data import Data
from ..label_tensor import LabelTensor
from ..graph import Graph
from .condition_interface import ConditionInterface
class InputTargetCondition(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", "target"]
_avail_input_cls = (torch.Tensor, LabelTensor, Data, Graph, list, tuple)
_avail_output_cls = (torch.Tensor, LabelTensor, Data, Graph, list, tuple)
def __new__(cls, input, target):
"""
Instanciate the correct subclass of InputTargetCondition by checking the
type of the input and target data.
:param input: torch.Tensor or Graph/Data object containing the input
:type input: torch.Tensor or Graph or Data
:param target: torch.Tensor or Graph/Data object containing the target
:type target: torch.Tensor or Graph or Data
:return: InputTargetCondition subclass
:rtype: TensorInputTensorTargetCondition or
TensorInputGraphTargetCondition or GraphInputTensorTargetCondition
or GraphInputGraphTargetCondition
"""
if cls != InputTargetCondition:
return super().__new__(cls)
if isinstance(input, (torch.Tensor, LabelTensor)) and isinstance(
target, (torch.Tensor, LabelTensor)
):
subclass = TensorInputTensorTargetCondition
return subclass.__new__(subclass, input, target)
if isinstance(input, (torch.Tensor, LabelTensor)) and isinstance(
target, (Graph, Data, list, tuple)
):
cls._check_graph_list_consistency(target)
subclass = TensorInputGraphTargetCondition
return subclass.__new__(subclass, input, target)
if isinstance(input, (Graph, Data, list, tuple)) and isinstance(
target, (torch.Tensor, LabelTensor)
):
cls._check_graph_list_consistency(input)
subclass = GraphInputTensorTargetCondition
return subclass.__new__(subclass, input, target)
if isinstance(input, (Graph, Data, list, tuple)) and isinstance(
target, (Graph, Data, list, tuple)
):
cls._check_graph_list_consistency(input)
cls._check_graph_list_consistency(target)
subclass = GraphInputGraphTargetCondition
return subclass.__new__(subclass, input, target)
raise ValueError(
"Invalid input/target types. "
"Please provide either Data, Graph, LabelTensor or torch.Tensor "
"objects."
)
def __init__(self, input, target):
"""
Initialize the InputTargetCondition, storing the input and target data.
:param input: torch.Tensor or Graph/Data object containing the input
:type input: torch.Tensor or Graph or Data
:param target: torch.Tensor or Graph/Data object containing the target
:type target: torch.Tensor or Graph or Data
"""
super().__init__()
self._check_input_target_len(input, target)
self.input = input
self.target = target
@staticmethod
def _check_input_target_len(input, target):
if isinstance(input, (Graph, Data)) or isinstance(
target, (Graph, Data)
):
return
if len(input) != len(target):
raise ValueError(
"The input and target lists must have the same length."
)
class TensorInputTensorTargetCondition(InputTargetCondition):
"""
InputTargetCondition subclass for torch.Tensor input and target data.
"""
class TensorInputGraphTargetCondition(InputTargetCondition):
"""
InputTargetCondition subclass for torch.Tensor input and Graph/Data target
data.
"""
class GraphInputTensorTargetCondition(InputTargetCondition):
"""
InputTargetCondition subclass for Graph/Data input and torch.Tensor target
data.
"""
class GraphInputGraphTargetCondition(InputTargetCondition):
"""
InputTargetCondition subclass for Graph/Data input and target data.
"""