add docs

examples/run_poisson_deeponet.py

@@ -5,11 +5,7 @@ import torch
 from problems.poisson import Poisson
 
 from pina import PINN, LabelTensor, Plotter
-from pina.model.deeponet import DeepONet, check_combos, spawn_combo_networks
-
-logging.basicConfig(
-    filename="poisson_deeponet.log", filemode="w", level=logging.INFO
-)
+from pina.model import DeepONet, FeedForward
 
 
 class SinFeature(torch.nn.Module):
@@ -36,27 +32,33 @@ class SinFeature(torch.nn.Module):
         return LabelTensor(t, [f"sin({self._label})"])
 
 
-def prepare_deeponet_model(args, problem, extra_feature_combo_func=None):
-    combos = tuple(map(lambda combo: combo.split("-"), args.combos.split(",")))
-    check_combos(combos, problem.input_variables)
+class myRBF(torch.nn.Module):
+    def __init__(self, input_):
 
-    extra_feature = extra_feature_combo_func if args.extra else None
-    networks = spawn_combo_networks(
-        combos=combos,
-        layers=list(map(int, args.layers.split(","))) if args.layers else [],
-        output_dimension=args.hidden * len(problem.output_variables),
-        func=torch.nn.Softplus,
-        extra_feature=extra_feature,
-        bias=not args.nobias,
-    )
+        super().__init__()
 
-    return DeepONet(
-        networks,
-        problem.output_variables,
-        aggregator=args.aggregator,
-        reduction=args.reduction,
-    )
+        self.input_variables = [input_]
+        self.a = torch.nn.Parameter(torch.tensor([-.3]))
+        # self.b = torch.nn.Parameter(torch.tensor([0.5]))
+        self.b = torch.tensor([0.5])
+        self.c = torch.nn.Parameter(torch.tensor([.5]))
 
+    def forward(self, x):
+        x = x.extract(self.input_variables)
+        result = self.a * torch.exp(-(x - self.b)**2/(self.c**2))
+        return result
+
+class myModel(torch.nn.Module):
+    def __init__(self):
+
+        super().__init__()
+        self.ffn_x = myRBF('x')
+        self.ffn_y = myRBF('y')
+
+    def forward(self, x):
+        result =  self.ffn_x(x) * self.ffn_y(x)
+        result.labels = ['u']
+        return result
 
 if __name__ == "__main__":
     parser = argparse.ArgumentParser(description="Run PINA")
@@ -65,49 +67,49 @@ if __name__ == "__main__":
     parser.add_argument("id_run", help="Run ID", type=int)
 
     parser.add_argument("--extra", help="Extra features", action="store_true")
-    parser.add_argument("--nobias", action="store_true")
-    parser.add_argument(
-        "--combos",
-        help="DeepONet internal network combinations",
-        type=str,
-        required=True,
-    )
-    parser.add_argument(
-        "--aggregator", help="Aggregator for DeepONet", type=str, default="*"
-    )
-    parser.add_argument(
-        "--reduction", help="Reduction for DeepONet", type=str, default="+"
-    )
-    parser.add_argument(
-        "--hidden",
-        help="Number of variables in the hidden DeepONet layer",
-        type=int,
-        required=True,
-    )
-    parser.add_argument(
-        "--layers",
-        help="Structure of the DeepONet partial layers",
-        type=str,
-        required=True,
-    )
-    cli_args = parser.parse_args()
+    args = parser.parse_args()
 
-    poisson_problem = Poisson()
+    problem = Poisson()
 
-    model = prepare_deeponet_model(
-        cli_args,
-        poisson_problem,
-        extra_feature_combo_func=lambda combo: [SinFeature(combo)],
-    )
-    pinn = PINN(poisson_problem, model, lr=0.01, regularizer=1e-8)
-    if cli_args.save:
+    # ffn_x = FeedForward(
+    #     input_variables=['x'], layers=[], output_variables=1,
+    #     func=torch.nn.Softplus,
+    #     extra_features=[SinFeature('x')]
+    # )
+    # ffn_y = FeedForward
+    #     input_variables=['y'], layers=[], output_variables=1,
+    #     func=torch.nn.Softplus,
+    #     extra_features=[SinFeature('y')]
+    # )
+    model = myModel()
+    test = torch.tensor([[0.0, 0.5]])
+    test.labels = ['x', 'y']
+    pinn = PINN(problem, model, lr=0.0001)
+
+    if args.save:
         pinn.span_pts(
             20, "grid", locations=["gamma1", "gamma2", "gamma3", "gamma4"]
         )
         pinn.span_pts(20, "grid", locations=["D"])
-        pinn.train(1.0e-10, 100)
-        pinn.save_state(f"pina.poisson_{cli_args.id_run}")
-    if cli_args.load:
-        pinn.load_state(f"pina.poisson_{cli_args.id_run}")
+        while True:
+            pinn.train(500, 50)
+            print(model.ffn_x.a)
+            print(model.ffn_x.b)
+            print(model.ffn_x.c)
+
+            xi = torch.linspace(0, 1, 64).reshape(-1, 1).as_subclass(LabelTensor)
+            xi.labels = ['x']
+            yi = model.ffn_x(xi)
+            y_truth = -torch.sin(xi*torch.pi)
+
+            import matplotlib.pyplot as plt
+            plt.plot(xi.detach().flatten(), yi.detach().flatten(), 'r-')
+            plt.plot(xi.detach().flatten(), y_truth.detach().flatten(), 'b-')
+            plt.plot(xi.detach().flatten(), -y_truth.detach().flatten(), 'b-')
+            plt.show()
+        pinn.save_state(f"pina.poisson_{args.id_run}")
+
+    if args.load:
+        pinn.load_state(f"pina.poisson_{args.id_run}")
         plotter = Plotter()
         plotter.plot(pinn)