move POD-NN model

pod_nn/experiments/pod_nn.yaml

@@ -0,0 +1,22 @@
+model:
+  model_class: model.pod_nn.PODNN
+  pod_rank: 128
+  layers: [16, 64, 128, 128, 64, 16]
+   
+optimizer:
+  optimizer_class: torch.optim.AdamW
+  lr: 0.001
+
+trainer: 
+  max_epochs: 10000
+  batch_size: 128
+  train_size: 0.9
+  val_size: 0.1
+  accelerator: 'cuda'
+  devices: 1
+  log_every_n_steps: 0
+  patience: 100
+  log_dir: logs/
+  name: pod_nn
+  # version: null
+  

pod_nn/model/pod_nn.py

@@ -0,0 +1,22 @@
+import torch
+from pina.model.block import PODBlock
+from pina.model import FeedForward
+
+class PODNN(torch.nn.Module):
+    def __init__(self, pod_rank, layers, func=torch.nn.Softplus):
+        super().__init__()
+        self.pod = PODBlock(pod_rank, scale_coefficients=False)
+        self.nn = FeedForward(
+            input_dimensions=3,
+            output_dimensions=pod_rank,
+            layers=layers,
+            func=func,
+        )
+        
+
+    def forward(self, p):
+        coefficients = self.nn(p)
+        return self.pod.expand(coefficients)
+
+    def fit_pod(self, x):
+        self.pod.fit(x)