Implementation of forward PINN on Polymerization reaction dataset

[AbdullahElyas]

Hi,
I am trying to implement forward PINN on Polymerization reaction dataset. The dataset is generated using CFD. There are two different domains defined:

1. Inner domain which has a rotating reference frame
   2)Outer domain has a stationary reference frame

The PDE equation is different for both the domains. I am trying to calculate the loss from PDE residuals for both the domains. I have tried to separately calculate the PDE residuals within the pde(x,y) function and then added residuals from both the domains to get net residuals at each training point. My question is: Is there a better way anyone can suggest to implement this. Currently I am not getting right results.

I have defined a dummy hypercube geometry of dimension 6 as my input size is 6. I have used anchors to define the training points. Moreover, I have used the pointsetBC for providing the training dataset.

I am attaching the code here .I would appreciate any assistance you can provide.

Thanks

loading data function

def load_training_data(num):

with open('..\data\data.pkl', 'rb') as f:
    data_X, _, data_Y, _,_,data_p = pickle.load(f)
idx = np.random.choice(data_X.shape[0], num, replace=False)
data_sampled_X = data_X[idx,:]
data_sampled_Y = data_Y[idx,:]
data_sampled_Y = data_Y[idx,:]
return data_sampled_X,data_sampled_Y,data_p

loading data

data_X,data_Y,data_p = load_training_data(5000)
print(data_X.shape)
print(data_Y.shape)
ob_u = data_Y[:,0:1]
ob_v = data_Y[:,1:2]
ob_w = data_Y[:,2:3]
ob_p = data_Y[:,3:4]
ob_T = data_Y[:,4:5]
ob_I = data_Y[:,5:6]
ob_M = data_Y[:,6:7]
ob_xyt = data_X
ob_xyt2 = data_p

del data_X,data_Y

pde function for pde residual calculation

def pde(x, y):

u = y[:, 0:1]
v = y[:, 1:2]
w = y[:, 2:3]
p = y[:, 3:4]
T = y[:, 4:5]
I = y[:, 5:6]


index = (x[:, 5:6] == 1).view(x.shape[0])
in_idx = torch.nonzero(index).flatten()

index = (x[:, 5:6] == 0).view(x.shape[0])
out_idx = torch.nonzero(index).flatten()



u_x = dde.grad.jacobian(y, x, i=0, j=0)
u_y = dde.grad.jacobian(y, x, i=0, j=1)
u_z = dde.grad.jacobian(y, x, i=0, j=2)

v_x = dde.grad.jacobian(y, x, i=1, j=0)
v_y = dde.grad.jacobian(y, x, i=1, j=1)
v_z = dde.grad.jacobian(y, x, i=1, j=2)

w_x = dde.grad.jacobian(y, x, i=2, j=0)
w_y = dde.grad.jacobian(y, x, i=2, j=1)
w_z = dde.grad.jacobian(y, x, i=2, j=2)

p_x = dde.grad.jacobian(y, x, i=3, j=0)
p_y = dde.grad.jacobian(y, x, i=3, j=1)
p_z = dde.grad.jacobian(y, x, i=3, j=2)

I_x = dde.grad.jacobian(y, x, i=4, j=0)
I_y = dde.grad.jacobian(y, x, i=4, j=1)
I_z = dde.grad.jacobian(y, x, i=4, j=2)

u_xx = dde.grad.hessian(y, x, component=0, i=0, j=0)
u_yy = dde.grad.hessian(y, x, component=0, i=1, j=1)
u_zz = dde.grad.hessian(y, x, component=0, i=2, j=2)

v_xx = dde.grad.hessian(y, x, component=1, i=0, j=0)
v_yy = dde.grad.hessian(y, x, component=1, i=1, j=1)
v_zz = dde.grad.hessian(y, x, component=1, i=2, j=2)

w_xx = dde.grad.hessian(y, x, component=2, i=0, j=0)
w_yy = dde.grad.hessian(y, x, component=2, i=1, j=1)
w_zz = dde.grad.hessian(y, x, component=2, i=2, j=2)






vel = 10.471975499999981

# Relative Velocity from MRF
ur = u - vel * x[:,1:2]
vr = v + vel * x[:,0:1]

ur_x = u_x
ur_y = u_y - vel
ur_z = u_z

vr_x = v_x + vel
vr_y = v_y
vr_z = v_z

# For calculation of residuals in the innder domain
f_cont_in = u_x + v_y + w_z

f_u_in = rho * (ur_x * ur + ur_y * vr + ur_z * w) + p_x - mu * (u_xx + u_yy + u_zz) - (vel * vel * x[:,0:1]) - (2 * vel * vr)
f_v_in = rho * (vr_x * ur + vr_y * vr + vr_z * w) + p_y - mu * (v_xx + v_yy + v_zz) - (vel * vel * x[:,1:2]) + (2 * vel * ur)
f_w_in = rho * (w_x * ur + w_y * vr + w_z * w) + p_z - mu * (w_xx + w_yy + w_zz)

f_I_in = (u * I_x + u_x * I + v * I_y + v_y * I + w * I_z + w_z * I) - k_i * torch.exp(-Ea / torch.abs(T))


#For the calculation of residuals in the outer domain

f_cont_out = u_x + v_y + w_z

f_u_out = rho * (u_x * u + u_y * v + u_z * w) + p_x - mu * (u_xx + u_yy + u_zz)
f_v_out = rho * (v_x * u + v_y * v + v_z * w) + p_y - mu * (v_xx + v_yy + v_zz)
f_w_out = rho * (w_x * u + w_y * v + w_z * w) + p_z - mu * (w_xx + w_yy + w_zz)

f_I_out = (u * I_x + u_x * I + v * I_y + v_y * I + w * I_z + w_z * I) - k_i * torch.exp(-Ea / torch.abs(T))



# Create masks 
mask_in = torch.zeros_like(x[:, 0:1], dtype=torch.float)
mask_in[in_idx] = 1.0

mask_out = torch.zeros_like(x[:, 0:1], dtype=torch.float)
mask_out[out_idx] = 1.0

# Apply masks to residuals
f_cont_in *= mask_in
f_u_in *= mask_in
f_v_in *= mask_in
f_w_in *= mask_in
f_I_in *= mask_in

f_cont_out *= mask_out
f_u_out *= mask_out
f_v_out *= mask_out
f_w_out *= mask_out
f_I_out *= mask_out



# add the tensors f_cont_in f_u_in f_v_in f_w_in f_I_in and f_cont_out f_u_out f_v_out f_w_out f_I_out 

f_cont = f_cont_in + f_cont_out
f_u = f_u_in + f_u_out
f_v= f_v_in + f_v_out
f_w = f_w_in + f_w_out
f_I = f_I_in + f_I_out

# add f_cont f_u f_v f_w f_I 




return f_cont, f_u, f_v, f_w, f_I

xmin = [0, 0, 0, 0, 0, 0]
xmax = [1, 1, 1, 1, 1, 1]

hypercube = dde.geometry.Hypercube(xmin, xmax)

observe_u = dde.icbc.PointSetBC(ob_xyt, ob_u, component=0)
observe_v = dde.icbc.PointSetBC(ob_xyt, ob_v, component=1)
observe_w = dde.icbc.PointSetBC(ob_xyt, ob_w, component=2)
observe_p = dde.icbc.PointSetBC(ob_xyt, ob_p, component=3)
observe_T = dde.icbc.PointSetBC(ob_xyt, ob_T, component=4)
observe_I = dde.icbc.PointSetBC(ob_xyt, ob_I, component=5)
observe_M = dde.icbc.PointSetBC(ob_xyt, ob_M, component=6)

data = dde.data.PDE(
hypercube,
pde,
[observe_u,observe_v,observe_w,observe_p,observe_T,observe_I,observe_M],
num_domain=0,
num_boundary=0,
anchors=ob_xyt2
)

Set up early stopping. Monitor the loss value, stop training when loss stops improving by a certain amount.

early_stopping = EarlyStopping(min_delta=1e-5, patience=50)

checkpointer = dde.callbacks.ModelCheckpoint(
"..\checkpoint\model", verbose=1, save_better_only=True
)

Neural Network setup

layer_size = [6] + [50] * 6 + [7]
activation = "tanh"
initializer = "Glorot uniform"
net = dde.maps.FNN(layer_size, activation, initializer)
resampler = dde.callbacks.PDEPointResampler(period=100)
model = dde.Model(data, net)

Compile, train and save model

model.compile("adam", lr=1e-4, loss_weights=[1, 0.01, 0.01, 0.01,1,1,1,1,0.01,0.001,1,1])

model.restore("model-200.pt")

loss_history, train_state = model.train(
iterations=10000, display_every=1, disregard_previous_best=True, callbacks=[resampler,checkpointer,early_stopping])

model.save(save_path = ".\model2")