Backend TensorFlow supports auxiliary variables (#348)

Author: ZongrenZou

deepxde/data/pde.py

@@ -171,7 +171,9 @@ def train_next_batch(self, batch_size=None):
             self.train_x = np.vstack((self.train_x, self.train_x_all))
         self.train_y = self.soln(self.train_x) if self.soln else None
         if self.auxiliary_var_fn is not None:
-            self.train_aux_vars = self.auxiliary_var_fn(self.train_x)
+            self.train_aux_vars = self.auxiliary_var_fn(self.train_x).astype(
+                config.real(np)
+            )
         return self.train_x, self.train_y, self.train_aux_vars
 
     @run_if_all_none("test_x", "test_y", "test_aux_vars")
@@ -182,7 +184,9 @@ def test(self):
             self.test_x = self.test_points()
         self.test_y = self.soln(self.test_x) if self.soln else None
         if self.auxiliary_var_fn is not None:
-            self.test_aux_vars = self.auxiliary_var_fn(self.test_x)
+            self.test_aux_vars = self.auxiliary_var_fn(self.test_x).astype(
+                config.real(np)
+            )
         return self.test_x, self.test_y, self.test_aux_vars
 
     def resample_train_points(self):
@@ -203,7 +207,9 @@ def add_anchors(self, anchors):
             self.train_x = np.vstack((self.train_x, self.train_x_all))
         self.train_y = self.soln(self.train_x) if self.soln else None
         if self.auxiliary_var_fn is not None:
-            self.train_aux_vars = self.auxiliary_var_fn(self.train_x)
+            self.train_aux_vars = self.auxiliary_var_fn(self.train_x).astype(
+                config.real(np)
+            )
 
     def train_points(self):
         X = np.empty((0, self.geom.dim), dtype=config.real(np))

deepxde/maps/tensorflow/nn.py

@@ -43,6 +43,10 @@ def auxiliary_vars(self):
         """Tensors: Any additional variables needed."""
         return self._auxiliary_vars
 
+    @auxiliary_vars.setter
+    def auxiliary_vars(self, value):
+        self._auxiliary_vars = value
+
     def apply_feature_transform(self, transform):
         """Compute the features by appling a transform to the network inputs, i.e.,
         features = transform(inputs). Then, outputs = network(features).

deepxde/model.py

@@ -142,20 +142,21 @@ def compute_losses(targets, outputs):
 
             # TODO: Avoid creating multiple graphs by using tf.TensorSpec.
             @tf.function
-            def outputs_losses(data_id, inputs, targets):
+            def outputs_losses(data_id, inputs, targets, auxiliary_vars=None):
                 self.net.data_id = data_id
                 self.net.inputs = inputs
                 self.net.targets = targets
+                self.net.auxiliary_vars = auxiliary_vars
                 outputs = self.net(inputs)
                 losses = compute_losses(targets, outputs)
                 return outputs, losses
 
             opt = optimizers.get(self.opt_name, learning_rate=lr, decay=decay)
 
             @tf.function
-            def train_step(data_id, inputs, targets):
+            def train_step(data_id, inputs, targets, auxiliary_vars=None):
                 with tf.GradientTape() as tape:
-                    _, losses = outputs_losses(data_id, inputs, targets)
+                    _, losses = outputs_losses(data_id, inputs, targets, auxiliary_vars)
                     total_loss = tf.math.reduce_sum(losses)
                 trainable_variables = (
                     self.net.trainable_variables + self.external_trainable_variables
@@ -504,8 +505,8 @@ def _run(
             )
             return self.sess.run(fetches, feed_dict=feed_dict)
         if backend_name == "tensorflow":
-            # TODO: Support training, dropout, auxiliary_vars
-            outs = fetches(data_id, inputs, targets)
+            # TODO: Support training, dropout
+            outs = fetches(data_id, inputs, targets, auxiliary_vars)
             return None if outs is None else [out.numpy() for out in outs]
         if backend_name == "pytorch":
             # TODO: Use torch.no_grad() in _test() and predict()