Backend tensorflow supports all diffusion examples

Author: lululxvi

deepxde/data/pde.py

@@ -107,7 +107,7 @@ def __init__(
                 )
             )
         self.train_distribution = train_distribution
-        self.anchors = anchors
+        self.anchors = None if anchors is None else anchors.astype(config.real(np))
         self.exclusions = exclusions
 
         self.soln = solution
@@ -192,6 +192,7 @@ def resample_train_points(self):
 
     def add_anchors(self, anchors):
         """Add new points for training PDE losses. The BC points will not be updated."""
+        anchors = anchors.astype(config.real(np))
         if self.anchors is None:
             self.anchors = anchors
         else:
@@ -236,7 +237,9 @@ def bc_points(self):
         x_bcs = [bc.collocation_points(self.train_x_all) for bc in self.bcs]
         self.num_bcs = list(map(len, x_bcs))
         self.train_x_bc = (
-            np.vstack(x_bcs) if x_bcs else np.empty([0, self.train_x_all.shape[-1]])
+            np.vstack(x_bcs)
+            if x_bcs
+            else np.empty([0, self.train_x_all.shape[-1]], dtype=config.real(np))
         )
         return self.train_x_bc
 

deepxde/icbcs/boundary_conditions.py

@@ -8,6 +8,7 @@
 
 import numpy as np
 
+from .. import config
 from .. import gradients as grad
 from ..backend import tf
 
@@ -146,12 +147,12 @@ class PointSetBC(object):
     """
 
     def __init__(self, points, values, component=0):
-        self.points = np.array(points)
+        self.points = np.array(points, dtype=config.real(np))
         if not isinstance(values, numbers.Number) and values.shape[1] != 1:
             raise RuntimeError(
                 "PointSetBC should output 1D values. Use argument 'component' for different components."
             )
-        self.values = values
+        self.values = values.astype(config.real(np))
         self.component = component
 
     def collocation_points(self, X):

deepxde/model.py

@@ -77,10 +77,11 @@ def compile(
                 weight the loss contributions. The loss value that will be minimized by
                 the model will then be the weighted sum of all individual losses,
                 weighted by the loss_weights coefficients.
-            external_trainable_variables: A trainable ``tf.Variable`` object or a list of
-                trainable ``tf.Variable`` objects. The unknown parameters in the physics
-                systems that need to be recovered. Trainable variables from the neural
-                networks and external_trainable_variables are trained together.
+            external_trainable_variables: A trainable ``tf.Variable`` object or a list
+                of trainable ``tf.Variable`` objects. The unknown parameters in the
+                physics systems that need to be recovered. If the backend is
+                tensorflow.compat.v1, `external_trainable_variables` is ignored, and all
+                trainable ``tf.Variable`` objects are automatically collected.
         """
         print("Compiling model...")
 
