From b8f5c22880988223c1a9befabca6e931f7a55cf3 Mon Sep 17 00:00:00 2001 From: Dario Coscia Date: Fri, 5 Apr 2024 17:33:29 +0200 Subject: [PATCH 01/14] fixing adaptive functions --- docs/source/_rst/_code.rst | 22 +- .../_rst/adaptive_functions/AdaptiveCELU.rst | 9 + .../_rst/adaptive_functions/AdaptiveELU.rst | 9 + .../_rst/adaptive_functions/AdaptiveExp.rst | 9 + .../AdaptiveFunctionInterface.rst | 8 + .../_rst/adaptive_functions/AdaptiveGELU.rst | 9 + .../_rst/adaptive_functions/AdaptiveMish.rst | 9 + .../_rst/adaptive_functions/AdaptiveReLU.rst | 9 + .../_rst/adaptive_functions/AdaptiveSIREN.rst | 9 + .../_rst/adaptive_functions/AdaptiveSiLU.rst | 9 + .../adaptive_functions/AdaptiveSigmoid.rst | 9 + .../adaptive_functions/AdaptiveSoftmax.rst | 9 + .../adaptive_functions/AdaptiveSoftmin.rst | 9 + .../_rst/adaptive_functions/AdaptiveTanh.rst | 9 + docs/source/_rst/layers/adaptive_func.rst | 7 - pina/adaptive_functions/__init__.py | 21 + pina/adaptive_functions/adaptive_func.py | 488 ++++++++++++++++++ .../adaptive_func_interface.py} | 74 ++- pina/model/layers/__init__.py | 4 +- tests/test_adaptive_functions.py | 62 +++ tests/test_layers/test_adaptive_func.py | 48 -- 21 files changed, 743 insertions(+), 99 deletions(-) create mode 100644 docs/source/_rst/adaptive_functions/AdaptiveCELU.rst create mode 100644 docs/source/_rst/adaptive_functions/AdaptiveELU.rst create mode 100644 docs/source/_rst/adaptive_functions/AdaptiveExp.rst create mode 100644 docs/source/_rst/adaptive_functions/AdaptiveFunctionInterface.rst create mode 100644 docs/source/_rst/adaptive_functions/AdaptiveGELU.rst create mode 100644 docs/source/_rst/adaptive_functions/AdaptiveMish.rst create mode 100644 docs/source/_rst/adaptive_functions/AdaptiveReLU.rst create mode 100644 docs/source/_rst/adaptive_functions/AdaptiveSIREN.rst create mode 100644 docs/source/_rst/adaptive_functions/AdaptiveSiLU.rst create mode 100644 docs/source/_rst/adaptive_functions/AdaptiveSigmoid.rst create mode 100644 docs/source/_rst/adaptive_functions/AdaptiveSoftmax.rst create mode 100644 docs/source/_rst/adaptive_functions/AdaptiveSoftmin.rst create mode 100644 docs/source/_rst/adaptive_functions/AdaptiveTanh.rst delete mode 100644 docs/source/_rst/layers/adaptive_func.rst create mode 100644 pina/adaptive_functions/__init__.py create mode 100644 pina/adaptive_functions/adaptive_func.py rename pina/{model/layers/adaptive_func.py => adaptive_functions/adaptive_func_interface.py} (68%) create mode 100644 tests/test_adaptive_functions.py delete mode 100644 tests/test_layers/test_adaptive_func.py diff --git a/docs/source/_rst/_code.rst b/docs/source/_rst/_code.rst index 74c78527d..77072a507 100644 --- a/docs/source/_rst/_code.rst +++ b/docs/source/_rst/_code.rst @@ -74,7 +74,27 @@ Layers Continuous convolution Proper Orthogonal Decomposition Periodic Boundary Condition embeddings - Adpative Activation Function + +Adaptive Activation Functions +------------------------------- + +.. toctree:: + :titlesonly: + + Adaptive Function Interface + Adaptive ReLU + Adaptive Sigmoid + Adaptive Tanh + Adaptive SiLU + Adaptive Mish + Adaptive ELU + Adaptive CELU + Adaptive GELU + Adaptive Softmin + Adaptive Softmax + Adaptive SIREN + Adaptive Exp + Equations and Operators ------------------------- diff --git a/docs/source/_rst/adaptive_functions/AdaptiveCELU.rst b/docs/source/_rst/adaptive_functions/AdaptiveCELU.rst new file mode 100644 index 000000000..9736ee631 --- /dev/null +++ b/docs/source/_rst/adaptive_functions/AdaptiveCELU.rst @@ -0,0 +1,9 @@ +AdaptiveCELU +============ + +.. currentmodule:: pina.adaptive_functions.adaptive_func + +.. autoclass:: AdaptiveCELU + :members: + :show-inheritance: + :inherited-members: AdaptiveActivationFunctionInterface diff --git a/docs/source/_rst/adaptive_functions/AdaptiveELU.rst b/docs/source/_rst/adaptive_functions/AdaptiveELU.rst new file mode 100644 index 000000000..ad04717f1 --- /dev/null +++ b/docs/source/_rst/adaptive_functions/AdaptiveELU.rst @@ -0,0 +1,9 @@ +AdaptiveELU +=========== + +.. currentmodule:: pina.adaptive_functions.adaptive_func + +.. autoclass:: AdaptiveELU + :members: + :show-inheritance: + :inherited-members: AdaptiveActivationFunctionInterface diff --git a/docs/source/_rst/adaptive_functions/AdaptiveExp.rst b/docs/source/_rst/adaptive_functions/AdaptiveExp.rst new file mode 100644 index 000000000..7d07cd52d --- /dev/null +++ b/docs/source/_rst/adaptive_functions/AdaptiveExp.rst @@ -0,0 +1,9 @@ +AdaptiveExp +=========== + +.. currentmodule:: pina.adaptive_functions.adaptive_func + +.. autoclass:: AdaptiveExp + :members: + :show-inheritance: + :inherited-members: AdaptiveActivationFunctionInterface diff --git a/docs/source/_rst/adaptive_functions/AdaptiveFunctionInterface.rst b/docs/source/_rst/adaptive_functions/AdaptiveFunctionInterface.rst new file mode 100644 index 000000000..7cdf754b7 --- /dev/null +++ b/docs/source/_rst/adaptive_functions/AdaptiveFunctionInterface.rst @@ -0,0 +1,8 @@ +AdaptiveActivationFunctionInterface +======================================= + +.. currentmodule:: pina.adaptive_functions.adaptive_func_interface + +.. automodule:: pina.adaptive_functions.adaptive_func_interface + :members: + :show-inheritance: diff --git a/docs/source/_rst/adaptive_functions/AdaptiveGELU.rst b/docs/source/_rst/adaptive_functions/AdaptiveGELU.rst new file mode 100644 index 000000000..86e587584 --- /dev/null +++ b/docs/source/_rst/adaptive_functions/AdaptiveGELU.rst @@ -0,0 +1,9 @@ +AdaptiveGELU +============ + +.. currentmodule:: pina.adaptive_functions.adaptive_func + +.. autoclass:: AdaptiveGELU + :members: + :show-inheritance: + :inherited-members: AdaptiveActivationFunctionInterface diff --git a/docs/source/_rst/adaptive_functions/AdaptiveMish.rst b/docs/source/_rst/adaptive_functions/AdaptiveMish.rst new file mode 100644 index 000000000..4e1e3b435 --- /dev/null +++ b/docs/source/_rst/adaptive_functions/AdaptiveMish.rst @@ -0,0 +1,9 @@ +AdaptiveMish +============ + +.. currentmodule:: pina.adaptive_functions.adaptive_func + +.. autoclass:: AdaptiveMish + :members: + :show-inheritance: + :inherited-members: AdaptiveActivationFunctionInterface diff --git a/docs/source/_rst/adaptive_functions/AdaptiveReLU.rst b/docs/source/_rst/adaptive_functions/AdaptiveReLU.rst new file mode 100644 index 000000000..ea08c29a9 --- /dev/null +++ b/docs/source/_rst/adaptive_functions/AdaptiveReLU.rst @@ -0,0 +1,9 @@ +AdaptiveReLU +============ + +.. currentmodule:: pina.adaptive_functions.adaptive_func + +.. autoclass:: AdaptiveReLU + :members: + :show-inheritance: + :inherited-members: AdaptiveActivationFunctionInterface diff --git a/docs/source/_rst/adaptive_functions/AdaptiveSIREN.rst b/docs/source/_rst/adaptive_functions/AdaptiveSIREN.rst new file mode 100644 index 000000000..96133bdd8 --- /dev/null +++ b/docs/source/_rst/adaptive_functions/AdaptiveSIREN.rst @@ -0,0 +1,9 @@ +AdaptiveSIREN +============= + +.. currentmodule:: pina.adaptive_functions.adaptive_func + +.. autoclass:: AdaptiveSIREN + :members: + :show-inheritance: + :inherited-members: AdaptiveActivationFunctionInterface diff --git a/docs/source/_rst/adaptive_functions/AdaptiveSiLU.rst b/docs/source/_rst/adaptive_functions/AdaptiveSiLU.rst new file mode 100644 index 000000000..2f359fded --- /dev/null +++ b/docs/source/_rst/adaptive_functions/AdaptiveSiLU.rst @@ -0,0 +1,9 @@ +AdaptiveSiLU +============ + +.. currentmodule:: pina.adaptive_functions.adaptive_func + +.. autoclass:: AdaptiveSiLU + :members: + :show-inheritance: + :inherited-members: AdaptiveActivationFunctionInterface diff --git a/docs/source/_rst/adaptive_functions/AdaptiveSigmoid.rst b/docs/source/_rst/adaptive_functions/AdaptiveSigmoid.rst new file mode 100644 index 000000000..6f495a8ed --- /dev/null +++ b/docs/source/_rst/adaptive_functions/AdaptiveSigmoid.rst @@ -0,0 +1,9 @@ +AdaptiveSigmoid +=============== + +.. currentmodule:: pina.adaptive_functions.adaptive_func + +.. autoclass:: AdaptiveSigmoid + :members: + :show-inheritance: + :inherited-members: AdaptiveActivationFunctionInterface diff --git a/docs/source/_rst/adaptive_functions/AdaptiveSoftmax.rst b/docs/source/_rst/adaptive_functions/AdaptiveSoftmax.rst new file mode 100644 index 000000000..5cab9c65c --- /dev/null +++ b/docs/source/_rst/adaptive_functions/AdaptiveSoftmax.rst @@ -0,0 +1,9 @@ +AdaptiveSoftmax +=============== + +.. currentmodule:: pina.adaptive_functions.adaptive_func + +.. autoclass:: AdaptiveSoftmax + :members: + :show-inheritance: + :inherited-members: AdaptiveActivationFunctionInterface diff --git a/docs/source/_rst/adaptive_functions/AdaptiveSoftmin.rst b/docs/source/_rst/adaptive_functions/AdaptiveSoftmin.rst new file mode 100644 index 000000000..a0e6c94ae --- /dev/null +++ b/docs/source/_rst/adaptive_functions/AdaptiveSoftmin.rst @@ -0,0 +1,9 @@ +AdaptiveSoftmin +=============== + +.. currentmodule:: pina.adaptive_functions.adaptive_func + +.. autoclass:: AdaptiveSoftmin + :members: + :show-inheritance: + :inherited-members: AdaptiveActivationFunctionInterface diff --git a/docs/source/_rst/adaptive_functions/AdaptiveTanh.rst b/docs/source/_rst/adaptive_functions/AdaptiveTanh.rst new file mode 100644 index 000000000..3e486512f --- /dev/null +++ b/docs/source/_rst/adaptive_functions/AdaptiveTanh.rst @@ -0,0 +1,9 @@ +AdaptiveTanh +============ + +.. currentmodule:: pina.adaptive_functions.adaptive_func + +.. autoclass:: AdaptiveTanh + :members: + :show-inheritance: + :inherited-members: AdaptiveActivationFunctionInterface diff --git a/docs/source/_rst/layers/adaptive_func.rst b/docs/source/_rst/layers/adaptive_func.rst deleted file mode 100644 index 278f9f7b9..000000000 --- a/docs/source/_rst/layers/adaptive_func.rst +++ /dev/null @@ -1,7 +0,0 @@ -AdaptiveActivationFunction -============================= -.. currentmodule:: pina.model.layers.adaptive_func - -.. autoclass:: AdaptiveActivationFunction - :members: - :show-inheritance: \ No newline at end of file diff --git a/pina/adaptive_functions/__init__.py b/pina/adaptive_functions/__init__.py new file mode 100644 index 000000000..0ab6053ed --- /dev/null +++ b/pina/adaptive_functions/__init__.py @@ -0,0 +1,21 @@ +__all__ = [ + 'AdaptiveActivationFunctionInterface', + 'AdaptiveReLU', + 'AdaptiveSigmoid', + 'AdaptiveTanh', + 'AdaptiveSiLU', + 'AdaptiveMish', + 'AdaptiveELU', + 'AdaptiveCELU', + 'AdaptiveGELU', + 'AdaptiveSoftmin', + 'AdaptiveSoftmax', + 'AdaptiveSIREN', + 'AdaptiveExp'] + +from .adaptive_func import (AdaptiveReLU, AdaptiveSigmoid, AdaptiveTanh, + AdaptiveSiLU, AdaptiveMish, AdaptiveELU, + AdaptiveCELU, AdaptiveGELU, AdaptiveSoftmin, + AdaptiveSoftmax, AdaptiveSIREN, AdaptiveExp) +from .adaptive_func_interface import AdaptiveActivationFunctionInterface + diff --git a/pina/adaptive_functions/adaptive_func.py b/pina/adaptive_functions/adaptive_func.py new file mode 100644 index 000000000..0ee22b2d0 --- /dev/null +++ b/pina/adaptive_functions/adaptive_func.py @@ -0,0 +1,488 @@ +""" Module for adaptive functions. """ + +import torch +from ..utils import check_consistency +from .adaptive_func_interface import AdaptiveActivationFunctionInterface + + +class AdaptiveReLU(AdaptiveActivationFunctionInterface): + r""" + Adaptive trainable :class:`~torch.nn.ReLU` activation function. + + Given the function :math:`\text{ReLU}:\mathbb{R}^n\rightarrow\mathbb{R}^n`, + the adaptive function + :math:`\text{ReLU}_{\text{adaptive}}:\mathbb{R}^n\rightarrow\mathbb{R}^n` + is defined as: + + .. math:: + \text{ReLU}_{\text{adaptive}}({x}) = \alpha\,\text{ReLU}(\beta{x}+\gamma), + + where :math:`\alpha,\,\beta,\,\gamma` are trainable parameters, and the + ReLU function is defined as: + + .. math:: + \text{ReLU}(x) = \max(0, x) + + .. seealso:: + + **Original reference**: Godfrey, Luke B., and Michael S. Gashler. + *A continuum among logarithmic, linear, and exponential functions, + and its potential to improve generalization in neural networks.* + 2015 7th international joint conference on knowledge discovery, + knowledge engineering and knowledge management (IC3K). + Vol. 1. IEEE, 2015. DOI: `arXiv preprint arXiv:1602.01321. + `_. + + Jagtap, Ameya D., Kenji Kawaguchi, and George Em Karniadakis. *Adaptive + activation functions accelerate convergence in deep and + physics-informed neural networks*. Journal of + Computational Physics 404 (2020): 109136. + DOI: `JCP 10.1016 + `_. + """ + def __init__(self, alpha=None, beta=None, gamma=None, fixed=None): + super().__init__(alpha, beta, gamma, fixed) + self._func = torch.nn.ReLU() + + +class AdaptiveSigmoid(AdaptiveActivationFunctionInterface): + r""" + Adaptive trainable :class:`~torch.nn.Sigmoid` activation function. + + Given the function :math:`\text{Sigmoid}:\mathbb{R}^n\rightarrow\mathbb{R}^n`, + the adaptive function + :math:`\text{Sigmoid}_{\text{adaptive}}:\mathbb{R}^n\rightarrow\mathbb{R}^n` + is defined as: + + .. math:: + \text{Sigmoid}_{\text{adaptive}}({x}) = \alpha\,\text{Sigmoid}(\beta{x}+\gamma), + + where :math:`\alpha,\,\beta,\,\gamma` are trainable parameters, and the + Sigmoid function is defined as: + + .. math:: + \text{Sigmoid}(x) = \frac{1}{1 + \exp(-x)} + + .. seealso:: + + **Original reference**: Godfrey, Luke B., and Michael S. Gashler. + *A continuum among logarithmic, linear, and exponential functions, + and its potential to improve generalization in neural networks.* + 2015 7th international joint conference on knowledge discovery, + knowledge engineering and knowledge management (IC3K). + Vol. 1. IEEE, 2015. DOI: `arXiv preprint arXiv:1602.01321. + `_. + + Jagtap, Ameya D., Kenji Kawaguchi, and George Em Karniadakis. *Adaptive + activation functions accelerate convergence in deep and + physics-informed neural networks*. Journal of + Computational Physics 404 (2020): 109136. + DOI: `JCP 10.1016 + `_. + """ + def __init__(self, alpha=None, beta=None, gamma=None, fixed=None): + super().__init__(alpha, beta, gamma, fixed) + self._func = torch.nn.Sigmoid() + + +class AdaptiveTanh(AdaptiveActivationFunctionInterface): + r""" + Adaptive trainable :class:`~torch.nn.Tanh` activation function. + + Given the function :math:`\text{Tanh}:\mathbb{R}^n\rightarrow\mathbb{R}^n`, + the adaptive function + :math:`\text{Tanh}_{\text{adaptive}}:\mathbb{R}^n\rightarrow\mathbb{R}^n` + is defined as: + + .. math:: + \text{Tanh}_{\text{adaptive}}({x}) = \alpha\,\text{Tanh}(\beta{x}+\gamma), + + where :math:`\alpha,\,\beta,\,\gamma` are trainable parameters, and the + Tanh function is defined as: + + .. math:: + \text{Tanh}(x) = \frac{\exp(x) - \exp(-x)} {\exp(x) + \exp(-x)} + + .. seealso:: + + **Original reference**: Godfrey, Luke B., and Michael S. Gashler. + *A continuum among logarithmic, linear, and exponential functions, + and its potential to improve generalization in neural networks.* + 2015 7th international joint conference on knowledge discovery, + knowledge engineering and knowledge management (IC3K). + Vol. 1. IEEE, 2015. DOI: `arXiv preprint arXiv:1602.01321. + `_. + + Jagtap, Ameya D., Kenji Kawaguchi, and George Em Karniadakis. *Adaptive + activation functions accelerate convergence in deep and + physics-informed neural networks*. Journal of + Computational Physics 404 (2020): 109136. + DOI: `JCP 10.1016 + `_. + """ + def __init__(self, alpha=None, beta=None, gamma=None, fixed=None): + super().__init__(alpha, beta, gamma, fixed) + self._func = torch.nn.Tanh() + + +class AdaptiveSiLU(AdaptiveActivationFunctionInterface): + r""" + Adaptive trainable :class:`~torch.nn.SiLU` activation function. + + Given the function :math:`\text{SiLU}:\mathbb{R}^n\rightarrow\mathbb{R}^n`, + the adaptive function + :math:`\text{SiLU}_{\text{adaptive}}:\mathbb{R}^n\rightarrow\mathbb{R}^n` + is defined as: + + .. math:: + \text{SiLU}_{\text{adaptive}}({x}) = \alpha\,\text{SiLU}(\beta{x}+\gamma), + + where :math:`\alpha,\,\beta,\,\gamma` are trainable parameters, and the + SiLU function is defined as: + + .. math:: + \text{SiLU}(x) = x * \sigma(x), \text{where }\sigma(x) + \text{ is the logistic sigmoid.} + + .. seealso:: + + **Original reference**: Godfrey, Luke B., and Michael S. Gashler. + *A continuum among logarithmic, linear, and exponential functions, + and its potential to improve generalization in neural networks.* + 2015 7th international joint conference on knowledge discovery, + knowledge engineering and knowledge management (IC3K). + Vol. 1. IEEE, 2015. DOI: `arXiv preprint arXiv:1602.01321. + `_. + + Jagtap, Ameya D., Kenji Kawaguchi, and George Em Karniadakis. *Adaptive + activation functions accelerate convergence in deep and + physics-informed neural networks*. Journal of + Computational Physics 404 (2020): 109136. + DOI: `JCP 10.1016 + `_. + """ + def __init__(self, alpha=None, beta=None, gamma=None, fixed=None): + super().__init__(alpha, beta, gamma, fixed) + self._func = torch.nn.SiLU() + + +class AdaptiveMish(AdaptiveActivationFunctionInterface): + r""" + Adaptive trainable :class:`~torch.nn.Mish` activation function. + + Given the function :math:`\text{Mish}:\mathbb{R}^n\rightarrow\mathbb{R}^n`, + the adaptive function + :math:`\text{Mish}_{\text{adaptive}}:\mathbb{R}^n\rightarrow\mathbb{R}^n` + is defined as: + + .. math:: + \text{Mish}_{\text{adaptive}}({x}) = \alpha\,\text{Mish}(\beta{x}+\gamma), + + where :math:`\alpha,\,\beta,\,\gamma` are trainable parameters, and the + Mish function is defined as: + + .. math:: + \text{Mish}(x) = x * \text{Tanh}(x) + + .. seealso:: + + **Original reference**: Godfrey, Luke B., and Michael S. Gashler. + *A continuum among logarithmic, linear, and exponential functions, + and its potential to improve generalization in neural networks.* + 2015 7th international joint conference on knowledge discovery, + knowledge engineering and knowledge management (IC3K). + Vol. 1. IEEE, 2015. DOI: `arXiv preprint arXiv:1602.01321. + `_. + + Jagtap, Ameya D., Kenji Kawaguchi, and George Em Karniadakis. *Adaptive + activation functions accelerate convergence in deep and + physics-informed neural networks*. Journal of + Computational Physics 404 (2020): 109136. + DOI: `JCP 10.1016 + `_. + """ + def __init__(self, alpha=None, beta=None, gamma=None, fixed=None): + super().__init__(alpha, beta, gamma, fixed) + self._func = torch.nn.Mish() + + +class AdaptiveELU(AdaptiveActivationFunctionInterface): + r""" + Adaptive trainable :class:`~torch.nn.ELU` activation function. + + Given the function :math:`\text{ELU}:\mathbb{R}^n\rightarrow\mathbb{R}^n`, + the adaptive function + :math:`\text{ELU}_{\text{adaptive}}:\mathbb{R}^n\rightarrow\mathbb{R}^n` + is defined as: + + .. math:: + \text{ELU}_{\text{adaptive}}({x}) = \alpha\,\text{ELU}(\beta{x}+\gamma), + + where :math:`\alpha,\,\beta,\,\gamma` are trainable parameters, and the + ELU function is defined as: + + .. math:: + \text{ELU}(x) = \begin{cases} + x, & \text{ if }x > 0\\ + \exp(x) - 1, & \text{ if }x \leq 0 + \end{cases} + + .. seealso:: + + **Original reference**: Godfrey, Luke B., and Michael S. Gashler. + *A continuum among logarithmic, linear, and exponential functions, + and its potential to improve generalization in neural networks.* + 2015 7th international joint conference on knowledge discovery, + knowledge engineering and knowledge management (IC3K). + Vol. 1. IEEE, 2015. DOI: `arXiv preprint arXiv:1602.01321. + `_. + + Jagtap, Ameya D., Kenji Kawaguchi, and George Em Karniadakis. *Adaptive + activation functions accelerate convergence in deep and + physics-informed neural networks*. Journal of + Computational Physics 404 (2020): 109136. + DOI: `JCP 10.1016 + `_. + """ + def __init__(self, alpha=None, beta=None, gamma=None, fixed=None): + super().__init__(alpha, beta, gamma, fixed) + self._func = torch.nn.ELU() + + +class AdaptiveCELU(AdaptiveActivationFunctionInterface): + r""" + Adaptive trainable :class:`~torch.nn.CELU` activation function. + + Given the function :math:`\text{CELU}:\mathbb{R}^n\rightarrow\mathbb{R}^n`, + the adaptive function + :math:`\text{CELU}_{\text{adaptive}}:\mathbb{R}^n\rightarrow\mathbb{R}^n` + is defined as: + + .. math:: + \text{CELU}_{\text{adaptive}}({x}) = \alpha\,\text{CELU}(\beta{x}+\gamma), + + where :math:`\alpha,\,\beta,\,\gamma` are trainable parameters, and the + CELU function is defined as: + + .. math:: + \text{CELU}(x) = \max(0,x) + \min(0, \alpha * (\exp(x) - 1)) + + .. seealso:: + + **Original reference**: Godfrey, Luke B., and Michael S. Gashler. + *A continuum among logarithmic, linear, and exponential functions, + and its potential to improve generalization in neural networks.* + 2015 7th international joint conference on knowledge discovery, + knowledge engineering and knowledge management (IC3K). + Vol. 1. IEEE, 2015. DOI: `arXiv preprint arXiv:1602.01321. + `_. + + Jagtap, Ameya D., Kenji Kawaguchi, and George Em Karniadakis. *Adaptive + activation functions accelerate convergence in deep and + physics-informed neural networks*. Journal of + Computational Physics 404 (2020): 109136. + DOI: `JCP 10.1016 + `_. + """ + def __init__(self, alpha=None, beta=None, gamma=None, fixed=None): + super().__init__(alpha, beta, gamma, fixed) + self._func = torch.nn.CELU() + +class AdaptiveGELU(AdaptiveActivationFunctionInterface): + r""" + Adaptive trainable :class:`~torch.nn.GELU` activation function. + + Given the function :math:`\text{GELU}:\mathbb{R}^n\rightarrow\mathbb{R}^n`, + the adaptive function + :math:`\text{GELU}_{\text{adaptive}}:\mathbb{R}^n\rightarrow\mathbb{R}^n` + is defined as: + + .. math:: + \text{GELU}_{\text{adaptive}}({x}) = \alpha\,\text{GELU}(\beta{x}+\gamma), + + where :math:`\alpha,\,\beta,\,\gamma` are trainable parameters, and the + GELU function is defined as: + + .. math:: + \text{GELU}(x) = 0.5 * x * (1 + \text{Tanh}(\sqrt{2 / \pi} * (x + 0.044715 * x^3))) + + + .. seealso:: + + **Original reference**: Godfrey, Luke B., and Michael S. Gashler. + *A continuum among logarithmic, linear, and exponential functions, + and its potential to improve generalization in neural networks.* + 2015 7th international joint conference on knowledge discovery, + knowledge engineering and knowledge management (IC3K). + Vol. 1. IEEE, 2015. DOI: `arXiv preprint arXiv:1602.01321. + `_. + + Jagtap, Ameya D., Kenji Kawaguchi, and George Em Karniadakis. *Adaptive + activation functions accelerate convergence in deep and + physics-informed neural networks*. Journal of + Computational Physics 404 (2020): 109136. + DOI: `JCP 10.1016 + `_. + """ + def __init__(self, alpha=None, beta=None, gamma=None, fixed=None): + super().__init__(alpha, beta, gamma, fixed) + self._func = torch.nn.GELU() + + +class AdaptiveSoftmin(AdaptiveActivationFunctionInterface): + r""" + Adaptive trainable :class:`~torch.nn.Softmin` activation function. + + Given the function :math:`\text{Softmin}:\mathbb{R}^n\rightarrow\mathbb{R}^n`, + the adaptive function + :math:`\text{Softmin}_{\text{adaptive}}:\mathbb{R}^n\rightarrow\mathbb{R}^n` + is defined as: + + .. math:: + \text{Softmin}_{\text{adaptive}}({x}) = \alpha\,\text{Softmin}(\beta{x}+\gamma), + + where :math:`\alpha,\,\beta,\,\gamma` are trainable parameters, and the + Softmin function is defined as: + + .. math:: + \text{Softmin}(x_{i}) = \frac{\exp(-x_i)}{\sum_j \exp(-x_j)} + + .. seealso:: + + **Original reference**: Godfrey, Luke B., and Michael S. Gashler. + *A continuum among logarithmic, linear, and exponential functions, + and its potential to improve generalization in neural networks.* + 2015 7th international joint conference on knowledge discovery, + knowledge engineering and knowledge management (IC3K). + Vol. 1. IEEE, 2015. DOI: `arXiv preprint arXiv:1602.01321. + `_. + + Jagtap, Ameya D., Kenji Kawaguchi, and George Em Karniadakis. *Adaptive + activation functions accelerate convergence in deep and + physics-informed neural networks*. Journal of + Computational Physics 404 (2020): 109136. + DOI: `JCP 10.1016 + `_. + """ + def __init__(self, alpha=None, beta=None, gamma=None, fixed=None): + super().__init__(alpha, beta, gamma, fixed) + self._func = torch.nn.Softmin() + + +class AdaptiveSoftmax(AdaptiveActivationFunctionInterface): + r""" + Adaptive trainable :class:`~torch.nn.Softmax` activation function. + + Given the function :math:`\text{Softmax}:\mathbb{R}^n\rightarrow\mathbb{R}^n`, + the adaptive function + :math:`\text{Softmax}_{\text{adaptive}}:\mathbb{R}^n\rightarrow\mathbb{R}^n` + is defined as: + + .. math:: + \text{Softmax}_{\text{adaptive}}({x}) = \alpha\,\text{Softmax}(\beta{x}+\gamma), + + where :math:`\alpha,\,\beta,\,\gamma` are trainable parameters, and the + Softmax function is defined as: + + .. math:: + \text{Softmax}(x_{i}) = \frac{\exp(x_i)}{\sum_j \exp(x_j)} + + .. seealso:: + + **Original reference**: Godfrey, Luke B., and Michael S. Gashler. + *A continuum among logarithmic, linear, and exponential functions, + and its potential to improve generalization in neural networks.* + 2015 7th international joint conference on knowledge discovery, + knowledge engineering and knowledge management (IC3K). + Vol. 1. IEEE, 2015. DOI: `arXiv preprint arXiv:1602.01321. + `_. + + Jagtap, Ameya D., Kenji Kawaguchi, and George Em Karniadakis. *Adaptive + activation functions accelerate convergence in deep and + physics-informed neural networks*. Journal of + Computational Physics 404 (2020): 109136. + DOI: `JCP 10.1016 + `_. + """ + def __init__(self, alpha=None, beta=None, gamma=None, fixed=None): + super().__init__(alpha, beta, gamma, fixed) + self._func = torch.nn.Softmax() + +class AdaptiveSIREN(AdaptiveActivationFunctionInterface): + r""" + Adaptive trainable :obj:`~torch.sin` function. + + Given the function :math:`\text{sin}:\mathbb{R}^n\rightarrow\mathbb{R}^n`, + the adaptive function + :math:`\text{sin}_{\text{adaptive}}:\mathbb{R}^n\rightarrow\mathbb{R}^n` + is defined as: + + .. math:: + \text{sin}_{\text{adaptive}}({x}) = \alpha\,\text{sin}(\beta{x}+\gamma), + + where :math:`\alpha,\,\beta,\,\gamma` are trainable parameters. + + .. seealso:: + + **Original reference**: Godfrey, Luke B., and Michael S. Gashler. + *A continuum among logarithmic, linear, and exponential functions, + and its potential to improve generalization in neural networks.* + 2015 7th international joint conference on knowledge discovery, + knowledge engineering and knowledge management (IC3K). + Vol. 1. IEEE, 2015. DOI: `arXiv preprint arXiv:1602.01321. + `_. + + Jagtap, Ameya D., Kenji Kawaguchi, and George Em Karniadakis. *Adaptive + activation functions accelerate convergence in deep and + physics-informed neural networks*. Journal of + Computational Physics 404 (2020): 109136. + DOI: `JCP 10.1016 + `_. + """ + def __init__(self, alpha=None, beta=None, gamma=None, fixed=None): + super().__init__(alpha, beta, gamma, fixed) + self._func = torch.sin + +class AdaptiveExp(AdaptiveActivationFunctionInterface): + r""" + Adaptive trainable :obj:`~torch.exp` function. + + Given the function :math:`\text{exp}:\mathbb{R}^n\rightarrow\mathbb{R}^n`, + the adaptive function + :math:`\text{exp}_{\text{adaptive}}:\mathbb{R}^n\rightarrow\mathbb{R}^n` + is defined as: + + .. math:: + \text{exp}_{\text{adaptive}}({x}) = \alpha\,\text{exp}(\beta{x}), + + where :math:`\alpha,\,\beta` are trainable parameters. + + .. seealso:: + + **Original reference**: Godfrey, Luke B., and Michael S. Gashler. + *A continuum among logarithmic, linear, and exponential functions, + and its potential to improve generalization in neural networks.* + 2015 7th international joint conference on knowledge discovery, + knowledge engineering and knowledge management (IC3K). + Vol. 1. IEEE, 2015. DOI: `arXiv preprint arXiv:1602.01321. + `_. + + Jagtap, Ameya D., Kenji Kawaguchi, and George Em Karniadakis. *Adaptive + activation functions accelerate convergence in deep and + physics-informed neural networks*. Journal of + Computational Physics 404 (2020): 109136. + DOI: `JCP 10.1016 + `_. + """ + def __init__(self, alpha=None, beta=None, fixed=None): + + # only alpha, and beta parameters (gamma=0 fixed) + if fixed is None: + fixed = ['gamma'] + else: + check_consistency(fixed, str) + fixed = list(fixed) + ['gamma'] + + # calling super + super().__init__(alpha, beta, 0., fixed) + self._func = torch.exp \ No newline at end of file diff --git a/pina/model/layers/adaptive_func.py b/pina/adaptive_functions/adaptive_func_interface.py similarity index 68% rename from pina/model/layers/adaptive_func.py rename to pina/adaptive_functions/adaptive_func_interface.py index 3bfc4be6c..b0522d52d 100644 --- a/pina/model/layers/adaptive_func.py +++ b/pina/adaptive_functions/adaptive_func_interface.py @@ -1,14 +1,18 @@ """ Module for adaptive functions. """ import torch + from pina.utils import check_consistency +from abc import ABCMeta -class AdaptiveActivationFunction(torch.nn.Module): +class AdaptiveActivationFunctionInterface(torch.nn.Module, metaclass=ABCMeta): r""" - The :class:`~pina.model.layers.adaptive_func.AdaptiveActivationFunction` + The + :class:`~pina.adaptive_functions.adaptive_func_interface.AdaptiveActivationFunctionInterface` class makes a :class:`torch.nn.Module` activation function into an adaptive - trainable activation function. + trainable activation function. If one wants to create an adpative activation + function, this class must be use as base class. Given a function :math:`f:\mathbb{R}^n\rightarrow\mathbb{R}^m`, the adaptive function :math:`f_{\text{adaptive}}:\mathbb{R}^n\rightarrow\mathbb{R}^m` @@ -19,28 +23,6 @@ class makes a :class:`torch.nn.Module` activation function into an adaptive where :math:`\alpha,\,\beta,\,\gamma` are trainable parameters. - :Example: - >>> import torch - >>> from pina.model.layers import AdaptiveActivationFunction - >>> - >>> # simple adaptive function with all trainable parameters - >>> AdaptiveTanh = AdaptiveActivationFunction(torch.nn.Tanh()) - >>> AdaptiveTanh(torch.rand(3)) - tensor([0.1084, 0.3931, 0.7294], grad_fn=) - >>> AdaptiveTanh.alpha - Parameter containing: - tensor(1., requires_grad=True) - >>> - >>> # simple adaptive function with trainable parameters fixed alpha - >>> AdaptiveTanh = AdaptiveActivationFunction(torch.nn.Tanh(), - ... fixed=['alpha']) - >>> AdaptiveTanh.alpha - tensor(1.) - >>> AdaptiveTanh.beta - Parameter containing: - tensor(1., requires_grad=True) - >>> - .. seealso:: **Original reference**: Godfrey, Luke B., and Michael S. Gashler. @@ -51,14 +33,18 @@ class makes a :class:`torch.nn.Module` activation function into an adaptive Vol. 1. IEEE, 2015. DOI: `arXiv preprint arXiv:1602.01321. `_. + Jagtap, Ameya D., Kenji Kawaguchi, and George Em Karniadakis. *Adaptive + activation functions accelerate convergence in deep and + physics-informed neural networks*. Journal of + Computational Physics 404 (2020): 109136. + DOI: `JCP 10.1016 + `_. """ - def __init__(self, func, alpha=None, beta=None, gamma=None, fixed=None): + def __init__(self, alpha=None, beta=None, gamma=None, fixed=None): """ - Initializes the AdaptiveActivationFunction module. + Initializes the Adaptive Function. - :param callable func: The original collable function. It could be an - initialized :meth:`torch.nn.Module`, or a python callable function. :param float | complex alpha: Scaling parameter alpha. Defaults to ``None``. When ``None`` is passed, the variable is initialized to 1. @@ -70,7 +56,7 @@ def __init__(self, func, alpha=None, beta=None, gamma=None, fixed=None): the variable is initialized to 1. :param list fixed: List of parameters to fix during training, i.e. not optimized (``requires_grad`` set to ``False``). - Options are ['alpha', 'beta', 'gamma']. Defaults to None. + Options are ``alpha``, ``beta``, ``gamma``. Defaults to None. """ super().__init__() @@ -94,8 +80,6 @@ def __init__(self, func, alpha=None, beta=None, gamma=None, fixed=None): check_consistency(alpha, (float, complex)) check_consistency(beta, (float, complex)) check_consistency(gamma, (float, complex)) - if not callable(func): - raise ValueError("Function must be a callable function.") # registering as tensors alpha = torch.tensor(alpha, requires_grad=False) @@ -119,34 +103,44 @@ def __init__(self, func, alpha=None, beta=None, gamma=None, fixed=None): self._gamma = torch.nn.Parameter(gamma, requires_grad=True) else: self.register_buffer("gamma", gamma) - - # registering function - self._func = func + + # storing the activation + self._func = None def forward(self, x): """ - Forward pass of the function. - Applies the function to the input elementwise. + Define the computation performed at every call. + The function to the input elementwise. + + :param x: The input tensor to evaluate the activation function. + :type x: torch.Tensor | LabelTensor """ return self.alpha * (self._func(self.beta * x + self.gamma)) @property def alpha(self): """ - The alpha variable + The alpha variable. """ return self._alpha @property def beta(self): """ - The alpha variable + The beta variable. """ return self._beta @property def gamma(self): """ - The alpha variable + The gamma variable. """ return self._gamma + + @property + def func(self): + """ + The callable activation function. + """ + return self._func \ No newline at end of file diff --git a/pina/model/layers/__init__.py b/pina/model/layers/__init__.py index 5f5a14f0c..de9686b58 100644 --- a/pina/model/layers/__init__.py +++ b/pina/model/layers/__init__.py @@ -12,7 +12,6 @@ "PeriodicBoundaryEmbedding", "AVNOBlock", "LowRankBlock", - "AdaptiveActivationFunction", ] from .convolution_2d import ContinuousConvBlock @@ -26,5 +25,4 @@ from .pod import PODBlock from .embedding import PeriodicBoundaryEmbedding from .avno_layer import AVNOBlock -from .lowrank_layer import LowRankBlock -from .adaptive_func import AdaptiveActivationFunction +from .lowrank_layer import LowRankBlock \ No newline at end of file diff --git a/tests/test_adaptive_functions.py b/tests/test_adaptive_functions.py new file mode 100644 index 000000000..43d9c1bc7 --- /dev/null +++ b/tests/test_adaptive_functions.py @@ -0,0 +1,62 @@ +import torch +import pytest + +from pina.adaptive_functions import (AdaptiveReLU, AdaptiveSigmoid, AdaptiveTanh, + AdaptiveSiLU, AdaptiveMish, AdaptiveELU, + AdaptiveCELU, AdaptiveGELU, AdaptiveSoftmin, + AdaptiveSoftmax, AdaptiveSIREN, AdaptiveExp) + + +adaptive_functions = (AdaptiveReLU, AdaptiveSigmoid, AdaptiveTanh, + AdaptiveSiLU, AdaptiveMish, AdaptiveELU, + AdaptiveCELU, AdaptiveGELU, AdaptiveSoftmin, + AdaptiveSoftmax, AdaptiveSIREN, AdaptiveExp) +x = torch.rand(10, requires_grad=True) + +@pytest.mark.parametrize("Func", adaptive_functions) +def test_constructor(Func): + if Func.__name__ == 'AdaptiveExp': + # simple + Func() + # setting values + af = Func(alpha=1., beta=2.) + assert af.alpha.requires_grad + assert af.beta.requires_grad + assert af.alpha == 1. + assert af.beta == 2. + else: + # simple + Func() + # setting values + af = Func(alpha=1., beta=2., gamma=3.) + assert af.alpha.requires_grad + assert af.beta.requires_grad + assert af.gamma.requires_grad + assert af.alpha == 1. + assert af.beta == 2. + assert af.gamma == 3. + + # fixed variables + af = Func(alpha=1., beta=2., fixed=['alpha']) + assert af.alpha.requires_grad is False + assert af.beta.requires_grad + assert af.alpha == 1. + assert af.beta == 2. + + with pytest.raises(TypeError): + Func(alpha=1., beta=2., fixed=['delta']) + + with pytest.raises(ValueError): + Func(alpha='s') + Func(alpha=1) + +@pytest.mark.parametrize("Func", adaptive_functions) +def test_forward(Func): + af = Func() + af(x) + +@pytest.mark.parametrize("Func", adaptive_functions) +def test_backward(Func): + af = Func() + y = af(x) + y.mean().backward() \ No newline at end of file diff --git a/tests/test_layers/test_adaptive_func.py b/tests/test_layers/test_adaptive_func.py deleted file mode 100644 index 3e3e6662c..000000000 --- a/tests/test_layers/test_adaptive_func.py +++ /dev/null @@ -1,48 +0,0 @@ -import torch -import pytest - -from pina.model.layers.adaptive_func import AdaptiveActivationFunction - -x = torch.rand(5) -torchfunc = torch.nn.Tanh() - -def test_constructor(): - # simple - AdaptiveActivationFunction(torchfunc) - - # setting values - af = AdaptiveActivationFunction(torchfunc, alpha=1., beta=2., gamma=3.) - assert af.alpha.requires_grad - assert af.beta.requires_grad - assert af.gamma.requires_grad - assert af.alpha == 1. - assert af.beta == 2. - assert af.gamma == 3. - - # fixed variables - af = AdaptiveActivationFunction(torchfunc, alpha=1., beta=2., - gamma=3., fixed=['alpha']) - assert af.alpha.requires_grad is False - assert af.beta.requires_grad - assert af.gamma.requires_grad - assert af.alpha == 1. - assert af.beta == 2. - assert af.gamma == 3. - - with pytest.raises(TypeError): - AdaptiveActivationFunction(torchfunc, alpha=1., beta=2., - gamma=3., fixed=['delta']) - - with pytest.raises(ValueError): - AdaptiveActivationFunction(torchfunc, alpha='s') - AdaptiveActivationFunction(torchfunc, alpha=1., fixed='alpha') - AdaptiveActivationFunction(torchfunc, alpha=1) - -def test_forward(): - af = AdaptiveActivationFunction(torchfunc) - af(x) - -def test_backward(): - af = AdaptiveActivationFunction(torchfunc) - y = af(x) - y.mean().backward() \ No newline at end of file From f101313a34a8b5181bfd2804156453ed29c79174 Mon Sep 17 00:00:00 2001 From: ndem0 Date: Mon, 8 Apr 2024 15:18:17 +0000 Subject: [PATCH 02/14] :art: Format Python code with psf/black --- pina/adaptive_functions/__init__.py | 46 ++++++++------ pina/adaptive_functions/adaptive_func.py | 63 ++++++++++++------- .../adaptive_func_interface.py | 8 +-- pina/model/layers/__init__.py | 2 +- 4 files changed, 72 insertions(+), 47 deletions(-) diff --git a/pina/adaptive_functions/__init__.py b/pina/adaptive_functions/__init__.py index 0ab6053ed..0fa0ecd9e 100644 --- a/pina/adaptive_functions/__init__.py +++ b/pina/adaptive_functions/__init__.py @@ -1,21 +1,31 @@ __all__ = [ - 'AdaptiveActivationFunctionInterface', - 'AdaptiveReLU', - 'AdaptiveSigmoid', - 'AdaptiveTanh', - 'AdaptiveSiLU', - 'AdaptiveMish', - 'AdaptiveELU', - 'AdaptiveCELU', - 'AdaptiveGELU', - 'AdaptiveSoftmin', - 'AdaptiveSoftmax', - 'AdaptiveSIREN', - 'AdaptiveExp'] + "AdaptiveActivationFunctionInterface", + "AdaptiveReLU", + "AdaptiveSigmoid", + "AdaptiveTanh", + "AdaptiveSiLU", + "AdaptiveMish", + "AdaptiveELU", + "AdaptiveCELU", + "AdaptiveGELU", + "AdaptiveSoftmin", + "AdaptiveSoftmax", + "AdaptiveSIREN", + "AdaptiveExp", +] -from .adaptive_func import (AdaptiveReLU, AdaptiveSigmoid, AdaptiveTanh, - AdaptiveSiLU, AdaptiveMish, AdaptiveELU, - AdaptiveCELU, AdaptiveGELU, AdaptiveSoftmin, - AdaptiveSoftmax, AdaptiveSIREN, AdaptiveExp) +from .adaptive_func import ( + AdaptiveReLU, + AdaptiveSigmoid, + AdaptiveTanh, + AdaptiveSiLU, + AdaptiveMish, + AdaptiveELU, + AdaptiveCELU, + AdaptiveGELU, + AdaptiveSoftmin, + AdaptiveSoftmax, + AdaptiveSIREN, + AdaptiveExp, +) from .adaptive_func_interface import AdaptiveActivationFunctionInterface - diff --git a/pina/adaptive_functions/adaptive_func.py b/pina/adaptive_functions/adaptive_func.py index 0ee22b2d0..30966f1fc 100644 --- a/pina/adaptive_functions/adaptive_func.py +++ b/pina/adaptive_functions/adaptive_func.py @@ -19,7 +19,7 @@ class AdaptiveReLU(AdaptiveActivationFunctionInterface): where :math:`\alpha,\,\beta,\,\gamma` are trainable parameters, and the ReLU function is defined as: - + .. math:: \text{ReLU}(x) = \max(0, x) @@ -36,10 +36,11 @@ class AdaptiveReLU(AdaptiveActivationFunctionInterface): Jagtap, Ameya D., Kenji Kawaguchi, and George Em Karniadakis. *Adaptive activation functions accelerate convergence in deep and physics-informed neural networks*. Journal of - Computational Physics 404 (2020): 109136. + Computational Physics 404 (2020): 109136. DOI: `JCP 10.1016 `_. """ + def __init__(self, alpha=None, beta=None, gamma=None, fixed=None): super().__init__(alpha, beta, gamma, fixed) self._func = torch.nn.ReLU() @@ -59,7 +60,7 @@ class AdaptiveSigmoid(AdaptiveActivationFunctionInterface): where :math:`\alpha,\,\beta,\,\gamma` are trainable parameters, and the Sigmoid function is defined as: - + .. math:: \text{Sigmoid}(x) = \frac{1}{1 + \exp(-x)} @@ -76,10 +77,11 @@ class AdaptiveSigmoid(AdaptiveActivationFunctionInterface): Jagtap, Ameya D., Kenji Kawaguchi, and George Em Karniadakis. *Adaptive activation functions accelerate convergence in deep and physics-informed neural networks*. Journal of - Computational Physics 404 (2020): 109136. + Computational Physics 404 (2020): 109136. DOI: `JCP 10.1016 `_. """ + def __init__(self, alpha=None, beta=None, gamma=None, fixed=None): super().__init__(alpha, beta, gamma, fixed) self._func = torch.nn.Sigmoid() @@ -99,7 +101,7 @@ class AdaptiveTanh(AdaptiveActivationFunctionInterface): where :math:`\alpha,\,\beta,\,\gamma` are trainable parameters, and the Tanh function is defined as: - + .. math:: \text{Tanh}(x) = \frac{\exp(x) - \exp(-x)} {\exp(x) + \exp(-x)} @@ -116,10 +118,11 @@ class AdaptiveTanh(AdaptiveActivationFunctionInterface): Jagtap, Ameya D., Kenji Kawaguchi, and George Em Karniadakis. *Adaptive activation functions accelerate convergence in deep and physics-informed neural networks*. Journal of - Computational Physics 404 (2020): 109136. + Computational Physics 404 (2020): 109136. DOI: `JCP 10.1016 `_. """ + def __init__(self, alpha=None, beta=None, gamma=None, fixed=None): super().__init__(alpha, beta, gamma, fixed) self._func = torch.nn.Tanh() @@ -139,7 +142,7 @@ class AdaptiveSiLU(AdaptiveActivationFunctionInterface): where :math:`\alpha,\,\beta,\,\gamma` are trainable parameters, and the SiLU function is defined as: - + .. math:: \text{SiLU}(x) = x * \sigma(x), \text{where }\sigma(x) \text{ is the logistic sigmoid.} @@ -157,10 +160,11 @@ class AdaptiveSiLU(AdaptiveActivationFunctionInterface): Jagtap, Ameya D., Kenji Kawaguchi, and George Em Karniadakis. *Adaptive activation functions accelerate convergence in deep and physics-informed neural networks*. Journal of - Computational Physics 404 (2020): 109136. + Computational Physics 404 (2020): 109136. DOI: `JCP 10.1016 `_. """ + def __init__(self, alpha=None, beta=None, gamma=None, fixed=None): super().__init__(alpha, beta, gamma, fixed) self._func = torch.nn.SiLU() @@ -180,7 +184,7 @@ class AdaptiveMish(AdaptiveActivationFunctionInterface): where :math:`\alpha,\,\beta,\,\gamma` are trainable parameters, and the Mish function is defined as: - + .. math:: \text{Mish}(x) = x * \text{Tanh}(x) @@ -197,10 +201,11 @@ class AdaptiveMish(AdaptiveActivationFunctionInterface): Jagtap, Ameya D., Kenji Kawaguchi, and George Em Karniadakis. *Adaptive activation functions accelerate convergence in deep and physics-informed neural networks*. Journal of - Computational Physics 404 (2020): 109136. + Computational Physics 404 (2020): 109136. DOI: `JCP 10.1016 `_. """ + def __init__(self, alpha=None, beta=None, gamma=None, fixed=None): super().__init__(alpha, beta, gamma, fixed) self._func = torch.nn.Mish() @@ -244,6 +249,7 @@ class AdaptiveELU(AdaptiveActivationFunctionInterface): DOI: `JCP 10.1016 `_. """ + def __init__(self, alpha=None, beta=None, gamma=None, fixed=None): super().__init__(alpha, beta, gamma, fixed) self._func = torch.nn.ELU() @@ -263,7 +269,7 @@ class AdaptiveCELU(AdaptiveActivationFunctionInterface): where :math:`\alpha,\,\beta,\,\gamma` are trainable parameters, and the CELU function is defined as: - + .. math:: \text{CELU}(x) = \max(0,x) + \min(0, \alpha * (\exp(x) - 1)) @@ -280,14 +286,16 @@ class AdaptiveCELU(AdaptiveActivationFunctionInterface): Jagtap, Ameya D., Kenji Kawaguchi, and George Em Karniadakis. *Adaptive activation functions accelerate convergence in deep and physics-informed neural networks*. Journal of - Computational Physics 404 (2020): 109136. + Computational Physics 404 (2020): 109136. DOI: `JCP 10.1016 `_. """ + def __init__(self, alpha=None, beta=None, gamma=None, fixed=None): super().__init__(alpha, beta, gamma, fixed) self._func = torch.nn.CELU() + class AdaptiveGELU(AdaptiveActivationFunctionInterface): r""" Adaptive trainable :class:`~torch.nn.GELU` activation function. @@ -302,7 +310,7 @@ class AdaptiveGELU(AdaptiveActivationFunctionInterface): where :math:`\alpha,\,\beta,\,\gamma` are trainable parameters, and the GELU function is defined as: - + .. math:: \text{GELU}(x) = 0.5 * x * (1 + \text{Tanh}(\sqrt{2 / \pi} * (x + 0.044715 * x^3))) @@ -320,10 +328,11 @@ class AdaptiveGELU(AdaptiveActivationFunctionInterface): Jagtap, Ameya D., Kenji Kawaguchi, and George Em Karniadakis. *Adaptive activation functions accelerate convergence in deep and physics-informed neural networks*. Journal of - Computational Physics 404 (2020): 109136. + Computational Physics 404 (2020): 109136. DOI: `JCP 10.1016 `_. """ + def __init__(self, alpha=None, beta=None, gamma=None, fixed=None): super().__init__(alpha, beta, gamma, fixed) self._func = torch.nn.GELU() @@ -343,7 +352,7 @@ class AdaptiveSoftmin(AdaptiveActivationFunctionInterface): where :math:`\alpha,\,\beta,\,\gamma` are trainable parameters, and the Softmin function is defined as: - + .. math:: \text{Softmin}(x_{i}) = \frac{\exp(-x_i)}{\sum_j \exp(-x_j)} @@ -360,10 +369,11 @@ class AdaptiveSoftmin(AdaptiveActivationFunctionInterface): Jagtap, Ameya D., Kenji Kawaguchi, and George Em Karniadakis. *Adaptive activation functions accelerate convergence in deep and physics-informed neural networks*. Journal of - Computational Physics 404 (2020): 109136. + Computational Physics 404 (2020): 109136. DOI: `JCP 10.1016 `_. """ + def __init__(self, alpha=None, beta=None, gamma=None, fixed=None): super().__init__(alpha, beta, gamma, fixed) self._func = torch.nn.Softmin() @@ -383,7 +393,7 @@ class AdaptiveSoftmax(AdaptiveActivationFunctionInterface): where :math:`\alpha,\,\beta,\,\gamma` are trainable parameters, and the Softmax function is defined as: - + .. math:: \text{Softmax}(x_{i}) = \frac{\exp(x_i)}{\sum_j \exp(x_j)} @@ -400,14 +410,16 @@ class AdaptiveSoftmax(AdaptiveActivationFunctionInterface): Jagtap, Ameya D., Kenji Kawaguchi, and George Em Karniadakis. *Adaptive activation functions accelerate convergence in deep and physics-informed neural networks*. Journal of - Computational Physics 404 (2020): 109136. + Computational Physics 404 (2020): 109136. DOI: `JCP 10.1016 `_. """ + def __init__(self, alpha=None, beta=None, gamma=None, fixed=None): super().__init__(alpha, beta, gamma, fixed) self._func = torch.nn.Softmax() + class AdaptiveSIREN(AdaptiveActivationFunctionInterface): r""" Adaptive trainable :obj:`~torch.sin` function. @@ -435,14 +447,16 @@ class AdaptiveSIREN(AdaptiveActivationFunctionInterface): Jagtap, Ameya D., Kenji Kawaguchi, and George Em Karniadakis. *Adaptive activation functions accelerate convergence in deep and physics-informed neural networks*. Journal of - Computational Physics 404 (2020): 109136. + Computational Physics 404 (2020): 109136. DOI: `JCP 10.1016 `_. """ + def __init__(self, alpha=None, beta=None, gamma=None, fixed=None): super().__init__(alpha, beta, gamma, fixed) self._func = torch.sin + class AdaptiveExp(AdaptiveActivationFunctionInterface): r""" Adaptive trainable :obj:`~torch.exp` function. @@ -470,19 +484,20 @@ class AdaptiveExp(AdaptiveActivationFunctionInterface): Jagtap, Ameya D., Kenji Kawaguchi, and George Em Karniadakis. *Adaptive activation functions accelerate convergence in deep and physics-informed neural networks*. Journal of - Computational Physics 404 (2020): 109136. + Computational Physics 404 (2020): 109136. DOI: `JCP 10.1016 `_. """ + def __init__(self, alpha=None, beta=None, fixed=None): # only alpha, and beta parameters (gamma=0 fixed) if fixed is None: - fixed = ['gamma'] + fixed = ["gamma"] else: check_consistency(fixed, str) - fixed = list(fixed) + ['gamma'] + fixed = list(fixed) + ["gamma"] # calling super - super().__init__(alpha, beta, 0., fixed) - self._func = torch.exp \ No newline at end of file + super().__init__(alpha, beta, 0.0, fixed) + self._func = torch.exp diff --git a/pina/adaptive_functions/adaptive_func_interface.py b/pina/adaptive_functions/adaptive_func_interface.py index b0522d52d..a12b78b67 100644 --- a/pina/adaptive_functions/adaptive_func_interface.py +++ b/pina/adaptive_functions/adaptive_func_interface.py @@ -36,7 +36,7 @@ class makes a :class:`torch.nn.Module` activation function into an adaptive Jagtap, Ameya D., Kenji Kawaguchi, and George Em Karniadakis. *Adaptive activation functions accelerate convergence in deep and physics-informed neural networks*. Journal of - Computational Physics 404 (2020): 109136. + Computational Physics 404 (2020): 109136. DOI: `JCP 10.1016 `_. """ @@ -103,7 +103,7 @@ def __init__(self, alpha=None, beta=None, gamma=None, fixed=None): self._gamma = torch.nn.Parameter(gamma, requires_grad=True) else: self.register_buffer("gamma", gamma) - + # storing the activation self._func = None @@ -137,10 +137,10 @@ def gamma(self): The gamma variable. """ return self._gamma - + @property def func(self): """ The callable activation function. """ - return self._func \ No newline at end of file + return self._func diff --git a/pina/model/layers/__init__.py b/pina/model/layers/__init__.py index de9686b58..5d2034040 100644 --- a/pina/model/layers/__init__.py +++ b/pina/model/layers/__init__.py @@ -25,4 +25,4 @@ from .pod import PODBlock from .embedding import PeriodicBoundaryEmbedding from .avno_layer import AVNOBlock -from .lowrank_layer import LowRankBlock \ No newline at end of file +from .lowrank_layer import LowRankBlock From 9de291528b5a01845e608d4572c3cbbe3bac2860 Mon Sep 17 00:00:00 2001 From: Giuseppe Alessio D'Inverno <66356297+AleDinve@users.noreply.github.com> Date: Thu, 18 Apr 2024 17:23:27 +0200 Subject: [PATCH 03/14] Equation Class Tutorial (#287) * Tutorial Equation 12 * .rst and readme fix for tutorial12 * small fix * modified rst --------- Co-authored-by: Dario Coscia --- docs/source/_rst/_tutorial.rst | 2 + .../_rst/tutorials/tutorial12/tutorial.rst | 162 ++++++++++++ tutorials/README.md | 1 + tutorials/tutorial12/tutorial.ipynb | 235 ++++++++++++++++++ tutorials/tutorial12/tutorial.py | 142 +++++++++++ 5 files changed, 542 insertions(+) create mode 100644 docs/source/_rst/tutorials/tutorial12/tutorial.rst create mode 100644 tutorials/tutorial12/tutorial.ipynb create mode 100644 tutorials/tutorial12/tutorial.py diff --git a/docs/source/_rst/_tutorial.rst b/docs/source/_rst/_tutorial.rst index 020dbe2d0..9a72dd6aa 100644 --- a/docs/source/_rst/_tutorial.rst +++ b/docs/source/_rst/_tutorial.rst @@ -10,9 +10,11 @@ Getting started with PINA :titlesonly: Introduction to PINA for Physics Informed Neural Networks training + Introduction to PINA Equation class PINA and PyTorch Lightning, training tips and visualizations Building custom geometries with PINA Location class + Physics Informed Neural Networks -------------------------------- .. toctree:: diff --git a/docs/source/_rst/tutorials/tutorial12/tutorial.rst b/docs/source/_rst/tutorials/tutorial12/tutorial.rst new file mode 100644 index 000000000..8dd4dcf4f --- /dev/null +++ b/docs/source/_rst/tutorials/tutorial12/tutorial.rst @@ -0,0 +1,162 @@ +Tutorial: The ``Equation`` Class +================================ + +In this tutorial, we will show how to use the ``Equation`` Class in +PINA. Specifically, we will see how use the Class and its inherited +classes to enforce residuals minimization in PINNs. + +Example: The Burgers 1D equation +-------------------------------- + +We will start implementing the viscous Burgers 1D problem Class, +described as follows: + +.. math:: + + + \begin{equation} + \begin{cases} + \frac{\partial u}{\partial t} + u \frac{\partial u}{\partial x} &= \nu \frac{\partial^2 u}{ \partial x^2}, \quad x\in(0,1), \quad t>0\\ + u(x,0) &= -\sin (\pi x)\\ + u(x,t) &= 0 \quad x = \pm 1\\ + \end{cases} + \end{equation} + +where we set :math:`\nu = \frac{0.01}{\pi}` . + +In the class that models this problem we will see in action the +``Equation`` class and one of its inherited classes, the ``FixedValue`` +class. + +.. code:: ipython3 + + #useful imports + from pina.problem import SpatialProblem, TimeDependentProblem + from pina.equation import Equation, FixedValue, FixedGradient, FixedFlux + from pina.geometry import CartesianDomain + import torch + from pina.operators import grad, laplacian + from pina import Condition + + + +.. code:: ipython3 + + class Burgers1D(TimeDependentProblem, SpatialProblem): + + # define the burger equation + def burger_equation(input_, output_): + du = grad(output_, input_) + ddu = grad(du, input_, components=['dudx']) + return ( + du.extract(['dudt']) + + output_.extract(['u'])*du.extract(['dudx']) - + (0.01/torch.pi)*ddu.extract(['ddudxdx']) + ) + + # define initial condition + def initial_condition(input_, output_): + u_expected = -torch.sin(torch.pi*input_.extract(['x'])) + return output_.extract(['u']) - u_expected + + # assign output/ spatial and temporal variables + output_variables = ['u'] + spatial_domain = CartesianDomain({'x': [-1, 1]}) + temporal_domain = CartesianDomain({'t': [0, 1]}) + + # problem condition statement + conditions = { + 'gamma1': Condition(location=CartesianDomain({'x': -1, 't': [0, 1]}), equation=FixedValue(0.)), + 'gamma2': Condition(location=CartesianDomain({'x': 1, 't': [0, 1]}), equation=FixedValue(0.)), + 't0': Condition(location=CartesianDomain({'x': [-1, 1], 't': 0}), equation=Equation(initial_condition)), + 'D': Condition(location=CartesianDomain({'x': [-1, 1], 't': [0, 1]}), equation=Equation(burger_equation)), + } + +The ``Equation`` class takes as input a function (in this case it +happens twice, with ``initial_condition`` and ``burger_equation``) which +computes a residual of an equation, such as a PDE. In a problem class +such as the one above, the ``Equation`` class with such a given input is +passed as a parameter in the specified ``Condition``. + +The ``FixedValue`` class takes as input a value of same dimensions of +the output functions; this class can be used to enforced a fixed value +for a specific condition, e.g. Dirichlet boundary conditions, as it +happens for instance in our example. + +Once the equations are set as above in the problem conditions, the PINN +solver will aim to minimize the residuals described in each equation in +the training phase. + +Available classes of equations include also: - ``FixedGradient`` and +``FixedFlux``: they work analogously to ``FixedValue`` class, where we +can require a constant value to be enforced, respectively, on the +gradient of the solution or the divergence of the solution; - +``Laplace``: it can be used to enforce the laplacian of the solution to +be zero; - ``SystemEquation``: we can enforce multiple conditions on the +same subdomain through this class, passing a list of residual equations +defined in the problem. + +Defining a new Equation class +----------------------------- + +``Equation`` classes can be also inherited to define a new class. As +example, we can see how to rewrite the above problem introducing a new +class ``Burgers1D``; during the class call, we can pass the viscosity +parameter :math:`\nu`: + +.. code:: ipython3 + + class Burgers1DEquation(Equation): + + def __init__(self, nu = 0.): + """ + Burgers1D class. This class can be + used to enforce the solution u to solve the viscous Burgers 1D Equation. + + :param torch.float32 nu: the viscosity coefficient. Default value is set to 0. + """ + self.nu = nu + + def equation(input_, output_): + return grad(output_, input_, d='x') +\ + output_*grad(output_, input_, d='t') -\ + self.nu*laplacian(output_, input_, d='x') + + + super().__init__(equation) + +Now we can just pass the above class as input for the last condition, +setting :math:`\nu= \frac{0.01}{\pi}`: + +.. code:: ipython3 + + class Burgers1D(TimeDependentProblem, SpatialProblem): + + # define initial condition + def initial_condition(input_, output_): + u_expected = -torch.sin(torch.pi*input_.extract(['x'])) + return output_.extract(['u']) - u_expected + + # assign output/ spatial and temporal variables + output_variables = ['u'] + spatial_domain = CartesianDomain({'x': [-1, 1]}) + temporal_domain = CartesianDomain({'t': [0, 1]}) + + # problem condition statement + conditions = { + 'gamma1': Condition(location=CartesianDomain({'x': -1, 't': [0, 1]}), equation=FixedValue(0.)), + 'gamma2': Condition(location=CartesianDomain({'x': 1, 't': [0, 1]}), equation=FixedValue(0.)), + 't0': Condition(location=CartesianDomain({'x': [-1, 1], 't': 0}), equation=Equation(initial_condition)), + 'D': Condition(location=CartesianDomain({'x': [-1, 1], 't': [0, 1]}), equation=Burgers1DEquation(0.01/torch.pi)), + } + +What’s next? +------------ + +Congratulations on completing the ``Equation`` class tutorial of +**PINA**! As we have seen, you can build new classes that inherits +``Equation`` to store more complex equations, as the Burgers 1D +equation, only requiring to pass the characteristic coefficients of the +problem. From now on, you can: - define additional complex equation +classes (e.g. ``SchrodingerEquation``, ``NavierStokeEquation``..) - +define more ``FixedOperator`` (e.g. ``FixedCurl``) diff --git a/tutorials/README.md b/tutorials/README.md index 10ba0c69e..1dfe601e1 100644 --- a/tutorials/README.md +++ b/tutorials/README.md @@ -7,6 +7,7 @@ In this folder we collect useful tutorials in order to understand the principles | Description | Tutorial | |---------------|-----------| Introduction to PINA for Physics Informed Neural Networks training|[[.ipynb](tutorial1/tutorial.ipynb), [.py](tutorial1/tutorial.py), [.html](http://mathlab.github.io/PINA/_rst/tutorials/tutorial1/tutorial.html)]| +Introduction to PINA `Equation` class|[[.ipynb](tutorial12/tutorial.ipynb), [.py](tutorial12/tutorial.py), [.html](http://mathlab.github.io/PINA/_rst/tutorials/tutorial12/tutorial.html)]| PINA and PyTorch Lightning, training tips and visualizations|[[.ipynb](tutorial11/tutorial.ipynb), [.py](tutorial11/tutorial.py), [.html](http://mathlab.github.io/PINA/_rst/tutorials/tutorial11/tutorial.html)]| Building custom geometries with PINA `Location` class|[[.ipynb](tutorial6/tutorial.ipynb), [.py](tutorial6/tutorial.py), [.html](http://mathlab.github.io/PINA/_rst/tutorials/tutorial6/tutorial.html)]| diff --git a/tutorials/tutorial12/tutorial.ipynb b/tutorials/tutorial12/tutorial.ipynb new file mode 100644 index 000000000..7c3d6118c --- /dev/null +++ b/tutorials/tutorial12/tutorial.ipynb @@ -0,0 +1,235 @@ +{ + "cells": [ + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "# Tutorial: The `Equation` Class" + ] + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "In this tutorial, we will show how to use the `Equation` Class in PINA. Specifically, we will see how use the Class and its inherited classes to enforce residuals minimization in PINNs." + ] + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "# Example: The Burgers 1D equation" + ] + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "We will start implementing the viscous Burgers 1D problem Class, described as follows:\n", + "\n", + "\n", + "$$\n", + "\\begin{equation}\n", + "\\begin{cases}\n", + "\\frac{\\partial u}{\\partial t} + u \\frac{\\partial u}{\\partial x} &= \\nu \\frac{\\partial^2 u}{ \\partial x^2}, \\quad x\\in(0,1), \\quad t>0\\\\\n", + "u(x,0) &= -\\sin (\\pi x)\\\\\n", + "u(x,t) &= 0 \\quad x = \\pm 1\\\\\n", + "\\end{cases}\n", + "\\end{equation}\n", + "$$\n", + "\n", + "where we set $ \\nu = \\frac{0.01}{\\pi}$.\n", + "\n", + "In the class that models this problem we will see in action the `Equation` class and one of its inherited classes, the `FixedValue` class. " + ] + }, + { + "cell_type": "code", + "execution_count": 7, + "metadata": {}, + "outputs": [], + "source": [ + "#useful imports\n", + "from pina.problem import SpatialProblem, TimeDependentProblem\n", + "from pina.equation import Equation, FixedValue, FixedGradient, FixedFlux\n", + "from pina.geometry import CartesianDomain\n", + "import torch\n", + "from pina.operators import grad, laplacian\n", + "from pina import Condition\n", + "\n" + ] + }, + { + "cell_type": "code", + "execution_count": 6, + "metadata": {}, + "outputs": [], + "source": [ + "class Burgers1D(TimeDependentProblem, SpatialProblem):\n", + "\n", + " # define the burger equation\n", + " def burger_equation(input_, output_):\n", + " du = grad(output_, input_)\n", + " ddu = grad(du, input_, components=['dudx'])\n", + " return (\n", + " du.extract(['dudt']) +\n", + " output_.extract(['u'])*du.extract(['dudx']) -\n", + " (0.01/torch.pi)*ddu.extract(['ddudxdx'])\n", + " )\n", + "\n", + " # define initial condition\n", + " def initial_condition(input_, output_):\n", + " u_expected = -torch.sin(torch.pi*input_.extract(['x']))\n", + " return output_.extract(['u']) - u_expected\n", + "\n", + " # assign output/ spatial and temporal variables\n", + " output_variables = ['u']\n", + " spatial_domain = CartesianDomain({'x': [-1, 1]})\n", + " temporal_domain = CartesianDomain({'t': [0, 1]})\n", + "\n", + " # problem condition statement\n", + " conditions = {\n", + " 'gamma1': Condition(location=CartesianDomain({'x': -1, 't': [0, 1]}), equation=FixedValue(0.)),\n", + " 'gamma2': Condition(location=CartesianDomain({'x': 1, 't': [0, 1]}), equation=FixedValue(0.)),\n", + " 't0': Condition(location=CartesianDomain({'x': [-1, 1], 't': 0}), equation=Equation(initial_condition)),\n", + " 'D': Condition(location=CartesianDomain({'x': [-1, 1], 't': [0, 1]}), equation=Equation(burger_equation)),\n", + " }" + ] + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "\n", + "The `Equation` class takes as input a function (in this case it happens twice, with `initial_condition` and `burger_equation`) which computes a residual of an equation, such as a PDE. In a problem class such as the one above, the `Equation` class with such a given input is passed as a parameter in the specified `Condition`. \n", + "\n", + "The `FixedValue` class takes as input a value of same dimensions of the output functions; this class can be used to enforced a fixed value for a specific condition, e.g. Dirichlet boundary conditions, as it happens for instance in our example.\n", + "\n", + "Once the equations are set as above in the problem conditions, the PINN solver will aim to minimize the residuals described in each equation in the training phase. " + ] + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "Available classes of equations include also:\n", + "- `FixedGradient` and `FixedFlux`: they work analogously to `FixedValue` class, where we can require a constant value to be enforced, respectively, on the gradient of the solution or the divergence of the solution;\n", + "- `Laplace`: it can be used to enforce the laplacian of the solution to be zero;\n", + "- `SystemEquation`: we can enforce multiple conditions on the same subdomain through this class, passing a list of residual equations defined in the problem.\n" + ] + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "# Defining a new Equation class" + ] + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "`Equation` classes can be also inherited to define a new class. As example, we can see how to rewrite the above problem introducing a new class `Burgers1D`; during the class call, we can pass the viscosity parameter $\\nu$:" + ] + }, + { + "cell_type": "code", + "execution_count": 13, + "metadata": {}, + "outputs": [], + "source": [ + "class Burgers1DEquation(Equation):\n", + " \n", + " def __init__(self, nu = 0.):\n", + " \"\"\"\n", + " Burgers1D class. This class can be\n", + " used to enforce the solution u to solve the viscous Burgers 1D Equation.\n", + " \n", + " :param torch.float32 nu: the viscosity coefficient. Default value is set to 0.\n", + " \"\"\"\n", + " self.nu = nu \n", + " \n", + " def equation(input_, output_):\n", + " return grad(output_, input_, d='x') +\\\n", + " output_*grad(output_, input_, d='t') -\\\n", + " self.nu*laplacian(output_, input_, d='x')\n", + "\n", + " \n", + " super().__init__(equation)" + ] + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "Now we can just pass the above class as input for the last condition, setting $\\nu= \\frac{0.01}{\\pi}$:" + ] + }, + { + "cell_type": "code", + "execution_count": 14, + "metadata": {}, + "outputs": [], + "source": [ + "class Burgers1D(TimeDependentProblem, SpatialProblem):\n", + "\n", + " # define initial condition\n", + " def initial_condition(input_, output_):\n", + " u_expected = -torch.sin(torch.pi*input_.extract(['x']))\n", + " return output_.extract(['u']) - u_expected\n", + "\n", + " # assign output/ spatial and temporal variables\n", + " output_variables = ['u']\n", + " spatial_domain = CartesianDomain({'x': [-1, 1]})\n", + " temporal_domain = CartesianDomain({'t': [0, 1]})\n", + "\n", + " # problem condition statement\n", + " conditions = {\n", + " 'gamma1': Condition(location=CartesianDomain({'x': -1, 't': [0, 1]}), equation=FixedValue(0.)),\n", + " 'gamma2': Condition(location=CartesianDomain({'x': 1, 't': [0, 1]}), equation=FixedValue(0.)),\n", + " 't0': Condition(location=CartesianDomain({'x': [-1, 1], 't': 0}), equation=Equation(initial_condition)),\n", + " 'D': Condition(location=CartesianDomain({'x': [-1, 1], 't': [0, 1]}), equation=Burgers1DEquation(0.01/torch.pi)),\n", + " }" + ] + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "# What's next?" + ] + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "Congratulations on completing the `Equation` class tutorial of **PINA**! As we have seen, you can build new classes that inherits `Equation` to store more complex equations, as the Burgers 1D equation, only requiring to pass the characteristic coefficients of the problem. \n", + "From now on, you can:\n", + "- define additional complex equation classes (e.g. `SchrodingerEquation`, `NavierStokeEquation`..)\n", + "- define more `FixedOperator` (e.g. `FixedCurl`)" + ] + } + ], + "metadata": { + "kernelspec": { + "display_name": "pina", + "language": "python", + "name": "python3" + }, + "language_info": { + "codemirror_mode": { + "name": "ipython", + "version": 3 + }, + "file_extension": ".py", + "mimetype": "text/x-python", + "name": "python", + "nbconvert_exporter": "python", + "pygments_lexer": "ipython3", + "version": "3.1.0" + }, + "orig_nbformat": 4 + }, + "nbformat": 4, + "nbformat_minor": 2 +} diff --git a/tutorials/tutorial12/tutorial.py b/tutorials/tutorial12/tutorial.py new file mode 100644 index 000000000..9b71eb4cb --- /dev/null +++ b/tutorials/tutorial12/tutorial.py @@ -0,0 +1,142 @@ +#!/usr/bin/env python +# coding: utf-8 + +# # Tutorial: The `Equation` Class + +# In this tutorial, we will show how to use the `Equation` Class in PINA. Specifically, we will see how use the Class and its inherited classes to enforce residuals minimization in PINNs. + +# # Example: The Burgers 1D equation + +# We will start implementing the viscous Burgers 1D problem Class, described as follows: +# +# +# $$ +# \begin{equation} +# \begin{cases} +# \frac{\partial u}{\partial t} + u \frac{\partial u}{\partial x} &= \nu \frac{\partial^2 u}{ \partial x^2}, \quad x\in(0,1), \quad t>0\\ +# u(x,0) &= -\sin (\pi x)\\ +# u(x,t) &= 0 \quad x = \pm 1\\ +# \end{cases} +# \end{equation} +# $$ +# +# where we set $ \nu = \frac{0.01}{\pi}$. +# +# In the class that models this problem we will see in action the `Equation` class and one of its inherited classes, the `FixedValue` class. + +# In[7]: + + +#useful imports +from pina.problem import SpatialProblem, TimeDependentProblem +from pina.equation import Equation, FixedValue, FixedGradient, FixedFlux +from pina.geometry import CartesianDomain +import torch +from pina.operators import grad, laplacian +from pina import Condition + + +# In[6]: + + +class Burgers1D(TimeDependentProblem, SpatialProblem): + + # define the burger equation + def burger_equation(input_, output_): + du = grad(output_, input_) + ddu = grad(du, input_, components=['dudx']) + return ( + du.extract(['dudt']) + + output_.extract(['u'])*du.extract(['dudx']) - + (0.01/torch.pi)*ddu.extract(['ddudxdx']) + ) + + # define initial condition + def initial_condition(input_, output_): + u_expected = -torch.sin(torch.pi*input_.extract(['x'])) + return output_.extract(['u']) - u_expected + + # assign output/ spatial and temporal variables + output_variables = ['u'] + spatial_domain = CartesianDomain({'x': [-1, 1]}) + temporal_domain = CartesianDomain({'t': [0, 1]}) + + # problem condition statement + conditions = { + 'gamma1': Condition(location=CartesianDomain({'x': -1, 't': [0, 1]}), equation=FixedValue(0.)), + 'gamma2': Condition(location=CartesianDomain({'x': 1, 't': [0, 1]}), equation=FixedValue(0.)), + 't0': Condition(location=CartesianDomain({'x': [-1, 1], 't': 0}), equation=Equation(initial_condition)), + 'D': Condition(location=CartesianDomain({'x': [-1, 1], 't': [0, 1]}), equation=Equation(burger_equation)), + } + + +# +# The `Equation` class takes as input a function (in this case it happens twice, with `initial_condition` and `burger_equation`) which computes a residual of an equation, such as a PDE. In a problem class such as the one above, the `Equation` class with such a given input is passed as a parameter in the specified `Condition`. +# +# The `FixedValue` class takes as input a value of same dimensions of the output functions; this class can be used to enforced a fixed value for a specific condition, e.g. Dirichlet boundary conditions, as it happens for instance in our example. +# +# Once the equations are set as above in the problem conditions, the PINN solver will aim to minimize the residuals described in each equation in the training phase. + +# Available classes of equations include also: +# - `FixedGradient` and `FixedFlux`: they work analogously to `FixedValue` class, where we can require a constant value to be enforced, respectively, on the gradient of the solution or the divergence of the solution; +# - `Laplace`: it can be used to enforce the laplacian of the solution to be zero; +# - `SystemEquation`: we can enforce multiple conditions on the same subdomain through this class, passing a list of residual equations defined in the problem. +# + +# # Defining a new Equation class + +# `Equation` classes can be also inherited to define a new class. As example, we can see how to rewrite the above problem introducing a new class `Burgers1D`; during the class call, we can pass the viscosity parameter $\nu$: + +# In[13]: + + +class Burgers1DEquation(Equation): + + def __init__(self, nu = 0.): + """ + Burgers1D class. This class can be + used to enforce the solution u to solve the viscous Burgers 1D Equation. + + :param torch.float32 nu: the viscosity coefficient. Default value is set to 0. + """ + self.nu = nu + + def equation(input_, output_): + return grad(output_, input_, d='x') + output_*grad(output_, input_, d='t') - self.nu*laplacian(output_, input_, d='x') + + + super().__init__(equation) + + +# Now we can just pass the above class as input for the last condition, setting $\nu= \frac{0.01}{\pi}$: + +# In[14]: + + +class Burgers1D(TimeDependentProblem, SpatialProblem): + + # define initial condition + def initial_condition(input_, output_): + u_expected = -torch.sin(torch.pi*input_.extract(['x'])) + return output_.extract(['u']) - u_expected + + # assign output/ spatial and temporal variables + output_variables = ['u'] + spatial_domain = CartesianDomain({'x': [-1, 1]}) + temporal_domain = CartesianDomain({'t': [0, 1]}) + + # problem condition statement + conditions = { + 'gamma1': Condition(location=CartesianDomain({'x': -1, 't': [0, 1]}), equation=FixedValue(0.)), + 'gamma2': Condition(location=CartesianDomain({'x': 1, 't': [0, 1]}), equation=FixedValue(0.)), + 't0': Condition(location=CartesianDomain({'x': [-1, 1], 't': 0}), equation=Equation(initial_condition)), + 'D': Condition(location=CartesianDomain({'x': [-1, 1], 't': [0, 1]}), equation=Burgers1DEquation(0.01/torch.pi)), + } + + +# # What's next? + +# Congratulations on completing the `Equation` class tutorial of **PINA**! As we have seen, you can build new classes that inherits `Equation` to store more complex equations, as the Burgers 1D equation, only requiring to pass the characteristic coefficients of the problem. +# From now on, you can: +# - define additional complex equation classes (e.g. `SchrodingerEquation`, `NavierStokeEquation`..) +# - define more `FixedOperator` (e.g. `FixedCurl`) From 3d988b34867bcfd40df9a0e05be0a5075eb22743 Mon Sep 17 00:00:00 2001 From: Nicola Demo Date: Tue, 30 Apr 2024 11:29:19 +0200 Subject: [PATCH 04/14] Update stokes.py --- examples/problems/stokes.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/examples/problems/stokes.py b/examples/problems/stokes.py index dbf40158f..f136d64ad 100644 --- a/examples/problems/stokes.py +++ b/examples/problems/stokes.py @@ -1,4 +1,4 @@ -""" Navier Stokes Problem """ +""" Steady Stokes Problem """ import torch from pina.problem import SpatialProblem From f21888d63d528c2d573daee80917db17ead4ba23 Mon Sep 17 00:00:00 2001 From: Dario Coscia <93731561+dario-coscia@users.noreply.github.com> Date: Tue, 30 Apr 2024 18:52:09 +0200 Subject: [PATCH 05/14] Update label_tensor.py cpu/gpu (#292) * Update label_tensor.py cpu/gpu * Update test_adaptive_refinment_callbacks.py * Update test_optimizer_callbacks.py --- pina/label_tensor.py | 4 ++-- tests/test_callbacks/test_adaptive_refinment_callbacks.py | 1 + tests/test_callbacks/test_optimizer_callbacks.py | 1 + 3 files changed, 4 insertions(+), 2 deletions(-) diff --git a/pina/label_tensor.py b/pina/label_tensor.py index fe8e1a850..c8a41f7b4 100644 --- a/pina/label_tensor.py +++ b/pina/label_tensor.py @@ -176,7 +176,7 @@ def cuda(self, *args, **kwargs): tmp = super().cuda(*args, **kwargs) new = self.__class__.clone(self) new.data = tmp.data - return tmp + return new def cpu(self, *args, **kwargs): """ @@ -185,7 +185,7 @@ def cpu(self, *args, **kwargs): tmp = super().cpu(*args, **kwargs) new = self.__class__.clone(self) new.data = tmp.data - return tmp + return new def extract(self, label_to_extract): """ diff --git a/tests/test_callbacks/test_adaptive_refinment_callbacks.py b/tests/test_callbacks/test_adaptive_refinment_callbacks.py index fb74367e6..214257d95 100644 --- a/tests/test_callbacks/test_adaptive_refinment_callbacks.py +++ b/tests/test_callbacks/test_adaptive_refinment_callbacks.py @@ -71,6 +71,7 @@ def test_r3refinment_routine(): # make the trainer trainer = Trainer(solver=solver, callbacks=[R3Refinement(sample_every=1)], + accelerator='cpu', max_epochs=5) trainer.train() diff --git a/tests/test_callbacks/test_optimizer_callbacks.py b/tests/test_callbacks/test_optimizer_callbacks.py index 6c167b600..0b0aabaab 100644 --- a/tests/test_callbacks/test_optimizer_callbacks.py +++ b/tests/test_callbacks/test_optimizer_callbacks.py @@ -84,5 +84,6 @@ def test_switch_optimizer_routine(): new_optimizers_kwargs={'lr': 0.01}, epoch_switch=3) ], + accelerator='cpu', max_epochs=5) trainer.train() From c724b38f901f8b0f77c31f3452a5ef02499ba02b Mon Sep 17 00:00:00 2001 From: Dario Coscia <93731561+dario-coscia@users.noreply.github.com> Date: Thu, 2 May 2024 17:41:56 +0200 Subject: [PATCH 06/14] Update monthly-tag.yml (#295) --- .github/workflows/monthly-tag.yml | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/.github/workflows/monthly-tag.yml b/.github/workflows/monthly-tag.yml index 94251d864..b77e42220 100644 --- a/.github/workflows/monthly-tag.yml +++ b/.github/workflows/monthly-tag.yml @@ -11,7 +11,7 @@ jobs: strategy: matrix: os: [windows-latest, macos-latest, ubuntu-latest] - python-version: [3.7, 3.8] + python-version: [3.8, 3.9, '3.10', '3.11', '3.12'] steps: - uses: actions/checkout@v2 - name: Set up Python From 190aad383c822955e8cec35dbc77bbaad7c3d365 Mon Sep 17 00:00:00 2001 From: Dario Coscia <93731561+dario-coscia@users.noreply.github.com> Date: Fri, 10 May 2024 14:07:01 +0200 Subject: [PATCH 07/14] PINN variants addition and Solvers Update (#263) * gpinn/basepinn new classes, pinn restructure * codacy fix gpinn/basepinn/pinn * inverse problem fix * Causal PINN (#267) * fix GPU training in inverse problem (#283) * Create a `compute_residual` attribute for `PINNInterface` * Modify dataloading in solvers (#286) * Modify PINNInterface by removing _loss_phys, _loss_data * Adding in PINNInterface a variable to track the current condition during training * Modify GPINN,PINN,CausalPINN to match changes in PINNInterface * Competitive Pinn Addition (#288) * fixing after rebase/ fix loss * fixing final issues --------- Co-authored-by: Dario Coscia * Modify min max formulation to max min for paper consistency * Adding SAPINN solver (#291) * rom solver * fix import --------- Co-authored-by: Dario Coscia Co-authored-by: Anna Ivagnes <75523024+annaivagnes@users.noreply.github.com> Co-authored-by: valc89 <103250118+valc89@users.noreply.github.com> Co-authored-by: Monthly Tag bot Co-authored-by: Nicola Demo --- docs/source/_rst/_code.rst | 6 + docs/source/_rst/solvers/basepinn.rst | 7 + docs/source/_rst/solvers/causalpinn.rst | 7 + docs/source/_rst/solvers/competitivepinn.rst | 7 + docs/source/_rst/solvers/gpinn.rst | 7 + docs/source/_rst/solvers/pinn.rst | 2 +- docs/source/_rst/solvers/rom.rst | 7 + docs/source/_rst/solvers/sapinn.rst | 7 + pina/model/avno.py | 4 +- pina/solvers/__init__.py | 21 +- pina/solvers/garom.py | 9 +- pina/solvers/pinn.py | 232 --------- pina/solvers/pinns/__init__.py | 15 + pina/solvers/pinns/basepinn.py | 247 ++++++++++ pina/solvers/pinns/causalpinn.py | 221 +++++++++ pina/solvers/pinns/competitive_pinn.py | 360 ++++++++++++++ pina/solvers/pinns/gpinn.py | 134 +++++ pina/solvers/pinns/pinn.py | 170 +++++++ pina/solvers/pinns/sapinn.py | 494 +++++++++++++++++++ pina/solvers/rom.py | 190 +++++++ pina/solvers/solver.py | 15 + pina/solvers/supervised.py | 54 +- pina/trainer.py | 7 + tests/test_solvers/test_causalpinn.py | 266 ++++++++++ tests/test_solvers/test_competitive_pinn.py | 418 ++++++++++++++++ tests/test_solvers/test_gpinn.py | 432 ++++++++++++++++ tests/test_solvers/test_pinn.py | 315 ++++++++---- tests/test_solvers/test_rom_solver.py | 105 ++++ tests/test_solvers/test_sapinn.py | 437 ++++++++++++++++ 29 files changed, 3838 insertions(+), 358 deletions(-) create mode 100644 docs/source/_rst/solvers/basepinn.rst create mode 100644 docs/source/_rst/solvers/causalpinn.rst create mode 100644 docs/source/_rst/solvers/competitivepinn.rst create mode 100644 docs/source/_rst/solvers/gpinn.rst create mode 100644 docs/source/_rst/solvers/rom.rst create mode 100644 docs/source/_rst/solvers/sapinn.rst delete mode 100644 pina/solvers/pinn.py create mode 100644 pina/solvers/pinns/__init__.py create mode 100644 pina/solvers/pinns/basepinn.py create mode 100644 pina/solvers/pinns/causalpinn.py create mode 100644 pina/solvers/pinns/competitive_pinn.py create mode 100644 pina/solvers/pinns/gpinn.py create mode 100644 pina/solvers/pinns/pinn.py create mode 100644 pina/solvers/pinns/sapinn.py create mode 100644 pina/solvers/rom.py create mode 100644 tests/test_solvers/test_causalpinn.py create mode 100644 tests/test_solvers/test_competitive_pinn.py create mode 100644 tests/test_solvers/test_gpinn.py create mode 100644 tests/test_solvers/test_rom_solver.py create mode 100644 tests/test_solvers/test_sapinn.py diff --git a/docs/source/_rst/_code.rst b/docs/source/_rst/_code.rst index 77072a507..d954920e9 100644 --- a/docs/source/_rst/_code.rst +++ b/docs/source/_rst/_code.rst @@ -35,8 +35,14 @@ Solvers :titlesonly: SolverInterface + PINNInterface PINN + GPINN + CausalPINN + CompetitivePINN + SAPINN Supervised solver + ReducedOrderModelSolver GAROM diff --git a/docs/source/_rst/solvers/basepinn.rst b/docs/source/_rst/solvers/basepinn.rst new file mode 100644 index 000000000..c6507953d --- /dev/null +++ b/docs/source/_rst/solvers/basepinn.rst @@ -0,0 +1,7 @@ +PINNInterface +================= +.. currentmodule:: pina.solvers.pinns.basepinn + +.. autoclass:: PINNInterface + :members: + :show-inheritance: \ No newline at end of file diff --git a/docs/source/_rst/solvers/causalpinn.rst b/docs/source/_rst/solvers/causalpinn.rst new file mode 100644 index 000000000..28f7f15ea --- /dev/null +++ b/docs/source/_rst/solvers/causalpinn.rst @@ -0,0 +1,7 @@ +CausalPINN +============== +.. currentmodule:: pina.solvers.pinns.causalpinn + +.. autoclass:: CausalPINN + :members: + :show-inheritance: \ No newline at end of file diff --git a/docs/source/_rst/solvers/competitivepinn.rst b/docs/source/_rst/solvers/competitivepinn.rst new file mode 100644 index 000000000..2bbe242b7 --- /dev/null +++ b/docs/source/_rst/solvers/competitivepinn.rst @@ -0,0 +1,7 @@ +CompetitivePINN +================= +.. currentmodule:: pina.solvers.pinns.competitive_pinn + +.. autoclass:: CompetitivePINN + :members: + :show-inheritance: \ No newline at end of file diff --git a/docs/source/_rst/solvers/gpinn.rst b/docs/source/_rst/solvers/gpinn.rst new file mode 100644 index 000000000..ee076a5d7 --- /dev/null +++ b/docs/source/_rst/solvers/gpinn.rst @@ -0,0 +1,7 @@ +GPINN +====== +.. currentmodule:: pina.solvers.pinns.gpinn + +.. autoclass:: GPINN + :members: + :show-inheritance: \ No newline at end of file diff --git a/docs/source/_rst/solvers/pinn.rst b/docs/source/_rst/solvers/pinn.rst index 3e9b2ef01..e1c2b59cd 100644 --- a/docs/source/_rst/solvers/pinn.rst +++ b/docs/source/_rst/solvers/pinn.rst @@ -1,6 +1,6 @@ PINN ====== -.. currentmodule:: pina.solvers.pinn +.. currentmodule:: pina.solvers.pinns.pinn .. autoclass:: PINN :members: diff --git a/docs/source/_rst/solvers/rom.rst b/docs/source/_rst/solvers/rom.rst new file mode 100644 index 000000000..3ee534bb5 --- /dev/null +++ b/docs/source/_rst/solvers/rom.rst @@ -0,0 +1,7 @@ +ReducedOrderModelSolver +========================== +.. currentmodule:: pina.solvers.rom + +.. autoclass:: ReducedOrderModelSolver + :members: + :show-inheritance: \ No newline at end of file diff --git a/docs/source/_rst/solvers/sapinn.rst b/docs/source/_rst/solvers/sapinn.rst new file mode 100644 index 000000000..b20891fff --- /dev/null +++ b/docs/source/_rst/solvers/sapinn.rst @@ -0,0 +1,7 @@ +SAPINN +====== +.. currentmodule:: pina.solvers.pinns.sapinn + +.. autoclass:: SAPINN + :members: + :show-inheritance: \ No newline at end of file diff --git a/pina/model/avno.py b/pina/model/avno.py index 878185bc9..2ac3b3f7e 100644 --- a/pina/model/avno.py +++ b/pina/model/avno.py @@ -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=2) + new_batch = concatenate((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=2) + new_batch = concatenate((new_batch, points_tmp), dim=-1) new_batch = self._projection_operator(new_batch) return new_batch diff --git a/pina/solvers/__init__.py b/pina/solvers/__init__.py index 0562dc2d1..2751e481c 100644 --- a/pina/solvers/__init__.py +++ b/pina/solvers/__init__.py @@ -1,6 +1,19 @@ -__all__ = ["PINN", "GAROM", "SupervisedSolver", "SolverInterface"] +__all__ = [ + "SolverInterface", + "PINNInterface", + "PINN", + "GPINN", + "CausalPINN", + "CompetitivePINN", + "SAPINN", + "SupervisedSolver", + "ReducedOrderModelSolver", + "GAROM", + ] -from .garom import GAROM -from .pinn import PINN -from .supervised import SupervisedSolver from .solver import SolverInterface +from .pinns import * +from .supervised import SupervisedSolver +from .rom import ReducedOrderModelSolver +from .garom import GAROM + diff --git a/pina/solvers/garom.py b/pina/solvers/garom.py index 08856704f..d6cd6246e 100644 --- a/pina/solvers/garom.py +++ b/pina/solvers/garom.py @@ -253,18 +253,11 @@ def training_step(self, batch, batch_idx): :rtype: LabelTensor """ - dataloader = self.trainer.train_dataloader condition_idx = batch["condition"] for condition_id in range(condition_idx.min(), condition_idx.max() + 1): - if sys.version_info >= (3, 8): - condition_name = dataloader.condition_names[condition_id] - else: - condition_name = dataloader.loaders.condition_names[ - condition_id - ] - + condition_name = self._dataloader.condition_names[condition_id] condition = self.problem.conditions[condition_name] pts = batch["pts"].detach() out = batch["output"] diff --git a/pina/solvers/pinn.py b/pina/solvers/pinn.py deleted file mode 100644 index 008034f36..000000000 --- a/pina/solvers/pinn.py +++ /dev/null @@ -1,232 +0,0 @@ -""" Module for PINN """ - -import torch - -try: - from torch.optim.lr_scheduler import LRScheduler # torch >= 2.0 -except ImportError: - from torch.optim.lr_scheduler import ( - _LRScheduler as LRScheduler, - ) # torch < 2.0 - -import sys -from torch.optim.lr_scheduler import ConstantLR - -from .solver import SolverInterface -from ..label_tensor import LabelTensor -from ..utils import check_consistency -from ..loss import LossInterface -from ..problem import InverseProblem -from torch.nn.modules.loss import _Loss - -torch.pi = torch.acos(torch.zeros(1)).item() * 2 # which is 3.1415927410125732 - - -class PINN(SolverInterface): - """ - PINN solver class. This class implements Physics Informed Neural - Network solvers, using a user specified ``model`` to solve a specific - ``problem``. It can be used for solving both forward and inverse problems. - - .. seealso:: - - **Original reference**: Karniadakis, G. E., Kevrekidis, I. G., Lu, L., - Perdikaris, P., Wang, S., & Yang, L. (2021). - Physics-informed machine learning. Nature Reviews Physics, 3(6), 422-440. - `_. - """ - - def __init__( - self, - problem, - model, - extra_features=None, - loss=torch.nn.MSELoss(), - optimizer=torch.optim.Adam, - optimizer_kwargs={"lr": 0.001}, - scheduler=ConstantLR, - scheduler_kwargs={"factor": 1, "total_iters": 0}, - ): - """ - :param AbstractProblem problem: The formulation of the problem. - :param torch.nn.Module model: The neural network model to use. - :param torch.nn.Module loss: The loss function used as minimizer, - default :class:`torch.nn.MSELoss`. - :param torch.nn.Module extra_features: The additional input - 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 torch.optim.LRScheduler scheduler: Learning - rate scheduler. - :param dict scheduler_kwargs: LR scheduler constructor keyword args. - """ - super().__init__( - models=[model], - problem=problem, - optimizers=[optimizer], - optimizers_kwargs=[optimizer_kwargs], - extra_features=extra_features, - ) - - # check consistency - check_consistency(scheduler, LRScheduler, subclass=True) - check_consistency(scheduler_kwargs, dict) - check_consistency(loss, (LossInterface, _Loss), subclass=False) - - # assign variables - self._scheduler = scheduler(self.optimizers[0], **scheduler_kwargs) - self._loss = loss - self._neural_net = self.models[0] - - # inverse problem handling - if isinstance(self.problem, InverseProblem): - self._params = self.problem.unknown_parameters - else: - self._params = None - - def forward(self, x): - """ - Forward pass implementation for the PINN - solver. - - :param torch.Tensor x: Input tensor. - :return: PINN solution. - :rtype: torch.Tensor - """ - return self.neural_net(x) - - def configure_optimizers(self): - """ - Optimizer configuration for the PINN - solver. - - :return: The optimizers and the schedulers - :rtype: tuple(list, list) - """ - # if the problem is an InverseProblem, add the unknown parameters - # to the parameters that the optimizer needs to optimize - if isinstance(self.problem, InverseProblem): - self.optimizers[0].add_param_group( - { - "params": [ - self._params[var] - for var in self.problem.unknown_variables - ] - } - ) - return self.optimizers, [self.scheduler] - - def _clamp_inverse_problem_params(self): - for v in self._params: - self._params[v].data.clamp_( - self.problem.unknown_parameter_domain.range_[v][0], - self.problem.unknown_parameter_domain.range_[v][1], - ) - - def _loss_data(self, input, output): - return self.loss(self.forward(input), output) - - def _loss_phys(self, samples, equation): - try: - residual = equation.residual(samples, self.forward(samples)) - except ( - TypeError - ): # this occurs when the function has three inputs, i.e. inverse problem - residual = equation.residual( - samples, self.forward(samples), self._params - ) - return self.loss( - torch.zeros_like(residual, requires_grad=True), residual - ) - - def training_step(self, batch, batch_idx): - """ - PINN solver training step. - - :param batch: The batch element in the dataloader. - :type batch: tuple - :param batch_idx: The batch index. - :type batch_idx: int - :return: The sum of the loss functions. - :rtype: LabelTensor - """ - - dataloader = self.trainer.train_dataloader - condition_losses = [] - - condition_idx = batch["condition"] - - for condition_id in range(condition_idx.min(), condition_idx.max() + 1): - - if sys.version_info >= (3, 8): - condition_name = dataloader.condition_names[condition_id] - else: - condition_name = dataloader.loaders.condition_names[ - condition_id - ] - condition = self.problem.conditions[condition_name] - pts = batch["pts"] - - if len(batch) == 2: - samples = pts[condition_idx == condition_id] - loss = self._loss_phys(samples, condition.equation) - elif len(batch) == 3: - samples = pts[condition_idx == condition_id] - ground_truth = batch["output"][condition_idx == condition_id] - loss = self._loss_data(samples, ground_truth) - else: - raise ValueError("Batch size not supported") - - # TODO for users this us hard to remember when creating a new solver, to fix in a smarter way - loss = loss.as_subclass(torch.Tensor) - - # # add condition losses and accumulate logging for each epoch - condition_losses.append(loss * condition.data_weight) - self.log( - condition_name + "_loss", - float(loss), - prog_bar=True, - logger=True, - on_epoch=True, - on_step=False, - ) - - # clamp unknown parameters of the InverseProblem to their domain ranges (if needed) - if isinstance(self.problem, InverseProblem): - self._clamp_inverse_problem_params() - - # TODO Fix the bug, tot_loss is a label tensor without labels - # we need to pass it as a torch tensor to make everything work - total_loss = sum(condition_losses) - self.log( - "mean_loss", - float(total_loss / len(condition_losses)), - prog_bar=True, - logger=True, - on_epoch=True, - on_step=False, - ) - - return total_loss - - @property - def scheduler(self): - """ - Scheduler for the PINN training. - """ - return self._scheduler - - @property - def neural_net(self): - """ - Neural network for the PINN training. - """ - return self._neural_net - - @property - def loss(self): - """ - Loss for the PINN training. - """ - return self._loss diff --git a/pina/solvers/pinns/__init__.py b/pina/solvers/pinns/__init__.py new file mode 100644 index 000000000..c8aa904c8 --- /dev/null +++ b/pina/solvers/pinns/__init__.py @@ -0,0 +1,15 @@ +__all__ = [ + "PINNInterface", + "PINN", + "GPINN", + "CausalPINN", + "CompetitivePINN", + "SAPINN", +] + +from .basepinn import PINNInterface +from .pinn import PINN +from .gpinn import GPINN +from .causalpinn import CausalPINN +from .competitive_pinn import CompetitivePINN +from .sapinn import SAPINN diff --git a/pina/solvers/pinns/basepinn.py b/pina/solvers/pinns/basepinn.py new file mode 100644 index 000000000..726cdf921 --- /dev/null +++ b/pina/solvers/pinns/basepinn.py @@ -0,0 +1,247 @@ +""" Module for PINN """ + +import sys +from abc import ABCMeta, abstractmethod +import torch + +from ...solvers.solver import SolverInterface +from pina.utils import check_consistency +from pina.loss import LossInterface +from pina.problem import InverseProblem +from torch.nn.modules.loss import _Loss + +torch.pi = torch.acos(torch.zeros(1)).item() * 2 # which is 3.1415927410125732 + +class PINNInterface(SolverInterface, metaclass=ABCMeta): + """ + Base PINN solver class. This class implements the Solver Interface + for Physics Informed Neural Network solvers. + + This class can be used to + define PINNs with multiple ``optimizers``, and/or ``models``. + By default it takes + an :class:`~pina.problem.abstract_problem.AbstractProblem`, so it is up + to the user to choose which problem the implemented solver inheriting from + this class is suitable for. + """ + + def __init__( + self, + models, + problem, + optimizers, + optimizers_kwargs, + extra_features, + loss, + ): + """ + :param models: Multiple torch neural network models instances. + :type models: list(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. + :param torch.nn.Module loss: The loss function used as minimizer, + default :class:`torch.nn.MSELoss`. + """ + super().__init__( + models=models, + problem=problem, + optimizers=optimizers, + optimizers_kwargs=optimizers_kwargs, + extra_features=extra_features, + ) + + # check consistency + check_consistency(loss, (LossInterface, _Loss), subclass=False) + + # assign variables + self._loss = loss + + # inverse problem handling + if isinstance(self.problem, InverseProblem): + self._params = self.problem.unknown_parameters + self._clamp_params = self._clamp_inverse_problem_params + else: + self._params = None + self._clamp_params = lambda : None + + # variable used internally to store residual losses at each epoch + # this variable save the residual at each iteration (not weighted) + self.__logged_res_losses = [] + + # variable used internally in pina for logging. This variable points to + # the current condition during the training step and returns the + # condition name. Whenever :meth:`store_log` is called the logged + # variable will be stored with name = self.__logged_metric + self.__logged_metric = None + + def training_step(self, batch, _): + """ + The Physics Informed Solver Training Step. This function takes care + of the physics informed training step, and it must not be override + if not intentionally. It handles the batching mechanism, the workload + division for the various conditions, the inverse problem clamping, + and loggers. + + :param tuple batch: The batch element in the dataloader. + :param int batch_idx: The batch index. + :return: The sum of the loss functions. + :rtype: LabelTensor + """ + + condition_losses = [] + condition_idx = batch["condition"] + + 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["pts"] + # condition name is logged (if logs enabled) + self.__logged_metric = condition_name + + if len(batch) == 2: + samples = pts[condition_idx == condition_id] + loss = self.loss_phys(samples, condition.equation) + elif len(batch) == 3: + samples = pts[condition_idx == condition_id] + ground_truth = batch["output"][condition_idx == condition_id] + loss = self.loss_data(samples, ground_truth) + else: + raise ValueError("Batch size not supported") + + # add condition losses for each epoch + condition_losses.append(loss * condition.data_weight) + + # clamp unknown parameters in InverseProblem (if needed) + self._clamp_params() + + # total loss (must be a torch.Tensor) + total_loss = sum(condition_losses) + return total_loss.as_subclass(torch.Tensor) + + def loss_data(self, input_tensor, output_tensor): + """ + The data loss for the PINN solver. It computes the loss between + 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 + network solution. + :return: The residual loss averaged on the input coordinates + :rtype: torch.Tensor + """ + loss_value = self.loss(self.forward(input_tensor), output_tensor) + self.store_log(loss_value=float(loss_value)) + return self.loss(self.forward(input_tensor), output_tensor) + + @abstractmethod + def loss_phys(self, samples, equation): + """ + Computes the physics loss for the physics informed solver based on given + samples and equation. This method must be override by all inherited + classes and it is the core to define a new physics informed solver. + + :param LabelTensor samples: The samples to evaluate the physics loss. + :param EquationInterface equation: The governing equation + representing the physics. + :return: The physics loss calculated based on given + samples and equation. + :rtype: LabelTensor + """ + pass + + def compute_residual(self, samples, equation): + """ + Compute the residual for Physics Informed learning. This function + returns the :obj:`~pina.equation.equation.Equation` specified in the + :obj:`~pina.condition.Condition` evaluated at the ``samples`` points. + + :param LabelTensor samples: The samples to evaluate the physics loss. + :param EquationInterface equation: The governing equation + representing the physics. + :return: The residual of the neural network solution. + :rtype: LabelTensor + """ + try: + residual = equation.residual(samples, self.forward(samples)) + except ( + TypeError + ): # this occurs when the function has three inputs, i.e. inverse problem + residual = equation.residual( + samples, self.forward(samples), self._params + ) + return residual + + def store_log(self, loss_value): + """ + Stores the loss value in the logger. This function should be + called for all conditions. It automatically handles the storing + conditions names. It must be used + anytime a specific variable wants to be stored for a specific condition. + A simple example is to use the variable to store the residual. + + :param str name: The name of the loss. + :param torch.Tensor loss_value: The value of the loss. + """ + self.log( + self.__logged_metric+'_loss', + loss_value, + prog_bar=True, + logger=True, + on_epoch=True, + on_step=False, + ) + self.__logged_res_losses.append(loss_value) + + def on_train_epoch_end(self): + """ + At the end of each epoch we free the stored losses. This function + should not be override if not intentionally. + """ + if self.__logged_res_losses: + # storing mean loss + self.__logged_metric = 'mean' + self.store_log( + sum(self.__logged_res_losses)/len(self.__logged_res_losses) + ) + # free the logged losses + self.__logged_res_losses = [] + return super().on_train_epoch_end() + + def _clamp_inverse_problem_params(self): + """ + Clamps the parameters of the inverse problem + solver to the specified ranges. + """ + for v in self._params: + self._params[v].data.clamp_( + self.problem.unknown_parameter_domain.range_[v][0], + self.problem.unknown_parameter_domain.range_[v][1], + ) + + @property + def loss(self): + """ + Loss used for training. + """ + return self._loss + + @property + def current_condition_name(self): + """ + Returns the condition name. This function can be used inside the + :meth:`loss_phys` to extract the condition at which the loss is + computed. + """ + return self.__logged_metric \ No newline at end of file diff --git a/pina/solvers/pinns/causalpinn.py b/pina/solvers/pinns/causalpinn.py new file mode 100644 index 000000000..fea0fe47b --- /dev/null +++ b/pina/solvers/pinns/causalpinn.py @@ -0,0 +1,221 @@ +""" Module for CausalPINN """ + +import torch + + +from torch.optim.lr_scheduler import ConstantLR + +from .pinn import PINN +from pina.problem import TimeDependentProblem +from pina.utils import check_consistency + + +class CausalPINN(PINN): + r""" + Causal Physics Informed Neural Network (PINN) solver class. + This class implements Causal Physics Informed Neural + Network solvers, using a user specified ``model`` to solve a specific + ``problem``. It can be used for solving both forward and inverse problems. + + The Causal Physics Informed Network aims to find + the solution :math:`\mathbf{u}:\Omega\rightarrow\mathbb{R}^m` + of the differential problem: + + .. math:: + + \begin{cases} + \mathcal{A}[\mathbf{u}](\mathbf{x})=0\quad,\mathbf{x}\in\Omega\\ + \mathcal{B}[\mathbf{u}](\mathbf{x})=0\quad, + \mathbf{x}\in\partial\Omega + \end{cases} + + minimizing the loss function + + .. math:: + \mathcal{L}_{\rm{problem}} = \frac{1}{N_t}\sum_{i=1}^{N_t} + \omega_{i}\mathcal{L}_r(t_i), + + where: + + .. math:: + \mathcal{L}_r(t) = \frac{1}{N}\sum_{i=1}^N + \mathcal{L}(\mathcal{A}[\mathbf{u}](\mathbf{x}_i, t)) + + \frac{1}{N}\sum_{i=1}^N + \mathcal{L}(\mathcal{B}[\mathbf{u}](\mathbf{x}_i, t)) + + and, + + .. math:: + \omega_i = \exp\left(\epsilon \sum_{k=1}^{i-1}\mathcal{L}_r(t_k)\right). + + :math:`\epsilon` is an hyperparameter, default set to :math:`100`, while + :math:`\mathcal{L}` is a specific loss function, + default Mean Square Error: + + .. math:: + \mathcal{L}(v) = \| v \|^2_2. + + + .. seealso:: + + **Original reference**: Wang, Sifan, Shyam Sankaran, and Paris + Perdikaris. "Respecting causality for training physics-informed + neural networks." Computer Methods in Applied Mechanics + and Engineering 421 (2024): 116813. + DOI `10.1016 `_. + + .. note:: + This class can only work for problems inheriting + from at least + :class:`~pina.problem.timedep_problem.TimeDependentProblem` class. + """ + + def __init__( + self, + problem, + model, + extra_features=None, + loss=torch.nn.MSELoss(), + optimizer=torch.optim.Adam, + optimizer_kwargs={"lr": 0.001}, + scheduler=ConstantLR, + scheduler_kwargs={"factor": 1, "total_iters": 0}, + eps=100, + ): + """ + :param AbstractProblem problem: The formulation of the problem. + :param torch.nn.Module model: The neural network model to use. + :param torch.nn.Module loss: The loss function used as minimizer, + default :class:`torch.nn.MSELoss`. + :param torch.nn.Module extra_features: The additional input + 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 torch.optim.LRScheduler scheduler: Learning + rate scheduler. + :param dict scheduler_kwargs: LR scheduler constructor keyword args. + :param int | float eps: The exponential decay parameter. Note that this + value is kept fixed during the training, but can be changed by means + of a callback, e.g. for annealing. + """ + super().__init__( + problem=problem, + model=model, + extra_features=extra_features, + loss=loss, + optimizer=optimizer, + optimizer_kwargs=optimizer_kwargs, + scheduler=scheduler, + scheduler_kwargs=scheduler_kwargs, + ) + + # checking consistency + check_consistency(eps, (int,float)) + self._eps = eps + if not isinstance(self.problem, TimeDependentProblem): + raise ValueError('Casual PINN works only for problems' + 'inheritig from TimeDependentProblem.') + + def loss_phys(self, samples, equation): + """ + Computes the physics loss for the Causal PINN solver based on given + samples and equation. + + :param LabelTensor samples: The samples to evaluate the physics loss. + :param EquationInterface equation: The governing equation + representing the physics. + :return: The physics loss calculated based on given + samples and equation. + :rtype: LabelTensor + """ + # split sequentially ordered time tensors into chunks + chunks, labels = self._split_tensor_into_chunks(samples) + # compute residuals - this correspond to ordered loss functions + # values for each time step. We apply `flatten` such that after + # concataning the residuals we obtain a tensor of shape #chunks + time_loss = [] + for chunk in chunks: + chunk.labels = labels + # classical PINN loss + residual = self.compute_residual(samples=chunk, equation=equation) + loss_val = self.loss( + torch.zeros_like(residual, requires_grad=True), residual + ) + time_loss.append(loss_val) + # store results + self.store_log(loss_value=float(sum(time_loss)/len(time_loss))) + # concatenate residuals + time_loss = torch.stack(time_loss) + # compute weights (without the gradient storing) + with torch.no_grad(): + weights = self._compute_weights(time_loss) + return (weights * time_loss).mean() + + @property + def eps(self): + """ + The exponential decay parameter. + """ + return self._eps + + @eps.setter + def eps(self, value): + """ + Setter method for the eps parameter. + + :param float value: The exponential decay parameter. + """ + check_consistency(value, float) + self._eps = value + + def _sort_label_tensor(self, tensor): + """ + Sorts the label tensor based on time variables. + + :param LabelTensor tensor: The label tensor to be sorted. + :return: The sorted label tensor based on time variables. + :rtype: LabelTensor + """ + # labels input tensors + labels = tensor.labels + # extract time tensor + time_tensor = tensor.extract(self.problem.temporal_domain.variables) + # sort the time tensors (this is very bad for GPU) + _, idx = torch.sort(time_tensor.tensor.flatten()) + tensor = tensor[idx] + tensor.labels = labels + return tensor + + def _split_tensor_into_chunks(self, tensor): + """ + Splits the label tensor into chunks based on time. + + :param LabelTensor tensor: The label tensor to be split. + :return: Tuple containing the chunks and the original labels. + :rtype: Tuple[List[LabelTensor], List] + """ + # labels input tensors + labels = tensor.labels + # labels input tensors + tensor = self._sort_label_tensor(tensor) + # extract time tensor + time_tensor = tensor.extract(self.problem.temporal_domain.variables) + # count unique tensors in time + _, idx_split = time_tensor.unique(return_counts=True) + # splitting + chunks = torch.split(tensor, tuple(idx_split)) + return chunks, labels # return chunks + + def _compute_weights(self, loss): + """ + Computes the weights for the physics loss based on the cumulative loss. + + :param LabelTensor loss: The physics loss values. + :return: The computed weights for the physics loss. + :rtype: LabelTensor + """ + # compute comulative loss and multiply by epsilos + cumulative_loss = self._eps * torch.cumsum(loss, dim=0) + # return the exponential of the weghited negative cumulative sum + return torch.exp(-cumulative_loss) diff --git a/pina/solvers/pinns/competitive_pinn.py b/pina/solvers/pinns/competitive_pinn.py new file mode 100644 index 000000000..6404c0bb4 --- /dev/null +++ b/pina/solvers/pinns/competitive_pinn.py @@ -0,0 +1,360 @@ +""" Module for CompetitivePINN """ + +import torch +import copy + +try: + from torch.optim.lr_scheduler import LRScheduler # torch >= 2.0 +except ImportError: + from torch.optim.lr_scheduler import ( + _LRScheduler as LRScheduler, + ) # torch < 2.0 + +from torch.optim.lr_scheduler import ConstantLR + +from .basepinn import PINNInterface +from pina.utils import check_consistency +from pina.problem import InverseProblem + + +class CompetitivePINN(PINNInterface): + r""" + Competitive Physics Informed Neural Network (PINN) solver class. + This class implements Competitive Physics Informed Neural + Network solvers, using a user specified ``model`` to solve a specific + ``problem``. It can be used for solving both forward and inverse problems. + + The Competitive Physics Informed Network aims to find + the solution :math:`\mathbf{u}:\Omega\rightarrow\mathbb{R}^m` + of the differential problem: + + .. math:: + + \begin{cases} + \mathcal{A}[\mathbf{u}](\mathbf{x})=0\quad,\mathbf{x}\in\Omega\\ + \mathcal{B}[\mathbf{u}](\mathbf{x})=0\quad, + \mathbf{x}\in\partial\Omega + \end{cases} + + with a minimization (on ``model`` parameters) maximation ( + on ``discriminator`` parameters) of the loss function + + .. math:: + \mathcal{L}_{\rm{problem}} = \frac{1}{N}\sum_{i=1}^N + \mathcal{L}(D(\mathbf{x}_i)\mathcal{A}[\mathbf{u}](\mathbf{x}_i))+ + \frac{1}{N}\sum_{i=1}^N + \mathcal{L}(D(\mathbf{x}_i)\mathcal{B}[\mathbf{u}](\mathbf{x}_i)) + + where :math:`D` is the discriminator network, which tries to find the points + where the network performs worst, and :math:`\mathcal{L}` is a specific loss + function, default Mean Square Error: + + .. math:: + \mathcal{L}(v) = \| v \|^2_2. + + .. seealso:: + + **Original reference**: Zeng, Qi, et al. + "Competitive physics informed networks." International Conference on + Learning Representations, ICLR 2022 + `OpenReview Preprint `_. + + .. warning:: + This solver does not currently support the possibility to pass + ``extra_feature``. + """ + + def __init__( + self, + problem, + model, + discriminator=None, + loss=torch.nn.MSELoss(), + optimizer_model=torch.optim.Adam, + optimizer_model_kwargs={"lr": 0.001}, + optimizer_discriminator=torch.optim.Adam, + optimizer_discriminator_kwargs={"lr": 0.001}, + scheduler_model=ConstantLR, + scheduler_model_kwargs={"factor": 1, "total_iters": 0}, + scheduler_discriminator=ConstantLR, + scheduler_discriminator_kwargs={"factor": 1, "total_iters": 0}, + ): + """ + :param AbstractProblem problem: The formualation of the problem. + :param torch.nn.Module model: The neural network model to use + for the model. + :param torch.nn.Module discriminator: The neural network model to use + for the discriminator. If ``None``, the discriminator network will + have the same architecture as the model network. + :param torch.nn.Module loss: The loss function used as minimizer, + default :class:`torch.nn.MSELoss`. + :param torch.optim.Optimizer optimizer_model: The neural + network optimizer to use for the model network + , default is `torch.optim.Adam`. + :param dict optimizer_model_kwargs: Optimizer constructor keyword + args. for the model. + :param torch.optim.Optimizer optimizer_discriminator: The neural + network optimizer to use for the discriminator network + , default is `torch.optim.Adam`. + :param dict optimizer_discriminator_kwargs: Optimizer constructor + keyword args. for the discriminator. + :param torch.optim.LRScheduler scheduler_model: Learning + rate scheduler for the model. + :param dict scheduler_model_kwargs: LR scheduler constructor + keyword args. + :param torch.optim.LRScheduler scheduler_discriminator: Learning + rate scheduler for the discriminator. + """ + if discriminator is None: + discriminator = copy.deepcopy(model) + + super().__init__( + models=[model, discriminator], + problem=problem, + optimizers=[optimizer_model, optimizer_discriminator], + optimizers_kwargs=[ + optimizer_model_kwargs, + optimizer_discriminator_kwargs, + ], + extra_features=None, # CompetitivePINN doesn't take extra features + loss=loss + ) + + # set automatic optimization for GANs + self.automatic_optimization = False + + # check consistency + check_consistency(scheduler_model, LRScheduler, subclass=True) + check_consistency(scheduler_model_kwargs, dict) + check_consistency(scheduler_discriminator, LRScheduler, subclass=True) + check_consistency(scheduler_discriminator_kwargs, dict) + + # assign schedulers + self._schedulers = [ + scheduler_model( + self.optimizers[0], **scheduler_model_kwargs + ), + scheduler_discriminator( + self.optimizers[1], **scheduler_discriminator_kwargs + ), + ] + + self._model = self.models[0] + self._discriminator = self.models[1] + + def forward(self, x): + r""" + Forward pass implementation for the PINN solver. It returns the function + evaluation :math:`\mathbf{u}(\mathbf{x})` at the control points + :math:`\mathbf{x}`. + + :param LabelTensor x: Input tensor for the PINN solver. It expects + a tensor :math:`N \times D`, where :math:`N` the number of points + in the mesh, :math:`D` the dimension of the problem, + :return: PINN solution evaluated at contro points. + :rtype: LabelTensor + """ + return self.neural_net(x) + + def loss_phys(self, samples, equation): + """ + Computes the physics loss for the Competitive PINN solver based on given + samples and equation. + + :param LabelTensor samples: The samples to evaluate the physics loss. + :param EquationInterface equation: The governing equation + representing the physics. + :return: The physics loss calculated based on given + samples and equation. + :rtype: LabelTensor + """ + # train one step of the model + with torch.no_grad(): + discriminator_bets = self.discriminator(samples) + loss_val = self._train_model(samples, equation, discriminator_bets) + self.store_log(loss_value=float(loss_val)) + # detaching samples from the computational graph to erase it and setting + # the gradient to true to create a new computational graph. + # In alternative set `retain_graph=True`. + samples = samples.detach() + samples.requires_grad = True + # train one step of discriminator + discriminator_bets = self.discriminator(samples) + self._train_discriminator(samples, equation, discriminator_bets) + return loss_val + + def loss_data(self, input_tensor, output_tensor): + """ + The data loss for the PINN solver. It computes the loss between the + network output against the true solution. + + :param LabelTensor input_tensor: The input to the neural networks. + :param LabelTensor output_tensor: The true solution to compare the + network solution. + :return: The computed data loss. + :rtype: torch.Tensor + """ + self.optimizer_model.zero_grad() + loss_val = super().loss_data( + input_tensor, output_tensor).as_subclass(torch.Tensor) + loss_val.backward() + self.optimizer_model.step() + return loss_val + + def configure_optimizers(self): + """ + Optimizer configuration for the Competitive PINN solver. + + :return: The optimizers and the schedulers + :rtype: tuple(list, list) + """ + # if the problem is an InverseProblem, add the unknown parameters + # to the parameters that the optimizer needs to optimize + if isinstance(self.problem, InverseProblem): + self.optimizer_model.add_param_group( + { + "params": [ + self._params[var] + for var in self.problem.unknown_variables + ] + } + ) + return self.optimizers, self._schedulers + + def on_train_batch_end(self,outputs, batch, batch_idx): + """ + This method is called at the end of each training batch, and ovverides + the PytorchLightining implementation for logging the checkpoints. + + :param torch.Tensor outputs: The output from the model for the + current batch. + :param tuple batch: The current batch of data. + :param int batch_idx: The index of the current batch. + :return: Whatever is returned by the parent + method ``on_train_batch_end``. + :rtype: Any + """ + # increase by one the counter of optimization to save loggers + self.trainer.fit_loop.epoch_loop.manual_optimization.optim_step_progress.total.completed += 1 + return super().on_train_batch_end(outputs, batch, batch_idx) + + def _train_discriminator(self, samples, equation, discriminator_bets): + """ + Trains the discriminator network of the Competitive PINN. + + :param LabelTensor samples: Input samples to evaluate the physics loss. + :param EquationInterface equation: The governing equation representing + the physics. + :param Tensor discriminator_bets: Predictions made by the discriminator + network. + """ + # manual optimization + self.optimizer_discriminator.zero_grad() + # compute residual, we detach because the weights of the generator + # model are fixed + residual = self.compute_residual(samples=samples, + equation=equation).detach() + # compute competitive residual, the minus is because we maximise + competitive_residual = residual * discriminator_bets + loss_val = - self.loss( + torch.zeros_like(competitive_residual, requires_grad=True), + competitive_residual + ).as_subclass(torch.Tensor) + # backprop + self.manual_backward(loss_val) + self.optimizer_discriminator.step() + return + + def _train_model(self, samples, equation, discriminator_bets): + """ + Trains the model network of the Competitive PINN. + + :param LabelTensor samples: Input samples to evaluate the physics loss. + :param EquationInterface equation: The governing equation representing + the physics. + :param Tensor discriminator_bets: Predictions made by the discriminator. + network. + :return: The computed data loss. + :rtype: torch.Tensor + """ + # manual optimization + self.optimizer_model.zero_grad() + # compute residual (detached for discriminator) and log + residual = self.compute_residual(samples=samples, equation=equation) + # store logging + with torch.no_grad(): + loss_residual = self.loss( + torch.zeros_like(residual), + residual + ) + # compute competitive residual, discriminator_bets are detached becase + # we optimize only the generator model + competitive_residual = residual * discriminator_bets.detach() + loss_val = self.loss( + torch.zeros_like(competitive_residual, requires_grad=True), + competitive_residual + ).as_subclass(torch.Tensor) + # backprop + self.manual_backward(loss_val) + self.optimizer_model.step() + return loss_residual + + @property + def neural_net(self): + """ + Returns the neural network model. + + :return: The neural network model. + :rtype: torch.nn.Module + """ + return self._model + + @property + def discriminator(self): + """ + Returns the discriminator model (if applicable). + + :return: The discriminator model. + :rtype: torch.nn.Module + """ + return self._discriminator + + @property + def optimizer_model(self): + """ + Returns the optimizer associated with the neural network model. + + :return: The optimizer for the neural network model. + :rtype: torch.optim.Optimizer + """ + return self.optimizers[0] + + @property + def optimizer_discriminator(self): + """ + Returns the optimizer associated with the discriminator (if applicable). + + :return: The optimizer for the discriminator. + :rtype: torch.optim.Optimizer + """ + return self.optimizers[1] + + @property + def scheduler_model(self): + """ + Returns the scheduler associated with the neural network model. + + :return: The scheduler for the neural network model. + :rtype: torch.optim.lr_scheduler._LRScheduler + """ + return self._schedulers[0] + + @property + def scheduler_discriminator(self): + """ + Returns the scheduler associated with the discriminator (if applicable). + + :return: The scheduler for the discriminator. + :rtype: torch.optim.lr_scheduler._LRScheduler + """ + return self._schedulers[1] \ No newline at end of file diff --git a/pina/solvers/pinns/gpinn.py b/pina/solvers/pinns/gpinn.py new file mode 100644 index 000000000..6eca1eac5 --- /dev/null +++ b/pina/solvers/pinns/gpinn.py @@ -0,0 +1,134 @@ +""" Module for GPINN """ + +import torch + + +from torch.optim.lr_scheduler import ConstantLR + +from .pinn import PINN +from pina.operators import grad +from pina.problem import SpatialProblem + + +class GPINN(PINN): + r""" + Gradient Physics Informed Neural Network (GPINN) solver class. + This class implements Gradient Physics Informed Neural + Network solvers, using a user specified ``model`` to solve a specific + ``problem``. It can be used for solving both forward and inverse problems. + + The Gradient Physics Informed Network aims to find + the solution :math:`\mathbf{u}:\Omega\rightarrow\mathbb{R}^m` + of the differential problem: + + .. math:: + + \begin{cases} + \mathcal{A}[\mathbf{u}](\mathbf{x})=0\quad,\mathbf{x}\in\Omega\\ + \mathcal{B}[\mathbf{u}](\mathbf{x})=0\quad, + \mathbf{x}\in\partial\Omega + \end{cases} + + minimizing the loss function + + .. math:: + \mathcal{L}_{\rm{problem}} =& \frac{1}{N}\sum_{i=1}^N + \mathcal{L}(\mathcal{A}[\mathbf{u}](\mathbf{x}_i)) + + \frac{1}{N}\sum_{i=1}^N + \mathcal{L}(\mathcal{B}[\mathbf{u}](\mathbf{x}_i)) + \\ + &\frac{1}{N}\sum_{i=1}^N + \nabla_{\mathbf{x}}\mathcal{L}(\mathcal{A}[\mathbf{u}](\mathbf{x}_i)) + + \frac{1}{N}\sum_{i=1}^N + \nabla_{\mathbf{x}}\mathcal{L}(\mathcal{B}[\mathbf{u}](\mathbf{x}_i)) + + + where :math:`\mathcal{L}` is a specific loss function, default Mean Square Error: + + .. math:: + \mathcal{L}(v) = \| v \|^2_2. + + .. seealso:: + + **Original reference**: Yu, Jeremy, et al. "Gradient-enhanced + physics-informed neural networks for forward and inverse + PDE problems." Computer Methods in Applied Mechanics + and Engineering 393 (2022): 114823. + DOI: `10.1016 `_. + + .. note:: + This class can only work for problems inheriting + from at least :class:`~pina.problem.spatial_problem.SpatialProblem` + class. + """ + + def __init__( + self, + problem, + model, + extra_features=None, + loss=torch.nn.MSELoss(), + optimizer=torch.optim.Adam, + optimizer_kwargs={"lr": 0.001}, + scheduler=ConstantLR, + scheduler_kwargs={"factor": 1, "total_iters": 0}, + ): + """ + :param AbstractProblem problem: The formulation of the problem. It must + inherit from at least + :class:`~pina.problem.spatial_problem.SpatialProblem` in order to + compute the gradient of the loss. + :param torch.nn.Module model: The neural network model to use. + :param torch.nn.Module loss: The loss function used as minimizer, + default :class:`torch.nn.MSELoss`. + :param torch.nn.Module extra_features: The additional input + 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 torch.optim.LRScheduler scheduler: Learning + rate scheduler. + :param dict scheduler_kwargs: LR scheduler constructor keyword args. + """ + super().__init__( + problem=problem, + model=model, + extra_features=extra_features, + loss=loss, + optimizer=optimizer, + optimizer_kwargs=optimizer_kwargs, + scheduler=scheduler, + scheduler_kwargs=scheduler_kwargs, + ) + if not isinstance(self.problem, SpatialProblem): + raise ValueError('Gradient PINN computes the gradient of the ' + 'PINN loss with respect to the spatial ' + 'coordinates, thus the PINA problem must be ' + 'a SpatialProblem.') + + + def loss_phys(self, samples, equation): + """ + Computes the physics loss for the GPINN solver based on given + samples and equation. + + :param LabelTensor samples: The samples to evaluate the physics loss. + :param EquationInterface equation: The governing equation + representing the physics. + :return: The physics loss calculated based on given + samples and equation. + :rtype: LabelTensor + """ + # classical PINN loss + residual = self.compute_residual(samples=samples, equation=equation) + loss_value = self.loss( + torch.zeros_like(residual, requires_grad=True), residual + ) + self.store_log(loss_value=float(loss_value)) + # gradient PINN loss + loss_value = loss_value.reshape(-1, 1) + loss_value.labels = ['__LOSS'] + loss_grad = grad(loss_value, samples, d=self.problem.spatial_variables) + g_loss_phys = self.loss( + torch.zeros_like(loss_grad, requires_grad=True), loss_grad + ) + return loss_value + g_loss_phys \ No newline at end of file diff --git a/pina/solvers/pinns/pinn.py b/pina/solvers/pinns/pinn.py new file mode 100644 index 000000000..318283a38 --- /dev/null +++ b/pina/solvers/pinns/pinn.py @@ -0,0 +1,170 @@ +""" Module for Physics Informed Neural Network. """ + +import torch + +try: + from torch.optim.lr_scheduler import LRScheduler # torch >= 2.0 +except ImportError: + from torch.optim.lr_scheduler import ( + _LRScheduler as LRScheduler, + ) # torch < 2.0 + +from torch.optim.lr_scheduler import ConstantLR + +from .basepinn import PINNInterface +from pina.utils import check_consistency +from pina.problem import InverseProblem + + +class PINN(PINNInterface): + r""" + Physics Informed Neural Network (PINN) solver class. + This class implements Physics Informed Neural + Network solvers, using a user specified ``model`` to solve a specific + ``problem``. It can be used for solving both forward and inverse problems. + + The Physics Informed Network aims to find + the solution :math:`\mathbf{u}:\Omega\rightarrow\mathbb{R}^m` + of the differential problem: + + .. math:: + + \begin{cases} + \mathcal{A}[\mathbf{u}](\mathbf{x})=0\quad,\mathbf{x}\in\Omega\\ + \mathcal{B}[\mathbf{u}](\mathbf{x})=0\quad, + \mathbf{x}\in\partial\Omega + \end{cases} + + minimizing the loss function + + .. math:: + \mathcal{L}_{\rm{problem}} = \frac{1}{N}\sum_{i=1}^N + \mathcal{L}(\mathcal{A}[\mathbf{u}](\mathbf{x}_i)) + + \frac{1}{N}\sum_{i=1}^N + \mathcal{L}(\mathcal{B}[\mathbf{u}](\mathbf{x}_i)) + + where :math:`\mathcal{L}` is a specific loss function, default Mean Square Error: + + .. math:: + \mathcal{L}(v) = \| v \|^2_2. + + .. seealso:: + + **Original reference**: Karniadakis, G. E., Kevrekidis, I. G., Lu, L., + Perdikaris, P., Wang, S., & Yang, L. (2021). + Physics-informed machine learning. Nature Reviews Physics, 3, 422-440. + DOI: `10.1038 `_. + """ + + def __init__( + self, + problem, + model, + extra_features=None, + loss=torch.nn.MSELoss(), + optimizer=torch.optim.Adam, + optimizer_kwargs={"lr": 0.001}, + scheduler=ConstantLR, + scheduler_kwargs={"factor": 1, "total_iters": 0}, + ): + """ + :param AbstractProblem problem: The formulation of the problem. + :param torch.nn.Module model: The neural network model to use. + :param torch.nn.Module loss: The loss function used as minimizer, + default :class:`torch.nn.MSELoss`. + :param torch.nn.Module extra_features: The additional input + 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 torch.optim.LRScheduler scheduler: Learning + rate scheduler. + :param dict scheduler_kwargs: LR scheduler constructor keyword args. + """ + super().__init__( + models=[model], + problem=problem, + optimizers=[optimizer], + optimizers_kwargs=[optimizer_kwargs], + extra_features=extra_features, + loss=loss + ) + + # check consistency + check_consistency(scheduler, LRScheduler, subclass=True) + check_consistency(scheduler_kwargs, dict) + + # assign variables + self._scheduler = scheduler(self.optimizers[0], **scheduler_kwargs) + self._neural_net = self.models[0] + + def forward(self, x): + r""" + Forward pass implementation for the PINN solver. It returns the function + evaluation :math:`\mathbf{u}(\mathbf{x})` at the control points + :math:`\mathbf{x}`. + + :param LabelTensor x: Input tensor for the PINN solver. It expects + a tensor :math:`N \times D`, where :math:`N` the number of points + in the mesh, :math:`D` the dimension of the problem, + :return: PINN solution evaluated at contro points. + :rtype: LabelTensor + """ + return self.neural_net(x) + + def loss_phys(self, samples, equation): + """ + Computes the physics loss for the PINN solver based on given + samples and equation. + + :param LabelTensor samples: The samples to evaluate the physics loss. + :param EquationInterface equation: The governing equation + representing the physics. + :return: The physics loss calculated based on given + samples and equation. + :rtype: LabelTensor + """ + residual = self.compute_residual(samples=samples, equation=equation) + loss_value = self.loss( + torch.zeros_like(residual, requires_grad=True), residual + ) + self.store_log(loss_value=float(loss_value)) + return loss_value + + + def configure_optimizers(self): + """ + Optimizer configuration for the PINN + solver. + + :return: The optimizers and the schedulers + :rtype: tuple(list, list) + """ + # if the problem is an InverseProblem, add the unknown parameters + # to the parameters that the optimizer needs to optimize + if isinstance(self.problem, InverseProblem): + self.optimizers[0].add_param_group( + { + "params": [ + self._params[var] + for var in self.problem.unknown_variables + ] + } + ) + return self.optimizers, [self.scheduler] + + + @property + def scheduler(self): + """ + Scheduler for the PINN training. + """ + return self._scheduler + + + @property + def neural_net(self): + """ + Neural network for the PINN training. + """ + return self._neural_net \ No newline at end of file diff --git a/pina/solvers/pinns/sapinn.py b/pina/solvers/pinns/sapinn.py new file mode 100644 index 000000000..8de2d14c1 --- /dev/null +++ b/pina/solvers/pinns/sapinn.py @@ -0,0 +1,494 @@ +import torch +from copy import deepcopy + +try: + from torch.optim.lr_scheduler import LRScheduler # torch >= 2.0 +except ImportError: + from torch.optim.lr_scheduler import ( + _LRScheduler as LRScheduler, + ) # torch < 2.0 + +from .basepinn import PINNInterface +from pina.utils import check_consistency +from pina.problem import InverseProblem + +from torch.optim.lr_scheduler import ConstantLR + +class Weights(torch.nn.Module): + """ + This class aims to implements the mask model for + self adaptive weights of the Self-Adaptive + PINN solver. + """ + + def __init__(self, func): + """ + :param torch.nn.Module func: the mask module of SAPINN + """ + super().__init__() + check_consistency(func, torch.nn.Module) + self.sa_weights = torch.nn.Parameter( + torch.Tensor() + ) + self.func = func + + def forward(self): + """ + Forward pass implementation for the mask module. + It returns the function on the weights + evaluation. + + :return: evaluation of self adaptive weights through the mask. + :rtype: torch.Tensor + """ + return self.func(self.sa_weights) + +class SAPINN(PINNInterface): + r""" + Self Adaptive Physics Informed Neural Network (SAPINN) solver class. + This class implements Self-Adaptive Physics Informed Neural + Network solvers, using a user specified ``model`` to solve a specific + ``problem``. It can be used for solving both forward and inverse problems. + + The Self Adapive Physics Informed Neural Network aims to find + the solution :math:`\mathbf{u}:\Omega\rightarrow\mathbb{R}^m` + of the differential problem: + + .. math:: + + \begin{cases} + \mathcal{A}[\mathbf{u}](\mathbf{x})=0\quad,\mathbf{x}\in\Omega\\ + \mathcal{B}[\mathbf{u}](\mathbf{x})=0\quad, + \mathbf{x}\in\partial\Omega + \end{cases} + + integrating the pointwise loss evaluation through a mask :math:`m` and + self adaptive weights that permit to focus the loss function on + specific training samples. + The loss function to solve the problem is + + .. math:: + + \mathcal{L}_{\rm{problem}} = \frac{1}{N} \sum_{i=1}^{N_\Omega} m + \left( \lambda_{\Omega}^{i} \right) \mathcal{L} \left( \mathcal{A} + [\mathbf{u}](\mathbf{x}) \right) + \frac{1}{N} + \sum_{i=1}^{N_{\partial\Omega}} + m \left( \lambda_{\partial\Omega}^{i} \right) \mathcal{L} + \left( \mathcal{B}[\mathbf{u}](\mathbf{x}) + \right), + + + denoting the self adaptive weights as + :math:`\lambda_{\Omega}^1, \dots, \lambda_{\Omega}^{N_\Omega}` and + :math:`\lambda_{\partial \Omega}^1, \dots, + \lambda_{\Omega}^{N_\partial \Omega}` + for :math:`\Omega` and :math:`\partial \Omega`, respectively. + + Self Adaptive Physics Informed Neural Network identifies the solution + and appropriate self adaptive weights by solving the following problem + + .. math:: + + \min_{w} \max_{\lambda_{\Omega}^k, \lambda_{\partial \Omega}^s} + \mathcal{L} , + + where :math:`w` denotes the network parameters, and + :math:`\mathcal{L}` is a specific loss + function, default Mean Square Error: + + .. math:: + \mathcal{L}(v) = \| v \|^2_2. + + .. seealso:: + **Original reference**: McClenny, Levi D., and Ulisses M. Braga-Neto. + "Self-adaptive physics-informed neural networks." + Journal of Computational Physics 474 (2023): 111722. + DOI: `10.1016/ + j.jcp.2022.111722 `_. + """ + + def __init__( + self, + problem, + model, + weights_function=torch.nn.Sigmoid(), + extra_features=None, + loss=torch.nn.MSELoss(), + optimizer_model=torch.optim.Adam, + optimizer_model_kwargs={"lr" : 0.001}, + optimizer_weights=torch.optim.Adam, + optimizer_weights_kwargs={"lr" : 0.001}, + scheduler_model=ConstantLR, + scheduler_model_kwargs={"factor" : 1, "total_iters" : 0}, + scheduler_weights=ConstantLR, + scheduler_weights_kwargs={"factor" : 1, "total_iters" : 0} + ): + """ + :param AbstractProblem problem: The formualation of the problem. + :param torch.nn.Module model: The neural network model to use + for the model. + :param torch.nn.Module weights_function: The neural network model + related to the mask of SAPINN. + default :obj:`~torch.nn.Sigmoid`. + :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. + :param torch.nn.Module loss: The loss function used as minimizer, + default :class:`torch.nn.MSELoss`. + :param torch.optim.Optimizer optimizer_model: The neural + network optimizer to use for the model network + , default is `torch.optim.Adam`. + :param dict optimizer_model_kwargs: Optimizer constructor keyword + args. for the model. + :param torch.optim.Optimizer optimizer_weights: The neural + network optimizer to use for mask model model, + default is `torch.optim.Adam`. + :param dict optimizer_weights_kwargs: Optimizer constructor + keyword args. for the mask module. + :param torch.optim.LRScheduler scheduler_model: Learning + rate scheduler for the model. + :param dict scheduler_model_kwargs: LR scheduler constructor + keyword args. + :param torch.optim.LRScheduler scheduler_weights: Learning + rate scheduler for the mask model. + :param dict scheduler_model_kwargs: LR scheduler constructor + keyword args. + """ + + # check consistency weitghs_function + check_consistency(weights_function, torch.nn.Module) + + # create models for weights + weights_dict = {} + for condition_name in problem.conditions: + weights_dict[condition_name] = Weights(weights_function) + weights_dict = torch.nn.ModuleDict(weights_dict) + + + super().__init__( + models=[model, weights_dict], + problem=problem, + optimizers=[optimizer_model, optimizer_weights], + optimizers_kwargs=[ + optimizer_model_kwargs, + optimizer_weights_kwargs + ], + extra_features=extra_features, + loss=loss + ) + + # set automatic optimization + self.automatic_optimization = False + + # check consistency + check_consistency(scheduler_model, LRScheduler, subclass=True) + check_consistency(scheduler_model_kwargs, dict) + check_consistency(scheduler_weights, LRScheduler, subclass=True) + check_consistency(scheduler_weights_kwargs, dict) + + # assign schedulers + self._schedulers = [ + scheduler_model( + self.optimizers[0], **scheduler_model_kwargs + ), + scheduler_weights( + self.optimizers[1], **scheduler_weights_kwargs + ), + ] + + self._model = self.models[0] + self._weights = self.models[1] + + self._vectorial_loss = deepcopy(loss) + self._vectorial_loss.reduction = "none" + + def forward(self, x): + """ + Forward pass implementation for the PINN + solver. It returns the function + evaluation :math:`\mathbf{u}(\mathbf{x})` at the control points + :math:`\mathbf{x}`. + + :param LabelTensor x: Input tensor for the SAPINN solver. It expects + a tensor :math:`N \\times D`, where :math:`N` the number of points + in the mesh, :math:`D` the dimension of the problem, + :return: PINN solution. + :rtype: LabelTensor + """ + return self.neural_net(x) + + def loss_phys(self, samples, equation): + """ + Computes the physics loss for the SAPINN solver based on given + samples and equation. + + :param LabelTensor samples: The samples to evaluate the physics loss. + :param EquationInterface equation: The governing equation + representing the physics. + :return: The physics loss calculated based on given + samples and equation. + :rtype: torch.Tensor + """ + # train weights + self.optimizer_weights.zero_grad() + weighted_loss, _ = self._loss_phys(samples, equation) + loss_value = - weighted_loss.as_subclass(torch.Tensor) + self.manual_backward(loss_value) + self.optimizer_weights.step() + + # detaching samples from the computational graph to erase it and setting + # the gradient to true to create a new computational graph. + # In alternative set `retain_graph=True`. + samples = samples.detach() + samples.requires_grad = True + + # train model + self.optimizer_model.zero_grad() + weighted_loss, loss = self._loss_phys(samples, equation) + loss_value = weighted_loss.as_subclass(torch.Tensor) + self.manual_backward(loss_value) + self.optimizer_model.step() + + # store loss without weights + self.store_log(loss_value=float(loss)) + return loss_value + + def loss_data(self, input_tensor, output_tensor): + """ + Computes the data loss for the SAPINN solver based on input and + output. It computes the loss between the + network output against the true solution. + + :param LabelTensor input_tensor: The input to the neural networks. + :param LabelTensor output_tensor: The true solution to compare the + network solution. + :return: The computed data loss. + :rtype: torch.Tensor + """ + # train weights + self.optimizer_weights.zero_grad() + weighted_loss, _ = self._loss_data(input_tensor, output_tensor) + loss_value = - weighted_loss.as_subclass(torch.Tensor) + self.manual_backward(loss_value) + self.optimizer_weights.step() + + # detaching samples from the computational graph to erase it and setting + # the gradient to true to create a new computational graph. + # In alternative set `retain_graph=True`. + input_tensor = input_tensor.detach() + input_tensor.requires_grad = True + + # train model + self.optimizer_model.zero_grad() + weighted_loss, loss = self._loss_data(input_tensor, output_tensor) + loss_value = weighted_loss.as_subclass(torch.Tensor) + self.manual_backward(loss_value) + self.optimizer_model.step() + + # store loss without weights + self.store_log(loss_value=float(loss)) + return loss_value + + def configure_optimizers(self): + """ + Optimizer configuration for the SAPINN + solver. + + :return: The optimizers and the schedulers + :rtype: tuple(list, list) + """ + # if the problem is an InverseProblem, add the unknown parameters + # to the parameters that the optimizer needs to optimize + if isinstance(self.problem, InverseProblem): + self.optimizers[0].add_param_group( + { + "params": [ + self._params[var] + for var in self.problem.unknown_variables + ] + } + ) + return self.optimizers, self._schedulers + + def on_train_batch_end(self,outputs, batch, batch_idx): + """ + This method is called at the end of each training batch, and ovverides + the PytorchLightining implementation for logging the checkpoints. + + :param torch.Tensor outputs: The output from the model for the + current batch. + :param tuple batch: The current batch of data. + :param int batch_idx: The index of the current batch. + :return: Whatever is returned by the parent + method ``on_train_batch_end``. + :rtype: Any + """ + # increase by one the counter of optimization to save loggers + self.trainer.fit_loop.epoch_loop.manual_optimization.optim_step_progress.total.completed += 1 + return super().on_train_batch_end(outputs, batch, batch_idx) + + def on_train_start(self): + """ + This method is called at the start of the training for setting + the self adaptive weights as parameters of the mask model. + + :return: Whatever is returned by the parent + method ``on_train_start``. + :rtype: Any + """ + device = torch.device( + self.trainer._accelerator_connector._accelerator_flag + ) + for condition_name, tensor in self.problem.input_pts.items(): + self.weights_dict.torchmodel[condition_name].sa_weights.data = torch.rand( + (tensor.shape[0], 1), + device = device + ) + return super().on_train_start() + + def on_load_checkpoint(self, checkpoint): + """ + Overriding the Pytorch Lightning ``on_load_checkpoint`` to handle + checkpoints for Self Adaptive Weights. This method should not be + overridden if not intentionally. + + :param dict checkpoint: Pytorch Lightning checkpoint dict. + """ + for condition_name, tensor in self.problem.input_pts.items(): + self.weights_dict.torchmodel[condition_name].sa_weights.data = torch.rand( + (tensor.shape[0], 1) + ) + return super().on_load_checkpoint(checkpoint) + + def _loss_phys(self, samples, equation): + """ + Elaboration of the physical loss for the SAPINN solver. + + :param LabelTensor samples: Input samples to evaluate the physics loss. + :param EquationInterface equation: the governing equation representing + the physics. + + :return: tuple with weighted and not weighted scalar loss + :rtype: List[LabelTensor, LabelTensor] + """ + residual = self.compute_residual(samples, equation) + return self._compute_loss(residual) + + def _loss_data(self, input_tensor, output_tensor): + """ + Elaboration of the loss related to data for the SAPINN solver. + + :param LabelTensor input_tensor: The input to the neural networks. + :param LabelTensor output_tensor: The true solution to compare the + network solution. + + :return: tuple with weighted and not weighted scalar loss + :rtype: List[LabelTensor, LabelTensor] + """ + residual = self.forward(input_tensor) - output_tensor + return self._compute_loss(residual) + + def _compute_loss(self, residual): + """ + Elaboration of the pointwise loss through the mask model and the + self adaptive weights + + :param LabelTensor residual: the matrix of residuals that have to + be weighted + + :return: tuple with weighted and not weighted loss + :rtype List[LabelTensor, LabelTensor] + """ + weights = self.weights_dict.torchmodel[ + self.current_condition_name].forward() + loss_value = self._vectorial_loss(torch.zeros_like( + residual, requires_grad=True), residual) + return ( + self._vect_to_scalar(weights * loss_value), + self._vect_to_scalar(loss_value) + ) + + def _vect_to_scalar(self, loss_value): + """ + Elaboration of the pointwise loss through the mask model and the + self adaptive weights + + :param LabelTensor loss_value: the matrix of pointwise loss + + :return: the scalar loss + :rtype LabelTensor + """ + if self.loss.reduction == "mean": + ret = torch.mean(loss_value) + elif self.loss.reduction == "sum": + ret = torch.sum(loss_value) + else: + raise RuntimeError(f"Invalid reduction, got {self.loss.reduction} " + "but expected mean or sum.") + return ret + + + @property + def neural_net(self): + """ + Returns the neural network model. + + :return: The neural network model. + :rtype: torch.nn.Module + """ + return self.models[0] + + @property + def weights_dict(self): + """ + Return the mask models associate to the application of + the mask to the self adaptive weights for each loss that + compones the global loss of the problem. + + :return: The ModuleDict for mask models. + :rtype: torch.nn.ModuleDict + """ + return self.models[1] + + @property + def scheduler_model(self): + """ + Returns the scheduler associated with the neural network model. + + :return: The scheduler for the neural network model. + :rtype: torch.optim.lr_scheduler._LRScheduler + """ + return self._scheduler[0] + + @property + def scheduler_weights(self): + """ + Returns the scheduler associated with the mask model (if applicable). + + :return: The scheduler for the mask model. + :rtype: torch.optim.lr_scheduler._LRScheduler + """ + return self._scheduler[1] + + @property + def optimizer_model(self): + """ + Returns the optimizer associated with the neural network model. + + :return: The optimizer for the neural network model. + :rtype: torch.optim.Optimizer + """ + return self.optimizers[0] + + @property + def optimizer_weights(self): + """ + Returns the optimizer associated with the mask model (if applicable). + + :return: The optimizer for the mask model. + :rtype: torch.optim.Optimizer + """ + return self.optimizers[1] \ No newline at end of file diff --git a/pina/solvers/rom.py b/pina/solvers/rom.py new file mode 100644 index 000000000..733d76f42 --- /dev/null +++ b/pina/solvers/rom.py @@ -0,0 +1,190 @@ +""" Module for ReducedOrderModelSolver """ + +import torch + +from pina.solvers import SupervisedSolver + +class ReducedOrderModelSolver(SupervisedSolver): + r""" + ReducedOrderModelSolver solver class. This class implements a + Reduced Order Model solver, using user specified ``reduction_network`` and + ``interpolation_network`` to solve a specific ``problem``. + + The Reduced Order Model approach aims to find + the solution :math:`\mathbf{u}:\Omega\rightarrow\mathbb{R}^m` + of the differential problem: + + .. math:: + + \begin{cases} + \mathcal{A}[\mathbf{u}(\mu)](\mathbf{x})=0\quad,\mathbf{x}\in\Omega\\ + \mathcal{B}[\mathbf{u}(\mu)](\mathbf{x})=0\quad, + \mathbf{x}\in\partial\Omega + \end{cases} + + This is done by using two neural networks. The ``reduction_network``, which + contains an encoder :math:`\mathcal{E}_{\rm{net}}`, a decoder + :math:`\mathcal{D}_{\rm{net}}`; and an ``interpolation_network`` + :math:`\mathcal{I}_{\rm{net}}`. The input is assumed to be discretised in + the spatial dimensions. + + The following loss function is minimized during training + + .. math:: + \mathcal{L}_{\rm{problem}} = \frac{1}{N}\sum_{i=1}^N + \mathcal{L}(\mathcal{E}_{\rm{net}}[\mathbf{u}(\mu_i)] - + \mathcal{I}_{\rm{net}}[\mu_i]) + + \mathcal{L}( + \mathcal{D}_{\rm{net}}[\mathcal{E}_{\rm{net}}[\mathbf{u}(\mu_i)]] - + \mathbf{u}(\mu_i)) + + where :math:`\mathcal{L}` is a specific loss function, default Mean Square Error: + + .. math:: + \mathcal{L}(v) = \| v \|^2_2. + + + .. seealso:: + + **Original reference**: Hesthaven, Jan S., and Stefano Ubbiali. + "Non-intrusive reduced order modeling of nonlinear problems + using neural networks." Journal of Computational + Physics 363 (2018): 55-78. + DOI `10.1016/j.jcp.2018.02.037 + `_. + + .. note:: + The specified ``reduction_network`` must contain two methods, + namely ``encode`` for input encoding and ``decode`` for decoding the + former result. The ``interpolation_network`` network ``forward`` output + represents the interpolation of the latent space obtain with + ``reduction_network.encode``. + + .. note:: + This solver uses the end-to-end training strategy, i.e. the + ``reduction_network`` and ``interpolation_network`` are trained + simultaneously. For reference on this trainig strategy look at: + Pichi, Federico, Beatriz Moya, and Jan S. Hesthaven. + "A graph convolutional autoencoder approach to model order reduction + for parametrized PDEs." Journal of + Computational Physics 501 (2024): 112762. + DOI + `10.1016/j.jcp.2024.112762 `_. + + .. warning:: + This solver works only for data-driven model. Hence in the ``problem`` + definition the codition must only contain ``input_points`` + (e.g. coefficient parameters, time parameters), and ``output_points``. + + .. warning:: + This solver does not currently support the possibility to pass + ``extra_feature``. + """ + + def __init__( + self, + problem, + reduction_network, + interpolation_network, + loss=torch.nn.MSELoss(), + optimizer=torch.optim.Adam, + optimizer_kwargs={"lr": 0.001}, + scheduler=torch.optim.lr_scheduler.ConstantLR, + scheduler_kwargs={"factor": 1, "total_iters": 0}, + ): + """ + :param AbstractProblem problem: The formualation of the problem. + :param torch.nn.Module reduction_network: The reduction network used + for reducing the input space. It must contain two methods, + namely ``encode`` for input encoding and ``decode`` for decoding the + former result. + :param torch.nn.Module interpolation_network: The interpolation network + for interpolating the control parameters to latent space obtain by + the ``reduction_network`` encoding. + :param torch.nn.Module loss: The loss function used as minimizer, + default :class:`torch.nn.MSELoss`. + :param torch.nn.Module extra_features: The additional input + 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. + """ + model = torch.nn.ModuleDict({ + 'reduction_network' : reduction_network, + 'interpolation_network' : interpolation_network}) + + super().__init__( + model=model, + problem=problem, + loss=loss, + optimizer=optimizer, + optimizer_kwargs=optimizer_kwargs, + scheduler=scheduler, + scheduler_kwargs=scheduler_kwargs + ) + + # assert reduction object contains encode/ decode + if not hasattr(self.neural_net['reduction_network'], 'encode'): + raise SyntaxError('reduction_network must have encode method. ' + 'The encode method should return a lower ' + 'dimensional representation of the input.') + if not hasattr(self.neural_net['reduction_network'], 'decode'): + raise SyntaxError('reduction_network must have decode method. ' + 'The decode method should return a high ' + 'dimensional representation of the encoding.') + + def forward(self, x): + """ + Forward pass implementation for the solver. It finds the encoder + representation by calling ``interpolation_network.forward`` on the + input, and maps this representation to output space by calling + ``reduction_network.decode``. + + :param torch.Tensor x: Input tensor. + :return: Solver solution. + :rtype: torch.Tensor + """ + reduction_network = self.neural_net['reduction_network'] + interpolation_network = self.neural_net['interpolation_network'] + return reduction_network.decode(interpolation_network(x)) + + def loss_data(self, input_pts, output_pts): + """ + The data loss for the ReducedOrderModelSolver solver. + It computes the loss between + 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 + network solution. + :return: The residual loss averaged on the input coordinates + :rtype: torch.Tensor + """ + # extract networks + reduction_network = self.neural_net['reduction_network'] + interpolation_network = self.neural_net['interpolation_network'] + # encoded representations loss + encode_repr_inter_net = interpolation_network(input_pts) + encode_repr_reduction_network = reduction_network.encode(output_pts) + loss_encode = self.loss(encode_repr_inter_net, + encode_repr_reduction_network) + # reconstruction loss + loss_reconstruction = self.loss( + reduction_network.decode(encode_repr_reduction_network), + output_pts) + + return loss_encode + loss_reconstruction + + @property + def neural_net(self): + """ + Neural network for training. It returns a :obj:`~torch.nn.ModuleDict` + containing the ``reduction_network`` and ``interpolation_network``. + """ + return self._neural_net.torchmodel diff --git a/pina/solvers/solver.py b/pina/solvers/solver.py index 324a023dd..729a9d485 100644 --- a/pina/solvers/solver.py +++ b/pina/solvers/solver.py @@ -6,6 +6,7 @@ from ..utils import check_consistency from ..problem import AbstractProblem import torch +import sys class SolverInterface(pytorch_lightning.LightningModule, metaclass=ABCMeta): @@ -141,6 +142,20 @@ 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): diff --git a/pina/solvers/supervised.py b/pina/solvers/supervised.py index c6a8a35bf..28a634b0b 100644 --- a/pina/solvers/supervised.py +++ b/pina/solvers/supervised.py @@ -1,7 +1,6 @@ """ Module for SupervisedSolver """ import torch -import sys try: from torch.optim.lr_scheduler import LRScheduler # torch >= 2.0 @@ -20,9 +19,32 @@ class SupervisedSolver(SolverInterface): - """ + r""" SupervisedSolver solver class. This class implements a SupervisedSolver, using a user specified ``model`` to solve a specific ``problem``. + + The Supervised Solver class aims to find + a map between the input :math:`\mathbf{s}:\Omega\rightarrow\mathbb{R}^m` + and the output :math:`\mathbf{u}:\Omega\rightarrow\mathbb{R}^m`. The input + can be discretised in space (as in :obj:`~pina.solvers.rom.ROMe2eSolver`), + or not (e.g. when training Neural Operators). + + Given a model :math:`\mathcal{M}`, the following loss function is + minimized during training: + + .. math:: + \mathcal{L}_{\rm{problem}} = \frac{1}{N}\sum_{i=1}^N + \mathcal{L}(\mathbf{u}_i - \mathcal{M}(\mathbf{v}_i)) + + where :math:`\mathcal{L}` is a specific loss function, + default Mean Square Error: + + .. math:: + \mathcal{L}(v) = \| v \|^2_2. + + In this context :math:`\mathbf{u}_i` and :math:`\mathbf{v}_i` means that + we are seeking to approximate multiple (discretised) functions given + multiple (discretised) input functions. """ def __init__( @@ -96,18 +118,12 @@ def training_step(self, batch, batch_idx): :return: The sum of the loss functions. :rtype: LabelTensor """ - - dataloader = self.trainer.train_dataloader + condition_idx = batch["condition"] for condition_id in range(condition_idx.min(), condition_idx.max() + 1): - if sys.version_info >= (3, 8): - condition_name = dataloader.condition_names[condition_id] - else: - condition_name = dataloader.loaders.condition_names[ - condition_id - ] + condition_name = self._dataloader.condition_names[condition_id] condition = self.problem.conditions[condition_name] pts = batch["pts"] out = batch["output"] @@ -118,14 +134,14 @@ def training_step(self, batch, batch_idx): # for data driven mode if not hasattr(condition, "output_points"): raise NotImplementedError( - "Supervised solver 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] loss = ( - self.loss(self.forward(input_pts), output_pts) + self.loss_data(input_pts=input_pts, output_pts=output_pts) * condition.data_weight ) loss = loss.as_subclass(torch.Tensor) @@ -133,6 +149,20 @@ def training_step(self, batch, batch_idx): self.log("mean_loss", float(loss), prog_bar=True, logger=True) return loss + def loss_data(self, input_pts, output_pts): + """ + The data loss for the Supervised solver. It computes the loss between + 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 + network solution. + :return: The residual loss averaged on the input coordinates + :rtype: torch.Tensor + """ + return self.loss(self.forward(input_pts), output_pts) + @property def scheduler(self): """ diff --git a/pina/trainer.py b/pina/trainer.py index 0acecaaa9..90779a6e9 100644 --- a/pina/trainer.py +++ b/pina/trainer.py @@ -1,5 +1,6 @@ """ Trainer module. """ +import torch import pytorch_lightning from .utils import check_consistency from .dataset import SamplePointDataset, SamplePointLoader, DataPointDataset @@ -63,6 +64,12 @@ def _create_or_update_loader(self): self._loader = SamplePointLoader( dataset_phys, dataset_data, batch_size=self.batch_size, shuffle=True ) + pb = self._model.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)) + + def train(self, **kwargs): """ diff --git a/tests/test_solvers/test_causalpinn.py b/tests/test_solvers/test_causalpinn.py new file mode 100644 index 000000000..58518ae41 --- /dev/null +++ b/tests/test_solvers/test_causalpinn.py @@ -0,0 +1,266 @@ +import torch +import pytest + +from pina.problem import TimeDependentProblem, InverseProblem, SpatialProblem +from pina.operators import grad +from pina.geometry import CartesianDomain +from pina import Condition, LabelTensor +from pina.solvers import CausalPINN +from pina.trainer import Trainer +from pina.model import FeedForward +from pina.equation.equation import Equation +from pina.equation.equation_factory import FixedValue +from pina.loss import LpLoss + + + +class FooProblem(SpatialProblem): + ''' + Foo problem formulation. + ''' + output_variables = ['u'] + conditions = {} + spatial_domain = None + + +class InverseDiffusionReactionSystem(TimeDependentProblem, SpatialProblem, InverseProblem): + + def diffusionreaction(input_, output_, params_): + x = input_.extract('x') + t = input_.extract('t') + u_t = grad(output_, input_, d='t') + u_x = grad(output_, input_, d='x') + u_xx = grad(u_x, input_, d='x') + r = torch.exp(-t) * (1.5 * torch.sin(2*x) + (8/3)*torch.sin(3*x) + + (15/4)*torch.sin(4*x) + (63/8)*torch.sin(8*x)) + return u_t - params_['mu']*u_xx - r + + def _solution(self, pts): + t = pts.extract('t') + x = pts.extract('x') + return torch.exp(-t) * (torch.sin(x) + (1/2)*torch.sin(2*x) + + (1/3)*torch.sin(3*x) + (1/4)*torch.sin(4*x) + + (1/8)*torch.sin(8*x)) + + # assign output/ spatial and temporal variables + output_variables = ['u'] + spatial_domain = CartesianDomain({'x': [-torch.pi, torch.pi]}) + temporal_domain = CartesianDomain({'t': [0, 1]}) + unknown_parameter_domain = CartesianDomain({'mu': [-1, 1]}) + + # problem condition statement + conditions = { + 'D': Condition(location=CartesianDomain({'x': [-torch.pi, torch.pi], + 't': [0, 1]}), + equation=Equation(diffusionreaction)), + 'data' : Condition(input_points=LabelTensor(torch.tensor([[0., 0.]]), ['x', 't']), + output_points=LabelTensor(torch.tensor([[0.]]), ['u'])), + } + +class DiffusionReactionSystem(TimeDependentProblem, SpatialProblem): + + def diffusionreaction(input_, output_): + x = input_.extract('x') + t = input_.extract('t') + u_t = grad(output_, input_, d='t') + u_x = grad(output_, input_, d='x') + u_xx = grad(u_x, input_, d='x') + r = torch.exp(-t) * (1.5 * torch.sin(2*x) + (8/3)*torch.sin(3*x) + + (15/4)*torch.sin(4*x) + (63/8)*torch.sin(8*x)) + return u_t - u_xx - r + + def _solution(self, pts): + t = pts.extract('t') + x = pts.extract('x') + return torch.exp(-t) * (torch.sin(x) + (1/2)*torch.sin(2*x) + + (1/3)*torch.sin(3*x) + (1/4)*torch.sin(4*x) + + (1/8)*torch.sin(8*x)) + + # assign output/ spatial and temporal variables + output_variables = ['u'] + spatial_domain = CartesianDomain({'x': [-torch.pi, torch.pi]}) + temporal_domain = CartesianDomain({'t': [0, 1]}) + + # problem condition statement + conditions = { + 'D': Condition(location=CartesianDomain({'x': [-torch.pi, torch.pi], + 't': [0, 1]}), + equation=Equation(diffusionreaction)), + } + +class myFeature(torch.nn.Module): + """ + Feature: sin(x) + """ + + def __init__(self): + super(myFeature, self).__init__() + + def forward(self, x): + t = (torch.sin(x.extract(['x']) * torch.pi)) + return LabelTensor(t, ['sin(x)']) + + +# make the problem +problem = DiffusionReactionSystem() +model = FeedForward(len(problem.input_variables), + len(problem.output_variables)) +model_extra_feats = FeedForward( + len(problem.input_variables) + 1, + len(problem.output_variables)) +extra_feats = [myFeature()] + + +def test_constructor(): + CausalPINN(problem=problem, model=model, extra_features=None) + + with pytest.raises(ValueError): + CausalPINN(FooProblem(), model=model, extra_features=None) + + +def test_constructor_extra_feats(): + model_extra_feats = FeedForward( + len(problem.input_variables) + 1, + len(problem.output_variables)) + CausalPINN(problem=problem, + model=model_extra_feats, + extra_features=extra_feats) + + +def test_train_cpu(): + problem = DiffusionReactionSystem() + boundaries = ['D'] + n = 10 + problem.discretise_domain(n, 'grid', locations=boundaries) + pinn = CausalPINN(problem = problem, + model=model, extra_features=None, loss=LpLoss()) + trainer = Trainer(solver=pinn, max_epochs=1, + accelerator='cpu', batch_size=20) + trainer.train() + + +def test_train_restore(): + tmpdir = "tests/tmp_restore" + problem = DiffusionReactionSystem() + boundaries = ['D'] + n = 10 + problem.discretise_domain(n, 'grid', locations=boundaries) + pinn = CausalPINN(problem=problem, + model=model, + extra_features=None, + loss=LpLoss()) + trainer = Trainer(solver=pinn, + max_epochs=5, + accelerator='cpu', + default_root_dir=tmpdir) + trainer.train() + ntrainer = Trainer(solver=pinn, 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" + problem = DiffusionReactionSystem() + boundaries = ['D'] + n = 10 + problem.discretise_domain(n, 'grid', locations=boundaries) + pinn = CausalPINN(problem=problem, + model=model, + extra_features=None, + loss=LpLoss()) + trainer = Trainer(solver=pinn, + max_epochs=15, + accelerator='cpu', + default_root_dir=tmpdir) + trainer.train() + new_pinn = CausalPINN.load_from_checkpoint( + f'{tmpdir}/lightning_logs/version_0/checkpoints/epoch=14-step=15.ckpt', + problem = problem, model=model) + test_pts = CartesianDomain({'x': [0, 1], 't': [0, 1]}).sample(10) + assert new_pinn.forward(test_pts).extract(['u']).shape == (10, 1) + assert new_pinn.forward(test_pts).extract( + ['u']).shape == pinn.forward(test_pts).extract(['u']).shape + torch.testing.assert_close( + new_pinn.forward(test_pts).extract(['u']), + pinn.forward(test_pts).extract(['u'])) + import shutil + shutil.rmtree(tmpdir) + +def test_train_inverse_problem_cpu(): + problem = InverseDiffusionReactionSystem() + boundaries = ['D'] + n = 100 + problem.discretise_domain(n, 'random', locations=boundaries) + pinn = CausalPINN(problem = problem, + model=model, extra_features=None, loss=LpLoss()) + trainer = Trainer(solver=pinn, max_epochs=1, + accelerator='cpu', batch_size=20) + trainer.train() + + +# # TODO does not currently work +# def test_train_inverse_problem_restore(): +# tmpdir = "tests/tmp_restore_inv" +# problem = InverseDiffusionReactionSystem() +# boundaries = ['D'] +# n = 100 +# problem.discretise_domain(n, 'random', locations=boundaries) +# pinn = CausalPINN(problem=problem, +# model=model, +# extra_features=None, +# loss=LpLoss()) +# trainer = Trainer(solver=pinn, +# max_epochs=5, +# accelerator='cpu', +# default_root_dir=tmpdir) +# trainer.train() +# ntrainer = Trainer(solver=pinn, max_epochs=5, 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_inverse_problem_load(): + tmpdir = "tests/tmp_load_inv" + problem = InverseDiffusionReactionSystem() + boundaries = ['D'] + n = 100 + problem.discretise_domain(n, 'random', locations=boundaries) + pinn = CausalPINN(problem=problem, + model=model, + extra_features=None, + loss=LpLoss()) + trainer = Trainer(solver=pinn, + max_epochs=15, + accelerator='cpu', + default_root_dir=tmpdir) + trainer.train() + new_pinn = CausalPINN.load_from_checkpoint( + f'{tmpdir}/lightning_logs/version_0/checkpoints/epoch=14-step=30.ckpt', + problem = problem, model=model) + test_pts = CartesianDomain({'x': [0, 1], 't': [0, 1]}).sample(10) + assert new_pinn.forward(test_pts).extract(['u']).shape == (10, 1) + assert new_pinn.forward(test_pts).extract( + ['u']).shape == pinn.forward(test_pts).extract(['u']).shape + torch.testing.assert_close( + new_pinn.forward(test_pts).extract(['u']), + pinn.forward(test_pts).extract(['u'])) + import shutil + shutil.rmtree(tmpdir) + + +def test_train_extra_feats_cpu(): + problem = DiffusionReactionSystem() + boundaries = ['D'] + n = 10 + problem.discretise_domain(n, 'grid', locations=boundaries) + pinn = CausalPINN(problem=problem, + model=model_extra_feats, + extra_features=extra_feats) + trainer = Trainer(solver=pinn, max_epochs=5, accelerator='cpu') + trainer.train() \ No newline at end of file diff --git a/tests/test_solvers/test_competitive_pinn.py b/tests/test_solvers/test_competitive_pinn.py new file mode 100644 index 000000000..9c6f6c8ec --- /dev/null +++ b/tests/test_solvers/test_competitive_pinn.py @@ -0,0 +1,418 @@ +import torch +import pytest + +from pina.problem import SpatialProblem, InverseProblem +from pina.operators import laplacian +from pina.geometry import CartesianDomain +from pina import Condition, LabelTensor +from pina.solvers import CompetitivePINN as PINN +from pina.trainer import Trainer +from pina.model import FeedForward +from pina.equation.equation import Equation +from pina.equation.equation_factory import FixedValue +from pina.loss import LpLoss + + +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) +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 InversePoisson(SpatialProblem, InverseProblem): + ''' + Problem definition for the Poisson equation. + ''' + output_variables = ['u'] + x_min = -2 + x_max = 2 + y_min = -2 + y_max = 2 + data_input = LabelTensor(torch.rand(10, 2), ['x', 'y']) + data_output = LabelTensor(torch.rand(10, 1), ['u']) + spatial_domain = CartesianDomain({'x': [x_min, x_max], 'y': [y_min, y_max]}) + # define the ranges for the parameters + unknown_parameter_domain = CartesianDomain({'mu1': [-1, 1], 'mu2': [-1, 1]}) + + def laplace_equation(input_, output_, params_): + ''' + Laplace equation with a force term. + ''' + force_term = torch.exp( + - 2*(input_.extract(['x']) - params_['mu1'])**2 + - 2*(input_.extract(['y']) - params_['mu2'])**2) + delta_u = laplacian(output_, input_, components=['u'], d=['x', 'y']) + + return delta_u - force_term + + # define the conditions for the loss (boundary conditions, equation, data) + conditions = { + 'gamma1': Condition(location=CartesianDomain({'x': [x_min, x_max], + 'y': y_max}), + equation=FixedValue(0.0, components=['u'])), + 'gamma2': Condition(location=CartesianDomain( + {'x': [x_min, x_max], 'y': y_min + }), + equation=FixedValue(0.0, components=['u'])), + 'gamma3': Condition(location=CartesianDomain( + {'x': x_max, 'y': [y_min, y_max] + }), + equation=FixedValue(0.0, components=['u'])), + 'gamma4': Condition(location=CartesianDomain( + {'x': x_min, 'y': [y_min, y_max] + }), + equation=FixedValue(0.0, components=['u'])), + 'D': Condition(location=CartesianDomain( + {'x': [x_min, x_max], 'y': [y_min, y_max] + }), + equation=Equation(laplace_equation)), + 'data': Condition(input_points=data_input.extract(['x', 'y']), + output_points=data_output) + } + + +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_) + } + + 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 + + +class myFeature(torch.nn.Module): + """ + Feature: sin(x) + """ + + def __init__(self): + super(myFeature, self).__init__() + + def forward(self, x): + t = (torch.sin(x.extract(['x']) * torch.pi) * + torch.sin(x.extract(['y']) * torch.pi)) + return LabelTensor(t, ['sin(x)sin(y)']) + + +# make the problem +poisson_problem = Poisson() +model = FeedForward(len(poisson_problem.input_variables), + len(poisson_problem.output_variables)) +model_extra_feats = FeedForward( + len(poisson_problem.input_variables) + 1, + len(poisson_problem.output_variables)) +extra_feats = [myFeature()] + + +def test_constructor(): + PINN(problem=poisson_problem, model=model) + PINN(problem=poisson_problem, model=model, discriminator = model) + + +def test_constructor_extra_feats(): + with pytest.raises(TypeError): + model_extra_feats = FeedForward( + len(poisson_problem.input_variables) + 1, + len(poisson_problem.output_variables)) + PINN(problem=poisson_problem, + model=model_extra_feats, + extra_features=extra_feats) + + +def test_train_cpu(): + poisson_problem = Poisson() + boundaries = ['gamma1', 'gamma2', 'gamma3', 'gamma4'] + n = 10 + poisson_problem.discretise_domain(n, 'grid', locations=boundaries) + pinn = PINN(problem = poisson_problem, model=model, loss=LpLoss()) + trainer = Trainer(solver=pinn, max_epochs=1, + accelerator='cpu', batch_size=20) + trainer.train() + + +def test_train_restore(): + tmpdir = "tests/tmp_restore" + poisson_problem = Poisson() + boundaries = ['gamma1', 'gamma2', 'gamma3', 'gamma4'] + n = 10 + poisson_problem.discretise_domain(n, 'grid', locations=boundaries) + pinn = PINN(problem=poisson_problem, + model=model, + loss=LpLoss()) + trainer = Trainer(solver=pinn, + max_epochs=5, + accelerator='cpu', + default_root_dir=tmpdir) + trainer.train() + ntrainer = Trainer(solver=pinn, max_epochs=15, accelerator='cpu') + t = ntrainer.train( + ckpt_path=f'{tmpdir}/lightning_logs/version_0/' + 'checkpoints/epoch=4-step=10.ckpt') + import shutil + shutil.rmtree(tmpdir) + + +def test_train_load(): + tmpdir = "tests/tmp_load" + poisson_problem = Poisson() + boundaries = ['gamma1', 'gamma2', 'gamma3', 'gamma4'] + n = 10 + poisson_problem.discretise_domain(n, 'grid', locations=boundaries) + pinn = PINN(problem=poisson_problem, + model=model, + loss=LpLoss()) + trainer = Trainer(solver=pinn, + max_epochs=15, + accelerator='cpu', + default_root_dir=tmpdir) + trainer.train() + new_pinn = PINN.load_from_checkpoint( + f'{tmpdir}/lightning_logs/version_0/checkpoints/epoch=14-step=30.ckpt', + problem = poisson_problem, model=model) + test_pts = CartesianDomain({'x': [0, 1], 'y': [0, 1]}).sample(10) + assert new_pinn.forward(test_pts).extract(['u']).shape == (10, 1) + assert new_pinn.forward(test_pts).extract( + ['u']).shape == pinn.forward(test_pts).extract(['u']).shape + torch.testing.assert_close( + new_pinn.forward(test_pts).extract(['u']), + pinn.forward(test_pts).extract(['u'])) + import shutil + shutil.rmtree(tmpdir) + +def test_train_inverse_problem_cpu(): + poisson_problem = InversePoisson() + boundaries = ['gamma1', 'gamma2', 'gamma3', 'gamma4', 'D'] + n = 100 + poisson_problem.discretise_domain(n, 'random', locations=boundaries) + pinn = PINN(problem = poisson_problem, model=model, loss=LpLoss()) + trainer = Trainer(solver=pinn, max_epochs=1, + accelerator='cpu', batch_size=20) + trainer.train() + + +# # TODO does not currently work +# def test_train_inverse_problem_restore(): +# tmpdir = "tests/tmp_restore_inv" +# poisson_problem = InversePoisson() +# boundaries = ['gamma1', 'gamma2', 'gamma3', 'gamma4', 'D'] +# n = 100 +# poisson_problem.discretise_domain(n, 'random', locations=boundaries) +# pinn = PINN(problem=poisson_problem, +# model=model, +# loss=LpLoss()) +# trainer = Trainer(solver=pinn, +# max_epochs=5, +# accelerator='cpu', +# default_root_dir=tmpdir) +# trainer.train() +# ntrainer = Trainer(solver=pinn, max_epochs=5, accelerator='cpu') +# t = ntrainer.train( +# ckpt_path=f'{tmpdir}/lightning_logs/version_0/checkpoints/epoch=4-step=10.ckpt') +# import shutil +# shutil.rmtree(tmpdir) + + +def test_train_inverse_problem_load(): + tmpdir = "tests/tmp_load_inv" + poisson_problem = InversePoisson() + boundaries = ['gamma1', 'gamma2', 'gamma3', 'gamma4', 'D'] + n = 100 + poisson_problem.discretise_domain(n, 'random', locations=boundaries) + pinn = PINN(problem=poisson_problem, + model=model, + loss=LpLoss()) + trainer = Trainer(solver=pinn, + max_epochs=15, + accelerator='cpu', + default_root_dir=tmpdir) + trainer.train() + new_pinn = PINN.load_from_checkpoint( + f'{tmpdir}/lightning_logs/version_0/checkpoints/epoch=14-step=30.ckpt', + problem = poisson_problem, model=model) + test_pts = CartesianDomain({'x': [0, 1], 'y': [0, 1]}).sample(10) + assert new_pinn.forward(test_pts).extract(['u']).shape == (10, 1) + assert new_pinn.forward(test_pts).extract( + ['u']).shape == pinn.forward(test_pts).extract(['u']).shape + torch.testing.assert_close( + new_pinn.forward(test_pts).extract(['u']), + pinn.forward(test_pts).extract(['u'])) + import shutil + shutil.rmtree(tmpdir) + +# # TODO fix asap. Basically sampling few variables +# # works only if both variables are in a range. +# # if one is fixed and the other not, this will +# # not work. This test also needs to be fixed and +# # insert in test problem not in test pinn. +# def test_train_cpu_sampling_few_vars(): +# poisson_problem = Poisson() +# boundaries = ['gamma1', 'gamma2', 'gamma3'] +# n = 10 +# poisson_problem.discretise_domain(n, 'grid', locations=boundaries) +# poisson_problem.discretise_domain(n, 'random', locations=['gamma4'], variables=['x']) +# poisson_problem.discretise_domain(n, 'random', locations=['gamma4'], variables=['y']) +# pinn = PINN(problem = poisson_problem, model=model, extra_features=None, loss=LpLoss()) +# trainer = Trainer(solver=pinn, kwargs={'max_epochs' : 5, 'accelerator':'cpu'}) +# trainer.train() + + +# TODO, fix GitHub actions to run also on GPU +# def test_train_gpu(): +# poisson_problem = Poisson() +# boundaries = ['gamma1', 'gamma2', 'gamma3', 'gamma4'] +# n = 10 +# poisson_problem.discretise_domain(n, 'grid', locations=boundaries) +# pinn = PINN(problem = poisson_problem, model=model, extra_features=None, loss=LpLoss()) +# trainer = Trainer(solver=pinn, kwargs={'max_epochs' : 5, 'accelerator':'gpu'}) +# trainer.train() + +# def test_train_gpu(): #TODO fix ASAP +# poisson_problem = Poisson() +# boundaries = ['gamma1', 'gamma2', 'gamma3', 'gamma4'] +# n = 10 +# poisson_problem.discretise_domain(n, 'grid', locations=boundaries) +# poisson_problem.conditions.pop('data') # The input/output pts are allocated on cpu +# pinn = PINN(problem = poisson_problem, model=model, extra_features=None, loss=LpLoss()) +# trainer = Trainer(solver=pinn, kwargs={'max_epochs' : 5, 'accelerator':'gpu'}) +# trainer.train() + +# def test_train_2(): +# boundaries = ['gamma1', 'gamma2', 'gamma3', 'gamma4'] +# n = 10 +# expected_keys = [[], list(range(0, 50, 3))] +# param = [0, 3] +# for i, truth_key in zip(param, expected_keys): +# pinn = PINN(problem, model) +# pinn.discretise_domain(n, 'grid', locations=boundaries) +# pinn.discretise_domain(n, 'grid', locations=['D']) +# pinn.train(50, save_loss=i) +# assert list(pinn.history_loss.keys()) == truth_key + + +# def test_train_extra_feats(): +# pinn = PINN(problem, model_extra_feat, [myFeature()]) +# boundaries = ['gamma1', 'gamma2', 'gamma3', 'gamma4'] +# n = 10 +# pinn.discretise_domain(n, 'grid', locations=boundaries) +# pinn.discretise_domain(n, 'grid', locations=['D']) +# pinn.train(5) + + +# def test_train_2_extra_feats(): +# boundaries = ['gamma1', 'gamma2', 'gamma3', 'gamma4'] +# n = 10 +# expected_keys = [[], list(range(0, 50, 3))] +# param = [0, 3] +# for i, truth_key in zip(param, expected_keys): +# pinn = PINN(problem, model_extra_feat, [myFeature()]) +# pinn.discretise_domain(n, 'grid', locations=boundaries) +# pinn.discretise_domain(n, 'grid', locations=['D']) +# pinn.train(50, save_loss=i) +# assert list(pinn.history_loss.keys()) == truth_key + + +# def test_train_with_optimizer_kwargs(): +# boundaries = ['gamma1', 'gamma2', 'gamma3', 'gamma4'] +# n = 10 +# expected_keys = [[], list(range(0, 50, 3))] +# param = [0, 3] +# for i, truth_key in zip(param, expected_keys): +# pinn = PINN(problem, model, optimizer_kwargs={'lr' : 0.3}) +# pinn.discretise_domain(n, 'grid', locations=boundaries) +# pinn.discretise_domain(n, 'grid', locations=['D']) +# pinn.train(50, save_loss=i) +# assert list(pinn.history_loss.keys()) == truth_key + + +# def test_train_with_lr_scheduler(): +# boundaries = ['gamma1', 'gamma2', 'gamma3', 'gamma4'] +# n = 10 +# expected_keys = [[], list(range(0, 50, 3))] +# param = [0, 3] +# for i, truth_key in zip(param, expected_keys): +# pinn = PINN( +# problem, +# model, +# lr_scheduler_type=torch.optim.lr_scheduler.CyclicLR, +# lr_scheduler_kwargs={'base_lr' : 0.1, 'max_lr' : 0.3, 'cycle_momentum': False} +# ) +# pinn.discretise_domain(n, 'grid', locations=boundaries) +# pinn.discretise_domain(n, 'grid', locations=['D']) +# pinn.train(50, save_loss=i) +# assert list(pinn.history_loss.keys()) == truth_key + + +# # def test_train_batch(): +# # pinn = PINN(problem, model, batch_size=6) +# # boundaries = ['gamma1', 'gamma2', 'gamma3', 'gamma4'] +# # n = 10 +# # pinn.discretise_domain(n, 'grid', locations=boundaries) +# # pinn.discretise_domain(n, 'grid', locations=['D']) +# # pinn.train(5) + + +# # def test_train_batch_2(): +# # boundaries = ['gamma1', 'gamma2', 'gamma3', 'gamma4'] +# # n = 10 +# # expected_keys = [[], list(range(0, 50, 3))] +# # param = [0, 3] +# # for i, truth_key in zip(param, expected_keys): +# # pinn = PINN(problem, model, batch_size=6) +# # pinn.discretise_domain(n, 'grid', locations=boundaries) +# # pinn.discretise_domain(n, 'grid', locations=['D']) +# # pinn.train(50, save_loss=i) +# # assert list(pinn.history_loss.keys()) == truth_key + + +# if torch.cuda.is_available(): + +# # def test_gpu_train(): +# # pinn = PINN(problem, model, batch_size=20, device='cuda') +# # boundaries = ['gamma1', 'gamma2', 'gamma3', 'gamma4'] +# # n = 100 +# # pinn.discretise_domain(n, 'grid', locations=boundaries) +# # pinn.discretise_domain(n, 'grid', locations=['D']) +# # pinn.train(5) + +# def test_gpu_train_nobatch(): +# pinn = PINN(problem, model, batch_size=None, device='cuda') +# boundaries = ['gamma1', 'gamma2', 'gamma3', 'gamma4'] +# n = 100 +# pinn.discretise_domain(n, 'grid', locations=boundaries) +# pinn.discretise_domain(n, 'grid', locations=['D']) +# pinn.train(5) + diff --git a/tests/test_solvers/test_gpinn.py b/tests/test_solvers/test_gpinn.py new file mode 100644 index 000000000..d00d3b4dc --- /dev/null +++ b/tests/test_solvers/test_gpinn.py @@ -0,0 +1,432 @@ +import torch + +from pina.problem import SpatialProblem, InverseProblem +from pina.operators import laplacian +from pina.geometry import CartesianDomain +from pina import Condition, LabelTensor +from pina.solvers import GPINN +from pina.trainer import Trainer +from pina.model import FeedForward +from pina.equation.equation import Equation +from pina.equation.equation_factory import FixedValue +from pina.loss import LpLoss + + +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) +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 InversePoisson(SpatialProblem, InverseProblem): + ''' + Problem definition for the Poisson equation. + ''' + output_variables = ['u'] + x_min = -2 + x_max = 2 + y_min = -2 + y_max = 2 + data_input = LabelTensor(torch.rand(10, 2), ['x', 'y']) + data_output = LabelTensor(torch.rand(10, 1), ['u']) + spatial_domain = CartesianDomain({'x': [x_min, x_max], 'y': [y_min, y_max]}) + # define the ranges for the parameters + unknown_parameter_domain = CartesianDomain({'mu1': [-1, 1], 'mu2': [-1, 1]}) + + def laplace_equation(input_, output_, params_): + ''' + Laplace equation with a force term. + ''' + force_term = torch.exp( + - 2*(input_.extract(['x']) - params_['mu1'])**2 + - 2*(input_.extract(['y']) - params_['mu2'])**2) + delta_u = laplacian(output_, input_, components=['u'], d=['x', 'y']) + + return delta_u - force_term + + # define the conditions for the loss (boundary conditions, equation, data) + conditions = { + 'gamma1': Condition(location=CartesianDomain({'x': [x_min, x_max], + 'y': y_max}), + equation=FixedValue(0.0, components=['u'])), + 'gamma2': Condition(location=CartesianDomain( + {'x': [x_min, x_max], 'y': y_min}), + equation=FixedValue(0.0, components=['u'])), + 'gamma3': Condition(location=CartesianDomain( + {'x': x_max, 'y': [y_min, y_max]}), + equation=FixedValue(0.0, components=['u'])), + 'gamma4': Condition(location=CartesianDomain( + {'x': x_min, 'y': [y_min, y_max] + }), + equation=FixedValue(0.0, components=['u'])), + 'D': Condition(location=CartesianDomain( + {'x': [x_min, x_max], 'y': [y_min, y_max] + }), + equation=Equation(laplace_equation)), + 'data': Condition( + input_points=data_input.extract(['x', 'y']), + output_points=data_output) + } + + +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_) + } + + 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 + + +class myFeature(torch.nn.Module): + """ + Feature: sin(x) + """ + + def __init__(self): + super(myFeature, self).__init__() + + def forward(self, x): + t = (torch.sin(x.extract(['x']) * torch.pi) * + torch.sin(x.extract(['y']) * torch.pi)) + return LabelTensor(t, ['sin(x)sin(y)']) + + +# make the problem +poisson_problem = Poisson() +model = FeedForward(len(poisson_problem.input_variables), + len(poisson_problem.output_variables)) +model_extra_feats = FeedForward( + len(poisson_problem.input_variables) + 1, + len(poisson_problem.output_variables)) +extra_feats = [myFeature()] + + +def test_constructor(): + GPINN(problem=poisson_problem, model=model, extra_features=None) + + +def test_constructor_extra_feats(): + model_extra_feats = FeedForward( + len(poisson_problem.input_variables) + 1, + len(poisson_problem.output_variables)) + GPINN(problem=poisson_problem, + model=model_extra_feats, + extra_features=extra_feats) + + +def test_train_cpu(): + poisson_problem = Poisson() + boundaries = ['gamma1', 'gamma2', 'gamma3', 'gamma4'] + n = 10 + poisson_problem.discretise_domain(n, 'grid', locations=boundaries) + pinn = GPINN(problem = poisson_problem, + model=model, extra_features=None, loss=LpLoss()) + trainer = Trainer(solver=pinn, max_epochs=1, + accelerator='cpu', batch_size=20) + trainer.train() + + +def test_train_restore(): + tmpdir = "tests/tmp_restore" + poisson_problem = Poisson() + boundaries = ['gamma1', 'gamma2', 'gamma3', 'gamma4'] + n = 10 + poisson_problem.discretise_domain(n, 'grid', locations=boundaries) + pinn = GPINN(problem=poisson_problem, + model=model, + extra_features=None, + loss=LpLoss()) + trainer = Trainer(solver=pinn, + max_epochs=5, + accelerator='cpu', + default_root_dir=tmpdir) + trainer.train() + ntrainer = Trainer(solver=pinn, max_epochs=15, accelerator='cpu') + t = ntrainer.train( + ckpt_path=f'{tmpdir}/lightning_logs/version_0/' + 'checkpoints/epoch=4-step=10.ckpt') + import shutil + shutil.rmtree(tmpdir) + + +def test_train_load(): + tmpdir = "tests/tmp_load" + poisson_problem = Poisson() + boundaries = ['gamma1', 'gamma2', 'gamma3', 'gamma4'] + n = 10 + poisson_problem.discretise_domain(n, 'grid', locations=boundaries) + pinn = GPINN(problem=poisson_problem, + model=model, + extra_features=None, + loss=LpLoss()) + trainer = Trainer(solver=pinn, + max_epochs=15, + accelerator='cpu', + default_root_dir=tmpdir) + trainer.train() + new_pinn = GPINN.load_from_checkpoint( + f'{tmpdir}/lightning_logs/version_0/checkpoints/epoch=14-step=30.ckpt', + problem = poisson_problem, model=model) + test_pts = CartesianDomain({'x': [0, 1], 'y': [0, 1]}).sample(10) + assert new_pinn.forward(test_pts).extract(['u']).shape == (10, 1) + assert new_pinn.forward(test_pts).extract( + ['u']).shape == pinn.forward(test_pts).extract(['u']).shape + torch.testing.assert_close( + new_pinn.forward(test_pts).extract(['u']), + pinn.forward(test_pts).extract(['u'])) + import shutil + shutil.rmtree(tmpdir) + +def test_train_inverse_problem_cpu(): + poisson_problem = InversePoisson() + boundaries = ['gamma1', 'gamma2', 'gamma3', 'gamma4', 'D'] + n = 100 + poisson_problem.discretise_domain(n, 'random', locations=boundaries) + pinn = GPINN(problem = poisson_problem, + model=model, extra_features=None, loss=LpLoss()) + trainer = Trainer(solver=pinn, max_epochs=1, + accelerator='cpu', batch_size=20) + trainer.train() + + +# # TODO does not currently work +# def test_train_inverse_problem_restore(): +# tmpdir = "tests/tmp_restore_inv" +# poisson_problem = InversePoisson() +# boundaries = ['gamma1', 'gamma2', 'gamma3', 'gamma4', 'D'] +# n = 100 +# poisson_problem.discretise_domain(n, 'random', locations=boundaries) +# pinn = GPINN(problem=poisson_problem, +# model=model, +# extra_features=None, +# loss=LpLoss()) +# trainer = Trainer(solver=pinn, +# max_epochs=5, +# accelerator='cpu', +# default_root_dir=tmpdir) +# trainer.train() +# ntrainer = Trainer(solver=pinn, max_epochs=5, accelerator='cpu') +# t = ntrainer.train( +# ckpt_path=f'{tmpdir}/lightning_logs/version_0/checkpoints/epoch=4-step=10.ckpt') +# import shutil +# shutil.rmtree(tmpdir) + + +def test_train_inverse_problem_load(): + tmpdir = "tests/tmp_load_inv" + poisson_problem = InversePoisson() + boundaries = ['gamma1', 'gamma2', 'gamma3', 'gamma4', 'D'] + n = 100 + poisson_problem.discretise_domain(n, 'random', locations=boundaries) + pinn = GPINN(problem=poisson_problem, + model=model, + extra_features=None, + loss=LpLoss()) + trainer = Trainer(solver=pinn, + max_epochs=15, + accelerator='cpu', + default_root_dir=tmpdir) + trainer.train() + new_pinn = GPINN.load_from_checkpoint( + f'{tmpdir}/lightning_logs/version_0/checkpoints/epoch=14-step=30.ckpt', + problem = poisson_problem, model=model) + test_pts = CartesianDomain({'x': [0, 1], 'y': [0, 1]}).sample(10) + assert new_pinn.forward(test_pts).extract(['u']).shape == (10, 1) + assert new_pinn.forward(test_pts).extract( + ['u']).shape == pinn.forward(test_pts).extract(['u']).shape + torch.testing.assert_close( + new_pinn.forward(test_pts).extract(['u']), + pinn.forward(test_pts).extract(['u'])) + import shutil + shutil.rmtree(tmpdir) + +# # TODO fix asap. Basically sampling few variables +# # works only if both variables are in a range. +# # if one is fixed and the other not, this will +# # not work. This test also needs to be fixed and +# # insert in test problem not in test pinn. +# def test_train_cpu_sampling_few_vars(): +# poisson_problem = Poisson() +# boundaries = ['gamma1', 'gamma2', 'gamma3'] +# n = 10 +# poisson_problem.discretise_domain(n, 'grid', locations=boundaries) +# poisson_problem.discretise_domain(n, 'random', locations=['gamma4'], variables=['x']) +# poisson_problem.discretise_domain(n, 'random', locations=['gamma4'], variables=['y']) +# pinn = GPINN(problem = poisson_problem, model=model, extra_features=None, loss=LpLoss()) +# trainer = Trainer(solver=pinn, kwargs={'max_epochs' : 5, 'accelerator':'cpu'}) +# trainer.train() + + +def test_train_extra_feats_cpu(): + poisson_problem = Poisson() + boundaries = ['gamma1', 'gamma2', 'gamma3', 'gamma4'] + n = 10 + poisson_problem.discretise_domain(n, 'grid', locations=boundaries) + pinn = GPINN(problem=poisson_problem, + model=model_extra_feats, + extra_features=extra_feats) + trainer = Trainer(solver=pinn, max_epochs=5, accelerator='cpu') + trainer.train() + + +# TODO, fix GitHub actions to run also on GPU +# def test_train_gpu(): +# poisson_problem = Poisson() +# boundaries = ['gamma1', 'gamma2', 'gamma3', 'gamma4'] +# n = 10 +# poisson_problem.discretise_domain(n, 'grid', locations=boundaries) +# pinn = GPINN(problem = poisson_problem, model=model, extra_features=None, loss=LpLoss()) +# trainer = Trainer(solver=pinn, kwargs={'max_epochs' : 5, 'accelerator':'gpu'}) +# trainer.train() + +# def test_train_gpu(): #TODO fix ASAP +# poisson_problem = Poisson() +# boundaries = ['gamma1', 'gamma2', 'gamma3', 'gamma4'] +# n = 10 +# poisson_problem.discretise_domain(n, 'grid', locations=boundaries) +# poisson_problem.conditions.pop('data') # The input/output pts are allocated on cpu +# pinn = GPINN(problem = poisson_problem, model=model, extra_features=None, loss=LpLoss()) +# trainer = Trainer(solver=pinn, kwargs={'max_epochs' : 5, 'accelerator':'gpu'}) +# trainer.train() + +# def test_train_2(): +# boundaries = ['gamma1', 'gamma2', 'gamma3', 'gamma4'] +# n = 10 +# expected_keys = [[], list(range(0, 50, 3))] +# param = [0, 3] +# for i, truth_key in zip(param, expected_keys): +# pinn = GPINN(problem, model) +# pinn.discretise_domain(n, 'grid', locations=boundaries) +# pinn.discretise_domain(n, 'grid', locations=['D']) +# pinn.train(50, save_loss=i) +# assert list(pinn.history_loss.keys()) == truth_key + + +# def test_train_extra_feats(): +# pinn = GPINN(problem, model_extra_feat, [myFeature()]) +# boundaries = ['gamma1', 'gamma2', 'gamma3', 'gamma4'] +# n = 10 +# pinn.discretise_domain(n, 'grid', locations=boundaries) +# pinn.discretise_domain(n, 'grid', locations=['D']) +# pinn.train(5) + + +# def test_train_2_extra_feats(): +# boundaries = ['gamma1', 'gamma2', 'gamma3', 'gamma4'] +# n = 10 +# expected_keys = [[], list(range(0, 50, 3))] +# param = [0, 3] +# for i, truth_key in zip(param, expected_keys): +# pinn = GPINN(problem, model_extra_feat, [myFeature()]) +# pinn.discretise_domain(n, 'grid', locations=boundaries) +# pinn.discretise_domain(n, 'grid', locations=['D']) +# pinn.train(50, save_loss=i) +# assert list(pinn.history_loss.keys()) == truth_key + + +# def test_train_with_optimizer_kwargs(): +# boundaries = ['gamma1', 'gamma2', 'gamma3', 'gamma4'] +# n = 10 +# expected_keys = [[], list(range(0, 50, 3))] +# param = [0, 3] +# for i, truth_key in zip(param, expected_keys): +# pinn = GPINN(problem, model, optimizer_kwargs={'lr' : 0.3}) +# pinn.discretise_domain(n, 'grid', locations=boundaries) +# pinn.discretise_domain(n, 'grid', locations=['D']) +# pinn.train(50, save_loss=i) +# assert list(pinn.history_loss.keys()) == truth_key + + +# def test_train_with_lr_scheduler(): +# boundaries = ['gamma1', 'gamma2', 'gamma3', 'gamma4'] +# n = 10 +# expected_keys = [[], list(range(0, 50, 3))] +# param = [0, 3] +# for i, truth_key in zip(param, expected_keys): +# pinn = GPINN( +# problem, +# model, +# lr_scheduler_type=torch.optim.lr_scheduler.CyclicLR, +# lr_scheduler_kwargs={'base_lr' : 0.1, 'max_lr' : 0.3, 'cycle_momentum': False} +# ) +# pinn.discretise_domain(n, 'grid', locations=boundaries) +# pinn.discretise_domain(n, 'grid', locations=['D']) +# pinn.train(50, save_loss=i) +# assert list(pinn.history_loss.keys()) == truth_key + + +# # def test_train_batch(): +# # pinn = GPINN(problem, model, batch_size=6) +# # boundaries = ['gamma1', 'gamma2', 'gamma3', 'gamma4'] +# # n = 10 +# # pinn.discretise_domain(n, 'grid', locations=boundaries) +# # pinn.discretise_domain(n, 'grid', locations=['D']) +# # pinn.train(5) + + +# # def test_train_batch_2(): +# # boundaries = ['gamma1', 'gamma2', 'gamma3', 'gamma4'] +# # n = 10 +# # expected_keys = [[], list(range(0, 50, 3))] +# # param = [0, 3] +# # for i, truth_key in zip(param, expected_keys): +# # pinn = GPINN(problem, model, batch_size=6) +# # pinn.discretise_domain(n, 'grid', locations=boundaries) +# # pinn.discretise_domain(n, 'grid', locations=['D']) +# # pinn.train(50, save_loss=i) +# # assert list(pinn.history_loss.keys()) == truth_key + + +# if torch.cuda.is_available(): + +# # def test_gpu_train(): +# # pinn = GPINN(problem, model, batch_size=20, device='cuda') +# # boundaries = ['gamma1', 'gamma2', 'gamma3', 'gamma4'] +# # n = 100 +# # pinn.discretise_domain(n, 'grid', locations=boundaries) +# # pinn.discretise_domain(n, 'grid', locations=['D']) +# # pinn.train(5) + +# def test_gpu_train_nobatch(): +# pinn = GPINN(problem, model, batch_size=None, device='cuda') +# boundaries = ['gamma1', 'gamma2', 'gamma3', 'gamma4'] +# n = 100 +# pinn.discretise_domain(n, 'grid', locations=boundaries) +# pinn.discretise_domain(n, 'grid', locations=['D']) +# pinn.train(5) + diff --git a/tests/test_solvers/test_pinn.py b/tests/test_solvers/test_pinn.py index 0a42410ca..ea3b077bc 100644 --- a/tests/test_solvers/test_pinn.py +++ b/tests/test_solvers/test_pinn.py @@ -1,6 +1,6 @@ import torch -from pina.problem import SpatialProblem +from pina.problem import SpatialProblem, InverseProblem from pina.operators import laplacian from pina.geometry import CartesianDomain from pina import Condition, LabelTensor @@ -26,6 +26,58 @@ def laplace_equation(input_, output_): out2_ = LabelTensor(torch.rand(60, 1), ['u']) +class InversePoisson(SpatialProblem, InverseProblem): + ''' + Problem definition for the Poisson equation. + ''' + output_variables = ['u'] + x_min = -2 + x_max = 2 + y_min = -2 + y_max = 2 + data_input = LabelTensor(torch.rand(10, 2), ['x', 'y']) + data_output = LabelTensor(torch.rand(10, 1), ['u']) + spatial_domain = CartesianDomain({'x': [x_min, x_max], 'y': [y_min, y_max]}) + # define the ranges for the parameters + unknown_parameter_domain = CartesianDomain({'mu1': [-1, 1], 'mu2': [-1, 1]}) + + def laplace_equation(input_, output_, params_): + ''' + Laplace equation with a force term. + ''' + force_term = torch.exp( + - 2*(input_.extract(['x']) - params_['mu1'])**2 + - 2*(input_.extract(['y']) - params_['mu2'])**2) + delta_u = laplacian(output_, input_, components=['u'], d=['x', 'y']) + + return delta_u - force_term + + # define the conditions for the loss (boundary conditions, equation, data) + conditions = { + 'gamma1': Condition(location=CartesianDomain({'x': [x_min, x_max], + 'y': y_max}), + equation=FixedValue(0.0, components=['u'])), + 'gamma2': Condition(location=CartesianDomain( + {'x': [x_min, x_max], 'y': y_min + }), + equation=FixedValue(0.0, components=['u'])), + 'gamma3': Condition(location=CartesianDomain( + {'x': x_max, 'y': [y_min, y_max] + }), + equation=FixedValue(0.0, components=['u'])), + 'gamma4': Condition(location=CartesianDomain( + {'x': x_min, 'y': [y_min, y_max] + }), + equation=FixedValue(0.0, components=['u'])), + 'D': Condition(location=CartesianDomain( + {'x': [x_min, x_max], 'y': [y_min, y_max] + }), + equation=Equation(laplace_equation)), + 'data': Condition(input_points=data_input.extract(['x', 'y']), + output_points=data_output) + } + + class Poisson(SpatialProblem): output_variables = ['u'] spatial_domain = CartesianDomain({'x': [0, 1], 'y': [0, 1]}) @@ -103,8 +155,10 @@ def test_train_cpu(): boundaries = ['gamma1', 'gamma2', 'gamma3', 'gamma4'] n = 10 poisson_problem.discretise_domain(n, 'grid', locations=boundaries) - pinn = PINN(problem = poisson_problem, model=model, extra_features=None, loss=LpLoss()) - trainer = Trainer(solver=pinn, max_epochs=1, accelerator='cpu', batch_size=20) + pinn = PINN(problem = poisson_problem, model=model, + extra_features=None, loss=LpLoss()) + trainer = Trainer(solver=pinn, max_epochs=1, + accelerator='cpu', batch_size=20) trainer.train() @@ -125,7 +179,8 @@ def test_train_restore(): trainer.train() ntrainer = Trainer(solver=pinn, max_epochs=15, accelerator='cpu') t = ntrainer.train( - ckpt_path=f'{tmpdir}/lightning_logs/version_0/checkpoints/epoch=4-step=10.ckpt') + ckpt_path=f'{tmpdir}/lightning_logs/version_0/' + 'checkpoints/epoch=4-step=10.ckpt') import shutil shutil.rmtree(tmpdir) @@ -158,6 +213,68 @@ def test_train_load(): import shutil shutil.rmtree(tmpdir) +def test_train_inverse_problem_cpu(): + poisson_problem = InversePoisson() + boundaries = ['gamma1', 'gamma2', 'gamma3', 'gamma4', 'D'] + n = 100 + poisson_problem.discretise_domain(n, 'random', locations=boundaries) + pinn = PINN(problem = poisson_problem, model=model, + extra_features=None, loss=LpLoss()) + trainer = Trainer(solver=pinn, max_epochs=1, + accelerator='cpu', batch_size=20) + trainer.train() + + +# # TODO does not currently work +# def test_train_inverse_problem_restore(): +# tmpdir = "tests/tmp_restore_inv" +# poisson_problem = InversePoisson() +# boundaries = ['gamma1', 'gamma2', 'gamma3', 'gamma4', 'D'] +# n = 100 +# poisson_problem.discretise_domain(n, 'random', locations=boundaries) +# pinn = PINN(problem=poisson_problem, +# model=model, +# extra_features=None, +# loss=LpLoss()) +# trainer = Trainer(solver=pinn, +# max_epochs=5, +# accelerator='cpu', +# default_root_dir=tmpdir) +# trainer.train() +# ntrainer = Trainer(solver=pinn, max_epochs=5, accelerator='cpu') +# t = ntrainer.train( +# ckpt_path=f'{tmpdir}/lightning_logs/version_0/checkpoints/epoch=4-step=10.ckpt') +# import shutil +# shutil.rmtree(tmpdir) + + +def test_train_inverse_problem_load(): + tmpdir = "tests/tmp_load_inv" + poisson_problem = InversePoisson() + boundaries = ['gamma1', 'gamma2', 'gamma3', 'gamma4', 'D'] + n = 100 + poisson_problem.discretise_domain(n, 'random', locations=boundaries) + pinn = PINN(problem=poisson_problem, + model=model, + extra_features=None, + loss=LpLoss()) + trainer = Trainer(solver=pinn, + max_epochs=15, + accelerator='cpu', + default_root_dir=tmpdir) + trainer.train() + new_pinn = PINN.load_from_checkpoint( + f'{tmpdir}/lightning_logs/version_0/checkpoints/epoch=14-step=30.ckpt', + problem = poisson_problem, model=model) + test_pts = CartesianDomain({'x': [0, 1], 'y': [0, 1]}).sample(10) + assert new_pinn.forward(test_pts).extract(['u']).shape == (10, 1) + assert new_pinn.forward(test_pts).extract( + ['u']).shape == pinn.forward(test_pts).extract(['u']).shape + torch.testing.assert_close( + new_pinn.forward(test_pts).extract(['u']), + pinn.forward(test_pts).extract(['u'])) + import shutil + shutil.rmtree(tmpdir) # # TODO fix asap. Basically sampling few variables # # works only if both variables are in a range. @@ -197,85 +314,32 @@ def test_train_extra_feats_cpu(): # pinn = PINN(problem = poisson_problem, model=model, extra_features=None, loss=LpLoss()) # trainer = Trainer(solver=pinn, kwargs={'max_epochs' : 5, 'accelerator':'gpu'}) # trainer.train() -""" -def test_train_gpu(): #TODO fix ASAP - poisson_problem = Poisson() - boundaries = ['gamma1', 'gamma2', 'gamma3', 'gamma4'] - n = 10 - poisson_problem.discretise_domain(n, 'grid', locations=boundaries) - poisson_problem.conditions.pop('data') # The input/output pts are allocated on cpu - pinn = PINN(problem = poisson_problem, model=model, extra_features=None, loss=LpLoss()) - trainer = Trainer(solver=pinn, kwargs={'max_epochs' : 5, 'accelerator':'gpu'}) - trainer.train() - -def test_train_2(): - boundaries = ['gamma1', 'gamma2', 'gamma3', 'gamma4'] - n = 10 - expected_keys = [[], list(range(0, 50, 3))] - param = [0, 3] - for i, truth_key in zip(param, expected_keys): - pinn = PINN(problem, model) - pinn.discretise_domain(n, 'grid', locations=boundaries) - pinn.discretise_domain(n, 'grid', locations=['D']) - pinn.train(50, save_loss=i) - assert list(pinn.history_loss.keys()) == truth_key - - -def test_train_extra_feats(): - pinn = PINN(problem, model_extra_feat, [myFeature()]) - boundaries = ['gamma1', 'gamma2', 'gamma3', 'gamma4'] - n = 10 - pinn.discretise_domain(n, 'grid', locations=boundaries) - pinn.discretise_domain(n, 'grid', locations=['D']) - pinn.train(5) - - -def test_train_2_extra_feats(): - boundaries = ['gamma1', 'gamma2', 'gamma3', 'gamma4'] - n = 10 - expected_keys = [[], list(range(0, 50, 3))] - param = [0, 3] - for i, truth_key in zip(param, expected_keys): - pinn = PINN(problem, model_extra_feat, [myFeature()]) - pinn.discretise_domain(n, 'grid', locations=boundaries) - pinn.discretise_domain(n, 'grid', locations=['D']) - pinn.train(50, save_loss=i) - assert list(pinn.history_loss.keys()) == truth_key +# def test_train_gpu(): #TODO fix ASAP +# poisson_problem = Poisson() +# boundaries = ['gamma1', 'gamma2', 'gamma3', 'gamma4'] +# n = 10 +# poisson_problem.discretise_domain(n, 'grid', locations=boundaries) +# poisson_problem.conditions.pop('data') # The input/output pts are allocated on cpu +# pinn = PINN(problem = poisson_problem, model=model, extra_features=None, loss=LpLoss()) +# trainer = Trainer(solver=pinn, kwargs={'max_epochs' : 5, 'accelerator':'gpu'}) +# trainer.train() -def test_train_with_optimizer_kwargs(): - boundaries = ['gamma1', 'gamma2', 'gamma3', 'gamma4'] - n = 10 - expected_keys = [[], list(range(0, 50, 3))] - param = [0, 3] - for i, truth_key in zip(param, expected_keys): - pinn = PINN(problem, model, optimizer_kwargs={'lr' : 0.3}) - pinn.discretise_domain(n, 'grid', locations=boundaries) - pinn.discretise_domain(n, 'grid', locations=['D']) - pinn.train(50, save_loss=i) - assert list(pinn.history_loss.keys()) == truth_key +# def test_train_2(): +# boundaries = ['gamma1', 'gamma2', 'gamma3', 'gamma4'] +# n = 10 +# expected_keys = [[], list(range(0, 50, 3))] +# param = [0, 3] +# for i, truth_key in zip(param, expected_keys): +# pinn = PINN(problem, model) +# pinn.discretise_domain(n, 'grid', locations=boundaries) +# pinn.discretise_domain(n, 'grid', locations=['D']) +# pinn.train(50, save_loss=i) +# assert list(pinn.history_loss.keys()) == truth_key -def test_train_with_lr_scheduler(): - boundaries = ['gamma1', 'gamma2', 'gamma3', 'gamma4'] - n = 10 - expected_keys = [[], list(range(0, 50, 3))] - param = [0, 3] - for i, truth_key in zip(param, expected_keys): - pinn = PINN( - problem, - model, - lr_scheduler_type=torch.optim.lr_scheduler.CyclicLR, - lr_scheduler_kwargs={'base_lr' : 0.1, 'max_lr' : 0.3, 'cycle_momentum': False} - ) - pinn.discretise_domain(n, 'grid', locations=boundaries) - pinn.discretise_domain(n, 'grid', locations=['D']) - pinn.train(50, save_loss=i) - assert list(pinn.history_loss.keys()) == truth_key - - -# def test_train_batch(): -# pinn = PINN(problem, model, batch_size=6) +# def test_train_extra_feats(): +# pinn = PINN(problem, model_extra_feat, [myFeature()]) # boundaries = ['gamma1', 'gamma2', 'gamma3', 'gamma4'] # n = 10 # pinn.discretise_domain(n, 'grid', locations=boundaries) @@ -283,34 +347,87 @@ def test_train_with_lr_scheduler(): # pinn.train(5) -# def test_train_batch_2(): +# def test_train_2_extra_feats(): +# boundaries = ['gamma1', 'gamma2', 'gamma3', 'gamma4'] +# n = 10 +# expected_keys = [[], list(range(0, 50, 3))] +# param = [0, 3] +# for i, truth_key in zip(param, expected_keys): +# pinn = PINN(problem, model_extra_feat, [myFeature()]) +# pinn.discretise_domain(n, 'grid', locations=boundaries) +# pinn.discretise_domain(n, 'grid', locations=['D']) +# pinn.train(50, save_loss=i) +# assert list(pinn.history_loss.keys()) == truth_key + + +# def test_train_with_optimizer_kwargs(): +# boundaries = ['gamma1', 'gamma2', 'gamma3', 'gamma4'] +# n = 10 +# expected_keys = [[], list(range(0, 50, 3))] +# param = [0, 3] +# for i, truth_key in zip(param, expected_keys): +# pinn = PINN(problem, model, optimizer_kwargs={'lr' : 0.3}) +# pinn.discretise_domain(n, 'grid', locations=boundaries) +# pinn.discretise_domain(n, 'grid', locations=['D']) +# pinn.train(50, save_loss=i) +# assert list(pinn.history_loss.keys()) == truth_key + + +# def test_train_with_lr_scheduler(): # boundaries = ['gamma1', 'gamma2', 'gamma3', 'gamma4'] # n = 10 # expected_keys = [[], list(range(0, 50, 3))] # param = [0, 3] # for i, truth_key in zip(param, expected_keys): -# pinn = PINN(problem, model, batch_size=6) +# pinn = PINN( +# problem, +# model, +# lr_scheduler_type=torch.optim.lr_scheduler.CyclicLR, +# lr_scheduler_kwargs={'base_lr' : 0.1, 'max_lr' : 0.3, 'cycle_momentum': False} +# ) # pinn.discretise_domain(n, 'grid', locations=boundaries) # pinn.discretise_domain(n, 'grid', locations=['D']) # pinn.train(50, save_loss=i) # assert list(pinn.history_loss.keys()) == truth_key -if torch.cuda.is_available(): - - # def test_gpu_train(): - # pinn = PINN(problem, model, batch_size=20, device='cuda') - # boundaries = ['gamma1', 'gamma2', 'gamma3', 'gamma4'] - # n = 100 - # pinn.discretise_domain(n, 'grid', locations=boundaries) - # pinn.discretise_domain(n, 'grid', locations=['D']) - # pinn.train(5) - - def test_gpu_train_nobatch(): - pinn = PINN(problem, model, batch_size=None, device='cuda') - boundaries = ['gamma1', 'gamma2', 'gamma3', 'gamma4'] - n = 100 - pinn.discretise_domain(n, 'grid', locations=boundaries) - pinn.discretise_domain(n, 'grid', locations=['D']) - pinn.train(5) -""" +# # def test_train_batch(): +# # pinn = PINN(problem, model, batch_size=6) +# # boundaries = ['gamma1', 'gamma2', 'gamma3', 'gamma4'] +# # n = 10 +# # pinn.discretise_domain(n, 'grid', locations=boundaries) +# # pinn.discretise_domain(n, 'grid', locations=['D']) +# # pinn.train(5) + + +# # def test_train_batch_2(): +# # boundaries = ['gamma1', 'gamma2', 'gamma3', 'gamma4'] +# # n = 10 +# # expected_keys = [[], list(range(0, 50, 3))] +# # param = [0, 3] +# # for i, truth_key in zip(param, expected_keys): +# # pinn = PINN(problem, model, batch_size=6) +# # pinn.discretise_domain(n, 'grid', locations=boundaries) +# # pinn.discretise_domain(n, 'grid', locations=['D']) +# # pinn.train(50, save_loss=i) +# # assert list(pinn.history_loss.keys()) == truth_key + + +# if torch.cuda.is_available(): + +# # def test_gpu_train(): +# # pinn = PINN(problem, model, batch_size=20, device='cuda') +# # boundaries = ['gamma1', 'gamma2', 'gamma3', 'gamma4'] +# # n = 100 +# # pinn.discretise_domain(n, 'grid', locations=boundaries) +# # pinn.discretise_domain(n, 'grid', locations=['D']) +# # pinn.train(5) + +# def test_gpu_train_nobatch(): +# pinn = PINN(problem, model, batch_size=None, device='cuda') +# boundaries = ['gamma1', 'gamma2', 'gamma3', 'gamma4'] +# n = 100 +# pinn.discretise_domain(n, 'grid', locations=boundaries) +# pinn.discretise_domain(n, 'grid', locations=['D']) +# pinn.train(5) + diff --git a/tests/test_solvers/test_rom_solver.py b/tests/test_solvers/test_rom_solver.py new file mode 100644 index 000000000..a16ffcaae --- /dev/null +++ b/tests/test_solvers/test_rom_solver.py @@ -0,0 +1,105 @@ +import torch +import pytest + +from pina.problem import AbstractProblem +from pina import Condition, LabelTensor +from pina.solvers import ReducedOrderModelSolver +from pina.trainer import Trainer +from pina.model import FeedForward +from pina.loss import LpLoss + + +class NeuralOperatorProblem(AbstractProblem): + input_variables = ['u_0', 'u_1'] + output_variables = [f'u_{i}' for i in range(100)] + conditions = {'data' : Condition(input_points= + LabelTensor(torch.rand(10, 2), + input_variables), + output_points= + LabelTensor(torch.rand(10, 100), + output_variables))} + + +# make the problem + extra feats +class AE(torch.nn.Module): + def __init__(self, input_dimensions, rank): + super().__init__() + self.encode = FeedForward(input_dimensions, rank, layers=[input_dimensions//4]) + self.decode = FeedForward(rank, input_dimensions, layers=[input_dimensions//4]) +class AE_missing_encode(torch.nn.Module): + def __init__(self, input_dimensions, rank): + super().__init__() + self.encode = FeedForward(input_dimensions, rank, layers=[input_dimensions//4]) +class AE_missing_decode(torch.nn.Module): + def __init__(self, input_dimensions, rank): + super().__init__() + self.decode = FeedForward(rank, input_dimensions, layers=[input_dimensions//4]) + +rank = 10 +problem = NeuralOperatorProblem() +interpolation_net = FeedForward(len(problem.input_variables), + rank) +reduction_net = AE(len(problem.output_variables), rank) + +def test_constructor(): + ReducedOrderModelSolver(problem=problem,reduction_network=reduction_net, + interpolation_network=interpolation_net) + with pytest.raises(SyntaxError): + ReducedOrderModelSolver(problem=problem, + reduction_network=AE_missing_encode( + len(problem.output_variables), rank), + interpolation_network=interpolation_net) + ReducedOrderModelSolver(problem=problem, + reduction_network=AE_missing_decode( + len(problem.output_variables), rank), + interpolation_network=interpolation_net) + + +def test_train_cpu(): + solver = ReducedOrderModelSolver(problem = problem,reduction_network=reduction_net, + interpolation_network=interpolation_net, loss=LpLoss()) + trainer = Trainer(solver=solver, max_epochs=3, accelerator='cpu', batch_size=20) + trainer.train() + + +def test_train_restore(): + tmpdir = "tests/tmp_restore" + solver = ReducedOrderModelSolver(problem=problem, + reduction_network=reduction_net, + interpolation_network=interpolation_net, + 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 = ReducedOrderModelSolver(problem=problem, + reduction_network=reduction_net, + interpolation_network=interpolation_net, + loss=LpLoss()) + trainer = Trainer(solver=solver, + max_epochs=15, + accelerator='cpu', + default_root_dir=tmpdir) + trainer.train() + new_solver = ReducedOrderModelSolver.load_from_checkpoint( + f'{tmpdir}/lightning_logs/version_0/checkpoints/epoch=14-step=15.ckpt', + problem = problem,reduction_network=reduction_net, + interpolation_network=interpolation_net) + test_pts = LabelTensor(torch.rand(20, 2), problem.input_variables) + assert new_solver.forward(test_pts).shape == (20, 100) + 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) \ No newline at end of file diff --git a/tests/test_solvers/test_sapinn.py b/tests/test_solvers/test_sapinn.py new file mode 100644 index 000000000..60c3094ce --- /dev/null +++ b/tests/test_solvers/test_sapinn.py @@ -0,0 +1,437 @@ +import torch +import pytest + +from pina.problem import SpatialProblem, InverseProblem +from pina.operators import laplacian +from pina.geometry import CartesianDomain +from pina import Condition, LabelTensor +from pina.solvers import SAPINN as PINN +from pina.trainer import Trainer +from pina.model import FeedForward +from pina.equation.equation import Equation +from pina.equation.equation_factory import FixedValue +from pina.loss import LpLoss + + +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) +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 InversePoisson(SpatialProblem, InverseProblem): + ''' + Problem definition for the Poisson equation. + ''' + output_variables = ['u'] + x_min = -2 + x_max = 2 + y_min = -2 + y_max = 2 + data_input = LabelTensor(torch.rand(10, 2), ['x', 'y']) + data_output = LabelTensor(torch.rand(10, 1), ['u']) + spatial_domain = CartesianDomain({'x': [x_min, x_max], 'y': [y_min, y_max]}) + # define the ranges for the parameters + unknown_parameter_domain = CartesianDomain({'mu1': [-1, 1], 'mu2': [-1, 1]}) + + def laplace_equation(input_, output_, params_): + ''' + Laplace equation with a force term. + ''' + force_term = torch.exp( + - 2*(input_.extract(['x']) - params_['mu1'])**2 + - 2*(input_.extract(['y']) - params_['mu2'])**2) + delta_u = laplacian(output_, input_, components=['u'], d=['x', 'y']) + + return delta_u - force_term + + # define the conditions for the loss (boundary conditions, equation, data) + conditions = { + 'gamma1': Condition(location=CartesianDomain({'x': [x_min, x_max], + 'y': y_max}), + equation=FixedValue(0.0, components=['u'])), + 'gamma2': Condition(location=CartesianDomain( + {'x': [x_min, x_max], 'y': y_min + }), + equation=FixedValue(0.0, components=['u'])), + 'gamma3': Condition(location=CartesianDomain( + {'x': x_max, 'y': [y_min, y_max] + }), + equation=FixedValue(0.0, components=['u'])), + 'gamma4': Condition(location=CartesianDomain( + {'x': x_min, 'y': [y_min, y_max] + }), + equation=FixedValue(0.0, components=['u'])), + 'D': Condition(location=CartesianDomain( + {'x': [x_min, x_max], 'y': [y_min, y_max] + }), + equation=Equation(laplace_equation)), + 'data': Condition(input_points=data_input.extract(['x', 'y']), + output_points=data_output) + } + + +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_) + } + + 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 + + +class myFeature(torch.nn.Module): + """ + Feature: sin(x) + """ + + def __init__(self): + super(myFeature, self).__init__() + + def forward(self, x): + t = (torch.sin(x.extract(['x']) * torch.pi) * + torch.sin(x.extract(['y']) * torch.pi)) + return LabelTensor(t, ['sin(x)sin(y)']) + + +# make the problem +poisson_problem = Poisson() +model = FeedForward(len(poisson_problem.input_variables), + len(poisson_problem.output_variables)) +model_extra_feats = FeedForward( + len(poisson_problem.input_variables) + 1, + len(poisson_problem.output_variables)) +extra_feats = [myFeature()] + + +def test_constructor(): + PINN(problem=poisson_problem, model=model, extra_features=None) + with pytest.raises(ValueError): + PINN(problem=poisson_problem, model=model, extra_features=None, + weights_function=1) + + +def test_constructor_extra_feats(): + model_extra_feats = FeedForward( + len(poisson_problem.input_variables) + 1, + len(poisson_problem.output_variables)) + PINN(problem=poisson_problem, + model=model_extra_feats, + extra_features=extra_feats) + + +def test_train_cpu(): + poisson_problem = Poisson() + boundaries = ['gamma1', 'gamma2', 'gamma3', 'gamma4'] + n = 10 + poisson_problem.discretise_domain(n, 'grid', locations=boundaries) + pinn = PINN(problem = poisson_problem, model=model, + extra_features=None, loss=LpLoss()) + trainer = Trainer(solver=pinn, max_epochs=1, + accelerator='cpu', batch_size=20) + trainer.train() + + +def test_train_restore(): + tmpdir = "tests/tmp_restore" + poisson_problem = Poisson() + boundaries = ['gamma1', 'gamma2', 'gamma3', 'gamma4'] + n = 10 + poisson_problem.discretise_domain(n, 'grid', locations=boundaries) + pinn = PINN(problem=poisson_problem, + model=model, + extra_features=None, + loss=LpLoss()) + trainer = Trainer(solver=pinn, + max_epochs=5, + accelerator='cpu', + default_root_dir=tmpdir) + trainer.train() + ntrainer = Trainer(solver=pinn, max_epochs=15, accelerator='cpu') + t = ntrainer.train( + ckpt_path=f'{tmpdir}/lightning_logs/version_0/' + 'checkpoints/epoch=4-step=10.ckpt') + import shutil + shutil.rmtree(tmpdir) + + +def test_train_load(): + tmpdir = "tests/tmp_load" + poisson_problem = Poisson() + boundaries = ['gamma1', 'gamma2', 'gamma3', 'gamma4'] + n = 10 + poisson_problem.discretise_domain(n, 'grid', locations=boundaries) + pinn = PINN(problem=poisson_problem, + model=model, + extra_features=None, + loss=LpLoss()) + trainer = Trainer(solver=pinn, + max_epochs=15, + accelerator='cpu', + default_root_dir=tmpdir) + trainer.train() + new_pinn = PINN.load_from_checkpoint( + f'{tmpdir}/lightning_logs/version_0/checkpoints/epoch=14-step=30.ckpt', + problem = poisson_problem, model=model) + test_pts = CartesianDomain({'x': [0, 1], 'y': [0, 1]}).sample(10) + assert new_pinn.forward(test_pts).extract(['u']).shape == (10, 1) + assert new_pinn.forward(test_pts).extract( + ['u']).shape == pinn.forward(test_pts).extract(['u']).shape + torch.testing.assert_close( + new_pinn.forward(test_pts).extract(['u']), + pinn.forward(test_pts).extract(['u'])) + import shutil + shutil.rmtree(tmpdir) + +def test_train_inverse_problem_cpu(): + poisson_problem = InversePoisson() + boundaries = ['gamma1', 'gamma2', 'gamma3', 'gamma4', 'D'] + n = 100 + poisson_problem.discretise_domain(n, 'random', locations=boundaries) + pinn = PINN(problem = poisson_problem, model=model, + extra_features=None, loss=LpLoss()) + trainer = Trainer(solver=pinn, max_epochs=1, + accelerator='cpu', batch_size=20) + trainer.train() + + +# # TODO does not currently work +# def test_train_inverse_problem_restore(): +# tmpdir = "tests/tmp_restore_inv" +# poisson_problem = InversePoisson() +# boundaries = ['gamma1', 'gamma2', 'gamma3', 'gamma4', 'D'] +# n = 100 +# poisson_problem.discretise_domain(n, 'random', locations=boundaries) +# pinn = PINN(problem=poisson_problem, +# model=model, +# extra_features=None, +# loss=LpLoss()) +# trainer = Trainer(solver=pinn, +# max_epochs=5, +# accelerator='cpu', +# default_root_dir=tmpdir) +# trainer.train() +# ntrainer = Trainer(solver=pinn, max_epochs=5, accelerator='cpu') +# t = ntrainer.train( +# ckpt_path=f'{tmpdir}/lightning_logs/version_0/checkpoints/epoch=4-step=10.ckpt') +# import shutil +# shutil.rmtree(tmpdir) + + +def test_train_inverse_problem_load(): + tmpdir = "tests/tmp_load_inv" + poisson_problem = InversePoisson() + boundaries = ['gamma1', 'gamma2', 'gamma3', 'gamma4', 'D'] + n = 100 + poisson_problem.discretise_domain(n, 'random', locations=boundaries) + pinn = PINN(problem=poisson_problem, + model=model, + extra_features=None, + loss=LpLoss()) + trainer = Trainer(solver=pinn, + max_epochs=15, + accelerator='cpu', + default_root_dir=tmpdir) + trainer.train() + new_pinn = PINN.load_from_checkpoint( + f'{tmpdir}/lightning_logs/version_0/checkpoints/epoch=14-step=30.ckpt', + problem = poisson_problem, model=model) + test_pts = CartesianDomain({'x': [0, 1], 'y': [0, 1]}).sample(10) + assert new_pinn.forward(test_pts).extract(['u']).shape == (10, 1) + assert new_pinn.forward(test_pts).extract( + ['u']).shape == pinn.forward(test_pts).extract(['u']).shape + torch.testing.assert_close( + new_pinn.forward(test_pts).extract(['u']), + pinn.forward(test_pts).extract(['u'])) + import shutil + shutil.rmtree(tmpdir) + +# # TODO fix asap. Basically sampling few variables +# # works only if both variables are in a range. +# # if one is fixed and the other not, this will +# # not work. This test also needs to be fixed and +# # insert in test problem not in test pinn. +# def test_train_cpu_sampling_few_vars(): +# poisson_problem = Poisson() +# boundaries = ['gamma1', 'gamma2', 'gamma3'] +# n = 10 +# poisson_problem.discretise_domain(n, 'grid', locations=boundaries) +# poisson_problem.discretise_domain(n, 'random', locations=['gamma4'], variables=['x']) +# poisson_problem.discretise_domain(n, 'random', locations=['gamma4'], variables=['y']) +# pinn = PINN(problem = poisson_problem, model=model, extra_features=None, loss=LpLoss()) +# trainer = Trainer(solver=pinn, kwargs={'max_epochs' : 5, 'accelerator':'cpu'}) +# trainer.train() + + +def test_train_extra_feats_cpu(): + poisson_problem = Poisson() + boundaries = ['gamma1', 'gamma2', 'gamma3', 'gamma4'] + n = 10 + poisson_problem.discretise_domain(n, 'grid', locations=boundaries) + pinn = PINN(problem=poisson_problem, + model=model_extra_feats, + extra_features=extra_feats) + trainer = Trainer(solver=pinn, max_epochs=5, accelerator='cpu') + trainer.train() + + +# TODO, fix GitHub actions to run also on GPU +# def test_train_gpu(): +# poisson_problem = Poisson() +# boundaries = ['gamma1', 'gamma2', 'gamma3', 'gamma4'] +# n = 10 +# poisson_problem.discretise_domain(n, 'grid', locations=boundaries) +# pinn = PINN(problem = poisson_problem, model=model, extra_features=None, loss=LpLoss()) +# trainer = Trainer(solver=pinn, kwargs={'max_epochs' : 5, 'accelerator':'gpu'}) +# trainer.train() + +# def test_train_gpu(): #TODO fix ASAP +# poisson_problem = Poisson() +# boundaries = ['gamma1', 'gamma2', 'gamma3', 'gamma4'] +# n = 10 +# poisson_problem.discretise_domain(n, 'grid', locations=boundaries) +# poisson_problem.conditions.pop('data') # The input/output pts are allocated on cpu +# pinn = PINN(problem = poisson_problem, model=model, extra_features=None, loss=LpLoss()) +# trainer = Trainer(solver=pinn, kwargs={'max_epochs' : 5, 'accelerator':'gpu'}) +# trainer.train() + +# def test_train_2(): +# boundaries = ['gamma1', 'gamma2', 'gamma3', 'gamma4'] +# n = 10 +# expected_keys = [[], list(range(0, 50, 3))] +# param = [0, 3] +# for i, truth_key in zip(param, expected_keys): +# pinn = PINN(problem, model) +# pinn.discretise_domain(n, 'grid', locations=boundaries) +# pinn.discretise_domain(n, 'grid', locations=['D']) +# pinn.train(50, save_loss=i) +# assert list(pinn.history_loss.keys()) == truth_key + + +# def test_train_extra_feats(): +# pinn = PINN(problem, model_extra_feat, [myFeature()]) +# boundaries = ['gamma1', 'gamma2', 'gamma3', 'gamma4'] +# n = 10 +# pinn.discretise_domain(n, 'grid', locations=boundaries) +# pinn.discretise_domain(n, 'grid', locations=['D']) +# pinn.train(5) + + +# def test_train_2_extra_feats(): +# boundaries = ['gamma1', 'gamma2', 'gamma3', 'gamma4'] +# n = 10 +# expected_keys = [[], list(range(0, 50, 3))] +# param = [0, 3] +# for i, truth_key in zip(param, expected_keys): +# pinn = PINN(problem, model_extra_feat, [myFeature()]) +# pinn.discretise_domain(n, 'grid', locations=boundaries) +# pinn.discretise_domain(n, 'grid', locations=['D']) +# pinn.train(50, save_loss=i) +# assert list(pinn.history_loss.keys()) == truth_key + + +# def test_train_with_optimizer_kwargs(): +# boundaries = ['gamma1', 'gamma2', 'gamma3', 'gamma4'] +# n = 10 +# expected_keys = [[], list(range(0, 50, 3))] +# param = [0, 3] +# for i, truth_key in zip(param, expected_keys): +# pinn = PINN(problem, model, optimizer_kwargs={'lr' : 0.3}) +# pinn.discretise_domain(n, 'grid', locations=boundaries) +# pinn.discretise_domain(n, 'grid', locations=['D']) +# pinn.train(50, save_loss=i) +# assert list(pinn.history_loss.keys()) == truth_key + + +# def test_train_with_lr_scheduler(): +# boundaries = ['gamma1', 'gamma2', 'gamma3', 'gamma4'] +# n = 10 +# expected_keys = [[], list(range(0, 50, 3))] +# param = [0, 3] +# for i, truth_key in zip(param, expected_keys): +# pinn = PINN( +# problem, +# model, +# lr_scheduler_type=torch.optim.lr_scheduler.CyclicLR, +# lr_scheduler_kwargs={'base_lr' : 0.1, 'max_lr' : 0.3, 'cycle_momentum': False} +# ) +# pinn.discretise_domain(n, 'grid', locations=boundaries) +# pinn.discretise_domain(n, 'grid', locations=['D']) +# pinn.train(50, save_loss=i) +# assert list(pinn.history_loss.keys()) == truth_key + + +# # def test_train_batch(): +# # pinn = PINN(problem, model, batch_size=6) +# # boundaries = ['gamma1', 'gamma2', 'gamma3', 'gamma4'] +# # n = 10 +# # pinn.discretise_domain(n, 'grid', locations=boundaries) +# # pinn.discretise_domain(n, 'grid', locations=['D']) +# # pinn.train(5) + + +# # def test_train_batch_2(): +# # boundaries = ['gamma1', 'gamma2', 'gamma3', 'gamma4'] +# # n = 10 +# # expected_keys = [[], list(range(0, 50, 3))] +# # param = [0, 3] +# # for i, truth_key in zip(param, expected_keys): +# # pinn = PINN(problem, model, batch_size=6) +# # pinn.discretise_domain(n, 'grid', locations=boundaries) +# # pinn.discretise_domain(n, 'grid', locations=['D']) +# # pinn.train(50, save_loss=i) +# # assert list(pinn.history_loss.keys()) == truth_key + + +# if torch.cuda.is_available(): + +# # def test_gpu_train(): +# # pinn = PINN(problem, model, batch_size=20, device='cuda') +# # boundaries = ['gamma1', 'gamma2', 'gamma3', 'gamma4'] +# # n = 100 +# # pinn.discretise_domain(n, 'grid', locations=boundaries) +# # pinn.discretise_domain(n, 'grid', locations=['D']) +# # pinn.train(5) + +# def test_gpu_train_nobatch(): +# pinn = PINN(problem, model, batch_size=None, device='cuda') +# boundaries = ['gamma1', 'gamma2', 'gamma3', 'gamma4'] +# n = 100 +# pinn.discretise_domain(n, 'grid', locations=boundaries) +# pinn.discretise_domain(n, 'grid', locations=['D']) +# pinn.train(5) + From 756241daae8e20403abcc39e672edecb3835d901 Mon Sep 17 00:00:00 2001 From: "github-actions[bot]" <41898282+github-actions[bot]@users.noreply.github.com> Date: Fri, 10 May 2024 14:08:01 +0200 Subject: [PATCH 08/14] :art: Format Python code with psf/black (#297) Co-authored-by: dario-coscia --- pina/solvers/__init__.py | 3 +- pina/solvers/pinns/basepinn.py | 33 +++---- pina/solvers/pinns/causalpinn.py | 34 +++---- pina/solvers/pinns/competitive_pinn.py | 39 +++++---- pina/solvers/pinns/gpinn.py | 31 +++---- pina/solvers/pinns/pinn.py | 7 +- pina/solvers/pinns/sapinn.py | 117 ++++++++++++------------- pina/solvers/rom.py | 51 ++++++----- pina/solvers/solver.py | 4 +- pina/solvers/supervised.py | 4 +- pina/trainer.py | 6 +- 11 files changed, 169 insertions(+), 160 deletions(-) diff --git a/pina/solvers/__init__.py b/pina/solvers/__init__.py index 2751e481c..6b7556614 100644 --- a/pina/solvers/__init__.py +++ b/pina/solvers/__init__.py @@ -9,11 +9,10 @@ "SupervisedSolver", "ReducedOrderModelSolver", "GAROM", - ] +] from .solver import SolverInterface from .pinns import * from .supervised import SupervisedSolver from .rom import ReducedOrderModelSolver from .garom import GAROM - diff --git a/pina/solvers/pinns/basepinn.py b/pina/solvers/pinns/basepinn.py index 726cdf921..53d4d3a90 100644 --- a/pina/solvers/pinns/basepinn.py +++ b/pina/solvers/pinns/basepinn.py @@ -12,11 +12,12 @@ torch.pi = torch.acos(torch.zeros(1)).item() * 2 # which is 3.1415927410125732 + class PINNInterface(SolverInterface, metaclass=ABCMeta): """ Base PINN solver class. This class implements the Solver Interface for Physics Informed Neural Network solvers. - + This class can be used to define PINNs with multiple ``optimizers``, and/or ``models``. By default it takes @@ -72,7 +73,7 @@ def __init__( self._clamp_params = self._clamp_inverse_problem_params else: self._params = None - self._clamp_params = lambda : None + self._clamp_params = lambda: None # variable used internally to store residual losses at each epoch # this variable save the residual at each iteration (not weighted) @@ -107,7 +108,7 @@ def training_step(self, batch, _): condition = self.problem.conditions[condition_name] pts = batch["pts"] # condition name is logged (if logs enabled) - self.__logged_metric = condition_name + self.__logged_metric = condition_name if len(batch) == 2: samples = pts[condition_idx == condition_id] @@ -160,7 +161,7 @@ def loss_phys(self, samples, equation): :rtype: LabelTensor """ pass - + def compute_residual(self, samples, equation): """ Compute the residual for Physics Informed learning. This function @@ -182,7 +183,7 @@ def compute_residual(self, samples, equation): samples, self.forward(samples), self._params ) return residual - + def store_log(self, loss_value): """ Stores the loss value in the logger. This function should be @@ -195,13 +196,13 @@ def store_log(self, loss_value): :param torch.Tensor loss_value: The value of the loss. """ self.log( - self.__logged_metric+'_loss', - loss_value, - prog_bar=True, - logger=True, - on_epoch=True, - on_step=False, - ) + self.__logged_metric + "_loss", + loss_value, + prog_bar=True, + logger=True, + on_epoch=True, + on_step=False, + ) self.__logged_res_losses.append(loss_value) def on_train_epoch_end(self): @@ -211,10 +212,10 @@ def on_train_epoch_end(self): """ if self.__logged_res_losses: # storing mean loss - self.__logged_metric = 'mean' + self.__logged_metric = "mean" self.store_log( - sum(self.__logged_res_losses)/len(self.__logged_res_losses) - ) + sum(self.__logged_res_losses) / len(self.__logged_res_losses) + ) # free the logged losses self.__logged_res_losses = [] return super().on_train_epoch_end() @@ -244,4 +245,4 @@ def current_condition_name(self): :meth:`loss_phys` to extract the condition at which the loss is computed. """ - return self.__logged_metric \ No newline at end of file + return self.__logged_metric diff --git a/pina/solvers/pinns/causalpinn.py b/pina/solvers/pinns/causalpinn.py index fea0fe47b..476e4c55c 100644 --- a/pina/solvers/pinns/causalpinn.py +++ b/pina/solvers/pinns/causalpinn.py @@ -97,25 +97,27 @@ def __init__( :param dict scheduler_kwargs: LR scheduler constructor keyword args. :param int | float eps: The exponential decay parameter. Note that this value is kept fixed during the training, but can be changed by means - of a callback, e.g. for annealing. + of a callback, e.g. for annealing. """ super().__init__( - problem=problem, - model=model, - extra_features=extra_features, - loss=loss, - optimizer=optimizer, - optimizer_kwargs=optimizer_kwargs, - scheduler=scheduler, - scheduler_kwargs=scheduler_kwargs, + problem=problem, + model=model, + extra_features=extra_features, + loss=loss, + optimizer=optimizer, + optimizer_kwargs=optimizer_kwargs, + scheduler=scheduler, + scheduler_kwargs=scheduler_kwargs, ) # checking consistency - check_consistency(eps, (int,float)) + check_consistency(eps, (int, float)) self._eps = eps if not isinstance(self.problem, TimeDependentProblem): - raise ValueError('Casual PINN works only for problems' - 'inheritig from TimeDependentProblem.') + raise ValueError( + "Casual PINN works only for problems" + "inheritig from TimeDependentProblem." + ) def loss_phys(self, samples, equation): """ @@ -144,14 +146,14 @@ def loss_phys(self, samples, equation): ) time_loss.append(loss_val) # store results - self.store_log(loss_value=float(sum(time_loss)/len(time_loss))) + self.store_log(loss_value=float(sum(time_loss) / len(time_loss))) # concatenate residuals time_loss = torch.stack(time_loss) # compute weights (without the gradient storing) with torch.no_grad(): weights = self._compute_weights(time_loss) return (weights * time_loss).mean() - + @property def eps(self): """ @@ -205,8 +207,8 @@ def _split_tensor_into_chunks(self, tensor): _, idx_split = time_tensor.unique(return_counts=True) # splitting chunks = torch.split(tensor, tuple(idx_split)) - return chunks, labels # return chunks - + return chunks, labels # return chunks + def _compute_weights(self, loss): """ Computes the weights for the physics loss based on the cumulative loss. diff --git a/pina/solvers/pinns/competitive_pinn.py b/pina/solvers/pinns/competitive_pinn.py index 6404c0bb4..5e011a473 100644 --- a/pina/solvers/pinns/competitive_pinn.py +++ b/pina/solvers/pinns/competitive_pinn.py @@ -117,7 +117,7 @@ def __init__( optimizer_discriminator_kwargs, ], extra_features=None, # CompetitivePINN doesn't take extra features - loss=loss + loss=loss, ) # set automatic optimization for GANs @@ -131,9 +131,7 @@ def __init__( # assign schedulers self._schedulers = [ - scheduler_model( - self.optimizers[0], **scheduler_model_kwargs - ), + scheduler_model(self.optimizers[0], **scheduler_model_kwargs), scheduler_discriminator( self.optimizers[1], **scheduler_discriminator_kwargs ), @@ -141,7 +139,7 @@ def __init__( self._model = self.models[0] self._discriminator = self.models[1] - + def forward(self, x): r""" Forward pass implementation for the PINN solver. It returns the function @@ -195,8 +193,11 @@ def loss_data(self, input_tensor, output_tensor): :rtype: torch.Tensor """ self.optimizer_model.zero_grad() - loss_val = super().loss_data( - input_tensor, output_tensor).as_subclass(torch.Tensor) + loss_val = ( + super() + .loss_data(input_tensor, output_tensor) + .as_subclass(torch.Tensor) + ) loss_val.backward() self.optimizer_model.step() return loss_val @@ -221,7 +222,7 @@ def configure_optimizers(self): ) return self.optimizers, self._schedulers - def on_train_batch_end(self,outputs, batch, batch_idx): + def on_train_batch_end(self, outputs, batch, batch_idx): """ This method is called at the end of each training batch, and ovverides the PytorchLightining implementation for logging the checkpoints. @@ -235,7 +236,9 @@ def on_train_batch_end(self,outputs, batch, batch_idx): :rtype: Any """ # increase by one the counter of optimization to save loggers - self.trainer.fit_loop.epoch_loop.manual_optimization.optim_step_progress.total.completed += 1 + self.trainer.fit_loop.epoch_loop.manual_optimization.optim_step_progress.total.completed += ( + 1 + ) return super().on_train_batch_end(outputs, batch, batch_idx) def _train_discriminator(self, samples, equation, discriminator_bets): @@ -252,13 +255,14 @@ def _train_discriminator(self, samples, equation, discriminator_bets): self.optimizer_discriminator.zero_grad() # compute residual, we detach because the weights of the generator # model are fixed - residual = self.compute_residual(samples=samples, - equation=equation).detach() + residual = self.compute_residual( + samples=samples, equation=equation + ).detach() # compute competitive residual, the minus is because we maximise competitive_residual = residual * discriminator_bets - loss_val = - self.loss( + loss_val = -self.loss( torch.zeros_like(competitive_residual, requires_grad=True), - competitive_residual + competitive_residual, ).as_subclass(torch.Tensor) # backprop self.manual_backward(loss_val) @@ -283,16 +287,13 @@ def _train_model(self, samples, equation, discriminator_bets): residual = self.compute_residual(samples=samples, equation=equation) # store logging with torch.no_grad(): - loss_residual = self.loss( - torch.zeros_like(residual), - residual - ) + loss_residual = self.loss(torch.zeros_like(residual), residual) # compute competitive residual, discriminator_bets are detached becase # we optimize only the generator model competitive_residual = residual * discriminator_bets.detach() loss_val = self.loss( torch.zeros_like(competitive_residual, requires_grad=True), - competitive_residual + competitive_residual, ).as_subclass(torch.Tensor) # backprop self.manual_backward(loss_val) @@ -357,4 +358,4 @@ def scheduler_discriminator(self): :return: The scheduler for the discriminator. :rtype: torch.optim.lr_scheduler._LRScheduler """ - return self._schedulers[1] \ No newline at end of file + return self._schedulers[1] diff --git a/pina/solvers/pinns/gpinn.py b/pina/solvers/pinns/gpinn.py index 6eca1eac5..5f259ca21 100644 --- a/pina/solvers/pinns/gpinn.py +++ b/pina/solvers/pinns/gpinn.py @@ -90,22 +90,23 @@ def __init__( :param dict scheduler_kwargs: LR scheduler constructor keyword args. """ super().__init__( - problem=problem, - model=model, - extra_features=extra_features, - loss=loss, - optimizer=optimizer, - optimizer_kwargs=optimizer_kwargs, - scheduler=scheduler, - scheduler_kwargs=scheduler_kwargs, + problem=problem, + model=model, + extra_features=extra_features, + loss=loss, + optimizer=optimizer, + optimizer_kwargs=optimizer_kwargs, + scheduler=scheduler, + scheduler_kwargs=scheduler_kwargs, ) if not isinstance(self.problem, SpatialProblem): - raise ValueError('Gradient PINN computes the gradient of the ' - 'PINN loss with respect to the spatial ' - 'coordinates, thus the PINA problem must be ' - 'a SpatialProblem.') + raise ValueError( + "Gradient PINN computes the gradient of the " + "PINN loss with respect to the spatial " + "coordinates, thus the PINA problem must be " + "a SpatialProblem." + ) - def loss_phys(self, samples, equation): """ Computes the physics loss for the GPINN solver based on given @@ -126,9 +127,9 @@ def loss_phys(self, samples, equation): self.store_log(loss_value=float(loss_value)) # gradient PINN loss loss_value = loss_value.reshape(-1, 1) - loss_value.labels = ['__LOSS'] + loss_value.labels = ["__LOSS"] loss_grad = grad(loss_value, samples, d=self.problem.spatial_variables) g_loss_phys = self.loss( torch.zeros_like(loss_grad, requires_grad=True), loss_grad ) - return loss_value + g_loss_phys \ No newline at end of file + return loss_value + g_loss_phys diff --git a/pina/solvers/pinns/pinn.py b/pina/solvers/pinns/pinn.py index 318283a38..15f908182 100644 --- a/pina/solvers/pinns/pinn.py +++ b/pina/solvers/pinns/pinn.py @@ -87,7 +87,7 @@ def __init__( optimizers=[optimizer], optimizers_kwargs=[optimizer_kwargs], extra_features=extra_features, - loss=loss + loss=loss, ) # check consistency @@ -131,7 +131,6 @@ def loss_phys(self, samples, equation): self.store_log(loss_value=float(loss_value)) return loss_value - def configure_optimizers(self): """ Optimizer configuration for the PINN @@ -153,7 +152,6 @@ def configure_optimizers(self): ) return self.optimizers, [self.scheduler] - @property def scheduler(self): """ @@ -161,10 +159,9 @@ def scheduler(self): """ return self._scheduler - @property def neural_net(self): """ Neural network for the PINN training. """ - return self._neural_net \ No newline at end of file + return self._neural_net diff --git a/pina/solvers/pinns/sapinn.py b/pina/solvers/pinns/sapinn.py index 8de2d14c1..751e21eff 100644 --- a/pina/solvers/pinns/sapinn.py +++ b/pina/solvers/pinns/sapinn.py @@ -14,6 +14,7 @@ from torch.optim.lr_scheduler import ConstantLR + class Weights(torch.nn.Module): """ This class aims to implements the mask model for @@ -27,11 +28,9 @@ def __init__(self, func): """ super().__init__() check_consistency(func, torch.nn.Module) - self.sa_weights = torch.nn.Parameter( - torch.Tensor() - ) + self.sa_weights = torch.nn.Parameter(torch.Tensor()) self.func = func - + def forward(self): """ Forward pass implementation for the mask module. @@ -43,6 +42,7 @@ def forward(self): """ return self.func(self.sa_weights) + class SAPINN(PINNInterface): r""" Self Adaptive Physics Informed Neural Network (SAPINN) solver class. @@ -106,22 +106,22 @@ class SAPINN(PINNInterface): DOI: `10.1016/ j.jcp.2022.111722 `_. """ - + def __init__( - self, - problem, - model, - weights_function=torch.nn.Sigmoid(), - extra_features=None, - loss=torch.nn.MSELoss(), - optimizer_model=torch.optim.Adam, - optimizer_model_kwargs={"lr" : 0.001}, - optimizer_weights=torch.optim.Adam, - optimizer_weights_kwargs={"lr" : 0.001}, - scheduler_model=ConstantLR, - scheduler_model_kwargs={"factor" : 1, "total_iters" : 0}, - scheduler_weights=ConstantLR, - scheduler_weights_kwargs={"factor" : 1, "total_iters" : 0} + self, + problem, + model, + weights_function=torch.nn.Sigmoid(), + extra_features=None, + loss=torch.nn.MSELoss(), + optimizer_model=torch.optim.Adam, + optimizer_model_kwargs={"lr": 0.001}, + optimizer_weights=torch.optim.Adam, + optimizer_weights_kwargs={"lr": 0.001}, + scheduler_model=ConstantLR, + scheduler_model_kwargs={"factor": 1, "total_iters": 0}, + scheduler_weights=ConstantLR, + scheduler_weights_kwargs={"factor": 1, "total_iters": 0}, ): """ :param AbstractProblem problem: The formualation of the problem. @@ -167,19 +167,18 @@ def __init__( weights_dict[condition_name] = Weights(weights_function) weights_dict = torch.nn.ModuleDict(weights_dict) - super().__init__( models=[model, weights_dict], problem=problem, optimizers=[optimizer_model, optimizer_weights], optimizers_kwargs=[ optimizer_model_kwargs, - optimizer_weights_kwargs + optimizer_weights_kwargs, ], extra_features=extra_features, - loss=loss + loss=loss, ) - + # set automatic optimization self.automatic_optimization = False @@ -191,12 +190,8 @@ def __init__( # assign schedulers self._schedulers = [ - scheduler_model( - self.optimizers[0], **scheduler_model_kwargs - ), - scheduler_weights( - self.optimizers[1], **scheduler_weights_kwargs - ), + scheduler_model(self.optimizers[0], **scheduler_model_kwargs), + scheduler_weights(self.optimizers[1], **scheduler_weights_kwargs), ] self._model = self.models[0] @@ -204,7 +199,7 @@ def __init__( self._vectorial_loss = deepcopy(loss) self._vectorial_loss.reduction = "none" - + def forward(self, x): """ Forward pass implementation for the PINN @@ -219,7 +214,7 @@ def forward(self, x): :rtype: LabelTensor """ return self.neural_net(x) - + def loss_phys(self, samples, equation): """ Computes the physics loss for the SAPINN solver based on given @@ -235,7 +230,7 @@ def loss_phys(self, samples, equation): # train weights self.optimizer_weights.zero_grad() weighted_loss, _ = self._loss_phys(samples, equation) - loss_value = - weighted_loss.as_subclass(torch.Tensor) + loss_value = -weighted_loss.as_subclass(torch.Tensor) self.manual_backward(loss_value) self.optimizer_weights.step() @@ -271,7 +266,7 @@ def loss_data(self, input_tensor, output_tensor): # train weights self.optimizer_weights.zero_grad() weighted_loss, _ = self._loss_data(input_tensor, output_tensor) - loss_value = - weighted_loss.as_subclass(torch.Tensor) + loss_value = -weighted_loss.as_subclass(torch.Tensor) self.manual_backward(loss_value) self.optimizer_weights.step() @@ -291,7 +286,7 @@ def loss_data(self, input_tensor, output_tensor): # store loss without weights self.store_log(loss_value=float(loss)) return loss_value - + def configure_optimizers(self): """ Optimizer configuration for the SAPINN @@ -312,8 +307,8 @@ def configure_optimizers(self): } ) return self.optimizers, self._schedulers - - def on_train_batch_end(self,outputs, batch, batch_idx): + + def on_train_batch_end(self, outputs, batch, batch_idx): """ This method is called at the end of each training batch, and ovverides the PytorchLightining implementation for logging the checkpoints. @@ -327,9 +322,11 @@ def on_train_batch_end(self,outputs, batch, batch_idx): :rtype: Any """ # increase by one the counter of optimization to save loggers - self.trainer.fit_loop.epoch_loop.manual_optimization.optim_step_progress.total.completed += 1 + self.trainer.fit_loop.epoch_loop.manual_optimization.optim_step_progress.total.completed += ( + 1 + ) return super().on_train_batch_end(outputs, batch, batch_idx) - + def on_train_start(self): """ This method is called at the start of the training for setting @@ -343,12 +340,11 @@ def on_train_start(self): self.trainer._accelerator_connector._accelerator_flag ) for condition_name, tensor in self.problem.input_pts.items(): - self.weights_dict.torchmodel[condition_name].sa_weights.data = torch.rand( - (tensor.shape[0], 1), - device = device + self.weights_dict.torchmodel[condition_name].sa_weights.data = ( + torch.rand((tensor.shape[0], 1), device=device) ) return super().on_train_start() - + def on_load_checkpoint(self, checkpoint): """ Overriding the Pytorch Lightning ``on_load_checkpoint`` to handle @@ -358,8 +354,8 @@ def on_load_checkpoint(self, checkpoint): :param dict checkpoint: Pytorch Lightning checkpoint dict. """ for condition_name, tensor in self.problem.input_pts.items(): - self.weights_dict.torchmodel[condition_name].sa_weights.data = torch.rand( - (tensor.shape[0], 1) + self.weights_dict.torchmodel[condition_name].sa_weights.data = ( + torch.rand((tensor.shape[0], 1)) ) return super().on_load_checkpoint(checkpoint) @@ -370,13 +366,13 @@ def _loss_phys(self, samples, equation): :param LabelTensor samples: Input samples to evaluate the physics loss. :param EquationInterface equation: the governing equation representing the physics. - + :return: tuple with weighted and not weighted scalar loss :rtype: List[LabelTensor, LabelTensor] """ residual = self.compute_residual(samples, equation) return self._compute_loss(residual) - + def _loss_data(self, input_tensor, output_tensor): """ Elaboration of the loss related to data for the SAPINN solver. @@ -384,7 +380,7 @@ def _loss_data(self, input_tensor, output_tensor): :param LabelTensor input_tensor: The input to the neural networks. :param LabelTensor output_tensor: The true solution to compare the network solution. - + :return: tuple with weighted and not weighted scalar loss :rtype: List[LabelTensor, LabelTensor] """ @@ -396,19 +392,21 @@ def _compute_loss(self, residual): Elaboration of the pointwise loss through the mask model and the self adaptive weights - :param LabelTensor residual: the matrix of residuals that have to + :param LabelTensor residual: the matrix of residuals that have to be weighted :return: tuple with weighted and not weighted loss :rtype List[LabelTensor, LabelTensor] """ weights = self.weights_dict.torchmodel[ - self.current_condition_name].forward() - loss_value = self._vectorial_loss(torch.zeros_like( - residual, requires_grad=True), residual) + self.current_condition_name + ].forward() + loss_value = self._vectorial_loss( + torch.zeros_like(residual, requires_grad=True), residual + ) return ( self._vect_to_scalar(weights * loss_value), - self._vect_to_scalar(loss_value) + self._vect_to_scalar(loss_value), ) def _vect_to_scalar(self, loss_value): @@ -426,10 +424,11 @@ def _vect_to_scalar(self, loss_value): elif self.loss.reduction == "sum": ret = torch.sum(loss_value) else: - raise RuntimeError(f"Invalid reduction, got {self.loss.reduction} " - "but expected mean or sum.") + raise RuntimeError( + f"Invalid reduction, got {self.loss.reduction} " + "but expected mean or sum." + ) return ret - @property def neural_net(self): @@ -440,7 +439,7 @@ def neural_net(self): :rtype: torch.nn.Module """ return self.models[0] - + @property def weights_dict(self): """ @@ -462,7 +461,7 @@ def scheduler_model(self): :rtype: torch.optim.lr_scheduler._LRScheduler """ return self._scheduler[0] - + @property def scheduler_weights(self): """ @@ -482,7 +481,7 @@ def optimizer_model(self): :rtype: torch.optim.Optimizer """ return self.optimizers[0] - + @property def optimizer_weights(self): """ @@ -491,4 +490,4 @@ def optimizer_weights(self): :return: The optimizer for the mask model. :rtype: torch.optim.Optimizer """ - return self.optimizers[1] \ No newline at end of file + return self.optimizers[1] diff --git a/pina/solvers/rom.py b/pina/solvers/rom.py index 733d76f42..ee4bcff43 100644 --- a/pina/solvers/rom.py +++ b/pina/solvers/rom.py @@ -4,6 +4,7 @@ from pina.solvers import SupervisedSolver + class ReducedOrderModelSolver(SupervisedSolver): r""" ReducedOrderModelSolver solver class. This class implements a @@ -114,10 +115,13 @@ def __init__( rate scheduler. :param dict scheduler_kwargs: LR scheduler constructor keyword args. """ - model = torch.nn.ModuleDict({ - 'reduction_network' : reduction_network, - 'interpolation_network' : interpolation_network}) - + model = torch.nn.ModuleDict( + { + "reduction_network": reduction_network, + "interpolation_network": interpolation_network, + } + ) + super().__init__( model=model, problem=problem, @@ -125,18 +129,22 @@ def __init__( optimizer=optimizer, optimizer_kwargs=optimizer_kwargs, scheduler=scheduler, - scheduler_kwargs=scheduler_kwargs + scheduler_kwargs=scheduler_kwargs, ) # assert reduction object contains encode/ decode - if not hasattr(self.neural_net['reduction_network'], 'encode'): - raise SyntaxError('reduction_network must have encode method. ' - 'The encode method should return a lower ' - 'dimensional representation of the input.') - if not hasattr(self.neural_net['reduction_network'], 'decode'): - raise SyntaxError('reduction_network must have decode method. ' - 'The decode method should return a high ' - 'dimensional representation of the encoding.') + if not hasattr(self.neural_net["reduction_network"], "encode"): + raise SyntaxError( + "reduction_network must have encode method. " + "The encode method should return a lower " + "dimensional representation of the input." + ) + if not hasattr(self.neural_net["reduction_network"], "decode"): + raise SyntaxError( + "reduction_network must have decode method. " + "The decode method should return a high " + "dimensional representation of the encoding." + ) def forward(self, x): """ @@ -149,8 +157,8 @@ def forward(self, x): :return: Solver solution. :rtype: torch.Tensor """ - reduction_network = self.neural_net['reduction_network'] - interpolation_network = self.neural_net['interpolation_network'] + reduction_network = self.neural_net["reduction_network"] + interpolation_network = self.neural_net["interpolation_network"] return reduction_network.decode(interpolation_network(x)) def loss_data(self, input_pts, output_pts): @@ -167,17 +175,18 @@ def loss_data(self, input_pts, output_pts): :rtype: torch.Tensor """ # extract networks - reduction_network = self.neural_net['reduction_network'] - interpolation_network = self.neural_net['interpolation_network'] + reduction_network = self.neural_net["reduction_network"] + interpolation_network = self.neural_net["interpolation_network"] # encoded representations loss encode_repr_inter_net = interpolation_network(input_pts) encode_repr_reduction_network = reduction_network.encode(output_pts) - loss_encode = self.loss(encode_repr_inter_net, - encode_repr_reduction_network) + loss_encode = self.loss( + encode_repr_inter_net, encode_repr_reduction_network + ) # reconstruction loss loss_reconstruction = self.loss( - reduction_network.decode(encode_repr_reduction_network), - output_pts) + reduction_network.decode(encode_repr_reduction_network), output_pts + ) return loss_encode + loss_reconstruction diff --git a/pina/solvers/solver.py b/pina/solvers/solver.py index 729a9d485..ec2f40c8d 100644 --- a/pina/solvers/solver.py +++ b/pina/solvers/solver.py @@ -142,13 +142,13 @@ 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): diff --git a/pina/solvers/supervised.py b/pina/solvers/supervised.py index 28a634b0b..425364614 100644 --- a/pina/solvers/supervised.py +++ b/pina/solvers/supervised.py @@ -118,7 +118,7 @@ def training_step(self, batch, batch_idx): :return: The sum of the loss functions. :rtype: LabelTensor """ - + condition_idx = batch["condition"] for condition_id in range(condition_idx.min(), condition_idx.max() + 1): @@ -162,7 +162,7 @@ def loss_data(self, input_pts, output_pts): :rtype: torch.Tensor """ return self.loss(self.forward(input_pts), output_pts) - + @property def scheduler(self): """ diff --git a/pina/trainer.py b/pina/trainer.py index 90779a6e9..40f4eb691 100644 --- a/pina/trainer.py +++ b/pina/trainer.py @@ -67,9 +67,9 @@ def _create_or_update_loader(self): pb = self._model.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] = torch.nn.Parameter( + pb.unknown_parameters[key].data.to(device) + ) def train(self, **kwargs): """ From ab3ebb8682b4e88b35912993f4c9811689ec6032 Mon Sep 17 00:00:00 2001 From: Nicola Demo Date: Tue, 18 Jun 2024 11:20:58 +0200 Subject: [PATCH 09/14] in progress --- pina/__init__.py | 2 +- pina/dataset.py | 2 +- pina/geometry/simplex.py | 2 +- pina/label_tensor.py | 638 +++++++++++++++++++++---------------- pina/model/fno.py | 2 +- pina/plotter.py | 2 +- tests/test_label_tensor.py | 278 ++++++++++------ 7 files changed, 548 insertions(+), 378 deletions(-) diff --git a/pina/__init__.py b/pina/__init__.py index 730b2ead4..c63440ba4 100644 --- a/pina/__init__.py +++ b/pina/__init__.py @@ -9,7 +9,7 @@ ] from .meta import * -from .label_tensor import LabelTensor +#from .label_tensor import LabelTensor from .solvers.solver import SolverInterface from .trainer import Trainer from .plotter import Plotter diff --git a/pina/dataset.py b/pina/dataset.py index c6a8d29e4..5f4ba5c7a 100644 --- a/pina/dataset.py +++ b/pina/dataset.py @@ -1,6 +1,6 @@ from torch.utils.data import Dataset import torch -from pina import LabelTensor +from .label_tensor import LabelTensor class SamplePointDataset(Dataset): diff --git a/pina/geometry/simplex.py b/pina/geometry/simplex.py index 2e78872aa..119d654f1 100644 --- a/pina/geometry/simplex.py +++ b/pina/geometry/simplex.py @@ -1,7 +1,7 @@ import torch from .location import Location from pina.geometry import CartesianDomain -from pina import LabelTensor +from ..label_tensor import LabelTensor from ..utils import check_consistency diff --git a/pina/label_tensor.py b/pina/label_tensor.py index c8a41f7b4..7cb2f71f0 100644 --- a/pina/label_tensor.py +++ b/pina/label_tensor.py @@ -5,6 +5,319 @@ from torch import Tensor + +# class LabelTensor(torch.Tensor): +# """Torch tensor with a label for any column.""" + +# @staticmethod +# def __new__(cls, x, labels, *args, **kwargs): +# return super().__new__(cls, x, *args, **kwargs) + +# def __init__(self, x, labels): +# """ +# Construct a `LabelTensor` by passing a tensor and a list of column +# labels. Such labels uniquely identify the columns of the tensor, +# allowing for an easier manipulation. + +# :param torch.Tensor x: The data tensor. +# :param labels: The labels of the columns. +# :type labels: str | list(str) | tuple(str) + +# :Example: +# >>> from pina import LabelTensor +# >>> tensor = LabelTensor(torch.rand((2000, 3)), ['a', 'b', 'c']) +# >>> tensor +# tensor([[6.7116e-02, 4.8892e-01, 8.9452e-01], +# [9.2392e-01, 8.2065e-01, 4.1986e-04], +# [8.9266e-01, 5.5446e-01, 6.3500e-01], +# ..., +# [5.8194e-01, 9.4268e-01, 4.1841e-01], +# [1.0246e-01, 9.5179e-01, 3.7043e-02], +# [9.6150e-01, 8.0656e-01, 8.3824e-01]]) +# >>> tensor.extract('a') +# tensor([[0.0671], +# [0.9239], +# [0.8927], +# ..., +# [0.5819], +# [0.1025], +# [0.9615]]) +# >>> tensor['a'] +# tensor([[0.0671], +# [0.9239], +# [0.8927], +# ..., +# [0.5819], +# [0.1025], +# [0.9615]]) +# >>> tensor.extract(['a', 'b']) +# tensor([[0.0671, 0.4889], +# [0.9239, 0.8207], +# [0.8927, 0.5545], +# ..., +# [0.5819, 0.9427], +# [0.1025, 0.9518], +# [0.9615, 0.8066]]) +# >>> tensor.extract(['b', 'a']) +# tensor([[0.4889, 0.0671], +# [0.8207, 0.9239], +# [0.5545, 0.8927], +# ..., +# [0.9427, 0.5819], +# [0.9518, 0.1025], +# [0.8066, 0.9615]]) +# """ +# if x.ndim == 1: +# x = x.reshape(-1, 1) + +# if isinstance(labels, str): +# labels = [labels] + +# if len(labels) != x.shape[-1]: +# raise ValueError( +# "the tensor has not the same number of columns of " +# "the passed labels." +# ) +# self._labels = labels + +# def __deepcopy__(self, __): +# """ +# Implements deepcopy for label tensor. By default it stores the +# current labels and use the :meth:`~torch._tensor.Tensor.__deepcopy__` +# method for creating a new :class:`pina.label_tensor.LabelTensor`. + +# :param __: Placeholder parameter. +# :type __: None +# :return: The deep copy of the :class:`pina.label_tensor.LabelTensor`. +# :rtype: LabelTensor +# """ +# labels = self.labels +# copy_tensor = deepcopy(self.tensor) +# return LabelTensor(copy_tensor, labels) + +# @property +# def labels(self): +# """Property decorator for labels + +# :return: labels of self +# :rtype: list +# """ +# return self._labels + +# @labels.setter +# def labels(self, labels): +# if len(labels) != self.shape[self.ndim - 1]: # small check +# raise ValueError( +# "The tensor has not the same number of columns of " +# "the passed labels." +# ) + +# self._labels = labels # assign the label + +# @staticmethod +# def vstack(label_tensors): +# """ +# Stack tensors vertically. For more details, see +# :meth:`torch.vstack`. + +# :param list(LabelTensor) label_tensors: the tensors to stack. They need +# to have equal labels. +# :return: the stacked tensor +# :rtype: LabelTensor +# """ +# if len(label_tensors) == 0: +# return [] + +# all_labels = [label for lt in label_tensors for label in lt.labels] +# if set(all_labels) != set(label_tensors[0].labels): +# raise RuntimeError("The tensors to stack have different labels") + +# labels = label_tensors[0].labels +# tensors = [lt.extract(labels) for lt in label_tensors] +# return LabelTensor(torch.vstack(tensors), labels) + +# def clone(self, *args, **kwargs): +# """ +# Clone the LabelTensor. For more details, see +# :meth:`torch.Tensor.clone`. + +# :return: A copy of the tensor. +# :rtype: LabelTensor +# """ +# # # used before merging +# # try: +# # out = LabelTensor(super().clone(*args, **kwargs), self.labels) +# # except: +# # out = super().clone(*args, **kwargs) +# out = LabelTensor(super().clone(*args, **kwargs), self.labels) +# return out + +# def to(self, *args, **kwargs): +# """ +# Performs Tensor dtype and/or device conversion. For more details, see +# :meth:`torch.Tensor.to`. +# """ +# tmp = super().to(*args, **kwargs) +# new = self.__class__.clone(self) +# new.data = tmp.data +# return new + +# def select(self, *args, **kwargs): +# """ +# Performs Tensor selection. For more details, see :meth:`torch.Tensor.select`. +# """ +# tmp = super().select(*args, **kwargs) +# tmp._labels = self._labels +# return tmp + +# def cuda(self, *args, **kwargs): +# """ +# Send Tensor to cuda. For more details, see :meth:`torch.Tensor.cuda`. +# """ +# tmp = super().cuda(*args, **kwargs) +# new = self.__class__.clone(self) +# new.data = tmp.data +# return new + +# def cpu(self, *args, **kwargs): +# """ +# Send Tensor to cpu. For more details, see :meth:`torch.Tensor.cpu`. +# """ +# tmp = super().cpu(*args, **kwargs) +# new = self.__class__.clone(self) +# new.data = tmp.data +# return new + +# def extract(self, label_to_extract): +# """ +# Extract the subset of the original tensor by returning all the columns +# corresponding to the passed ``label_to_extract``. + +# :param label_to_extract: The label(s) to extract. +# :type label_to_extract: str | list(str) | tuple(str) +# :raises TypeError: Labels are not ``str``. +# :raises ValueError: Label to extract is not in the labels ``list``. +# """ + +# if isinstance(label_to_extract, str): +# label_to_extract = [label_to_extract] +# elif isinstance(label_to_extract, (tuple, list)): # TODO +# pass +# else: +# raise TypeError( +# "`label_to_extract` should be a str, or a str iterator" +# ) + +# indeces = [] +# for f in label_to_extract: +# try: +# indeces.append(self.labels.index(f)) +# except ValueError: +# raise ValueError(f"`{f}` not in the labels list") + +# new_data = super(Tensor, self.T).__getitem__(indeces).T +# new_labels = [self.labels[idx] for idx in indeces] + +# extracted_tensor = new_data.as_subclass(LabelTensor) +# extracted_tensor.labels = new_labels + +# return extracted_tensor + +# def detach(self): +# detached = super().detach() +# if hasattr(self, "_labels"): +# detached._labels = self._labels +# return detached + +# def requires_grad_(self, mode=True): +# lt = super().requires_grad_(mode) +# lt.labels = self.labels +# return lt + +# def append(self, lt, mode="std"): +# """ +# Return a copy of the merged tensors. + +# :param LabelTensor lt: The tensor to merge. +# :param str mode: {'std', 'first', 'cross'} +# :return: The merged tensors. +# :rtype: LabelTensor +# """ +# if set(self.labels).intersection(lt.labels): +# raise RuntimeError("The tensors to merge have common labels") + +# new_labels = self.labels + lt.labels +# if mode == "std": +# new_tensor = torch.cat((self, lt), dim=1) +# elif mode == "first": +# raise NotImplementedError +# elif mode == "cross": +# tensor1 = self +# tensor2 = lt +# 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_tensor = torch.cat((tensor1, tensor2), dim=1) + +# new_tensor = new_tensor.as_subclass(LabelTensor) +# new_tensor.labels = new_labels +# return new_tensor + +# def __getitem__(self, index): +# """ +# Return a copy of the selected tensor. +# """ + +# if isinstance(index, str) or ( +# isinstance(index, (tuple, list)) +# and all(isinstance(a, str) for a in index) +# ): +# return self.extract(index) + +# selected_lt = super(Tensor, self).__getitem__(index) + +# try: +# len_index = len(index) +# except TypeError: +# len_index = 1 + +# if isinstance(index, int) or len_index == 1: +# if selected_lt.ndim == 1: +# selected_lt = selected_lt.reshape(1, -1) +# if hasattr(self, "labels"): +# selected_lt.labels = self.labels +# elif len_index == 2: +# if selected_lt.ndim == 1: +# selected_lt = selected_lt.reshape(-1, 1) +# if hasattr(self, "labels"): +# if isinstance(index[1], list): +# selected_lt.labels = [self.labels[i] for i in index[1]] +# else: +# selected_lt.labels = self.labels[index[1]] +# else: +# selected_lt.labels = self.labels + +# return selected_lt + + +# def __str__(self): +# if hasattr(self, "labels"): +# s = f"labels({str(self.labels)})\n" +# else: +# s = "no labels\n" +# s += super().__str__() +# return s + +def issubset(a, b): + """ + Check if a is a subset of b. + """ + return set(a).issubset(set(b)) + class LabelTensor(torch.Tensor): """Torch tensor with a label for any column.""" @@ -12,180 +325,35 @@ class LabelTensor(torch.Tensor): def __new__(cls, x, labels, *args, **kwargs): return super().__new__(cls, x, *args, **kwargs) - def __init__(self, x, labels): - """ - Construct a `LabelTensor` by passing a tensor and a list of column - labels. Such labels uniquely identify the columns of the tensor, - allowing for an easier manipulation. - - :param torch.Tensor x: The data tensor. - :param labels: The labels of the columns. - :type labels: str | list(str) | tuple(str) - - :Example: - >>> from pina import LabelTensor - >>> tensor = LabelTensor(torch.rand((2000, 3)), ['a', 'b', 'c']) - >>> tensor - tensor([[6.7116e-02, 4.8892e-01, 8.9452e-01], - [9.2392e-01, 8.2065e-01, 4.1986e-04], - [8.9266e-01, 5.5446e-01, 6.3500e-01], - ..., - [5.8194e-01, 9.4268e-01, 4.1841e-01], - [1.0246e-01, 9.5179e-01, 3.7043e-02], - [9.6150e-01, 8.0656e-01, 8.3824e-01]]) - >>> tensor.extract('a') - tensor([[0.0671], - [0.9239], - [0.8927], - ..., - [0.5819], - [0.1025], - [0.9615]]) - >>> tensor['a'] - tensor([[0.0671], - [0.9239], - [0.8927], - ..., - [0.5819], - [0.1025], - [0.9615]]) - >>> tensor.extract(['a', 'b']) - tensor([[0.0671, 0.4889], - [0.9239, 0.8207], - [0.8927, 0.5545], - ..., - [0.5819, 0.9427], - [0.1025, 0.9518], - [0.9615, 0.8066]]) - >>> tensor.extract(['b', 'a']) - tensor([[0.4889, 0.0671], - [0.8207, 0.9239], - [0.5545, 0.8927], - ..., - [0.9427, 0.5819], - [0.9518, 0.1025], - [0.8066, 0.9615]]) - """ - if x.ndim == 1: - x = x.reshape(-1, 1) - - if isinstance(labels, str): - labels = [labels] - - if len(labels) != x.shape[-1]: - raise ValueError( - "the tensor has not the same number of columns of " - "the passed labels." - ) - self._labels = labels - - def __deepcopy__(self, __): - """ - Implements deepcopy for label tensor. By default it stores the - current labels and use the :meth:`~torch._tensor.Tensor.__deepcopy__` - method for creating a new :class:`pina.label_tensor.LabelTensor`. - - :param __: Placeholder parameter. - :type __: None - :return: The deep copy of the :class:`pina.label_tensor.LabelTensor`. - :rtype: LabelTensor - """ - labels = self.labels - copy_tensor = deepcopy(self.tensor) - return LabelTensor(copy_tensor, labels) - @property - def labels(self): - """Property decorator for labels - - :return: labels of self - :rtype: list - """ - return self._labels - - @labels.setter - def labels(self, labels): - if len(labels) != self.shape[self.ndim - 1]: # small check - raise ValueError( - "The tensor has not the same number of columns of " - "the passed labels." - ) - - self._labels = labels # assign the label - - @staticmethod - def vstack(label_tensors): - """ - Stack tensors vertically. For more details, see - :meth:`torch.vstack`. - - :param list(LabelTensor) label_tensors: the tensors to stack. They need - to have equal labels. - :return: the stacked tensor - :rtype: LabelTensor - """ - if len(label_tensors) == 0: - return [] - - all_labels = [label for lt in label_tensors for label in lt.labels] - if set(all_labels) != set(label_tensors[0].labels): - raise RuntimeError("The tensors to stack have different labels") - - labels = label_tensors[0].labels - tensors = [lt.extract(labels) for lt in label_tensors] - return LabelTensor(torch.vstack(tensors), labels) + def tensor(self): + return self.as_subclass(Tensor) - def clone(self, *args, **kwargs): - """ - Clone the LabelTensor. For more details, see - :meth:`torch.Tensor.clone`. + def __len__(self) -> int: + return super().__len__() - :return: A copy of the tensor. - :rtype: LabelTensor - """ - # # used before merging - # try: - # out = LabelTensor(super().clone(*args, **kwargs), self.labels) - # except: - # out = super().clone(*args, **kwargs) - out = LabelTensor(super().clone(*args, **kwargs), self.labels) - return out - - def to(self, *args, **kwargs): - """ - Performs Tensor dtype and/or device conversion. For more details, see - :meth:`torch.Tensor.to`. + def __init__(self, x, labels): """ - tmp = super().to(*args, **kwargs) - new = self.__class__.clone(self) - new.data = tmp.data - return new + Construct a `LabelTensor` by passing a dict of the labels - def select(self, *args, **kwargs): - """ - Performs Tensor selection. For more details, see :meth:`torch.Tensor.select`. + :Example: + >>> from pina import LabelTensor + >>> tensor = LabelTensor( + >>> torch.rand((2000, 3)), + {1: {"name": "space"['a', 'b', 'c']) + """ - tmp = super().select(*args, **kwargs) - tmp._labels = self._labels - return tmp + from .utils import check_consistency + check_consistency(labels, dict) - def cuda(self, *args, **kwargs): - """ - Send Tensor to cuda. For more details, see :meth:`torch.Tensor.cuda`. - """ - tmp = super().cuda(*args, **kwargs) - new = self.__class__.clone(self) - new.data = tmp.data - return new + self.labels = { + idx_: { + 'dof': range(x.shape[idx_]), + 'name': idx_ + } for idx_ in range(x.ndim) + } + self.labels.update(labels) - def cpu(self, *args, **kwargs): - """ - Send Tensor to cpu. For more details, see :meth:`torch.Tensor.cpu`. - """ - tmp = super().cpu(*args, **kwargs) - new = self.__class__.clone(self) - new.data = tmp.data - return new def extract(self, label_to_extract): """ @@ -197,122 +365,52 @@ def extract(self, label_to_extract): :raises TypeError: Labels are not ``str``. :raises ValueError: Label to extract is not in the labels ``list``. """ - - if isinstance(label_to_extract, str): + if isinstance(label_to_extract, (int, str)): label_to_extract = [label_to_extract] - elif isinstance(label_to_extract, (tuple, list)): # TODO - pass - else: - raise TypeError( - "`label_to_extract` should be a str, or a str iterator" - ) + if isinstance(label_to_extract, (tuple, list)): + + for k, v in self.labels.items(): + if issubset(label_to_extract, v['dof']): + break - indeces = [] - for f in label_to_extract: - try: - indeces.append(self.labels.index(f)) - except ValueError: - raise ValueError(f"`{f}` not in the labels list") + label_to_extract = {v['name']: label_to_extract} - new_data = super(Tensor, self.T).__getitem__(indeces).T - new_labels = [self.labels[idx] for idx in indeces] + for k, v in label_to_extract.items(): + if isinstance(v, (int, str)): + label_to_extract[k] = [v] - extracted_tensor = new_data.as_subclass(LabelTensor) - extracted_tensor.labels = new_labels + indeces = [] + for dim in range(self.ndim): - return extracted_tensor + boolean_idx = [True] * self.shape[dim] - def detach(self): - detached = super().detach() - if hasattr(self, "_labels"): - detached._labels = self._labels - return detached + for dim_to_extract, dof_to_extract in label_to_extract.items(): + if dim_to_extract == self.labels[dim]['name']: + boolean_idx = [False] * self.shape[dim] + for label in dof_to_extract: + idx_to_keep = self.labels[dim]['dof'].index(label) + boolean_idx[idx_to_keep] = True - def requires_grad_(self, mode=True): - lt = super().requires_grad_(mode) - lt.labels = self.labels - return lt + boolean_idx = torch.Tensor(boolean_idx).bool() - def append(self, lt, mode="std"): - """ - Return a copy of the merged tensors. + indeces.append(boolean_idx) - :param LabelTensor lt: The tensor to merge. - :param str mode: {'std', 'first', 'cross'} - :return: The merged tensors. - :rtype: LabelTensor - """ - if set(self.labels).intersection(lt.labels): - raise RuntimeError("The tensors to merge have common labels") - - new_labels = self.labels + lt.labels - if mode == "std": - new_tensor = torch.cat((self, lt), dim=1) - elif mode == "first": - raise NotImplementedError - elif mode == "cross": - tensor1 = self - tensor2 = lt - 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_tensor = torch.cat((tensor1, tensor2), dim=1) - - new_tensor = new_tensor.as_subclass(LabelTensor) - new_tensor.labels = new_labels - return new_tensor - - def __getitem__(self, index): - """ - Return a copy of the selected tensor. - """ + final_shapes = [sum(idx) for idx in indeces] + grids = torch.meshgrid(*indeces) - if isinstance(index, str) or ( - isinstance(index, (tuple, list)) - and all(isinstance(a, str) for a in index) - ): - return self.extract(index) - - selected_lt = super(Tensor, self).__getitem__(index) - - try: - len_index = len(index) - except TypeError: - len_index = 1 - - if isinstance(index, int) or len_index == 1: - if selected_lt.ndim == 1: - selected_lt = selected_lt.reshape(1, -1) - if hasattr(self, "labels"): - selected_lt.labels = self.labels - elif len_index == 2: - if selected_lt.ndim == 1: - selected_lt = selected_lt.reshape(-1, 1) - if hasattr(self, "labels"): - if isinstance(index[1], list): - selected_lt.labels = [self.labels[i] for i in index[1]] - else: - selected_lt.labels = self.labels[index[1]] - else: - selected_lt.labels = self.labels - - return selected_lt + ii = grids[0] + for grid in grids[1:]: + ii = torch.logical_and(ii, grid) - @property - def tensor(self): - return self.as_subclass(Tensor) + new_tensor = self.tensor[ii].reshape(*final_shapes) - def __len__(self) -> int: - return super().__len__() + return LabelTensor(new_tensor, label_to_extract) + def __str__(self): - if hasattr(self, "labels"): - s = f"labels({str(self.labels)})\n" - else: - s = "no labels\n" + s = '' + for key, value in self.labels.items(): + s += f"{key}: {value}\n" + s += '\n' s += super().__str__() - return s + return s \ No newline at end of file diff --git a/pina/model/fno.py b/pina/model/fno.py index 910b41603..92ca18361 100644 --- a/pina/model/fno.py +++ b/pina/model/fno.py @@ -4,7 +4,7 @@ import torch import torch.nn as nn -from pina import LabelTensor +from ..label_tensor import LabelTensor import warnings from ..utils import check_consistency from .layers.fourier import FourierBlock1D, FourierBlock2D, FourierBlock3D diff --git a/pina/plotter.py b/pina/plotter.py index 63d86468b..0d430ff10 100644 --- a/pina/plotter.py +++ b/pina/plotter.py @@ -3,7 +3,7 @@ import matplotlib.pyplot as plt import torch from pina.callbacks import MetricTracker -from pina import LabelTensor +from .label_tensor import LabelTensor class Plotter: diff --git a/tests/test_label_tensor.py b/tests/test_label_tensor.py index 05dace5e3..5d5693d50 100644 --- a/tests/test_label_tensor.py +++ b/tests/test_label_tensor.py @@ -1,119 +1,191 @@ import torch import pytest -from pina import LabelTensor +from pina.label_tensor import LabelTensor +#import pina data = torch.rand((20, 3)) -labels = ['a', 'b', 'c'] - - -def test_constructor(): +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']) - -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'] +@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(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(): + 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", [2, [2], {'samples': [2]}]) +def test_extract_row(labels, labels_te): 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(): + 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) - 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(): + 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) - 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]) + 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_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)) +# tensor_view3 = tensor[:, 2] +# assert tensor_view3.labels == labels[2] +# assert torch.allclose(tensor_view3, data[:, 2].reshape(-1, 1)) From eec0249c91eb14c7c12d83b9d9751a75c655f0d2 Mon Sep 17 00:00:00 2001 From: Nicola Demo Date: Fri, 21 Jun 2024 14:37:55 +0200 Subject: [PATCH 10/14] optimizer and scheduler classes --- pina/__init__.py | 6 +++++- pina/optimizer.py | 21 +++++++++++++++++++++ pina/scheduler.py | 29 +++++++++++++++++++++++++++++ tests/test_optimizer.py | 20 ++++++++++++++++++++ tests/test_scheduler.py | 27 +++++++++++++++++++++++++++ 5 files changed, 102 insertions(+), 1 deletion(-) create mode 100644 pina/optimizer.py create mode 100644 pina/scheduler.py create mode 100644 tests/test_optimizer.py create mode 100644 tests/test_scheduler.py diff --git a/pina/__init__.py b/pina/__init__.py index c63440ba4..7c7253338 100644 --- a/pina/__init__.py +++ b/pina/__init__.py @@ -6,13 +6,17 @@ "Condition", "SamplePointDataset", "SamplePointLoader", + "TorchOptimizer", + "TorchScheduler", ] from .meta import * -#from .label_tensor import LabelTensor +from .label_tensor import LabelTensor from .solvers.solver import SolverInterface from .trainer import Trainer from .plotter import Plotter from .condition import Condition from .dataset import SamplePointDataset from .dataset import SamplePointLoader +from .optimizer import TorchOptimizer +from .scheduler import TorchScheduler \ No newline at end of file diff --git a/pina/optimizer.py b/pina/optimizer.py new file mode 100644 index 000000000..d400e82f3 --- /dev/null +++ b/pina/optimizer.py @@ -0,0 +1,21 @@ +""" Module for PINA Optimizer """ + +import torch +from .utils import check_consistency + +class Optimizer: # TODO improve interface + pass + + +class TorchOptimizer(Optimizer): + + def __init__(self, optimizer_class, **kwargs): + check_consistency(optimizer_class, torch.optim.Optimizer, subclass=True) + + self.optimizer_class = optimizer_class + self.kwargs = kwargs + + def hook(self, parameters): + self.optimizer_instance = self.optimizer_class( + parameters, **self.kwargs + ) diff --git a/pina/scheduler.py b/pina/scheduler.py new file mode 100644 index 000000000..563f829c5 --- /dev/null +++ b/pina/scheduler.py @@ -0,0 +1,29 @@ +""" Module for PINA Scheduler """ + +try: + from torch.optim.lr_scheduler import LRScheduler # torch >= 2.0 +except ImportError: + from torch.optim.lr_scheduler import ( + _LRScheduler as LRScheduler, + ) # torch < 2.0 +from .optimizer import Optimizer +from .utils import check_consistency + + +class Scheduler: # TODO improve interface + pass + + +class TorchScheduler(Scheduler): + + def __init__(self, scheduler_class, **kwargs): + check_consistency(scheduler_class, LRScheduler, subclass=True) + + self.scheduler_class = scheduler_class + self.kwargs = kwargs + + def hook(self, optimizer): + check_consistency(optimizer, Optimizer) + self.scheduler_instance = self.scheduler_class( + optimizer.optimizer_instance, **self.kwargs + ) diff --git a/tests/test_optimizer.py b/tests/test_optimizer.py new file mode 100644 index 000000000..489bbdc05 --- /dev/null +++ b/tests/test_optimizer.py @@ -0,0 +1,20 @@ + +import torch +import pytest +from pina import TorchOptimizer + +opt_list = [ + 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 diff --git a/tests/test_scheduler.py b/tests/test_scheduler.py new file mode 100644 index 000000000..4cde13efb --- /dev/null +++ b/tests/test_scheduler.py @@ -0,0 +1,27 @@ + +import torch +import pytest +from pina import TorchOptimizer, TorchScheduler + +opt_list = [ + torch.optim.Adam, + torch.optim.AdamW, + torch.optim.SGD, + torch.optim.RMSprop +] + +sch_list = [ + torch.optim.lr_scheduler.ConstantLR +] + +@pytest.mark.parametrize("scheduler_class", sch_list) +def test_constructor(scheduler_class): + TorchScheduler(scheduler_class) + +@pytest.mark.parametrize("optimizer_class", opt_list) +@pytest.mark.parametrize("scheduler_class", sch_list) +def test_hook(optimizer_class, scheduler_class): + opt = TorchOptimizer(optimizer_class, lr=1e-3) + opt.hook(torch.nn.Linear(10, 10).parameters()) + sch = TorchScheduler(scheduler_class) + sch.hook(opt) \ No newline at end of file From 560f6d097027583740be624a1a3bf23cdada9a49 Mon Sep 17 00:00:00 2001 From: Nicola Demo Date: Thu, 1 Aug 2024 16:30:35 +0200 Subject: [PATCH 11/14] minor --- pina/optimizer.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/pina/optimizer.py b/pina/optimizer.py index d400e82f3..08631e625 100644 --- a/pina/optimizer.py +++ b/pina/optimizer.py @@ -18,4 +18,4 @@ def __init__(self, optimizer_class, **kwargs): def hook(self, parameters): self.optimizer_instance = self.optimizer_class( parameters, **self.kwargs - ) + ) \ No newline at end of file From 252002ff524172f8d05b71a20b3791ff0135fa62 Mon Sep 17 00:00:00 2001 From: Nicola Demo Date: Mon, 5 Aug 2024 17:34:34 +0200 Subject: [PATCH 12/14] refact --- pina/__init__.py | 5 +- pina/condition/__init__.py | 10 + pina/{ => condition}/condition.py | 42 +-- pina/condition/condition_interface.py | 15 ++ pina/condition/domain_equation_condition.py | 28 ++ pina/condition/input_equation_condition.py | 23 ++ pina/condition/input_output_condition.py | 35 +++ pina/data/__init__.py | 0 pina/{ => data}/dataset.py | 2 +- pina/optim/__init__.py | 11 + pina/optim/optimizer_interface.py | 7 + pina/optim/scheduler_interface.py | 7 + pina/optim/torch_optimizer.py | 19 ++ pina/optim/torch_scheduler.py | 27 ++ pina/solvers/solver.py | 255 +++++++++++++------ pina/solvers/supervised.py | 63 ++--- pina/trainer.py | 2 +- tests/test_dataset.py | 2 +- tests/test_solvers/test_supervised_solver.py | 109 ++++---- 19 files changed, 486 insertions(+), 176 deletions(-) create mode 100644 pina/condition/__init__.py rename pina/{ => condition}/condition.py (73%) create mode 100644 pina/condition/condition_interface.py create mode 100644 pina/condition/domain_equation_condition.py create mode 100644 pina/condition/input_equation_condition.py create mode 100644 pina/condition/input_output_condition.py create mode 100644 pina/data/__init__.py rename pina/{ => data}/dataset.py (99%) create mode 100644 pina/optim/__init__.py create mode 100644 pina/optim/optimizer_interface.py create mode 100644 pina/optim/scheduler_interface.py create mode 100644 pina/optim/torch_optimizer.py create mode 100644 pina/optim/torch_scheduler.py diff --git a/pina/__init__.py b/pina/__init__.py index 7c7253338..e45a4af84 100644 --- a/pina/__init__.py +++ b/pina/__init__.py @@ -19,4 +19,7 @@ from .dataset import SamplePointDataset from .dataset import SamplePointLoader from .optimizer import TorchOptimizer -from .scheduler import TorchScheduler \ No newline at end of file +from .scheduler import TorchScheduler +from .condition.condition import Condition +from .data.dataset import SamplePointDataset +from .data.dataset import SamplePointLoader diff --git a/pina/condition/__init__.py b/pina/condition/__init__.py new file mode 100644 index 000000000..56d1ee4ad --- /dev/null +++ b/pina/condition/__init__.py @@ -0,0 +1,10 @@ +__all__ = [ + 'Condition', + 'ConditionInterface', + 'InputOutputCondition', + 'InputEquationCondition' + 'LocationEquationCondition', +] + +from .condition_interface import ConditionInterface +from .input_output_condition import InputOutputCondition \ No newline at end of file diff --git a/pina/condition.py b/pina/condition/condition.py similarity index 73% rename from pina/condition.py rename to pina/condition/condition.py index 5125fe084..da3c6f647 100644 --- a/pina/condition.py +++ b/pina/condition/condition.py @@ -1,8 +1,8 @@ """ Condition module. """ -from .label_tensor import LabelTensor -from .geometry import Location -from .equation.equation import Equation +from ..label_tensor import LabelTensor +from ..geometry import Location +from ..equation.equation import Equation def dummy(a): @@ -59,24 +59,32 @@ class Condition: "data_weight", ] - 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 + # 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_) + # return isinstance(dict_[key_], class_) - def __init__(self, *args, **kwargs): - """ - Constructor for the `Condition` class. - """ - self.data_weight = kwargs.pop("data_weight", 1.0) + # 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__}." - ) + # if len(args) != 0: + # raise ValueError( + # f"Condition takes only the following keyword arguments: {Condition.__slots__}." + # ) + from . import InputOutputCondition + def __new__(cls, *args, **kwargs): + + if sorted(kwargs.keys()) == sorted(["input_points", "output_points"]): + return InputOutputCondition(**kwargs) + else: + raise ValueError(f"Invalid keyword arguments {kwargs.keys()}.") + if ( sorted(kwargs.keys()) != sorted(["input_points", "output_points"]) and sorted(kwargs.keys()) != sorted(["location", "equation"]) diff --git a/pina/condition/condition_interface.py b/pina/condition/condition_interface.py new file mode 100644 index 000000000..bb43293cd --- /dev/null +++ b/pina/condition/condition_interface.py @@ -0,0 +1,15 @@ + +from abc import ABCMeta, abstractmethod + + +class ConditionInterface(metaclass=ABCMeta): + + @abstractmethod + def residual(self, model): + """ + Compute the residual of the condition. + + :param model: The model to evaluate the condition. + :return: The residual of the condition. + """ + pass \ No newline at end of file diff --git a/pina/condition/domain_equation_condition.py b/pina/condition/domain_equation_condition.py new file mode 100644 index 000000000..8f45f8f4a --- /dev/null +++ b/pina/condition/domain_equation_condition.py @@ -0,0 +1,28 @@ +from .condition_interface import ConditionInterface + +class DomainEquationCondition(ConditionInterface): + """ + Condition for input/output data. + """ + + __slots__ = ["domain", "equation"] + + def __init__(self, domain, equation): + """ + Constructor for the `InputOutputCondition` class. + """ + super().__init__() + self.domain = domain + self.equation = 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_points: The input points. + :param torch.Tensor output_points: The output points. + """ + return equation.residual(model(input_pts)) \ No newline at end of file diff --git a/pina/condition/input_equation_condition.py b/pina/condition/input_equation_condition.py new file mode 100644 index 000000000..1c57ed78c --- /dev/null +++ b/pina/condition/input_equation_condition.py @@ -0,0 +1,23 @@ + +from . import ConditionInterface + +class InputOutputCondition(ConditionInterface): + """ + Condition for input/output data. + """ + + __slots__ = ["input_points", "output_points"] + + def __init__(self, input_points, output_points): + """ + Constructor for the `InputOutputCondition` class. + """ + super().__init__() + self.input_points = input_points + self.output_points = output_points + + def residual(self, model): + """ + Compute the residual of the condition. + """ + return self.output_points - model(self.input_points) \ No newline at end of file diff --git a/pina/condition/input_output_condition.py b/pina/condition/input_output_condition.py new file mode 100644 index 000000000..d8040d64b --- /dev/null +++ b/pina/condition/input_output_condition.py @@ -0,0 +1,35 @@ + +from . import ConditionInterface + +class InputOutputCondition(ConditionInterface): + """ + Condition for input/output data. + """ + + __slots__ = ["input_points", "output_points"] + + def __init__(self, input_points, output_points): + """ + Constructor for the `InputOutputCondition` class. + """ + super().__init__() + self.input_points = input_points + self.output_points = output_points + + def residual(self, model): + """ + Compute the residual of the condition. + """ + return self.batch_residual(model, self.input_points, 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/data/__init__.py b/pina/data/__init__.py new file mode 100644 index 000000000..e69de29bb diff --git a/pina/dataset.py b/pina/data/dataset.py similarity index 99% rename from pina/dataset.py rename to pina/data/dataset.py index 5f4ba5c7a..bf874e1bf 100644 --- a/pina/dataset.py +++ b/pina/data/dataset.py @@ -1,6 +1,6 @@ from torch.utils.data import Dataset import torch -from .label_tensor import LabelTensor +from ..label_tensor import LabelTensor class SamplePointDataset(Dataset): diff --git a/pina/optim/__init__.py b/pina/optim/__init__.py new file mode 100644 index 000000000..699706c16 --- /dev/null +++ b/pina/optim/__init__.py @@ -0,0 +1,11 @@ +__all__ = [ + "Optimizer", + "TorchOptimizer", + "Scheduler", + "TorchScheduler", +] + +from .optimizer_interface import Optimizer +from .torch_optimizer import TorchOptimizer +from .scheduler_interface import Scheduler +from .torch_scheduler import TorchScheduler \ No newline at end of file diff --git a/pina/optim/optimizer_interface.py b/pina/optim/optimizer_interface.py new file mode 100644 index 000000000..c25506254 --- /dev/null +++ b/pina/optim/optimizer_interface.py @@ -0,0 +1,7 @@ +""" Module for PINA Optimizer """ + +from abc import ABCMeta + + +class Optimizer(metaclass=ABCMeta): # TODO improve interface + pass \ No newline at end of file diff --git a/pina/optim/scheduler_interface.py b/pina/optim/scheduler_interface.py new file mode 100644 index 000000000..dbc0ca82d --- /dev/null +++ b/pina/optim/scheduler_interface.py @@ -0,0 +1,7 @@ +""" Module for PINA Optimizer """ + +from abc import ABCMeta + + +class Scheduler(metaclass=ABCMeta): # TODO improve interface + pass \ No newline at end of file diff --git a/pina/optim/torch_optimizer.py b/pina/optim/torch_optimizer.py new file mode 100644 index 000000000..239819a4f --- /dev/null +++ b/pina/optim/torch_optimizer.py @@ -0,0 +1,19 @@ +""" Module for PINA Torch Optimizer """ + +import torch + +from ..utils import check_consistency +from .optimizer_interface import Optimizer + +class TorchOptimizer(Optimizer): + + def __init__(self, optimizer_class, **kwargs): + check_consistency(optimizer_class, torch.optim.Optimizer, subclass=True) + + self.optimizer_class = optimizer_class + self.kwargs = kwargs + + def hook(self, parameters): + self.optimizer_instance = self.optimizer_class( + parameters, **self.kwargs + ) \ No newline at end of file diff --git a/pina/optim/torch_scheduler.py b/pina/optim/torch_scheduler.py new file mode 100644 index 000000000..50e1d91f7 --- /dev/null +++ b/pina/optim/torch_scheduler.py @@ -0,0 +1,27 @@ +""" Module for PINA Torch Optimizer """ + +import torch +try: + from torch.optim.lr_scheduler import LRScheduler # torch >= 2.0 +except ImportError: + from torch.optim.lr_scheduler import ( + _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): + check_consistency(scheduler_class, LRScheduler, subclass=True) + + self.scheduler_class = scheduler_class + self.kwargs = 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 diff --git a/pina/solvers/solver.py b/pina/solvers/solver.py index ec2f40c8d..0112c860a 100644 --- a/pina/solvers/solver.py +++ b/pina/solvers/solver.py @@ -5,10 +5,173 @@ import pytorch_lightning from ..utils import check_consistency from ..problem import AbstractProblem +from ..optim import Optimizer, Scheduler import torch 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 @@ -18,45 +181,36 @@ class SolverInterface(pytorch_lightning.LightningModule, metaclass=ABCMeta): def __init__( self, - models, + model, problem, - optimizers, - optimizers_kwargs, - extra_features=None, + optimizer, + scheduler, ): """ - :param models: A torch neural network model instance. - :type models: torch.nn.Module + :param model: A torch neural network model instance. + :type model: 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. + :param list(torch.optim.Optimizer) optimizer: A list of neural network + optimizers to use. """ super().__init__() # check consistency of the inputs - check_consistency(models, torch.nn.Module) + check_consistency(model, torch.nn.Module) check_consistency(problem, AbstractProblem) - check_consistency(optimizers, torch.optim.Optimizer, subclass=True) - check_consistency(optimizers_kwargs, dict) + check_consistency(optimizer, Optimizer) + check_consistency(scheduler, Scheduler) # 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] + if not isinstance(model, list): + model = [model] + if not isinstance(optimizer, list): + optimizer = [optimizer] # number of models and optimizers - len_model = len(models) - len_optimizer = len(optimizers) - len_optimizer_kwargs = len(optimizers_kwargs) + len_model = len(model) + len_optimizer = len(optimizer) # check length consistency optimizers if len_model != len_optimizer: @@ -66,52 +220,11 @@ def __init__( " 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 + self._pina_model = model + self._pina_optimizer = optimizer + self._pina_scheduler = scheduler @abstractmethod def forward(self, *args, **kwargs): @@ -129,13 +242,13 @@ def configure_optimizers(self): def models(self): """ The torch model.""" - return self._pina_models + return self._pina_model @property def optimizers(self): """ The torch model.""" - return self._pina_optimizers + return self._pina_optimizer @property def problem(self): diff --git a/pina/solvers/supervised.py b/pina/solvers/supervised.py index 425364614..0285096ba 100644 --- a/pina/solvers/supervised.py +++ b/pina/solvers/supervised.py @@ -1,21 +1,14 @@ """ Module for SupervisedSolver """ import torch +from torch.nn.modules.loss import _Loss -try: - from torch.optim.lr_scheduler import LRScheduler # torch >= 2.0 -except ImportError: - from torch.optim.lr_scheduler import ( - _LRScheduler as LRScheduler, - ) # torch < 2.0 - -from torch.optim.lr_scheduler import ConstantLR +from ..optim import Optimizer, Scheduler, TorchOptimizer, TorchScheduler from .solver import SolverInterface from ..label_tensor import LabelTensor from ..utils import check_consistency from ..loss import LossInterface -from torch.nn.modules.loss import _Loss class SupervisedSolver(SolverInterface): @@ -51,12 +44,9 @@ def __init__( self, problem, model, - extra_features=None, - loss=torch.nn.MSELoss(), - optimizer=torch.optim.Adam, - optimizer_kwargs={"lr": 0.001}, - scheduler=ConstantLR, - scheduler_kwargs={"factor": 1, "total_iters": 0}, + loss=None, + optimizer=None, + scheduler=None, ): """ :param AbstractProblem problem: The formualation of the problem. @@ -73,24 +63,26 @@ def __init__( rate scheduler. :param dict scheduler_kwargs: LR scheduler constructor keyword args. """ + if loss is None: + loss = torch.nn.MSELoss() + + if optimizer is None: + optimizer = TorchOptimizer(torch.optim.Adam, lr=0.001) + + if scheduler is None: + scheduler = TorchScheduler( + torch.optim.lr_scheduler.ConstantLR) + super().__init__( - models=[model], + model=model, problem=problem, - optimizers=[optimizer], - optimizers_kwargs=[optimizer_kwargs], - extra_features=extra_features, + optimizer=optimizer, + scheduler=scheduler, ) # check consistency - check_consistency(scheduler, LRScheduler, subclass=True) - check_consistency(scheduler_kwargs, dict) check_consistency(loss, (LossInterface, _Loss), subclass=False) - # assign variables - self._scheduler = scheduler(self.optimizers[0], **scheduler_kwargs) - self._loss = loss - self._neural_net = self.models[0] - def forward(self, x): """Forward pass implementation for the solver. @@ -98,7 +90,7 @@ def forward(self, x): :return: Solver solution. :rtype: torch.Tensor """ - return self.neural_net(x) + return self._pina_model(x) def configure_optimizers(self): """Optimizer configuration for the solver. @@ -106,7 +98,9 @@ def configure_optimizers(self): :return: The optimizers and the schedulers :rtype: tuple(list, list) """ - return self.optimizers, [self.scheduler] + self._pina_optimizer.hook(self._pina_model.parameters()) + self._pina_scheduler.hook(self._pina_optimizer) + return self._pina_optimizer, self._pina_scheduler def training_step(self, batch, batch_idx): """Solver training step. @@ -168,14 +162,21 @@ def scheduler(self): """ Scheduler for training. """ - return self._scheduler + return self._pina_scheduler + + @property + def optimizer(self): + """ + Optimizer for training. + """ + return self._pina_optimizer @property - def neural_net(self): + def model(self): """ Neural network for training. """ - return self._neural_net + return self._pina_model @property def loss(self): diff --git a/pina/trainer.py b/pina/trainer.py index 40f4eb691..25d21b77c 100644 --- a/pina/trainer.py +++ b/pina/trainer.py @@ -3,7 +3,7 @@ import torch import pytorch_lightning from .utils import check_consistency -from .dataset import SamplePointDataset, SamplePointLoader, DataPointDataset +from .data.dataset import SamplePointDataset, SamplePointLoader, DataPointDataset from .solvers.solver import SolverInterface diff --git a/tests/test_dataset.py b/tests/test_dataset.py index ff1b6c228..cb6a9e4b0 100644 --- a/tests/test_dataset.py +++ b/tests/test_dataset.py @@ -1,7 +1,7 @@ import torch import pytest -from pina.dataset import SamplePointDataset, SamplePointLoader, DataPointDataset +from pina.data.dataset import SamplePointDataset, SamplePointLoader, DataPointDataset from pina import LabelTensor, Condition from pina.equation import Equation from pina.geometry import CartesianDomain diff --git a/tests/test_solvers/test_supervised_solver.py b/tests/test_solvers/test_supervised_solver.py index dfe0bd867..d9cbea3ce 100644 --- a/tests/test_solvers/test_supervised_solver.py +++ b/tests/test_solvers/test_supervised_solver.py @@ -11,8 +11,11 @@ class NeuralOperatorProblem(AbstractProblem): input_variables = ['u_0', 'u_1'] output_variables = ['u'] - conditions = {'data' : Condition(input_points=LabelTensor(torch.rand(100, 2), input_variables), - output_points=LabelTensor(torch.rand(100, 1), output_variables))} + conditions = { + # 'data' : Condition( + # input_points=LabelTensor(torch.rand(100, 2), input_variables), + # output_points=LabelTensor(torch.rand(100, 1), output_variables)) + } class myFeature(torch.nn.Module): """ @@ -39,63 +42,63 @@ def forward(self, x): def test_constructor(): - SupervisedSolver(problem=problem, model=model, extra_features=None) + SupervisedSolver(problem=problem, model=model) -def test_constructor_extra_feats(): - SupervisedSolver(problem=problem, model=model_extra_feats, extra_features=extra_feats) +# def test_constructor_extra_feats(): +# SupervisedSolver(problem=problem, model=model_extra_feats, extra_features=extra_feats) def test_train_cpu(): - solver = SupervisedSolver(problem = problem, model=model, extra_features=None, loss=LpLoss()) + solver = SupervisedSolver(problem = problem, model=model, loss=LpLoss()) trainer = Trainer(solver=solver, max_epochs=3, accelerator='cpu', batch_size=20) trainer.train() -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() \ No newline at end of file +# 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() \ No newline at end of file From 2f61e2f93b81696ef197f84c5c035a283b400a17 Mon Sep 17 00:00:00 2001 From: Nicola Demo Date: Thu, 8 Aug 2024 16:19:52 +0200 Subject: [PATCH 13/14] supervised working --- examples/problems/burgers.py | 2 +- examples/problems/first_order_ode.py | 2 +- .../parametric_elliptic_optimal_control.py | 2 +- examples/problems/parametric_poisson.py | 2 +- examples/problems/poisson.py | 2 +- examples/problems/stokes.py | 2 +- examples/problems/wave.py | 2 +- pina/__init__.py | 7 +- pina/condition/__init__.py | 8 +- pina/condition/condition.py | 22 +-- pina/condition/condition_interface.py | 3 + ...ondition.py => domain_output_condition.py} | 18 ++- pina/condition/input_equation_condition.py | 2 +- pina/data/__init__.py | 7 + pina/data/data_dataset.py | 41 +++++ pina/data/pina_batch.py | 36 +++++ pina/data/{dataset.py => pina_dataloader.py} | 142 ++++++------------ pina/data/sample_dataset.py | 43 ++++++ pina/{geometry => domain}/__init__.py | 4 +- pina/{geometry => domain}/cartesian.py | 4 +- .../{geometry => domain}/difference_domain.py | 0 .../domain_interface.py} | 4 +- pina/{geometry => domain}/ellipsoid.py | 4 +- pina/{geometry => domain}/exclusion_domain.py | 0 .../intersection_domain.py | 0 .../operation_interface.py | 6 +- pina/{geometry => domain}/simplex.py | 6 +- pina/{geometry => domain}/union_domain.py | 0 pina/label_tensor.py | 26 +++- pina/problem/abstract_problem.py | 88 ++++------- pina/solvers/solver.py | 2 + pina/solvers/supervised.py | 37 ++++- pina/trainer.py | 46 +++--- .../test_adaptive_refinment_callbacks.py | 2 +- .../test_optimizer_callbacks.py | 2 +- tests/test_condition.py | 2 +- tests/test_dataset.py | 2 +- tests/test_geometry/test_cartesian.py | 2 +- tests/test_geometry/test_difference.py | 2 +- tests/test_geometry/test_ellipsoid.py | 2 +- tests/test_geometry/test_exclusion.py | 2 +- tests/test_geometry/test_intersection.py | 2 +- tests/test_geometry/test_simplex.py | 2 +- tests/test_geometry/test_union.py | 2 +- tests/test_plotter.py | 2 +- tests/test_problem.py | 2 +- tests/test_solvers/test_causalpinn.py | 2 +- tests/test_solvers/test_competitive_pinn.py | 2 +- tests/test_solvers/test_gpinn.py | 2 +- tests/test_solvers/test_pinn.py | 2 +- tests/test_solvers/test_sapinn.py | 2 +- tests/test_solvers/test_supervised_solver.py | 15 +- tests/test_utils.py | 4 +- tutorials/tutorial1/tutorial.py | 4 +- tutorials/tutorial11/tutorial.py | 2 +- tutorials/tutorial12/tutorial.py | 2 +- tutorials/tutorial2/tutorial.py | 2 +- tutorials/tutorial3/tutorial.py | 2 +- tutorials/tutorial6/tutorial.py | 2 +- tutorials/tutorial7/tutorial.py | 2 +- tutorials/tutorial8/tutorial.py | 2 +- tutorials/tutorial9/tutorial.py | 2 +- 62 files changed, 379 insertions(+), 266 deletions(-) rename pina/condition/{input_output_condition.py => domain_output_condition.py} (65%) create mode 100644 pina/data/data_dataset.py create mode 100644 pina/data/pina_batch.py rename pina/data/{dataset.py => pina_dataloader.py} (67%) create mode 100644 pina/data/sample_dataset.py rename pina/{geometry => domain}/__init__.py (87%) rename pina/{geometry => domain}/cartesian.py (99%) rename pina/{geometry => domain}/difference_domain.py (100%) rename pina/{geometry/location.py => domain/domain_interface.py} (90%) rename pina/{geometry => domain}/ellipsoid.py (99%) rename pina/{geometry => domain}/exclusion_domain.py (100%) rename pina/{geometry => domain}/intersection_domain.py (100%) rename pina/{geometry => domain}/operation_interface.py (92%) rename pina/{geometry => domain}/simplex.py (98%) rename pina/{geometry => domain}/union_domain.py (100%) diff --git a/examples/problems/burgers.py b/examples/problems/burgers.py index 5da9ccb47..33c1b5db2 100644 --- a/examples/problems/burgers.py +++ b/examples/problems/burgers.py @@ -15,7 +15,7 @@ import torch -from pina.geometry import CartesianDomain +from pina.domain import CartesianDomain from pina import Condition from pina.problem import TimeDependentProblem, SpatialProblem from pina.operators import grad diff --git a/examples/problems/first_order_ode.py b/examples/problems/first_order_ode.py index 802b85bfe..be1d88c42 100644 --- a/examples/problems/first_order_ode.py +++ b/examples/problems/first_order_ode.py @@ -17,7 +17,7 @@ from pina.problem import SpatialProblem from pina import Condition -from pina.geometry import CartesianDomain +from pina.domain import CartesianDomain from pina.operators import grad from pina.equation import Equation, FixedValue import torch diff --git a/examples/problems/parametric_elliptic_optimal_control.py b/examples/problems/parametric_elliptic_optimal_control.py index 3fd26c45b..d5925ebd6 100644 --- a/examples/problems/parametric_elliptic_optimal_control.py +++ b/examples/problems/parametric_elliptic_optimal_control.py @@ -2,7 +2,7 @@ from pina import Condition -from pina.geometry import CartesianDomain +from pina.domain import CartesianDomain from pina.equation import SystemEquation, FixedValue from pina.problem import SpatialProblem, ParametricProblem from pina.operators import laplacian diff --git a/examples/problems/parametric_poisson.py b/examples/problems/parametric_poisson.py index 58867d5bb..64dfdaaee 100644 --- a/examples/problems/parametric_poisson.py +++ b/examples/problems/parametric_poisson.py @@ -14,7 +14,7 @@ # ===================================================== # -from pina.geometry import CartesianDomain +from pina.domain import CartesianDomain from pina.problem import SpatialProblem, ParametricProblem from pina.operators import laplacian from pina.equation import FixedValue, Equation diff --git a/examples/problems/poisson.py b/examples/problems/poisson.py index c817787bd..e4a6cf98e 100644 --- a/examples/problems/poisson.py +++ b/examples/problems/poisson.py @@ -13,7 +13,7 @@ import torch -from pina.geometry import CartesianDomain +from pina.domain import CartesianDomain from pina import Condition from pina.problem import SpatialProblem from pina.operators import laplacian diff --git a/examples/problems/stokes.py b/examples/problems/stokes.py index f136d64ad..c7b13873c 100644 --- a/examples/problems/stokes.py +++ b/examples/problems/stokes.py @@ -4,7 +4,7 @@ from pina.problem import SpatialProblem from pina.operators import laplacian, grad, div from pina import Condition, LabelTensor -from pina.geometry import CartesianDomain +from pina.domain import CartesianDomain from pina.equation import SystemEquation, Equation # ===================================================== # diff --git a/examples/problems/wave.py b/examples/problems/wave.py index cce94da68..124a62d4b 100644 --- a/examples/problems/wave.py +++ b/examples/problems/wave.py @@ -2,7 +2,7 @@ import torch -from pina.geometry import CartesianDomain +from pina.domain import CartesianDomain from pina import Condition from pina.problem import SpatialProblem, TimeDependentProblem from pina.operators import laplacian, grad diff --git a/pina/__init__.py b/pina/__init__.py index e45a4af84..835232b12 100644 --- a/pina/__init__.py +++ b/pina/__init__.py @@ -15,11 +15,8 @@ from .solvers.solver import SolverInterface from .trainer import Trainer from .plotter import Plotter -from .condition import Condition -from .dataset import SamplePointDataset -from .dataset import SamplePointLoader from .optimizer import TorchOptimizer from .scheduler import TorchScheduler from .condition.condition import Condition -from .data.dataset import SamplePointDataset -from .data.dataset import SamplePointLoader +from .data import SamplePointDataset +from .data import SamplePointLoader diff --git a/pina/condition/__init__.py b/pina/condition/__init__.py index 56d1ee4ad..ff329c1bc 100644 --- a/pina/condition/__init__.py +++ b/pina/condition/__init__.py @@ -1,10 +1,10 @@ __all__ = [ 'Condition', 'ConditionInterface', - 'InputOutputCondition', - 'InputEquationCondition' - 'LocationEquationCondition', + 'DomainOutputCondition', + 'DomainEquationCondition' ] from .condition_interface import ConditionInterface -from .input_output_condition import InputOutputCondition \ No newline at end of file +from .domain_output_condition import DomainOutputCondition +from .domain_equation_condition import DomainEquationCondition \ No newline at end of file diff --git a/pina/condition/condition.py b/pina/condition/condition.py index da3c6f647..eec523c92 100644 --- a/pina/condition/condition.py +++ b/pina/condition/condition.py @@ -1,9 +1,11 @@ """ Condition module. """ from ..label_tensor import LabelTensor -from ..geometry import Location +from ..domain import DomainInterface from ..equation.equation import Equation +from . import DomainOutputCondition, DomainEquationCondition + def dummy(a): """Dummy function for testing purposes.""" @@ -51,14 +53,6 @@ class Condition: """ - __slots__ = [ - "input_points", - "output_points", - "location", - "equation", - "data_weight", - ] - # 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(): @@ -77,11 +71,17 @@ class Condition: # f"Condition takes only the following keyword arguments: {Condition.__slots__}." # ) - from . import InputOutputCondition def __new__(cls, *args, **kwargs): if sorted(kwargs.keys()) == sorted(["input_points", "output_points"]): - return InputOutputCondition(**kwargs) + 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"]): + return DomainEquationCondition(**kwargs) else: raise ValueError(f"Invalid keyword arguments {kwargs.keys()}.") diff --git a/pina/condition/condition_interface.py b/pina/condition/condition_interface.py index bb43293cd..f6b51bf96 100644 --- a/pina/condition/condition_interface.py +++ b/pina/condition/condition_interface.py @@ -4,6 +4,9 @@ class ConditionInterface(metaclass=ABCMeta): + def __init__(self) -> None: + self._problem = None + @abstractmethod def residual(self, model): """ diff --git a/pina/condition/input_output_condition.py b/pina/condition/domain_output_condition.py similarity index 65% rename from pina/condition/input_output_condition.py rename to pina/condition/domain_output_condition.py index d8040d64b..49a0cb6ed 100644 --- a/pina/condition/input_output_condition.py +++ b/pina/condition/domain_output_condition.py @@ -1,26 +1,34 @@ from . import ConditionInterface -class InputOutputCondition(ConditionInterface): +class DomainOutputCondition(ConditionInterface): """ Condition for input/output data. """ - __slots__ = ["input_points", "output_points"] + __slots__ = ["domain", "output_points"] - def __init__(self, input_points, output_points): + def __init__(self, domain, output_points): """ Constructor for the `InputOutputCondition` class. """ super().__init__() - self.input_points = input_points + 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.input_points, self.output_points) + return self.batch_residual(model, self.domain, self.output_points) @staticmethod def batch_residual(model, input_points, output_points): diff --git a/pina/condition/input_equation_condition.py b/pina/condition/input_equation_condition.py index 1c57ed78c..288022c00 100644 --- a/pina/condition/input_equation_condition.py +++ b/pina/condition/input_equation_condition.py @@ -1,7 +1,7 @@ from . import ConditionInterface -class InputOutputCondition(ConditionInterface): +class InputEquationCondition(ConditionInterface): """ Condition for input/output data. """ diff --git a/pina/data/__init__.py b/pina/data/__init__.py index e69de29bb..fba19b92c 100644 --- a/pina/data/__init__.py +++ b/pina/data/__init__.py @@ -0,0 +1,7 @@ +__all__ = [ +] + +from .pina_dataloader import SamplePointLoader +from .data_dataset import DataPointDataset +from .sample_dataset import SamplePointDataset +from .pina_batch import Batch \ No newline at end of file diff --git a/pina/data/data_dataset.py b/pina/data/data_dataset.py new file mode 100644 index 000000000..9dff2d7ed --- /dev/null +++ b/pina/data/data_dataset.py @@ -0,0 +1,41 @@ +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/pina_batch.py b/pina/data/pina_batch.py new file mode 100644 index 000000000..cb1296ede --- /dev/null +++ b/pina/data/pina_batch.py @@ -0,0 +1,36 @@ + + +class Batch: + """ + This class is used to create a dataset of sample points. + """ + + 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] + + elif type_ == "data": + + 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 diff --git a/pina/data/dataset.py b/pina/data/pina_dataloader.py similarity index 67% rename from pina/data/dataset.py rename to pina/data/pina_dataloader.py index bf874e1bf..1b71a46a3 100644 --- a/pina/data/dataset.py +++ b/pina/data/pina_dataloader.py @@ -1,84 +1,8 @@ -from torch.utils.data import Dataset import torch -from ..label_tensor import LabelTensor - - -class SamplePointDataset(Dataset): - """ - This class is used to create a dataset of sample 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.vstack(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] - - -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.vstack(input_list) - self.output_pts = LabelTensor.vstack(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] +from .sample_dataset import SamplePointDataset +from .data_dataset import DataPointDataset +from .pina_batch import Batch class SamplePointLoader: """ @@ -133,6 +57,8 @@ def __init__( 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. @@ -169,7 +95,7 @@ def _prepare_data_dataset(self, dataset, batch_size, shuffle): 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 @@ -216,6 +142,29 @@ def _prepare_sample_dataset(self, dataset, batch_size, shuffle): self.tensor_conditions, batch_num ) + def _prepare_batches(self): + """ + Prepare the batches. + """ + 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) + else: + batch = Batch( + "data", idx_, + self.batch_input_pts, + self.batch_output_pts, + self.batch_data_conditions) + print(batch.input.labels) + + self.batches.append(batch) + def __iter__(self): """ Return an iterator over the points. Any element of the iterator is a @@ -233,21 +182,24 @@ def __iter__(self): :rtype: iter """ # for i in self.random_idx: - 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 + 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 def __len__(self): """ diff --git a/pina/data/sample_dataset.py b/pina/data/sample_dataset.py new file mode 100644 index 000000000..84af2920f --- /dev/null +++ b/pina/data/sample_dataset.py @@ -0,0 +1,43 @@ +from torch.utils.data import Dataset +import torch + +from ..label_tensor import LabelTensor + + +class SamplePointDataset(Dataset): + """ + This class is used to create a dataset of sample 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 diff --git a/pina/geometry/__init__.py b/pina/domain/__init__.py similarity index 87% rename from pina/geometry/__init__.py rename to pina/domain/__init__.py index 963136a3e..e5a327b4c 100644 --- a/pina/geometry/__init__.py +++ b/pina/domain/__init__.py @@ -1,5 +1,5 @@ __all__ = [ - "Location", + "DomainInterface", "CartesianDomain", "EllipsoidDomain", "Union", @@ -10,7 +10,7 @@ "SimplexDomain", ] -from .location import Location +from .domain_interface import DomainInterface from .cartesian import CartesianDomain from .ellipsoid import EllipsoidDomain from .exclusion_domain import Exclusion diff --git a/pina/geometry/cartesian.py b/pina/domain/cartesian.py similarity index 99% rename from pina/geometry/cartesian.py rename to pina/domain/cartesian.py index 6008589b9..04f3abefe 100644 --- a/pina/geometry/cartesian.py +++ b/pina/domain/cartesian.py @@ -1,11 +1,11 @@ import torch -from .location import Location +from .domain_interface import DomainInterface from ..label_tensor import LabelTensor from ..utils import torch_lhs, chebyshev_roots -class CartesianDomain(Location): +class CartesianDomain(DomainInterface): """PINA implementation of Hypercube domain.""" def __init__(self, cartesian_dict): diff --git a/pina/geometry/difference_domain.py b/pina/domain/difference_domain.py similarity index 100% rename from pina/geometry/difference_domain.py rename to pina/domain/difference_domain.py diff --git a/pina/geometry/location.py b/pina/domain/domain_interface.py similarity index 90% rename from pina/geometry/location.py rename to pina/domain/domain_interface.py index a22dfe13f..5906d2851 100644 --- a/pina/geometry/location.py +++ b/pina/domain/domain_interface.py @@ -1,9 +1,9 @@ -"""Module for Location class.""" +"""Module for the DomainInterface class.""" from abc import ABCMeta, abstractmethod -class Location(metaclass=ABCMeta): +class DomainInterface(metaclass=ABCMeta): """ Abstract Location class. Any geometry entity should inherit from this class. diff --git a/pina/geometry/ellipsoid.py b/pina/domain/ellipsoid.py similarity index 99% rename from pina/geometry/ellipsoid.py rename to pina/domain/ellipsoid.py index 0764d81e1..dcbc55543 100644 --- a/pina/geometry/ellipsoid.py +++ b/pina/domain/ellipsoid.py @@ -1,11 +1,11 @@ import torch -from .location import Location +from .domain_interface import DomainInterface from ..label_tensor import LabelTensor from ..utils import check_consistency -class EllipsoidDomain(Location): +class EllipsoidDomain(DomainInterface): """PINA implementation of Ellipsoid domain.""" def __init__(self, ellipsoid_dict, sample_surface=False): diff --git a/pina/geometry/exclusion_domain.py b/pina/domain/exclusion_domain.py similarity index 100% rename from pina/geometry/exclusion_domain.py rename to pina/domain/exclusion_domain.py diff --git a/pina/geometry/intersection_domain.py b/pina/domain/intersection_domain.py similarity index 100% rename from pina/geometry/intersection_domain.py rename to pina/domain/intersection_domain.py diff --git a/pina/geometry/operation_interface.py b/pina/domain/operation_interface.py similarity index 92% rename from pina/geometry/operation_interface.py rename to pina/domain/operation_interface.py index 4f7709b9a..acd4cf44b 100644 --- a/pina/geometry/operation_interface.py +++ b/pina/domain/operation_interface.py @@ -1,11 +1,11 @@ """ Module for OperationInterface class. """ -from .location import Location +from .domain_interface import DomainInterface from ..utils import check_consistency from abc import ABCMeta, abstractmethod -class OperationInterface(Location, metaclass=ABCMeta): +class OperationInterface(DomainInterface, metaclass=ABCMeta): def __init__(self, geometries): """ @@ -15,7 +15,7 @@ def __init__(self, geometries): such as ``EllipsoidDomain`` or ``CartesianDomain``. """ # check consistency geometries - check_consistency(geometries, Location) + check_consistency(geometries, DomainInterface) # check we are passing always different # geometries with the same labels. diff --git a/pina/geometry/simplex.py b/pina/domain/simplex.py similarity index 98% rename from pina/geometry/simplex.py rename to pina/domain/simplex.py index 119d654f1..8d26422ae 100644 --- a/pina/geometry/simplex.py +++ b/pina/domain/simplex.py @@ -1,11 +1,11 @@ import torch -from .location import Location -from pina.geometry import CartesianDomain +from .domain_interface import DomainInterface +from pina.domain import CartesianDomain from ..label_tensor import LabelTensor from ..utils import check_consistency -class SimplexDomain(Location): +class SimplexDomain(DomainInterface): """PINA implementation of a Simplex.""" def __init__(self, simplex_matrix, sample_surface=False): diff --git a/pina/geometry/union_domain.py b/pina/domain/union_domain.py similarity index 100% rename from pina/geometry/union_domain.py rename to pina/domain/union_domain.py diff --git a/pina/label_tensor.py b/pina/label_tensor.py index 7cb2f71f0..12d2182f5 100644 --- a/pina/label_tensor.py +++ b/pina/label_tensor.py @@ -229,10 +229,6 @@ # detached._labels = self._labels # return detached -# def requires_grad_(self, mode=True): -# lt = super().requires_grad_(mode) -# lt.labels = self.labels -# return lt # def append(self, lt, mode="std"): # """ @@ -406,11 +402,29 @@ def extract(self, label_to_extract): return LabelTensor(new_tensor, label_to_extract) - def __str__(self): s = '' for key, value in self.labels.items(): s += f"{key}: {value}\n" s += '\n' s += super().__str__() - return s \ No newline at end of file + return s + + @staticmethod + def stack(tensors): + """ + """ + if len(tensors) == 0: + return [] + + if len(tensors) == 1: + return tensors[0] + + raise NotImplementedError + labels = [tensor.labels for tensor in tensors] + print(labels) + + def requires_grad_(self, mode=True): + lt = super().requires_grad_(mode) + lt.labels = self.labels + return lt \ No newline at end of file diff --git a/pina/problem/abstract_problem.py b/pina/problem/abstract_problem.py index 6ebba6c5e..05b335670 100644 --- a/pina/problem/abstract_problem.py +++ b/pina/problem/abstract_problem.py @@ -5,6 +5,8 @@ from copy import deepcopy import torch +from .. import LabelTensor + class AbstractProblem(metaclass=ABCMeta): """ @@ -18,17 +20,26 @@ class AbstractProblem(metaclass=ABCMeta): def __init__(self): - # variable storing all points - self.input_pts = {} - # 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 = {} + self._discretized_domains = {} + + for name, domain in self.domains.items(): + if isinstance(domain, (torch.Tensor, LabelTensor)): + self._discretized_domains[name] = domain + for condition_name in self.conditions: - self._have_sampled_points[condition_name] = False + self.conditions[condition_name]._problem = self + # # variable storing all points + # self.input_pts = {} + + # # 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._have_sampled_points[condition_name] = False - # put in self.input_pts all the points that we don't need to sample - self._span_condition_points() + # # put in self.input_pts all the points that we don't need to sample + # self._span_condition_points() def __deepcopy__(self, memo): """ @@ -63,15 +74,20 @@ def input_variables(self): variables += self.spatial_variables if hasattr(self, "temporal_variable"): variables += self.temporal_variable - if hasattr(self, "parameters"): + if hasattr(self, "unknown_parameters"): variables += self.parameters if hasattr(self, "custom_variables"): variables += self.custom_variables return variables + @input_variables.setter + def input_variables(self, variables): + raise RuntimeError + @property - def domain(self): + @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 @@ -80,27 +96,7 @@ def domain(self): :return: the domain(s) of ``self`` :rtype: list[Location] """ - domains = [ - getattr(self, f"{t}_domain") - for t in ["spatial", "temporal", "parameter"] - if hasattr(self, f"{t}_domain") - ] - - if len(domains) == 1: - return domains[0] - elif len(domains) == 0: - raise RuntimeError - - if len(set(map(type, domains))) == 1: - domain = domains[0].__class__({}) - [domain.update(d) for d in domains] - return domain - else: - raise RuntimeError("different domains") - - @input_variables.setter - def input_variables(self, variables): - raise RuntimeError + pass @property @abstractmethod @@ -116,7 +112,9 @@ def conditions(self): """ The conditions of the problem. """ - pass + return self._conditions + + def _span_condition_points(self): """ @@ -272,28 +270,4 @@ def add_points(self, new_points): # merging merged_pts = torch.vstack([old_pts, new_points[location]]) merged_pts.labels = old_pts.labels - self.input_pts[location] = merged_pts - - @property - def have_sampled_points(self): - """ - Check if all points for - ``Location`` are sampled. - """ - return all(self._have_sampled_points.values()) - - @property - def not_sampled_points(self): - """ - Check which points are - not sampled. - """ - # variables which are not sampled - not_sampled = None - if self.have_sampled_points is False: - # check which one are not sampled: - not_sampled = [] - for condition_name, is_sample in self._have_sampled_points.items(): - if not is_sample: - not_sampled.append(condition_name) - return not_sampled + self.input_pts[location] = merged_pts \ No newline at end of file diff --git a/pina/solvers/solver.py b/pina/solvers/solver.py index 0112c860a..a27e93641 100644 --- a/pina/solvers/solver.py +++ b/pina/solvers/solver.py @@ -205,6 +205,8 @@ def __init__( # 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] diff --git a/pina/solvers/supervised.py b/pina/solvers/supervised.py index 0285096ba..00a9c8fa0 100644 --- a/pina/solvers/supervised.py +++ b/pina/solvers/supervised.py @@ -82,6 +82,7 @@ def __init__( # check consistency check_consistency(loss, (LossInterface, _Loss), subclass=False) + self.loss = loss def forward(self, x): """Forward pass implementation for the solver. @@ -90,7 +91,16 @@ def forward(self, x): :return: Solver solution. :rtype: torch.Tensor """ - return self._pina_model(x) + + output = self._pina_model[0](x) + + output.labels = { + 1: { + "name": "output", + "dof": self.problem.output_variables + } + } + return output def configure_optimizers(self): """Optimizer configuration for the solver. @@ -98,9 +108,12 @@ def configure_optimizers(self): :return: The optimizers and the schedulers :rtype: tuple(list, list) """ - self._pina_optimizer.hook(self._pina_model.parameters()) - self._pina_scheduler.hook(self._pina_optimizer) - return self._pina_optimizer, self._pina_scheduler + self._pina_optimizer[0].hook(self._pina_model[0].parameters()) + self._pina_scheduler[0].hook(self._pina_optimizer[0]) + return ( + [self._pina_optimizer[0].optimizer_instance], + [self._pina_scheduler[0].scheduler_instance] + ) def training_step(self, batch, batch_idx): """Solver training step. @@ -113,14 +126,16 @@ def training_step(self, batch, batch_idx): :rtype: LabelTensor """ - condition_idx = batch["condition"] + condition_idx = batch.condition 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["pts"] - out = batch["output"] + pts = batch.input + out = batch.output + print(out) + print(pts) if condition_name not in self.problem.conditions: raise RuntimeError("Something wrong happened.") @@ -134,9 +149,11 @@ def training_step(self, batch, batch_idx): 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) - * condition.data_weight ) loss = loss.as_subclass(torch.Tensor) @@ -155,6 +172,10 @@ def loss_data(self, input_pts, output_pts): :return: The residual loss averaged on the input coordinates :rtype: torch.Tensor """ + print(input_pts) + print(output_pts) + print(self.loss) + print(self.forward(input_pts)) return self.loss(self.forward(input_pts), output_pts) @property diff --git a/pina/trainer.py b/pina/trainer.py index 25d21b77c..758bbaaf0 100644 --- a/pina/trainer.py +++ b/pina/trainer.py @@ -3,7 +3,7 @@ import torch import pytorch_lightning from .utils import check_consistency -from .data.dataset import SamplePointDataset, SamplePointLoader, DataPointDataset +from .data import SamplePointDataset, SamplePointLoader, DataPointDataset from .solvers.solver import SolverInterface @@ -35,19 +35,33 @@ def __init__(self, solver, batch_size=None, **kwargs): self._model = 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 _create_or_update_loader(self): + # 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 + 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) + ) + + def _create_loader(self): """ This method is used here because is resampling is needed during training, there is no need to define to touch the @@ -64,12 +78,6 @@ def _create_or_update_loader(self): self._loader = SamplePointLoader( dataset_phys, dataset_data, batch_size=self.batch_size, shuffle=True ) - pb = self._model.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) - ) def train(self, **kwargs): """ diff --git a/tests/test_callbacks/test_adaptive_refinment_callbacks.py b/tests/test_callbacks/test_adaptive_refinment_callbacks.py index 214257d95..0a1f87826 100644 --- a/tests/test_callbacks/test_adaptive_refinment_callbacks.py +++ b/tests/test_callbacks/test_adaptive_refinment_callbacks.py @@ -4,7 +4,7 @@ from pina.problem import SpatialProblem from pina.operators import laplacian -from pina.geometry import CartesianDomain +from pina.domain import CartesianDomain from pina import Condition, LabelTensor from pina.solvers import PINN from pina.trainer import Trainer diff --git a/tests/test_callbacks/test_optimizer_callbacks.py b/tests/test_callbacks/test_optimizer_callbacks.py index 0b0aabaab..898d3f502 100644 --- a/tests/test_callbacks/test_optimizer_callbacks.py +++ b/tests/test_callbacks/test_optimizer_callbacks.py @@ -4,7 +4,7 @@ from pina.problem import SpatialProblem from pina.operators import laplacian -from pina.geometry import CartesianDomain +from pina.domain import CartesianDomain from pina import Condition, LabelTensor from pina.solvers import PINN from pina.trainer import Trainer diff --git a/tests/test_condition.py b/tests/test_condition.py index 23c9d126b..5f1c6236b 100644 --- a/tests/test_condition.py +++ b/tests/test_condition.py @@ -3,7 +3,7 @@ from pina import LabelTensor, Condition from pina.solvers import PINN -from pina.geometry import CartesianDomain +from pina.domain import CartesianDomain from pina.problem import SpatialProblem from pina.model import FeedForward from pina.operators import laplacian diff --git a/tests/test_dataset.py b/tests/test_dataset.py index cb6a9e4b0..40f219228 100644 --- a/tests/test_dataset.py +++ b/tests/test_dataset.py @@ -4,7 +4,7 @@ from pina.data.dataset import SamplePointDataset, SamplePointLoader, DataPointDataset from pina import LabelTensor, Condition from pina.equation import Equation -from pina.geometry import CartesianDomain +from pina.domain import CartesianDomain from pina.problem import SpatialProblem from pina.model import FeedForward from pina.operators import laplacian diff --git a/tests/test_geometry/test_cartesian.py b/tests/test_geometry/test_cartesian.py index 3e7a8c900..65026c332 100644 --- a/tests/test_geometry/test_cartesian.py +++ b/tests/test_geometry/test_cartesian.py @@ -1,7 +1,7 @@ import torch from pina import LabelTensor -from pina.geometry import CartesianDomain +from pina.domain import CartesianDomain def test_constructor(): CartesianDomain({'x': [0, 1], 'y': [0, 1]}) diff --git a/tests/test_geometry/test_difference.py b/tests/test_geometry/test_difference.py index b165fa710..c5300aae0 100644 --- a/tests/test_geometry/test_difference.py +++ b/tests/test_geometry/test_difference.py @@ -1,7 +1,7 @@ import torch from pina import LabelTensor -from pina.geometry import Difference, EllipsoidDomain, CartesianDomain +from pina.domain import Difference, EllipsoidDomain, CartesianDomain def test_constructor_two_CartesianDomains(): diff --git a/tests/test_geometry/test_ellipsoid.py b/tests/test_geometry/test_ellipsoid.py index 9ab0989ba..fa776f9fc 100644 --- a/tests/test_geometry/test_ellipsoid.py +++ b/tests/test_geometry/test_ellipsoid.py @@ -2,7 +2,7 @@ import pytest from pina import LabelTensor -from pina.geometry import EllipsoidDomain +from pina.domain import EllipsoidDomain def test_constructor(): diff --git a/tests/test_geometry/test_exclusion.py b/tests/test_geometry/test_exclusion.py index b6400cde6..f11fa7f06 100644 --- a/tests/test_geometry/test_exclusion.py +++ b/tests/test_geometry/test_exclusion.py @@ -1,7 +1,7 @@ import torch from pina import LabelTensor -from pina.geometry import Exclusion, EllipsoidDomain, CartesianDomain +from pina.domain import Exclusion, EllipsoidDomain, CartesianDomain def test_constructor_two_CartesianDomains(): diff --git a/tests/test_geometry/test_intersection.py b/tests/test_geometry/test_intersection.py index 61061072f..4929cacad 100644 --- a/tests/test_geometry/test_intersection.py +++ b/tests/test_geometry/test_intersection.py @@ -1,7 +1,7 @@ import torch from pina import LabelTensor -from pina.geometry import Intersection, EllipsoidDomain, CartesianDomain +from pina.domain import Intersection, EllipsoidDomain, CartesianDomain def test_constructor_two_CartesianDomains(): diff --git a/tests/test_geometry/test_simplex.py b/tests/test_geometry/test_simplex.py index 1f59585c6..7fc34ce2d 100644 --- a/tests/test_geometry/test_simplex.py +++ b/tests/test_geometry/test_simplex.py @@ -2,7 +2,7 @@ import pytest from pina import LabelTensor -from pina.geometry import SimplexDomain +from pina.domain import SimplexDomain def test_constructor(): diff --git a/tests/test_geometry/test_union.py b/tests/test_geometry/test_union.py index 16f8bca2e..acde89542 100644 --- a/tests/test_geometry/test_union.py +++ b/tests/test_geometry/test_union.py @@ -1,7 +1,7 @@ import torch from pina import LabelTensor -from pina.geometry import Union, EllipsoidDomain, CartesianDomain +from pina.domain import Union, EllipsoidDomain, CartesianDomain def test_constructor_two_CartesianDomains(): diff --git a/tests/test_plotter.py b/tests/test_plotter.py index 99f99bc7e..98eb08845 100644 --- a/tests/test_plotter.py +++ b/tests/test_plotter.py @@ -1,4 +1,4 @@ -from pina.geometry import CartesianDomain +from pina.domain import CartesianDomain from pina import Condition, Plotter from matplotlib.testing.decorators import image_comparison import matplotlib.pyplot as plt diff --git a/tests/test_problem.py b/tests/test_problem.py index 8dcd49900..3e1e1ee23 100644 --- a/tests/test_problem.py +++ b/tests/test_problem.py @@ -4,7 +4,7 @@ from pina.problem import SpatialProblem from pina.operators import laplacian from pina import LabelTensor, Condition -from pina.geometry import CartesianDomain +from pina.domain import CartesianDomain from pina.equation.equation import Equation from pina.equation.equation_factory import FixedValue diff --git a/tests/test_solvers/test_causalpinn.py b/tests/test_solvers/test_causalpinn.py index 58518ae41..3b1f4b8b4 100644 --- a/tests/test_solvers/test_causalpinn.py +++ b/tests/test_solvers/test_causalpinn.py @@ -3,7 +3,7 @@ from pina.problem import TimeDependentProblem, InverseProblem, SpatialProblem from pina.operators import grad -from pina.geometry import CartesianDomain +from pina.domain import CartesianDomain from pina import Condition, LabelTensor from pina.solvers import CausalPINN from pina.trainer import Trainer diff --git a/tests/test_solvers/test_competitive_pinn.py b/tests/test_solvers/test_competitive_pinn.py index 9c6f6c8ec..97815f5ee 100644 --- a/tests/test_solvers/test_competitive_pinn.py +++ b/tests/test_solvers/test_competitive_pinn.py @@ -3,7 +3,7 @@ from pina.problem import SpatialProblem, InverseProblem from pina.operators import laplacian -from pina.geometry import CartesianDomain +from pina.domain import CartesianDomain from pina import Condition, LabelTensor from pina.solvers import CompetitivePINN as PINN from pina.trainer import Trainer diff --git a/tests/test_solvers/test_gpinn.py b/tests/test_solvers/test_gpinn.py index d00d3b4dc..b4bdc73a9 100644 --- a/tests/test_solvers/test_gpinn.py +++ b/tests/test_solvers/test_gpinn.py @@ -2,7 +2,7 @@ from pina.problem import SpatialProblem, InverseProblem from pina.operators import laplacian -from pina.geometry import CartesianDomain +from pina.domain import CartesianDomain from pina import Condition, LabelTensor from pina.solvers import GPINN from pina.trainer import Trainer diff --git a/tests/test_solvers/test_pinn.py b/tests/test_solvers/test_pinn.py index ea3b077bc..87afae7b0 100644 --- a/tests/test_solvers/test_pinn.py +++ b/tests/test_solvers/test_pinn.py @@ -2,7 +2,7 @@ from pina.problem import SpatialProblem, InverseProblem from pina.operators import laplacian -from pina.geometry import CartesianDomain +from pina.domain import CartesianDomain from pina import Condition, LabelTensor from pina.solvers import PINN from pina.trainer import Trainer diff --git a/tests/test_solvers/test_sapinn.py b/tests/test_solvers/test_sapinn.py index 60c3094ce..1117a0688 100644 --- a/tests/test_solvers/test_sapinn.py +++ b/tests/test_solvers/test_sapinn.py @@ -3,7 +3,7 @@ from pina.problem import SpatialProblem, InverseProblem from pina.operators import laplacian -from pina.geometry import CartesianDomain +from pina.domain import CartesianDomain from pina import Condition, LabelTensor from pina.solvers import SAPINN as PINN from pina.trainer import Trainer diff --git a/tests/test_solvers/test_supervised_solver.py b/tests/test_solvers/test_supervised_solver.py index d9cbea3ce..1b2b7d4b7 100644 --- a/tests/test_solvers/test_supervised_solver.py +++ b/tests/test_solvers/test_supervised_solver.py @@ -11,10 +11,17 @@ 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( - # input_points=LabelTensor(torch.rand(100, 2), input_variables), - # output_points=LabelTensor(torch.rand(100, 1), output_variables)) + 'data' : Condition( + domain='pts', + output_points=LabelTensor( + torch.rand(100, 1), + labels={1: {'name': 'output', 'dof': ['u']}} + ) + ) } class myFeature(torch.nn.Module): @@ -31,8 +38,8 @@ def forward(self, x): return LabelTensor(t, ['sin(x)sin(y)']) -# make the problem + extra feats problem = NeuralOperatorProblem() +# make the problem + extra feats extra_feats = [myFeature()] model = FeedForward(len(problem.input_variables), len(problem.output_variables)) diff --git a/tests/test_utils.py b/tests/test_utils.py index 46305f647..46a5083e6 100644 --- a/tests/test_utils.py +++ b/tests/test_utils.py @@ -3,10 +3,10 @@ from pina.utils import merge_tensors from pina.label_tensor import LabelTensor from pina import LabelTensor -from pina.geometry import EllipsoidDomain, CartesianDomain +from pina.domain import EllipsoidDomain, CartesianDomain from pina.utils import check_consistency import pytest -from pina.geometry import Location +from pina.domain import Location def test_merge_tensors(): diff --git a/tutorials/tutorial1/tutorial.py b/tutorials/tutorial1/tutorial.py index 455ddb5a3..4b41ff74d 100644 --- a/tutorials/tutorial1/tutorial.py +++ b/tutorials/tutorial1/tutorial.py @@ -54,7 +54,7 @@ from pina.problem import SpatialProblem, TimeDependentProblem -from pina.geometry import CartesianDomain +from pina.domain import CartesianDomain class TimeSpaceODE(SpatialProblem, TimeDependentProblem): @@ -86,7 +86,7 @@ class TimeSpaceODE(SpatialProblem, TimeDependentProblem): from pina.problem import SpatialProblem from pina.operators import grad from pina import Condition -from pina.geometry import CartesianDomain +from pina.domain import CartesianDomain from pina.equation import Equation, FixedValue import torch diff --git a/tutorials/tutorial11/tutorial.py b/tutorials/tutorial11/tutorial.py index 1f51f4aca..c0bee1528 100644 --- a/tutorials/tutorial11/tutorial.py +++ b/tutorials/tutorial11/tutorial.py @@ -22,7 +22,7 @@ from pina.model import FeedForward from pina.problem import SpatialProblem from pina.operators import grad -from pina.geometry import CartesianDomain +from pina.domain import CartesianDomain from pina.equation import Equation, FixedValue class SimpleODE(SpatialProblem): diff --git a/tutorials/tutorial12/tutorial.py b/tutorials/tutorial12/tutorial.py index 9b71eb4cb..a5c742297 100644 --- a/tutorials/tutorial12/tutorial.py +++ b/tutorials/tutorial12/tutorial.py @@ -30,7 +30,7 @@ #useful imports from pina.problem import SpatialProblem, TimeDependentProblem from pina.equation import Equation, FixedValue, FixedGradient, FixedFlux -from pina.geometry import CartesianDomain +from pina.domain import CartesianDomain import torch from pina.operators import grad, laplacian from pina import Condition diff --git a/tutorials/tutorial2/tutorial.py b/tutorials/tutorial2/tutorial.py index afc2410e8..9b080d839 100644 --- a/tutorials/tutorial2/tutorial.py +++ b/tutorials/tutorial2/tutorial.py @@ -19,7 +19,7 @@ from pina.solvers import PINN from pina.trainer import Trainer from pina.plotter import Plotter -from pina.geometry import CartesianDomain +from pina.domain import CartesianDomain from pina.equation import Equation, FixedValue from pina import Condition, LabelTensor from pina.callbacks import MetricTracker diff --git a/tutorials/tutorial3/tutorial.py b/tutorials/tutorial3/tutorial.py index 5933b2e26..80ca8d039 100644 --- a/tutorials/tutorial3/tutorial.py +++ b/tutorials/tutorial3/tutorial.py @@ -14,7 +14,7 @@ from pina.problem import SpatialProblem, TimeDependentProblem from pina.operators import laplacian, grad -from pina.geometry import CartesianDomain +from pina.domain import CartesianDomain from pina.solvers import PINN from pina.trainer import Trainer from pina.equation import Equation diff --git a/tutorials/tutorial6/tutorial.py b/tutorials/tutorial6/tutorial.py index 7955a5594..30f958d96 100644 --- a/tutorials/tutorial6/tutorial.py +++ b/tutorials/tutorial6/tutorial.py @@ -15,7 +15,7 @@ import matplotlib.pyplot as plt -from pina.geometry import EllipsoidDomain, Difference, CartesianDomain, Union, SimplexDomain +from pina.domain import EllipsoidDomain, Difference, CartesianDomain, Union, SimplexDomain from pina.label_tensor import LabelTensor def plot_scatter(ax, pts, title): diff --git a/tutorials/tutorial7/tutorial.py b/tutorials/tutorial7/tutorial.py index 3c58fd349..ba652dd25 100644 --- a/tutorials/tutorial7/tutorial.py +++ b/tutorials/tutorial7/tutorial.py @@ -35,7 +35,7 @@ from pina.equation import Equation, FixedValue from pina import Condition, Trainer from pina.solvers import PINN -from pina.geometry import CartesianDomain +from pina.domain import CartesianDomain # Then, we import the pre-saved data, for ($\mu_1$, $\mu_2$)=($0.5$, $0.5$). These two values are the optimal parameters that we want to find through the neural network training. In particular, we import the `input_points`(the spatial coordinates), and the `output_points` (the corresponding $u$ values evaluated at the `input_points`). diff --git a/tutorials/tutorial8/tutorial.py b/tutorials/tutorial8/tutorial.py index 49b0b9340..3e4e2c04a 100644 --- a/tutorials/tutorial8/tutorial.py +++ b/tutorials/tutorial8/tutorial.py @@ -23,7 +23,7 @@ import torch import pina -from pina.geometry import CartesianDomain +from pina.domain import CartesianDomain from pina.problem import ParametricProblem from pina.model.layers import PODBlock diff --git a/tutorials/tutorial9/tutorial.py b/tutorials/tutorial9/tutorial.py index 1c6d72d4c..5bc243a71 100644 --- a/tutorials/tutorial9/tutorial.py +++ b/tutorials/tutorial9/tutorial.py @@ -24,7 +24,7 @@ from pina.model.layers import PeriodicBoundaryEmbedding # The PBC module from pina.solvers import PINN from pina.trainer import Trainer -from pina.geometry import CartesianDomain +from pina.domain import CartesianDomain from pina.equation import Equation From 88ce55ce438775024587c3313031ca4c37e6c295 Mon Sep 17 00:00:00 2001 From: Nicola Demo Date: Mon, 9 Sep 2024 10:50:54 +0200 Subject: [PATCH 14/14] refact --- pina/graph.py | 118 +++++++++++ pina/label_tensor.py | 6 +- pina/loss.py | 209 ------------------- pina/loss/__init__.py | 9 + pina/loss/loss_interface.py | 61 ++++++ pina/loss/lp_loss.py | 78 +++++++ pina/loss/power_loss.py | 79 +++++++ pina/loss/weighted_aggregation.py | 35 ++++ pina/loss/weightning_interface.py | 24 +++ pina/model/layers/messa_passing.py | 11 + pina/solvers/pinns/basepinn.py | 2 +- pina/solvers/supervised.py | 6 +- tests/test_loss/test_lploss.py | 2 +- tests/test_loss/test_powerloss.py | 2 +- tests/test_solvers/test_causalpinn.py | 2 +- tests/test_solvers/test_competitive_pinn.py | 2 +- tests/test_solvers/test_gpinn.py | 2 +- tests/test_solvers/test_pinn.py | 2 +- tests/test_solvers/test_rom_solver.py | 2 +- tests/test_solvers/test_sapinn.py | 2 +- tests/test_solvers/test_supervised_solver.py | 51 ++++- tutorials/tutorial10/tutorial.py | 2 +- tutorials/tutorial5/tutorial.py | 2 +- 23 files changed, 480 insertions(+), 229 deletions(-) create mode 100644 pina/graph.py delete mode 100644 pina/loss.py create mode 100644 pina/loss/__init__.py create mode 100644 pina/loss/loss_interface.py create mode 100644 pina/loss/lp_loss.py create mode 100644 pina/loss/power_loss.py create mode 100644 pina/loss/weighted_aggregation.py create mode 100644 pina/loss/weightning_interface.py create mode 100644 pina/model/layers/messa_passing.py diff --git a/pina/graph.py b/pina/graph.py new file mode 100644 index 000000000..97b2770e6 --- /dev/null +++ b/pina/graph.py @@ -0,0 +1,118 @@ +""" Module for Loss class """ + +import logging +from torch_geometric.nn import MessagePassing, InstanceNorm, radius_graph +from torch_geometric.data import Data +import torch + +class Graph: + """ + PINA Graph managing the PyG Data class. + """ + def __init__(self, data): + self.data = data + + @staticmethod + def _build_triangulation(**kwargs): + logging.debug("Creating graph with triangulation mode.") + + # check for mandatory arguments + if "nodes_coordinates" not in kwargs: + raise ValueError("Nodes coordinates must be provided in the kwargs.") + if "nodes_data" not in kwargs: + raise ValueError("Nodes data must be provided in the kwargs.") + if "triangles" not in kwargs: + raise ValueError("Triangles must be provided in the kwargs.") + + nodes_coordinates = kwargs["nodes_coordinates"] + nodes_data = kwargs["nodes_data"] + triangles = kwargs["triangles"] + + + + def less_first(a, b): + return [a, b] if a < b else [b, a] + + list_of_edges = [] + + for triangle in triangles: + for e1, e2 in [[0, 1], [1, 2], [2, 0]]: + list_of_edges.append(less_first(triangle[e1],triangle[e2])) + + array_of_edges = torch.unique(torch.Tensor(list_of_edges), dim=0) # remove duplicates + array_of_edges = array_of_edges.t().contiguous() + print(array_of_edges) + + # list_of_lengths = [] + + # for p1,p2 in array_of_edges: + # x1, y1 = tri.points[p1] + # x2, y2 = tri.points[p2] + # list_of_lengths.append((x1-x2)**2 + (y1-y2)**2) + + # array_of_lengths = np.sqrt(np.array(list_of_lengths)) + + # return array_of_edges, array_of_lengths + + return Data( + x=nodes_data, + pos=nodes_coordinates.T, + + edge_index=array_of_edges, + ) + + @staticmethod + def _build_radius(**kwargs): + logging.debug("Creating graph with radius mode.") + + # check for mandatory arguments + if "nodes_coordinates" not in kwargs: + raise ValueError("Nodes coordinates must be provided in the kwargs.") + if "nodes_data" not in kwargs: + raise ValueError("Nodes data must be provided in the kwargs.") + if "radius" not in kwargs: + raise ValueError("Radius must be provided in the kwargs.") + + nodes_coordinates = kwargs["nodes_coordinates"] + nodes_data = kwargs["nodes_data"] + radius = kwargs["radius"] + + edges_data = kwargs.get("edge_data", None) + loop = kwargs.get("loop", False) + batch = kwargs.get("batch", None) + + logging.debug(f"radius: {radius}, loop: {loop}, " + f"batch: {batch}") + + edge_index = radius_graph( + x=nodes_coordinates.tensor, + r=radius, + loop=loop, + batch=batch, + ) + + logging.debug(f"edge_index computed") + return Data( + x=nodes_data, + pos=nodes_coordinates, + edge_index=edge_index, + edge_attr=edges_data, + ) + + @staticmethod + def build(mode, **kwargs): + """ + Constructor for the `Graph` class. + """ + if mode == "radius": + graph = Graph._build_radius(**kwargs) + elif mode == "triangulation": + graph = Graph._build_triangulation(**kwargs) + else: + raise ValueError(f"Mode {mode} not recognized") + + return Graph(graph) + + + def __repr__(self): + return f"Graph(data={self.data})" \ No newline at end of file diff --git a/pina/label_tensor.py b/pina/label_tensor.py index 12d2182f5..0a5515365 100644 --- a/pina/label_tensor.py +++ b/pina/label_tensor.py @@ -427,4 +427,8 @@ def stack(tensors): def requires_grad_(self, mode=True): lt = super().requires_grad_(mode) lt.labels = self.labels - return lt \ No newline at end of file + return lt + + @property + def dtype(self): + return super().dtype \ No newline at end of file diff --git a/pina/loss.py b/pina/loss.py deleted file mode 100644 index 3cbf88880..000000000 --- a/pina/loss.py +++ /dev/null @@ -1,209 +0,0 @@ -""" Module for Loss class """ - -from abc import ABCMeta, abstractmethod -from torch.nn.modules.loss import _Loss -import torch -from .utils import check_consistency - -__all__ = ["LossInterface", "LpLoss", "PowerLoss"] - - -class LossInterface(_Loss, metaclass=ABCMeta): - """ - The abstract ``LossInterface`` class. All the class defining a PINA Loss - should be inheritied from this class. - """ - - def __init__(self, reduction="mean"): - """ - :param str reduction: Specifies the reduction to apply to the output: - ``none`` | ``mean`` | ``sum``. When ``none``: no reduction - will be applied, ``mean``: the sum of the output will be divided - by the number of elements in the output, ``sum``: the output will - be summed. Note: ``size_average`` and ``reduce`` are in the - process of being deprecated, and in the meantime, specifying either of - those two args will override ``reduction``. Default: ``mean``. - """ - super().__init__(reduction=reduction, size_average=None, reduce=None) - - @abstractmethod - def forward(self, input, target): - """Forward method for loss function. - - :param torch.Tensor input: Input tensor from real data. - :param torch.Tensor target: Model tensor output. - :return: Loss evaluation. - :rtype: torch.Tensor - """ - pass - - def _reduction(self, loss): - """Simple helper function to check reduction - - :param reduction: Specifies the reduction to apply to the output: - ``none`` | ``mean`` | ``sum``. When ``none``: no reduction - will be applied, ``mean``: the sum of the output will be divided - by the number of elements in the output, ``sum``: the output will - be summed. Note: ``size_average`` and ``reduce`` are in the - process of being deprecated, and in the meantime, specifying either of - those two args will override ``reduction``. Default: ``mean``. - :type reduction: str - :param loss: Loss tensor for each element. - :type loss: torch.Tensor - :return: Reduced loss. - :rtype: torch.Tensor - """ - if self.reduction == "none": - ret = loss - elif self.reduction == "mean": - ret = torch.mean(loss, keepdim=True, dim=-1) - elif self.reduction == "sum": - ret = torch.sum(loss, keepdim=True, dim=-1) - else: - raise ValueError(self.reduction + " is not valid") - return ret - - -class LpLoss(LossInterface): - r""" - The Lp loss implementation class. Creates a criterion that measures - the Lp error between each element in the input :math:`x` and - target :math:`y`. - - The unreduced (i.e. with ``reduction`` set to ``none``) loss can - be described as: - - .. math:: - \ell(x, y) = L = \{l_1,\dots,l_N\}^\top, \quad - l_n = \left[\sum_{i=1}^{D} \left| x_n^i - y_n^i \right|^p \right], - - If ``'relative'`` is set to true: - - .. math:: - \ell(x, y) = L = \{l_1,\dots,l_N\}^\top, \quad - l_n = \frac{ [\sum_{i=1}^{D} | x_n^i - y_n^i|^p] }{[\sum_{i=1}^{D}|y_n^i|^p]}, - - where :math:`N` is the batch size. If ``reduction`` is not ``none`` - (default ``mean``), then: - - .. math:: - \ell(x, y) = - \begin{cases} - \operatorname{mean}(L), & \text{if reduction} = \text{`mean';}\\ - \operatorname{sum}(L), & \text{if reduction} = \text{`sum'.} - \end{cases} - - :math:`x` and :math:`y` are tensors of arbitrary shapes with a total - of :math:`n` elements each. - - The sum operation still operates over all the elements, and divides by :math:`n`. - - The division by :math:`n` can be avoided if one sets ``reduction`` to ``sum``. - """ - - def __init__(self, p=2, reduction="mean", relative=False): - """ - :param int p: Degree of Lp norm. It specifies the type of norm to - be calculated. See `list of possible orders in torch linalg - `_ to - for possible degrees. Default 2 (euclidean norm). - :param str reduction: Specifies the reduction to apply to the output: - ``none`` | ``mean`` | ``sum``. ``none``: no reduction - will be applied, ``mean``: the sum of the output will be divided - by the number of elements in the output, ``sum``: the output will - be summed. - :param bool relative: Specifies if relative error should be computed. - """ - super().__init__(reduction=reduction) - - # check consistency - check_consistency(p, (str, int, float)) - check_consistency(relative, bool) - - self.p = p - self.relative = relative - - def forward(self, input, target): - """Forward method for loss function. - - :param torch.Tensor input: Input tensor from real data. - :param torch.Tensor target: Model tensor output. - :return: Loss evaluation. - :rtype: torch.Tensor - """ - loss = torch.linalg.norm((input - target), ord=self.p, dim=-1) - if self.relative: - loss = loss / torch.linalg.norm(input, ord=self.p, dim=-1) - return self._reduction(loss) - - -class PowerLoss(LossInterface): - r""" - The PowerLoss loss implementation class. Creates a criterion that measures - the error between each element in the input :math:`x` and - target :math:`y` powered to a specific integer. - - The unreduced (i.e. with ``reduction`` set to ``none``) loss can - be described as: - - .. math:: - \ell(x, y) = L = \{l_1,\dots,l_N\}^\top, \quad - l_n = \frac{1}{D}\left[\sum_{i=1}^{D} \left| x_n^i - y_n^i \right|^p \right], - - If ``'relative'`` is set to true: - - .. math:: - \ell(x, y) = L = \{l_1,\dots,l_N\}^\top, \quad - l_n = \frac{ \sum_{i=1}^{D} | x_n^i - y_n^i|^p }{\sum_{i=1}^{D}|y_n^i|^p}, - - where :math:`N` is the batch size. If ``reduction`` is not ``none`` - (default ``mean``), then: - - .. math:: - \ell(x, y) = - \begin{cases} - \operatorname{mean}(L), & \text{if reduction} = \text{`mean';}\\ - \operatorname{sum}(L), & \text{if reduction} = \text{`sum'.} - \end{cases} - - :math:`x` and :math:`y` are tensors of arbitrary shapes with a total - of :math:`n` elements each. - - The sum operation still operates over all the elements, and divides by :math:`n`. - - The division by :math:`n` can be avoided if one sets ``reduction`` to ``sum``. - """ - - def __init__(self, p=2, reduction="mean", relative=False): - """ - :param int p: Degree of Lp norm. It specifies the type of norm to - be calculated. See `list of possible orders in torch linalg - `_ to - see the possible degrees. Default 2 (euclidean norm). - :param str reduction: Specifies the reduction to apply to the output: - ``none`` | ``mean`` | ``sum``. When ``none``: no reduction - will be applied, ``mean``: the sum of the output will be divided - by the number of elements in the output, ``sum``: the output will - be summed. - :param bool relative: Specifies if relative error should be computed. - """ - super().__init__(reduction=reduction) - - # check consistency - check_consistency(p, (str, int, float)) - self.p = p - check_consistency(relative, bool) - self.relative = relative - - def forward(self, input, target): - """Forward method for loss function. - - :param torch.Tensor input: Input tensor from real data. - :param torch.Tensor target: Model tensor output. - :return: Loss evaluation. - :rtype: torch.Tensor - """ - loss = torch.abs((input - target)).pow(self.p).mean(-1) - if self.relative: - loss = loss / torch.abs(input).pow(self.p).mean(-1) - return self._reduction(loss) diff --git a/pina/loss/__init__.py b/pina/loss/__init__.py new file mode 100644 index 000000000..35138125d --- /dev/null +++ b/pina/loss/__init__.py @@ -0,0 +1,9 @@ +__all__ = [ + 'LpLoss', + +] + +from .loss_interface import LossInterface +from .power_loss import PowerLoss +from .lp_loss import LpLoss +from .weightning_interface import weightningInterface \ No newline at end of file diff --git a/pina/loss/loss_interface.py b/pina/loss/loss_interface.py new file mode 100644 index 000000000..5093a65d4 --- /dev/null +++ b/pina/loss/loss_interface.py @@ -0,0 +1,61 @@ +""" Module for Loss Interface """ + +from abc import ABCMeta, abstractmethod +from torch.nn.modules.loss import _Loss +import torch + + +class LossInterface(_Loss, metaclass=ABCMeta): + """ + The abstract ``LossInterface`` class. All the class defining a PINA Loss + should be inheritied from this class. + """ + + def __init__(self, reduction="mean"): + """ + :param str reduction: Specifies the reduction to apply to the output: + ``none`` | ``mean`` | ``sum``. When ``none``: no reduction + will be applied, ``mean``: the sum of the output will be divided + by the number of elements in the output, ``sum``: the output will + be summed. Note: ``size_average`` and ``reduce`` are in the + process of being deprecated, and in the meantime, specifying either of + those two args will override ``reduction``. Default: ``mean``. + """ + super().__init__(reduction=reduction, size_average=None, reduce=None) + + @abstractmethod + def forward(self, input, target): + """Forward method for loss function. + + :param torch.Tensor input: Input tensor from real data. + :param torch.Tensor target: Model tensor output. + :return: Loss evaluation. + :rtype: torch.Tensor + """ + pass + + def _reduction(self, loss): + """Simple helper function to check reduction + + :param reduction: Specifies the reduction to apply to the output: + ``none`` | ``mean`` | ``sum``. When ``none``: no reduction + will be applied, ``mean``: the sum of the output will be divided + by the number of elements in the output, ``sum``: the output will + be summed. Note: ``size_average`` and ``reduce`` are in the + process of being deprecated, and in the meantime, specifying either of + those two args will override ``reduction``. Default: ``mean``. + :type reduction: str + :param loss: Loss tensor for each element. + :type loss: torch.Tensor + :return: Reduced loss. + :rtype: torch.Tensor + """ + if self.reduction == "none": + ret = loss + elif self.reduction == "mean": + ret = torch.mean(loss, keepdim=True, dim=-1) + elif self.reduction == "sum": + ret = torch.sum(loss, keepdim=True, dim=-1) + else: + raise ValueError(self.reduction + " is not valid") + return ret \ No newline at end of file diff --git a/pina/loss/lp_loss.py b/pina/loss/lp_loss.py new file mode 100644 index 000000000..978efa853 --- /dev/null +++ b/pina/loss/lp_loss.py @@ -0,0 +1,78 @@ +""" Module for LpLoss class """ + +import torch + +from ..utils import check_consistency +from .loss_interface import LossInterface + +class LpLoss(LossInterface): + r""" + The Lp loss implementation class. Creates a criterion that measures + the Lp error between each element in the input :math:`x` and + target :math:`y`. + + The unreduced (i.e. with ``reduction`` set to ``none``) loss can + be described as: + + .. math:: + \ell(x, y) = L = \{l_1,\dots,l_N\}^\top, \quad + l_n = \left[\sum_{i=1}^{D} \left| x_n^i - y_n^i \right|^p \right], + + If ``'relative'`` is set to true: + + .. math:: + \ell(x, y) = L = \{l_1,\dots,l_N\}^\top, \quad + l_n = \frac{ [\sum_{i=1}^{D} | x_n^i - y_n^i|^p] }{[\sum_{i=1}^{D}|y_n^i|^p]}, + + where :math:`N` is the batch size. If ``reduction`` is not ``none`` + (default ``mean``), then: + + .. math:: + \ell(x, y) = + \begin{cases} + \operatorname{mean}(L), & \text{if reduction} = \text{`mean';}\\ + \operatorname{sum}(L), & \text{if reduction} = \text{`sum'.} + \end{cases} + + :math:`x` and :math:`y` are tensors of arbitrary shapes with a total + of :math:`n` elements each. + + The sum operation still operates over all the elements, and divides by :math:`n`. + + The division by :math:`n` can be avoided if one sets ``reduction`` to ``sum``. + """ + + def __init__(self, p=2, reduction="mean", relative=False): + """ + :param int p: Degree of Lp norm. It specifies the type of norm to + be calculated. See `list of possible orders in torch linalg + `_ to + for possible degrees. Default 2 (euclidean norm). + :param str reduction: Specifies the reduction to apply to the output: + ``none`` | ``mean`` | ``sum``. ``none``: no reduction + will be applied, ``mean``: the sum of the output will be divided + by the number of elements in the output, ``sum``: the output will + be summed. + :param bool relative: Specifies if relative error should be computed. + """ + super().__init__(reduction=reduction) + + # check consistency + check_consistency(p, (str, int, float)) + check_consistency(relative, bool) + + self.p = p + self.relative = relative + + def forward(self, input, target): + """Forward method for loss function. + + :param torch.Tensor input: Input tensor from real data. + :param torch.Tensor target: Model tensor output. + :return: Loss evaluation. + :rtype: torch.Tensor + """ + loss = torch.linalg.norm((input - target), ord=self.p, dim=-1) + if self.relative: + loss = loss / torch.linalg.norm(input, ord=self.p, dim=-1) + return self._reduction(loss) \ No newline at end of file diff --git a/pina/loss/power_loss.py b/pina/loss/power_loss.py new file mode 100644 index 000000000..4f3fc65c7 --- /dev/null +++ b/pina/loss/power_loss.py @@ -0,0 +1,79 @@ +""" Module for PowerLoss class """ + +import torch + +from ..utils import check_consistency +from .loss_interface import LossInterface + + +class PowerLoss(LossInterface): + r""" + The PowerLoss loss implementation class. Creates a criterion that measures + the error between each element in the input :math:`x` and + target :math:`y` powered to a specific integer. + + The unreduced (i.e. with ``reduction`` set to ``none``) loss can + be described as: + + .. math:: + \ell(x, y) = L = \{l_1,\dots,l_N\}^\top, \quad + l_n = \frac{1}{D}\left[\sum_{i=1}^{D} \left| x_n^i - y_n^i \right|^p \right], + + If ``'relative'`` is set to true: + + .. math:: + \ell(x, y) = L = \{l_1,\dots,l_N\}^\top, \quad + l_n = \frac{ \sum_{i=1}^{D} | x_n^i - y_n^i|^p }{\sum_{i=1}^{D}|y_n^i|^p}, + + where :math:`N` is the batch size. If ``reduction`` is not ``none`` + (default ``mean``), then: + + .. math:: + \ell(x, y) = + \begin{cases} + \operatorname{mean}(L), & \text{if reduction} = \text{`mean';}\\ + \operatorname{sum}(L), & \text{if reduction} = \text{`sum'.} + \end{cases} + + :math:`x` and :math:`y` are tensors of arbitrary shapes with a total + of :math:`n` elements each. + + The sum operation still operates over all the elements, and divides by :math:`n`. + + The division by :math:`n` can be avoided if one sets ``reduction`` to ``sum``. + """ + + def __init__(self, p=2, reduction="mean", relative=False): + """ + :param int p: Degree of Lp norm. It specifies the type of norm to + be calculated. See `list of possible orders in torch linalg + `_ to + see the possible degrees. Default 2 (euclidean norm). + :param str reduction: Specifies the reduction to apply to the output: + ``none`` | ``mean`` | ``sum``. When ``none``: no reduction + will be applied, ``mean``: the sum of the output will be divided + by the number of elements in the output, ``sum``: the output will + be summed. + :param bool relative: Specifies if relative error should be computed. + """ + super().__init__(reduction=reduction) + + # check consistency + check_consistency(p, (str, int, float)) + check_consistency(relative, bool) + + self.p = p + self.relative = relative + + def forward(self, input, target): + """Forward method for loss function. + + :param torch.Tensor input: Input tensor from real data. + :param torch.Tensor target: Model tensor output. + :return: Loss evaluation. + :rtype: torch.Tensor + """ + loss = torch.abs((input - target)).pow(self.p).mean(-1) + if self.relative: + loss = loss / torch.abs(input).pow(self.p).mean(-1) + return self._reduction(loss) \ No newline at end of file diff --git a/pina/loss/weighted_aggregation.py b/pina/loss/weighted_aggregation.py new file mode 100644 index 000000000..357ebf5b9 --- /dev/null +++ b/pina/loss/weighted_aggregation.py @@ -0,0 +1,35 @@ +""" Module for Loss Interface """ + +from .weightning_interface import weightningInterface + + +class WeightedAggregation(WeightningInterface): + """ + TODO + """ + def __init__(self, aggr='mean', weights=None): + self.aggr = aggr + self.weights = weights + + def aggregate(self, losses): + """ + Aggregate the losses. + + :param dict(torch.Tensor) input: The dictionary of losses. + :return: The losses aggregation. It should be a scalar Tensor. + :rtype: torch.Tensor + """ + if self.weights: + weighted_losses = { + condition: self.weights[condition] * losses[condition] + for condition in losses + } + else: + weighted_losses = losses + + if self.aggr == 'mean': + return sum(weighted_losses.values()) / len(weighted_losses) + elif self.aggr == 'sum': + return sum(weighted_losses.values()) + else: + raise ValueError(self.aggr + " is not valid for aggregation.") diff --git a/pina/loss/weightning_interface.py b/pina/loss/weightning_interface.py new file mode 100644 index 000000000..904a819a8 --- /dev/null +++ b/pina/loss/weightning_interface.py @@ -0,0 +1,24 @@ +""" Module for Loss Interface """ + +from abc import ABCMeta, abstractmethod + + +class weightningInterface(metaclass=ABCMeta): + """ + The ``weightingInterface`` class. TODO + """ + + @abstractmethod + def __init__(self, *args, **kwargs): + pass + + @abstractmethod + def aggregate(self, losses): + """ + Aggregate the losses. + + :param list(torch.Tensor) input: The list + :return: The losses aggregation. It should be a scalar Tensor. + :rtype: torch.Tensor + """ + pass \ No newline at end of file diff --git a/pina/model/layers/messa_passing.py b/pina/model/layers/messa_passing.py new file mode 100644 index 000000000..8599734df --- /dev/null +++ b/pina/model/layers/messa_passing.py @@ -0,0 +1,11 @@ +""" Module for Averaging Neural Operator Layer class. """ + +from torch import nn, mean +from torch_geometric.nn import MessagePassing, InstanceNorm, radius_graph + +from pina.utils import check_consistency + + +class MessagePassingBlock(nn.Module): + + diff --git a/pina/solvers/pinns/basepinn.py b/pina/solvers/pinns/basepinn.py index 53d4d3a90..9e841d65d 100644 --- a/pina/solvers/pinns/basepinn.py +++ b/pina/solvers/pinns/basepinn.py @@ -6,7 +6,7 @@ from ...solvers.solver import SolverInterface from pina.utils import check_consistency -from pina.loss import LossInterface +from pina.loss.loss_interface import LossInterface from pina.problem import InverseProblem from torch.nn.modules.loss import _Loss diff --git a/pina/solvers/supervised.py b/pina/solvers/supervised.py index 00a9c8fa0..4f3d497fe 100644 --- a/pina/solvers/supervised.py +++ b/pina/solvers/supervised.py @@ -8,7 +8,7 @@ from .solver import SolverInterface from ..label_tensor import LabelTensor from ..utils import check_consistency -from ..loss import LossInterface +from ..loss.loss_interface import LossInterface class SupervisedSolver(SolverInterface): @@ -172,10 +172,6 @@ def loss_data(self, input_pts, output_pts): :return: The residual loss averaged on the input coordinates :rtype: torch.Tensor """ - print(input_pts) - print(output_pts) - print(self.loss) - print(self.forward(input_pts)) return self.loss(self.forward(input_pts), output_pts) @property diff --git a/tests/test_loss/test_lploss.py b/tests/test_loss/test_lploss.py index 3743970df..3ddd503ee 100644 --- a/tests/test_loss/test_lploss.py +++ b/tests/test_loss/test_lploss.py @@ -1,7 +1,7 @@ import torch import pytest -from pina.loss import * +from pina.loss.loss_interface import * input = torch.tensor([[3.], [1.], [-8.]]) target = torch.tensor([[6.], [4.], [2.]]) diff --git a/tests/test_loss/test_powerloss.py b/tests/test_loss/test_powerloss.py index 7ea26755d..fd5fddee8 100644 --- a/tests/test_loss/test_powerloss.py +++ b/tests/test_loss/test_powerloss.py @@ -1,7 +1,7 @@ import torch import pytest -from pina.loss import PowerLoss +from pina.loss.loss_interface import PowerLoss input = torch.tensor([[3.], [1.], [-8.]]) target = torch.tensor([[6.], [4.], [2.]]) diff --git a/tests/test_solvers/test_causalpinn.py b/tests/test_solvers/test_causalpinn.py index 3b1f4b8b4..887529f24 100644 --- a/tests/test_solvers/test_causalpinn.py +++ b/tests/test_solvers/test_causalpinn.py @@ -10,7 +10,7 @@ from pina.model import FeedForward from pina.equation.equation import Equation from pina.equation.equation_factory import FixedValue -from pina.loss import LpLoss +from pina.loss.loss_interface import LpLoss diff --git a/tests/test_solvers/test_competitive_pinn.py b/tests/test_solvers/test_competitive_pinn.py index 97815f5ee..eb177cd61 100644 --- a/tests/test_solvers/test_competitive_pinn.py +++ b/tests/test_solvers/test_competitive_pinn.py @@ -10,7 +10,7 @@ from pina.model import FeedForward from pina.equation.equation import Equation from pina.equation.equation_factory import FixedValue -from pina.loss import LpLoss +from pina.loss.loss_interface import LpLoss def laplace_equation(input_, output_): diff --git a/tests/test_solvers/test_gpinn.py b/tests/test_solvers/test_gpinn.py index b4bdc73a9..de6555ff8 100644 --- a/tests/test_solvers/test_gpinn.py +++ b/tests/test_solvers/test_gpinn.py @@ -9,7 +9,7 @@ from pina.model import FeedForward from pina.equation.equation import Equation from pina.equation.equation_factory import FixedValue -from pina.loss import LpLoss +from pina.loss.loss_interface import LpLoss def laplace_equation(input_, output_): diff --git a/tests/test_solvers/test_pinn.py b/tests/test_solvers/test_pinn.py index 87afae7b0..8ee9d6124 100644 --- a/tests/test_solvers/test_pinn.py +++ b/tests/test_solvers/test_pinn.py @@ -9,7 +9,7 @@ from pina.model import FeedForward from pina.equation.equation import Equation from pina.equation.equation_factory import FixedValue -from pina.loss import LpLoss +from pina.loss.loss_interface import LpLoss def laplace_equation(input_, output_): diff --git a/tests/test_solvers/test_rom_solver.py b/tests/test_solvers/test_rom_solver.py index a16ffcaae..2a54d6504 100644 --- a/tests/test_solvers/test_rom_solver.py +++ b/tests/test_solvers/test_rom_solver.py @@ -6,7 +6,7 @@ from pina.solvers import ReducedOrderModelSolver from pina.trainer import Trainer from pina.model import FeedForward -from pina.loss import LpLoss +from pina.loss.loss_interface import LpLoss class NeuralOperatorProblem(AbstractProblem): diff --git a/tests/test_solvers/test_sapinn.py b/tests/test_solvers/test_sapinn.py index 1117a0688..4a9be259f 100644 --- a/tests/test_solvers/test_sapinn.py +++ b/tests/test_solvers/test_sapinn.py @@ -10,7 +10,7 @@ from pina.model import FeedForward from pina.equation.equation import Equation from pina.equation.equation_factory import FixedValue -from pina.loss import LpLoss +from pina.loss.loss_interface import LpLoss def laplace_equation(input_, output_): diff --git a/tests/test_solvers/test_supervised_solver.py b/tests/test_solvers/test_supervised_solver.py index 1b2b7d4b7..1bb812027 100644 --- a/tests/test_solvers/test_supervised_solver.py +++ b/tests/test_solvers/test_supervised_solver.py @@ -5,14 +5,17 @@ from pina.solvers import SupervisedSolver from pina.trainer import Trainer from pina.model import FeedForward -from pina.loss import LpLoss +from pina.loss.loss_interface import LpLoss 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']}}) + 'pts': LabelTensor( + torch.rand(100, 2), + labels={1: {'name': 'space', 'dof': ['u_0', 'u_1']}} + ) } conditions = { 'data' : Condition( @@ -56,9 +59,51 @@ def test_constructor(): # 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[0](input.data, edge_index=input.data.edge_index) + g.labels = {1: {'name': 'output', 'dof': ['u']}} + return g + du_dt_new = LabelTensor(self.model[0](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 = AutoSolver(problem = problem, model=GraphModel(2, 1), loss=LpLoss()) + trainer = Trainer(solver=solver, max_epochs=30, accelerator='cpu', batch_size=20) + trainer.train() + assert False + + def test_train_cpu(): solver = SupervisedSolver(problem = problem, model=model, loss=LpLoss()) - trainer = Trainer(solver=solver, max_epochs=3, accelerator='cpu', batch_size=20) + trainer = Trainer(solver=solver, max_epochs=300, accelerator='cpu', batch_size=20) trainer.train() diff --git a/tutorials/tutorial10/tutorial.py b/tutorials/tutorial10/tutorial.py index ea247f6dd..0e7a125f1 100644 --- a/tutorials/tutorial10/tutorial.py +++ b/tutorials/tutorial10/tutorial.py @@ -226,7 +226,7 @@ class NeuralOperatorProblem(AbstractProblem): # In[8]: -from pina.loss import PowerLoss +from pina.loss.loss_interface import PowerLoss error_metric = PowerLoss(p=2) # we use the MSE loss diff --git a/tutorials/tutorial5/tutorial.py b/tutorials/tutorial5/tutorial.py index 3b28a9f12..1b5fd166d 100644 --- a/tutorials/tutorial5/tutorial.py +++ b/tutorials/tutorial5/tutorial.py @@ -99,7 +99,7 @@ class NeuralOperatorSolver(AbstractProblem): # In[19]: -from pina.loss import LpLoss +from pina.loss.loss_interface import LpLoss # make the metric metric_err = LpLoss(relative=True)