-
Notifications
You must be signed in to change notification settings - Fork 6
/
estimate.py
74 lines (57 loc) · 2.33 KB
/
estimate.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
import numpy as np
class ParameterEstimate:
def __init__(self, dim_x, x_id, interpolator_lst, basis_test, ode_functional, t_int, t_weight):
# number of equations / variables
self.dim_x = dim_x
# the index of variables being estimated
self.x_id = x_id
assert x_id < dim_x
# interpolator_lst: list of list
# top level list of N samples
# second level list of P variables
self.B = len(interpolator_lst)
self.interpolator_lst = interpolator_lst
for i in interpolator_lst:
assert len(i) == dim_x
self.basis_test = basis_test
# take theta, generate function f(x)
# f(x): T, B, D -> T, B
self.ode_functional = ode_functional
# grid and weights for integration: returned by generate integration grid
self.t_int = t_int
self.t_weight = t_weight
self.t_len = len(t_int)
# cache
self.x_hat = np.zeros((self.t_len, self.B, self.dim_x))
self.constants = None
self.basis_design = None
self.compute_x_hat_cache()
self.compute_test_basis_cache()
self.compute_constant_cache()
def compute_x_hat_cache(self):
# x_hat
# T, B, D
for b in range(self.B):
for d in range(self.dim_x):
interpolator = self.interpolator_lst[b][d]
self.x_hat[:, b, d] = interpolator.x_hat(self.t_int)
def compute_test_basis_cache(self):
# generate design matrix of basis function
# T, L
self.basis_design = self.basis_test.design_matrix(self.t_int)
def compute_constant_cache(self):
# T, B
x_hat_single = self.x_hat[:, :, self.x_id]
# generate design matrix of basis function
# T, L
basis_derivative_design = self.basis_test.design_matrix(self.t_int, derivative=True)
# B, L
self.constants = np.matmul((x_hat_single * self.t_weight[:, None]).T, basis_derivative_design)
def compute_loss_theta(self, theta):
f = self.ode_functional(theta)
# (T ; T, B, D) -> T, B
vec_field = f(self.t_int, self.x_hat)[:, :, self.x_id]
# B, L
theta_pred = np.matmul((vec_field * self.t_weight[:, None]).T, self.basis_design)
loss = np.sum((self.constants + theta_pred) ** 2)
return loss