From 8a1f07c8ae94aaef7d67de080a1c815af0601d67 Mon Sep 17 00:00:00 2001 From: Nicola Demo Date: Mon, 7 Mar 2022 10:09:40 +0100 Subject: [PATCH] Vectorial output --- pina/condition.py | 3 +++ pina/label_tensor.py | 3 +++ pina/model/feed_forward.py | 2 ++ pina/pinn.py | 15 ++++++++--- pina/plotter.py | 54 +++++++++++++++++++++++++++++++++++--- pina/span.py | 1 + 6 files changed, 71 insertions(+), 7 deletions(-) diff --git a/pina/condition.py b/pina/condition.py index 398c387..27eb42b 100644 --- a/pina/condition.py +++ b/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 diff --git a/pina/label_tensor.py b/pina/label_tensor.py index 16eec26..c9e219b 100644 --- a/pina/label_tensor.py +++ b/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: diff --git a/pina/model/feed_forward.py b/pina/model/feed_forward.py index 6095a97..96e82bb 100644 --- a/pina/model/feed_forward.py +++ b/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) diff --git a/pina/pinn.py b/pina/pinn.py index 916f2be..fc3d233 100644 --- a/pina/pinn.py +++ b/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) - losses.append(self._compute_norm(residuals)) + if isinstance(condition.function, list): + 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() diff --git a/pina/plotter.py b/pina/plotter.py index 15add45..c70d956 100644 --- a/pina/plotter.py +++ b/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[:, 1].reshape(res, res), + pts.tensor[:, 0].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() diff --git a/pina/span.py b/pina/span.py index c8f33a6..f3c55e5 100644 --- a/pina/span.py +++ b/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])