In [ ]:

# ML Packages
import jax
import jax.numpy as jnp
import equinox as eqx  # NNs and other useful bits --> https://docs.kidger.site/equinox/
import optax

# Other packages
import os
import time
from tqdm import tqdm

# ML Packages import jax import jax.numpy as jnp import equinox as eqx # NNs and other useful bits --> https://docs.kidger.site/equinox/ import optax # Other packages import os import time from tqdm import tqdm

In [ ]:

os.environ["CUDA_VISIBLE_DEVICES"] = "1"

os.environ["CUDA_VISIBLE_DEVICES"] = "1"

In [ ]:

# Enable 64-bit precision
jax.config.update("jax_enable_x64", True)
# JAX Key to enforce reproducibility
key = jax.random.PRNGKey(1234)

# Enable 64-bit precision jax.config.update("jax_enable_x64", True) # JAX Key to enforce reproducibility key = jax.random.PRNGKey(1234)

Neural Operator Training¶

In this notebook, we will solve the 2D heat equation described previously using a DeepONet model. In this notebook we will learn the following:

• How to format our dataset for DeepONet training
• How to train a DeepONet
• How to define a DeepONet module in JAX and Equinox
• How to define a training algorithm for training ML models
• How to visualize results

Recap: The Heat Equation¶

The heat equation, as described in our first example, will be solved for a 2D domain, $\Omega = (0, 10) \times (0, 10)$ and $t > 0$. Recall that this transient PDE is given as: \begin{equation} \tag{1} \frac{\partial u}{\partial t} - \alpha (\frac{\partial^2 u}{\partial x^2} + \frac{\partial^2 u}{\partial y^2}) = 0 \quad x \in [0, 100] \quad y \in [0, 100] \quad t \in [0, 20] \end{equation} where $u [K]$ is the temperature, $x, y [m]$ are the spatial coordinates, $t [s]$ is time, and $\alpha$ is the thermal diffusivity of the domain.

It has Dirichlet BCs on the bottom, top and left sides given as: $$ u(x, 0, t) = u(x, 1, t) = u(0, y, t) = 0 K $$

With the right side given as: $$ u(1, y, t) = 100 K $$

It has an IC inside throughout the domain given by: $$ u(x, y, 0) = 0 K \quad x \in [0, 100] \quad y \in [0, 100] $$

Training our DeepONet¶

Following our discussion in the Neural-Operators notebook, a standard DeepONet is trained using pairwise samples of input and output functions, $\{a^{(i)}, u^{(i)}\}^{N}_{i}$. The main steps involved in training is described as follows:

• $N$ representative input functions, $a^{(i)}$, are selected at random from the set $1 \le i \le N$. These functions are then evaluated at $m$ sensor locations, i.e. $a_{j}^{(i)} = a^{(i)}(x_{j})$ from the set $1 \le j \le m$.
• For each $a^{(i)}$, the corresponding solution functions, $u^{(i)}$, is determined (here using finite differences for the heat equation).
• We the sample $u^{(i)}$ at $R$ random locations, i.e. $u_{k}^{(i)} = u^{(i)}(y_{k})$ from the set $1 \le k \le R$.
• The training set, $S$, can then be built such that: \begin{equation} \tag{2} S = { (a_{j}^{(i)}, y^{(k)}, u_{k}^{(i)})} \quad 1 \le i \le N, 1 \le j \le M, 1 \le k \le R \end{equation} which will contain $N \times R$ training samples.
• The loss function can then be defined to determine the differences between the predictions, $\hat{u}_{k}^{(i)}$, with the true value, $u_{k}^{(i)}$.
• Training is then performed to minimized this loss function and find the optimal parameters, $\theta$ for our DeepONet model.

Our Datasets¶

The input function, $a$, for our NO will be based on different $\alpha$ values. We will generate multiple heat PDE solutions, that vary based on a constant $\alpha$ throughout the 2D domain. This is done in the Data-Generation notebook.

The dataset is described as follows:

• We sample $i=10$ input functions, $a^{(i)}$, for evaluations.
• Each $a^{(i)}$ is represented as a straight horizontal line, i.e. a constant function. This allows us to only need to sample $M=1$ sensor points.
• Our $a^{(i)}$ dataset is therefore only $10$ samples.
• The corresponding $u^{(i)}$ is then represented on a 2D mesh at $100 \times 100$ nodal locations for $20$ time steps, which will be randomly sampled during training.
• This represents a $u^{(i)}$ dataset of $10 \times 100 \times 100 \times 20$ samples.

Our aim is to predict the solution, $u^{(i)}$, throughout our 2D domain at any time step for different diffusivity, $a^{(i)}$.

In [ ]:

dataset = jnp.load("./data/data.npz", allow_pickle=True)
a = dataset["a"]
u = dataset["u_noise"]
mesh = dataset["mesh"]
info = dataset["info"].item()

for k, v in info.items():
    print(f"{k}: {v}")

dataset = jnp.load("./data/data.npz", allow_pickle=True) a = dataset["a"] u = dataset["u_noise"] mesh = dataset["mesh"] info = dataset["info"].item() for k, v in info.items(): print(f"{k}: {v}")

In [ ]:

# Extract Max-Min
t_max, t_min = info["t_max"], info["t_min"]
x_max, x_min = info["x_max"], info["x_min"]
y_max, y_min = info["y_max"], info["y_min"]

# Extract Max-Min t_max, t_min = info["t_max"], info["t_min"] x_max, x_min = info["x_max"], info["x_min"] y_max, y_min = info["y_max"], info["y_min"]

In [ ]:

def norm_data(data, max, min, **kwargs):
    """Min-Max Normalization between -1, 1

    Args:
        data (array): Dataset to normalize
        max (float): Maximum of dataset
        min (float): Minimum of dataset

    Returns:
        array: Normalized dataset
    """
    data = 2 * ((data - min) / (max - min)) - 1
    return data

def norm_data(data, max, min, **kwargs): """Min-Max Normalization between -1, 1 Args: data (array): Dataset to normalize max (float): Maximum of dataset min (float): Minimum of dataset Returns: array: Normalized dataset """ data = 2 * ((data - min) / (max - min)) - 1 return data

In [ ]:

mesh[:, 0] = norm_data(data=mesh[:, 0], max=t_max, min=t_min)
mesh[:, 1] = norm_data(data=mesh[:, 1], max=x_max, min=x_min)
mesh[:, 2] = norm_data(data=mesh[:, 2], max=y_max, min=y_min)

mesh[:, 0] = norm_data(data=mesh[:, 0], max=t_max, min=t_min) mesh[:, 1] = norm_data(data=mesh[:, 1], max=x_max, min=x_min) mesh[:, 2] = norm_data(data=mesh[:, 2], max=y_max, min=y_min)

In [ ]:

b_mesh = jnp.repeat(mesh[jnp.newaxis, :], repeats=a.shape[0], axis=0)

b_mesh = jnp.repeat(mesh[jnp.newaxis, :], repeats=a.shape[0], axis=0)

In [ ]:

# a.shape -> BATCH, 1
# u.shape -> BATCH, TIME STEPS, NX, NY
# mesh.shape -> BATCH, TIME STEPS * NX * NY, TXY
a.shape, u.shape, b_mesh.shape

# a.shape -> BATCH, 1 # u.shape -> BATCH, TIME STEPS, NX, NY # mesh.shape -> BATCH, TIME STEPS * NX * NY, TXY a.shape, u.shape, b_mesh.shape

We will use a training and test dataset split of 20%.

In [ ]:

split = int(0.20 * a.shape[0])
a_train, a_test = a[split:], a[:split]
u_train, u_test = u[split:], u[:split]
mesh_train, mesh_test = b_mesh[split:], b_mesh[:split]

split = int(0.20 * a.shape[0]) a_train, a_test = a[split:], a[:split] u_train, u_test = u[split:], u[:split] mesh_train, mesh_test = b_mesh[split:], b_mesh[:split]

In [ ]:

b_train, t_train, y_train = a_train, mesh_train, u_train
b_test, t_test, y_test = a_test, mesh_test, u_test

b_train, t_train, y_train = a_train, mesh_train, u_train b_test, t_test, y_test = a_test, mesh_test, u_test

In [ ]:

(
    a_train.shape,
    a_test.shape,
    u_train.shape,
    u_test.shape,
    mesh_train.shape,
    mesh_test.shape,
)

( a_train.shape, a_test.shape, u_train.shape, u_test.shape, mesh_train.shape, mesh_test.shape, )

In [ ]:

