In [ ]:

import os
import shutil
import numpy as np

import os import shutil import numpy as np

Our Neural Operator Dataset¶

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] $$

Our domain, $\Omega$, for storing the PDE solution, $u$, with our BCs and IC is created below.

In [ ]:

def init_domain(x_len, y_len, delta_x, delta_y, max_time):
    """Initialize 2D domain

    Args:
        x_len (int): Domain length in x-direction [m]
        y_len (int): Domain length in y-direction [m]
        delta_x (int): Spacing between mesh nodes in x-direction [m]
        delta_y (int): Spacing between mesh nodes in y-direction [m]
        max_time (int): Number of time steps

    Returns:
        array: Empty domain
    """
    # Number of spatial mesh nodes
    nx, ny = int(x_len / delta_x), int(y_len / delta_y)

    # Initialize solution: u(j, i, n)
    u0 = np.zeros((nx + 1, ny + 1, max_time + 1))

    # Domain info
    info = {
        "x_len": x_len,
        "y_len": y_len,
        "delta_x": delta_x,
        "delta_y": delta_y,
        "nx": nx,
        "ny": ny,
        "max_time": max_time,
    }

    return u0, info

def init_domain(x_len, y_len, delta_x, delta_y, max_time): """Initialize 2D domain Args: x_len (int): Domain length in x-direction [m] y_len (int): Domain length in y-direction [m] delta_x (int): Spacing between mesh nodes in x-direction [m] delta_y (int): Spacing between mesh nodes in y-direction [m] max_time (int): Number of time steps Returns: array: Empty domain """ # Number of spatial mesh nodes nx, ny = int(x_len / delta_x), int(y_len / delta_y) # Initialize solution: u(j, i, n) u0 = np.zeros((nx + 1, ny + 1, max_time + 1)) # Domain info info = { "x_len": x_len, "y_len": y_len, "delta_x": delta_x, "delta_y": delta_y, "nx": nx, "ny": ny, "max_time": max_time, } return u0, info

In [ ]:

def init_u(domain, u_initial, u_right, u_left, u_bottom, u_top, info):
    """Initialize temperature distribution with Dirichlet BCs

    Args:
        domain (array): Empty domain
        u_initial (float): Initial temperature of domain at time step 0
        u_right (float): Temperature on right side
        u_left (float): Temperature on left side
        u_bottom (float): Temperature on bottom side
        u_top (float): Temperature on top side
        info (dict): Domain description

    Returns:
        array: Initial temperature distribution
    """
    u = domain
    nx, ny = info["nx"], info["nx"]

    # Set the initial condition
    u[:, :, 0] = u_initial

    # Set the boundary conditions
    u[:, nx:, :] = u_right
    u[:, :1, :] = u_left
    u[:1, 1:nx, :] = u_bottom
    u[ny:, 1:nx, :] = u_top

    return u

def init_u(domain, u_initial, u_right, u_left, u_bottom, u_top, info): """Initialize temperature distribution with Dirichlet BCs Args: domain (array): Empty domain u_initial (float): Initial temperature of domain at time step 0 u_right (float): Temperature on right side u_left (float): Temperature on left side u_bottom (float): Temperature on bottom side u_top (float): Temperature on top side info (dict): Domain description Returns: array: Initial temperature distribution """ u = domain nx, ny = info["nx"], info["nx"] # Set the initial condition u[:, :, 0] = u_initial # Set the boundary conditions u[:, nx:, :] = u_right u[:, :1, :] = u_left u[:1, 1:nx, :] = u_bottom u[ny:, 1:nx, :] = u_top return u

In [ ]:

def calc_u_fd(u, alpha, info):
    """Finite difference Method for 2D Heat equation

    Args:
        u (array): Initial temperature distribution at nodes (j, i, n)
        alpha (float): Thermal diffusivity
        info (dict): Domain description

    Returns:
        array: Final temperature distribution
    """
    delta_x, delta_y, max_time = info["delta_x"], info["delta_x"], info["max_time"]

    # Stability calcs
    delta_t = min(
        (delta_x**2 * delta_y**2) / (2 * alpha * (delta_x**2 + delta_y**2)), 0.5
    )

    # Evaluate solution
    for n in range(0, max_time, 1):
        u0 = u.copy()
        u_xx = (u0[2:, 1:-1, n] - 2 * u0[1:-1, 1:-1, n] + u0[:-2, 1:-1, n]) / (
            delta_x**2
        )
        u_yy = (u0[1:-1, 2:, n] - 2 * u0[1:-1, 1:-1, n] + u0[1:-1, :-2, n]) / (
            delta_y**2
        )
        u[1:-1, 1:-1, n + 1] = u0[1:-1, 1:-1, n] + (alpha * delta_t) * (u_xx + u_yy)

    return u

def calc_u_fd(u, alpha, info): """Finite difference Method for 2D Heat equation Args: u (array): Initial temperature distribution at nodes (j, i, n) alpha (float): Thermal diffusivity info (dict): Domain description Returns: array: Final temperature distribution """ delta_x, delta_y, max_time = info["delta_x"], info["delta_x"], info["max_time"] # Stability calcs delta_t = min( (delta_x**2 * delta_y**2) / (2 * alpha * (delta_x**2 + delta_y**2)), 0.5 ) # Evaluate solution for n in range(0, max_time, 1): u0 = u.copy() u_xx = (u0[2:, 1:-1, n] - 2 * u0[1:-1, 1:-1, n] + u0[:-2, 1:-1, n]) / ( delta_x**2 ) u_yy = (u0[1:-1, 2:, n] - 2 * u0[1:-1, 1:-1, n] + u0[1:-1, :-2, n]) / ( delta_y**2 ) u[1:-1, 1:-1, n + 1] = u0[1:-1, 1:-1, n] + (alpha * delta_t) * (u_xx + u_yy) return u

