Merge branch 'lululxvi:master' into fix-variable-compile

Author: bonneted

.github/workflows/build.yml

@@ -19,7 +19,7 @@ jobs:
     strategy:
       fail-fast: false
       matrix:
-        python-version: ["3.9", "3.10", "3.11"]
+        python-version: ["3.9", "3.10", "3.11", "3.12"]
         os: [ubuntu-latest, macos-latest, windows-latest]
 
     steps:

README.md

@@ -74,7 +74,7 @@ DeepXDE requires one of the following backend-specific dependencies to be instal
 
 - TensorFlow 1.x: [TensorFlow](https://www.tensorflow.org)>=2.7.0
 - TensorFlow 2.x: [TensorFlow](https://www.tensorflow.org)>=2.3.0, [TensorFlow Probability](https://www.tensorflow.org/probability)>=0.11.0
-- PyTorch: [PyTorch](https://pytorch.org)>=1.9.0
+- PyTorch: [PyTorch](https://pytorch.org)>=2.0.0
 - JAX: [JAX](https://jax.readthedocs.io), [Flax](https://flax.readthedocs.io), [Optax](https://optax.readthedocs.io)
 - PaddlePaddle: [PaddlePaddle](https://www.paddlepaddle.org.cn/en)>=2.6.0
 

deepxde/backend/pytorch/tensor.py

@@ -4,8 +4,8 @@
 import torch
 
 
-if Version(torch.__version__) < Version("1.9.0"):
-    raise RuntimeError("DeepXDE requires PyTorch>=1.9.0.")
+if Version(torch.__version__) < Version("2.0.0"):
+    raise RuntimeError("DeepXDE requires PyTorch>=2.0.0.")
 
 # To write device-agnostic (CPU or GPU) code, a common pattern is to first determine
 # torch.device and then use it for all the tensors.

deepxde/data/mf.py

@@ -1,7 +1,8 @@
 import numpy as np
 
 from .data import Data
-from ..backend import tf
+from .. import backend as bkd
+from .. import config
 from ..utils import run_if_any_none, standardize
 
 
@@ -83,20 +84,20 @@ def __init__(
         standardize=False,
     ):
         if X_lo_train is not None:
-            self.X_lo_train = X_lo_train
-            self.X_hi_train = X_hi_train
-            self.y_lo_train = y_lo_train
-            self.y_hi_train = y_hi_train
-            self.X_hi_test = X_hi_test
-            self.y_hi_test = y_hi_test
+            self.X_lo_train = X_lo_train.astype(config.real(np))
+            self.X_hi_train = X_hi_train.astype(config.real(np))
+            self.y_lo_train = y_lo_train.astype(config.real(np))
+            self.y_hi_train = y_hi_train.astype(config.real(np))
+            self.X_hi_test = X_hi_test.astype(config.real(np))
+            self.y_hi_test = y_hi_test.astype(config.real(np))
         elif fname_lo_train is not None:
-            data = np.loadtxt(fname_lo_train)
+            data = np.loadtxt(fname_lo_train).astype(config.real(np))
             self.X_lo_train = data[:, col_x]
             self.y_lo_train = data[:, col_y]
-            data = np.loadtxt(fname_hi_train)
+            data = np.loadtxt(fname_hi_train).astype(config.real(np))
             self.X_hi_train = data[:, col_x]
             self.y_hi_train = data[:, col_y]
-            data = np.loadtxt(fname_hi_test)
+            data = np.loadtxt(fname_hi_test).astype(config.real(np))
             self.X_hi_test = data[:, col_x]
             self.y_hi_test = data[:, col_y]
         else:
@@ -116,7 +117,10 @@ def losses_train(self, targets, outputs, loss_fn, inputs, model, aux=None):
         return [loss_lo, loss_hi]
 
     def losses_test(self, targets, outputs, loss_fn, inputs, model, aux=None):
-        return [0, loss_fn(targets[1], outputs[1])]
+        return [
+            bkd.as_tensor(0, dtype=config.real(bkd.lib)),
+            loss_fn(targets[1], outputs[1]),
+        ]
 
     @run_if_any_none("X_train", "y_train")
     def train_next_batch(self, batch_size=None):

deepxde/data/pde_operator.py

@@ -237,23 +237,59 @@ def __init__(
         self.train_next_batch()
         self.test()
 
-    def _losses(self, outputs, loss_fn, inputs, model, num_func):
+    def _losses(self, outputs, loss_fn, inputs, model, num_func, aux=None):
         bcs_start = np.cumsum([0] + self.pde.num_bcs)
 
         losses = []
-        for i in range(num_func):
-            out = outputs[i]
-            # Single output
-            if bkd.ndim(out) == 1:
-                out = out[:, None]
+        # PDE loss
+        if config.autodiff == "reverse":  # reverse mode AD
+            for i in range(num_func):
+                out = outputs[i]
+                # Single output
+                if bkd.ndim(out) == 1:
+                    out = out[:, None]
+                f = []
+                if self.pde.pde is not None:
+                    f = self.pde.pde(
+                        inputs[1], out, model.net.auxiliary_vars[i][:, None]
+                    )
+                    if not isinstance(f, (list, tuple)):
+                        f = [f]
+                error_f = [fi[bcs_start[-1] :] for fi in f]
+                losses_i = [loss_fn(bkd.zeros_like(error), error) for error in error_f]
+                losses.append(losses_i)
+
+            losses = zip(*losses)
+            # Use stack instead of as_tensor to keep the gradients.
+            losses = [bkd.reduce_mean(bkd.stack(loss, 0)) for loss in losses]
+        elif config.autodiff == "forward":  # forward mode AD
+
+            def forward_call(trunk_input):
+                return aux[0]((inputs[0], trunk_input))
+
             f = []
             if self.pde.pde is not None:
-                f = self.pde.pde(inputs[1], out, model.net.auxiliary_vars[i][:, None])
+                # Each f has the shape (N1, N2)
+                f = self.pde.pde(
+                    inputs[1], (outputs, forward_call), model.net.auxiliary_vars
+                )
                 if not isinstance(f, (list, tuple)):
                     f = [f]
-            error_f = [fi[bcs_start[-1] :] for fi in f]
-            losses_i = [loss_fn(bkd.zeros_like(error), error) for error in error_f]
-
+            # Each error has the shape (N1, ~N2)
+            error_f = [fi[:, bcs_start[-1] :] for fi in f]
+            for error in error_f:
+                error_i = []
+                for i in range(num_func):
+                    error_i.append(loss_fn(bkd.zeros_like(error[i]), error[i]))
+                losses.append(bkd.reduce_mean(bkd.stack(error_i, 0)))
+
+        # BC loss
+        losses_bc = []
+        for i in range(num_func):
+            losses_i = []
+            out = outputs[i]
+            if bkd.ndim(out) == 1:
+                out = out[:, None]
             for j, bc in enumerate(self.pde.bcs):
                 beg, end = bcs_start[j], bcs_start[j + 1]
                 # The same BC points are used for training and testing.
@@ -267,19 +303,21 @@ def _losses(self, outputs, loss_fn, inputs, model, num_func):
                 )
                 losses_i.append(loss_fn(bkd.zeros_like(error), error))
 
-            losses.append(losses_i)
+            losses_bc.append(losses_i)
 
-        losses = zip(*losses)
-        # Use stack instead of as_tensor to keep the gradients.
-        losses = [bkd.reduce_mean(bkd.stack(loss, 0)) for loss in losses]
+        losses_bc = zip(*losses_bc)
+        losses_bc = [bkd.reduce_mean(bkd.stack(loss, 0)) for loss in losses_bc]
+        losses.extend(losses_bc)
         return losses
 
     def losses_train(self, targets, outputs, loss_fn, inputs, model, aux=None):
         num_func = self.num_func if self.batch_size is None else self.batch_size
-        return self._losses(outputs, loss_fn, inputs, model, num_func)
+        return self._losses(outputs, loss_fn, inputs, model, num_func, aux=aux)
 
     def losses_test(self, targets, outputs, loss_fn, inputs, model, aux=None):
-        return self._losses(outputs, loss_fn, inputs, model, len(self.test_x[0]))
+        return self._losses(
+            outputs, loss_fn, inputs, model, len(self.test_x[0]), aux=aux
+        )
 
     def train_next_batch(self, batch_size=None):
         if self.train_x is None:

deepxde/model.py

@@ -518,7 +518,11 @@ def outputs_losses_test(inputs, targets, auxiliary_vars):
             list(self.net.parameters()) + self.external_trainable_variables
         )
         self.opt = optimizers.get(
-            trainable_variables, self.opt_name, learning_rate=lr, decay=decay
+            trainable_variables,
+            self.opt_name,
+            learning_rate=lr,
+            decay=decay,
+            weight_decay=self.net.regularizer,
         )
 
         def train_step(inputs, targets, auxiliary_vars):

deepxde/nn/paddle/__init__.py

@@ -4,11 +4,13 @@
     "DeepONet",
     "DeepONetCartesianProd",
     "FNN",
+    "MfNN",
     "MsFFN",
     "PFNN",
     "STMsFFN",
 ]
 
 from .deeponet import DeepONet, DeepONetCartesianProd
 from .fnn import FNN, PFNN
+from .mfnn import MfNN
 from .msffn import MsFFN, STMsFFN

deepxde/nn/paddle/fnn.py

@@ -3,12 +3,20 @@
 from .nn import NN
 from .. import activations
 from .. import initializers
+from .. import regularizers
 
 
 class FNN(NN):
     """Fully-connected neural network."""
 
-    def __init__(self, layer_sizes, activation, kernel_initializer):
+    def __init__(
+        self,
+        layer_sizes,
+        activation,
+        kernel_initializer,
+        regularization=None,
+        dropout_rate=0,
+    ):
         super().__init__()
         if isinstance(activation, list):
             if not (len(layer_sizes) - 1) == len(activation):
@@ -20,6 +28,13 @@ def __init__(self, layer_sizes, activation, kernel_initializer):
             self.activation = activations.get(activation)
         initializer = initializers.get(kernel_initializer)
         initializer_zero = initializers.get("zeros")
+        self.regularizer = regularizers.get(regularization)
+        self.dropout_rate = dropout_rate
+        if dropout_rate > 0:
+            self.dropouts = [
+                paddle.nn.Dropout(p=dropout_rate)
+                for _ in range(1, len(layer_sizes) - 1)
+            ]
 
         self.linears = paddle.nn.LayerList()
         for i in range(1, len(layer_sizes)):
@@ -37,6 +52,8 @@ def forward(self, inputs):
                 if isinstance(self.activation, list)
                 else self.activation(linear(x))
             )
+            if self.dropout_rate > 0:
+                x = self.dropouts[j](x)
         x = self.linears[-1](x)
         if self._output_transform is not None:
             x = self._output_transform(inputs, x)
@@ -58,11 +75,14 @@ class PFNN(NN):
         kernel_initializer: Initializer for the kernel weights matrix.
     """
 
-    def __init__(self, layer_sizes, activation, kernel_initializer):
+    def __init__(
+        self, layer_sizes, activation, kernel_initializer, regularization=None
+    ):
         super().__init__()
         self.activation = activations.get(activation)
         initializer = initializers.get(kernel_initializer)
         initializer_zero = initializers.get("zeros")
+        self.regularizer = regularizers.get(regularization)
 
         if len(layer_sizes) <= 1:
             raise ValueError("must specify input and output sizes")
@@ -73,7 +93,6 @@ def __init__(self, layer_sizes, activation, kernel_initializer):
 
         n_output = layer_sizes[-1]
 
-
         def make_linear(n_input, n_output):
             linear = paddle.nn.Linear(n_input, n_output)
             initializer(linear.weight)
@@ -92,18 +111,22 @@ def make_linear(n_input, n_output):
                 if isinstance(prev_layer_size, (list, tuple)):
                     # e.g. [8, 8, 8] -> [16, 16, 16]
                     self.layers.append(
-                        paddle.nn.LayerList([
-                            make_linear(prev_layer_size[j], curr_layer_size[j])
-                            for j in range(n_output)
-                        ])
+                        paddle.nn.LayerList(
+                            [
+                                make_linear(prev_layer_size[j], curr_layer_size[j])
+                                for j in range(n_output)
+                            ]
+                        )
                     )
                 else:
                     # e.g. 64 -> [8, 8, 8]
                     self.layers.append(
-                        paddle.nn.LayerList([
-                            make_linear(prev_layer_size, curr_layer_size[j])
-                            for j in range(n_output)
-                        ])
+                        paddle.nn.LayerList(
+                            [
+                                make_linear(prev_layer_size, curr_layer_size[j])
+                                for j in range(n_output)
+                            ]
+                        )
                     )
             else:  # e.g. 64 -> 64
                 if not isinstance(prev_layer_size, int):
@@ -115,10 +138,9 @@ def make_linear(n_input, n_output):
         # output layers
         if isinstance(layer_sizes[-2], (list, tuple)):  # e.g. [3, 3, 3] -> 3
             self.layers.append(
-                paddle.nn.LayerList([
-                    make_linear(layer_sizes[-2][j], 1)
-                    for j in range(n_output)
-                ])
+                paddle.nn.LayerList(
+                    [make_linear(layer_sizes[-2][j], 1) for j in range(n_output)]
+                )
             )
         else:
             self.layers.append(make_linear(layer_sizes[-2], n_output))