-
Notifications
You must be signed in to change notification settings - Fork 0
/
lbfgs_annulus_main.py
103 lines (84 loc) · 4.06 KB
/
lbfgs_annulus_main.py
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
import numpy as np
import tensorflow as tf
import tensorflow_probability as tfp
from base_classes.pinn import PINN_Elastic2D
from utils.plot_functions import *
import sys
# gauss points
gauss_points = scipy.io.loadmat('mesh/annulus_gauss_pt_wt.mat')['gauss_pt_wt']
weights = tf.convert_to_tensor(np.prod(gauss_points[:,2:], axis=1), dtype=tf.float64)
gauss_points = tf.convert_to_tensor(gauss_points[:,:2], dtype=tf.float64)
plot_gauss_points(gauss_points, title='Mesh')
# inner boundary
inner_boundary = tf.stack((tf.cos(np.linspace(0, np.pi/2, 1000)),tf.sin(np.linspace(0, np.pi/2,1000))), axis=1)
# plot nodes
plot_X, plot_Y = tf.meshgrid(tf.linspace(0,4,800), tf.linspace(0,4,800))
plot_nodes = tf.cast(tf.stack((tf.reshape(plot_X, (-1,)), tf.reshape(plot_Y, (-1,))), axis=1), dtype=tf.float64)
class Hybrid(PINN_Elastic2D):
def __init__(self, E, nu, layer_sizes, lb, ub, training_nodes, weights, activation, boundary):
super().__init__(E,
nu,
layer_sizes,
lb,
ub,
weights,
activation)
self.training_nodes = training_nodes
self.boundary = boundary
def set_other_params(self, P=10):
self.P = P
def traction_work_done(self, x):
work_done = tf.reduce_mean(tf.reduce_sum(self.P*self.boundary*self(self.boundary), axis=1))*np.pi/2
return work_done
def dirichlet_bc(self, x, y):
return x*y
def train(self, adam_steps=100, lbfgs=False, max_iterations=100, num_correction_pairs=10, max_line_search_iterations=50):
# Optimisation steps for Adam
for i in range(adam_steps):
self.train_step(self.training_nodes)
# Optimisation steps for lbfgs
if lbfgs:
self.lbfgs_setup()
lbfgs_func = self.lbfgs_function()
# convert initial model parameters to a 1D tf.Tensor
init_params = tf.dynamic_stitch(self.idx, self.trainable_weights)
# train the model with L-BFGS solver
results = tfp.optimizer.lbfgs_minimize(value_and_gradients_function=lbfgs_func,
initial_position=init_params,
max_iterations=max_iterations,
max_line_search_iterations=max_line_search_iterations,
num_correction_pairs=num_correction_pairs)
# assign back the update parameters
self.assign_new_model_parameters(results.position)
if __name__=="__main__":
tf.keras.backend.set_floatx("float64")
pinn = Hybrid(E=1.0E5,
nu=0.3,
layer_sizes=[2,30,30,30,2],
lb = tf.reduce_min(gauss_points, axis=0),
ub = tf.reduce_max(gauss_points, axis=0),
training_nodes=gauss_points,
weights=weights,
activation = tf.nn.tanh,
boundary=inner_boundary)
pinn.set_other_params(P=10)
pinn.train(adam_steps=int(sys.argv[1]),
lbfgs=True,
max_iterations=int(sys.argv[1]),
max_line_search_iterations=50,
num_correction_pairs=100)
plt.figure(figsize=(5,4))
plt.plot(range(len(pinn.adam_history)), pinn.adam_history, label='Adam')
plt.plot(range(len(pinn.adam_history), len(pinn.adam_history)+len(pinn.lbfgs_history)), pinn.lbfgs_history, label='L-BFGS')
ll = min(min(pinn.adam_history), min(pinn.lbfgs_history))
ul = max(pinn.adam_history)
plt.ylim([-0.002, 0.1])
plt.legend()
plt.tight_layout()
plt.savefig('loss.png')
plt.show()
u = pinn(plot_nodes).numpy()
condition1 = tf.norm(plot_nodes, axis=1) > 4
condition2 = tf.norm(plot_nodes, axis=1) < 1
plot_scaler_field(u[:,0], title='ux', shape=plot_X.shape, conditions=[condition1, condition2])
plot_scaler_field(u[:,1], title='uy', shape=plot_X.shape, conditions=[condition1, condition2])