@@ -92,12 +93,18 @@ def compile(
                 self.saver = tf.train.Saver(max_to_keep=None)
 
         self.opt_name = optimizer
-        if external_trainable_variables is not None:
+        if external_trainable_variables is None:
+            self.external_trainable_variables = []
+        else:
+            if backend_name == "tensorflow.compat.v1":
+                print(
+                    "Warning: For the backend tensorflow.compat.v1, "
+                    "`external_trainable_variables` is ignored, and all trainable "
+                    "``tf.Variable`` objects are automatically collected."
+                )
             if not isinstance(external_trainable_variables, list):
                 external_trainable_variables = [external_trainable_variables]
             self.external_trainable_variables = external_trainable_variables
-        else:
-            self.external_trainable_variables = []
 
         loss_fn = losses_module.get(loss)
         if backend_name == "tensorflow.compat.v1":
@@ -135,6 +142,7 @@ def compute_losses(targets, outputs):
 
             opt = optimizers.get(self.opt_name, learning_rate=lr, decay=decay)
 
+            # TODO: Avoid creating multiple graphs by using tf.TensorSpec.
             @tf.function
             def outputs_losses(data_id, inputs, targets):
                 self.net.data_id = data_id

examples/diffusion_1d.py

@@ -1,4 +1,4 @@
-"""Backend supported: tensorflow.compat.v1"""
+"""Backend supported: tensorflow.compat.v1, tensorflow"""
 import deepxde as dde
 import numpy as np
 from deepxde.backend import tf
@@ -19,36 +19,31 @@ def func(x):
     return np.sin(np.pi * x[:, 0:1]) * np.exp(-x[:, 1:])
 
 
-def main():
-    geom = dde.geometry.Interval(-1, 1)
-    timedomain = dde.geometry.TimeDomain(0, 1)
-    geomtime = dde.geometry.GeometryXTime(geom, timedomain)
-
-    bc = dde.DirichletBC(geomtime, func, lambda _, on_boundary: on_boundary)
-    ic = dde.IC(geomtime, func, lambda _, on_initial: on_initial)
-    data = dde.data.TimePDE(
-        geomtime,
-        pde,
-        [bc, ic],
-        num_domain=40,
-        num_boundary=20,
-        num_initial=10,
-        solution=func,
-        num_test=10000,
-    )
-
-    layer_size = [2] + [32] * 3 + [1]
-    activation = "tanh"
-    initializer = "Glorot uniform"
-    net = dde.maps.FNN(layer_size, activation, initializer)
+geom = dde.geometry.Interval(-1, 1)
+timedomain = dde.geometry.TimeDomain(0, 1)
+geomtime = dde.geometry.GeometryXTime(geom, timedomain)
 
-    model = dde.Model(data, net)
+bc = dde.DirichletBC(geomtime, func, lambda _, on_boundary: on_boundary)
+ic = dde.IC(geomtime, func, lambda _, on_initial: on_initial)
+data = dde.data.TimePDE(
+    geomtime,
+    pde,
+    [bc, ic],
+    num_domain=40,
+    num_boundary=20,
+    num_initial=10,
+    solution=func,
+    num_test=10000,
+)
 
-    model.compile("adam", lr=0.001, metrics=["l2 relative error"])
-    losshistory, train_state = model.train(epochs=10000)
+layer_size = [2] + [32] * 3 + [1]
+activation = "tanh"
+initializer = "Glorot uniform"
+net = dde.maps.FNN(layer_size, activation, initializer)
 
-    dde.saveplot(losshistory, train_state, issave=True, isplot=True)
+model = dde.Model(data, net)
 
+model.compile("adam", lr=0.001, metrics=["l2 relative error"])
+losshistory, train_state = model.train(epochs=10000)
 
-if __name__ == "__main__":
-    main()
+dde.saveplot(losshistory, train_state, issave=True, isplot=True)

examples/diffusion_1d_exactBC.py

@@ -1,4 +1,4 @@
-"""Backend supported: tensorflow.compat.v1"""
+"""Backend supported: tensorflow.compat.v1, tensorflow"""
 import deepxde as dde
 import numpy as np
 from deepxde.backend import tf
@@ -19,35 +19,23 @@ def func(x):
     return np.sin(np.pi * x[:, 0:1]) * np.exp(-x[:, 1:])
 
 
-def main():
-    geom = dde.geometry.Interval(-1, 1)
-    timedomain = dde.geometry.TimeDomain(0, 1)
-    geomtime = dde.geometry.GeometryXTime(geom, timedomain)
+geom = dde.geometry.Interval(-1, 1)
+timedomain = dde.geometry.TimeDomain(0, 1)
+geomtime = dde.geometry.GeometryXTime(geom, timedomain)
 
-    data = dde.data.TimePDE(
-        geomtime,
-        pde,
-        [],
-        num_domain=40,
-        solution=func,
-        num_test=10000,
-    )
-
-    layer_size = [2] + [32] * 3 + [1]
-    activation = "tanh"
-    initializer = "Glorot uniform"
-    net = dde.maps.FNN(layer_size, activation, initializer)
-    net.apply_output_transform(
-        lambda x, y: x[:, 1:2] * (1 - x[:, 0:1] ** 2) * y + tf.sin(np.pi * x[:, 0:1])
-    )
-
-    model = dde.Model(data, net)
+data = dde.data.TimePDE(geomtime, pde, [], num_domain=40, solution=func, num_test=10000)
 
-    model.compile("adam", lr=0.001, metrics=["l2 relative error"])
-    losshistory, train_state = model.train(epochs=10000)
+layer_size = [2] + [32] * 3 + [1]
+activation = "tanh"
+initializer = "Glorot uniform"
+net = dde.maps.FNN(layer_size, activation, initializer)
+net.apply_output_transform(
+    lambda x, y: x[:, 1:2] * (1 - x[:, 0:1] ** 2) * y + tf.sin(np.pi * x[:, 0:1])
+)
 
-    dde.saveplot(losshistory, train_state, issave=True, isplot=True)
+model = dde.Model(data, net)
 
+model.compile("adam", lr=0.001, metrics=["l2 relative error"])
+losshistory, train_state = model.train(epochs=10000)
 
-if __name__ == "__main__":
-    main()
+dde.saveplot(losshistory, train_state, issave=True, isplot=True)

examples/diffusion_1d_inverse.py

@@ -1,4 +1,4 @@
-"""Backend supported: tensorflow.compat.v1"""
+"""Backend supported: tensorflow.compat.v1, tensorflow"""
 import deepxde as dde
 import numpy as np
 from deepxde.backend import tf
@@ -22,42 +22,39 @@ def func(x):
     return np.sin(np.pi * x[:, 0:1]) * np.exp(-x[:, 1:])
 
 
-def main():
-    geom = dde.geometry.Interval(-1, 1)
-    timedomain = dde.geometry.TimeDomain(0, 1)
-    geomtime = dde.geometry.GeometryXTime(geom, timedomain)
+geom = dde.geometry.Interval(-1, 1)
+timedomain = dde.geometry.TimeDomain(0, 1)
+geomtime = dde.geometry.GeometryXTime(geom, timedomain)
 
-    bc = dde.DirichletBC(geomtime, func, lambda _, on_boundary: on_boundary)
-    ic = dde.IC(geomtime, func, lambda _, on_initial: on_initial)
+bc = dde.DirichletBC(geomtime, func, lambda _, on_boundary: on_boundary)
+ic = dde.IC(geomtime, func, lambda _, on_initial: on_initial)
 
-    observe_x = np.vstack((np.linspace(-1, 1, num=10), np.full((10), 1))).T
-    observe_y = dde.PointSetBC(observe_x, func(observe_x), component=0)
+observe_x = np.vstack((np.linspace(-1, 1, num=10), np.full((10), 1))).T
+observe_y = dde.PointSetBC(observe_x, func(observe_x), component=0)
 
-    data = dde.data.TimePDE(
-        geomtime,
-        pde,
-        [bc, ic, observe_y],
-        num_domain=40,
-        num_boundary=20,
-        num_initial=10,
-        anchors=observe_x,
-        solution=func,
-        num_test=10000,
-    )
-
-    layer_size = [2] + [32] * 3 + [1]
-    activation = "tanh"
-    initializer = "Glorot uniform"
-    net = dde.maps.FNN(layer_size, activation, initializer)
-
-    model = dde.Model(data, net)
+data = dde.data.TimePDE(
+    geomtime,
+    pde,
+    [bc, ic, observe_y],
+    num_domain=40,
+    num_boundary=20,
+    num_initial=10,
+    anchors=observe_x,
+    solution=func,
+    num_test=10000,
+)
 
-    model.compile("adam", lr=0.001, metrics=["l2 relative error"])
-    variable = dde.callbacks.VariableValue(C, period=1000)
-    losshistory, train_state = model.train(epochs=50000, callbacks=[variable])
+layer_size = [2] + [32] * 3 + [1]
+activation = "tanh"
+initializer = "Glorot uniform"
+net = dde.maps.FNN(layer_size, activation, initializer)
 
-    dde.saveplot(losshistory, train_state, issave=True, isplot=True)
+model = dde.Model(data, net)
 
+model.compile(
+    "adam", lr=0.001, metrics=["l2 relative error"], external_trainable_variables=C
+)
+variable = dde.callbacks.VariableValue(C, period=1000)
+losshistory, train_state = model.train(epochs=50000, callbacks=[variable])
 
-if __name__ == "__main__":
-    main()
+dde.saveplot(losshistory, train_state, issave=True, isplot=True)

examples/diffusion_1d_resample.py

@@ -1,4 +1,4 @@
-"""Backend supported: tensorflow.compat.v1"""
+"""Backend supported: tensorflow.compat.v1, tensorflow"""
 import deepxde as dde
 import numpy as np
 from deepxde.backend import tf
@@ -19,38 +19,33 @@ def func(x):
     return np.sin(np.pi * x[:, 0:1]) * np.exp(-x[:, 1:])
 
 
-def main():
-    geom = dde.geometry.Interval(-1, 1)
-    timedomain = dde.geometry.TimeDomain(0, 1)
-    geomtime = dde.geometry.GeometryXTime(geom, timedomain)
-
-    bc = dde.DirichletBC(geomtime, func, lambda _, on_boundary: on_boundary)
-    ic = dde.IC(geomtime, func, lambda _, on_initial: on_initial)
-    data = dde.data.TimePDE(
-        geomtime,
-        pde,
-        [bc, ic],
-        num_domain=40,
-        num_boundary=20,
-        num_initial=10,
-        train_distribution="pseudo",
-        solution=func,
-        num_test=10000,
-    )
-
-    layer_size = [2] + [32] * 3 + [1]
-    activation = "tanh"
-    initializer = "Glorot uniform"
-    net = dde.maps.FNN(layer_size, activation, initializer)
-
-    model = dde.Model(data, net)
-
-    resampler = dde.callbacks.PDEResidualResampler(period=100)
-    model.compile("adam", lr=0.001, metrics=["l2 relative error"])
-    losshistory, train_state = model.train(epochs=2000, callbacks=[resampler])
-
-    dde.saveplot(losshistory, train_state, issave=True, isplot=True)
-
-
-if __name__ == "__main__":
-    main()
+geom = dde.geometry.Interval(-1, 1)
+timedomain = dde.geometry.TimeDomain(0, 1)
+geomtime = dde.geometry.GeometryXTime(geom, timedomain)
+
+bc = dde.DirichletBC(geomtime, func, lambda _, on_boundary: on_boundary)
+ic = dde.IC(geomtime, func, lambda _, on_initial: on_initial)
+data = dde.data.TimePDE(
+    geomtime,
+    pde,
+    [bc, ic],
+    num_domain=40,
+    num_boundary=20,
+    num_initial=10,
+    train_distribution="pseudo",
+    solution=func,
+    num_test=10000,
+)
+
+layer_size = [2] + [32] * 3 + [1]
+activation = "tanh"
+initializer = "Glorot uniform"
+net = dde.maps.FNN(layer_size, activation, initializer)
+
+model = dde.Model(data, net)
+
+resampler = dde.callbacks.PDEResidualResampler(period=100)
+model.compile("adam", lr=0.001, metrics=["l2 relative error"])
+losshistory, train_state = model.train(epochs=2000, callbacks=[resampler])
+
+dde.saveplot(losshistory, train_state, issave=True, isplot=True)