Vectorial output

This commit is contained in:
Nicola Demo
2022-03-07 10:09:40 +01:00
parent 1812ddb8d9
commit 8a1f07c8ae
6 changed files with 71 additions and 7 deletions

View File

@@ -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

View File

@@ -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:

View File

@@ -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)

View File

@@ -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()

View File

@@ -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()

View File

@@ -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])