Bug fix: PointSetBC supports multi-component outputs (#1058)

Author: mitchelldaneker

deepxde/icbc/boundary_conditions.py

@@ -171,8 +171,9 @@ class PointSetBC:
     Args:
         points: An array of points where the corresponding target values are known and
             used for training.
-        values: A 2D-array of values that gives the exact solution of the problem.
+        values: A scalar or a 2D-array of values that gives the exact solution of the problem.
         component: Integer or a list of integers. The output components satisfying this BC.
+            List of integers only supported for the backend PyTorch.
         batch_size: The number of points per minibatch, or `None` to return all points.
             This is only supported for the backend PyTorch.
         shuffle: Randomize the order on each pass through the data when batching.
@@ -183,10 +184,12 @@ def __init__(
     ):
         self.points = np.array(points, dtype=config.real(np))
         self.values = bkd.as_tensor(values, dtype=config.real(bkd.lib))
-        if isinstance(component, numbers.Number):
-            self.component = [component]
-        else:
-            self.component = component
+        self.component = component
+        if isinstance(component, list) and backend_name != "pytorch":
+            # TODO: Add support for multiple components in other backends
+            raise RuntimeError(
+                "multiple components only implemented for pytorch backend"
+            )
         self.batch_size = batch_size
 
         if batch_size is not None:  # batch iterator and state
@@ -210,10 +213,31 @@ def collocation_points(self, X):
 
     def error(self, X, inputs, outputs, beg, end, aux_var=None):
         if self.batch_size is not None:
+            if isinstance(self.component, numbers.Number):
+                return (
+                    outputs[beg:end, self.component : self.component + 1]
+                    - self.values[self.batch_indices]
+                )
             return (
                 outputs[beg:end, self.component]
                 - self.values[self.batch_indices]
             )
+        if isinstance(self.component, numbers.Number):
+            return (
+                outputs[beg:end, self.component : self.component + 1]
+                - self.values
+            )
+        # When a concat is provided, the following code works 'fast' in paddle cpu,
+        # and slow in both tensorflow backends, jax untested.
+        # tf.gather can be used instead of for loop but is also slow
+        # if len(self.component) > 1:
+        #    calculated_error = outputs[beg:end, self.component[0]] - self.values[:,0]
+        #    for i in range(1,len(self.component)):
+        #        tmp = outputs[beg:end, self.component[i]] - self.values[:,i]
+        #        calculated_error = bkd.lib.concat([calculated_error,tmp],axis=0)
+        # else:
+        #    calculated_error = outputs[beg:end, self.component[0]] - self.values
+        # return calculated_error
         return outputs[beg:end, self.component] - self.values