In [ ]:

# Initialize domain
u0, info = init_domain(x_len=20, y_len=20, max_time=10, delta_x=1, delta_y=1)

# Apply ICs & BCs
u = init_u(
    domain=u0,
    u_initial=0.0,
    u_right=100.0,
    u_left=0.0,
    u_bottom=0.0,
    u_top=0.0,
    info=info,
)

# Initialize domain u0, info = init_domain(x_len=20, y_len=20, max_time=10, delta_x=1, delta_y=1) # Apply ICs & BCs u = init_u( domain=u0, u_initial=0.0, u_right=100.0, u_left=0.0, u_bottom=0.0, u_top=0.0, info=info, )

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 below:

In [ ]:

# Thermal diffusivity
alphas = np.random.rand(100) * 2

# Thermal diffusivity alphas = np.random.rand(100) * 2

In [ ]:

u_all, a_all = [], []
for i, alpha in enumerate(alphas):
    # Calculate solution for different diffusivity
    u0 = u.copy()
    u_final = calc_u_fd(u0, alpha, info=info)

    # Add random gaussian noise to data
    mu, sigma = 0, 0.1  # mean and standard deviation
    noise = np.random.normal(mu, sigma, u_final.shape) * 0.05
    u_final_noise = u_final + noise

    # Collect data
    u_all.append(u_final_noise)
    a_all.append(alpha)

u_fn = np.array(u_all)
a_fn = np.array(a_all).reshape(-1, 1)

u_all, a_all = [], [] for i, alpha in enumerate(alphas): # Calculate solution for different diffusivity u0 = u.copy() u_final = calc_u_fd(u0, alpha, info=info) # Add random gaussian noise to data mu, sigma = 0, 0.1 # mean and standard deviation noise = np.random.normal(mu, sigma, u_final.shape) * 0.05 u_final_noise = u_final + noise # Collect data u_all.append(u_final_noise) a_all.append(alpha) u_fn = np.array(u_all) a_fn = np.array(a_all).reshape(-1, 1)

The nodal coordinates for the 2D domain is also needed for evaluating the performance of the DeepONet model. These coordinates will be used as input for the trunk network. This is done below:

In [ ]:

# Create Mesh
x = np.linspace(0, info["x_len"], num=info["nx"] + 1)
y = np.linspace(0, info["y_len"], num=info["ny"] + 1)
t = np.linspace(0, info["max_time"], num=info["max_time"] + 1)
xx, yy, tt = np.meshgrid(x, y, t, indexing="ij")

# Create Mesh x = np.linspace(0, info["x_len"], num=info["nx"] + 1) y = np.linspace(0, info["y_len"], num=info["ny"] + 1) t = np.linspace(0, info["max_time"], num=info["max_time"] + 1) xx, yy, tt = np.meshgrid(x, y, t, indexing="ij")

In [ ]:

tmesh = np.concatenate(
    (
        tt.flatten()[:, np.newaxis],
        xx.flatten()[:, np.newaxis],
        yy.flatten()[:, np.newaxis],
    ),
    axis=1,
)

tmesh = np.concatenate( ( tt.flatten()[:, np.newaxis], xx.flatten()[:, np.newaxis], yy.flatten()[:, np.newaxis], ), axis=1, )

In [ ]:

# Determine min & max values for normalizing
t_max, t_min = tmesh[:, 0].max(), tmesh[:, 0].min()
x_max, x_min = tmesh[:, 1].max(), tmesh[:, 1].min()
y_max, y_min = tmesh[:, 2].max(), tmesh[:, 2].min()

# Determine min & max values for normalizing t_max, t_min = tmesh[:, 0].max(), tmesh[:, 0].min() x_max, x_min = tmesh[:, 1].max(), tmesh[:, 1].min() y_max, y_min = tmesh[:, 2].max(), tmesh[:, 2].min()

In [ ]:

# Add to info
info["t_max"], info["t_min"] = t_max, t_min
info["x_max"], info["x_min"] = x_max, x_min
info["y_max"], info["y_min"] = y_max, y_min

# Add to info info["t_max"], info["t_min"] = t_max, t_min info["x_max"], info["x_min"] = x_max, x_min info["y_max"], info["y_min"] = y_max, y_min

In [ ]:

# Save results to file
if os.path.exists("./data"):
    shutil.rmtree("./data")
    os.makedirs("./data")
elif not os.path.exists("./data"):
    os.makedirs("./data")
np.savez(
    file="./data/data.npz",
    allow_pickle=True,
    u_noise=u_fn,
    a=a_fn,
    mesh=tmesh,
    info=info,
)

# Save results to file if os.path.exists("./data"): shutil.rmtree("./data") os.makedirs("./data") elif not os.path.exists("./data"): os.makedirs("./data") np.savez( file="./data/data.npz", allow_pickle=True, u_noise=u_fn, a=a_fn, mesh=tmesh, info=info, )