feat (add examples and tests for 2D monitoring)

Author: fvalle1

examples/pinn_forward/diffusion_1d_with_monitor.py

@@ -0,0 +1,103 @@
+"""Backend supported: tensorflow.compat.v1, tensorflow, pytorch, jax, paddle
+
+Same problem as ``diffusion_1d.py``, but using
+``dde.callbacks.TrainingMonitor`` to watch the predicted solution over the
+2D (x, t) domain and the loss history update live during training, instead
+of only inspecting a plot after training finishes.
+"""
+import deepxde as dde
+import numpy as np
+# Backend tensorflow.compat.v1 or tensorflow
+from deepxde.backend import tf
+# Backend pytorch
+# import torch
+# Backend jax
+# import jax.numpy as jnp
+# Backend paddle
+# import paddle
+
+
+def pde(x, y):
+    # Most backends
+    dy_t = dde.grad.jacobian(y, x, j=1)
+    dy_xx = dde.grad.hessian(y, x, j=0)
+    # Backend jax
+    # dy_t, _ = dde.grad.jacobian(y, x, j=1)
+    # dy_xx, _ = dde.grad.hessian(y, x, j=0)
+    # Backend tensorflow.compat.v1 or tensorflow
+    return (
+        dy_t
+        - dy_xx
+        + tf.exp(-x[:, 1:])
+        * (tf.sin(np.pi * x[:, 0:1]) - np.pi ** 2 * tf.sin(np.pi * x[:, 0:1]))
+    )
+    # Backend pytorch
+    # return (
+    #     dy_t
+    #     - dy_xx
+    #     + torch.exp(-x[:, 1:])
+    #     * (torch.sin(np.pi * x[:, 0:1]) - np.pi ** 2 * torch.sin(np.pi * x[:, 0:1]))
+    # )
+    # Backend jax
+    # return (
+    #     dy_t
+    #     - dy_xx
+    #     + jnp.exp(-x[:, 1:])
+    #     * (jnp.sin(np.pi * x[..., 0:1]) - np.pi ** 2 * jnp.sin(np.pi * x[..., 0:1]))
+    # )
+    # Backend paddle
+    # return (
+    #     dy_t
+    #     - dy_xx
+    #     + paddle.exp(-x[:, 1:])
+    #     * (paddle.sin(np.pi * x[:, 0:1]) - np.pi ** 2 * paddle.sin(np.pi * x[:, 0:1]))
+    # )
+
+
+def func(x):
+    return np.sin(np.pi * x[:, 0:1]) * np.exp(-x[:, 1:])
+
+
+geom = dde.geometry.Interval(-1, 1)
+timedomain = dde.geometry.TimeDomain(0, 1)
+geomtime = dde.geometry.GeometryXTime(geom, timedomain)
+
+bc = dde.icbc.DirichletBC(geomtime, func, lambda _, on_boundary: on_boundary)
+ic = dde.icbc.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.nn.FNN(layer_size, activation, initializer)
+
+model = dde.Model(data, net)
+
+model.compile("adam", lr=0.001, metrics=["l2 relative error"])
+
+# Points (x, t) at which the live plot evaluates and shows the predicted
+# solution as a color scatter, against the reference `func` for comparison.
+x_line = np.linspace(-1, 1, 40)
+t_line = np.linspace(0, 1, 40)
+x_grid, t_grid = np.meshgrid(x_line, t_line)
+x_plot = np.vstack((x_grid.ravel(), t_grid.ravel())).T
+
+monitor = dde.callbacks.TrainingMonitor(
+    period=200,
+    x_plot=x_plot,
+    y_reference=func,
+    show_loss=True,
+)
+
+losshistory, train_state = model.train(iterations=10000, callbacks=[monitor])
+
+dde.saveplot(losshistory, train_state, issave=True, isplot=True)

tests/test_callbacks.py

@@ -44,6 +44,14 @@ def _make_monitor(period=3, show_loss=True):
     )
 
 
+def _make_2d_monitor(period=3, y_reference=True, show_loss=True):
+    x_plot = np.random.rand(20, 2)
+    y_reference = (lambda x: x[:, 0:1] + x[:, 1:2]) if y_reference else None
+    return TrainingMonitor(
+        period=period, x_plot=x_plot, y_reference=y_reference, show_loss=show_loss
+    )
+
+
 def test_training_monitor_requires_x_plot():
     try:
         TrainingMonitor()
@@ -101,5 +109,62 @@ def spy():
     assert monitor.enabled is True
 
 
+def test_training_monitor_rejects_invalid_x_plot_dim():
+    x_plot = np.random.rand(10, 3)
+    try:
+        TrainingMonitor(x_plot=x_plot)
+    except ValueError:
+        pass
+    else:
+        raise AssertionError("`x_plot` with 3 columns should raise ValueError")
+
+
+def test_training_monitor_2d_triggers_and_draws():
+    # Space-time PDEs (e.g. diffusion_1d.py) evaluate the solution over
+    # (x, t) pairs, so `x_plot` has 2 columns instead of 1. This exercises
+    # the scatter/colorbar drawing path used in that case.
+    period = 3
+    monitor = _make_2d_monitor(period=period)
+    monitor._has_display = lambda matplotlib_module: True
+
+    model = _FakeModel()
+    monitor.set_model(model)
+    monitor.on_train_begin()
+    assert monitor.enabled is True
+    assert monitor.dim == 2
+    assert monitor.ax_ref is not None
+
+    calls = []
+    original_update_plot = monitor._update_plot
+
+    def spy():
+        calls.append(model.train_state.iteration)
+        original_update_plot()
+
+    monitor._update_plot = spy
+
+    num_epochs = 10
+    for _ in range(num_epochs):
+        model.train_state.iteration += 1
+        monitor.on_epoch_end()
+
+    assert calls == [period, 2 * period, 3 * period]
+    assert monitor.enabled is True
+
+
+def test_training_monitor_2d_without_reference():
+    monitor = _make_2d_monitor(period=1, y_reference=False)
+    monitor._has_display = lambda matplotlib_module: True
+
+    model = _FakeModel()
+    monitor.set_model(model)
+    monitor.on_train_begin()
+    assert monitor.ax_ref is None
+
+    model.train_state.iteration += 1
+    monitor.on_epoch_end()
+    assert monitor.enabled is True
+
+
 def test_training_monitor_accessible_from_dde_callbacks():
     assert dde.callbacks.TrainingMonitor is TrainingMonitor