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

add multiple outputs possibility in DeepONet

Browse Source
This commit is contained in:
Anna Ivagnes
2024-10-03 17:00:23 +02:00
committed by Nicola Demo
parent bb44c022e9
commit 5d2ca62e65
2 changed files with 36 additions and 3 deletions
Download Patch File Download Diff File Expand all files Collapse all files

6
pina/model/deeponet.py
View File

@@ -172,6 +172,7 @@ class MIONet(torch.nn.Module):
raise ValueError(f"Unsupported aggregation: {str(aggregator)}")
self._aggregator = aggregator_func
self._aggregator_type = aggregator
def _init_reduction(self, reduction):
reduction_funcs = DeepONet._symbol_functions(dim=-1)
@@ -183,6 +184,7 @@ class MIONet(torch.nn.Module):
raise ValueError(f"Unsupported reduction: {reduction}")
self._reduction = reduction_func
self._reduction_type = reduction
def _get_vars(self, x, indeces):
if isinstance(indeces[0], str):
@@ -222,7 +224,9 @@ class MIONet(torch.nn.Module):
aggregated = self._aggregator(torch.dstack(output_))
# reduce
output_ = self._reduction(aggregated).reshape(-1, 1)
output_ = self._reduction(aggregated)
if self._reduction_type in DeepONet._symbol_functions(dim=-1):
output_ = output_.reshape(-1, 1)
# scale and translate
output_ *= self._scale