From a05adea4e30091e6e395d77f790c5d332b9bcf5b Mon Sep 17 00:00:00 2001 From: Your Name Date: Wed, 20 Jul 2022 17:23:53 +0200 Subject: [PATCH] minor fix --- pina/label_tensor.py | 31 +++- pina/model/feed_forward.py | 7 +- pina/operators.py | 17 ++- pina/pinn.py | 6 +- pina/plotter.py | 221 +++++++++++++---------------- pina/problem/abstract_problem.py | 25 +++- pina/problem/parametric_problem.py | 7 +- pina/problem/spatial_problem.py | 6 +- pina/problem/timedep_problem.py | 7 +- pina/span.py | 107 ++++++-------- 10 files changed, 231 insertions(+), 203 deletions(-) diff --git a/pina/label_tensor.py b/pina/label_tensor.py index 139678e..334d740 100644 --- a/pina/label_tensor.py +++ b/pina/label_tensor.py @@ -107,10 +107,12 @@ class LabelTensor(torch.Tensor): raise TypeError( '`label_to_extract` should be a str, or a str iterator') - try: - indeces = [self.labels.index(f) for f in label_to_extract] - except ValueError: - raise ValueError('`label_to_extract` not in the labels list') + indeces = [] + for f in label_to_extract: + try: + indeces.append(self.labels.index(f)) + except ValueError: + raise ValueError(f'`{f}` not in the labels list') new_data = self[:, indeces].float() new_labels = [self.labels[idx] for idx in indeces] @@ -118,11 +120,12 @@ class LabelTensor(torch.Tensor): return extracted_tensor - def append(self, lt): + def append(self, lt, mode='std'): """ Return a copy of the merged tensors. :param LabelTensor lt: the tensor to merge. + :param str mode: {'std', 'first', 'cross'} :return: the merged tensors :rtype: LabelTensor """ @@ -130,7 +133,23 @@ class LabelTensor(torch.Tensor): raise RuntimeError('The tensors to merge have common labels') new_labels = self.labels + lt.labels - new_tensor = torch.cat((self, lt), dim=1) + if mode == 'std': + new_tensor = torch.cat((self, lt), dim=1) + elif mode == 'first': + raise NotImplementedError + elif mode == 'cross': + tensor1 = self + tensor2 = lt + n1 = tensor1.shape[0] + n2 = tensor2.shape[0] + + tensor1 = LabelTensor( + tensor1.repeat(n2, 1), + labels=tensor1.labels) + tensor2 = LabelTensor( + tensor2.repeat_interleave(n1, dim=0), + labels=tensor2.labels) + new_tensor = torch.cat((tensor1, tensor2), dim=1) return LabelTensor(new_tensor, new_labels) def __str__(self): diff --git a/pina/model/feed_forward.py b/pina/model/feed_forward.py index f3d887a..7639bf2 100644 --- a/pina/model/feed_forward.py +++ b/pina/model/feed_forward.py @@ -90,13 +90,16 @@ class FeedForward(torch.nn.Module): :return: the output computed by the model. :rtype: LabelTensor """ + if self.input_variables: x = x.extract(self.input_variables) for i, feature in enumerate(self.extra_features): x = x.append(feature(x)) + output = self.model(x) + if self.output_variables: - return LabelTensor(self.model(x), self.output_variables) + return LabelTensor(output, self.output_variables) else: - return self.model(x) + return output diff --git a/pina/operators.py b/pina/operators.py index b87132a..cd1456f 100644 --- a/pina/operators.py +++ b/pina/operators.py @@ -83,14 +83,19 @@ def div(output_, input_, components=None, d=None): raise ValueError grad_output = grad(output_, input_, components, d) - div = torch.empty(input_.shape[0], len(components)) + div = torch.zeros(input_.shape[0], 1) + # print(grad_output) + # print('empty', div) labels = [None] * len(components) - for i, c in enumerate(components): - c_fields = [f'd{c}d{di}' for di in d] - div[:, i] = grad_output.extract(c_fields).sum(axis=1) - labels[i] = '+'.join(c_fields) + for i, (c, d) in enumerate(zip(components, d)): + c_fields = f'd{c}d{d}' + # print(c_fields) + div[:, 0] += grad_output.extract(c_fields).sum(axis=1) + labels[i] = c_fields + # print('full', div) + # print(labels) - return LabelTensor(div, labels) + return LabelTensor(div, ['+'.join(labels)]) def nabla(output_, input_, components=None, d=None, method='std'): diff --git a/pina/pinn.py b/pina/pinn.py index 9a1b24e..8a2527b 100644 --- a/pina/pinn.py +++ b/pina/pinn.py @@ -119,7 +119,6 @@ class PINN(object): return self - def span_pts(self, *args, **kwargs): """ >>> pinn.span_pts(n=10, mode='grid') @@ -141,8 +140,11 @@ class PINN(object): tensor2.repeat_interleave(n1, dim=0), labels=tensor2.labels) return tensor1.append(tensor2) - elif len(tensors): + elif len(tensors) == 1: return tensors[0] + else: + recursive_result = merge_tensors(tensors[1:]) + return merge_tensors([tensors[0], recursive_result]) if isinstance(args[0], int) and isinstance(args[1], str): argument = {} diff --git a/pina/plotter.py b/pina/plotter.py index 88dd2db..f22ac6c 100644 --- a/pina/plotter.py +++ b/pina/plotter.py @@ -13,130 +13,113 @@ from .problem import SpatialProblem, TimeDependentProblem class Plotter: - def _plot_2D(self, obj, method='contourf'): - """ - """ - if not isinstance(obj, PINN): - raise RuntimeError + def plot_samples(self, pinn, variables=None): - res = 256 - pts = obj.problem.spatial_domain.discretize(res, 'grid') - grids_container = [ - pts[:, 0].reshape(res, res), - pts[:, 1].reshape(res, res), - ] - pts = LabelTensor(torch.tensor(pts), obj.problem.input_variables) - predicted_output = obj.model(pts.tensor) + if variables is None: + variables = pinn.problem.domain.variables + elif variables == 'spatial': + variables = pinn.problem.spatial_domain.variables + elif variables == 'temporal': + variables = pinn.problem.temporal_domain.variables - 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)) + if len(variables) not in [1, 2, 3]: + raise ValueError - cb = getattr(axes[0], method)(*grids_container, predicted_output.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.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.reshape(res, res).detach()) - fig.colorbar(cb, ax=axes) + fig = plt.figure() + proj = '3d' if len(variables) == 3 else None + ax = fig.add_subplot(projection=proj) + for location in pinn.input_pts: + coords = pinn.input_pts[location].extract(variables).T.detach() + if coords.shape[0] == 1: # 1D samples + ax.plot(coords[0], torch.zeros(coords[0].shape), '.', + label=location) + else: + ax.plot(*coords, '.', label=location) + ax.set_xlabel(variables[0]) + try: + ax.set_ylabel(variables[1]) + except: + pass - def _plot_1D_TimeDep(self, obj, method='contourf'): - """ - """ - if not isinstance(obj, PINN): - raise RuntimeError - - res = 256 - grids_container = np.meshgrid( - obj.problem.spatial_domain.discretize(res, 'grid'), - obj.problem.temporal_domain.discretize(res, 'grid'), - ) - pts = np.hstack([ - grids_container[0].reshape(-1, 1), - grids_container[1].reshape(-1, 1), - ]) - pts = LabelTensor(torch.tensor(pts), obj.problem.input_variables) - predicted_output = obj.model(pts.tensor) - - 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.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.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.reshape(res, res).detach()) - fig.colorbar(cb, ax=axes) - - - - def plot(self, obj, method='contourf', component='u', parametric=False, params_value=1.5, filename=None): - """ - """ - res = 256 - pts = obj.problem.domain.sample(res, 'grid') - if parametric: - pts_params = torch.ones(pts.shape[0], len(obj.problem.parameters), dtype=pts.dtype)*params_value - pts_params = LabelTensor(pts_params, obj.problem.parameters) - pts = pts.append(pts_params) - grids_container = [ - pts[:, 0].reshape(res, res), - pts[:, 1].reshape(res, res), - ] - ind_dict = {} - all_locations = [condition for condition in obj.problem.conditions] - for location in all_locations: - if hasattr(obj.problem.conditions[location], 'location'): - keys_range_ = obj.problem.conditions[location].location.range_.keys() - if ('x' in keys_range_) and ('y' in keys_range_): - range_x = obj.problem.conditions[location].location.range_['x'] - range_y = obj.problem.conditions[location].location.range_['y'] - ind_x = np.where(np.logical_or(pts[:, 0]range_x[1])) - ind_y = np.where(np.logical_or(pts[:, 1]range_y[1])) - ind_to_exclude = np.union1d(ind_x, ind_y) - ind_dict[location] = ind_to_exclude - import functools - from functools import reduce - - final_inds = reduce(np.intersect1d, ind_dict.values()) - predicted_output = obj.model(pts) - predicted_output = predicted_output.extract([component]) - predicted_output[final_inds] = np.nan - if hasattr(obj.problem, 'truth_solution'): - truth_output = obj.problem.truth_solution(*pts.T).float() - fig, axes = plt.subplots(nrows=1, ncols=3, figsize=(16, 6)) - - cb = getattr(axes[0], method)(*grids_container, predicted_output.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.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.reshape(res, res).detach(), levels=32) - fig.colorbar(cb, ax=axes) - - if filename: - plt.title('Output {} with parameter {}'.format(component, params_value)) - plt.savefig(filename) - else: - plt.show() - - def plot_samples(self, obj): - for location in obj.input_pts: - pts_x = obj.input_pts[location].extract(['x']) - pts_y = obj.input_pts[location].extract(['y']) - plt.plot(pts_x.detach(), pts_y.detach(), '.', label=location) + try: + ax.set_zlabel(variables[2]) + except: + pass plt.legend() plt.show() + + def _1d_plot(self, pts, pred, method, truth_solution=None): + """ + """ + fig, ax = plt.subplots(nrows=1, ncols=1, figsize=(8, 8)) + + ax.plot(pts, pred.detach()) + + if truth_solution: + truth_output = truth_solution(pts).float() + ax.plot(pts, truth_output.detach()) + + plt.xlabel(pts.labels[0]) + plt.ylabel(pred.labels[0]) + plt.show() + + def _2d_plot(self, pts, pred, v, res, method, truth_solution=None): + """ + """ + + grids = [p_.reshape(res, res) for p_ in pts.extract(v).T] + + pred_output = pred.reshape(res, res) + if truth_solution: + truth_output = truth_solution(*pts.T).float().reshape(res, res) + fig, ax = plt.subplots(nrows=1, ncols=3, figsize=(16, 6)) + + cb = getattr(ax[0], method)(*grids, pred_output.detach()) + fig.colorbar(cb, ax=ax[0]) + cb = getattr(ax[1], method)(*grids, truth_output.detach()) + fig.colorbar(cb, ax=ax[1]) + cb = getattr(ax[2], method)(*grids, + (truth_output-pred_output).detach()) + fig.colorbar(cb, ax=ax[2]) + else: + fig, ax = plt.subplots(nrows=1, ncols=1, figsize=(8, 6)) + cb = getattr(ax, method)(*grids, pred_output.detach()) + fig.colorbar(cb, ax=ax) + + + def plot(self, pinn, components, fixed_variables={}, method='contourf', + res=256, filename=None): + """ + """ + + v = [ + var for var in pinn.problem.input_variables + if var not in fixed_variables.keys() + ] + pts = pinn.problem.domain.sample(res, 'grid', variables=v) + + for variable, value in fixed_variables.items(): + new = LabelTensor(torch.ones(pts.shape[0], 1)*value, [variable]) + pts = pts.append(new) + + predicted_output = pinn.model(pts) + if isinstance(components, str): + predicted_output = predicted_output.extract(components) + elif callable(components): + predicted_output = components(predicted_output) + + truth_solution = getattr(pinn.problem, 'truth_solution', None) + if len(v) == 1: + self._1d_plot(pts, predicted_output, method, truth_solution) + elif len(v) == 2: + self._2d_plot(pts, predicted_output, v, res, method, + truth_solution) + + if filename: + plt.title('Output {} with parameter {}'.format(components, + fixed_variables)) + plt.savefig(filename) + else: + plt.show() diff --git a/pina/problem/abstract_problem.py b/pina/problem/abstract_problem.py index e9a0f76..0364873 100644 --- a/pina/problem/abstract_problem.py +++ b/pina/problem/abstract_problem.py @@ -18,9 +18,32 @@ class AbstractProblem(metaclass=ABCMeta): return variables + @property + def domain(self): + + domains = [ + getattr(self, f'{t}_domain') + for t in ['spatial', 'temporal', 'parameter'] + if hasattr(self, f'{t}_domain') + ] + + if len(domains) == 1: + return domains[0] + elif len(domains) == 0: + raise RuntimeError + + if len(set(map(type, domains))) == 1: + domain = domains[0].__class__({}) + [domain.update(d) for d in domains] + return domain + else: + raise RuntimeError('different domains') + + + @input_variables.setter def input_variables(self, variables): - raise NotImplementedError + raise RuntimeError @property @abstractmethod diff --git a/pina/problem/parametric_problem.py b/pina/problem/parametric_problem.py index 67a1cb5..bc11fdb 100644 --- a/pina/problem/parametric_problem.py +++ b/pina/problem/parametric_problem.py @@ -5,7 +5,10 @@ from .abstract_problem import AbstractProblem class ParametricProblem(AbstractProblem): - @property @abstractmethod - def parameters(self): + def parameter_domain(self): pass + + @property + def parameters(self): + return self.parameter_domain.variables diff --git a/pina/problem/spatial_problem.py b/pina/problem/spatial_problem.py index 0e93809..3d236a1 100644 --- a/pina/problem/spatial_problem.py +++ b/pina/problem/spatial_problem.py @@ -6,5 +6,9 @@ from .abstract_problem import AbstractProblem class SpatialProblem(AbstractProblem): @abstractmethod - def spatial_variables(self): + def spatial_domain(self): pass + + @property + def spatial_variables(self): + return self.spatial_domain.variables diff --git a/pina/problem/timedep_problem.py b/pina/problem/timedep_problem.py index 79fee87..2b80fe1 100644 --- a/pina/problem/timedep_problem.py +++ b/pina/problem/timedep_problem.py @@ -5,7 +5,10 @@ from .abstract_problem import AbstractProblem class TimeDependentProblem(AbstractProblem): - @property @abstractmethod - def temporal_variable(self): + def temporal_domain(self): pass + + @property + def temporal_variables(self): + return self.temporal_domain.variables diff --git a/pina/span.py b/pina/span.py index f8e2053..9f83ce7 100644 --- a/pina/span.py +++ b/pina/span.py @@ -20,72 +20,55 @@ class Span(Location): else: raise TypeError + @property + def variables(self): + return list(self.fixed_.keys()) + list(self.range_.keys()) + + def update(self, new_span): + self.fixed_.update(new_span.fixed_) + self.range_.update(new_span.range_) + + def _sample_range(self, n, mode, bounds): + """ + """ + if mode == 'random': + pts = np.random.uniform(size=(n, 1)) + elif mode == 'chebyshev': + pts = np.array([chebyshev_roots(n) * .5 + .5]).reshape(-1, 1) + elif mode == 'grid': + pts = np.linspace(0, 1, n).reshape(-1, 1) + elif mode == 'lh' or mode == 'latin': + from scipy.stats import qmc + sampler = qmc.LatinHypercube(d=1) + pts = sampler.random(n) + + pts *= bounds[1] - bounds[0] + pts += bounds[0] + + pts = pts.astype(np.float32) + return pts + def sample(self, n, mode='random', variables='all'): if variables == 'all': - spatial_range_ = list(self.range_.keys()) - spatial_fixed_ = list(self.fixed_.keys()) - bounds = np.array(list(self.range_.values())) - fixed = np.array(list(self.fixed_.values())) - else: - bounds = [] - spatial_range_ = [] - spatial_fixed_ = [] - fixed = [] - for variable in variables: - if variable in self.range_.keys(): - spatial_range_.append(variable) - bounds.append(list(self.range_[variable])) - elif variable in self.fixed_.keys(): - spatial_fixed_.append(variable) - fixed.append(int(self.fixed_[variable])) - fixed = torch.Tensor(fixed) - bounds = np.array(bounds) + variables = list(self.range_.keys()) + list(self.fixed_.keys()) - if mode == 'random': - pts = np.random.uniform(size=(n, bounds.shape[0])) - elif mode == 'chebyshev': - pts = np.array([ - chebyshev_roots(n) * .5 + .5 - for _ in range(bounds.shape[0])]) - grids = np.meshgrid(*pts) - pts = np.hstack([grid.reshape(-1, 1) for grid in grids]) - elif mode == 'grid': - pts = np.array([ - np.linspace(0, 1, n) - for _ in range(bounds.shape[0])]) - grids = np.meshgrid(*pts) - pts = np.hstack([grid.reshape(-1, 1) for grid in grids]) - elif mode == 'lh' or mode == 'latin': - from scipy.stats import qmc - sampler = qmc.LatinHypercube(d=bounds.shape[0]) - pts = sampler.random(n) + result = None + for variable in variables: + if variable in self.range_.keys(): + bound = np.asarray(self.range_[variable]) + pts_variable = self._sample_range(n, mode, bound) + pts_variable = LabelTensor( + torch.from_numpy(pts_variable), [variable]) - # Scale pts - pts *= bounds[:, 1] - bounds[:, 0] - pts += bounds[:, 0] + elif variable in self.fixed_.keys(): + value = self.fixed_[variable] + pts_variable = LabelTensor(torch.ones(n, 1)*value, [variable]) - pts = pts.astype(np.float32) - pts = torch.from_numpy(pts) + if result is None: + result = pts_variable + else: + intersect = 'std' if mode == 'random' else 'cross' + result = result.append(pts_variable, intersect) - pts_range_ = LabelTensor(pts, spatial_range_) - - if not len(spatial_fixed_)==0: - pts_fixed_ = torch.ones(pts.shape[0], len(spatial_fixed_), - dtype=pts.dtype) * fixed - pts_fixed_ = pts_fixed_.float() - pts_fixed_ = LabelTensor(pts_fixed_, spatial_fixed_) - pts_range_ = pts_range_.append(pts_fixed_) - - return pts_range_ - - - def meshgrid(self, n): - pts = np.array([ - np.linspace(0, 1, n) - for _ in range(self.bound.shape[0])]) - - pts *= self.bound[:, 1] - self.bound[:, 0] - pts += self.bound[:, 0] - - return np.meshgrid(*pts) + return result