224 lines
7.7 KiB
Python
224 lines
7.7 KiB
Python
""" 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 <https://doi.org/10.1016/j.cma.2024.116813>`_.
|
|
|
|
.. 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)
|