first commit

Helmholtz/Helmholtz2D.py

@@ -0,0 +1,227 @@
+import tensorflow as tf
+import numpy as np
+import matplotlib.pyplot as plt
+from scipy.interpolate import griddata
+import seaborn as sns
+from Helmholtz2D_model_tf import Sampler, Helmholtz2D
+
+if __name__ == '__main__':
+
+ a_1 = 1
+ a_2 = 4
+
+ def u(x, a_1, a_2):
+ return np.sin(a_1 * np.pi * x[:, 0:1]) * np.sin(a_2 * np.pi * x[:, 1:2])
+
+ def u_xx(x, a_1, a_2):
+ return - (a_1 * np.pi) ** 2 * np.sin(a_1 * np.pi * x[:, 0:1]) * np.sin(a_2 * np.pi * x[:, 1:2])
+
+ def u_yy(x, a_1, a_2):
+ return - (a_2 * np.pi) ** 2 * np.sin(a_1 * np.pi * x[:, 0:1]) * np.sin(a_2 * np.pi * x[:, 1:2])
+
+ # Forcing
+ def f(x, a_1, a_2, lam):
+ return u_xx(x, a_1, a_2) + u_yy(x, a_1, a_2) + lam * u(x, a_1, a_2)
+
+ def operator(u, x1, x2, lam, sigma_x1=1.0, sigma_x2=1.0):
+ u_x1 = tf.gradients(u, x1)[0] / sigma_x1
+ u_x2 = tf.gradients(u, x2)[0] / sigma_x2
+ u_xx1 = tf.gradients(u_x1, x1)[0] / sigma_x1
+ u_xx2 = tf.gradients(u_x2, x2)[0] / sigma_x2
+ residual = u_xx1 + u_xx2 + lam * u
+ return residual
+
+ # Parameter
+ lam = 1.0
+
+ # Domain boundaries
+ bc1_coords = np.array([[-1.0, -1.0],
+ [1.0, -1.0]])
+ bc2_coords = np.array([[1.0, -1.0],
+ [1.0, 1.0]])
+ bc3_coords = np.array([[1.0, 1.0],
+ [-1.0, 1.0]])
+ bc4_coords = np.array([[-1.0, 1.0],
+ [-1.0, -1.0]])
+
+ dom_coords = np.array([[-1.0, -1.0],
+ [1.0, 1.0]])
+
+ # Create initial conditions samplers
+ ics_sampler = None
+
+ # Create boundary conditions samplers
+ bc1 = Sampler(2, bc1_coords, lambda x: u(x, a_1, a_2), name='Dirichlet BC1')
+ bc2 = Sampler(2, bc2_coords, lambda x: u(x, a_1, a_2), name='Dirichlet BC2')
+ bc3 = Sampler(2, bc3_coords, lambda x: u(x, a_1, a_2), name='Dirichlet BC3')
+ bc4 = Sampler(2, bc4_coords, lambda x: u(x, a_1, a_2), name='Dirichlet BC4')
+ bcs_sampler = [bc1, bc2, bc3, bc4]
+
+ # Create residual sampler
+ res_sampler = Sampler(2, dom_coords, lambda x: f(x, a_1, a_2, lam), name='Forcing')
+
+ # Define model
+ mode = 'M1' # Method: 'M1', 'M2', 'M3', 'M4'
+ stiff_ratio = False # Log the eigenvalues of Hessian of losses
+
+ layers = [2, 50, 50, 50, 1]
+ model = Helmholtz2D(layers, operator, ics_sampler, bcs_sampler, res_sampler, lam, mode, stiff_ratio)
+
+ # Train model
+ model.train(nIter=40001, batch_size=128)
+
+ # Test data
+ nn = 100
+ x1 = np.linspace(dom_coords[0, 0], dom_coords[1, 0], nn)[:, None]
+ x2 = np.linspace(dom_coords[0, 1], dom_coords[1, 1], nn)[:, None]
+ x1, x2 = np.meshgrid(x1, x2)
+ X_star = np.hstack((x1.flatten()[:, None], x2.flatten()[:, None]))
+
+ # Exact solution
+ u_star = u(X_star, a_1, a_2)
+ f_star = f(X_star, a_1, a_2, lam)
+
+ # Predictions
+ u_pred = model.predict_u(X_star)
+ f_pred = model.predict_r(X_star)
+
+ # Relative error
+ error_u = np.linalg.norm(u_star - u_pred, 2) / np.linalg.norm(u_star, 2)
+ error_f = np.linalg.norm(f_star - f_pred, 2) / np.linalg.norm(f_star, 2)
+
+ print('Relative L2 error_u: {:.2e}'.format(error_u))
+ print('Relative L2 error_u: {:.2e}'.format(error_f))
+
+ ### Plot ###
+
+ # Exact solution & Predicted solution
+ # Exact soluton
+ U_star = griddata(X_star, u_star.flatten(), (x1, x2), method='cubic')
+ F_star = griddata(X_star, f_star.flatten(), (x1, x2), method='cubic')
+
+ # Predicted solution
+ U_pred = griddata(X_star, u_pred.flatten(), (x1, x2), method='cubic')
+ F_pred = griddata(X_star, f_pred.flatten(), (x1, x2), method='cubic')
+
+ fig_1 = plt.figure(1, figsize=(18, 5))
+ plt.subplot(1, 3, 1)
+ plt.pcolor(x1, x2, U_star, cmap='jet')
+ plt.colorbar()
+ plt.xlabel(r'$x_1$')
+ plt.ylabel(r'$x_2$')
+ plt.title('Exact $u(x)$')
+
+ plt.subplot(1, 3, 2)
+ plt.pcolor(x1, x2, U_pred, cmap='jet')
+ plt.colorbar()
+ plt.xlabel(r'$x_1$')
+ plt.ylabel(r'$x_2$')
+ plt.title('Predicted $u(x)$')
+
+ plt.subplot(1, 3, 3)
+ plt.pcolor(x1, x2, np.abs(U_star - U_pred), cmap='jet')
+ plt.colorbar()
+ plt.xlabel(r'$x_1$')
+ plt.ylabel(r'$x_2$')
+ plt.title('Absolute error')
+ plt.tight_layout()
+ plt.show()
+
+ # Residual loss & Boundary loss
+ loss_res = model.loss_res_log
+ loss_bcs = model.loss_bcs_log
+
+ fig_2 = plt.figure(2)
+ ax = fig_2.add_subplot(1, 1, 1)
+ ax.plot(loss_res, label='$\mathcal{L}_{r}$')
+ ax.plot(loss_bcs, label='$\mathcal{L}_{u_b}$')
+ ax.set_yscale('log')
+ ax.set_xlabel('iterations')
+ ax.set_ylabel('Loss')
+ plt.legend()
+ plt.tight_layout()
+ plt.show()
+
+ # Adaptive Constant
+ adaptive_constant = model.adpative_constant_log
+
+ fig_3 = plt.figure(3)
+ ax = fig_3.add_subplot(1, 1, 1)
+ ax.plot(adaptive_constant, label='$\lambda_{u_b}$')
+ ax.set_xlabel('iterations')
+ plt.legend()
+ plt.tight_layout()
+ plt.show()
+
+ # Gradients at the end of training
+ data_gradients_res = model.dict_gradients_res_layers
+ data_gradients_bcs = model.dict_gradients_bcs_layers
+
+ gradients_res_list = []
+ gradients_bcs_list = []
+
+ num_hidden_layers = len(layers) - 1
+ for j in range(num_hidden_layers):
+ gradient_res = data_gradients_res['layer_' + str(j + 1)][-1]
+ gradient_bcs = data_gradients_bcs['layer_' + str(j + 1)][-1]
+
+ gradients_res_list.append(gradient_res)
+ gradients_bcs_list.append(gradient_bcs)
+
+ cnt = 1
+ fig_4 = plt.figure(4, figsize=(13, 4))
+ for j in range(num_hidden_layers):
+ ax = plt.subplot(1, 4, cnt)
+ ax.set_title('Layer {}'.format(j + 1))
+ ax.set_yscale('symlog')
+ gradients_res = data_gradients_res['layer_' + str(j + 1)][-1]
+ gradients_bcs = data_gradients_bcs['layer_' + str(j + 1)][-1]
+ sns.distplot(gradients_bcs, hist=False,
+ kde_kws={"shade": False},
+ norm_hist=True, label=r'$\nabla_\theta \lambda_{u_b} \mathcal{L}_{u_b}$')
+ sns.distplot(gradients_res, hist=False,
+ kde_kws={"shade": False},
+ norm_hist=True, label=r'$\nabla_\theta \mathcal{L}_r$')
+
+ ax.get_legend().remove()
+ ax.set_xlim([-3.0, 3.0])
+ ax.set_ylim([0,100])
+ cnt += 1
+ handles, labels = ax.get_legend_handles_labels()
+ fig_4.legend(handles, labels, loc="upper left", bbox_to_anchor=(0.35, -0.01),
+ borderaxespad=0, bbox_transform=fig_4.transFigure, ncol=2)
+ plt.tight_layout()
+ plt.show()
+
+ # Eigenvalues if applicable
+ if stiff_ratio:
+ eigenvalues_list = model.eigenvalue_log
+ eigenvalues_bcs_list = model.eigenvalue_bcs_log
+ eigenvalues_res_list = model.eigenvalue_res_log
+ eigenvalues_res = eigenvalues_res_list[-1]
+ eigenvalues_bcs = eigenvalues_bcs_list[-1]
+
+ fig_5 = plt.figure(5)
+ ax = fig_5.add_subplot(1, 1, 1)
+ ax.plot(eigenvalues_res, label='$\mathcal{L}_r$')
+ ax.plot(eigenvalues_bcs, label='$\mathcal{L}_{u_b}$')
+ ax.set_xlabel('index')
+ ax.set_ylabel('eigenvalue')
+ ax.set_yscale('symlog')
+ plt.legend()
+ plt.tight_layout()
+ plt.show()
+
+
+
+
+
+
+
+
+
+
+
+
+
+