def dataloader(datasets, batch_size, key):
    """Load training batch

    Args:
        datasets (Array): Training & test sets
        batch_size (_int): Number of samples to load

    Yields:
        Iterator: Training batch samples
    """
    n_samples = datasets[0].shape[0]
    assert all(d.shape[0] == n_samples for d in datasets)
    n_batches = int(jnp.ceil(n_samples / batch_size))

    # Return the permuted array 'n_samples'
    perm = jax.random.permutation(key, n_samples)

    # Loop over dataset
    if batch_size < 0:
        yield tuple(d[perm] for d in datasets)
    else:
        for batch_id in range(n_batches):
            start = batch_id * batch_size
            end = min((batch_id + 1) * batch_size, n_samples)
            batch_indices = perm[start:end]

            yield tuple(d[batch_indices] for d in datasets)

def dataloader(datasets, batch_size, key): """Load training batch Args: datasets (Array): Training & test sets batch_size (_int): Number of samples to load Yields: Iterator: Training batch samples """ n_samples = datasets[0].shape[0] assert all(d.shape[0] == n_samples for d in datasets) n_batches = int(jnp.ceil(n_samples / batch_size)) # Return the permuted array 'n_samples' perm = jax.random.permutation(key, n_samples) # Loop over dataset if batch_size < 0: yield tuple(d[perm] for d in datasets) else: for batch_id in range(n_batches): start = batch_id * batch_size end = min((batch_id + 1) * batch_size, n_samples) batch_indices = perm[start:end] yield tuple(d[batch_indices] for d in datasets)

Defining our DeepONet¶

In [ ]:

class DeepONet2d(eqx.Module):
    """2D DeepONet model definition where the Branch & Trunk nets are configured as MLPs"""

    branch_net: eqx.nn.MLP
    trunk_net: eqx.nn.MLP

    def __init__(
        self,
        in_size_branch,
        in_size_trunk,
        out_size,
        width_size,
        depth,
        activation,
        *,
        key,
    ):
        b_key, t_key = jax.random.split(key, num=2)
        self.branch_net = eqx.nn.MLP(
            in_size=in_size_branch,
            out_size=out_size,
            width_size=width_size,
            depth=depth,
            activation=activation,
            key=b_key,
        )
        self.trunk_net = eqx.nn.MLP(
            in_size=in_size_trunk,
            out_size=out_size,
            width_size=width_size,
            depth=depth,
            activation=activation,
            final_activation=activation,
            key=t_key,
        )

    def __call__(self, x_branch, x_trunk):
        """
        x_branch.shape = (in_size_branch,)
        x_trunk.shape = (in_size_trunk,)
        return u_pred: Prediction for u
        """
        branch_out = self.branch_net(x_branch)
        # VMAP -> Loop over (t, x, y)
        trunk_out = jax.vmap(self.trunk_net)(x_trunk)

        # Compute dot product per time step
        inner_product = jnp.sum(branch_out * trunk_out, keepdims=True, axis=1)

        # Arrange for MSE
        inner_product = inner_product.reshape(11, 21, 21, 1)
        inner_product = inner_product.swapaxes(
            0, -1
        ).squeeze()  # time evaluate first in trunk
        inner_product = inner_product.swapaxes(
            0, 1
        ).squeeze()  # y evaluate second in trunk

        return inner_product

    @staticmethod
    def save_model(filename, model):
        """Save hyperparameters and leaves of model in the same file

        Args:
            filename (string): Eg. "model.eqx"
            model (eqx.model)
        """
        # Create save path && directory
        MODEL_DIR = "results/trained"
        MODEL_PATH = os.path.join(MODEL_DIR, filename)
        os.makedirs(MODEL_DIR, exist_ok=True)

        with open(MODEL_PATH, "wb") as f:
            eqx.tree_serialise_leaves(f, model)

