2d_neuralnetworkintegratorexamples.py

# -*- coding: utf-8 -*-
"""2D - NeuralNetworkIntegratorExamples.ipynb

Automatically generated by Colab.

Original file is located at
    https://colab.research.google.com/drive/1xuWb_OSU2c77TL1Q5D-D_ByOe_p3qWE-
"""

!pip install tensorflow_addons

import numpy as np
import tensorflow as tf
from tensorflow.keras.layers import Dense
from tensorflow.keras.models import Model
from tensorflow.keras.optimizers import Adam
from tensorflow.keras.initializers import RandomNormal
from tensorflow.keras import regularizers

tf.random.set_seed(28)
np.random.seed(28)

batch_size = 32
epochs = 10
learning_rate = 0.001

initDev = 0.05

inputs = tf.keras.Input(shape=(2,))
x = inputs
for i in range(1, 22):
    x = Dense(2048, activation='gelu', kernel_initializer=RandomNormal(mean=0.0, stddev=initDev, seed=i))(x)
output = Dense(1, activation='linear', kernel_initializer=RandomNormal(mean=0.0, stddev=initDev, seed=100))(x)
model = Model(inputs, output)

optimizer = Adam(learning_rate=learning_rate)

def f_tbi(x, y):
    mu = 0.0
    sigma = 1.0
    return (1 / (2 * np.pi * sigma ** 2)) * np.exp(-0.5 * ((x - mu) ** 2 + (y - mu) ** 2) / (2 * sigma ** 2))

def f_tbi(x, y):
    return x**2 + y**2

initial_points = np.array([[0.0, 0.0], [5, 0], [-5, 0], [0, 5], [0, -5], [-5, 5], [-5, -5], [5, -5], [5, 5]], dtype=np.float32)
initial_values = np.array([[0], [41.66666666], [-41.66666666], [0], [0], [-166.66666666], [-166.66666666], [166.66666666], [166.66666666]], dtype=np.float32)
#initial_values = np.array([[0], [0.2819799927292551], [-0.2819799927292551], [0], [0], [-0.0005443494432976496], [-0.0005443494432976496], [0.0005443494432976496], [0.0005443494432976496]], dtype=np.float32)
#[-2.5, -2.5], [-2.5, 2.5], [2.5, 2.5], [2.5, -2.5]
#[-20.833333333333336], [-20.833333333333336], [20.833333333333336], [20.833333333333336]

x = np.linspace(-5, 5, 26, dtype=np.float32)
y = np.linspace(-5, 5, 26, dtype=np.float32)
x_grid, y_grid = np.meshgrid(x, y)
z = f_tbi(x_grid, y_grid).reshape(-1, 1).astype(np.float32)
xy = np.column_stack([x_grid.flatten(), y_grid.flatten()]).astype(np.float32)

dataset = tf.data.Dataset.from_tensor_slices((xy, z)).shuffle(buffer_size=128*128).batch(batch_size)

best_loss = float('inf')
best_epoch = -1

@tf.function
def train_step(inputs, targets, init_points, init_values):
    with tf.GradientTape() as tape:
        tape.watch(inputs)
        predictions = model(inputs, training=True)
        grads = tape.gradient(predictions, inputs)[:, 0]
    with tf.GradientTape() as tape_loss:
        tape_loss.watch(model.trainable_variables)
        gradient_loss = tf.reduce_mean(tf.square(grads - targets))
        init_predictions = model(init_points, training=True)
        init_loss = tf.reduce_mean(tf.square(init_predictions - init_values))
        loss = 0.001 * gradient_loss + 0.0001 * init_loss
        gradients = tape_loss.gradient(loss, model.trainable_variables)
        optimizer.apply_gradients(zip(gradients, model.trainable_variables))
    return loss

for epoch in range(epochs):
    epoch_loss = 0
    steps = 0
    for inputs, targets in dataset:
        loss = train_step(inputs, targets, initial_points, initial_values)
        epoch_loss += loss.numpy()
        steps += 1
    epoch_loss /= steps
    if epoch_loss < best_loss:
        best_loss = epoch_loss
        best_epoch = epoch
        model.save_weights('best_model.h5')  # Save the best model
    print(f'Epoch {epoch + 1}, Loss: {epoch_loss}')

print(f'Best model saved from Epoch {best_epoch + 1} with Loss: {best_loss}')

# Load the best model saved
model.load_weights('best_model.h5')

