Backend JAX supports more examples on forward problems with PINNs (#1595)

Author: ZongrenZou

examples/pinn_forward/Poisson_Dirichlet_1d.py

@@ -1,21 +1,28 @@
-"""Backend supported: tensorflow.compat.v1, tensorflow, pytorch, paddle"""
+"""Backend supported: tensorflow.compat.v1, tensorflow, pytorch, jax, paddle"""
 import deepxde as dde
 import matplotlib.pyplot as plt
 import numpy as np
 # Import tf if using backend tensorflow.compat.v1 or tensorflow
 from deepxde.backend import tf
 # Import torch if using backend pytorch
 # import torch
+# Import jax.numpy if using backend jax
+# import jax.numpy as jnp
 # Import paddle if using backend paddle
 # import paddle
 
 
 def pde(x, y):
+    # Most backends
     dy_xx = dde.grad.hessian(y, x)
+    # Backend jax
+    # dy_xx, _ = dde.grad.hessian(y, x)
     # Use tf.sin for backend tensorflow.compat.v1 or tensorflow
     return -dy_xx - np.pi ** 2 * tf.sin(np.pi * x)
     # Use torch.sin for backend pytorch
     # return -dy_xx - np.pi ** 2 * torch.sin(np.pi * x)
+    # Use jax.numpy.sin for backend jax
+    # return -dy_xx - np.pi ** 2 * jnp.sin(np.pi * x)
     # Use paddle.sin for backend paddle
     # return -dy_xx - np.pi ** 2 * paddle.sin(np.pi * x)
 

examples/pinn_forward/Poisson_Dirichlet_1d_exactBC.py

@@ -1,4 +1,4 @@
-"""Backend supported: tensorflow.compat.v1, tensorflow, paddle"""
+"""Backend supported: tensorflow.compat.v1, tensorflow, jax, paddle"""
 import deepxde as dde
 import numpy as np
 
@@ -10,14 +10,21 @@
     from deepxde.backend import tf
 
     sin = tf.sin
+elif dde.backend.backend_name == "jax":
+    import jax
+
+    sin = jax.numpy.sin
 elif dde.backend.backend_name == "paddle":
     import paddle
 
     sin = paddle.sin
 
 
 def pde(x, y):
+    # Most backends
     dy_xx = dde.grad.hessian(y, x)
+    # Backend jax
+    # dy_xx, _ = dde.grad.hessian(y, x)
     summation = sum([i * sin(i * x) for i in range(1, 5)])
     return -dy_xx - summation - 8 * sin(8 * x)
 

examples/pinn_forward/Poisson_periodic_1d.py

@@ -1,20 +1,27 @@
-"""Backend supported: tensorflow.compat.v1, tensorflow, pytorch, paddle"""
+"""Backend supported: tensorflow.compat.v1, tensorflow, pytorch, jax, paddle"""
 import deepxde as dde
 import numpy as np
 # Import tf if using backend tensorflow.compat.v1 or tensorflow
 from deepxde.backend import tf
 # Import torch if using backend pytorch
 # import torch
+# Import jax.numpy if using backend jax
+# import jax.numpy as jnp
 # Import paddle if using backend paddle
 # import paddle
 
 
 def pde(x, y):
+    # Most backends
     dy_xx = dde.grad.hessian(y, x)
+    # Backend jax
+    # dy_xx, _ = dde.grad.hessian(y, x)
     # Use tf.sin for backend tensorflow.compat.v1 or tensorflow
     return -dy_xx - np.pi**2 * tf.sin(np.pi * x)
     # Use torch.sin for backend pytorch
     # return -dy_xx - np.pi ** 2 * torch.sin(np.pi * x)
+    # Use jax.numpy.sin for backend jax
+    # return -dy_xx - np.pi ** 2 * jnp.sin(np.pi * x)
     # Use paddle.sin for backend paddle
     # return -dy_xx - np.pi ** 2 * paddle.sin(np.pi * x)