deeponet tutorial

Author: dario-coscia

pina/model/deeponet.py

@@ -52,7 +52,8 @@ def __init__(
         :param reduction: The reduction to be used to reduce the aggregated
             result of the modules in ``networks`` to the desired output
             dimension. Available reductions include: sum: ``+``, product: ``*``,
-            mean: ``mean``, min: ``min``, max: ``max``. Default is ``+``.
+            mean: ``mean``, min: ``min``, max: ``max``, identity: "id".
+            Default is ``+``.
         :type reduction: str or Callable
         :param bool scale: If ``True``, the final output is scaled before being
             returned in the forward pass. Default is ``True``.
@@ -122,19 +123,6 @@ def __init__(
         check_consistency(scale, bool)
         check_consistency(translation, bool)
 
-        # check trunk branch nets consistency
-        shapes = []
-        for key, value in networks.items():
-            check_consistency(value, (str, int))
-            check_consistency(key, torch.nn.Module)
-            input_ = torch.rand(10, len(value))
-            shapes.append(key(input_).shape[-1])
-
-        if not all(map(lambda x: x == shapes[0], shapes)):
-            raise ValueError(
-                "The passed networks have not the same output dimension."
-            )
-
         # assign trunk and branch net with their input indeces
         self.models = torch.nn.ModuleList(networks.keys())
         self._indeces = networks.values()
@@ -171,6 +159,7 @@ def _symbol_functions(**kwargs):
             "mean": partial(torch.mean, **kwargs),
             "min": lambda x: torch.min(x, **kwargs).values,
             "max": lambda x: torch.max(x, **kwargs).values,
+            "id": lambda x: x,
         }
 
     def _init_aggregator(self, aggregator):
@@ -181,7 +170,7 @@ def _init_aggregator(self, aggregator):
         :type aggregator: str or Callable
         :raises ValueError: If the aggregator is not supported.
         """
-        aggregator_funcs = self._symbol_functions(dim=2)
+        aggregator_funcs = self._symbol_functions(dim=-1)
         if aggregator in aggregator_funcs:
             aggregator_func = aggregator_funcs[aggregator]
         elif isinstance(aggregator, nn.Module) or is_function(aggregator):
@@ -264,13 +253,9 @@ def forward(self, x):
         # reduce
         output_ = self._reduction(aggregated)
         if self._reduction_type in self._symbol_functions(dim=-1):
-            output_ = output_.reshape(-1, 1)
-
-        # scale and translate
-        output_ *= self._scale
-        output_ += self._trasl
+            output_ = output_.reshape(*output_.shape, 1)
 
-        return output_
+        return self._scale * output_ + self._trasl
 
     @property
     def aggregator(self):

tests/test_model/test_deeponet.py

@@ -9,7 +9,7 @@
 data = torch.rand((20, 3))
 input_vars = ["a", "b", "c"]
 input_ = LabelTensor(data, input_vars)
-symbol_funcs_red = DeepONet._symbol_functions(dim=-1)
+symbol_funcs_red = DeepONet._symbol_functions()
 output_dims = [1, 5, 10, 20]
 
 
@@ -26,20 +26,6 @@ def test_constructor():
     )
 
 
-def test_constructor_fails_when_invalid_inner_layer_size():
-    branch_net = FeedForward(input_dimensions=1, output_dimensions=10)
-    trunk_net = FeedForward(input_dimensions=2, output_dimensions=8)
-    with pytest.raises(ValueError):
-        DeepONet(
-            branch_net=branch_net,
-            trunk_net=trunk_net,
-            input_indeces_branch_net=["a"],
-            input_indeces_trunk_net=["b", "c"],
-            reduction="+",
-            aggregator="*",
-        )
-
-
 def test_forward_extract_str():
     branch_net = FeedForward(input_dimensions=1, output_dimensions=10)
     trunk_net = FeedForward(input_dimensions=2, output_dimensions=10)

tests/test_model/test_mionet.py

@@ -18,15 +18,6 @@ def test_constructor():
     MIONet(networks=networks, reduction="+", aggregator="*")
 
 
-def test_constructor_fails_when_invalid_inner_layer_size():
-    branch_net1 = FeedForward(input_dimensions=1, output_dimensions=10)
-    branch_net2 = FeedForward(input_dimensions=2, output_dimensions=10)
-    trunk_net = FeedForward(input_dimensions=1, output_dimensions=12)
-    networks = {branch_net1: ["x"], branch_net2: ["x", "y"], trunk_net: ["z"]}
-    with pytest.raises(ValueError):
-        MIONet(networks=networks, reduction="+", aggregator="*")
-
-
 def test_forward_extract_str():
     branch_net1 = FeedForward(input_dimensions=1, output_dimensions=10)
     branch_net2 = FeedForward(input_dimensions=1, output_dimensions=10)