Update model.py (#441)

Author: yadhukartha

deepxde/model.py

@@ -528,13 +528,22 @@ def predict(self, x, operator=None, callbacks=None):
         if operator is None:
             y = self._outputs(False, x)
         else:
-            # TODO: predict operator with auxiliary_vars
             if backend_name == "tensorflow.compat.v1":
                 if utils.get_num_args(operator) == 2:
                     op = operator(self.net.inputs, self.net.outputs)
+                    feed_dict = self.net.feed_dict(False, x)
                 elif utils.get_num_args(operator) == 3:
-                    op = operator(self.net.inputs, self.net.outputs, x)
-                y = self.sess.run(op, feed_dict=self.net.feed_dict(False, x))
+                    op = operator(
+                        self.net.inputs, self.net.outputs, self.net.auxiliary_vars
+                    )
+                    feed_dict = self.net.feed_dict(
+                        False,
+                        x,
+                        auxiliary_vars=self.data.auxiliary_var_fn(x).astype(
+                            config.real(np)
+                        ),
+                    )
+                y = self.sess.run(op, feed_dict=feed_dict)
             elif backend_name == "tensorflow":
                 if utils.get_num_args(operator) == 2:
 
@@ -548,7 +557,11 @@ def op(inputs):
                     @tf.function
                     def op(inputs):
                         y = self.net(inputs)
-                        return operator(inputs, y, x)
+                        return operator(
+                            inputs,
+                            y,
+                            self.data.auxiliary_var_fn(x).astype(config.real(np)),
+                        )
 
                 y = op(x)
                 y = utils.to_numpy(y)
@@ -559,7 +572,11 @@ def op(inputs):
                 if utils.get_num_args(operator) == 2:
                     y = operator(inputs, outputs)
                 elif utils.get_num_args(operator) == 3:
-                    y = operator(inputs, outputs, x)
+                    # TODO: Pytorch backend Implementation of Auxiliary variables.
+                    raise NotImplementedError(
+                        "pytorch auxiliary variable not been implemented for this backend."
+                    )
+                    # y = operator(inputs, outputs, torch.as_tensor(self.data.auxiliary_var_fn(x).astype(config.real(np))))
                 y = utils.to_numpy(y)
         self.callbacks.on_predict_end()
         return y