x_test = np.linspace(-5, 5, 26)
y_test = np.linspace(-5, 5, 26)
X_test, Y_test = np.meshgrid(x_test, y_test)
XY_test = np.stack([X_test.flatten(), Y_test.flatten()], axis=1)
Z_pred = model.predict(XY_test).reshape(X_test.shape)

import numpy as np
import matplotlib.pyplot as plt
from scipy.integrate import quad
from mpl_toolkits.mplot3d import Axes3D

# Define the function f(a, y) = a^2 + y^2
def integrand(a, y):
    return f_tbi(a, y)

# Compute the integral of f from 0 to x for each (x, y) in the meshgrid
def compute_integral(x, y):
    # Function to compute integral from 0 to x
    integral, _ = quad(integrand, 0, x, args=(y))
    return integral

# Create a meshgrid for x and y
x_values = np.linspace(-5, 5, 50)  # More points for smoother plot
y_values = np.linspace(-5, 5, 50)
X, Y = np.meshgrid(x_values, y_values)

# Compute the integral values at each point in the meshgrid
Z_integral = np.vectorize(compute_integral)(X, Y)

# Create a 3D plot
fig = plt.figure(figsize=(14, 9))
ax = plt.axes(projection='3d')

# Plot the neural network predictions surface (Z_pred and XY_test from your previous plot)
# Note: You need to have these values defined from your neural network predictions
# If not defined here, ensure Z_pred and XY_test are calculated prior to this script
ax.plot_surface(X, Y, Z_integral, cmap='viridis', alpha=0.6, label='Integral Surface')

# Add your neural network output surface
ax.plot_surface(X_test, Y_test, Z_pred, color='grey', alpha=0.6, label='NN Output Surface')

# Adding initial points
ax.scatter(initial_points[:, 0], initial_points[:, 1], initial_values, color='red', s=50)

# Setting labels
ax.set_xlabel('X axis')
ax.set_ylabel('Y axis')
ax.set_zlabel('Z axis')

# Customize the view angle for better visualization
ax.view_init(elev=15, azim=90)

# Show plot
plt.show()

print(compute_integral(0, 0))
print(compute_integral(5, 0))
print(compute_integral(-5, 0))
print(compute_integral(0, 5))
print(compute_integral(0, -5))
print(compute_integral(-5, 5))
print(compute_integral(-5, -5))
print(compute_integral(5, -5))
print(compute_integral(5, 5))

import numpy as np
import tensorflow as tf
from tensorflow.keras.layers import Dense
from tensorflow.keras.models import Model
from tensorflow.keras.optimizers import Adam
from tensorflow.keras.initializers import RandomNormal
from matplotlib.colors import LinearSegmentedColormap
import matplotlib.pyplot as plt
from mpl_toolkits.mplot3d import Axes3D  # for 3D plotting
from scipy.integrate import quad
import imageio
from io import BytesIO

sampling = 500

# Set seeds for reproducibility
tf.random.set_seed(28)
np.random.seed(28)

# Define custom colormap for the NN predicted surface
blue_color   = (69/255, 202/255, 1)
purple_color = (181/255, 71/255, 225/255)
pink_color   = (1, 27/255, 107/255)
custom_cmap  = LinearSegmentedColormap.from_list("custom", [blue_color, purple_color, pink_color])

# Training parameters
batch_size = 600
epochs = 100  # Increase this (e.g., to 50) for full training
learning_rate = 0.01
initDev = 0.03

# Build the neural network model
inputs = tf.keras.Input(shape=(2,))
x = inputs
for i in range(1, 4):
    x = Dense(1024, activation='gelu',
              kernel_initializer=RandomNormal(mean=0.0, stddev=initDev, seed=i))(x)
output = Dense(1, activation='linear',
               kernel_initializer=RandomNormal(mean=0.0, stddev=initDev, seed=100))(x)
model = Model(inputs, output)
optimizer = Adam(learning_rate=learning_rate)

# Define the function f(x,y) = x^2 + y^2 (the target gradient)
def f_tbi(x, y):
    return x**2 + y**2

# For a given y, compute the antiderivative with respect to x:
def integrand(a, y):
    return f_tbi(a, y)

def compute_integral(x, y):
    integral, _ = quad(integrand, 0, x, args=(y))
    return integral

# Define exactly 9 anchor points (in the (x,y) plane), each mapping uniquely to a z-value.
initial_points = np.array([[0.0, 0.0],
                           [5, 0],
                           [-5, 0],
                           [0, 5],
                           [0, -5],
                           [-5, 5],
                           [-5, -5],
                           [5, -5],
                           [5, 5]], dtype=np.float32)
# Compute each anchor point's corresponding z-value (the true antiderivative)
initial_values = np.array([compute_integral(pt[0], pt[1]) for pt in initial_points],
                            dtype=np.float32).reshape(-1, 1)

# Create training dataset on a 10x10 grid over [-5,5] x [-5,5]
x_lin = np.linspace(-5, 5, 10, dtype=np.float32)
y_lin = np.linspace(-5, 5, 10, dtype=np.float32)
x_grid, y_grid = np.meshgrid(x_lin, y_lin)
# For training, we use f_tbi(x,y) so that the network’s derivative (with respect to x) matches f_tbi.
z = f_tbi(x_grid, y_grid).reshape(-1, 1).astype(np.float32)
xy = np.column_stack([x_grid.flatten(), y_grid.flatten()]).astype(np.float32)
dataset = tf.data.Dataset.from_tensor_slices((xy, z)).shuffle(sampling*sampling).batch(batch_size)

# Precompute the true antiderivative surface (used for visualization) on a 50x50 grid.
x_values = np.linspace(-5, 5, 50)
y_values = np.linspace(-5, 5, 50)
X, Y = np.meshgrid(x_values, y_values)
Z_integral = np.vectorize(compute_integral)(X, Y)

# Prepare test grid for NN predictions (using the same grid as training)
x_test = np.linspace(-5, 5, sampling)
y_test = np.linspace(-5, 5, sampling)
X_test, Y_test = np.meshgrid(x_test, y_test)
XY_test = np.stack([X_test.flatten(), Y_test.flatten()], axis=1).astype(np.float32)

# Training step:
# 1. Gradient loss: Enforce that the derivative of u(x,y) with respect to x matches f(x,y)=x^2+y^2.
# 2. Anchor loss: Enforce that u(x,y) at the 9 anchor points matches the true integrated values.
@tf.function
def train_step(inputs, targets, init_points, init_values):
    with tf.GradientTape(persistent=True) as tape:
        tape.watch(inputs)
        predictions = model(inputs, training=True)
        # Compute partial derivative with respect to x.
        grads = tape.gradient(predictions, inputs)[:, 0]
        gradient_loss = tf.reduce_mean(tf.square(grads - targets))
        # Evaluate network at anchor points (9 unique points).
        init_predictions = model(init_points, training=True)
        init_loss = tf.reduce_mean(tf.square(init_predictions - init_values))
        # Total loss is a weighted sum.
        loss = 0.1 * gradient_loss + 120 * init_loss
    gradients = tape.gradient(loss, model.trainable_variables)
    optimizer.apply_gradients(zip(gradients, model.trainable_variables))
    return loss

best_loss = float('inf')
best_epoch = -1
# Set view angle: elevate 15° and rotate to the right (azimuth=110°)
view_elev = 15
view_azim = 110
frames = []

for epoch in range(epochs):
    epoch_loss = 0.0
    steps = 0
    for inputs_batch, targets_batch in dataset:
        loss = train_step(inputs_batch, targets_batch, initial_points, initial_values)
        epoch_loss += loss.numpy()
        steps += 1
    epoch_loss /= steps
    print(f'Epoch {epoch+1}, Loss: {epoch_loss:.6f}')

    # Generate NN prediction for the current epoch.
    Z_pred_epoch = model(tf.constant(XY_test), training=False).numpy().reshape(X_test.shape)

    fig = plt.figure(figsize=(8,6))
    ax = fig.add_subplot(111, projection='3d')

    # Plot the true antiderivative surface (gray, alpha=0.5).
    ax.plot_surface(X, Y, Z_integral, color='gray', alpha=0.5, edgecolor='none')

    # Plot the NN predicted surface (using the custom colormap).
    ax.plot_surface(X_test, Y_test, Z_pred_epoch, cmap=custom_cmap, alpha=0.6, edgecolor='none')

    # Plot the anchor points: exactly 9 unique dots.
    ax.scatter(initial_points[:,0],
               initial_points[:,1],
               initial_values.flatten(),   # flatten to get shape (9,)
               color=blue_color, s=50)

    # Set the view angle.
    ax.view_init(elev=view_elev, azim=view_azim)

    # Remove axis ticks.
    ax.set_xticks([])
    ax.set_yticks([])
    ax.set_zticks([])

    # Remove axes entirely by hiding panes and setting tick colors to transparent.
    ax.xaxis.pane.set_visible(False)
    ax.yaxis.pane.set_visible(False)
    ax.zaxis.pane.set_visible(False)
    ax.xaxis._axinfo["tick"]["color"] = (0,0,0,0)
    ax.yaxis._axinfo["tick"]["color"] = (0,0,0,0)
    ax.zaxis._axinfo["tick"]["color"] = (0,0,0,0)

    # Capture the current frame at higher resolution.
    buf = BytesIO()
    plt.savefig(buf, format='png', bbox_inches='tight', dpi=300)
    buf.seek(0)
    frames.append(imageio.imread(buf))
    plt.close(fig)

# Save the GIF (1 second per frame).
imageio.mimsave('training_progress.gif', frames, duration=1)

import plotly.graph_objects as go

# Create the true surface trace (displayed in gray with 50% opacity)
trace_true = go.Surface(
    x=X,
    y=Y,
    z=Z_integral,
    colorscale='Greys',
    opacity=0.5,
    showscale=False,
    name='True Surface'
)

# Create the NN predicted surface trace (using a Viridis colorscale here)
trace_pred = go.Surface(
    x=X_test,
    y=Y_test,
    z=Z_pred_epoch,
    colorscale='Viridis',
    opacity=0.6,
    name='NN Predicted Surface'
)

# Create the anchor points trace (exactly 9 dots)
trace_anchors = go.Scatter3d(
    x=initial_points[:, 0],
    y=initial_points[:, 1],
    z=initial_values.flatten(),  # flatten to ensure a 1D array of 9 elements
    mode='markers',
    marker=dict(
        size=5,
        color='blue'
    ),
    name='Anchors'
)

# Combine all traces into one figure.
fig = go.Figure(data=[trace_true, trace_pred, trace_anchors])

# Update layout with axis titles and an adjusted camera view.
fig.update_layout(
    title='Final Interactive Plot',
    scene=dict(
        xaxis_title='X',
        yaxis_title='Y',
        zaxis_title='Z',
        camera=dict(eye=dict(x=1.5, y=1.5, z=1.5))
    )
)

fig.show()

model.save("my_model.h5")
#loaded_model = tf.keras.models.load_model("my_model.h5")

import tensorflow as tf
import numpy as np
import plotly.graph_objects as go
from tensorflow.keras.optimizers import Adam

# Load the saved model.
loaded_model = tf.keras.models.load_model("my_model.h5")
# Since the training configuration wasn’t saved, compile the model manually.
loaded_model.compile(optimizer=Adam(), loss='mse')

# Force execution on CPU.
with tf.device('/CPU:0'):
    # Convert XY_test to a Tensor and get predictions directly.
    Z_pred = loaded_model(tf.constant(XY_test, dtype=tf.float32), training=False).numpy()

# Reshape predictions to match the grid shape of X_test and Y_test.
Z_pred = Z_pred.reshape(X_test.shape)

# Define a custom colorscale (pink-purple-blue).
custom_colorscale = [
    [0.0, 'rgb(255,27,107)'],   # pink
    [0.5, 'rgb(181,71,225)'],   # purple
    [1.0, 'rgb(69,202,255)']     # blue
]

# Create the true surface trace (displayed in gray with 50% opacity).
trace_true = go.Surface(
    x=X,
    y=Y,
    z=Z_integral,
    colorscale='Greys',
    opacity=0.5,
    showscale=False,
    name='True Surface'
)

# Create the NN predicted surface trace using the custom colorscale.
trace_pred = go.Surface(
    x=X_test,
    y=Y_test,
    z=Z_pred,
    colorscale=custom_colorscale,
    opacity=0.6,
    name='NN Predicted Surface'
)

# Create the anchor points trace (exactly 9 dots).
trace_anchors = go.Scatter3d(
    x=initial_points[:, 0],
    y=initial_points[:, 1],
    z=initial_values.flatten(),  # ensure a 1D array of 9 elements
    mode='markers',
    marker=dict(
        size=5,
        color='rgb(255,27,107)'  # using pink for the anchor points
    ),
    name='Anchors'
)

# Combine all traces into one figure.
fig = go.Figure(data=[trace_true, trace_pred, trace_anchors])

# Update layout with axis titles and an adjusted camera view.
fig.update_layout(
    title='Final Interactive Plot',
    scene=dict(
        xaxis_title='X',
        yaxis_title='Y',
        zaxis_title='Z',
        camera=dict(eye=dict(x=1.5, y=1.5, z=1.5))
    )
)

fig.show()

# Define the input for which you want to predict the output
input_value1 = np.array([[5, -5]], dtype=np.float32)
input_value2 = np.array([[4.5, -5]], dtype=np.float32)

# Use the model to predict the output for the given input
print(model.predict(input_value1)-model.predict(input_value2))