minor fix

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

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

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'):

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 = {}

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[0], pts[:, 0]>range_x[1]))
-                    ind_y = np.where(np.logical_or(pts[:, 1]<range_y[0], 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()

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

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

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

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

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