├── Damped_Sine ├── create_nice_plot.py ├── example.py ├── neural_ode.py └── plotting.py ├── Glycolysis_noisy ├── HMC.py ├── create_nice_plot.py ├── neural_ode.py ├── plotting.py └── pre_train.py ├── Hybrid_Sine_NN ├── HMC.py ├── Show_NN_fit.py ├── create_nice_plot_hybrid.py ├── neural_ode.py ├── plotting.py └── pre_train.py ├── LV_noisy ├── create_nice_plot.py ├── example.py ├── neural_ode.py └── plotting.py ├── README.md └── Spiral_noisy ├── create_nice_plot.py ├── example.py ├── neural_ode.py └── plotting.py /Damped_Sine/create_nice_plot.py: -------------------------------------------------------------------------------- 1 | import os 2 | os.environ['KMP_DUPLICATE_LIB_OK']='True' 3 | 4 | import numpy as np 5 | import numpy.random as npr 6 | import tensorflow as tf 7 | import matplotlib.pyplot as plt 8 | plt.switch_backend('agg') 9 | 10 | from scipy.integrate import odeint 11 | from scipy import stats 12 | 13 | 14 | from plotting import newfig, savefig 15 | 16 | np.random.seed(1234) 17 | tf.set_random_seed(1234) 18 | 19 | if __name__ == "__main__": 20 | 21 | def Damped_pendulum(z, t, alpha, beta): 22 | x, y = z 23 | dzdt = [y, - alpha * y - beta * np.sin(x)] 24 | return dzdt 25 | 26 | Order = 2 27 | def Damped_pendulum_BNODE(z, t, W, Order): 28 | x, y = z 29 | x = x[None,None] 30 | y = y[None,None] 31 | temp = x ** 0 32 | for m in range(1, Order): 33 | for n in range(m + 1): 34 | temp = np.concatenate([temp, x**n * y**(m - n)], 1) 35 | 36 | candidate = np.matmul(temp, W) 37 | dzdt = [candidate[0,0], candidate[0,1]] 38 | return dzdt 39 | 40 | 41 | data_size = 8000 42 | batch_time = 200 43 | batch_size = 1000 44 | 45 | data_size_true = 2000 46 | 47 | 48 | N_samples = 100 49 | N_total = 500 50 | 51 | alpha = 0.2 52 | beta = 8.91 53 | 54 | 55 | z0 = [-1.193, -3.876] 56 | t_grid_train = np.linspace(0, 20, data_size) 57 | t_grid_true = np.linspace(0, 40, data_size_true) 58 | 59 | y_train = odeint(Damped_pendulum, z0, t_grid_train, args=(alpha, beta)) 60 | idx = np.random.choice(np.arange(data_size - batch_time - 1, dtype=np.int64), batch_size, replace=False) 61 | y_train = y_train[idx] 62 | t_grid_train = t_grid_train[idx] 63 | 64 | y_true = odeint(Damped_pendulum, z0, t_grid_true, args=(alpha, beta)) 65 | 66 | sigma_normal1 = np.std(y_train[:,0:1]) 67 | sigma_normal2 = np.std(y_train[:,1:2]) 68 | 69 | sigma_normal = np.asarray([sigma_normal1, sigma_normal2]) 70 | 71 | 72 | parameters = np.load("parameters.npy") 73 | precision = np.load("loggammalist.npy") 74 | loglikelihood = np.load("loglikelihood.npy") 75 | precision = np.exp(precision) 76 | print("precision", precision) 77 | print(parameters.shape) 78 | 79 | loglikelihood = loglikelihood[-N_total:] 80 | num_samples = parameters.shape[0] 81 | length_dict = parameters.shape[1] 82 | num_dim = parameters.shape[2] 83 | 84 | idx_MAP = np.argmin(loglikelihood) 85 | MAP = parameters[idx_MAP, :, :] 86 | 87 | y_MAP = odeint(Damped_pendulum_BNODE, z0, t_grid_true, args=(MAP, Order)) 88 | 89 | 90 | mu_pred = y_MAP 91 | 92 | y_BNODE = np.zeros((t_grid_true.shape[0], num_dim, N_samples)) 93 | 94 | for k in range(N_samples): 95 | print(k) 96 | idx_1 = np.random.randint(N_total) 97 | idx_2 = np.random.randint(N_total) 98 | W_sample = parameters[-idx_1, :, :] 99 | precision_here = precision[-idx_2] * num_dim 100 | y_BNODE[:,:,k] = odeint(Damped_pendulum_BNODE, z0, t_grid_true, args=(W_sample, Order)) 101 | Sigma_data = np.ones_like(mu_pred) / np.sqrt(precision_here) 102 | y_BNODE[:,:,k] = y_BNODE[:,:,k] + Sigma_data * sigma_normal * np.random.normal() 103 | 104 | mu_pred = np.mean(y_BNODE, axis = 2) 105 | Sigma_pred = np.var(y_BNODE, axis = 2) 106 | 107 | 108 | 109 | plt.rcParams.update(plt.rcParamsDefault) 110 | plt.rc('font', family='serif') 111 | plt.rcParams.update({'font.size': 16, 112 | 'lines.linewidth': 2, 113 | 'axes.labelsize': 20, # fontsize for x and y labels (was 10) 114 | 'axes.titlesize': 20, 115 | 'xtick.labelsize': 16, 116 | 'ytick.labelsize': 16, 117 | 'legend.fontsize': 20, 118 | 'axes.linewidth': 2, 119 | "pgf.texsystem": "pdflatex", # change this if using xetex or lautex 120 | "text.usetex": True, # use LaTeX to write all text 121 | }) 122 | 123 | plt.figure(1, figsize=(6,6)) 124 | plt.xticks(fontsize=14) 125 | plt.yticks(fontsize=14) 126 | plt.plot(y_train[:,0], y_train[:,1], 'ro', label = "Training data") 127 | plt.plot(y_true[:,0], y_true[:,1], 'b-', label = "Exact Trajectory") 128 | plt.xlabel('$x_1$',fontsize=18) 129 | plt.ylabel('$x_2$',fontsize=18) 130 | plt.xlim((-2.3, 2.3)) 131 | plt.ylim((-5., 6.5)) 132 | plt.legend(loc='upper right', frameon=False, prop={'size': 14}) 133 | plt.savefig('./Training_data.png', dpi = 300) 134 | 135 | 136 | plt.figure(4, figsize=(12,6.5)) 137 | plt.xticks(fontsize=22) 138 | plt.yticks(fontsize=22) 139 | plt.plot(t_grid_train, y_train[:,0], 'ro', label = "Training data of $x_1(t)$") 140 | plt.plot(t_grid_true, y_true[:,0], 'r-', label = "True Trajectory of $x_1(t)$") 141 | plt.plot(t_grid_true, y_MAP[:,0], 'g--', label = "MAP Trajectory of $x_1(t)$") 142 | lower_0 = mu_pred[:,0] - 2.0*np.sqrt(Sigma_pred[:,0]) 143 | upper_0 = mu_pred[:,0] + 2.0*np.sqrt(Sigma_pred[:,0]) 144 | plt.fill_between(t_grid_true.flatten(), lower_0.flatten(), upper_0.flatten(), 145 | facecolor='orange', alpha=0.5, label="Two std band") 146 | plt.xlabel('$t$',fontsize=26) 147 | plt.ylabel('$x_1(t)$',fontsize=26) 148 | plt.ylim((-3., 3.)) 149 | plt.legend(loc='upper right', frameon=False, prop={'size': 20}) 150 | plt.savefig('./BNODE_Prediction_x1.png', dpi = 300) 151 | 152 | 153 | plt.figure(5, figsize=(12,6.5)) 154 | plt.xticks(fontsize=22) 155 | plt.yticks(fontsize=22) 156 | plt.plot(t_grid_train, y_train[:,1], 'ro', label = "Training data of $x_2(t)$") 157 | plt.plot(t_grid_true, y_true[:,1], 'r-', label = "True Trajectory of $x_2(t)$") 158 | plt.plot(t_grid_true, y_MAP[:,1], 'g--', label = "MAP Trajectory of $x_2(t)$") 159 | lower_1 = mu_pred[:,1] - 2.0*np.sqrt(Sigma_pred[:,1]) 160 | upper_1 = mu_pred[:,1] + 2.0*np.sqrt(Sigma_pred[:,1]) 161 | plt.fill_between(t_grid_true.flatten(), lower_1.flatten(), upper_1.flatten(), 162 | facecolor='orange', alpha=0.5, label="Two std band") 163 | plt.xlabel('$t$',fontsize=26) 164 | plt.ylabel('$x_2(t)$',fontsize=26) 165 | plt.legend(loc='upper right', frameon=False, prop={'size': 20}) 166 | plt.savefig('./BNODE_Prediction_x2.png', dpi = 300) 167 | 168 | 169 | 170 | 171 | -------------------------------------------------------------------------------- /Damped_Sine/example.py: -------------------------------------------------------------------------------- 1 | import os 2 | os.environ['KMP_DUPLICATE_LIB_OK']='True' 3 | 4 | import numpy as np 5 | import numpy.random as npr 6 | import tensorflow as tf 7 | import matplotlib.pyplot as plt 8 | plt.switch_backend('agg') 9 | import tensorflow.contrib.eager as tfe 10 | from tensorflow.keras import layers, initializers 11 | from scipy.integrate import odeint 12 | 13 | 14 | keras = tf.keras 15 | tf.enable_eager_execution() 16 | 17 | from neural_ode import NeuralODE 18 | 19 | np.random.seed(1234) 20 | tf.set_random_seed(1234) 21 | 22 | if __name__ == "__main__": 23 | 24 | def Damped_pendulum(z, t, alpha, beta): 25 | x, y = z 26 | dzdt = [y, - alpha * y - beta * np.sin(x)] 27 | return dzdt 28 | 29 | def plot_spiral(trajectories, width = 1.): 30 | x = trajectories[:,0] 31 | y = trajectories[:,1] 32 | plt.plot(x, y, linewidth=width) 33 | 34 | 35 | data_size = 8001 36 | batch_time = 200 37 | niters = 1500 38 | batch_size = 1000 39 | 40 | 41 | alpha = 0.2 42 | beta = 8.91 43 | 44 | 45 | z0 = [-1.193, -3.876] 46 | true_y0 = tf.to_float([[-1.193, -3.876]]) 47 | t_grid = np.linspace(0, 20, data_size) 48 | 49 | true_yy = odeint(Damped_pendulum, z0, t_grid, args=(alpha, beta)) 50 | print(true_yy.shape) 51 | 52 | true_y = true_yy + 0.0 * np.random.randn(true_yy.shape[0], true_yy.shape[1]) 53 | 54 | sigma_normal1 = np.std(true_y[:,0:1]) 55 | sigma_normal2 = np.std(true_y[:,1:2]) 56 | 57 | true_y[:,0:1] = true_y[:,0:1]/ sigma_normal1 58 | true_y[:,1:2] = true_y[:,1:2]/ sigma_normal2 59 | 60 | true_y0 = tf.to_float([[-1.193/sigma_normal1, -3.876/sigma_normal2]]) 61 | 62 | 63 | def get_batch(): 64 | """Returns initial point and last point over sampled frament of trajectory""" 65 | starts = np.random.choice(np.arange(data_size - batch_time - 1, dtype=np.int64), batch_size, replace=False) 66 | batch_y0 = true_y[starts] 67 | batch_yN = true_y[starts + batch_time] 68 | return tf.to_float(batch_y0), tf.to_float(batch_yN) 69 | 70 | 71 | Order = 2 72 | num_param = 1 73 | for m in range(1, Order): 74 | for n in range(m + 1): 75 | num_param += 1 76 | 77 | num_param = num_param * 2 78 | 79 | 80 | t0 = t_grid[:batch_time][0] 81 | t1 = t_grid[:batch_time][-1] 82 | t_in = np.linspace(t0, t1, 20) 83 | 84 | 85 | batch_y0, batch_yN = get_batch() 86 | 87 | ######################################### 88 | ########## precondition start ########### 89 | ######################################### 90 | niters_pre = 1000 91 | 92 | class ODEModel_pre(tf.keras.Model): 93 | def __init__(self): 94 | super(ODEModel_pre, self).__init__() 95 | self.Weights = tfe.Variable(tf.random_normal([num_param//2, 2], dtype=tf.float32)*0.01, dtype=tf.float32) 96 | 97 | def call(self, inputs, **kwargs): 98 | t, y = inputs 99 | h = y 100 | x1 = h[:,0:1] * sigma_normal1 101 | x2 = h[:,1:2] * sigma_normal2 102 | temp = x1 ** 0 103 | for m in range(1, Order): 104 | for n in range(m + 1): 105 | temp = tf.concat([temp, x1**n * x2**(m - n)], 1) 106 | 107 | h_out = tf.matmul(temp, self.Weights) / np.asarray([sigma_normal1, sigma_normal2]) 108 | 109 | return h_out 110 | 111 | 112 | model_pre = ODEModel_pre() 113 | neural_ode_pre = NeuralODE(model_pre, t_in) 114 | optimizer = tf.train.AdamOptimizer(3e-2) 115 | 116 | 117 | def compute_gradients_and_update_pre(batch_y0, batch_yN): 118 | """Takes start positions (x0, y0) and final positions (xN, yN)""" 119 | pred_y = neural_ode_pre.forward(batch_y0) 120 | with tf.GradientTape() as g_pre: 121 | g_pre.watch(pred_y) 122 | loss = tf.reduce_mean((pred_y - batch_yN)**2) + tf.reduce_sum(tf.abs(model_pre.Weights)) 123 | 124 | dLoss = g_pre.gradient(loss, pred_y) 125 | h_start, dfdh0, dWeights = neural_ode_pre.backward(pred_y, dLoss) 126 | optimizer.apply_gradients(zip(dWeights, model_pre.weights)) 127 | return loss, dWeights 128 | 129 | # Compile EAGER graph to static (this will be much faster) 130 | compute_gradients_and_update_pre = tfe.defun(compute_gradients_and_update_pre) 131 | 132 | parameters_pre = np.zeros((niters_pre, num_param)) 133 | 134 | for step in range(niters_pre): 135 | print(step) 136 | loss, dWeights = compute_gradients_and_update_pre(batch_y0, batch_yN) 137 | 138 | model_parameters_pre = model_pre.trainable_weights[0].numpy().flatten() 139 | for k in range(num_param): 140 | parameters_pre[step, k] = model_parameters_pre[k] 141 | 142 | print(parameters_pre[step,:]) 143 | 144 | ######################################### 145 | ########## precondition end ############# 146 | ######################################### 147 | 148 | initial_weight = model_pre.trainable_weights[0].numpy() 149 | print(initial_weight.shape, "here") 150 | 151 | class ODEModel(tf.keras.Model): 152 | def __init__(self, initial_weight): 153 | super(ODEModel, self).__init__() 154 | self.Weights = tfe.Variable(initial_weight, dtype=tf.float32) 155 | 156 | def call(self, inputs, **kwargs): 157 | t, y = inputs 158 | h = y 159 | x1 = h[:,0:1] * sigma_normal1 160 | x2 = h[:,1:2] * sigma_normal2 161 | temp = x1 ** 0 162 | for m in range(1, Order): 163 | for n in range(m + 1): 164 | temp = tf.concat([temp, x1**n * x2**(m - n)], 1) 165 | 166 | h_out = tf.matmul(temp, self.Weights) / np.asarray([sigma_normal1, sigma_normal2]) 167 | return h_out 168 | 169 | 170 | model = ODEModel(initial_weight) 171 | neural_ode = NeuralODE(model, t = t_in) 172 | 173 | def compute_gradients_and_update(batch_y0, batch_yN): 174 | """Takes start positions (x0, y0) and final positions (xN, yN)""" 175 | pred_y = neural_ode.forward(batch_y0) 176 | with tf.GradientTape() as g: 177 | g.watch(pred_y) 178 | loss = tf.reduce_sum((pred_y - batch_yN)**2) 179 | 180 | dLoss = g.gradient(loss, pred_y) 181 | h_start, dfdh0, dWeights = neural_ode.backward(pred_y, dLoss) 182 | 183 | return loss, dWeights 184 | 185 | # Compile EAGER graph to static (this will be much faster) 186 | compute_gradients_and_update = tfe.defun(compute_gradients_and_update) 187 | 188 | 189 | # function to compute the kinetic energy 190 | def kinetic_energy(V, loggamma_v, loglambda_v): 191 | q = (np.sum(-V**2) - loggamma_v**2 - loglambda_v**2)/2.0 192 | return q 193 | 194 | def compute_gradient_param(dWeights, loggamma, loglambda, batch_size, para_num): 195 | WW = model.trainable_weights[0].numpy() 196 | dWeights = np.exp(loggamma)/2.0 * dWeights + np.exp(loglambda) * np.sign(WW) 197 | return dWeights 198 | 199 | def compute_gradient_hyper(loss, weights, loggamma, loglambda, batch_size, para_num): 200 | grad_loggamma = np.exp(loggamma) * (loss/2.0 + 1.0) - (batch_size/2.0 + 1.0) 201 | grad_loglambda = np.exp(loglambda) * (np.sum(np.abs(weights)) + 1.0) - (para_num + 1.0) 202 | 203 | return grad_loggamma, grad_loglambda 204 | 205 | def compute_Hamiltonian(loss, weights, loggamma, loglambda, batch_size, para_num): 206 | H = np.exp(loggamma)*(loss/2.0 + 1.0) + np.exp(loglambda)*(np.sum(np.abs(weights)) + 1.0)\ 207 | - (batch_size/2.0 + 1.0) * loggamma - (para_num + 1.0) * loglambda 208 | return H 209 | 210 | def leap_frog(v_in, w_in, loggamma_in, loglambda_in, loggamma_v_in, loglambda_v_in): 211 | # assign weights from the previous step to the model 212 | model.trainable_weights[0].assign(w_in) 213 | 214 | print(model.trainable_weights) 215 | 216 | v_new = v_in 217 | loggamma_v_new = loggamma_v_in 218 | loglambda_v_new = loglambda_v_in 219 | 220 | loggamma_new = loggamma_in 221 | loglambda_new = loglambda_in 222 | w_new = w_in 223 | 224 | for m in range(L): 225 | 226 | loss, dWeights = compute_gradients_and_update(batch_y0, batch_yN) # evaluate the gradient 227 | print(loss) 228 | 229 | dWeights = np.asarray(dWeights[0]) # make the gradient to be numpy array 230 | dWeights = compute_gradient_param(dWeights, loggamma_new, loglambda_new, batch_size, num_param) 231 | grad_loggamma, grad_loglambda = compute_gradient_hyper(loss, w_new, loggamma_new, loglambda_new, batch_size, num_param) 232 | 233 | loggamma_v_new = loggamma_v_new - epsilon/2*grad_loggamma 234 | loglambda_v_new = loglambda_v_new - epsilon/2*grad_loglambda 235 | v_new = v_new - epsilon/2*(dWeights) 236 | w_new = model.trainable_weights[0].numpy() + epsilon * v_new 237 | model.trainable_weights[0].assign(w_new) 238 | loggamma_new = loggamma_new + epsilon * loggamma_v_new 239 | loglambda_new = loglambda_new + epsilon * loglambda_v_new 240 | 241 | # Second half of the leap frog 242 | loss, dWeights = compute_gradients_and_update(batch_y0, batch_yN) 243 | dWeights = np.asarray(dWeights[0]) 244 | dWeights = compute_gradient_param(dWeights, loggamma_new, loglambda_new, batch_size, num_param) 245 | grad_loggamma, grad_loglambda = compute_gradient_hyper(loss, w_new, loggamma_new, loglambda_new, batch_size, num_param) 246 | 247 | v_new = v_new - epsilon/2*(dWeights) 248 | loggamma_v_new = loggamma_v_new - epsilon/2*grad_loggamma 249 | loglambda_v_new = loglambda_v_new - epsilon/2*grad_loglambda 250 | 251 | # print(dWeights) 252 | print(dWeights) 253 | print(np.exp(loggamma_new)) 254 | print(np.exp(loglambda_new)) 255 | 256 | return v_new, w_new, loggamma_new, loglambda_new, loggamma_v_new, loglambda_v_new 257 | 258 | 259 | 260 | 261 | 262 | neural_ode_test = NeuralODE(model, t=t_grid[0:data_size:20]) 263 | parameters = np.zeros((niters, num_param//2, 2)) # book keeping the parameters 264 | loggammalist = np.zeros((niters, 1)) # book keeping the loggamma 265 | loglambdalist = np.zeros((niters, 1)) # book keeping the loggamma 266 | loglikelihood = np.zeros((niters, 1)) # book keeping the loggamma 267 | L = 10 # leap frog step number 268 | epsilon = 0.001 # leap frog step size 269 | 270 | 271 | # initial weight 272 | w_temp = initial_weight 273 | loggamma_temp = 4. + np.random.normal() 274 | loglambda_temp = np.random.normal() 275 | 276 | # training steps 277 | for step in range(niters): 278 | print(step) 279 | v_initial = np.random.randn(num_param//2, 2) # initialize the velocity 280 | loggamma_v_initial = np.random.normal() 281 | loglambda_v_initial = np.random.normal() 282 | 283 | loss_initial, _ = compute_gradients_and_update(batch_y0, batch_yN) # compute the initial Hamiltonian 284 | loss_initial = compute_Hamiltonian(loss_initial, w_temp, loggamma_temp, loglambda_temp, batch_size, num_param) 285 | 286 | v_new, w_new, loggamma_new, loglambda_new, loggamma_v_new, loglambda_v_new = \ 287 | leap_frog(v_initial, w_temp, loggamma_temp, loglambda_temp, loggamma_v_initial, loglambda_v_initial) 288 | 289 | # compute the final Hamiltonian 290 | loss_finial, _ = compute_gradients_and_update(batch_y0, batch_yN) 291 | loss_finial = compute_Hamiltonian(loss_finial, w_new, loggamma_new, loglambda_new, batch_size, num_param) 292 | 293 | # making decisions 294 | p_temp = np.exp(-loss_finial + loss_initial + \ 295 | kinetic_energy(v_new, loggamma_v_new, loglambda_v_new) - kinetic_energy(v_initial, loggamma_v_initial, loglambda_v_initial)) 296 | 297 | p = min(1, p_temp) 298 | p_decision = np.random.uniform() 299 | if p > p_decision: 300 | parameters[step:step+1, :, :] = w_new 301 | w_temp = w_new 302 | loggammalist[step, 0] = loggamma_new 303 | loglambdalist[step, 0] = loglambda_new 304 | loglikelihood[step, 0] = loss_finial 305 | loggamma_temp = loggamma_new 306 | loglambda_temp = loglambda_new 307 | else: 308 | parameters[step:step+1, :, :] = w_temp 309 | model.trainable_weights[0].assign(w_temp) 310 | loggammalist[step, 0] = loggamma_temp 311 | loglambdalist[step, 0] = loglambda_temp 312 | loglikelihood[step, 0] = loss_initial 313 | 314 | print('probability', p) 315 | print(p > p_decision) 316 | 317 | 318 | 319 | np.save('parameters', parameters) 320 | np.save('loggammalist', loggammalist) 321 | np.save('loglikelihood', loglikelihood) 322 | 323 | 324 | 325 | 326 | -------------------------------------------------------------------------------- /Damped_Sine/neural_ode.py: -------------------------------------------------------------------------------- 1 | from typing import Optional, List 2 | import numpy as np 3 | import tensorflow as tf 4 | from tensorflow.python.framework.ops import EagerTensor 5 | import tensorflow.contrib.eager as tfe 6 | 7 | keras = tf.keras 8 | 9 | 10 | def zip_map(zipped, update_op): 11 | return [update_op(*elems) for elems in zipped] 12 | 13 | 14 | def euler_update(h_list, dh_list, dt): 15 | return zip_map(zip(h_list, dh_list), lambda h, dh: h + tf.cast(dt, h.dtype) * dh) 16 | 17 | 18 | def euler_step(func, dt, state): 19 | return euler_update(state, func(state), dt) 20 | 21 | 22 | def rk2_step(func, dt, state): 23 | k1 = func(state) 24 | k2 = func(euler_update(state, k1, dt)) 25 | return zip_map(zip(state, k1, k2), 26 | lambda h, dk1, dk2: h + tf.cast(dt, h.dtype) * (dk1 + dk2) / 2) 27 | 28 | 29 | def rk4_step(func, dt, state): 30 | k1 = func(state) 31 | k2 = func(euler_update(state, k1, dt / 2)) 32 | k3 = func(euler_update(state, k2, dt / 2)) 33 | k4 = func(euler_update(state, k3, dt)) 34 | 35 | return zip_map( 36 | zip(state, k1, k2, k3, k4), 37 | lambda h, dk1, dk2, dk3, dk4: h + tf.cast(dt, h.dtype) * ( 38 | dk1 + 2 * dk2 + 2 * dk3 + dk4) / 6, 39 | ) 40 | 41 | 42 | class NeuralODE: 43 | def __init__( 44 | self, model: tf.keras.Model, t=np.linspace(0, 1, 40), 45 | solver=rk4_step 46 | ): 47 | self._t = t 48 | self._model = model 49 | self._solver = solver 50 | self._deltas_t = t[1:] - t[:-1] 51 | 52 | def forward(self, inputs: tf.Tensor, return_states: Optional[str] = None): 53 | 54 | def _forward_dynamics(_state): 55 | """Used in solver _state == (time, tensor)""" 56 | return [1.0, self._model(inputs=_state)] 57 | 58 | states = [] 59 | 60 | def _append_state(_state): 61 | tensors = _state[1] 62 | if return_states == "numpy": 63 | states.append(tensors.numpy()) 64 | elif return_states == "tf": 65 | states.append(tensors) 66 | 67 | with tf.name_scope("forward"): 68 | t0 = tf.to_float(self._t[0]) 69 | state = [t0, inputs] 70 | _append_state(state) 71 | for dt in self._deltas_t: 72 | state = self._solver( 73 | func=_forward_dynamics, dt=tf.to_float(dt), state=state 74 | ) 75 | _append_state(state) 76 | 77 | outputs = state[1] 78 | if return_states: 79 | return outputs, states 80 | return outputs 81 | 82 | def _backward_dynamics(self, state): 83 | t = state[0] 84 | ht = state[1] 85 | at = -state[2] 86 | 87 | with tf.GradientTape() as g: 88 | g.watch(ht) 89 | ht_new = self._model(inputs=[t, ht]) 90 | 91 | gradients = g.gradient( 92 | target=ht_new, sources=[ht] + self._model.weights, 93 | output_gradients=at 94 | ) 95 | 96 | return [1.0, ht_new, *gradients] 97 | 98 | def backward(self, outputs: tf.Tensor, 99 | output_gradients: Optional[tf.Tensor] = None): 100 | 101 | with tf.name_scope("backward"): 102 | grad_weights = [tf.zeros_like(w) for w in self._model.weights] 103 | t0 = tf.to_float(self._t[-1]) 104 | 105 | if output_gradients is None: 106 | output_gradients = tf.zeros_like(outputs) 107 | 108 | state = [t0, outputs, output_gradients, *grad_weights] 109 | for dt in self._deltas_t[::-1]: 110 | state = self._solver( 111 | self._backward_dynamics, dt=-tf.to_float(dt), state=state 112 | ) 113 | 114 | inputs = state[1] 115 | dLdInputs = state[2] 116 | dLdWeights = state[3:] 117 | return inputs, dLdInputs, dLdWeights 118 | -------------------------------------------------------------------------------- /Damped_Sine/plotting.py: -------------------------------------------------------------------------------- 1 | import numpy as np 2 | import matplotlib as mpl 3 | #mpl.use('pgf') 4 | 5 | def figsize(scale, nplots = 1): 6 | fig_width_pt = 390.0 # Get this from LaTeX using \the\textwidth 7 | inches_per_pt = 1.0/72.27 # Convert pt to inch 8 | golden_mean = (np.sqrt(5.0)-1.0)/2.0 # Aesthetic ratio (you could change this) 9 | fig_width = fig_width_pt*inches_per_pt*scale # width in inches 10 | fig_height = nplots*fig_width*golden_mean # height in inches 11 | fig_size = [fig_width,fig_height] 12 | return fig_size 13 | 14 | pgf_with_latex = { # setup matplotlib to use latex for output 15 | "pgf.texsystem": "pdflatex", # change this if using xetex or lautex 16 | "text.usetex": True, # use LaTeX to write all text 17 | "font.family": "serif", 18 | "font.serif": [], # blank entries should cause plots to inherit fonts from the document 19 | "font.sans-serif": [], 20 | "font.monospace": [], 21 | "axes.labelsize": 10, # LaTeX default is 10pt font. 22 | "font.size": 10, 23 | "legend.fontsize": 8, # Make the legend/label fonts a little smaller 24 | "xtick.labelsize": 8, 25 | "ytick.labelsize": 8, 26 | "figure.figsize": figsize(1.0), # default fig size of 0.9 textwidth 27 | "pgf.preamble": [ 28 | r"\usepackage[utf8x]{inputenc}", # use utf8 fonts becasue your computer can handle it :) 29 | r"\usepackage[T1]{fontenc}", # plots will be generated using this preamble 30 | ] 31 | } 32 | mpl.rcParams.update(pgf_with_latex) 33 | 34 | import matplotlib.pyplot as plt 35 | 36 | # I make my own newfig and savefig functions 37 | def newfig(width, nplots = 1): 38 | fig = plt.figure(figsize=figsize(width, nplots)) 39 | ax = fig.add_subplot(111) 40 | return fig, ax 41 | 42 | def savefig(filename, crop = True): 43 | if crop == True: 44 | # plt.savefig('{}.pgf'.format(filename), bbox_inches='tight', pad_inches=0) 45 | plt.savefig('{}.pdf'.format(filename), bbox_inches='tight', pad_inches=0) 46 | plt.savefig('{}.eps'.format(filename), bbox_inches='tight', pad_inches=0) 47 | else: 48 | # plt.savefig('{}.pgf'.format(filename)) 49 | plt.savefig('{}.pdf'.format(filename)) 50 | plt.savefig('{}.eps'.format(filename)) 51 | 52 | ## Simple plot 53 | #fig, ax = newfig(1.0) 54 | # 55 | #def ema(y, a): 56 | # s = [] 57 | # s.append(y[0]) 58 | # for t in range(1, len(y)): 59 | # s.append(a * y[t] + (1-a) * s[t-1]) 60 | # return np.array(s) 61 | # 62 | #y = [0]*200 63 | #y.extend([20]*(1000-len(y))) 64 | #s = ema(y, 0.01) 65 | # 66 | #ax.plot(s) 67 | #ax.set_xlabel('X Label') 68 | #ax.set_ylabel('EMA') 69 | # 70 | #savefig('ema') -------------------------------------------------------------------------------- /Glycolysis_noisy/HMC.py: -------------------------------------------------------------------------------- 1 | import os 2 | os.environ['KMP_DUPLICATE_LIB_OK']='True' 3 | 4 | import numpy as np 5 | import numpy.random as npr 6 | import tensorflow as tf 7 | import matplotlib.pyplot as plt 8 | import matplotlib as mpl 9 | from mpl_toolkits.mplot3d import Axes3D 10 | plt.switch_backend('agg') 11 | import tensorflow.contrib.eager as tfe 12 | from tensorflow.keras import layers, initializers 13 | from scipy.integrate import odeint 14 | 15 | 16 | keras = tf.keras 17 | tf.enable_eager_execution() 18 | 19 | from neural_ode import NeuralODE 20 | 21 | np.random.seed(1234) 22 | tf.set_random_seed(1234) 23 | 24 | 25 | if __name__ == "__main__": 26 | 27 | def Glycolysis(z, t, J0, k1, k2, k3, k4, k5, k6, k, ka, q, KI, phi, N, A): 28 | x1, x2, x3, x4, x5, x6, x7 = z 29 | 30 | J = ka * (x4 - x7) 31 | N1 = N - x5 32 | A2 = A - x6 33 | 34 | v1 = k1 * x1 * x6 * (1 + (x6 / KI) ** q) ** (-1) 35 | v2 = k2 * x2 * N1 36 | v3 = k3 * x3 * A2 37 | v4 = k4 * x4 * x5 38 | v5 = k5 * x6 39 | v6 = k6 * x2 * x5 40 | v7 = k * x7 41 | 42 | r1 = J0 - v1 43 | r2 = 2*v1 - v2 - v6 44 | r3 = v2 - v3 45 | r4 = v3 - v4 - J 46 | r5 = v2 - v4 - v6 47 | r6 = - 2 * v1 + 2 * v3 - v5 48 | r7 = phi * J - v7 49 | dzdt = [r1, r2, r3, r4, r5, r6, r7] 50 | return dzdt 51 | 52 | 53 | def plot_spiral(trajectories, width = 1.): 54 | x1 = trajectories[:,0] 55 | x2 = trajectories[:,1] 56 | x3 = trajectories[:,2] 57 | x4 = trajectories[:,3] 58 | x5 = trajectories[:,4] 59 | x6 = trajectories[:,5] 60 | x7 = trajectories[:,6] 61 | plt.plot(x1, linewidth = width) 62 | plt.plot(x2, linewidth = width) 63 | plt.plot(x3, linewidth = width) 64 | plt.plot(x4, linewidth = width) 65 | plt.plot(x5, linewidth = width) 66 | plt.plot(x6, linewidth = width) 67 | plt.plot(x7, linewidth = width) 68 | 69 | 70 | data_size = 4000 71 | batch_time = 60 72 | niters = 5000 73 | batch_size = 1000 74 | 75 | J0 = 2.5 76 | k1 = 100. 77 | k2 = 6. 78 | k3 = 16. 79 | k4 = 100. 80 | k5 = 1.28 81 | k6 = 12. 82 | k = 1.8 83 | ka = 13. 84 | q = 4. 85 | KI = 0.52 86 | phi = 0.1 87 | N = 1. 88 | A = 4. 89 | 90 | 91 | t_grid = np.linspace(0, 5, data_size) 92 | z0 = [0.5, 1.9, 0.18, 0.15, 0.16, 0.1, 0.064] 93 | true_yy = odeint(Glycolysis, z0, t_grid, args=(J0, k1, k2, k3, k4, k5, k6, k, ka, q, KI, phi, N, A)) 94 | 95 | true_y = np.load("true_y.npy") 96 | batch_y0 = np.load("batch_y0.npy") 97 | batch_yN = np.load("batch_yN.npy") 98 | 99 | batch_y0 = tf.to_float(batch_y0) 100 | batch_yN = tf.to_float(batch_yN) 101 | 102 | # normalizing data and system 103 | sigma_normal1 = np.std(true_y[:,0:1]) 104 | sigma_normal2 = np.std(true_y[:,1:2]) 105 | sigma_normal3 = np.std(true_y[:,2:3]) 106 | sigma_normal4 = np.std(true_y[:,3:4]) 107 | sigma_normal5 = np.std(true_y[:,4:5]) 108 | sigma_normal6 = np.std(true_y[:,5:6]) 109 | sigma_normal7 = np.std(true_y[:,6:7]) 110 | 111 | 112 | num_param = 14 113 | para_num = num_param 114 | 115 | t0 = t_grid[:batch_time][0] 116 | t1 = t_grid[:batch_time][-1] 117 | t_in = np.linspace(t0, t1, 10) 118 | 119 | parameters_pre = np.load("parameters_pre.npy") 120 | print(np.exp(parameters_pre)) 121 | 122 | ######################################### 123 | ########## precondition end ############# 124 | ######################################### 125 | 126 | initial_weight = parameters_pre 127 | print(initial_weight.shape, "here") 128 | 129 | class ODEModel(tf.keras.Model): 130 | def __init__(self): 131 | super(ODEModel, self).__init__() 132 | self.Weights = tfe.Variable(tf.random_normal([num_param, 1], dtype=tf.float32)*0.01, dtype=tf.float32) 133 | 134 | def call(self, inputs, **kwargs): 135 | t, y = inputs 136 | h = y 137 | x1 = h[:,0:1] * sigma_normal1 138 | x2 = h[:,1:2] * sigma_normal2 139 | x3 = h[:,2:3] * sigma_normal3 140 | x4 = h[:,3:4] * sigma_normal4 141 | x5 = h[:,4:5] * sigma_normal5 142 | x6 = h[:,5:6] * sigma_normal6 143 | x7 = h[:,6:7] * sigma_normal7 144 | 145 | J0 = self.Weights[0] 146 | k1 = self.Weights[1] 147 | k2 = self.Weights[2] 148 | k3 = self.Weights[3] 149 | k4 = self.Weights[4] 150 | k5 = self.Weights[5] 151 | k6 = self.Weights[6] 152 | k = self.Weights[7] 153 | ka = self.Weights[8] 154 | q = self.Weights[9] 155 | KI = self.Weights[10] 156 | phi = self.Weights[11] 157 | N = self.Weights[12] 158 | A = self.Weights[13] 159 | 160 | ka = tf.exp(ka) 161 | q = tf.exp(q) 162 | N = tf.exp(N) 163 | A = tf.exp(A) 164 | k1 = tf.exp(k1) 165 | k2 = tf.exp(k2) 166 | k3 = tf.exp(k3) 167 | k4 = tf.exp(k4) 168 | k5 = tf.exp(k5) 169 | k6 = tf.exp(k6) 170 | k = tf.exp(k) 171 | KI = tf.exp(KI) 172 | J0 = tf.exp(J0) 173 | phi = tf.exp(phi) 174 | 175 | J = ka * (x4 - x7) 176 | N1 = N - x5 177 | A2 = A - x6 178 | 179 | v1 = k1 * x1 * x6 * (1 + (x6 / KI) ** q) ** (-1) 180 | v2 = k2 * x2 * N1 181 | v3 = k3 * x3 * A2 182 | v4 = k4 * x4 * x5 183 | v5 = k5 * x6 184 | v6 = k6 * x2 * x5 185 | v7 = k * x7 186 | 187 | r1 = J0 - v1 188 | r2 = 2*v1 - v2 - v6 189 | r3 = v2 - v3 190 | r4 = v3 - v4 - J 191 | r5 = v2 - v4 - v6 192 | r6 = - 2 * v1 + 2 * v3 - v5 193 | r7 = phi * J - v7 194 | 195 | h_out = tf.concat([r1/sigma_normal1, r2/sigma_normal2, r3/sigma_normal3, r4/sigma_normal4, r5/sigma_normal5, r6/sigma_normal6, r7/sigma_normal7], 1) 196 | return h_out 197 | 198 | 199 | model = ODEModel() 200 | neural_ode = NeuralODE(model, t = t_in) 201 | 202 | def compute_gradients_and_update(batch_y0, batch_yN): 203 | """Takes start positions (x0, y0) and final positions (xN, yN)""" 204 | pred_y = neural_ode.forward(batch_y0) 205 | with tf.GradientTape() as g: 206 | g.watch(pred_y) 207 | loss = tf.reduce_sum((pred_y - batch_yN)**2) 208 | 209 | dLoss = g.gradient(loss, pred_y) 210 | h_start, dfdh0, dWeights = neural_ode.backward(pred_y, dLoss) 211 | 212 | return loss, dWeights 213 | 214 | # Compile EAGER graph to static (this will be much faster) 215 | compute_gradients_and_update = tfe.defun(compute_gradients_and_update) 216 | 217 | 218 | # function to compute the kinetic energy 219 | def kinetic_energy(V, loggamma_v, loglambda_v): 220 | q = (np.sum(-V**2) - loggamma_v**2 - loglambda_v**2)/2.0 221 | return q 222 | 223 | def compute_gradient_param(dWeights, loggamma, loglambda, batch_size, para_num): 224 | WW = model.trainable_weights[0].numpy() 225 | dWeights = np.exp(loggamma)/2.0 * dWeights + np.exp(loglambda) * np.sign(WW) 226 | return dWeights 227 | 228 | def compute_gradient_hyper(loss, weights, loggamma, loglambda, batch_size, para_num): 229 | grad_loggamma = np.exp(loggamma) * (loss/2.0 + 1.0) - (batch_size/2.0 + 1.0) 230 | grad_loglambda = np.exp(loglambda) * (np.sum(np.abs(weights)) + 1.0) - (para_num + 1.0) 231 | 232 | return grad_loggamma, grad_loglambda 233 | 234 | def compute_Hamiltonian(loss, weights, loggamma, loglambda, batch_size, para_num): 235 | H = np.exp(loggamma)*(loss/2.0 + 1.0) + np.exp(loglambda)*(np.sum(np.abs(weights)) + 1.0)\ 236 | - (batch_size/2.0 + 1.0) * loggamma - (para_num + 1.0) * loglambda 237 | return H 238 | 239 | def leap_frog(v_in, w_in, loggamma_in, loglambda_in, loggamma_v_in, loglambda_v_in): 240 | # assign weights from the previous step to the model 241 | model.trainable_weights[0].assign(w_in) 242 | v_new = v_in 243 | loggamma_v_new = loggamma_v_in 244 | loglambda_v_new = loglambda_v_in 245 | 246 | loggamma_new = loggamma_in 247 | loglambda_new = loglambda_in 248 | w_new = w_in 249 | 250 | for m in range(L): 251 | loss, dWeights = compute_gradients_and_update(batch_y0, batch_yN) # evaluate the gradient 252 | 253 | dWeights = np.asarray(dWeights[0]) # make the gradient to be numpy array 254 | dWeights = compute_gradient_param(dWeights, loggamma_new, loglambda_new, batch_size, para_num) 255 | grad_loggamma, grad_loglambda = compute_gradient_hyper(loss, w_new, loggamma_new, loglambda_new, batch_size, para_num) 256 | 257 | loggamma_v_new = loggamma_v_new - epsilon/2*grad_loggamma 258 | loglambda_v_new = loglambda_v_new - epsilon/2*grad_loglambda 259 | v_new = v_new - epsilon/2*(dWeights) 260 | w_new = model.trainable_weights[0].numpy() + epsilon * v_new 261 | model.trainable_weights[0].assign(w_new) 262 | loggamma_new = loggamma_new + epsilon * loggamma_v_new 263 | loglambda_new = loglambda_new + epsilon * loglambda_v_new 264 | 265 | # Second half of the leap frog 266 | loss, dWeights = compute_gradients_and_update(batch_y0, batch_yN) 267 | dWeights = np.asarray(dWeights[0]) 268 | dWeights = compute_gradient_param(dWeights, loggamma_new, loglambda_new, batch_size, para_num) 269 | grad_loggamma, grad_loglambda = compute_gradient_hyper(loss, w_new, loggamma_new, loglambda_new, batch_size, para_num) 270 | 271 | v_new = v_new - epsilon/2*(dWeights) 272 | loggamma_v_new = loggamma_v_new - epsilon/2*grad_loggamma 273 | loglambda_v_new = loglambda_v_new - epsilon/2*grad_loglambda 274 | 275 | print(np.exp(loggamma_new)) 276 | print(np.exp(loglambda_new)) 277 | 278 | return v_new, w_new, loggamma_new, loglambda_new, loggamma_v_new, loglambda_v_new 279 | 280 | 281 | 282 | neural_ode_test = NeuralODE(model, t=t_grid[0:data_size:20]) 283 | parameters = np.zeros((niters, para_num)) # book keeping the parameters 284 | loggammalist = np.zeros((niters, 1)) # book keeping the loggamma 285 | loglambdalist = np.zeros((niters, 1)) # book keeping the loggamma 286 | loglikelihood = np.zeros((niters, 1)) # book keeping the loggamma 287 | L = 10 # leap frog step number 288 | epsilon = 0.001 # leap frog step size 289 | epsilon_max = 0.0002 # max 0.001 290 | epsilon_min = 0.0002 # max 0.001 291 | 292 | 293 | def compute_epsilon(step): 294 | coefficient = np.log(epsilon_max/epsilon_min) 295 | return epsilon_max * np.exp( - step * coefficient / niters) 296 | 297 | 298 | # initial weight 299 | w_temp = initial_weight 300 | print("initial_w", np.exp(w_temp)) 301 | loggamma_temp = 4. + np.random.normal() 302 | loglambda_temp = np.random.normal() 303 | 304 | model.trainable_weights[0].assign(w_temp) 305 | loss_original, _ = compute_gradients_and_update(batch_y0, batch_yN) # compute the initial Hamiltonian 306 | 307 | loggamma_temp = np.log(batch_size / loss_original) 308 | print("This is initial guess", loggamma_temp, "with loss", loss_original) 309 | if loggamma_temp > 6.: 310 | loggamma_temp = 6. 311 | epsilon_max = 0.0001 312 | epsilon_min = 0.0001 313 | 314 | 315 | # training steps 316 | for step in range(niters): 317 | 318 | epsilon = compute_epsilon(step) 319 | 320 | print(step) 321 | v_initial = np.random.randn(para_num, 1) # initialize the velocity 322 | loggamma_v_initial = np.random.normal() 323 | loglambda_v_initial = np.random.normal() 324 | 325 | loss_initial, _ = compute_gradients_and_update(batch_y0, batch_yN) # compute the initial Hamiltonian 326 | loss_initial = compute_Hamiltonian(loss_initial, w_temp, loggamma_temp, loglambda_temp, batch_size, para_num) 327 | 328 | v_new, w_new, loggamma_new, loglambda_new, loggamma_v_new, loglambda_v_new = \ 329 | leap_frog(v_initial, w_temp, loggamma_temp, loglambda_temp, loggamma_v_initial, loglambda_v_initial) 330 | 331 | # compute the final Hamiltonian 332 | loss_finial, _ = compute_gradients_and_update(batch_y0, batch_yN) 333 | loss_finial = compute_Hamiltonian(loss_finial, w_new, loggamma_new, loglambda_new, batch_size, para_num) 334 | 335 | # making decisions 336 | p_temp = np.exp(-loss_finial + loss_initial + \ 337 | kinetic_energy(v_new, loggamma_v_new, loglambda_v_new) - kinetic_energy(v_initial, loggamma_v_initial, loglambda_v_initial)) 338 | 339 | p = min(1, p_temp) 340 | p_decision = np.random.uniform() 341 | if p > p_decision: 342 | parameters[step:step+1, :] = np.transpose(w_new) 343 | w_temp = w_new 344 | loggammalist[step, 0] = loggamma_new 345 | loglambdalist[step, 0] = loglambda_new 346 | loglikelihood[step, 0] = loss_finial 347 | loggamma_temp = loggamma_new 348 | loglambda_temp = loglambda_new 349 | else: 350 | parameters[step:step+1, :] = np.transpose(w_temp) 351 | model.trainable_weights[0].assign(w_temp) 352 | loggammalist[step, 0] = loggamma_temp 353 | loglambdalist[step, 0] = loglambda_temp 354 | loglikelihood[step, 0] = loss_initial 355 | 356 | print('probability', p) 357 | print(p > p_decision) 358 | 359 | 360 | np.save('parameters', parameters) 361 | np.save('loggammalist', loggammalist) 362 | np.save('loglikelihood', loglikelihood) 363 | -------------------------------------------------------------------------------- /Glycolysis_noisy/create_nice_plot.py: -------------------------------------------------------------------------------- 1 | import os 2 | os.environ['KMP_DUPLICATE_LIB_OK']='True' 3 | 4 | import numpy as np 5 | import numpy.random as npr 6 | import tensorflow as tf 7 | import matplotlib.pyplot as plt 8 | plt.switch_backend('agg') 9 | 10 | from scipy.integrate import odeint 11 | from scipy import stats 12 | 13 | from plotting import newfig, savefig 14 | 15 | np.random.seed(1234) 16 | tf.set_random_seed(1234) 17 | 18 | 19 | if __name__ == "__main__": 20 | 21 | def Glycolysis(z, t, J0, k1, k2, k3, k4, k5, k6, k, ka, q, KI, phi, N, A): 22 | x1, x2, x3, x4, x5, x6, x7 = z 23 | 24 | J = ka * (x4 - x7) 25 | N1 = N - x5 26 | A2 = A - x6 27 | 28 | v1 = k1 * x1 * x6 * (1 + (x6 / KI) ** q) ** (-1) 29 | v2 = k2 * x2 * N1 30 | v3 = k3 * x3 * A2 31 | v4 = k4 * x4 * x5 32 | v5 = k5 * x6 33 | v6 = k6 * x2 * x5 34 | v7 = k * x7 35 | 36 | r1 = J0 - v1 37 | r2 = 2*v1 - v2 - v6 38 | r3 = v2 - v3 39 | r4 = v3 - v4 - J 40 | r5 = v2 - v4 - v6 41 | r6 = - 2 * v1 + 2 * v3 - v5 42 | r7 = phi * J - v7 43 | dzdt = [r1, r2, r3, r4, r5, r6, r7] 44 | return dzdt 45 | 46 | 47 | def plot_spiral(trajectories, width = 1.): 48 | x1 = trajectories[:,0] 49 | x2 = trajectories[:,1] 50 | x3 = trajectories[:,2] 51 | x4 = trajectories[:,3] 52 | x5 = trajectories[:,4] 53 | x6 = trajectories[:,5] 54 | x7 = trajectories[:,6] 55 | plt.plot(x1, linewidth = width) 56 | plt.plot(x2, linewidth = width) 57 | plt.plot(x3, linewidth = width) 58 | plt.plot(x4, linewidth = width) 59 | plt.plot(x5, linewidth = width) 60 | plt.plot(x6, linewidth = width) 61 | plt.plot(x7, linewidth = width) 62 | 63 | Order = 3 64 | def Glycolysis_BNODE(z, t, W, Order): 65 | x1, x2, x3, x4, x5, x6, x7 = z 66 | 67 | J0 = W[0] 68 | k1 = W[1] 69 | k2 = W[2] 70 | k3 = W[3] 71 | k4 = W[4] 72 | k5 = W[5] 73 | k6 = W[6] 74 | k = W[7] 75 | ka = W[8] 76 | q = W[9] 77 | KI = W[10] 78 | phi = W[11] 79 | N = W[12] 80 | A = W[13] 81 | 82 | J = ka * (x4 - x7) 83 | N1 = N - x5 84 | A2 = A - x6 85 | 86 | v1 = k1 * x1 * x6 * (1 + (x6 / KI) ** q) ** (-1) 87 | v2 = k2 * x2 * N1 88 | v3 = k3 * x3 * A2 89 | v4 = k4 * x4 * x5 90 | v5 = k5 * x6 91 | v6 = k6 * x2 * x5 92 | v7 = k * x7 93 | 94 | r1 = J0 - v1 95 | r2 = 2*v1 - v2 - v6 96 | r3 = v2 - v3 97 | r4 = v3 - v4 - J 98 | r5 = v2 - v4 - v6 99 | r6 = - 2 * v1 + 2 * v3 - v5 100 | r7 = phi * J - v7 101 | dzdt = [r1, r2, r3, r4, r5, r6, r7] 102 | return dzdt 103 | 104 | 105 | 106 | data_size = 4000 107 | batch_time = 60 108 | batch_size = 1000 109 | 110 | data_size_true = 4000 111 | 112 | 113 | N_samples = 100 114 | N_total = 1500 115 | 116 | 117 | J0 = 2.5 118 | k1 = 100. 119 | k2 = 6. 120 | k3 = 16. 121 | k4 = 100. 122 | k5 = 1.28 123 | k6 = 12. 124 | k = 1.8 125 | ka = 13. 126 | q = 4. 127 | KI = 0.52 128 | phi = 0.1 129 | N = 1. 130 | A = 4. 131 | 132 | z0 = [0.5, 1.9, 0.18, 0.15, 0.16, 0.1, 0.064] 133 | t_grid_train = np.linspace(0, 5, data_size) 134 | t_grid_true = np.linspace(0, 10, data_size_true) 135 | 136 | y_train = odeint(Glycolysis, z0, t_grid_train, args=(J0, k1, k2, k3, k4, k5, k6, k, ka, q, KI, phi, N, A)) 137 | idx = np.random.choice(np.arange(data_size - batch_time - 1, dtype=np.int64), batch_size, replace=False) 138 | y_train = y_train[idx] 139 | t_grid_train = t_grid_train[idx] 140 | 141 | sigma_normal1 = np.std(y_train[:,0:1]) 142 | sigma_normal2 = np.std(y_train[:,1:2]) 143 | sigma_normal3 = np.std(y_train[:,2:3]) 144 | sigma_normal4 = np.std(y_train[:,3:4]) 145 | sigma_normal5 = np.std(y_train[:,4:5]) 146 | sigma_normal6 = np.std(y_train[:,5:6]) 147 | sigma_normal7 = np.std(y_train[:,6:7]) 148 | 149 | noise_level = 0.02 150 | 151 | y_train[:,0:1] = y_train[:,0:1]/ sigma_normal1 + noise_level * np.random.randn(y_train[:,0:1].shape[0], y_train[:,0:1].shape[1]) 152 | y_train[:,1:2] = y_train[:,1:2]/ sigma_normal2 + noise_level * np.random.randn(y_train[:,1:2].shape[0], y_train[:,1:2].shape[1]) 153 | y_train[:,2:3] = y_train[:,2:3]/ sigma_normal3 + noise_level * np.random.randn(y_train[:,2:3].shape[0], y_train[:,2:3].shape[1]) 154 | y_train[:,3:4] = y_train[:,3:4]/ sigma_normal4 + noise_level * np.random.randn(y_train[:,3:4].shape[0], y_train[:,3:4].shape[1]) 155 | y_train[:,4:5] = y_train[:,4:5]/ sigma_normal5 + noise_level * np.random.randn(y_train[:,4:5].shape[0], y_train[:,4:5].shape[1]) 156 | y_train[:,5:6] = y_train[:,5:6]/ sigma_normal6 + noise_level * np.random.randn(y_train[:,5:6].shape[0], y_train[:,5:6].shape[1]) 157 | y_train[:,6:7] = y_train[:,6:7]/ sigma_normal7 + noise_level * np.random.randn(y_train[:,6:7].shape[0], y_train[:,6:7].shape[1]) 158 | 159 | y_train[:,0:1] = y_train[:,0:1] * sigma_normal1 160 | y_train[:,1:2] = y_train[:,1:2] * sigma_normal2 161 | y_train[:,2:3] = y_train[:,2:3] * sigma_normal3 162 | y_train[:,3:4] = y_train[:,3:4] * sigma_normal4 163 | y_train[:,4:5] = y_train[:,4:5] * sigma_normal5 164 | y_train[:,5:6] = y_train[:,5:6] * sigma_normal6 165 | y_train[:,6:7] = y_train[:,6:7] * sigma_normal7 166 | 167 | 168 | sigma_normal = np.asarray([sigma_normal1, sigma_normal2, sigma_normal3, sigma_normal4, sigma_normal5, sigma_normal6, sigma_normal7]) 169 | 170 | y_true = odeint(Glycolysis, z0, t_grid_true, args=(J0, k1, k2, k3, k4, k5, k6, k, ka, q, KI, phi, N, A)) 171 | 172 | num_dim = 7 173 | parameters = np.load("parameters.npy") 174 | precision = np.load("loggammalist.npy") 175 | loglikelihood = np.load("loglikelihood.npy") 176 | precision = np.exp(precision) * num_dim 177 | print("precision", precision) 178 | print(parameters.shape) 179 | parameters = np.exp(parameters) 180 | num_samples = parameters.shape[0] 181 | length_dict = parameters.shape[1] 182 | 183 | loglikelihood = loglikelihood[-N_total:] 184 | idx_MAP = np.argmin(loglikelihood) 185 | print(idx_MAP, "index") 186 | print(loglikelihood) 187 | MAP = parameters[idx_MAP, :] 188 | print(MAP) 189 | 190 | 191 | y_MAP = odeint(Glycolysis_BNODE, z0, t_grid_true, args=(MAP, Order)) 192 | 193 | mu_pred = y_MAP 194 | 195 | y_BNODE = np.zeros((t_grid_true.shape[0], num_dim, N_samples)) 196 | 197 | for k in range(N_samples): 198 | print(k) 199 | idx_1 = np.random.randint(N_total) 200 | idx_2 = np.random.randint(N_total) 201 | W_sample = parameters[-idx_1, :] 202 | precision_here = precision[-idx_2] 203 | y_BNODE[:,:,k] = odeint(Glycolysis_BNODE, z0, t_grid_true, args=(W_sample, Order)) 204 | Sigma_data = np.ones_like(mu_pred) / np.sqrt(precision_here) 205 | y_BNODE[:,:,k] = y_BNODE[:,:,k] + Sigma_data * sigma_normal * np.random.normal() 206 | print((Sigma_data * sigma_normal).shape) 207 | 208 | 209 | 210 | mu_pred = np.mean(y_BNODE, axis = 2) 211 | Sigma_pred = np.var(y_BNODE, axis = 2) 212 | 213 | 214 | plt.rcParams.update(plt.rcParamsDefault) 215 | plt.rc('font', family='serif') 216 | plt.rcParams.update({'font.size': 16, 217 | 'lines.linewidth': 2, 218 | 'axes.labelsize': 20, # fontsize for x and y labels (was 10) 219 | 'axes.titlesize': 20, 220 | 'xtick.labelsize': 16, 221 | 'ytick.labelsize': 16, 222 | 'legend.fontsize': 20, 223 | 'axes.linewidth': 2, 224 | "pgf.texsystem": "pdflatex", # change this if using xetex or lautex 225 | "text.usetex": True, # use LaTeX to write all text 226 | }) 227 | 228 | 229 | 230 | plt.figure(4, figsize=(12,6.5)) 231 | plt.xticks(fontsize=22) 232 | plt.yticks(fontsize=22) 233 | plt.plot(t_grid_true, y_BNODE[:,0,0], '-', color = "gray", alpha = 0.8, label = "Sample Trajectory of $S_1(t)$", linewidth = 0.5) 234 | plt.plot(t_grid_true, y_BNODE[:,0,:], '-', color = "gray", alpha = 0.8, linewidth = 0.5) 235 | plt.plot(t_grid_train, y_train[:,0], 'ro', label = "Training data of $S_1(t)$") 236 | plt.plot(t_grid_true, y_true[:,0], 'r-', label = "True Trajectory of $S_1(t)$") 237 | plt.plot(t_grid_true, y_MAP[:,0], 'b--', label = "MAP Trajectory of $S_1(t)$") 238 | plt.xlabel('$t$',fontsize=26) 239 | plt.ylabel('$S_1(t)$',fontsize=26) 240 | plt.ylim((-0.1, 2.8)) 241 | plt.legend(loc='upper right', frameon=False, prop={'size': 20}) 242 | plt.savefig('./BNODE_Prediction_x1.png', dpi = 300) 243 | 244 | 245 | plt.figure(5, figsize=(12,6.5)) 246 | plt.xticks(fontsize=22) 247 | plt.yticks(fontsize=22) 248 | plt.plot(t_grid_true, y_BNODE[:,1,0], '-', color = "gray", alpha = 0.8, label = "Sample Trajectory of $S_2(t)$", linewidth = 0.5) 249 | plt.plot(t_grid_true, y_BNODE[:,1,:], '-', color = "gray", alpha = 0.8, linewidth = 0.5) 250 | plt.plot(t_grid_train, y_train[:,1], 'ro', label = "Training data of $S_2(t)$") 251 | plt.plot(t_grid_true, y_true[:,1], 'r-', label = "True Trajectory of $S_2(t)$") 252 | plt.plot(t_grid_true, y_MAP[:,1], 'b--', label = "MAP Trajectory of $S_2(t)$") 253 | plt.xlabel('$t$',fontsize=26) 254 | plt.ylabel('$S_2(t)$',fontsize=26) 255 | plt.ylim((-0.1, 3.5)) 256 | plt.legend(loc='upper right', frameon=False, prop={'size': 20}) 257 | plt.savefig('./BNODE_Prediction_x2.png', dpi = 300) 258 | 259 | 260 | plt.figure(6, figsize=(12,6.5)) 261 | plt.xticks(fontsize=22) 262 | plt.yticks(fontsize=22) 263 | plt.plot(t_grid_true, y_BNODE[:,2,0], '-', color = "gray", alpha = 0.8, label = "Sample Trajectory of $S_3(t)$", linewidth = 0.5) 264 | plt.plot(t_grid_true, y_BNODE[:,2,:], '-', color = "gray", alpha = 0.8, linewidth = 0.5) 265 | plt.plot(t_grid_train, y_train[:,2], 'ro', label = "Training data of $S_3(t)$") 266 | plt.plot(t_grid_true, y_true[:,2], 'r-', label = "True Trajectory of $S_3(t)$") 267 | plt.plot(t_grid_true, y_MAP[:,2], 'b--', label = "MAP Trajectory of $S_3(t)$") 268 | plt.xlabel('$t$',fontsize=26) 269 | plt.ylabel('$S_3(t)$',fontsize=26) 270 | plt.ylim((0., 0.3)) 271 | plt.legend(loc='upper right', frameon=False, prop={'size': 20}) 272 | plt.savefig('./BNODE_Prediction_x3.png', dpi = 300) 273 | 274 | 275 | plt.figure(7, figsize=(12,6.5)) 276 | plt.xticks(fontsize=22) 277 | plt.yticks(fontsize=22) 278 | plt.plot(t_grid_true, y_BNODE[:,3,0], '-', color = "gray", alpha = 0.8, label = "Sample Trajectory of $S_4(t)$", linewidth = 0.5) 279 | plt.plot(t_grid_true, y_BNODE[:,3,:], '-', color = "gray", alpha = 0.8, linewidth = 0.5) 280 | plt.plot(t_grid_train, y_train[:,3], 'ro', label = "Training data of $S_4(t)$") 281 | plt.plot(t_grid_true, y_true[:,3], 'r-', label = "True Trajectory of $S_4(t)$") 282 | plt.plot(t_grid_true, y_MAP[:,3], 'b--', label = "MAP Trajectory of $S_4(t)$") 283 | plt.xlabel('$t$',fontsize=26) 284 | plt.ylabel('$S_4(t)$',fontsize=26) 285 | plt.ylim((0.0, 0.6)) 286 | plt.legend(loc='upper right', frameon=False, prop={'size': 20}) 287 | plt.savefig('./BNODE_Prediction_x4.png', dpi = 300) 288 | 289 | 290 | plt.figure(8, figsize=(12,6.5)) 291 | plt.xticks(fontsize=22) 292 | plt.yticks(fontsize=22) 293 | plt.plot(t_grid_true, y_BNODE[:,4,0], '-', color = "gray", alpha = 0.8, label = "Sample Trajectory of $N_2(t)$", linewidth = 0.5) 294 | plt.plot(t_grid_true, y_BNODE[:,4,:], '-', color = "gray", alpha = 0.8, linewidth = 0.5) 295 | plt.plot(t_grid_train, y_train[:,4], 'ro', label = "Training data of $N_2(t)$") 296 | plt.plot(t_grid_true, y_true[:,4], 'r-', label = "True Trajectory of $N_2(t)$") 297 | plt.plot(t_grid_true, y_MAP[:,4], 'b--', label = "MAP Trajectory of $N_2(t)$") 298 | plt.xlabel('$t$',fontsize=26) 299 | plt.ylabel('$N_2(t)$',fontsize=26) 300 | plt.ylim((0.05, 0.3)) 301 | plt.legend(loc='upper right', frameon=False, prop={'size': 20}) 302 | plt.savefig('./BNODE_Prediction_x5.png', dpi = 300) 303 | 304 | 305 | plt.figure(9, figsize=(12,6.5)) 306 | plt.xticks(fontsize=22) 307 | plt.yticks(fontsize=22) 308 | plt.plot(t_grid_true, y_BNODE[:,5,0], '-', color = "gray", alpha = 0.8, label = "Sample Trajectory of $A_3(t)$", linewidth = 0.5) 309 | plt.plot(t_grid_true, y_BNODE[:,5,:], '-', color = "gray", alpha = 0.8, linewidth = 0.5) 310 | plt.plot(t_grid_train, y_train[:,5], 'ro', label = "Training data of $A_3(t)$") 311 | plt.plot(t_grid_true, y_true[:,5], 'r-', label = "True Trajectory of $A_3(t)$") 312 | plt.plot(t_grid_true, y_MAP[:,5], 'b--', label = "MAP Trajectory of $A_3(t)$") 313 | plt.xlabel('$t$',fontsize=26) 314 | plt.ylabel('$A_3(t)$',fontsize=26) 315 | plt.ylim((-0.1, 4.5)) 316 | plt.legend(loc='upper right', frameon=False, prop={'size': 20}) 317 | plt.savefig('./BNODE_Prediction_x6.png', dpi = 300) 318 | 319 | 320 | plt.figure(10, figsize=(12,6.5)) 321 | plt.xticks(fontsize=22) 322 | plt.yticks(fontsize=22) 323 | plt.plot(t_grid_true, y_BNODE[:,6,0], '-', color = "gray", alpha = 0.8, label = "Sample Trajectory of $S_4^{ex}(t)$", linewidth = 0.5) 324 | plt.plot(t_grid_true, y_BNODE[:,6,:], '-', color = "gray", alpha = 0.8, linewidth = 0.5) 325 | plt.plot(t_grid_train, y_train[:,6], 'ro', label = "Training data of $S_4^{ex}(t)$") 326 | plt.plot(t_grid_true, y_true[:,6], 'r-', label = "True Trajectory of $S_4^{ex}(t)$") 327 | plt.plot(t_grid_true, y_MAP[:,6], 'b--', label = "MAP Trajectory of $S_4^{ex}(t)$") 328 | plt.xlabel('$t$',fontsize=26) 329 | plt.ylabel('$S_4^{ex}(t)$',fontsize=26) 330 | plt.ylim((0.04, 0.14)) 331 | plt.legend(loc='upper right', frameon=False, prop={'size': 20}) 332 | plt.savefig('./BNODE_Prediction_x7.png', dpi = 300) 333 | 334 | 335 | 336 | 337 | 338 | 339 | -------------------------------------------------------------------------------- /Glycolysis_noisy/neural_ode.py: -------------------------------------------------------------------------------- 1 | from typing import Optional, List 2 | import numpy as np 3 | import tensorflow as tf 4 | from tensorflow.python.framework.ops import EagerTensor 5 | import tensorflow.contrib.eager as tfe 6 | 7 | keras = tf.keras 8 | 9 | 10 | def zip_map(zipped, update_op): 11 | return [update_op(*elems) for elems in zipped] 12 | 13 | 14 | def euler_update(h_list, dh_list, dt): 15 | return zip_map(zip(h_list, dh_list), lambda h, dh: h + tf.cast(dt, h.dtype) * dh) 16 | 17 | 18 | def euler_step(func, dt, state): 19 | return euler_update(state, func(state), dt) 20 | 21 | 22 | def rk2_step(func, dt, state): 23 | k1 = func(state) 24 | k2 = func(euler_update(state, k1, dt)) 25 | return zip_map(zip(state, k1, k2), 26 | lambda h, dk1, dk2: h + tf.cast(dt, h.dtype) * (dk1 + dk2) / 2) 27 | 28 | 29 | def rk4_step(func, dt, state): 30 | k1 = func(state) 31 | k2 = func(euler_update(state, k1, dt / 2)) 32 | k3 = func(euler_update(state, k2, dt / 2)) 33 | k4 = func(euler_update(state, k3, dt)) 34 | 35 | return zip_map( 36 | zip(state, k1, k2, k3, k4), 37 | lambda h, dk1, dk2, dk3, dk4: h + tf.cast(dt, h.dtype) * ( 38 | dk1 + 2 * dk2 + 2 * dk3 + dk4) / 6, 39 | ) 40 | 41 | 42 | class NeuralODE: 43 | def __init__( 44 | self, model: tf.keras.Model, t=np.linspace(0, 1, 40), 45 | solver=rk4_step 46 | ): 47 | self._t = t 48 | self._model = model 49 | self._solver = solver 50 | self._deltas_t = t[1:] - t[:-1] 51 | 52 | def forward(self, inputs: tf.Tensor, return_states: Optional[str] = None): 53 | 54 | def _forward_dynamics(_state): 55 | """Used in solver _state == (time, tensor)""" 56 | return [1.0, self._model(inputs=_state)] 57 | 58 | states = [] 59 | 60 | def _append_state(_state): 61 | tensors = _state[1] 62 | if return_states == "numpy": 63 | states.append(tensors.numpy()) 64 | elif return_states == "tf": 65 | states.append(tensors) 66 | 67 | with tf.name_scope("forward"): 68 | t0 = tf.to_float(self._t[0]) 69 | state = [t0, inputs] 70 | _append_state(state) 71 | for dt in self._deltas_t: 72 | state = self._solver( 73 | func=_forward_dynamics, dt=tf.to_float(dt), state=state 74 | ) 75 | _append_state(state) 76 | 77 | outputs = state[1] 78 | if return_states: 79 | return outputs, states 80 | return outputs 81 | 82 | def _backward_dynamics(self, state): 83 | t = state[0] 84 | ht = state[1] 85 | at = -state[2] 86 | 87 | with tf.GradientTape() as g: 88 | g.watch(ht) 89 | ht_new = self._model(inputs=[t, ht]) 90 | 91 | gradients = g.gradient( 92 | target=ht_new, sources=[ht] + self._model.weights, 93 | output_gradients=at 94 | ) 95 | 96 | return [1.0, ht_new, *gradients] 97 | 98 | def backward(self, outputs: tf.Tensor, 99 | output_gradients: Optional[tf.Tensor] = None): 100 | 101 | with tf.name_scope("backward"): 102 | grad_weights = [tf.zeros_like(w) for w in self._model.weights] 103 | t0 = tf.to_float(self._t[-1]) 104 | 105 | if output_gradients is None: 106 | output_gradients = tf.zeros_like(outputs) 107 | 108 | state = [t0, outputs, output_gradients, *grad_weights] 109 | for dt in self._deltas_t[::-1]: 110 | state = self._solver( 111 | self._backward_dynamics, dt=-tf.to_float(dt), state=state 112 | ) 113 | 114 | inputs = state[1] 115 | dLdInputs = state[2] 116 | dLdWeights = state[3:] 117 | return inputs, dLdInputs, dLdWeights 118 | -------------------------------------------------------------------------------- /Glycolysis_noisy/plotting.py: -------------------------------------------------------------------------------- 1 | import numpy as np 2 | import matplotlib as mpl 3 | #mpl.use('pgf') 4 | 5 | def figsize(scale, nplots = 1): 6 | fig_width_pt = 390.0 # Get this from LaTeX using \the\textwidth 7 | inches_per_pt = 1.0/72.27 # Convert pt to inch 8 | golden_mean = (np.sqrt(5.0)-1.0)/2.0 # Aesthetic ratio (you could change this) 9 | fig_width = fig_width_pt*inches_per_pt*scale # width in inches 10 | fig_height = nplots*fig_width*golden_mean # height in inches 11 | fig_size = [fig_width,fig_height] 12 | return fig_size 13 | 14 | pgf_with_latex = { # setup matplotlib to use latex for output 15 | "pgf.texsystem": "pdflatex", # change this if using xetex or lautex 16 | "text.usetex": True, # use LaTeX to write all text 17 | "font.family": "serif", 18 | "font.serif": [], # blank entries should cause plots to inherit fonts from the document 19 | "font.sans-serif": [], 20 | "font.monospace": [], 21 | "axes.labelsize": 10, # LaTeX default is 10pt font. 22 | "font.size": 10, 23 | "legend.fontsize": 8, # Make the legend/label fonts a little smaller 24 | "xtick.labelsize": 8, 25 | "ytick.labelsize": 8, 26 | "figure.figsize": figsize(1.0), # default fig size of 0.9 textwidth 27 | "pgf.preamble": [ 28 | r"\usepackage[utf8x]{inputenc}", # use utf8 fonts becasue your computer can handle it :) 29 | r"\usepackage[T1]{fontenc}", # plots will be generated using this preamble 30 | ] 31 | } 32 | mpl.rcParams.update(pgf_with_latex) 33 | 34 | import matplotlib.pyplot as plt 35 | 36 | # I make my own newfig and savefig functions 37 | def newfig(width, nplots = 1): 38 | fig = plt.figure(figsize=figsize(width, nplots)) 39 | ax = fig.add_subplot(111) 40 | return fig, ax 41 | 42 | def savefig(filename, crop = True): 43 | if crop == True: 44 | # plt.savefig('{}.pgf'.format(filename), bbox_inches='tight', pad_inches=0) 45 | plt.savefig('{}.pdf'.format(filename), bbox_inches='tight', pad_inches=0) 46 | plt.savefig('{}.eps'.format(filename), bbox_inches='tight', pad_inches=0) 47 | else: 48 | # plt.savefig('{}.pgf'.format(filename)) 49 | plt.savefig('{}.pdf'.format(filename)) 50 | plt.savefig('{}.eps'.format(filename)) 51 | 52 | ## Simple plot 53 | #fig, ax = newfig(1.0) 54 | # 55 | #def ema(y, a): 56 | # s = [] 57 | # s.append(y[0]) 58 | # for t in range(1, len(y)): 59 | # s.append(a * y[t] + (1-a) * s[t-1]) 60 | # return np.array(s) 61 | # 62 | #y = [0]*200 63 | #y.extend([20]*(1000-len(y))) 64 | #s = ema(y, 0.01) 65 | # 66 | #ax.plot(s) 67 | #ax.set_xlabel('X Label') 68 | #ax.set_ylabel('EMA') 69 | # 70 | #savefig('ema') -------------------------------------------------------------------------------- /Glycolysis_noisy/pre_train.py: -------------------------------------------------------------------------------- 1 | import os 2 | os.environ['KMP_DUPLICATE_LIB_OK']='True' 3 | 4 | import numpy as np 5 | import numpy.random as npr 6 | import tensorflow as tf 7 | import matplotlib.pyplot as plt 8 | import matplotlib as mpl 9 | from mpl_toolkits.mplot3d import Axes3D 10 | plt.switch_backend('agg') 11 | import tensorflow.contrib.eager as tfe 12 | from tensorflow.keras import layers, initializers 13 | from scipy.integrate import odeint 14 | 15 | 16 | keras = tf.keras 17 | tf.enable_eager_execution() 18 | 19 | from neural_ode import NeuralODE 20 | 21 | np.random.seed(1234) 22 | tf.set_random_seed(1234) 23 | 24 | 25 | if __name__ == "__main__": 26 | 27 | def Glycolysis(z, t, J0, k1, k2, k3, k4, k5, k6, k, ka, q, KI, phi, N, A): 28 | x1, x2, x3, x4, x5, x6, x7 = z 29 | 30 | J = ka * (x4 - x7) 31 | N1 = N - x5 32 | A2 = A - x6 33 | 34 | v1 = k1 * x1 * x6 * (1 + (x6 / KI) ** q) ** (-1) 35 | v2 = k2 * x2 * N1 36 | v3 = k3 * x3 * A2 37 | v4 = k4 * x4 * x5 38 | v5 = k5 * x6 39 | v6 = k6 * x2 * x5 40 | v7 = k * x7 41 | 42 | r1 = J0 - v1 43 | r2 = 2*v1 - v2 - v6 44 | r3 = v2 - v3 45 | r4 = v3 - v4 - J 46 | r5 = v2 - v4 - v6 47 | r6 = - 2 * v1 + 2 * v3 - v5 48 | r7 = phi * J - v7 49 | dzdt = [r1, r2, r3, r4, r5, r6, r7] 50 | return dzdt 51 | 52 | 53 | def plot_spiral(trajectories, width = 1.): 54 | x1 = trajectories[:,0] 55 | x2 = trajectories[:,1] 56 | x3 = trajectories[:,2] 57 | x4 = trajectories[:,3] 58 | x5 = trajectories[:,4] 59 | x6 = trajectories[:,5] 60 | x7 = trajectories[:,6] 61 | plt.plot(x1, linewidth = width) 62 | plt.plot(x2, linewidth = width) 63 | plt.plot(x3, linewidth = width) 64 | plt.plot(x4, linewidth = width) 65 | plt.plot(x5, linewidth = width) 66 | plt.plot(x6, linewidth = width) 67 | plt.plot(x7, linewidth = width) 68 | 69 | 70 | data_size = 4000 71 | batch_time = 60 72 | batch_size = 1000 73 | 74 | J0 = 2.5 75 | k1 = 100. 76 | k2 = 6. 77 | k3 = 16. 78 | k4 = 100. 79 | k5 = 1.28 80 | k6 = 12. 81 | k = 1.8 82 | ka = 13. 83 | q = 4. 84 | KI = 0.52 85 | phi = 0.1 86 | N = 1. 87 | A = 4. 88 | 89 | 90 | t_grid = np.linspace(0, 5, data_size) 91 | true_y0 = tf.to_float([[0.5, 1.9, 0.18, 0.15, 0.16, 0.1, 0.064]]) 92 | z0 = [0.5, 1.9, 0.18, 0.15, 0.16, 0.1, 0.064] 93 | true_yy = odeint(Glycolysis, z0, t_grid, args=(J0, k1, k2, k3, k4, k5, k6, k, ka, q, KI, phi, N, A)) 94 | 95 | true_y = true_yy + 0. * np.random.randn(true_yy.shape[0], true_yy.shape[1]) 96 | 97 | # normalizing data and system 98 | sigma_normal1 = np.std(true_y[:,0:1]) 99 | sigma_normal2 = np.std(true_y[:,1:2]) 100 | sigma_normal3 = np.std(true_y[:,2:3]) 101 | sigma_normal4 = np.std(true_y[:,3:4]) 102 | sigma_normal5 = np.std(true_y[:,4:5]) 103 | sigma_normal6 = np.std(true_y[:,5:6]) 104 | sigma_normal7 = np.std(true_y[:,6:7]) 105 | 106 | np.save('true_y', true_y) 107 | 108 | # scale for the isotropic noise 109 | noise_level = 0.02 110 | 111 | true_y[:,0:1] = true_y[:,0:1]/ sigma_normal1 + noise_level * np.random.randn(true_y[:,0:1].shape[0], true_y[:,0:1].shape[1]) 112 | true_y[:,1:2] = true_y[:,1:2]/ sigma_normal2 + noise_level * np.random.randn(true_y[:,1:2].shape[0], true_y[:,1:2].shape[1]) 113 | true_y[:,2:3] = true_y[:,2:3]/ sigma_normal3 + noise_level * np.random.randn(true_y[:,2:3].shape[0], true_y[:,2:3].shape[1]) 114 | true_y[:,3:4] = true_y[:,3:4]/ sigma_normal4 + noise_level * np.random.randn(true_y[:,3:4].shape[0], true_y[:,3:4].shape[1]) 115 | true_y[:,4:5] = true_y[:,4:5]/ sigma_normal5 + noise_level * np.random.randn(true_y[:,4:5].shape[0], true_y[:,4:5].shape[1]) 116 | true_y[:,5:6] = true_y[:,5:6]/ sigma_normal6 + noise_level * np.random.randn(true_y[:,5:6].shape[0], true_y[:,5:6].shape[1]) 117 | true_y[:,6:7] = true_y[:,6:7]/ sigma_normal7 + noise_level * np.random.randn(true_y[:,6:7].shape[0], true_y[:,6:7].shape[1]) 118 | 119 | 120 | 121 | plt.figure(1, figsize=(12,8)) 122 | plot_spiral(true_y, 1) 123 | plot_spiral(true_yy, 1) 124 | plt.xlabel('$t$',fontsize=13) 125 | plt.ylabel('$x$',fontsize=13) 126 | plt.legend(loc='upper left', frameon=False, prop={'size': 13}) 127 | plt.savefig('./True_trajectory.png', dpi = 600) 128 | 129 | 130 | def get_batch(): 131 | """Returns initial point and last point over sampled frament of trajectory""" 132 | starts = np.random.choice(np.arange(data_size - batch_time - 1, dtype=np.int64), batch_size, replace=False) 133 | batch_y0 = true_y[starts] 134 | batch_yN = true_y[starts + batch_time] 135 | return tf.to_float(batch_y0), tf.to_float(batch_yN) 136 | 137 | 138 | num_param = 14 139 | para_num = num_param 140 | 141 | t0 = t_grid[:batch_time][0] 142 | t1 = t_grid[:batch_time][-1] 143 | t_in = np.linspace(t0, t1, 10) 144 | 145 | batch_y0, batch_yN = get_batch() 146 | 147 | np.save('batch_y0', batch_y0) 148 | np.save('batch_yN', batch_yN) 149 | 150 | ######################################### 151 | ########## precondition start ########### 152 | ######################################### 153 | niters_pre = 10000 154 | 155 | class ODEModel_pre(tf.keras.Model): 156 | def __init__(self): 157 | super(ODEModel_pre, self).__init__() 158 | self.Weights = tfe.Variable(tf.random_normal([num_param, 1], dtype=tf.float32)*0.01, dtype=tf.float32) 159 | 160 | def call(self, inputs, **kwargs): 161 | t, y = inputs 162 | h = y 163 | x1 = h[:,0:1] * sigma_normal1 164 | x2 = h[:,1:2] * sigma_normal2 165 | x3 = h[:,2:3] * sigma_normal3 166 | x4 = h[:,3:4] * sigma_normal4 167 | x5 = h[:,4:5] * sigma_normal5 168 | x6 = h[:,5:6] * sigma_normal6 169 | x7 = h[:,6:7] * sigma_normal7 170 | 171 | J0 = self.Weights[0] 172 | k1 = self.Weights[1] 173 | k2 = self.Weights[2] 174 | k3 = self.Weights[3] 175 | k4 = self.Weights[4] 176 | k5 = self.Weights[5] 177 | k6 = self.Weights[6] 178 | k = self.Weights[7] 179 | ka = self.Weights[8] 180 | q = self.Weights[9] 181 | KI = self.Weights[10] 182 | phi = self.Weights[11] 183 | N = self.Weights[12] 184 | A = self.Weights[13] 185 | 186 | ka = tf.exp(ka) 187 | q = tf.exp(q) 188 | N = tf.exp(N) 189 | A = tf.exp(A) 190 | k1 = tf.exp(k1) 191 | k2 = tf.exp(k2) 192 | k3 = tf.exp(k3) 193 | k4 = tf.exp(k4) 194 | k5 = tf.exp(k5) 195 | k6 = tf.exp(k6) 196 | k = tf.exp(k) 197 | KI = tf.exp(KI) 198 | J0 = tf.exp(J0) 199 | phi = tf.exp(phi) 200 | 201 | J = ka * (x4 - x7) 202 | N1 = N - x5 203 | A2 = A - x6 204 | 205 | v1 = k1 * x1 * x6 * (1 + (x6 / KI) ** q) ** (-1) 206 | v2 = k2 * x2 * N1 207 | v3 = k3 * x3 * A2 208 | v4 = k4 * x4 * x5 209 | v5 = k5 * x6 210 | v6 = k6 * x2 * x5 211 | v7 = k * x7 212 | 213 | r1 = J0 - v1 214 | r2 = 2*v1 - v2 - v6 215 | r3 = v2 - v3 216 | r4 = v3 - v4 - J 217 | r5 = v2 - v4 - v6 218 | r6 = - 2 * v1 + 2 * v3 - v5 219 | r7 = phi * J - v7 220 | 221 | h_out = tf.concat([r1/sigma_normal1, r2/sigma_normal2, r3/sigma_normal3, r4/sigma_normal4, r5/sigma_normal5, r6/sigma_normal6, r7/sigma_normal7], 1) 222 | return h_out 223 | 224 | 225 | model_pre = ODEModel_pre() 226 | neural_ode_pre = NeuralODE(model_pre, t_in) 227 | optimizer = tf.train.AdamOptimizer(5e-2) 228 | 229 | 230 | def compute_gradients_and_update_pre(batch_y0, batch_yN): 231 | """Takes start positions (x0, y0) and final positions (xN, yN)""" 232 | pred_y = neural_ode_pre.forward(batch_y0) 233 | with tf.GradientTape() as g_pre: 234 | g_pre.watch(pred_y) 235 | loss = tf.reduce_mean((pred_y - batch_yN)**2) + tf.reduce_sum(tf.abs(model_pre.trainable_weights[0])) 236 | loss_star = tf.reduce_sum((pred_y - batch_yN)**2) 237 | 238 | dLoss = g_pre.gradient(loss, pred_y) 239 | h_start, dfdh0, dWeights = neural_ode_pre.backward(pred_y, dLoss) 240 | optimizer.apply_gradients(zip(dWeights, model_pre.weights)) 241 | return loss, dWeights, loss_star 242 | 243 | # Compile EAGER graph to static (this will be much faster) 244 | compute_gradients_and_update_pre = tfe.defun(compute_gradients_and_update_pre) 245 | 246 | parameters_pre = np.zeros((para_num, 1)) 247 | 248 | for step in range(niters_pre): 249 | print(step) 250 | loss, dWeights, loss_star = compute_gradients_and_update_pre(batch_y0, batch_yN) 251 | print("loss is", loss_star) 252 | 253 | parameters_pre = model_pre.trainable_weights[0].numpy() 254 | 255 | print(np.exp(parameters_pre)) 256 | 257 | np.save('parameters_pre', parameters_pre) 258 | 259 | ######################################### 260 | ########## precondition end ############# 261 | ######################################### 262 | 263 | -------------------------------------------------------------------------------- /Hybrid_Sine_NN/HMC.py: -------------------------------------------------------------------------------- 1 | import os 2 | os.environ['KMP_DUPLICATE_LIB_OK']='True' 3 | 4 | import numpy as np 5 | import numpy.random as npr 6 | import tensorflow as tf 7 | import matplotlib.pyplot as plt 8 | plt.switch_backend('agg') 9 | import tensorflow.contrib.eager as tfe 10 | from tensorflow.keras import layers, initializers 11 | from scipy.integrate import odeint 12 | 13 | 14 | keras = tf.keras 15 | tf.enable_eager_execution() 16 | 17 | from neural_ode import NeuralODE 18 | 19 | np.random.seed(1234) 20 | tf.set_random_seed(1234) 21 | 22 | if __name__ == "__main__": 23 | 24 | def Damped_pendulum(z, t, alpha, beta): 25 | x, y = z 26 | dzdt = [y, - alpha * y - beta * np.sin(x)] 27 | return dzdt 28 | 29 | def plot_spiral(trajectories, width = 1.): 30 | x = trajectories[:,0] 31 | y = trajectories[:,1] 32 | plt.plot(x, y, linewidth=width) 33 | 34 | 35 | data_size = 8001 36 | batch_time = 200 37 | batch_size = 1000 38 | niters = 1500 39 | 40 | 41 | alpha = 0.2 42 | beta = 8.91 43 | 44 | 45 | z0 = [-1.193, -3.876] 46 | true_y0 = tf.to_float([[-1.193, -3.876]]) 47 | t_grid = np.linspace(0, 20, data_size) 48 | 49 | true_yy = odeint(Damped_pendulum, z0, t_grid, args=(alpha, beta)) 50 | print(true_yy.shape) 51 | 52 | 53 | true_y = np.load("true_y.npy") 54 | batch_y0 = np.load("batch_y0.npy") 55 | batch_yN = np.load("batch_yN.npy") 56 | 57 | batch_y0 = tf.to_float(batch_y0) 58 | batch_yN = tf.to_float(batch_yN) 59 | 60 | # normalizing data and system 61 | sigma_normal1 = np.std(true_y[:,0:1]) 62 | sigma_normal2 = np.std(true_y[:,1:2]) 63 | 64 | 65 | t0 = t_grid[:batch_time][0] 66 | t1 = t_grid[:batch_time][-1] 67 | t_in = np.linspace(t0, t1, 20) 68 | 69 | 70 | 71 | ## NN structure 72 | layers = [1, 20, 20, 1] 73 | def get_para_num(): 74 | L = len(layers) 75 | para_num = 0 76 | for k in range(L-1): 77 | para_num = para_num + layers[k] * layers[k+1] + layers[k+1] 78 | return para_num 79 | 80 | num_param = 9 81 | num_param_NN = get_para_num() 82 | 83 | 84 | para_num = num_param + num_param_NN 85 | 86 | parameters_pre = np.load("parameters_pre.npy") 87 | print(parameters_pre) 88 | 89 | ######################################### 90 | ########## precondition end ############# 91 | ######################################### 92 | 93 | initial_weight = parameters_pre 94 | print(initial_weight.shape, "here") 95 | 96 | class ODEModel(tf.keras.Model): 97 | def __init__(self): 98 | super(ODEModel, self).__init__() 99 | self.Weights = tfe.Variable(tf.random_normal([num_param + num_param_NN, 1], dtype=tf.float32)*0.01, dtype=tf.float32) 100 | 101 | def forward_pass(self, H, layers, W): 102 | num_layers = len(layers) 103 | W_seq = W 104 | 105 | for k in range(0,num_layers-2): 106 | W = W_seq[0:layers[k] * layers[k+1]] 107 | W = tf.reshape(W, (layers[k], layers[k+1])) 108 | W_seq = W_seq[layers[k] * layers[k+1]:] 109 | b = W_seq[0:layers[k+1]] 110 | b = tf.reshape(b, (1, layers[k+1])) 111 | W_seq = W_seq[layers[k+1]:] 112 | H = tf.tanh(tf.add(tf.matmul(H, W), b)) 113 | 114 | W = W_seq[0:layers[num_layers-2] * layers[num_layers-1]] 115 | W = tf.reshape(W, (layers[num_layers-2], layers[num_layers-1])) 116 | W_seq = W_seq[layers[num_layers-2] * layers[num_layers-1]:] 117 | b = W_seq[0:layers[num_layers-1]] 118 | b = tf.reshape(b, (1, layers[num_layers-1])) 119 | H = tf.add(tf.matmul(H, W), b) 120 | return H 121 | 122 | def call(self, inputs, **kwargs): 123 | t, y = inputs 124 | h = y 125 | h1 = h[:,0:1] 126 | h2 = h[:,1:2] 127 | 128 | p11 = self.Weights[0] 129 | p12 = self.Weights[1] 130 | p13 = self.Weights[2] 131 | p14 = self.Weights[3] 132 | p15 = self.Weights[4] 133 | p16 = self.Weights[5] 134 | p23 = self.Weights[6] 135 | p25 = self.Weights[7] 136 | p26 = self.Weights[8] 137 | 138 | w_NN = self.Weights[num_param:,:] 139 | 140 | NN_out = self.forward_pass(sigma_normal1 * h1, layers, w_NN) 141 | 142 | h_out1 = p11 + sigma_normal1 * p12 *h1 + sigma_normal2 * p13 *h2 + sigma_normal1**2 * p14*h1**2\ 143 | + sigma_normal1 * sigma_normal2 * p15*h1*h2 + sigma_normal2**2 * p16*h2**2 144 | h_out2 = sigma_normal2 * p23 *h2 \ 145 | + sigma_normal1 * sigma_normal2 * p25*h1*h2 + sigma_normal2**2 * p26*h2**2 + NN_out 146 | h_out = tf.concat([h_out1/sigma_normal1, h_out2/sigma_normal2], 1) 147 | return h_out 148 | 149 | 150 | model = ODEModel() 151 | neural_ode = NeuralODE(model, t = t_in) 152 | 153 | def compute_gradients_and_update(batch_y0, batch_yN): 154 | """Takes start positions (x0, y0) and final positions (xN, yN)""" 155 | pred_y = neural_ode.forward(batch_y0) 156 | with tf.GradientTape() as g: 157 | g.watch(pred_y) 158 | loss = tf.reduce_sum((pred_y - batch_yN)**2) 159 | 160 | dLoss = g.gradient(loss, pred_y) 161 | h_start, dfdh0, dWeights = neural_ode.backward(pred_y, dLoss) 162 | 163 | return loss, dWeights 164 | 165 | # Compile EAGER graph to static (this will be much faster) 166 | compute_gradients_and_update = tfe.defun(compute_gradients_and_update) 167 | 168 | 169 | # function to compute the kinetic energy 170 | def kinetic_energy(V, loggamma_v, loglambda_v): 171 | q = (np.sum(-V**2) - loggamma_v**2 - loglambda_v**2)/2.0 172 | return q 173 | 174 | def compute_gradient_param(dWeights, loggamma, loglambda, batch_size, para_num): 175 | WW = model.trainable_weights[0].numpy() 176 | dWeights = np.exp(loggamma)/2.0 * dWeights + np.exp(loglambda) * np.sign(WW) 177 | return dWeights 178 | 179 | def compute_gradient_hyper(loss, weights, loggamma, loglambda, batch_size, para_num): 180 | grad_loggamma = np.exp(loggamma) * (loss/2.0 + 1.0) - (batch_size/2.0 + 1.0) 181 | grad_loglambda = np.exp(loglambda) * (np.sum(np.abs(weights)) + 1.0) - (para_num + 1.0) 182 | 183 | return grad_loggamma, grad_loglambda 184 | 185 | def compute_Hamiltonian(loss, weights, loggamma, loglambda, batch_size, para_num): 186 | H = np.exp(loggamma)*(loss/2.0 + 1.0) + np.exp(loglambda)*(np.sum(np.abs(weights)) + 1.0)\ 187 | - (batch_size/2.0 + 1.0) * loggamma - (para_num + 1.0) * loglambda 188 | return H 189 | 190 | def leap_frog(v_in, w_in, loggamma_in, loglambda_in, loggamma_v_in, loglambda_v_in): 191 | # assign weights from the previous step to the model 192 | model.trainable_weights[0].assign(w_in) 193 | 194 | v_new = v_in 195 | loggamma_v_new = loggamma_v_in 196 | loglambda_v_new = loglambda_v_in 197 | 198 | loggamma_new = loggamma_in 199 | loglambda_new = loglambda_in 200 | w_new = w_in 201 | 202 | for m in range(L): 203 | 204 | loss, dWeights = compute_gradients_and_update(batch_y0, batch_yN) # evaluate the gradient 205 | 206 | dWeights = np.asarray(dWeights[0]) # make the gradient to be numpy array 207 | # print(dWeights) 208 | dWeights = compute_gradient_param(dWeights, loggamma_new, loglambda_new, batch_size, para_num) 209 | grad_loggamma, grad_loglambda = compute_gradient_hyper(loss, w_new, loggamma_new, loglambda_new, batch_size, para_num) 210 | 211 | loggamma_v_new = loggamma_v_new - epsilon/2*grad_loggamma 212 | loglambda_v_new = loglambda_v_new - epsilon/2*grad_loglambda 213 | v_new = v_new - epsilon/2*(dWeights) 214 | w_new = model.trainable_weights[0].numpy() + epsilon * v_new 215 | model.trainable_weights[0].assign(w_new) 216 | loggamma_new = loggamma_new + epsilon * loggamma_v_new 217 | loglambda_new = loglambda_new + epsilon * loglambda_v_new 218 | 219 | # Second half of the leap frog 220 | loss, dWeights = compute_gradients_and_update(batch_y0, batch_yN) 221 | dWeights = np.asarray(dWeights[0]) 222 | dWeights = compute_gradient_param(dWeights, loggamma_new, loglambda_new, batch_size, para_num) 223 | grad_loggamma, grad_loglambda = compute_gradient_hyper(loss, w_new, loggamma_new, loglambda_new, batch_size, para_num) 224 | 225 | v_new = v_new - epsilon/2*(dWeights) 226 | loggamma_v_new = loggamma_v_new - epsilon/2*grad_loggamma 227 | loglambda_v_new = loglambda_v_new - epsilon/2*grad_loglambda 228 | 229 | print(np.exp(loggamma_new)) 230 | print(np.exp(loglambda_new)) 231 | 232 | return v_new, w_new, loggamma_new, loglambda_new, loggamma_v_new, loglambda_v_new 233 | 234 | 235 | 236 | neural_ode_test = NeuralODE(model, t=t_grid[0:data_size:20]) 237 | parameters = np.zeros((niters, para_num)) # book keeping the parameters 238 | loggammalist = np.zeros((niters, 1)) # book keeping the loggamma 239 | loglambdalist = np.zeros((niters, 1)) # book keeping the loggamma 240 | loglikelihood = np.zeros((niters, 1)) # book keeping the loggamma 241 | L = 10 # leap frog step number 242 | epsilon = 0.001 # leap frog step size 243 | epsilon_max = 0.001 # max 0.001 244 | epsilon_min = 0.001 # max 0.001 245 | 246 | 247 | def compute_epsilon(step): 248 | coefficient = np.log(epsilon_max/epsilon_min) 249 | return epsilon_max * np.exp( - step * coefficient / niters) 250 | 251 | 252 | # initial weight 253 | w_temp = initial_weight 254 | print("initial_w", w_temp) 255 | loggamma_temp = 4. + np.random.normal() 256 | loglambda_temp = np.random.normal() 257 | 258 | model.trainable_weights[0].assign(w_temp) 259 | loss_original, _ = compute_gradients_and_update(batch_y0, batch_yN) # compute the initial Hamiltonian 260 | 261 | loggamma_temp = np.log(batch_size / loss_original) 262 | print("This is initial guess", loggamma_temp, "with loss", loss_original) 263 | if loggamma_temp > 6.: 264 | loggamma_temp = 6. 265 | epsilon_max = 0.0005 266 | epsilon_min = 0.0005 267 | 268 | 269 | # training steps 270 | for step in range(niters): 271 | 272 | epsilon = compute_epsilon(step) 273 | 274 | print(step) 275 | v_initial = np.random.randn(para_num, 1) # initialize the velocity 276 | loggamma_v_initial = np.random.normal() 277 | loglambda_v_initial = np.random.normal() 278 | 279 | loss_initial, _ = compute_gradients_and_update(batch_y0, batch_yN) # compute the initial Hamiltonian 280 | loss_initial = compute_Hamiltonian(loss_initial, w_temp, loggamma_temp, loglambda_temp, batch_size, para_num) 281 | 282 | v_new, w_new, loggamma_new, loglambda_new, loggamma_v_new, loglambda_v_new = \ 283 | leap_frog(v_initial, w_temp, loggamma_temp, loglambda_temp, loggamma_v_initial, loglambda_v_initial) 284 | 285 | # compute the final Hamiltonian 286 | loss_finial, _ = compute_gradients_and_update(batch_y0, batch_yN) 287 | loss_finial = compute_Hamiltonian(loss_finial, w_new, loggamma_new, loglambda_new, batch_size, para_num) 288 | 289 | # making decisions 290 | p_temp = np.exp(-loss_finial + loss_initial + \ 291 | kinetic_energy(v_new, loggamma_v_new, loglambda_v_new) - kinetic_energy(v_initial, loggamma_v_initial, loglambda_v_initial)) 292 | 293 | p = min(1, p_temp) 294 | p_decision = np.random.uniform() 295 | if p > p_decision: 296 | parameters[step:step+1, :] = np.transpose(w_new) 297 | w_temp = w_new 298 | loggammalist[step, 0] = loggamma_new 299 | loglambdalist[step, 0] = loglambda_new 300 | loglikelihood[step, 0] = loss_finial 301 | loggamma_temp = loggamma_new 302 | loglambda_temp = loglambda_new 303 | else: 304 | parameters[step:step+1, :] = np.transpose(w_temp) 305 | model.trainable_weights[0].assign(w_temp) 306 | loggammalist[step, 0] = loggamma_temp 307 | loglambdalist[step, 0] = loglambda_temp 308 | loglikelihood[step, 0] = loss_initial 309 | 310 | print('probability', p) 311 | print(p > p_decision) 312 | 313 | 314 | np.save('parameters', parameters) 315 | np.save('loggammalist', loggammalist) 316 | np.save('loglikelihood', loglikelihood) 317 | 318 | -------------------------------------------------------------------------------- /Hybrid_Sine_NN/Show_NN_fit.py: -------------------------------------------------------------------------------- 1 | import os 2 | os.environ['KMP_DUPLICATE_LIB_OK']='True' 3 | 4 | import numpy as np 5 | import numpy.random as npr 6 | import tensorflow as tf 7 | import matplotlib.pyplot as plt 8 | plt.switch_backend('agg') 9 | 10 | from scipy.integrate import odeint 11 | from scipy import stats 12 | 13 | 14 | from plotting import newfig, savefig 15 | 16 | np.random.seed(1234) 17 | tf.set_random_seed(1234) 18 | 19 | if __name__ == "__main__": 20 | 21 | layers = [1, 20, 20, 1] 22 | def get_para_num(): 23 | L = len(layers) 24 | para_num = 0 25 | for k in range(L-1): 26 | para_num = para_num + layers[k] * layers[k+1] + layers[k+1] 27 | return para_num 28 | 29 | true_y = np.load("true_y.npy") 30 | batch_y0 = np.load("batch_y0.npy") 31 | 32 | sigma_normal1 = np.std(true_y[:,0:1]) 33 | 34 | 35 | x1 = true_y[:,0] 36 | x1_data = batch_y0[:,0] * sigma_normal1 37 | 38 | 39 | print(x1.shape) 40 | 41 | num_param = 9 42 | num_param_NN = get_para_num() 43 | 44 | 45 | def NN_forward_pass(H, layers, W): 46 | num_layers = len(layers) 47 | W_seq = W 48 | 49 | for k in range(0,num_layers-2): 50 | W = W_seq[0:layers[k] * layers[k+1]] 51 | W = np.reshape(W, (layers[k], layers[k+1])) 52 | W_seq = W_seq[layers[k] * layers[k+1]:] 53 | b = W_seq[0:layers[k+1]] 54 | b = np.reshape(b, (1, layers[k+1])) 55 | W_seq = W_seq[layers[k+1]:] 56 | H = np.tanh(np.add(np.matmul(H, W), b)) 57 | 58 | W = W_seq[0:layers[num_layers-2] * layers[num_layers-1]] 59 | W = np.reshape(W, (layers[num_layers-2], layers[num_layers-1])) 60 | W_seq = W_seq[layers[num_layers-2] * layers[num_layers-1]:] 61 | b = W_seq[0:layers[num_layers-1]] 62 | b = np.reshape(b, (1, layers[num_layers-1])) 63 | H = np.add(np.matmul(H, W), b) 64 | return H 65 | 66 | 67 | N_samples = 100 68 | beta = 8.91 69 | N_total = 500 70 | 71 | X_test = np.linspace(-np.pi, np.pi, 500)[:,None] 72 | Y_test = - beta * np.sin(X_test) 73 | 74 | 75 | parameters = np.load("parameters.npy") 76 | precision = np.load("loggammalist.npy") 77 | loglikelihood = np.load("loglikelihood.npy") 78 | print(parameters.shape) 79 | parameters = parameters[:,num_param:] 80 | print(parameters.shape) 81 | loglikelihood = loglikelihood[-N_total:] 82 | num_samples = parameters.shape[0] 83 | length_dict = parameters.shape[1] 84 | num_dim = 2 85 | 86 | 87 | idx_MAP = np.argmin(loglikelihood) 88 | MAP = parameters[idx_MAP, :] 89 | 90 | y_MAP = NN_forward_pass(X_test, layers, MAP) 91 | 92 | 93 | y_BNODE = np.zeros((X_test.shape[0], N_samples)) 94 | 95 | for k in range(N_samples): 96 | print(k) 97 | idx_1 = np.random.randint(N_total) 98 | W_sample = parameters[-idx_1, :] 99 | y_BNODE[:,k:k+1] = NN_forward_pass(X_test, layers, W_sample) 100 | 101 | 102 | mu_pred = np.mean(y_BNODE, axis = 1) 103 | Sigma_pred = np.var(y_BNODE, axis = 1) 104 | 105 | 106 | plt.rcParams.update(plt.rcParamsDefault) 107 | plt.rc('font', family='serif') 108 | plt.rcParams.update({'font.size': 16, 109 | 'lines.linewidth': 2, 110 | 'axes.labelsize': 20, # fontsize for x and y labels (was 10) 111 | 'axes.titlesize': 20, 112 | 'xtick.labelsize': 16, 113 | 'ytick.labelsize': 16, 114 | 'legend.fontsize': 20, 115 | 'axes.linewidth': 2, 116 | "pgf.texsystem": "pdflatex", # change this if using xetex or lautex 117 | "text.usetex": True, # use LaTeX to write all text 118 | }) 119 | 120 | 121 | plt.figure(4, figsize=(9,6)) 122 | plt.xticks(fontsize=20) 123 | plt.yticks(fontsize=20) 124 | plt.plot(x1_data, - beta * np.sin(x1_data), 'ro', alpha = 0.5, markersize = 4, label = "Training data") 125 | plt.plot(X_test, Y_test, 'r-', label = r"$-\beta\sin(x_1)$") 126 | plt.plot(X_test, y_MAP, 'g--', label = r"MAP curve of NN output") 127 | lower_0 = mu_pred - 2.0*np.sqrt(Sigma_pred) 128 | upper_0 = mu_pred + 2.0*np.sqrt(Sigma_pred) 129 | plt.fill_between(X_test.flatten(), lower_0.flatten(), upper_0.flatten(), 130 | facecolor='orange', alpha=0.5, label="Two std band") 131 | plt.xlabel('$x_1$',fontsize=22) 132 | plt.ylabel(r'$-\beta\sin(x_1)$',fontsize=20) 133 | plt.axvspan(min(x1), max(x1), alpha=0.1, color='blue') 134 | plt.xlim((-2.2, 2.)) 135 | plt.legend(loc='upper right', frameon=False, prop={'size': 18}) 136 | plt.savefig('./BNODE_Sine.png', dpi = 300) 137 | 138 | 139 | 140 | 141 | 142 | -------------------------------------------------------------------------------- /Hybrid_Sine_NN/create_nice_plot_hybrid.py: -------------------------------------------------------------------------------- 1 | import os 2 | os.environ['KMP_DUPLICATE_LIB_OK']='True' 3 | 4 | import numpy as np 5 | import numpy.random as npr 6 | import tensorflow as tf 7 | import matplotlib.pyplot as plt 8 | plt.switch_backend('agg') 9 | 10 | from scipy.integrate import odeint 11 | from scipy import stats 12 | 13 | 14 | from plotting import newfig, savefig 15 | 16 | 17 | np.random.seed(1234) 18 | tf.set_random_seed(1234) 19 | 20 | 21 | if __name__ == "__main__": 22 | 23 | layers = [1, 20, 20, 1] 24 | def get_para_num(): 25 | L = len(layers) 26 | para_num = 0 27 | for k in range(L-1): 28 | para_num = para_num + layers[k] * layers[k+1] + layers[k+1] 29 | return para_num 30 | 31 | 32 | num_param = 9 33 | num_param_NN = get_para_num() 34 | 35 | para_num = num_param + num_param_NN 36 | 37 | 38 | def Damped_pendulum(z, t, alpha, beta): 39 | x, y = z 40 | dzdt = [y, - alpha * y - beta * np.sin(x)] 41 | return dzdt 42 | 43 | def NN_forward_pass(H, layers, W): 44 | num_layers = len(layers) 45 | W_seq = W 46 | 47 | for k in range(0,num_layers-2): 48 | W = W_seq[0:layers[k] * layers[k+1]] 49 | W = np.reshape(W, (layers[k], layers[k+1])) 50 | W_seq = W_seq[layers[k] * layers[k+1]:] 51 | b = W_seq[0:layers[k+1]] 52 | b = np.reshape(b, (1, layers[k+1])) 53 | W_seq = W_seq[layers[k+1]:] 54 | H = np.tanh(np.add(np.matmul(H, W), b)) 55 | 56 | W = W_seq[0:layers[num_layers-2] * layers[num_layers-1]] 57 | W = np.reshape(W, (layers[num_layers-2], layers[num_layers-1])) 58 | W_seq = W_seq[layers[num_layers-2] * layers[num_layers-1]:] 59 | b = W_seq[0:layers[num_layers-1]] 60 | b = np.reshape(b, (1, layers[num_layers-1])) 61 | H = np.add(np.matmul(H, W), b) 62 | return H 63 | 64 | Order = 2 65 | def Damped_pendulum_BNODE(z, t, W, Order): 66 | x, y = z 67 | 68 | p11 = W[0] 69 | p12 = W[1] 70 | p13 = W[2] 71 | p14 = W[3] 72 | p15 = W[4] 73 | p16 = W[5] 74 | p23 = W[6] 75 | p25 = W[7] 76 | p26 = W[8] 77 | 78 | w_NN = W[num_param:][:,None] 79 | 80 | x_in = [x] 81 | NN_out = NN_forward_pass(x_in, layers, w_NN) 82 | 83 | r1 = p11 + p12 * x + p13 * y + p14 *x**2 + p15 * x * y * p16 * y**2 84 | r2 = p23 * y + p25 * x * y * p26 * y**2 + NN_out 85 | dzdt = [r1, r2] 86 | 87 | return dzdt 88 | 89 | 90 | data_size = 8000 91 | batch_time = 200 92 | batch_size = 1000 93 | 94 | data_size_true = 2000 95 | 96 | alpha = 0.2 97 | beta = 8.91 98 | 99 | N_samples = 100 100 | N_total = 500 101 | 102 | 103 | z0 = [-1.193, -3.876] 104 | t_grid_train = np.linspace(0, 20, data_size) 105 | t_grid_true = np.linspace(0, 40, data_size_true) 106 | 107 | y_train = odeint(Damped_pendulum, z0, t_grid_train, args=(alpha, beta)) 108 | idx = np.random.choice(np.arange(data_size - batch_time - 1, dtype=np.int64), batch_size, replace=False) 109 | y_train = y_train[idx] 110 | t_grid_train = t_grid_train[idx] 111 | 112 | y_true = odeint(Damped_pendulum, z0, t_grid_true, args=(alpha, beta)) 113 | 114 | sigma_normal1 = np.std(y_train[:,0:1]) 115 | sigma_normal2 = np.std(y_train[:,1:2]) 116 | 117 | sigma_normal = np.asarray([sigma_normal1, sigma_normal2]) 118 | 119 | 120 | parameters = np.load("parameters.npy") 121 | precision = np.load("loggammalist.npy") 122 | loglikelihood = np.load("loglikelihood.npy") 123 | precision = np.exp(precision) 124 | print("precision", precision) 125 | print(parameters.shape) 126 | loglikelihood = loglikelihood[- N_total:] 127 | num_samples = parameters.shape[0] 128 | length_dict = parameters.shape[1] 129 | num_dim = 2 130 | 131 | idx_MAP = np.argmin(loglikelihood) 132 | MAP = parameters[idx_MAP, :] 133 | 134 | y_MAP = odeint(Damped_pendulum_BNODE, z0, t_grid_true, args=(MAP, Order)) 135 | 136 | 137 | mu_pred = y_MAP 138 | 139 | y_BNODE = np.zeros((t_grid_true.shape[0], num_dim, N_samples)) 140 | 141 | for k in range(N_samples): 142 | print(k) 143 | idx_1 = np.random.randint(N_total) 144 | idx_2 = np.random.randint(N_total) 145 | W_sample = parameters[-idx_1, :] 146 | precision_here = precision[-idx_2] * num_dim 147 | y_BNODE[:,:,k] = odeint(Damped_pendulum_BNODE, z0, t_grid_true, args=(W_sample, Order)) 148 | Sigma_data = np.ones_like(mu_pred) / np.sqrt(precision_here) 149 | y_BNODE[:,:,k] = y_BNODE[:,:,k] + Sigma_data * sigma_normal * np.random.normal() 150 | 151 | mu_pred = np.mean(y_BNODE, axis = 2) 152 | Sigma_pred = np.var(y_BNODE, axis = 2) 153 | 154 | 155 | 156 | plt.rcParams.update(plt.rcParamsDefault) 157 | plt.rc('font', family='serif') 158 | plt.rcParams.update({'font.size': 16, 159 | 'lines.linewidth': 2, 160 | 'axes.labelsize': 20, # fontsize for x and y labels (was 10) 161 | 'axes.titlesize': 20, 162 | 'xtick.labelsize': 16, 163 | 'ytick.labelsize': 16, 164 | 'legend.fontsize': 20, 165 | 'axes.linewidth': 2, 166 | "pgf.texsystem": "pdflatex", # change this if using xetex or lautex 167 | "text.usetex": True, # use LaTeX to write all text 168 | }) 169 | 170 | plt.figure(1, figsize=(6,6)) 171 | plt.xticks(fontsize=14) 172 | plt.yticks(fontsize=14) 173 | plt.plot(y_train[:,0], y_train[:,1], 'ro', label = "Training data") 174 | plt.plot(y_true[:,0], y_true[:,1], 'b-', label = "Exact Trajectory") 175 | plt.xlabel('$x_1$',fontsize=18) 176 | plt.ylabel('$x_2$',fontsize=18) 177 | plt.xlim((-2.3, 2.3)) 178 | plt.ylim((-5., 6.5)) 179 | plt.legend(loc='upper right', frameon=False, prop={'size': 14}) 180 | plt.savefig('./Training_data.png', dpi = 300) 181 | 182 | 183 | plt.figure(4, figsize=(12,6.5)) 184 | plt.xticks(fontsize=22) 185 | plt.yticks(fontsize=22) 186 | plt.plot(t_grid_train, y_train[:,0], 'ro', label = "Training data of $x_1(t)$") 187 | plt.plot(t_grid_true, y_true[:,0], 'r-', label = "True Trajectory of $x_1(t)$") 188 | plt.plot(t_grid_true, y_MAP[:,0], 'g--', label = "MAP Trajectory of $x_1(t)$") 189 | lower_0 = mu_pred[:,0] - 2.0*np.sqrt(Sigma_pred[:,0]) 190 | upper_0 = mu_pred[:,0] + 2.0*np.sqrt(Sigma_pred[:,0]) 191 | plt.fill_between(t_grid_true.flatten(), lower_0.flatten(), upper_0.flatten(), 192 | facecolor='orange', alpha=0.5, label="Two std band") 193 | plt.xlabel('$t$',fontsize=26) 194 | plt.ylabel('$x_1(t)$',fontsize=26) 195 | plt.ylim((-3., 3.)) 196 | plt.legend(loc='upper right', frameon=False, prop={'size': 20}) 197 | plt.savefig('./BNODE_Prediction_x1.png', dpi = 300) 198 | 199 | 200 | plt.figure(5, figsize=(12,6.5)) 201 | plt.xticks(fontsize=22) 202 | plt.yticks(fontsize=22) 203 | plt.plot(t_grid_train, y_train[:,1], 'ro', label = "Training data of $x_2(t)$") 204 | plt.plot(t_grid_true, y_true[:,1], 'r-', label = "True Trajectory of $x_2(t)$") 205 | plt.plot(t_grid_true, y_MAP[:,1], 'g--', label = "MAP Trajectory of $x_2(t)$") 206 | lower_1 = mu_pred[:,1] - 2.0*np.sqrt(Sigma_pred[:,1]) 207 | upper_1 = mu_pred[:,1] + 2.0*np.sqrt(Sigma_pred[:,1]) 208 | plt.fill_between(t_grid_true.flatten(), lower_1.flatten(), upper_1.flatten(), 209 | facecolor='orange', alpha=0.5, label="Two std band") 210 | plt.xlabel('$t$',fontsize=26) 211 | plt.ylabel('$x_2(t)$',fontsize=26) 212 | plt.legend(loc='upper right', frameon=False, prop={'size': 20}) 213 | plt.savefig('./BNODE_Prediction_x2.png', dpi = 300) 214 | 215 | 216 | 217 | 218 | 219 | -------------------------------------------------------------------------------- /Hybrid_Sine_NN/neural_ode.py: -------------------------------------------------------------------------------- 1 | from typing import Optional, List 2 | import numpy as np 3 | import tensorflow as tf 4 | from tensorflow.python.framework.ops import EagerTensor 5 | import tensorflow.contrib.eager as tfe 6 | 7 | keras = tf.keras 8 | 9 | 10 | def zip_map(zipped, update_op): 11 | return [update_op(*elems) for elems in zipped] 12 | 13 | 14 | def euler_update(h_list, dh_list, dt): 15 | return zip_map(zip(h_list, dh_list), lambda h, dh: h + tf.cast(dt, h.dtype) * dh) 16 | 17 | 18 | def euler_step(func, dt, state): 19 | return euler_update(state, func(state), dt) 20 | 21 | 22 | def rk2_step(func, dt, state): 23 | k1 = func(state) 24 | k2 = func(euler_update(state, k1, dt)) 25 | return zip_map(zip(state, k1, k2), 26 | lambda h, dk1, dk2: h + tf.cast(dt, h.dtype) * (dk1 + dk2) / 2) 27 | 28 | 29 | def rk4_step(func, dt, state): 30 | k1 = func(state) 31 | k2 = func(euler_update(state, k1, dt / 2)) 32 | k3 = func(euler_update(state, k2, dt / 2)) 33 | k4 = func(euler_update(state, k3, dt)) 34 | 35 | return zip_map( 36 | zip(state, k1, k2, k3, k4), 37 | lambda h, dk1, dk2, dk3, dk4: h + tf.cast(dt, h.dtype) * ( 38 | dk1 + 2 * dk2 + 2 * dk3 + dk4) / 6, 39 | ) 40 | 41 | 42 | class NeuralODE: 43 | def __init__( 44 | self, model: tf.keras.Model, t=np.linspace(0, 1, 40), 45 | solver=rk4_step 46 | ): 47 | self._t = t 48 | self._model = model 49 | self._solver = solver 50 | self._deltas_t = t[1:] - t[:-1] 51 | 52 | def forward(self, inputs: tf.Tensor, return_states: Optional[str] = None): 53 | 54 | def _forward_dynamics(_state): 55 | """Used in solver _state == (time, tensor)""" 56 | return [1.0, self._model(inputs=_state)] 57 | 58 | states = [] 59 | 60 | def _append_state(_state): 61 | tensors = _state[1] 62 | if return_states == "numpy": 63 | states.append(tensors.numpy()) 64 | elif return_states == "tf": 65 | states.append(tensors) 66 | 67 | with tf.name_scope("forward"): 68 | t0 = tf.to_float(self._t[0]) 69 | state = [t0, inputs] 70 | _append_state(state) 71 | for dt in self._deltas_t: 72 | state = self._solver( 73 | func=_forward_dynamics, dt=tf.to_float(dt), state=state 74 | ) 75 | _append_state(state) 76 | 77 | outputs = state[1] 78 | if return_states: 79 | return outputs, states 80 | return outputs 81 | 82 | def _backward_dynamics(self, state): 83 | t = state[0] 84 | ht = state[1] 85 | at = -state[2] 86 | 87 | with tf.GradientTape() as g: 88 | g.watch(ht) 89 | ht_new = self._model(inputs=[t, ht]) 90 | 91 | gradients = g.gradient( 92 | target=ht_new, sources=[ht] + self._model.weights, 93 | output_gradients=at 94 | ) 95 | 96 | return [1.0, ht_new, *gradients] 97 | 98 | def backward(self, outputs: tf.Tensor, 99 | output_gradients: Optional[tf.Tensor] = None): 100 | 101 | with tf.name_scope("backward"): 102 | grad_weights = [tf.zeros_like(w) for w in self._model.weights] 103 | t0 = tf.to_float(self._t[-1]) 104 | 105 | if output_gradients is None: 106 | output_gradients = tf.zeros_like(outputs) 107 | 108 | state = [t0, outputs, output_gradients, *grad_weights] 109 | for dt in self._deltas_t[::-1]: 110 | state = self._solver( 111 | self._backward_dynamics, dt=-tf.to_float(dt), state=state 112 | ) 113 | 114 | inputs = state[1] 115 | dLdInputs = state[2] 116 | dLdWeights = state[3:] 117 | return inputs, dLdInputs, dLdWeights 118 | -------------------------------------------------------------------------------- /Hybrid_Sine_NN/plotting.py: -------------------------------------------------------------------------------- 1 | import numpy as np 2 | import matplotlib as mpl 3 | #mpl.use('pgf') 4 | 5 | def figsize(scale, nplots = 1): 6 | fig_width_pt = 390.0 # Get this from LaTeX using \the\textwidth 7 | inches_per_pt = 1.0/72.27 # Convert pt to inch 8 | golden_mean = (np.sqrt(5.0)-1.0)/2.0 # Aesthetic ratio (you could change this) 9 | fig_width = fig_width_pt*inches_per_pt*scale # width in inches 10 | fig_height = nplots*fig_width*golden_mean # height in inches 11 | fig_size = [fig_width,fig_height] 12 | return fig_size 13 | 14 | pgf_with_latex = { # setup matplotlib to use latex for output 15 | "pgf.texsystem": "pdflatex", # change this if using xetex or lautex 16 | "text.usetex": True, # use LaTeX to write all text 17 | "font.family": "serif", 18 | "font.serif": [], # blank entries should cause plots to inherit fonts from the document 19 | "font.sans-serif": [], 20 | "font.monospace": [], 21 | "axes.labelsize": 10, # LaTeX default is 10pt font. 22 | "font.size": 10, 23 | "legend.fontsize": 8, # Make the legend/label fonts a little smaller 24 | "xtick.labelsize": 8, 25 | "ytick.labelsize": 8, 26 | "figure.figsize": figsize(1.0), # default fig size of 0.9 textwidth 27 | "pgf.preamble": [ 28 | r"\usepackage[utf8x]{inputenc}", # use utf8 fonts becasue your computer can handle it :) 29 | r"\usepackage[T1]{fontenc}", # plots will be generated using this preamble 30 | ] 31 | } 32 | mpl.rcParams.update(pgf_with_latex) 33 | 34 | import matplotlib.pyplot as plt 35 | 36 | # I make my own newfig and savefig functions 37 | def newfig(width, nplots = 1): 38 | fig = plt.figure(figsize=figsize(width, nplots)) 39 | ax = fig.add_subplot(111) 40 | return fig, ax 41 | 42 | def savefig(filename, crop = True): 43 | if crop == True: 44 | # plt.savefig('{}.pgf'.format(filename), bbox_inches='tight', pad_inches=0) 45 | plt.savefig('{}.pdf'.format(filename), bbox_inches='tight', pad_inches=0) 46 | plt.savefig('{}.eps'.format(filename), bbox_inches='tight', pad_inches=0) 47 | else: 48 | # plt.savefig('{}.pgf'.format(filename)) 49 | plt.savefig('{}.pdf'.format(filename)) 50 | plt.savefig('{}.eps'.format(filename)) 51 | 52 | ## Simple plot 53 | #fig, ax = newfig(1.0) 54 | # 55 | #def ema(y, a): 56 | # s = [] 57 | # s.append(y[0]) 58 | # for t in range(1, len(y)): 59 | # s.append(a * y[t] + (1-a) * s[t-1]) 60 | # return np.array(s) 61 | # 62 | #y = [0]*200 63 | #y.extend([20]*(1000-len(y))) 64 | #s = ema(y, 0.01) 65 | # 66 | #ax.plot(s) 67 | #ax.set_xlabel('X Label') 68 | #ax.set_ylabel('EMA') 69 | # 70 | #savefig('ema') -------------------------------------------------------------------------------- /Hybrid_Sine_NN/pre_train.py: -------------------------------------------------------------------------------- 1 | import os 2 | os.environ['KMP_DUPLICATE_LIB_OK']='True' 3 | 4 | import numpy as np 5 | import numpy.random as npr 6 | import tensorflow as tf 7 | import matplotlib.pyplot as plt 8 | plt.switch_backend('agg') 9 | import tensorflow.contrib.eager as tfe 10 | from tensorflow.keras import layers, initializers 11 | from scipy.integrate import odeint 12 | 13 | 14 | keras = tf.keras 15 | tf.enable_eager_execution() 16 | 17 | from neural_ode import NeuralODE 18 | 19 | np.random.seed(1234) 20 | tf.set_random_seed(1234) 21 | 22 | 23 | if __name__ == "__main__": 24 | 25 | def Damped_pendulum(z, t, alpha, beta): 26 | x, y = z 27 | dzdt = [y, - alpha * y - beta * np.sin(x)] 28 | return dzdt 29 | 30 | def plot_spiral(trajectories, width = 1.): 31 | x = trajectories[:,0] 32 | y = trajectories[:,1] 33 | plt.plot(x, y, linewidth=width) 34 | 35 | 36 | data_size = 8001 37 | batch_time = 200 38 | batch_size = 1000 39 | 40 | 41 | alpha = 0.2 42 | beta = 8.91 43 | 44 | 45 | z0 = [-1.193, -3.876] 46 | true_y0 = tf.to_float([[-1.193, -3.876]]) 47 | t_grid = np.linspace(0, 20, data_size) 48 | 49 | true_yy = odeint(Damped_pendulum, z0, t_grid, args=(alpha, beta)) 50 | print(true_yy.shape) 51 | 52 | true_y = true_yy + 0.0 * np.random.randn(true_yy.shape[0], true_yy.shape[1]) 53 | 54 | np.save('true_y', true_y) 55 | 56 | 57 | sigma_normal1 = np.std(true_y[:,0:1]) 58 | sigma_normal2 = np.std(true_y[:,1:2]) 59 | 60 | true_y[:,0:1] = true_y[:,0:1]/ sigma_normal1 61 | true_y[:,1:2] = true_y[:,1:2]/ sigma_normal2 62 | 63 | true_y0 = tf.to_float([[-1.193/sigma_normal1, -3.876/sigma_normal2]]) 64 | 65 | 66 | def get_batch(): 67 | """Returns initial point and last point over sampled frament of trajectory""" 68 | starts = np.random.choice(np.arange(data_size - batch_time - 1, dtype=np.int64), batch_size, replace=False) 69 | batch_y0 = true_y[starts] 70 | batch_yN = true_y[starts + batch_time] 71 | return tf.to_float(batch_y0), tf.to_float(batch_yN) 72 | 73 | 74 | 75 | layers = [1, 20, 20, 1] 76 | def get_para_num(): 77 | L = len(layers) 78 | para_num = 0 79 | for k in range(L-1): 80 | para_num = para_num + layers[k] * layers[k+1] + layers[k+1] 81 | return para_num 82 | 83 | 84 | num_param = 9 85 | num_param_NN = get_para_num() 86 | 87 | 88 | para_num = num_param + num_param_NN 89 | 90 | print("total parameters", para_num, "NN parameters", num_param_NN) 91 | 92 | 93 | t0 = t_grid[:batch_time][0] 94 | t1 = t_grid[:batch_time][-1] 95 | t_in = np.linspace(t0, t1, 20) 96 | 97 | 98 | batch_y0, batch_yN = get_batch() 99 | 100 | np.save('batch_y0', batch_y0) 101 | np.save('batch_yN', batch_yN) 102 | 103 | 104 | ######################################### 105 | ########## precondition start ########### 106 | ######################################### 107 | niters_pre = 2000 108 | 109 | class ODEModel_pre(tf.keras.Model): 110 | def __init__(self): 111 | super(ODEModel_pre, self).__init__() 112 | self.Weights = tfe.Variable(tf.random_normal([num_param + num_param_NN, 1], dtype=tf.float32)*0.01, dtype=tf.float32) 113 | 114 | def forward_pass(self, H, layers, W): 115 | num_layers = len(layers) 116 | W_seq = W 117 | 118 | for k in range(0,num_layers-2): 119 | W = W_seq[0:layers[k] * layers[k+1]] 120 | W = tf.reshape(W, (layers[k], layers[k+1])) 121 | W_seq = W_seq[layers[k] * layers[k+1]:] 122 | b = W_seq[0:layers[k+1]] 123 | b = tf.reshape(b, (1, layers[k+1])) 124 | W_seq = W_seq[layers[k+1]:] 125 | H = tf.tanh(tf.add(tf.matmul(H, W), b)) 126 | 127 | W = W_seq[0:layers[num_layers-2] * layers[num_layers-1]] 128 | W = tf.reshape(W, (layers[num_layers-2], layers[num_layers-1])) 129 | W_seq = W_seq[layers[num_layers-2] * layers[num_layers-1]:] 130 | b = W_seq[0:layers[num_layers-1]] 131 | b = tf.reshape(b, (1, layers[num_layers-1])) 132 | H = tf.add(tf.matmul(H, W), b) 133 | return H 134 | 135 | def call(self, inputs, **kwargs): 136 | t, y = inputs 137 | h = y 138 | h1 = h[:,0:1] 139 | h2 = h[:,1:2] 140 | 141 | p11 = self.Weights[0] 142 | p12 = self.Weights[1] 143 | p13 = self.Weights[2] 144 | p14 = self.Weights[3] 145 | p15 = self.Weights[4] 146 | p16 = self.Weights[5] 147 | p23 = self.Weights[6] 148 | p25 = self.Weights[7] 149 | p26 = self.Weights[8] 150 | 151 | w_NN = self.Weights[num_param:,:] 152 | 153 | NN_out = self.forward_pass(sigma_normal1 * h1, layers, w_NN) 154 | 155 | h_out1 = p11 + sigma_normal1 * p12 *h1 + sigma_normal2 * p13 *h2 + sigma_normal1**2 * p14*h1**2\ 156 | + sigma_normal1 * sigma_normal2 * p15*h1*h2 + sigma_normal2**2 * p16*h2**2 157 | h_out2 = sigma_normal2 * p23 *h2 \ 158 | + sigma_normal1 * sigma_normal2 * p25*h1*h2 + sigma_normal2**2 * p26*h2**2 + NN_out 159 | h_out = tf.concat([h_out1/sigma_normal1, h_out2/sigma_normal2], 1) 160 | return h_out 161 | 162 | 163 | model_pre = ODEModel_pre() 164 | neural_ode_pre = NeuralODE(model_pre, t_in) 165 | optimizer = tf.train.AdamOptimizer(1e-1) 166 | 167 | 168 | def compute_gradients_and_update_pre(batch_y0, batch_yN): 169 | """Takes start positions (x0, y0) and final positions (xN, yN)""" 170 | pred_y = neural_ode_pre.forward(batch_y0) 171 | with tf.GradientTape() as g_pre: 172 | g_pre.watch(pred_y) 173 | loss = tf.reduce_mean((pred_y - batch_yN)**2) + tf.reduce_sum(tf.abs(model_pre.trainable_weights[0])) 174 | 175 | dLoss = g_pre.gradient(loss, pred_y) 176 | h_start, dfdh0, dWeights = neural_ode_pre.backward(pred_y, dLoss) 177 | optimizer.apply_gradients(zip(dWeights, model_pre.weights)) 178 | return loss, dWeights 179 | 180 | # Compile EAGER graph to static (this will be much faster) 181 | compute_gradients_and_update_pre = tfe.defun(compute_gradients_and_update_pre) 182 | 183 | parameters_pre = np.zeros((para_num, 1)) 184 | 185 | for step in range(niters_pre): 186 | print(step) 187 | loss, dWeights = compute_gradients_and_update_pre(batch_y0, batch_yN) 188 | parameters_pre = model_pre.trainable_weights[0].numpy() 189 | 190 | print(parameters_pre) 191 | 192 | np.save('parameters_pre', parameters_pre) 193 | 194 | ######################################### 195 | ########## precondition end ############# 196 | ######################################### 197 | 198 | -------------------------------------------------------------------------------- /LV_noisy/create_nice_plot.py: -------------------------------------------------------------------------------- 1 | import os 2 | os.environ['KMP_DUPLICATE_LIB_OK']='True' 3 | 4 | import numpy as np 5 | import numpy.random as npr 6 | import tensorflow as tf 7 | import matplotlib.pyplot as plt 8 | plt.switch_backend('agg') 9 | 10 | from scipy.integrate import odeint 11 | from scipy import stats 12 | 13 | from plotting import newfig, savefig 14 | 15 | np.random.seed(1234) 16 | tf.set_random_seed(1234) 17 | 18 | if __name__ == "__main__": 19 | 20 | def VP(z, t, alpha, beta, gamma, sigma): 21 | x, y = z 22 | dzdt = [alpha * x - beta * x * y, - gamma * y + sigma * x*y] 23 | return dzdt 24 | 25 | Order = 3 26 | def VP_BNODE(z, t, W, Order): 27 | x, y = z 28 | alpha = W[0] 29 | beta = W[1] 30 | gamma = W[2] 31 | sigma = W[3] 32 | dzdt = [alpha * x + beta * x * y, gamma * y + sigma * x*y] 33 | return dzdt 34 | 35 | 36 | data_size = 16001 37 | batch_time = 320 38 | batch_size = 1000 39 | 40 | data_size_true = 2000 41 | 42 | 43 | N_samples = 100 44 | N_total = 500 45 | 46 | 47 | alpha = 1. 48 | beta = 0.1 49 | gamma = 1.5 50 | sigma = 0.75 51 | 52 | 53 | z0 = [5., 5.] 54 | t_grid_train = np.linspace(0, 25, data_size) 55 | t_grid_true = np.linspace(0, 45, data_size_true) 56 | 57 | y_train = odeint(VP, z0, t_grid_train, args=(alpha, beta, gamma, sigma)) 58 | idx = np.random.choice(np.arange(data_size - batch_time - 1, dtype=np.int64), batch_size, replace=False) 59 | y_train = y_train[idx] 60 | 61 | 62 | y_train = y_train + 0.0 * np.random.randn(y_train.shape[0], y_train.shape[1]) 63 | t_grid_train = t_grid_train[idx] 64 | 65 | 66 | sigma_normal1 = np.std(y_train[:,0:1]) 67 | sigma_normal2 = np.std(y_train[:,1:2]) 68 | 69 | noise_level = 0.03 70 | 71 | y_train[:,0:1] = y_train[:,0:1]/ sigma_normal1 + noise_level * np.random.randn(y_train[:,0:1].shape[0], y_train[:,0:1].shape[1]) 72 | y_train[:,1:2] = y_train[:,1:2]/ sigma_normal2 + noise_level * np.random.randn(y_train[:,1:2].shape[0], y_train[:,1:2].shape[1]) 73 | 74 | y_train[:,0:1] = y_train[:,0:1] * sigma_normal1 75 | y_train[:,1:2] = y_train[:,1:2] * sigma_normal2 76 | 77 | sigma_normal = np.asarray([sigma_normal1, sigma_normal2]) 78 | 79 | 80 | y_true = odeint(VP, z0, t_grid_true, args=(alpha, beta, gamma, sigma)) 81 | 82 | num_dim = 2 83 | parameters = np.load("parameters.npy") 84 | precision = np.load("loggammalist.npy") 85 | loglikelihood = np.load("loglikelihood.npy") 86 | precision = np.exp(precision) 87 | print("precision", precision) 88 | print(parameters.shape) 89 | loglikelihood = loglikelihood[-N_total:] 90 | num_samples = parameters.shape[0] 91 | length_dict = parameters.shape[1] 92 | 93 | idx_MAP = np.argmin(loglikelihood) 94 | MAP = parameters[idx_MAP, :] 95 | 96 | 97 | y_MAP = odeint(VP_BNODE, z0, t_grid_true, args=(MAP, Order)) 98 | 99 | mu_pred = y_MAP 100 | 101 | y_BNODE = np.zeros((t_grid_true.shape[0], num_dim, N_samples)) 102 | 103 | for k in range(N_samples): 104 | print(k) 105 | idx_1 = np.random.randint(N_total) 106 | idx_2 = np.random.randint(N_total) 107 | W_sample = parameters[-idx_1, :] 108 | precision_here = precision[-idx_2] * num_dim 109 | y_BNODE[:,:,k] = odeint(VP_BNODE, z0, t_grid_true, args=(W_sample, Order)) 110 | Sigma_data = np.ones_like(mu_pred) / np.sqrt(precision_here) 111 | y_BNODE[:,:,k] = y_BNODE[:,:,k] + Sigma_data * np.random.normal() 112 | 113 | 114 | mu_pred = np.mean(y_BNODE, axis = 2) 115 | Sigma_pred = np.var(y_BNODE, axis = 2) 116 | 117 | 118 | plt.rcParams.update(plt.rcParamsDefault) 119 | plt.rc('font', family='serif') 120 | plt.rcParams.update({'font.size': 16, 121 | 'lines.linewidth': 2, 122 | 'axes.labelsize': 20, # fontsize for x and y labels (was 10) 123 | 'axes.titlesize': 20, 124 | 'xtick.labelsize': 16, 125 | 'ytick.labelsize': 16, 126 | 'legend.fontsize': 20, 127 | 'axes.linewidth': 2, 128 | "pgf.texsystem": "pdflatex", # change this if using xetex or lautex 129 | "text.usetex": True, # use LaTeX to write all text 130 | }) 131 | 132 | 133 | plt.figure(1, figsize=(6,6)) 134 | plt.xticks(fontsize=14) 135 | plt.yticks(fontsize=14) 136 | plt.plot(y_train[:,0], y_train[:,1], 'ro', label = "Training data") 137 | plt.plot(y_true[:,0], y_true[:,1], 'b-', label = "Exact Trajectory") 138 | plt.xlabel('$x_1$',fontsize=18) 139 | plt.ylabel('$x_2$',fontsize=18) 140 | plt.legend(loc='upper right', frameon=False, prop={'size': 13}) 141 | plt.savefig('./Training_data.png', dpi = 300) 142 | 143 | 144 | plt.figure(4, figsize=(12,6.5)) 145 | plt.xticks(fontsize=22) 146 | plt.yticks(fontsize=22) 147 | plt.plot(t_grid_true, y_BNODE[:,0,0], '-', color = "gray", alpha = 0.8, label = "Sample Trajectory of $x_1(t)$", linewidth = 0.5) 148 | plt.plot(t_grid_true, y_BNODE[:,0,:], '-', color = "gray", alpha = 0.8, linewidth = 0.5) 149 | plt.plot(t_grid_train, y_train[:,0], 'ro', label = "Training data of $x_1(t)$") 150 | plt.plot(t_grid_true, y_true[:,0], 'r-', label = "True Trajectory of $x_1(t)$") 151 | plt.plot(t_grid_true, y_MAP[:,0], 'b--', label = "MAP Trajectory of $x_1(t)$") plt.xlabel('$t$',fontsize=26) 152 | plt.ylabel('$x_1(t)$',fontsize=26) 153 | plt.ylim((-1.1, 9.5)) 154 | plt.legend(loc='upper right', frameon=False, prop={'size': 20}) 155 | plt.savefig('./BNODE_Prediction_x1.png', dpi = 300) 156 | 157 | 158 | plt.figure(5, figsize=(12,6.5)) 159 | plt.xticks(fontsize=22) 160 | plt.yticks(fontsize=22) 161 | plt.plot(t_grid_true, y_BNODE[:,1,0], '-', color = "gray", alpha = 0.8, label = "Sample Trajectory of $x_2(t)$", linewidth = 0.5) 162 | plt.plot(t_grid_true, y_BNODE[:,1,:], '-', color = "gray", alpha = 0.8, linewidth = 0.5) 163 | plt.plot(t_grid_train, y_train[:,1], 'ro', label = "Training data of $x_2(t)$") 164 | plt.plot(t_grid_true, y_true[:,1], 'r-', label = "True Trajectory of $x_2(t)$") 165 | plt.plot(t_grid_true, y_MAP[:,1], 'b--', label = "MAP Trajectory of $x_2(t)$") 166 | plt.xlabel('$t$',fontsize=26) 167 | plt.ylabel('$x_2(t)$',fontsize=26) 168 | plt.ylim((-5.1, 55.5)) 169 | plt.legend(loc='upper right', frameon=False, prop={'size': 20}) 170 | plt.savefig('./BNODE_Prediction_x2.png', dpi = 300) 171 | 172 | 173 | 174 | 175 | 176 | -------------------------------------------------------------------------------- /LV_noisy/example.py: -------------------------------------------------------------------------------- 1 | import os 2 | os.environ['KMP_DUPLICATE_LIB_OK']='True' 3 | 4 | import numpy as np 5 | import numpy.random as npr 6 | import tensorflow as tf 7 | import matplotlib.pyplot as plt 8 | import matplotlib as mpl 9 | from mpl_toolkits.mplot3d import Axes3D 10 | plt.switch_backend('agg') 11 | import tensorflow.contrib.eager as tfe 12 | from tensorflow.keras import layers, initializers 13 | from scipy.integrate import odeint 14 | 15 | 16 | keras = tf.keras 17 | tf.enable_eager_execution() 18 | 19 | from neural_ode import NeuralODE 20 | 21 | np.random.seed(1234) 22 | tf.set_random_seed(1234) 23 | 24 | if __name__ == "__main__": 25 | 26 | def VP(z, t, alpha, beta, gamma, sigma): 27 | x, y = z 28 | dzdt = [alpha * x - beta * x * y, - gamma * y + sigma * x*y] 29 | return dzdt 30 | 31 | 32 | def plot_spiral(trajectories, width = 1.): 33 | x = trajectories[:,0] 34 | y = trajectories[:,1] 35 | plt.plot(x, y, linewidth = width) 36 | 37 | 38 | data_size = 16001 39 | batch_time = 320 40 | niters = 5000 41 | batch_size = 1000 42 | 43 | alpha = 1. 44 | beta = 0.1 45 | gamma = 1.5 46 | sigma = 0.75 47 | 48 | t_grid = np.linspace(0, 25, data_size) 49 | z0 = [5., 5.] 50 | true_yy = odeint(VP, z0, t_grid, args=(alpha, beta, gamma, sigma)) 51 | 52 | 53 | 54 | true_y = true_yy 55 | 56 | # normalizing data and system 57 | sigma_x = np.std(true_yy[:,0:1]) 58 | sigma_y = np.std(true_yy[:,1:2]) 59 | 60 | noise_level = 0.03 61 | sigma_normal = max(sigma_x, sigma_y) 62 | 63 | 64 | true_y[:,0:1] = true_y[:,0:1]/ sigma_x + noise_level * np.random.randn(true_y[:,0:1].shape[0], true_y[:,0:1].shape[1]) 65 | true_y[:,1:2] = true_y[:,1:2]/ sigma_y + noise_level * np.random.randn(true_y[:,1:2].shape[0], true_y[:,1:2].shape[1]) 66 | 67 | 68 | 69 | 70 | def get_batch(): 71 | """Returns initial point and last point over sampled frament of trajectory""" 72 | starts = np.random.choice(np.arange(data_size - batch_time - 1, dtype=np.int64), batch_size, replace=False) 73 | batch_y0 = true_y[starts] 74 | batch_yN = true_y[starts + batch_time] 75 | return tf.to_float(batch_y0), tf.to_float(batch_yN) 76 | 77 | 78 | num_param = 4 79 | para_num = num_param 80 | 81 | t0 = t_grid[:batch_time][0] 82 | t1 = t_grid[:batch_time][-1] 83 | t_in = np.linspace(t0, t1, 20) 84 | 85 | 86 | batch_y0, batch_yN = get_batch() 87 | 88 | 89 | ######################################### 90 | ########## precondition start ########### 91 | ######################################### 92 | niters_pre = 500 93 | 94 | class ODEModel_pre(tf.keras.Model): 95 | def __init__(self): 96 | super(ODEModel_pre, self).__init__() 97 | self.Weights = tfe.Variable(tf.random_normal([num_param, 1], dtype=tf.float32)*0.01, dtype=tf.float32) 98 | 99 | 100 | def call(self, inputs, **kwargs): 101 | t, y = inputs 102 | h = y 103 | h1 = h[:,0:1] 104 | h2 = h[:,1:2] 105 | 106 | p1 = self.Weights[0] 107 | p2 = self.Weights[1] 108 | p3 = self.Weights[2] 109 | p4 = self.Weights[3] 110 | 111 | h_out1 = p1 * h1 + sigma_y * p2 * h2*h1 112 | h_out2 = p3 * h2 + sigma_x * p4 * h2*h1 113 | h_out = tf.concat([h_out1, h_out2], 1) 114 | return h_out 115 | 116 | 117 | model_pre = ODEModel_pre() 118 | neural_ode_pre = NeuralODE(model_pre, t_in) 119 | optimizer = tf.train.AdamOptimizer(3e-2) 120 | 121 | 122 | def compute_gradients_and_update_pre(batch_y0, batch_yN): 123 | """Takes start positions (x0, y0) and final positions (xN, yN)""" 124 | pred_y = neural_ode_pre.forward(batch_y0) 125 | with tf.GradientTape() as g_pre: 126 | g_pre.watch(pred_y) 127 | loss = tf.reduce_mean((pred_y - batch_yN)**2) + tf.reduce_sum(tf.abs(model_pre.trainable_weights[0])) 128 | 129 | dLoss = g_pre.gradient(loss, pred_y) 130 | h_start, dfdh0, dWeights = neural_ode_pre.backward(pred_y, dLoss) 131 | optimizer.apply_gradients(zip(dWeights, model_pre.weights)) 132 | return loss, dWeights 133 | 134 | # Compile EAGER graph to static (this will be much faster) 135 | compute_gradients_and_update_pre = tfe.defun(compute_gradients_and_update_pre) 136 | 137 | parameters_pre = np.zeros((para_num, 1)) 138 | 139 | 140 | for step in range(niters_pre): 141 | print(step) 142 | loss, dWeights = compute_gradients_and_update_pre(batch_y0, batch_yN) 143 | parameters_pre = model_pre.trainable_weights[0].numpy() 144 | 145 | 146 | print(parameters_pre) 147 | 148 | ######################################### 149 | ########## precondition end ############# 150 | ######################################### 151 | 152 | initial_weight = parameters_pre 153 | print(initial_weight.shape, "here") 154 | 155 | 156 | class ODEModel(tf.keras.Model): 157 | def __init__(self): 158 | super(ODEModel, self).__init__() 159 | self.Weights = tfe.Variable(tf.random_normal([num_param, 1], dtype=tf.float32)*0.01, dtype=tf.float32) 160 | 161 | 162 | def call(self, inputs, **kwargs): 163 | t, y = inputs 164 | h = y 165 | h1 = h[:,0:1] 166 | h2 = h[:,1:2] 167 | 168 | p1 = self.Weights[0] 169 | p2 = self.Weights[1] 170 | p3 = self.Weights[2] 171 | p4 = self.Weights[3] 172 | 173 | h_out1 = p1 * h1 + sigma_y * p2 * h2*h1 174 | h_out2 = p3 * h2 + sigma_x * p4 * h2*h1 175 | h_out = tf.concat([h_out1, h_out2], 1) 176 | return h_out 177 | 178 | 179 | model = ODEModel() 180 | neural_ode = NeuralODE(model, t = t_in) 181 | 182 | def compute_gradients_and_update(batch_y0, batch_yN): 183 | """Takes start positions (x0, y0) and final positions (xN, yN)""" 184 | pred_y = neural_ode.forward(batch_y0) 185 | with tf.GradientTape() as g: 186 | g.watch(pred_y) 187 | loss = tf.reduce_sum((pred_y - batch_yN)**2) 188 | 189 | dLoss = g.gradient(loss, pred_y) 190 | h_start, dfdh0, dWeights = neural_ode.backward(pred_y, dLoss) 191 | 192 | return loss, dWeights 193 | 194 | # Compile EAGER graph to static (this will be much faster) 195 | compute_gradients_and_update = tfe.defun(compute_gradients_and_update) 196 | 197 | 198 | # function to compute the kinetic energy 199 | def kinetic_energy(V, loggamma_v, loglambda_v): 200 | q = (np.sum(-V**2) - loggamma_v**2 - loglambda_v**2)/2.0 201 | return q 202 | 203 | def compute_gradient_param(dWeights, loggamma, loglambda, batch_size, para_num): 204 | WW = model.trainable_weights[0].numpy() 205 | dWeights = np.exp(loggamma)/2.0 * dWeights + np.exp(loglambda) * np.sign(WW) 206 | return dWeights 207 | 208 | def compute_gradient_hyper(loss, weights, loggamma, loglambda, batch_size, para_num): 209 | grad_loggamma = np.exp(loggamma) * (loss/2.0 + 1.0) - (batch_size/2.0 + 1.0) 210 | grad_loglambda = np.exp(loglambda) * (np.sum(np.abs(weights)) + 1.0) - (para_num + 1.0) 211 | 212 | return grad_loggamma, grad_loglambda 213 | 214 | def compute_Hamiltonian(loss, weights, loggamma, loglambda, batch_size, para_num): 215 | H = np.exp(loggamma)*(loss/2.0 + 1.0) + np.exp(loglambda)*(np.sum(np.abs(weights)) + 1.0)\ 216 | - (batch_size/2.0 + 1.0) * loggamma - (para_num + 1.0) * loglambda 217 | return H 218 | 219 | def leap_frog(v_in, w_in, loggamma_in, loglambda_in, loggamma_v_in, loglambda_v_in): 220 | # assign weights from the previous step to the model 221 | model.trainable_weights[0].assign(w_in) 222 | v_new = v_in 223 | loggamma_v_new = loggamma_v_in 224 | loglambda_v_new = loglambda_v_in 225 | 226 | loggamma_new = loggamma_in 227 | loglambda_new = loglambda_in 228 | w_new = w_in 229 | 230 | for m in range(L): 231 | loss, dWeights = compute_gradients_and_update(batch_y0, batch_yN) # evaluate the gradient 232 | 233 | dWeights = np.asarray(dWeights[0]) # make the gradient to be numpy array 234 | dWeights = compute_gradient_param(dWeights, loggamma_new, loglambda_new, batch_size, para_num) 235 | grad_loggamma, grad_loglambda = compute_gradient_hyper(loss, w_new, loggamma_new, loglambda_new, batch_size, para_num) 236 | 237 | loggamma_v_new = loggamma_v_new - epsilon/2*grad_loggamma 238 | loglambda_v_new = loglambda_v_new - epsilon/2*grad_loglambda 239 | v_new = v_new - epsilon/2*(dWeights) 240 | w_new = model.trainable_weights[0].numpy() + epsilon * v_new 241 | model.trainable_weights[0].assign(w_new) 242 | loggamma_new = loggamma_new + epsilon * loggamma_v_new 243 | loglambda_new = loglambda_new + epsilon * loglambda_v_new 244 | 245 | # Second half of the leap frog 246 | loss, dWeights = compute_gradients_and_update(batch_y0, batch_yN) 247 | dWeights = np.asarray(dWeights[0]) 248 | dWeights = compute_gradient_param(dWeights, loggamma_new, loglambda_new, batch_size, para_num) 249 | grad_loggamma, grad_loglambda = compute_gradient_hyper(loss, w_new, loggamma_new, loglambda_new, batch_size, para_num) 250 | 251 | v_new = v_new - epsilon/2*(dWeights) 252 | loggamma_v_new = loggamma_v_new - epsilon/2*grad_loggamma 253 | loglambda_v_new = loglambda_v_new - epsilon/2*grad_loglambda 254 | 255 | print(np.exp(loggamma_new)) 256 | print(np.exp(loglambda_new)) 257 | 258 | return v_new, w_new, loggamma_new, loglambda_new, loggamma_v_new, loglambda_v_new 259 | 260 | 261 | 262 | neural_ode_test = NeuralODE(model, t=t_grid[0:data_size:20]) 263 | parameters = np.zeros((niters, para_num)) # book keeping the parameters 264 | loggammalist = np.zeros((niters, 1)) # book keeping the loggamma 265 | loglambdalist = np.zeros((niters, 1)) # book keeping the loggamma 266 | loglikelihood = np.zeros((niters, 1)) # book keeping the loggamma 267 | L = 10 # leap frog step number 268 | epsilon = 0.001 # leap frog step size 269 | epsilon_max = 0.0002 # max 0.001 270 | epsilon_min = 0.0002 # max 0.001 271 | 272 | 273 | def compute_epsilon(step): 274 | coefficient = np.log(epsilon_max/epsilon_min) 275 | return epsilon_max * np.exp( - step * coefficient / niters) 276 | 277 | 278 | # initial weight 279 | w_temp = initial_weight 280 | print("initial_w", w_temp) 281 | loggamma_temp = 4. + np.random.normal() 282 | loglambda_temp = np.random.normal() 283 | 284 | model.trainable_weights[0].assign(w_temp) 285 | loss_original, _ = compute_gradients_and_update(batch_y0, batch_yN) # compute the initial Hamiltonian 286 | 287 | 288 | loggamma_temp = np.log(batch_size / loss_original) 289 | print("This is initial guess", loggamma_temp, "with loss", loss_original) 290 | if loggamma_temp > 6.: 291 | loggamma_temp = 6. 292 | epsilon_max = 0.0002 293 | epsilon_min = 0.0002 294 | 295 | 296 | # training steps 297 | for step in range(niters): 298 | 299 | epsilon = compute_epsilon(step) 300 | 301 | print(step) 302 | v_initial = np.random.randn(para_num, 1) # initialize the velocity 303 | loggamma_v_initial = np.random.normal() 304 | loglambda_v_initial = np.random.normal() 305 | 306 | loss_initial, _ = compute_gradients_and_update(batch_y0, batch_yN) # compute the initial Hamiltonian 307 | loss_initial = compute_Hamiltonian(loss_initial, w_temp, loggamma_temp, loglambda_temp, batch_size, para_num) 308 | 309 | v_new, w_new, loggamma_new, loglambda_new, loggamma_v_new, loglambda_v_new = \ 310 | leap_frog(v_initial, w_temp, loggamma_temp, loglambda_temp, loggamma_v_initial, loglambda_v_initial) 311 | 312 | # compute the final Hamiltonian 313 | loss_finial, _ = compute_gradients_and_update(batch_y0, batch_yN) 314 | loss_finial = compute_Hamiltonian(loss_finial, w_new, loggamma_new, loglambda_new, batch_size, para_num) 315 | 316 | # making decisions 317 | p_temp = np.exp(-loss_finial + loss_initial + \ 318 | kinetic_energy(v_new, loggamma_v_new, loglambda_v_new) - kinetic_energy(v_initial, loggamma_v_initial, loglambda_v_initial)) 319 | 320 | p = min(1, p_temp) 321 | p_decision = np.random.uniform() 322 | if p > p_decision: 323 | parameters[step:step+1, :] = np.transpose(w_new) 324 | w_temp = w_new 325 | loggammalist[step, 0] = loggamma_new 326 | loglambdalist[step, 0] = loglambda_new 327 | loglikelihood[step, 0] = loss_finial 328 | loggamma_temp = loggamma_new 329 | loglambda_temp = loglambda_new 330 | else: 331 | parameters[step:step+1, :] = np.transpose(w_temp) 332 | model.trainable_weights[0].assign(w_temp) 333 | loggammalist[step, 0] = loggamma_temp 334 | loglambdalist[step, 0] = loglambda_temp 335 | loglikelihood[step, 0] = loss_initial 336 | 337 | print('probability', p) 338 | print(p > p_decision) 339 | 340 | 341 | np.save('parameters', parameters) 342 | np.save('loggammalist', loggammalist) 343 | np.save('loglikelihood', loglikelihood) 344 | -------------------------------------------------------------------------------- /LV_noisy/neural_ode.py: -------------------------------------------------------------------------------- 1 | from typing import Optional, List 2 | import numpy as np 3 | import tensorflow as tf 4 | from tensorflow.python.framework.ops import EagerTensor 5 | import tensorflow.contrib.eager as tfe 6 | 7 | keras = tf.keras 8 | 9 | 10 | def zip_map(zipped, update_op): 11 | return [update_op(*elems) for elems in zipped] 12 | 13 | 14 | def euler_update(h_list, dh_list, dt): 15 | return zip_map(zip(h_list, dh_list), lambda h, dh: h + tf.cast(dt, h.dtype) * dh) 16 | 17 | 18 | def euler_step(func, dt, state): 19 | return euler_update(state, func(state), dt) 20 | 21 | 22 | def rk2_step(func, dt, state): 23 | k1 = func(state) 24 | k2 = func(euler_update(state, k1, dt)) 25 | return zip_map(zip(state, k1, k2), 26 | lambda h, dk1, dk2: h + tf.cast(dt, h.dtype) * (dk1 + dk2) / 2) 27 | 28 | 29 | def rk4_step(func, dt, state): 30 | k1 = func(state) 31 | k2 = func(euler_update(state, k1, dt / 2)) 32 | k3 = func(euler_update(state, k2, dt / 2)) 33 | k4 = func(euler_update(state, k3, dt)) 34 | 35 | return zip_map( 36 | zip(state, k1, k2, k3, k4), 37 | lambda h, dk1, dk2, dk3, dk4: h + tf.cast(dt, h.dtype) * ( 38 | dk1 + 2 * dk2 + 2 * dk3 + dk4) / 6, 39 | ) 40 | 41 | 42 | class NeuralODE: 43 | def __init__( 44 | self, model: tf.keras.Model, t=np.linspace(0, 1, 40), 45 | solver=rk4_step 46 | ): 47 | self._t = t 48 | self._model = model 49 | self._solver = solver 50 | self._deltas_t = t[1:] - t[:-1] 51 | 52 | def forward(self, inputs: tf.Tensor, return_states: Optional[str] = None): 53 | 54 | def _forward_dynamics(_state): 55 | """Used in solver _state == (time, tensor)""" 56 | return [1.0, self._model(inputs=_state)] 57 | 58 | states = [] 59 | 60 | def _append_state(_state): 61 | tensors = _state[1] 62 | if return_states == "numpy": 63 | states.append(tensors.numpy()) 64 | elif return_states == "tf": 65 | states.append(tensors) 66 | 67 | with tf.name_scope("forward"): 68 | t0 = tf.to_float(self._t[0]) 69 | state = [t0, inputs] 70 | _append_state(state) 71 | for dt in self._deltas_t: 72 | state = self._solver( 73 | func=_forward_dynamics, dt=tf.to_float(dt), state=state 74 | ) 75 | _append_state(state) 76 | 77 | outputs = state[1] 78 | if return_states: 79 | return outputs, states 80 | return outputs 81 | 82 | def _backward_dynamics(self, state): 83 | t = state[0] 84 | ht = state[1] 85 | at = -state[2] 86 | 87 | with tf.GradientTape() as g: 88 | g.watch(ht) 89 | ht_new = self._model(inputs=[t, ht]) 90 | 91 | gradients = g.gradient( 92 | target=ht_new, sources=[ht] + self._model.weights, 93 | output_gradients=at 94 | ) 95 | 96 | return [1.0, ht_new, *gradients] 97 | 98 | def backward(self, outputs: tf.Tensor, 99 | output_gradients: Optional[tf.Tensor] = None): 100 | 101 | with tf.name_scope("backward"): 102 | grad_weights = [tf.zeros_like(w) for w in self._model.weights] 103 | t0 = tf.to_float(self._t[-1]) 104 | 105 | if output_gradients is None: 106 | output_gradients = tf.zeros_like(outputs) 107 | 108 | state = [t0, outputs, output_gradients, *grad_weights] 109 | for dt in self._deltas_t[::-1]: 110 | state = self._solver( 111 | self._backward_dynamics, dt=-tf.to_float(dt), state=state 112 | ) 113 | 114 | inputs = state[1] 115 | dLdInputs = state[2] 116 | dLdWeights = state[3:] 117 | return inputs, dLdInputs, dLdWeights 118 | -------------------------------------------------------------------------------- /LV_noisy/plotting.py: -------------------------------------------------------------------------------- 1 | import numpy as np 2 | import matplotlib as mpl 3 | #mpl.use('pgf') 4 | 5 | def figsize(scale, nplots = 1): 6 | fig_width_pt = 390.0 # Get this from LaTeX using \the\textwidth 7 | inches_per_pt = 1.0/72.27 # Convert pt to inch 8 | golden_mean = (np.sqrt(5.0)-1.0)/2.0 # Aesthetic ratio (you could change this) 9 | fig_width = fig_width_pt*inches_per_pt*scale # width in inches 10 | fig_height = nplots*fig_width*golden_mean # height in inches 11 | fig_size = [fig_width,fig_height] 12 | return fig_size 13 | 14 | pgf_with_latex = { # setup matplotlib to use latex for output 15 | "pgf.texsystem": "pdflatex", # change this if using xetex or lautex 16 | "text.usetex": True, # use LaTeX to write all text 17 | "font.family": "serif", 18 | "font.serif": [], # blank entries should cause plots to inherit fonts from the document 19 | "font.sans-serif": [], 20 | "font.monospace": [], 21 | "axes.labelsize": 10, # LaTeX default is 10pt font. 22 | "font.size": 10, 23 | "legend.fontsize": 8, # Make the legend/label fonts a little smaller 24 | "xtick.labelsize": 8, 25 | "ytick.labelsize": 8, 26 | "figure.figsize": figsize(1.0), # default fig size of 0.9 textwidth 27 | "pgf.preamble": [ 28 | r"\usepackage[utf8x]{inputenc}", # use utf8 fonts becasue your computer can handle it :) 29 | r"\usepackage[T1]{fontenc}", # plots will be generated using this preamble 30 | ] 31 | } 32 | mpl.rcParams.update(pgf_with_latex) 33 | 34 | import matplotlib.pyplot as plt 35 | 36 | # I make my own newfig and savefig functions 37 | def newfig(width, nplots = 1): 38 | fig = plt.figure(figsize=figsize(width, nplots)) 39 | ax = fig.add_subplot(111) 40 | return fig, ax 41 | 42 | def savefig(filename, crop = True): 43 | if crop == True: 44 | # plt.savefig('{}.pgf'.format(filename), bbox_inches='tight', pad_inches=0) 45 | plt.savefig('{}.pdf'.format(filename), bbox_inches='tight', pad_inches=0) 46 | plt.savefig('{}.eps'.format(filename), bbox_inches='tight', pad_inches=0) 47 | else: 48 | # plt.savefig('{}.pgf'.format(filename)) 49 | plt.savefig('{}.pdf'.format(filename)) 50 | plt.savefig('{}.eps'.format(filename)) 51 | 52 | ## Simple plot 53 | #fig, ax = newfig(1.0) 54 | # 55 | #def ema(y, a): 56 | # s = [] 57 | # s.append(y[0]) 58 | # for t in range(1, len(y)): 59 | # s.append(a * y[t] + (1-a) * s[t-1]) 60 | # return np.array(s) 61 | # 62 | #y = [0]*200 63 | #y.extend([20]*(1000-len(y))) 64 | #s = ema(y, 0.01) 65 | # 66 | #ax.plot(s) 67 | #ax.set_xlabel('X Label') 68 | #ax.set_ylabel('EMA') 69 | # 70 | #savefig('ema') -------------------------------------------------------------------------------- /README.md: -------------------------------------------------------------------------------- 1 | # Bayesian differential programming for robust systems identification under uncertainty 2 | 3 | Code and data accompanying the manuscript titled "Bayesian differential programming for robust systems identification under uncertainty", authored by Y. Yang, A. Bhouri and P. Perdikaris. Will become available after the peer review process is complete. 4 | 5 | ### Abstract 6 | 7 | This paper presents a machine learning framework for Bayesian systems identification from noisy, sparse and irregular observations of nonlinear dynamical systems. The proposed method takes advantage of recent developments in differentiable programming to propagate gradient information through ordinary differential equation solvers and perform Bayesian inference with respect to unknown model parameters using Hamiltonian Monte Carlo sampling. This allows an efficient inference of the posterior distributions over plausible models with quantified uncertainty, while the use of sparsity-promoting priors enables the discovery of interpretable and parsimonious representations for the underlying latent dynamics. A series of numerical studies is presented to demonstrate the effectiveness of the proposed methods including nonlinear oscillators, predator-prey systems and examples from systems biology. Taken all together, our findings put forth a flexible and robust workflow for data-driven model discovery under uncertainty. 8 | 9 | ### Citation 10 | 11 | @article{yang2020bayesian, 12 | title={Bayesian differential programming for robust systems identification under uncertainty}, 13 | author={Yang, Yibo and Bhouri, Mohamed Aziz and Perdikaris, Paris}, 14 | journal={arXiv preprint arXiv:2004.06843}, 15 | year={2020} 16 | } 17 | -------------------------------------------------------------------------------- /Spiral_noisy/create_nice_plot.py: -------------------------------------------------------------------------------- 1 | import os 2 | os.environ['KMP_DUPLICATE_LIB_OK']='True' 3 | 4 | import numpy as np 5 | import numpy.random as npr 6 | import tensorflow as tf 7 | import matplotlib.pyplot as plt 8 | plt.switch_backend('agg') 9 | 10 | from scipy.integrate import odeint 11 | from scipy import stats 12 | 13 | 14 | from plotting import newfig, savefig 15 | 16 | 17 | np.random.seed(1234) 18 | tf.set_random_seed(1234) 19 | 20 | if __name__ == "__main__": 21 | 22 | def Spiral(z, t, alpha, beta, gamma, theta): 23 | x, y = z 24 | dzdt = [alpha * x**3 + beta * y**3, gamma * x**3 + theta * y**3] 25 | return dzdt 26 | 27 | 28 | def Spiral_SINDy(z, t, alpha, beta, gamma, theta, addition): 29 | x, y = z 30 | dzdt = [alpha * x**3 + beta * y**3, gamma * x**3 + theta * y**3 + addition * x**2*y] 31 | return dzdt 32 | 33 | 34 | Order = 4 35 | def Spiral_BNODE(z, t, W, Order): 36 | x, y = z 37 | x = x[None,None] 38 | y = y[None,None] 39 | temp = x ** 0 40 | for m in range(1, Order): 41 | for n in range(m + 1): 42 | temp = np.concatenate([temp, x**n * y**(m - n)], 1) 43 | 44 | candidate = np.matmul(temp, W) 45 | dzdt = [candidate[0,0], candidate[0,1]] 46 | return dzdt 47 | 48 | 49 | data_size = 1000 50 | data_size_true = 2000 51 | 52 | N_samples = 100 53 | N_total = 500 54 | 55 | alpha = -0.1 56 | beta = 2.0 57 | gamma = -2.0 58 | theta = -0.1 59 | 60 | 61 | SINDy_alpha = -0.13949865 62 | SINDy_gamma = -1.974 63 | SINDy_beta = 2.09897442 64 | SINDy_theta = 0 65 | SINDy_add = -0.18839257 66 | 67 | 68 | z0 = [2., 0.] 69 | t_grid_train = np.linspace(0, 20, data_size) 70 | t_grid_true = np.linspace(0, 40, data_size_true) 71 | 72 | y_train = odeint(Spiral, z0, t_grid_train, args=(alpha, beta, gamma, theta)) 73 | y_true = odeint(Spiral, z0, t_grid_true, args=(alpha, beta, gamma, theta)) 74 | y_SINDy = odeint(Spiral_SINDy, z0, t_grid_true, args=(SINDy_alpha, SINDy_beta, SINDy_gamma, SINDy_theta, SINDy_add)) 75 | 76 | 77 | y_train = y_train + 0.02 * np.random.randn(y_train.shape[0], y_train.shape[1]) 78 | 79 | parameters = np.load("parameters.npy") 80 | precision = np.load("loggammalist.npy") 81 | loglikelihood = np.load("loglikelihood.npy") 82 | print(parameters.shape) 83 | loglikelihood = loglikelihood[-N_total:] 84 | num_samples = parameters.shape[0] 85 | length_dict = parameters.shape[1] 86 | num_dim = parameters.shape[2] 87 | precision = np.exp(precision) * num_dim 88 | print("precision", precision) 89 | 90 | 91 | 92 | idx_MAP = np.argmin(loglikelihood) 93 | MAP = parameters[idx_MAP, :, :] 94 | 95 | y_MAP = odeint(Spiral_BNODE, z0, t_grid_true, args=(MAP, Order)) 96 | 97 | mu_pred = y_MAP 98 | y_BNODE = np.zeros((t_grid_true.shape[0], num_dim, N_samples)) 99 | 100 | for k in range(N_samples): 101 | print(k) 102 | idx_1 = np.random.randint(N_total) 103 | idx_2 = np.random.randint(N_total) 104 | W_sample = parameters[-idx_1, :, :] 105 | precision_here = precision[-idx_2] 106 | y_BNODE[:,:,k] = odeint(Spiral_BNODE, z0, t_grid_true, args=(W_sample, Order)) 107 | Sigma_data = np.ones_like(mu_pred) / np.sqrt(precision_here) 108 | y_BNODE[:,:,k] = y_BNODE[:,:,k] + Sigma_data * np.random.normal() 109 | print((Sigma_data).shape) 110 | 111 | 112 | mu_pred = np.mean(y_BNODE, axis = 2) 113 | Sigma_pred = np.var(y_BNODE, axis = 2) 114 | 115 | 116 | plt.rcParams.update(plt.rcParamsDefault) 117 | plt.rc('font', family='serif') 118 | plt.rcParams.update({'font.size': 16, 119 | 'lines.linewidth': 2, 120 | 'axes.labelsize': 20, # fontsize for x and y labels (was 10) 121 | 'axes.titlesize': 20, 122 | 'xtick.labelsize': 16, 123 | 'ytick.labelsize': 16, 124 | 'legend.fontsize': 20, 125 | 'axes.linewidth': 2, 126 | "pgf.texsystem": "pdflatex", # change this if using xetex or lautex 127 | "text.usetex": True, # use LaTeX to write all text 128 | }) 129 | 130 | plt.figure(1, figsize=(6,6)) 131 | plt.xticks(fontsize=14) 132 | plt.yticks(fontsize=14) 133 | plt.plot(y_train[:,0], y_train[:,1], 'ro', label = "Training data") 134 | plt.plot(y_true[:,0], y_true[:,1], 'b-', label = "Exact Trajectory") 135 | plt.xlabel('$x_1$',fontsize=18) 136 | plt.ylabel('$x_2$',fontsize=18) 137 | plt.xlim((-2.3, 2.3)) 138 | plt.ylim((-2.3, 2.3)) 139 | plt.legend(loc='upper right', frameon=False, prop={'size': 14}) 140 | plt.savefig('./Training_data.png', dpi = 300) 141 | 142 | 143 | 144 | plt.figure(2, figsize=(12,6.5)) 145 | plt.xticks(fontsize=22) 146 | plt.yticks(fontsize=22) 147 | plt.plot(t_grid_train, y_train[:,0], 'ro', label = "Training data of $x_1(t)$") 148 | plt.plot(t_grid_true, y_true[:,0], 'r-', label = "True Trajectory of $x_1(t)$") 149 | plt.plot(t_grid_true, y_SINDy[:,0], 'b--', label = "SINDy Trajectory of $x_1(t)$") 150 | plt.xlabel('$t$',fontsize=26) 151 | plt.ylabel('$x_1(t)$',fontsize=26) 152 | plt.ylim((-2.3, 2.5)) 153 | plt.legend(loc='upper right', frameon=False, prop={'size': 20}) 154 | plt.savefig('./SINDy_Prediction_x1.png', dpi = 300) 155 | 156 | 157 | plt.figure(3, figsize=(12,6.5)) 158 | plt.xticks(fontsize=22) 159 | plt.yticks(fontsize=22) 160 | plt.plot(t_grid_train, y_train[:,1], 'ro', label = "Training data of $x_2(t)$") 161 | plt.plot(t_grid_true, y_true[:,1], 'r-', label = "True Trajectory of $x_2(t)$") 162 | plt.plot(t_grid_true, y_SINDy[:,1], 'b--', label = "SINDy Trajectory of $x_2(t)$") 163 | plt.xlabel('$t$',fontsize=26) 164 | plt.ylabel('$x_2(t)$',fontsize=26) 165 | plt.ylim((-2.3, 2.5)) 166 | plt.legend(loc='upper right', frameon=False, prop={'size': 20}) 167 | plt.savefig('./SINDy_Prediction_x2.png', dpi = 300) 168 | 169 | 170 | plt.figure(4, figsize=(12,6.5)) 171 | plt.xticks(fontsize=22) 172 | plt.yticks(fontsize=22) 173 | plt.plot(t_grid_train, y_train[:,0], 'ro', label = "Training data of $x_1(t)$") 174 | plt.plot(t_grid_true, y_true[:,0], 'r-', label = "True Trajectory of $x_1(t)$") 175 | plt.plot(t_grid_true, y_MAP[:,0], 'g--', label = "MAP Trajectory of $x_1(t)$") 176 | lower_0 = mu_pred[:,0] - 2.0*np.sqrt(Sigma_pred[:,0]) 177 | upper_0 = mu_pred[:,0] + 2.0*np.sqrt(Sigma_pred[:,0]) 178 | plt.fill_between(t_grid_true.flatten(), lower_0.flatten(), upper_0.flatten(), 179 | facecolor='orange', alpha=0.5, label="Two std band") 180 | plt.xlabel('$t$',fontsize=26) 181 | plt.ylabel('$x_1(t)$',fontsize=26) 182 | plt.ylim((-2.1, 2.5)) 183 | plt.legend(loc='upper right', frameon=False, prop={'size': 20}) 184 | plt.savefig('./BNODE_Prediction_x1.png', dpi = 300) 185 | 186 | 187 | plt.figure(5, figsize=(12,6.5)) 188 | plt.xticks(fontsize=22) 189 | plt.yticks(fontsize=22) 190 | plt.plot(t_grid_train, y_train[:,1], 'ro', label = "Training data of $x_2(t)$") 191 | plt.plot(t_grid_true, y_true[:,1], 'r-', label = "True Trajectory of $x_2(t)$") 192 | plt.plot(t_grid_true, y_MAP[:,1], 'g--', label = "MAP Trajectory of $x_2(t)$") 193 | lower_1 = mu_pred[:,1] - 2.0*np.sqrt(Sigma_pred[:,1]) 194 | upper_1 = mu_pred[:,1] + 2.0*np.sqrt(Sigma_pred[:,1]) 195 | plt.fill_between(t_grid_true.flatten(), lower_1.flatten(), upper_1.flatten(), 196 | facecolor='orange', alpha=0.5, label="Two std band") 197 | plt.xlabel('$t$',fontsize=26) 198 | plt.ylabel('$x_2(t)$',fontsize=26) 199 | plt.ylim((-2.1, 2.5)) 200 | plt.legend(loc='upper right', frameon=False, prop={'size': 20}) 201 | plt.savefig('./BNODE_Prediction_x2.png', dpi = 300) 202 | 203 | 204 | 205 | -------------------------------------------------------------------------------- /Spiral_noisy/example.py: -------------------------------------------------------------------------------- 1 | import os 2 | os.environ['KMP_DUPLICATE_LIB_OK']='True' 3 | 4 | import numpy as np 5 | import numpy.random as npr 6 | import tensorflow as tf 7 | import matplotlib.pyplot as plt 8 | plt.switch_backend('agg') 9 | import tensorflow.contrib.eager as tfe 10 | from tensorflow.keras import layers, initializers 11 | from scipy.integrate import odeint 12 | 13 | 14 | keras = tf.keras 15 | tf.enable_eager_execution() 16 | 17 | from neural_ode import NeuralODE 18 | 19 | np.random.seed(1234) 20 | tf.set_random_seed(1234) 21 | 22 | if __name__ == "__main__": 23 | 24 | def Spiral(z, t, alpha, beta, gamma, theta): 25 | x, y = z 26 | dzdt = [alpha * x**3 + beta * y**3, gamma * x**3 + theta * y**3] 27 | return dzdt 28 | 29 | def plot_spiral(trajectories, width = 1.): 30 | x = trajectories[:,0] 31 | y = trajectories[:,1] 32 | plt.plot(x, y, linewidth=width) 33 | 34 | 35 | data_size = 1001 36 | batch_time = 80 37 | niters = 1500 38 | batch_size = 920 39 | 40 | 41 | alpha = -0.1 42 | beta = 2.0 43 | gamma = -2.0 44 | theta = -0.1 45 | 46 | z0 = [2., 0.] 47 | true_y0 = tf.to_float([[2., 0.]]) 48 | t_grid = np.linspace(0, 20, data_size) 49 | 50 | true_yy = odeint(Spiral, z0, t_grid, args=(alpha, beta, gamma, theta)) 51 | true_yy.shape 52 | 53 | true_yy.shape 54 | true_y = true_yy + 0.02 * np.random.randn(true_yy.shape[0], true_yy.shape[1]) 55 | 56 | 57 | def get_batch(): 58 | """Returns initial point and last point over sampled frament of trajectory""" 59 | starts = np.random.choice(np.arange(data_size - batch_time - 1, dtype=np.int64), batch_size, replace=False) 60 | batch_y0 = true_y[starts] 61 | batch_yN = true_y[starts + batch_time] 62 | return tf.to_float(batch_y0), tf.to_float(batch_yN) 63 | 64 | 65 | Order = 4 66 | num_param = 1 67 | for m in range(1, Order): 68 | for n in range(m + 1): 69 | num_param += 1 70 | 71 | num_param = num_param * 2 72 | 73 | 74 | t0 = t_grid[:batch_time][0] 75 | t1 = t_grid[:batch_time][-1] 76 | t_in = np.linspace(t0, t1, 20) 77 | 78 | 79 | batch_y0, batch_yN = get_batch() 80 | 81 | ######################################### 82 | ########## precondition start ########### 83 | ######################################### 84 | niters_pre = 1000 85 | 86 | class ODEModel_pre(tf.keras.Model): 87 | def __init__(self): 88 | super(ODEModel_pre, self).__init__() 89 | self.Weights = tfe.Variable(tf.random_normal([num_param//2, 2], dtype=tf.float32)*0.01, dtype=tf.float32) 90 | 91 | def call(self, inputs, **kwargs): 92 | t, y = inputs 93 | h = y #**3 # batch_size x 2 94 | x1 = h[:,0:1] 95 | x2 = h[:,1:2] 96 | temp = x1 ** 0 97 | for m in range(1, Order): 98 | for n in range(m + 1): 99 | temp = tf.concat([temp, x1**n * x2**(m - n)], 1) 100 | 101 | h_out = tf.matmul(temp, self.Weights) 102 | return h_out 103 | 104 | 105 | model_pre = ODEModel_pre() 106 | neural_ode_pre = NeuralODE(model_pre, t_in) 107 | optimizer = tf.train.AdamOptimizer(3e-2) 108 | 109 | 110 | def compute_gradients_and_update_pre(batch_y0, batch_yN): 111 | """Takes start positions (x0, y0) and final positions (xN, yN)""" 112 | pred_y = neural_ode_pre.forward(batch_y0) 113 | with tf.GradientTape() as g_pre: 114 | g_pre.watch(pred_y) 115 | loss = tf.reduce_mean((pred_y - batch_yN)**2) + tf.reduce_sum(tf.abs(model_pre.Weights)) 116 | 117 | dLoss = g_pre.gradient(loss, pred_y) 118 | h_start, dfdh0, dWeights = neural_ode_pre.backward(pred_y, dLoss) 119 | optimizer.apply_gradients(zip(dWeights, model_pre.weights)) 120 | return loss, dWeights 121 | 122 | # Compile EAGER graph to static (this will be much faster) 123 | compute_gradients_and_update_pre = tfe.defun(compute_gradients_and_update_pre) 124 | 125 | parameters_pre = np.zeros((niters_pre, num_param)) 126 | 127 | for step in range(niters_pre): 128 | print(step) 129 | loss, dWeights = compute_gradients_and_update_pre(batch_y0, batch_yN) 130 | 131 | model_parameters_pre = model_pre.trainable_weights[0].numpy().flatten() 132 | for k in range(num_param): 133 | parameters_pre[step, k] = model_parameters_pre[k] 134 | 135 | print(parameters_pre[step,:]) 136 | 137 | ######################################### 138 | ########## precondition end ############# 139 | ######################################### 140 | 141 | initial_weight = model_pre.trainable_weights[0].numpy() 142 | print(initial_weight.shape, "here") 143 | 144 | class ODEModel(tf.keras.Model): 145 | def __init__(self, initial_weight): 146 | super(ODEModel, self).__init__() 147 | self.Weights = tfe.Variable(initial_weight, dtype=tf.float32) 148 | 149 | def call(self, inputs, **kwargs): 150 | t, y = inputs 151 | h = y 152 | x1 = h[:,0:1] 153 | x2 = h[:,1:2] 154 | temp = x1 ** 0 155 | for m in range(1, Order): 156 | for n in range(m + 1): 157 | temp = tf.concat([temp, x1**n * x2**(m - n)], 1) 158 | 159 | h_out = tf.matmul(temp, self.Weights) 160 | return h_out 161 | 162 | 163 | model = ODEModel(initial_weight) 164 | neural_ode = NeuralODE(model, t = t_in) 165 | 166 | def compute_gradients_and_update(batch_y0, batch_yN): 167 | """Takes start positions (x0, y0) and final positions (xN, yN)""" 168 | pred_y = neural_ode.forward(batch_y0) 169 | with tf.GradientTape() as g: 170 | g.watch(pred_y) 171 | loss = tf.reduce_sum((pred_y - batch_yN)**2) 172 | 173 | dLoss = g.gradient(loss, pred_y) 174 | h_start, dfdh0, dWeights = neural_ode.backward(pred_y, dLoss) 175 | 176 | return loss, dWeights 177 | 178 | # Compile EAGER graph to static (this will be much faster) 179 | compute_gradients_and_update = tfe.defun(compute_gradients_and_update) 180 | 181 | 182 | # function to compute the kinetic energy 183 | def kinetic_energy(V, loggamma_v, loglambda_v): 184 | q = (np.sum(-V**2) - loggamma_v**2 - loglambda_v**2)/2.0 185 | return q 186 | 187 | def compute_gradient_param(dWeights, loggamma, loglambda, batch_size, para_num): 188 | WW = model.trainable_weights[0].numpy() 189 | dWeights = np.exp(loggamma)/2.0 * dWeights + np.exp(loglambda) * np.sign(WW) 190 | return dWeights 191 | 192 | def compute_gradient_hyper(loss, weights, loggamma, loglambda, batch_size, para_num): 193 | grad_loggamma = np.exp(loggamma) * (loss/2.0 + 1.0) - (batch_size/2.0 + 1.0) 194 | grad_loglambda = np.exp(loglambda) * (np.sum(np.abs(weights)) + 1.0) - (para_num + 1.0) 195 | 196 | return grad_loggamma, grad_loglambda 197 | 198 | def compute_Hamiltonian(loss, weights, loggamma, loglambda, batch_size, para_num): 199 | H = np.exp(loggamma)*(loss/2.0 + 1.0) + np.exp(loglambda)*(np.sum(np.abs(weights)) + 1.0)\ 200 | - (batch_size/2.0 + 1.0) * loggamma - (para_num + 1.0) * loglambda 201 | return H 202 | 203 | def leap_frog(v_in, w_in, loggamma_in, loglambda_in, loggamma_v_in, loglambda_v_in): 204 | # assign weights from the previous step to the model 205 | model.trainable_weights[0].assign(w_in) 206 | 207 | print(model.trainable_weights) 208 | 209 | v_new = v_in 210 | loggamma_v_new = loggamma_v_in 211 | loglambda_v_new = loglambda_v_in 212 | 213 | loggamma_new = loggamma_in 214 | loglambda_new = loglambda_in 215 | w_new = w_in 216 | 217 | for m in range(L): 218 | 219 | loss, dWeights = compute_gradients_and_update(batch_y0, batch_yN) # evaluate the gradient 220 | print(loss) 221 | 222 | dWeights = np.asarray(dWeights[0]) # make the gradient to be numpy array 223 | dWeights = compute_gradient_param(dWeights, loggamma_new, loglambda_new, batch_size, num_param) 224 | grad_loggamma, grad_loglambda = compute_gradient_hyper(loss, w_new, loggamma_new, loglambda_new, batch_size, num_param) 225 | 226 | loggamma_v_new = loggamma_v_new - epsilon/2*grad_loggamma 227 | loglambda_v_new = loglambda_v_new - epsilon/2*grad_loglambda 228 | v_new = v_new - epsilon/2*(dWeights) 229 | w_new = model.trainable_weights[0].numpy() + epsilon * v_new 230 | model.trainable_weights[0].assign(w_new) 231 | loggamma_new = loggamma_new + epsilon * loggamma_v_new 232 | loglambda_new = loglambda_new + epsilon * loglambda_v_new 233 | 234 | # Second half of the leap frog 235 | loss, dWeights = compute_gradients_and_update(batch_y0, batch_yN) 236 | dWeights = np.asarray(dWeights[0]) 237 | dWeights = compute_gradient_param(dWeights, loggamma_new, loglambda_new, batch_size, num_param) 238 | grad_loggamma, grad_loglambda = compute_gradient_hyper(loss, w_new, loggamma_new, loglambda_new, batch_size, num_param) 239 | 240 | v_new = v_new - epsilon/2*(dWeights) 241 | loggamma_v_new = loggamma_v_new - epsilon/2*grad_loggamma 242 | loglambda_v_new = loglambda_v_new - epsilon/2*grad_loglambda 243 | 244 | print(dWeights) 245 | print(np.exp(loggamma_new)) 246 | print(np.exp(loglambda_new)) 247 | 248 | return v_new, w_new, loggamma_new, loglambda_new, loggamma_v_new, loglambda_v_new 249 | 250 | 251 | 252 | 253 | neural_ode_test = NeuralODE(model, t=t_grid) 254 | parameters = np.zeros((niters, num_param//2, 2)) # book keeping the parameters 255 | loggammalist = np.zeros((niters, 1)) # book keeping the loggamma 256 | loglambdalist = np.zeros((niters, 1)) # book keeping the loggamma 257 | loglikelihood = np.zeros((niters, 1)) # book keeping the loggamma 258 | L = 10 # leap frog step number 259 | epsilon = 0.001 # leap frog step size 260 | epsilon_max = 0.0002 # max 0.001 261 | epsilon_min = 0.0002 # max 0.001 262 | 263 | 264 | def compute_epsilon(step): 265 | coefficient = np.log(epsilon_max/epsilon_min) 266 | return epsilon_max * np.exp( - step * coefficient / niters) 267 | 268 | 269 | # initial weight 270 | w_temp = initial_weight 271 | print("initial_w", w_temp) 272 | loggamma_temp = 4. + np.random.normal() 273 | loglambda_temp = np.random.normal() 274 | 275 | model.trainable_weights[0].assign(w_temp) 276 | loss_original, _ = compute_gradients_and_update(batch_y0, batch_yN) # compute the initial Hamiltonian 277 | 278 | loggamma_temp = np.log(batch_size / loss_original) 279 | print("This is initial guess", loggamma_temp, "with loss", loss_original) 280 | if loggamma_temp > 6.: 281 | loggamma_temp = 6. 282 | epsilon_max = 0.0001 283 | epsilon_min = 0.0001 284 | 285 | 286 | 287 | # training steps 288 | for step in range(niters): 289 | 290 | epsilon = compute_epsilon(step) 291 | 292 | print(step) 293 | v_initial = np.random.randn(num_param//2, 2) # initialize the velocity 294 | loggamma_v_initial = np.random.normal() 295 | loglambda_v_initial = np.random.normal() 296 | 297 | loss_initial, _ = compute_gradients_and_update(batch_y0, batch_yN) # compute the initial Hamiltonian 298 | loss_initial = compute_Hamiltonian(loss_initial, w_temp, loggamma_temp, loglambda_temp, batch_size, num_param) 299 | 300 | v_new, w_new, loggamma_new, loglambda_new, loggamma_v_new, loglambda_v_new = \ 301 | leap_frog(v_initial, w_temp, loggamma_temp, loglambda_temp, loggamma_v_initial, loglambda_v_initial) 302 | 303 | # compute the final Hamiltonian 304 | loss_finial, _ = compute_gradients_and_update(batch_y0, batch_yN) 305 | loss_finial = compute_Hamiltonian(loss_finial, w_new, loggamma_new, loglambda_new, batch_size, num_param) 306 | 307 | # making decisions 308 | p_temp = np.exp(-loss_finial + loss_initial + \ 309 | kinetic_energy(v_new, loggamma_v_new, loglambda_v_new) - kinetic_energy(v_initial, loggamma_v_initial, loglambda_v_initial)) 310 | 311 | p = min(1, p_temp) 312 | p_decision = np.random.uniform() 313 | if p > p_decision: 314 | parameters[step:step+1, :, :] = w_new 315 | w_temp = w_new 316 | loggammalist[step, 0] = loggamma_new 317 | loglambdalist[step, 0] = loglambda_new 318 | loglikelihood[step, 0] = loss_finial 319 | loggamma_temp = loggamma_new 320 | loglambda_temp = loglambda_new 321 | else: 322 | parameters[step:step+1, :, :] = w_temp 323 | model.trainable_weights[0].assign(w_temp) 324 | loggammalist[step, 0] = loggamma_temp 325 | loglambdalist[step, 0] = loglambda_temp 326 | loglikelihood[step, 0] = loss_initial 327 | 328 | 329 | print('probability', p) 330 | print(p > p_decision) 331 | 332 | 333 | np.save('parameters', parameters) 334 | np.save('loggammalist', loggammalist) 335 | np.save('loglikelihood', loglikelihood) 336 | 337 | 338 | 339 | 340 | -------------------------------------------------------------------------------- /Spiral_noisy/neural_ode.py: -------------------------------------------------------------------------------- 1 | from typing import Optional, List 2 | import numpy as np 3 | import tensorflow as tf 4 | from tensorflow.python.framework.ops import EagerTensor 5 | import tensorflow.contrib.eager as tfe 6 | 7 | keras = tf.keras 8 | 9 | 10 | def zip_map(zipped, update_op): 11 | return [update_op(*elems) for elems in zipped] 12 | 13 | 14 | def euler_update(h_list, dh_list, dt): 15 | return zip_map(zip(h_list, dh_list), lambda h, dh: h + tf.cast(dt, h.dtype) * dh) 16 | 17 | 18 | def euler_step(func, dt, state): 19 | return euler_update(state, func(state), dt) 20 | 21 | 22 | def rk2_step(func, dt, state): 23 | k1 = func(state) 24 | k2 = func(euler_update(state, k1, dt)) 25 | return zip_map(zip(state, k1, k2), 26 | lambda h, dk1, dk2: h + tf.cast(dt, h.dtype) * (dk1 + dk2) / 2) 27 | 28 | 29 | def rk4_step(func, dt, state): 30 | k1 = func(state) 31 | k2 = func(euler_update(state, k1, dt / 2)) 32 | k3 = func(euler_update(state, k2, dt / 2)) 33 | k4 = func(euler_update(state, k3, dt)) 34 | 35 | return zip_map( 36 | zip(state, k1, k2, k3, k4), 37 | lambda h, dk1, dk2, dk3, dk4: h + tf.cast(dt, h.dtype) * ( 38 | dk1 + 2 * dk2 + 2 * dk3 + dk4) / 6, 39 | ) 40 | 41 | 42 | class NeuralODE: 43 | def __init__( 44 | self, model: tf.keras.Model, t=np.linspace(0, 1, 40), 45 | solver=rk4_step 46 | ): 47 | self._t = t 48 | self._model = model 49 | self._solver = solver 50 | self._deltas_t = t[1:] - t[:-1] 51 | 52 | def forward(self, inputs: tf.Tensor, return_states: Optional[str] = None): 53 | 54 | def _forward_dynamics(_state): 55 | """Used in solver _state == (time, tensor)""" 56 | return [1.0, self._model(inputs=_state)] 57 | 58 | states = [] 59 | 60 | def _append_state(_state): 61 | tensors = _state[1] 62 | if return_states == "numpy": 63 | states.append(tensors.numpy()) 64 | elif return_states == "tf": 65 | states.append(tensors) 66 | 67 | with tf.name_scope("forward"): 68 | t0 = tf.to_float(self._t[0]) 69 | state = [t0, inputs] 70 | _append_state(state) 71 | for dt in self._deltas_t: 72 | state = self._solver( 73 | func=_forward_dynamics, dt=tf.to_float(dt), state=state 74 | ) 75 | _append_state(state) 76 | 77 | outputs = state[1] 78 | if return_states: 79 | return outputs, states 80 | return outputs 81 | 82 | def _backward_dynamics(self, state): 83 | t = state[0] 84 | ht = state[1] 85 | at = -state[2] 86 | 87 | with tf.GradientTape() as g: 88 | g.watch(ht) 89 | ht_new = self._model(inputs=[t, ht]) 90 | 91 | gradients = g.gradient( 92 | target=ht_new, sources=[ht] + self._model.weights, 93 | output_gradients=at 94 | ) 95 | 96 | return [1.0, ht_new, *gradients] 97 | 98 | def backward(self, outputs: tf.Tensor, 99 | output_gradients: Optional[tf.Tensor] = None): 100 | 101 | with tf.name_scope("backward"): 102 | grad_weights = [tf.zeros_like(w) for w in self._model.weights] 103 | t0 = tf.to_float(self._t[-1]) 104 | 105 | if output_gradients is None: 106 | output_gradients = tf.zeros_like(outputs) 107 | 108 | state = [t0, outputs, output_gradients, *grad_weights] 109 | for dt in self._deltas_t[::-1]: 110 | state = self._solver( 111 | self._backward_dynamics, dt=-tf.to_float(dt), state=state 112 | ) 113 | 114 | inputs = state[1] 115 | dLdInputs = state[2] 116 | dLdWeights = state[3:] 117 | return inputs, dLdInputs, dLdWeights 118 | -------------------------------------------------------------------------------- /Spiral_noisy/plotting.py: -------------------------------------------------------------------------------- 1 | import numpy as np 2 | import matplotlib as mpl 3 | #mpl.use('pgf') 4 | 5 | def figsize(scale, nplots = 1): 6 | fig_width_pt = 390.0 # Get this from LaTeX using \the\textwidth 7 | inches_per_pt = 1.0/72.27 # Convert pt to inch 8 | golden_mean = (np.sqrt(5.0)-1.0)/2.0 # Aesthetic ratio (you could change this) 9 | fig_width = fig_width_pt*inches_per_pt*scale # width in inches 10 | fig_height = nplots*fig_width*golden_mean # height in inches 11 | fig_size = [fig_width,fig_height] 12 | return fig_size 13 | 14 | pgf_with_latex = { # setup matplotlib to use latex for output 15 | "pgf.texsystem": "pdflatex", # change this if using xetex or lautex 16 | "text.usetex": True, # use LaTeX to write all text 17 | "font.family": "serif", 18 | "font.serif": [], # blank entries should cause plots to inherit fonts from the document 19 | "font.sans-serif": [], 20 | "font.monospace": [], 21 | "axes.labelsize": 10, # LaTeX default is 10pt font. 22 | "font.size": 10, 23 | "legend.fontsize": 8, # Make the legend/label fonts a little smaller 24 | "xtick.labelsize": 8, 25 | "ytick.labelsize": 8, 26 | "figure.figsize": figsize(1.0), # default fig size of 0.9 textwidth 27 | "pgf.preamble": [ 28 | r"\usepackage[utf8x]{inputenc}", # use utf8 fonts becasue your computer can handle it :) 29 | r"\usepackage[T1]{fontenc}", # plots will be generated using this preamble 30 | ] 31 | } 32 | mpl.rcParams.update(pgf_with_latex) 33 | 34 | import matplotlib.pyplot as plt 35 | 36 | # I make my own newfig and savefig functions 37 | def newfig(width, nplots = 1): 38 | fig = plt.figure(figsize=figsize(width, nplots)) 39 | ax = fig.add_subplot(111) 40 | return fig, ax 41 | 42 | def savefig(filename, crop = True): 43 | if crop == True: 44 | # plt.savefig('{}.pgf'.format(filename), bbox_inches='tight', pad_inches=0) 45 | plt.savefig('{}.pdf'.format(filename), bbox_inches='tight', pad_inches=0) 46 | plt.savefig('{}.eps'.format(filename), bbox_inches='tight', pad_inches=0) 47 | else: 48 | # plt.savefig('{}.pgf'.format(filename)) 49 | plt.savefig('{}.pdf'.format(filename)) 50 | plt.savefig('{}.eps'.format(filename)) 51 | 52 | ## Simple plot 53 | #fig, ax = newfig(1.0) 54 | # 55 | #def ema(y, a): 56 | # s = [] 57 | # s.append(y[0]) 58 | # for t in range(1, len(y)): 59 | # s.append(a * y[t] + (1-a) * s[t-1]) 60 | # return np.array(s) 61 | # 62 | #y = [0]*200 63 | #y.extend([20]*(1000-len(y))) 64 | #s = ema(y, 0.01) 65 | # 66 | #ax.plot(s) 67 | #ax.set_xlabel('X Label') 68 | #ax.set_ylabel('EMA') 69 | # 70 | #savefig('ema') --------------------------------------------------------------------------------