-
Notifications
You must be signed in to change notification settings - Fork 0
/
simpleHarmonicExperimentalData.py
166 lines (127 loc) · 5.27 KB
/
simpleHarmonicExperimentalData.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
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
from typing import Callable
import argparse
import matplotlib.pyplot as plt
import torch
from torch import nn
import numpy as np
import torchopt
from matplotlib.animation import FuncAnimation
from pinn import make_forward_fn
#Based off of https://github.com/madagra/basic-pinn.
# Oscillator motion parameters
m = 1.0 # mass
k = 1.0 # spring constant
x0 = 1.0 # initial displacement
v0 = 0.0 # initial velocity
# Boundary conditions
X_BOUNDARY_1 = 0.0
F_BOUNDARY_1 = x0
X_BOUNDARY_2 = 0.0
F_BOUNDARY_2 = v0
def make_loss_fn(f: Callable, d2fdx2: Callable, x_data: torch.Tensor, y_data: torch.Tensor) -> Callable:
def loss_fn(params: torch.Tensor, x: torch.Tensor):
# interior loss
f_value = f(x, params)
interior = d2fdx2(x, params) + k / m * f_value
# boundary losses
x_boundary_1 = torch.tensor([X_BOUNDARY_1])
f_boundary_1 = torch.tensor([F_BOUNDARY_1])
x_boundary_2 = torch.tensor([X_BOUNDARY_2])
f_boundary_2 = torch.tensor([F_BOUNDARY_2])
boundary_1 = f(x_boundary_1, params) - f_boundary_1
boundary_2 = d2fdx2(x_boundary_2, params) - f_boundary_2
#data loss
f_data = f(x_data, params)
loss = nn.MSELoss()
data_loss = loss(f_data, y_data)
# Weighting of the loss
weight_interior = 0.25
weight_boundary_1 = 0.37
weight_boundary_2 = 0.37
weight_data = 1.0
loss_value = (
weight_interior * loss(interior, torch.zeros_like(interior))
+ weight_boundary_1 * loss(boundary_1, torch.zeros_like(boundary_1))
+ weight_boundary_2 * loss(boundary_2, torch.zeros_like(boundary_2))
+ weight_data * data_loss
)
return loss_value
return loss_fn
if __name__ == "__main__":
# make it reproducible
torch.manual_seed(2)
# parse input from user
parser = argparse.ArgumentParser()
parser.add_argument("-n", "--num-hidden", type=int, default=4)
parser.add_argument("-d", "--dim-hidden", type=int, default=20)
parser.add_argument("-b", "--batch-size", type=int, default=100)
parser.add_argument("-lr", "--learning-rate", type=float, default=1e-2)
parser.add_argument("-e", "--num-epochs", type=int, default=500)
args = parser.parse_args()
# configuration
num_hidden = args.num_hidden
dim_hidden = args.dim_hidden
batch_size = args.batch_size
num_iter = args.num_epochs
tolerance = 1e-8
learning_rate = args.learning_rate
domain = (-5.0, 5.0)
analytical_sol_fn = (
lambda x: x0 * np.cos(np.sqrt(k / m) * x) + v0 / np.sqrt(k / m) * np.sin(np.sqrt(k / m) * x)
)
num_datapoints = 30
x_data = torch.FloatTensor(num_datapoints).uniform_(domain[0], domain[1]).reshape(-1, 1)
y_data = torch.tensor([analytical_sol_fn(x_i) for x_i in x_data]).reshape(-1, 1)
noise = torch.randn_like(y_data) * 0.1
y_data += noise
# function versions of model forward, gradient and loss
fmodel, params, funcs = make_forward_fn(
num_hidden=num_hidden, dim_hidden=dim_hidden, derivative_order=2
)
f = funcs[0]
d2fdx2 = funcs[2]
loss_fn = make_loss_fn(f, d2fdx2, x_data, y_data)
# choose optimizer with functional API using functorch
optimizer = torchopt.FuncOptimizer(torchopt.adam(lr=learning_rate))
# train the model
loss_evolution = []
fig, ax = plt.subplots()
x_eval = torch.linspace(domain[0], domain[1], steps=100).reshape(-1, 1)
x_eval_np = x_eval.detach().numpy()
def update(i):
global params
ax.clear()
# Sample points in the domain randomly for each epoch
x = torch.FloatTensor(batch_size).uniform_(domain[0], domain[1])
# Update the parameters
loss = loss_fn(params, x)
params = optimizer.step(loss, params)
loss_evolution.append(loss.item())
print("Epoch: {}, Loss: {}".format(i, loss.item()))
x_sample_np = torch.FloatTensor(batch_size).uniform_(domain[0], domain[1]).detach().numpy()
f_eval = f(x_eval, params)
# ax.scatter(x_sample_np, analytical_sol_fn(x_sample_np), color="red", label="Sample training points")
ax.plot(x_eval_np, f_eval.detach().numpy(), label="PINN solution at iter {}".format(i))
ax.plot(
x_eval_np,
analytical_sol_fn(x_eval_np),
label=f"Analytic solution",
color="green",
alpha=0.75,
)
ax.scatter(x_data, y_data, color="red", label="Sparse data points")
#set the title to the differential equation d2f/dt2 = -k / m * f(t), k=1,m=1, f(0)=1, f'(0)=0
# ax.set(title="Simple Harmonic Oscillator equation solved with PINNs (iter {})".format(i), xlabel="t", ylabel="f(t)")
ax.set(title="Sparse Data Simple Harmonic Oscillator PINN Solution\n"+r"$\frac{d^2f}{dt^2} = -\frac{k}{m}f(t),\ k=1,\ m=1,\ f(0)=1,\ f'(0)=0$", xlabel="t", ylabel="f(t)")
ax.set_ylim(-2,5)
ax.legend()
anim = FuncAnimation(fig, update, frames=num_iter, interval=10, repeat=False)
#save to gif
anim.save('HarmonicSparseData.gif', dpi=80, writer='imagemagick')
plt.show()
# # plot loss evolution
# plt.plot(loss_evolution)
# plt.title("Loss evolution")
# plt.xlabel("Epoch")
# plt.ylabel("Loss")
# plt.show()