Inverse problems variables failed to converge

[AliaNajwaMY]

Hello,

I'm currently trying to solve the problem $u_tt - a u_xx - b u_x -c u =0$ when $x \in [0,1]$ and $t\in [0.1]$. The boundary condition is $u(x, 0) = u(x,1)=0$ and the initial condition is $u(x, 0)= \sin(\pi x)$. In my code, I'm trying to recover a, b and c (variables) when the a, b and c that I originally wanted is non 0. I managed to get a, b and c recovered when $b=c =0$ but has been having trouble when either $b$ or c is nonzero. Below is my code; please suggest anything I can do to improve my results.

import deepxde as dde
import numpy as np
from aim import Run
from functools import partial
import concurrent.futures
import matplotlib as plt
import plotly.express as px


def heat_eq_exact_solution(x, t):
        """Returns the exact solution for a given x and t (for sinusoidal initial conditions).

        Parameters
        ----------
        x : np.ndarray
        t : np.ndarray
        """
        # return np.exp(j * np.pi * x) * np.exp((-(self.a) * (np.pi**2) + self.b * j * np.pi + self.c) * t) - np.exp(-(j * np.pi * x)) * np.exp((-(self.a) * (np.pi**2) - self.b * j * np.pi + self.c) * t) 
        # return np.exp(-(n**2 * np.pi**2 * self.a * t) / (L**2)) * np.sin(n * np.pi * x / L)
        # return ((np.exp(j * np.pi * x)* np.exp((- (self.a)*np.pi**2+self.b * j * np.pi + self.c)*t) - np.exp(-j * np.pi * x)* np.exp((- (self.a)*np.pi**2 - self.b * j * np.pi + self.c)*t)))/(2*j)
        return np.sin(np.pi * x + b_ori * np.pi  * t) * np.exp((-a_ori * np.pi ** 2  + c_ori) *t)


def gen_exact_solution(eq):
    """Generates exact solution for the heat equation for the given values of x and t."""
    # Number of points in each dimension:
    x_dim, t_dim = (256, 201)

    # Bounds of 'x' and 't':
    x_min, t_min = (0, 0.0)
    x_max, t_max = (L, 1.0)

    # Create tensors:
    t = np.linspace(t_min, t_max, num=t_dim).reshape(t_dim, 1)
    x = np.linspace(x_min, x_max, num=x_dim).reshape(x_dim, 1)
    usol = np.zeros((x_dim, t_dim), dtype=np.complex128).reshape(x_dim, t_dim)


    # Obtain the value of the exact solution for each generated point:
    for i in range(x_dim):
        for k in range(t_dim):
            usol[i][k] = heat_eq_exact_solution(x[i], t[k])


    # Save solution:
    np.savez(f"/home/yusaini.a/PINN/nonzero_param/heat_eq_data_{eq}", x=x, t=t, usol=usol)


def gen_testdata(eq):
    """Import and preprocess the dataset with the exact solution."""
    # Load the data:
    data = np.load(f"/home/yusaini.a/PINN/nonzero_param/heat_eq_data_{eq}.npz")
    # Obtain the values for t, x, and the excat solution:
    t, x, exact = data["t"], data["x"], data["usol"].T
    # Process the data and flatten it out (like labels and features):
    xx, tt = np.meshgrid(x, t)
    X = np.vstack((np.ravel(xx), np.ravel(tt))).T
    y = exact.flatten()[:, None]
    return X, y
    
class pde_heat_i:
    def __init__(self, m, a, b, c):
        self.m = m
        self.a = a
        self.b = b 
        self.c = c 
    
    def gen_traindata(self):
        data = np.load(f"/home/yusaini.a/PINN/nonzero_param/heat_eq_data_{self.m}.npz")
        t, x, exact = data["t"], data["x"], data["usol"].T
        # Process the data and flatten it out (like labels and features):
        xx, tt = np.meshgrid(x, t)
        X = np.vstack((np.ravel(xx), np.ravel(tt))).T
        y = exact.flatten()[:, None]
        return X, y
    
    def pde(self, x, y):
        """Expresses the PDE residual of the heat equation."""
        dy_t = dde.grad.jacobian(y, x, i=0, j=1)
        dy_x = dde.grad.jacobian(y, x, i=0, j=0)
        dy_xx = dde.grad.hessian(y, x, i=0, j=0)
        return dy_t - self.a * dy_xx - self.b * dy_x - self.c * y
    