class DeepONet2d(eqx.Module): """2D DeepONet model definition where the Branch & Trunk nets are configured as MLPs""" branch_net: eqx.nn.MLP trunk_net: eqx.nn.MLP def __init__( self, in_size_branch, in_size_trunk, out_size, width_size, depth, activation, *, key, ): b_key, t_key = jax.random.split(key, num=2) self.branch_net = eqx.nn.MLP( in_size=in_size_branch, out_size=out_size, width_size=width_size, depth=depth, activation=activation, key=b_key, ) self.trunk_net = eqx.nn.MLP( in_size=in_size_trunk, out_size=out_size, width_size=width_size, depth=depth, activation=activation, final_activation=activation, key=t_key, ) def __call__(self, x_branch, x_trunk): """ x_branch.shape = (in_size_branch,) x_trunk.shape = (in_size_trunk,) return u_pred: Prediction for u """ branch_out = self.branch_net(x_branch) # VMAP -> Loop over (t, x, y) trunk_out = jax.vmap(self.trunk_net)(x_trunk) # Compute dot product per time step inner_product = jnp.sum(branch_out * trunk_out, keepdims=True, axis=1) # Arrange for MSE inner_product = inner_product.reshape(11, 21, 21, 1) inner_product = inner_product.swapaxes( 0, -1 ).squeeze() # time evaluate first in trunk inner_product = inner_product.swapaxes( 0, 1 ).squeeze() # y evaluate second in trunk return inner_product @staticmethod def save_model(filename, model): """Save hyperparameters and leaves of model in the same file Args: filename (string): Eg. "model.eqx" model (eqx.model) """ # Create save path && directory MODEL_DIR = "results/trained" MODEL_PATH = os.path.join(MODEL_DIR, filename) os.makedirs(MODEL_DIR, exist_ok=True) with open(MODEL_PATH, "wb") as f: eqx.tree_serialise_leaves(f, model)

In [ ]:

@eqx.filter_value_and_grad
def compute_loss(model, branch_x, trunk_x, y):
    """MSE Loss

    Args:
        model (eqx.Module): Equinox module
        branch_x (eqx.Module): Branch input function
        trunk_x (eqx.Module): Trunk solution coordinates
        y (Array): Ground-truth

    Returns:
        mse: MSE Loss
    """
    pred_y = jax.vmap(model)(branch_x, trunk_x)
    sqrd_diff = jnp.square(pred_y - y)
    mse = jnp.mean(sqrd_diff)
    return jnp.sqrt(mse)

@eqx.filter_value_and_grad def compute_loss(model, branch_x, trunk_x, y): """MSE Loss Args: model (eqx.Module): Equinox module branch_x (eqx.Module): Branch input function trunk_x (eqx.Module): Trunk solution coordinates y (Array): Ground-truth Returns: mse: MSE Loss """ pred_y = jax.vmap(model)(branch_x, trunk_x) sqrd_diff = jnp.square(pred_y - y) mse = jnp.mean(sqrd_diff) return jnp.sqrt(mse)

In [ ]:

@eqx.filter_jit
def make_step(model, branch_x, trunk_x, y, optimizer, opt_state):
    """One training step

    Args:
        model (eqx.Module): Equinox module
        branch_x (eqx.Module): Branch input function
        trunk_x (eqx.Module): Trunk solution coordinates
        y (Array): Ground-truth
        optimizer (NamedTuple): Optax optimizer
        opt_state (Array): Optimizer state at previous training step
    """
    loss, grad = compute_loss(model, branch_x, trunk_x, y)
    updates, new_state = optimizer.update(grad, opt_state)
    new_model = eqx.apply_updates(model, updates)
    return loss, new_model, new_state

@eqx.filter_jit def make_step(model, branch_x, trunk_x, y, optimizer, opt_state): """One training step Args: model (eqx.Module): Equinox module branch_x (eqx.Module): Branch input function trunk_x (eqx.Module): Trunk solution coordinates y (Array): Ground-truth optimizer (NamedTuple): Optax optimizer opt_state (Array): Optimizer state at previous training step """ loss, grad = compute_loss(model, branch_x, trunk_x, y) updates, new_state = optimizer.update(grad, opt_state) new_model = eqx.apply_updates(model, updates) return loss, new_model, new_state

In [ ]:

def train(model, train_data_full, test_data_full, steps, batch_size, optimizer, key):
    b_train, t_train, y_train = train_data_full
    b_test, t_test, y_test = test_data_full

    # Initialize optimiser && model params
    model_params = eqx.filter(model, eqx.is_array)
    opt_state = optimizer.init(model_params)

    # ------------------------------- Track results
    results = {}
    results.get("opt_state", opt_state)
    results.get("model_params", model_params)

    _metrics = {"train_loss": [], "test_loss": []}
    metrics = results.get("metrics", _metrics)

    _best_loss = jnp.inf
    best_loss = results.get("best_loss", _best_loss)

    # Training loop
    start_time = time.time()

    for step in tqdm(range(steps)):
        # ------------------------------- Training step
        train_data = dataloader((b_train, t_train, y_train), batch_size, key)

        for branch_x, trunk_x, y in train_data:
            train_loss, model, opt_state = make_step(
                model, branch_x, trunk_x, y, optimizer, opt_state
            )
            train_loss = train_loss.item()
            metrics["train_loss"].append(train_loss)

        # --------------------------------- Test step -
        test_data = dataloader((b_test, t_test, y_test), -1, key)

        for branch_x, trunk_x, y in test_data:
            test_loss = compute_loss(model, branch_x, trunk_x, y)
            test_loss = test_loss[0]
            metrics["test_loss"].append(test_loss)

        # ------------------------------- Track results
        if test_loss < best_loss:
            best_loss = test_loss
            model_params = model

        time_taken = time.time() - start_time

        if (((step + 1) % 100) == 0) or (step == steps):
            print(
                f"Epoch={step + 1} | Train_loss={train_loss:.8f} | Best_loss={best_loss:.8f} | Test_loss:{test_loss:.8f} | Time:{time_taken:.3f}s"
            )

            # Save model
            DeepONet2d.save_model(filename=f"dnet-{step}.eqx", model=model)

            results = {
                "metrics": metrics,
                "best_loss": best_loss,
                "model_params": model_params,
                "last_epoch": step + 1,
                "time_taken": time_taken,
            }

    return model, results

def train(model, train_data_full, test_data_full, steps, batch_size, optimizer, key): b_train, t_train, y_train = train_data_full b_test, t_test, y_test = test_data_full # Initialize optimiser && model params model_params = eqx.filter(model, eqx.is_array) opt_state = optimizer.init(model_params) # ------------------------------- Track results results = {} results.get("opt_state", opt_state) results.get("model_params", model_params) _metrics = {"train_loss": [], "test_loss": []} metrics = results.get("metrics", _metrics) _best_loss = jnp.inf best_loss = results.get("best_loss", _best_loss) # Training loop start_time = time.time() for step in tqdm(range(steps)): # ------------------------------- Training step train_data = dataloader((b_train, t_train, y_train), batch_size, key) for branch_x, trunk_x, y in train_data: train_loss, model, opt_state = make_step( model, branch_x, trunk_x, y, optimizer, opt_state ) train_loss = train_loss.item() metrics["train_loss"].append(train_loss) # --------------------------------- Test step - test_data = dataloader((b_test, t_test, y_test), -1, key) for branch_x, trunk_x, y in test_data: test_loss = compute_loss(model, branch_x, trunk_x, y) test_loss = test_loss[0] metrics["test_loss"].append(test_loss) # ------------------------------- Track results if test_loss < best_loss: best_loss = test_loss model_params = model time_taken = time.time() - start_time if (((step + 1) % 100) == 0) or (step == steps): print( f"Epoch={step + 1} | Train_loss={train_loss:.8f} | Best_loss={best_loss:.8f} | Test_loss:{test_loss:.8f} | Time:{time_taken:.3f}s" ) # Save model DeepONet2d.save_model(filename=f"dnet-{step}.eqx", model=model) results = { "metrics": metrics, "best_loss": best_loss, "model_params": model_params, "last_epoch": step + 1, "time_taken": time_taken, } return model, results

In [ ]:

# Model Parameters
LEARNING_RATE = 1e-3
ACTIVATION = jax.nn.tanh
BATCH_SIZE = 10
STEPS = 1000

# Model Parameters LEARNING_RATE = 1e-3 ACTIVATION = jax.nn.tanh BATCH_SIZE = 10 STEPS = 1000

In [ ]:

model_key, train_key, sub_key = jax.random.split(key=key, num=3)

dnet = DeepONet2d(
    in_size_branch=1,
    in_size_trunk=3,
    width_size=128,
    depth=3,
    out_size=128,
    activation=ACTIVATION,
    key=model_key,
)

optimizer = optax.adam(LEARNING_RATE)

model_key, train_key, sub_key = jax.random.split(key=key, num=3) dnet = DeepONet2d( in_size_branch=1, in_size_trunk=3, width_size=128, depth=3, out_size=128, activation=ACTIVATION, key=model_key, ) optimizer = optax.adam(LEARNING_RATE)

In [ ]:

dnet, results = train(
    model=dnet,
    train_data_full=(b_train, t_train, y_train),
    test_data_full=(b_test, t_test, y_test),
    steps=STEPS,
    batch_size=BATCH_SIZE,
    optimizer=optimizer,
    key=train_key,
)

dnet, results = train( model=dnet, train_data_full=(b_train, t_train, y_train), test_data_full=(b_test, t_test, y_test), steps=STEPS, batch_size=BATCH_SIZE, optimizer=optimizer, key=train_key, )

In [ ]:

u_pred = jax.vmap(dnet)(b_test, t_test)

u_pred = jax.vmap(dnet)(b_test, t_test)

In [ ]:

import matplotlib.pyplot as plt
import matplotlib.animation as mplanim
from matplotlib import (
    rc,
)  # Configure matplotlib to render animations in Jupyter notebooks

rc("animation", html="html5")

import matplotlib.pyplot as plt import matplotlib.animation as mplanim from matplotlib import ( rc, ) # Configure matplotlib to render animations in Jupyter notebooks rc("animation", html="html5")

In [ ]:

def update_u(n, u):
    """Animation showing heat diffusion throughout domain

    Args:
        n (int): Current time step
        u (array): Final temperature distribution
    """
    ax.set_title(rf"$u(x, y, {n})$, at $t={n}$s")
    hm = ax.pcolormesh(u[0, :, :, n], shading="auto", cmap=plt.cm.jet)
    return hm

def update_u(n, u): """Animation showing heat diffusion throughout domain Args: n (int): Current time step u (array): Final temperature distribution """ ax.set_title(rf"$u(x, y, {n})$, at $t={n}$s") hm = ax.pcolormesh(u[0, :, :, n], shading="auto", cmap=plt.cm.jet) return hm

In [ ]:

u_test.shape, u_pred.shape

u_test.shape, u_pred.shape

In [ ]:

# Visualize domain at t=0
fig, ax = plt.subplots(layout="constrained")
ax.set_title(rf"$u(x, y, {0})$, at $t={0}$s")
hm = ax.pcolormesh(
    u_pred[0, :, :, 0],
    shading="auto",
    cmap=plt.cm.jet,
    vmin=u_pred[0, :, :, 0].min(),
    vmax=u_pred[0, :, :, 0].max(),
)
fig.colorbar(mappable=hm, ax=ax)

# Visualize domain at t=0 fig, ax = plt.subplots(layout="constrained") ax.set_title(rf"$u(x, y, {0})$, at $t={0}$s") hm = ax.pcolormesh( u_pred[0, :, :, 0], shading="auto", cmap=plt.cm.jet, vmin=u_pred[0, :, :, 0].min(), vmax=u_pred[0, :, :, 0].max(), ) fig.colorbar(mappable=hm, ax=ax)

In [ ]:

ani = mplanim.FuncAnimation(
    fig=fig, func=update_u, fargs=(u_pred,), interval=200, frames=10, repeat=True
)
ani

ani = mplanim.FuncAnimation( fig=fig, func=update_u, fargs=(u_pred,), interval=200, frames=10, repeat=True ) ani

In [ ]:

# Visualize domain at t=0
fig, ax = plt.subplots(layout="constrained")
ax.set_title(rf"$u(x, y, {0})$, at $t={0}$s")
hm = ax.pcolormesh(
    u_test[0, :, :, 0],
    shading="auto",
    cmap=plt.cm.jet,
    vmin=u_test[0, :, :, 0].min(),
    vmax=u_test[0, :, :, 0].max(),
)
fig.colorbar(mappable=hm, ax=ax)

# Visualize domain at t=0 fig, ax = plt.subplots(layout="constrained") ax.set_title(rf"$u(x, y, {0})$, at $t={0}$s") hm = ax.pcolormesh( u_test[0, :, :, 0], shading="auto", cmap=plt.cm.jet, vmin=u_test[0, :, :, 0].min(), vmax=u_test[0, :, :, 0].max(), ) fig.colorbar(mappable=hm, ax=ax)

In [ ]:

ani = mplanim.FuncAnimation(
    fig=fig, func=update_u, fargs=(u_test,), interval=200, frames=10, repeat=True
)
ani

ani = mplanim.FuncAnimation( fig=fig, func=update_u, fargs=(u_test,), interval=200, frames=10, repeat=True ) ani