Vectorial output

pina/condition.py

@@ -15,6 +15,9 @@ class Condition:
             elif isinstance(args[0], Location) and callable(args[1]):
                 self.location = args[0]
                 self.function = args[1]
+            elif isinstance(args[0], Location) and isinstance(args[1], list):
+                self.location = args[0]
+                self.function = args[1]
             else:
                 raise ValueError
 

pina/label_tensor.py

@@ -12,6 +12,9 @@ class LabelTensor():
         self.tensor = x
 
     def __getitem__(self, key):
+        if isinstance(key, (tuple, list)):
+            indeces = [self.labels.index(k) for k in key]
+            return LabelTensor(self.tensor[:, indeces], [self.labels[idx] for idx in indeces])
         if key in self.labels:
             return self.tensor[:, self.labels.index(key)]
         else:

pina/model/feed_forward.py

@@ -52,6 +52,8 @@ class FeedForward(torch.nn.Module):
     def forward(self, x):
         """
         """
+
+        x = x[self.input_variables]
         nf = len(self.extra_features)
         if nf == 0:
             return LabelTensor(self.model(x.tensor), self.output_variables)

pina/pinn.py

@@ -65,8 +65,8 @@ class PINN(object):
         self.model = model
         self.model.to(dtype=self.dtype, device=self.device)
 
-        self.input_pts = {}
         self.truth_values = {}
+        self.input_pts = {}
 
 
         self.trained_epoch = 0
@@ -171,13 +171,15 @@ class PINN(object):
             except:
                 pts = condition.input_points
 
+            print(location, pts)
+
             self.input_pts[location] = pts
-            print(pts.tensor.shape)
             self.input_pts[location].tensor.to(dtype=self.dtype, device=self.device)
             self.input_pts[location].tensor.requires_grad_(True)
             self.input_pts[location].tensor.retain_grad()
 
 
+
     def plot_pts(self, locations='all'):
         import matplotlib
         matplotlib.use('GTK3Agg')
@@ -209,8 +211,13 @@ class PINN(object):
 
                 predicted = self.model(pts)
 
-                residuals = condition.function(pts, predicted)
+                if isinstance(condition.function, list):
-                losses.append(self._compute_norm(residuals))
+                    for function in condition.function:
+                        residuals = function(pts, predicted)
+                        losses.append(self._compute_norm(residuals))
+                else:
+                    residuals = condition.function(pts, predicted)
+                    losses.append(self._compute_norm(residuals))
 
             self.optimizer.zero_grad()
             sum(losses).backward()

pina/plotter.py

@@ -84,11 +84,13 @@ class Plotter:
         """
         res = 256
         pts = obj.problem.domain.sample(res, 'grid')
+        print(pts)
         grids_container = [
-            pts[:, 0].reshape(res, res),
+            pts.tensor[:, 0].reshape(res, res),
-            pts[:, 1].reshape(res, res),
+            pts.tensor[:, 1].reshape(res, res),
         ]
         predicted_output = obj.model(pts)
+        predicted_output = predicted_output['p']
 
         if hasattr(obj.problem, 'truth_solution'):
             truth_output = obj.problem.truth_solution(*pts.tensor.T).float()
@@ -102,10 +104,56 @@ class Plotter:
             fig.colorbar(cb, ax=axes[2])
         else:
             fig, axes = plt.subplots(nrows=1, ncols=1, figsize=(8, 6))
-            cb = getattr(axes, method)(*grids_container, predicted_output.tensor.reshape(res, res).detach())
+            # cb = getattr(axes, method)(*grids_container, predicted_output.tensor.reshape(res, res).detach())
+            cb = getattr(axes, method)(*grids_container, predicted_output.reshape(res, res).detach())
             fig.colorbar(cb, ax=axes)
 
         if filename:
             plt.savefig(filename)
         else:
             plt.show()
+
+
+    def plot(self, obj, method='contourf', filename=None):
+        """
+        """
+        res = 256
+        pts = obj.problem.domain.sample(res, 'grid')
+        print(pts)
+        grids_container = [
+            pts.tensor[:, 0].reshape(res, res),
+            pts.tensor[:, 1].reshape(res, res),
+        ]
+        predicted_output = obj.model(pts)
+        predicted_output = predicted_output['ux']
+
+        if hasattr(obj.problem, 'truth_solution'):
+            truth_output = obj.problem.truth_solution(*pts.tensor.T).float()
+            fig, axes = plt.subplots(nrows=1, ncols=3, figsize=(16, 6))
+
+            cb = getattr(axes[0], method)(*grids_container, predicted_output.tensor.reshape(res, res).detach())
+            fig.colorbar(cb, ax=axes[0])
+            cb = getattr(axes[1], method)(*grids_container, truth_output.reshape(res, res).detach())
+            fig.colorbar(cb, ax=axes[1])
+            cb = getattr(axes[2], method)(*grids_container, (truth_output-predicted_output.tensor.float().flatten()).detach().reshape(res, res))
+            fig.colorbar(cb, ax=axes[2])
+        else:
+            fig, axes = plt.subplots(nrows=1, ncols=1, figsize=(8, 6))
+            # cb = getattr(axes, method)(*grids_container, predicted_output.tensor.reshape(res, res).detach())
+            cb = getattr(axes, method)(*grids_container, predicted_output.reshape(res, res).detach())
+            fig.colorbar(cb, ax=axes)
+
+        if filename:
+            plt.savefig(filename)
+        else:
+            plt.show()
+    
+
+
+    def plot_samples(self, obj):
+
+        for location in obj.input_pts:
+            plt.plot(*obj.input_pts[location].tensor.T.detach(), '.', label=location)
+
+        plt.legend()
+        plt.show()

pina/span.py

@@ -37,6 +37,7 @@ class Span(Location):
                 for _ in range(bounds.shape[0])])
             grids = np.meshgrid(*pts)
             pts = np.hstack([grid.reshape(-1, 1) for grid in grids])
+            print(pts)
         elif mode == 'lh' or mode == 'latin':
             from scipy.stats import qmc
             sampler = qmc.LatinHypercube(d=bounds.shape[0])