def inverse_heat(m, a_start, b_start, c_start, parameter):
    run = Run()

    params = {"a":parameter[m][0], "b":parameter[m][1], "c" : parameter[m][2]}
    run["pde_params"] = params
    run["Experiment name"] = "provided solution"
    
    gen_exact_solution(m)

    
    run.name = f"20000 iterations, inverse problem {parameter[m]}, increased test points"

    a = dde.Variable(a_start)
    b = dde.Variable(b_start)
    c = dde.Variable(c_start)

    hpde = pde_heat_i(m, a, b, c)
    observe_x, y = hpde.gen_traindata()
    observe_y = dde.icbc.PointSetBC(observe_x, y, component=0)

    
    data_settings = dde.data.TimePDE(
        geomtime,
        hpde.pde,
        [bc, ic, observe_y],
        num_domain=2000,
        num_boundary=100,
        num_initial=100,
        anchors=observe_x,
        num_test=2000,
    )

    net = dde.nn.FNN([2] + [20] * 3 + [1], "tanh", "Glorot uniform")

    model = dde.Model(data_settings , net)
    
    model.compile("adam", lr=0.001, external_trainable_variables=[a, b, c]) #changed from 0.01
    

    variable = dde.callbacks.VariableValue([a, b, c], period=1000, filename=f"/home/yusaini.a/PINN/nonzero_param/variables_{m}.dat")

    losshistory, train_state = model.train(iterations=20000, callbacks=[variable])
    
    loss_train = np.array([np.sum(loss) for loss in losshistory.loss_train])
    loss_test = np.array([np.sum(loss) for loss in losshistory.loss_test])
    for i in range(len(loss_train)):
        run.track(loss_train[i], name="Training Loss")

    for i in range(len(loss_test)):
        run.track(loss_test[i], name="Test Loss")
    
    output_dir = "/home/yusaini.a/PINN/nonzero_param"

    return losshistory, train_state, model, hpde



def plot_pde(model, hpde):
    # print(f"Loss history plot saved as {output_dir}/loss_history.png")

    X, y_true = hpde.gen_traindata()
    y_pred = model.predict(X)
    
    # print(X, y_true.real)
    # print(y_true.real.shape, y_pred.shape)
    y_diff = y_true.real.flatten() - y_pred.flatten()


    fig1 = px.scatter_3d(x=X[:,0], y=X[:,1], z=y_true.real.flatten())
    fig1.show()
    
    fig2 = px.scatter_3d(x=X[:,0], y=X[:,1], z=y_pred.flatten())
    fig2.show()
    fig = px.scatter_3d(x=X[:,0], y=X[:,1], z=y_diff.flatten())
    fig.show()


param = np.array([[2.0, 0, 0], [2.0, np.pi, 0], [1.0, 3.0, 4.0 ]])


L = 1
n = 1
geom = dde.geometry.Interval(0, L)
timedomain = dde.geometry.TimeDomain(0, 1)
geomtime = dde.geometry.GeometryXTime(geom, timedomain)


bc = dde.icbc.DirichletBC(geomtime, lambda x: 0, lambda _, on_boundary: on_boundary)
ic = dde.icbc.IC(
    geomtime,
    lambda x: np.sin(n * np.pi * x[:, 0:1] / L),
    lambda _, on_initial: on_initial,
)

# partial_inverse_heat = partial(inverse_heat, a_start = 0.5, b_start = 0.5, c_start = 0.5, parameter = param )
# with concurrent.futures.ProcessPoolExecutor() as executor: 
#     index_param = range(param.shape[0])
#     results = executor.map(partial_inverse_heat, index_param) # run do_something with every item of the list secs
for p in range(param.shape[0]):
    
    a_ori = param[p][0]
    b_ori = param[p][1]
    c_ori = param[p][2]

    losshistory, train_state, model, hpde = inverse_heat(p, 0.5, 0.5, 0.5, param)
    
    dde.saveplot(losshistory, train_state, issave=True, isplot=True, output_dir="/home/yusaini.a/PINN/nonzero_param")
    # Now explicitly plot and save the loss history as an image
    # plt.figure()
    # Plot the training and testing loss history
    # dde.utils.plot_loss_history(losshistory)
    # # Save the plot as a PNG file
    # plt.savefig(f"{output_dir}/loss_history_{m}.png")
    # plt.close()

    plot_pde(model, hpde)