Backend JAX: Fix input and output transform (#1705)

Author: bonneted

deepxde/backend/jax/tensor.py

@@ -61,6 +61,14 @@ def to_numpy(input_tensor):
     return np.asarray(input_tensor)
 
 
+def concat(values, axis):
+    return jnp.concatenate(values, axis=axis)
+
+
+def stack(values, axis):
+    return jnp.stack(values, axis=axis)
+
+
 def elu(x):
     return jax.nn.elu(x)
 
@@ -85,6 +93,10 @@ def sin(x):
     return jnp.sin(x)
 
 
+def cos(x):
+    return jnp.cos(x)
+
+
 def square(x):
     return jnp.square(x)
 

deepxde/nn/jax/nn.py

@@ -13,10 +13,30 @@ 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).
         """
-        self._input_transform = transform
+
+        def transform_handling_flat(x):
+            """Handle inputs of shape (n,)"""
+            # TODO: Support tuple or list inputs.
+            if isinstance(x, (list, tuple)):
+                return transform(x)
+            if x.ndim == 1:
+                return transform(x.reshape(1, -1)).squeeze()
+            return transform(x)
+
+        self._input_transform = transform_handling_flat
 
     def apply_output_transform(self, transform):
         """Apply a transform to the network outputs, i.e.,
         outputs = transform(inputs, outputs).
         """
-        self._output_transform = transform
+
+        def transform_handling_flat(inputs, outputs):
+            """Handle inputs of shape (n,)"""
+            # TODO: Support tuple or list inputs.
+            if isinstance(inputs, (list, tuple)):
+                return transform(inputs, outputs)
+            if inputs.ndim == 1:
+                return transform(inputs.reshape(1, -1), outputs.reshape(1, -1)).squeeze()
+            return transform(inputs, outputs)
+
+        self._output_transform = transform_handling_flat

examples/pinn_forward/elasticity_plate.py

@@ -11,24 +11,13 @@
 mu = 0.5
 Q = 4.0
 
-# Define function
-if dde.backend.backend_name == "pytorch":
-    import torch
-
-    sin = torch.sin
-    cos = torch.cos
-elif dde.backend.backend_name == "paddle":
-    import paddle
-
-    sin = paddle.sin
-    cos = paddle.cos
-elif dde.backend.backend_name == "jax":
-    import jax.numpy as jnp
-
-    sin = jnp.sin
-    cos = jnp.cos
+# Define functions
+sin = dde.backend.sin
+cos = dde.backend.cos
+stack = dde.backend.stack
 
 geom = dde.geometry.Rectangle([0, 0], [1, 1])
+BC_type = ["hard", "soft"][0]
 
 
 def boundary_left(x, on_boundary):
@@ -66,6 +55,7 @@ def func(x):
     return np.hstack((ux, uy, Sxx, Syy, Sxy))
 
 
+# Soft Boundary Conditions
 ux_top_bc = dde.icbc.DirichletBC(geom, lambda x: 0, boundary_top, component=0)
 ux_bottom_bc = dde.icbc.DirichletBC(geom, lambda x: 0, boundary_bottom, component=0)
 uy_left_bc = dde.icbc.DirichletBC(geom, lambda x: 0, boundary_left, component=1)
@@ -81,6 +71,17 @@ def func(x):
 )
 
 
+# Hard Boundary Conditions
+def hard_BC(x, f):
+    Ux = f[:, 0] * x[:, 1] * (1 - x[:, 1])
+    Uy = f[:, 1] * x[:, 0] * (1 - x[:, 0]) * x[:, 1]
+
+    Sxx = f[:, 2] * x[:, 0] * (1 - x[:, 0])
+    Syy = f[:, 3] * (1 - x[:, 1]) + (lmbd + 2 * mu) * Q * sin(np.pi * x[:, 0])
+    Sxy = f[:, 4]
+    return stack((Ux, Uy, Sxx, Syy, Sxy), axis=1)
+
+
 def fx(x):
     return (
         -lmbd
@@ -147,10 +148,10 @@ def pde(x, f):
     return [momentum_x, momentum_y, stress_x, stress_y, stress_xy]
 
 
-data = dde.data.PDE(
-    geom,
-    pde,
-    [
+if BC_type == "hard":
+    bcs = []
+else:
+    bcs = [
         ux_top_bc,
         ux_bottom_bc,
         uy_left_bc,
@@ -159,7 +160,12 @@ def pde(x, f):
         sxx_left_bc,
         sxx_right_bc,
         syy_top_bc,
-    ],
+    ]
+
+data = dde.data.PDE(
+    geom,
+    pde,
+    bcs,
     num_domain=500,
     num_boundary=500,
     solution=func,
@@ -170,6 +176,8 @@ def pde(x, f):
 activation = "tanh"
 initializer = "Glorot uniform"
 net = dde.nn.PFNN(layers, activation, initializer)
+if BC_type == "hard":
+    net.apply_output_transform(hard_BC)
 
 model = dde.Model(data, net)
 model.compile("adam", lr=0.001, metrics=["l2 relative error"])