In this section, we want to use an MLP to fit the function $z = \sin(\sqrt{(x^2 + y^2)})$ .

Essentially, after training, we use $x$, and $y$, we want to predict the $z$.

This is fitting problem, we assume we do not know the function structure, but just use the data generated from the function.

(Part 1) Import Library

import torch
import torch.nn as nn
import torch.optim as optim
import numpy as np

(Part 2) Data Generation

num_points = 100

x = np.linspace(-6, 6, num_points)
y = np.linspace(-6, 6, num_points)
x, y = np.meshgrid(x, y)
z = np.sin(np.sqrt(x**2 + y**2))

x = torch.tensor(x.reshape(-1, 1), dtype=torch.float32)
y = torch.tensor(y.reshape(-1, 1), dtype=torch.float32)
z = torch.tensor(z.reshape(-1, 1), dtype=torch.float32)

xy = torch.cat([x, y], axis=-1)

Let’s explore the $x, y, z$ relationships using Plotly.

import plotly.graph_objects as go

fig = go.Figure()

fig.add_trace(go.Surface(x=x.reshape(num_points, num_points), 
                         y=y.reshape(num_points, num_points), 
                         z=z.reshape(num_points, num_points)))

fig.show()

(Part 3) Model Definition:

# Step 2: Model definition
class SimpleMLP(nn.Module):
    def __init__(self):
        super(SimpleMLP, self).__init__()
        self.fc1 = nn.Linear(2, 50)  # Input layer
        self.fc2 = nn.Linear(50, 50) # Hidden layer
        self.fc3 = nn.Linear(50, 1)  # Output layer
        self.relu = nn.ReLU()

    def forward(self, x):
        x = self.relu(self.fc1(x))
        x = self.relu(self.fc2(x))
        return self.fc3(x)

(Part 4) Create Model Instance, Loss Function, and Optimizer:

# Model, loss, and optimizer setup
model = SimpleMLP()
criterion = nn.MSELoss()
optimizer = optim.Adam(model.parameters(), lr=0.01)

(Part 5) Training Loos:

epochs = 1000
model.train()
for epoch in range(epochs):
    optimizer.zero_grad()
    z_hat = model(xy)
    loss = criterion(z_hat, z)
    loss.backward()
    optimizer.step()
    if (epoch+1) % 100 == 0:
        print(f'Epoch {epoch+1}, Loss: {loss.item()}')

(Part 6) Using the Training Data to See How Neural Network fits Function:

# Step 4: Visualization of results, using training data
model.eval()
with torch.no_grad():
    z_hat = model(xy).reshape(num_points, num_points).numpy()

(Part 7) Visualization of Results:

from plotly.subplots import make_subplots
fig = make_subplots(rows=1, cols=2,
                    specs=[[{"type": "scatter3d"}, {"type": "scatter3d"}]])

fig.add_trace(go.Surface(x=x.reshape(num_points, num_points), 
                         y=y.reshape(num_points, num_points), 
                         z=z.reshape(num_points, num_points), showscale=False), row=1, col=1)

fig.add_trace(go.Surface(x=x.reshape(num_points, num_points), 
                         y=y.reshape(num_points, num_points), 
                         z=z_hat.reshape(num_points, num_points).detach(), showscale=False), row=1, col=2)

fig.show()

(Part 8) Check Interpolation: