Changing the Equation

[ShaikhaTheGreen]

Hello, and thank you for the fantastic framework!

I tried changing plugging in a different equation:
f_u = u_tt - (c**2) * u_xx

with IC:
u(x,0) = 0
and BC:
u(0,t) = sin(4 * pi * freq * t)
u(1,t) = 0

The following is my code implementation:

import numpy as np
from numpy.polynomial.hermite import hermgauss
import tensorflow as tf
import time
from pyDOE import lhs
from swarm.optimizers.pso_adam import pso
from swarm.utils import multilayer_perceptron, decode
import matplotlib.pyplot as plt
import matplotlib.animation as animation

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

data_file = 'Wave1D_0to1'
data = np.load('dataset/'+data_file+'.npz')
t = np.array(data["t_input"]).flatten()[:, None]
x = np.array(data["x_input"]).flatten()[:, None]
Exact_u = np.array(data["u_output"])



uxn = len(x)
utn = len(t)

xlo = 0.0
xhi = 1.0
ux = np.linspace(xlo, xhi, uxn)

tlo = 0.0
thi = 1.0
ut = np.linspace(tlo, thi, utn)


layer_sizes = [2] + 5 * [15] + [1]
pop_size = 10
n_iter = 6000

# collocation points
def collocation_points(size):
    X = lhs(2, size)
    xcl = tf.expand_dims(
        tf.convert_to_tensor(xlo + (xhi - xlo) * X[:, 0], dtype=tf.float32), -1
    )
    tcl = tf.expand_dims(
        tf.convert_to_tensor(tlo + (thi - tlo) * X[:, 1], dtype=tf.float32), -1
    )
    return xcl, tcl

freq = 1
# initial condition points
x0 = tf.expand_dims(tf.convert_to_tensor(ux, dtype=tf.float32), -1)
t0 = tf.zeros(tf.shape(x0), dtype=tf.float32)
u0 = tf.zeros(tf.shape(x0), dtype=tf.float32) #-tf.math.sin(np.pi * x0)

# Dirichlet boundary condition points
xlb = tf.expand_dims(xlo * tf.ones(tf.shape(ut), dtype=tf.float32), -1)
xub = tf.expand_dims(xhi * tf.ones(tf.shape(ut), dtype=tf.float32), -1)
tlb = tf.expand_dims(tf.convert_to_tensor(ut, dtype=tf.float32), -1)
tub = tf.expand_dims(tf.convert_to_tensor(ut, dtype=tf.float32), -1)

ulb = tf.expand_dims(tf.math.sin(4 * np.pi * freq * ut), -1)
uub = tf.expand_dims(tf.zeros(tf.shape(ut), dtype=tf.float32), -1)

c_s = 1
@tf.function
def f_model(w, b): 
    x, t = collocation_points(500)
    u = multilayer_perceptron(w, b, tf.concat([x, t], 1))
    u_x = tf.gradients(u, x)[0]
    u_xx = tf.gradients(u_x, x)[0]
    u_t = tf.gradients(u, t)[0]
    u_tt = tf.gradients(u_t, t)[0]
    #f_u = u_tt - [(c_s ** 2)* element for element in u_xx]
    f_u = u_tt - (c_s ** 2) * u_xx
    print(f_u.shape)
    print(u_tt.shape)
    print(u_xx.shape)
    return f_u


@tf.function
def loss(w, b):
    u0_pred = multilayer_perceptron(w, b, tf.concat([x0, t0], 1))
    ulb_pred = multilayer_perceptron(w, b, tf.concat([xlb, tlb], 1))
    uub_pred = multilayer_perceptron(w, b, tf.concat([xub, tub], 1))
    f_pred = f_model(w, b)

    # IC loss
    mse_0 = tf.reduce_mean(tf.pow(u0 - u0_pred, 2))

    # Dirichlet BC loss
    mse_lb = tf.reduce_mean(tf.pow(ulb_pred - ulb, 2))
    mse_ub = tf.reduce_mean(tf.pow(uub_pred - uub, 2))

    # Residual loss
    mse_f = tf.reduce_mean(tf.pow(f_pred, 2))

    return mse_0 + mse_f + mse_lb + mse_ub


def loss_grad():
    def _loss(w, b):
        with tf.GradientTape(persistent=True) as tape:
            tape.watch(w)
            tape.watch(b)
            loss_value = loss(w, b)
        trainable_variables = w + b
        grads = tape.gradient(loss_value, trainable_variables)
        return loss_value, grads

    return _loss


def run_swarm(swarm, X):
    new_swarm = []
    for particle in swarm:
        w, b = decode(particle, layer_sizes)
        new_swarm.append(
            multilayer_perceptron(w, b, X_flat.astype(np.float32))
        )
    return tf.convert_to_tensor(new_swarm, dtype=tf.float32)


opt = pso(
    loss_grad(),
    layer_sizes,
    n_iter,
    pop_size,
    0.999,
    8e-4,
    5e-3,
    initialization_method="xavier",
    verbose=True,
    gd_alpha=1e-4,
)




start = time.time()
opt.train()
end = time.time()
print("\nTime elapsed: ", end - start)



X, T = np.meshgrid(x, t)
X_flat = np.hstack((X.flatten()[:, None], T.flatten()[:, None]))
u_star = Exact_u.T.flatten()
x_ = np.squeeze(x)
swarm = opt.get_swarm()
preds = run_swarm(swarm, X_flat.astype(np.float32))
mean = tf.squeeze(tf.reduce_mean(preds, axis=0))
var = tf.squeeze(tf.math.reduce_std(preds, axis=0))
print("Last Loss: ", opt.loss_history[-1])

time_steps = utn - 1  # total number of time steps in animation
fps = 15  # frames/second of animation
error_u = np.linalg.norm(u_star - mean, 2) / np.linalg.norm(u_star, 2)
print("Error u: %e" % (error_u))


def snapshot(i):
    l_ind = i * uxn
    u_ind = (i + 1) * uxn
    plt.clf()
    for k in range(len(preds)):
        plt.plot(x_, preds[k][l_ind:u_ind], linewidth=0.3)
    plt.scatter(x_d, u_d_flat[i * data_size : (i + 1) * data_size])
    plt.plot(x_, u_star[l_ind:u_ind], "b-", linewidth=3, label="Exact")
    plt.plot(
        x_,
        mean[l_ind:u_ind],
        "r--",
        linewidth=3,
        label=opt.name,
    )
    plt.fill_between(
        x_,
        mean[l_ind:u_ind] - 3 * var[l_ind:u_ind],
        mean[l_ind:u_ind] + 3 * var[l_ind:u_ind],
        color="gray",
        alpha=0.2,
    )

    plt.xlabel("$x$", fontsize="xx-large")
    plt.ylabel("$u(t,x)$", fontsize="xx-large")
    plt.xlim(-1.0, 1.0)
    plt.ylim(-1.02, 1.02)
    plt.grid()
    plt.legend(fontsize="x-large")

fig = plt.figure(figsize=(8, 8), dpi=150)
# Call the animator:
anim = animation.FuncAnimation(fig, snapshot, frames=time_steps)
# Save the animation as an mp4. This requires ffmpeg to be installed.
anim.save("wave_demo.gif", fps=fps)

I encountered what seems like a simple error, but I cannot resolve it. The error script is:

---> 24 opt = pso(
     25     loss_grad(),
     26     layer_sizes,
     27     n_iter,
     28     pop_size,
     29     0.999,
     30     8e-4,
     31     5e-3,
     32     initialization_method="xavier",
     33     verbose=True,
     34     gd_alpha=1e-4,
     35 )

File c:\users\user\deep learning\pso-pinn-main\src\swarm\optimizers\pso_adam.py:52, in pso.__init__(self, loss_op, layer_sizes, n_iter, pop_size, b, c1, c2, gd_alpha, initialization_method, verbose, c_decrease, pre_trained_x, beta_1, beta_2, epsilon)
     48 self.initialization_method = initialization_method
     49 self.x = (
     50     pre_trained_x if pre_trained_x is not None else self.build_swarm()
     51 )
...
    File "C:\Users\User\AppData\Local\Temp\ipykernel_5196\1050840099.py", line 29, in loss  *
        mse_lb = tf.reduce_mean(tf.pow(ulb_pred - ulb, 2))

    TypeError: Input 'y' of 'Sub' Op has type float64 that does not match type float32 of argument 'x'.

Any ideas?

[caio-davi]

According to your error, it seems you have a type mismatch. Try to force ulb to tf.float32:

ulb = tf.expand_dims(tf.math.sin(4 * np.pi * tf.convert_to_tensor(freq * ut, dtype=tf.float32)), -1)

I'm not sure, but it should do the trick.

[ShaikhaTheGreen]

Update:
It worked great! Thank you for your suggestion! I'm also trying to test two cases of this equation:

Case1: An increased frequency. So, instead of sin(4pit), I'll use sin(16pit).

frequency = 4 shows:

However, I did not succeed in obtaining the expected results when changing the frequency to 16 (looks reversed for some reason):

Case2: A heterogenous domain. In the original equation, I'd like to assign the value of c as a function dependent on the location. For instance, I defined it within f_model() to be as such:

c1 = 1 
    c2 = 0.5
    # Implementing mask c = c1 * (x<=0.5) + c2 *(x>0.5)
    c_s = tf.where(x[:, 0:1] <= 0.5, x[:, 0:1] *0 + c1, x[:, 0:1] * 0 + c2)

but the results were not accurate enough:

The fixes I tried so far for case1 are:

• Increasing pop_size
• Increasing collocation points
• Modifying loss weights. Tried increasing the weight for mse_lb with values (1.5, 5, and 10).
  These have improved the prediction of the curve partially, but it did not fix the reverse propagation.

The fixes I tried so far for case2 are:

• Increasing n_iter
• Increasing collocation points
• Modifying loss weights. Tried increasing the weight for mse_f with values (1.5, 5, and 10).
  These fixes did not improve the accuracy.

Any suggestions would be helpful. Thanks again!

[caio-davi]

I'm sorry for the delay, it's have been a busy week.
Did you only change the `lower boundary condition between cases 1 and 2?
Strange indeed. Especially for this "inverse" movement.
I think I'll have some free time during the weekend, so I can review your example and try to figure out if there is only a parametrization problem or something else.
Did you try it with a traditional PINN?

[caio-davi]

Can you share your experiments? I tried to reproduce it based on the shared code but didn't get much success.

[ShaikhaTheGreen]

I apologize for the delayed response. Thank you for looking into this. Not much of a change in the other two experiments.

In case 1:
I only changed this line:
ulb = tf.expand_dims(tf.math.sin(4 * np.pi * tf.convert_to_tensor(freq * ut, dtype=tf.float32)), -1)
to
ulb = tf.expand_dims(tf.math.sin(16 * np.pi * tf.convert_to_tensor(freq * ut, dtype=tf.float32)), -1)

In case 2:
I added this code to f_model()

    c1 = 1
    c2 = 0.5
    # Implementing mask c = c1 * (x>0.5) + c2 *(x<=0.5)
    c_s = tf.where(x[:, 0:1] <= 0.5, x[:, 0:1] *0 + c1, x[:, 0:1] * 0 + c2)
    f_u = u_tt - (c_s ** 2) * u_xx

Please let me know if you're able to reproduce the problem.