Added linear interp. and model test scripts

EMRI_DET/nn/model_creation.py

@@ -45,6 +45,8 @@ def create_mlp(input_features, output_features, neurons, layers, activation, mod
     return model
 
 
-def load_mlp(model_name):
+def load_mlp(model_name, get_state_dict=False):
     model = pickle.load(open(get_script_path()+'/../models/'+model_name+'_function.pickle', "rb"))  # load blank model
+    if get_state_dict:
+        model.load_state_dict(torch.load(get_script_path()+'/../models/'+model_name+'_model.pth'))
     return model

mock_data/mock_data_function_1/scripts/linear_interpolation.py

+import numpy as np
+from scipy.interpolate import RegularGridInterpolator
+import cupy as cp
+from cupyx.scipy.ndimage import map_coordinates
+import time
+import matplotlib.pyplot as plt
+import torch
+from EMRI_DET.nn.model_creation import load_mlp
+from EMRI_DET.utilities import norm, unnorm_inputs, norm_inputs, unnorm
+
+
+xgrid = np.load('../data/xgrid.npy')
+ygrid = np.load('../data/ygrid.npy')
+
+xdata = np.load('../data/xtest.npy')
+ydata = np.load('../data/ytest.npy')
+
+nM = 8
+nq = 5
+na = 4
+ne = 4
+nY = 4
+nTh = 5
+nPh = 8
+nt = 5
+
+xdata_in = ()
+for i in range(xgrid.shape[1]):
+    xdata_in += (np.unique(xgrid[:,i]),)
+
+y_in = ygrid.reshape((nM,nq,na,ne,nY,nTh,nPh,nTh,nt))
+
+xdata_in2 = xgrid.copy()
+
+print('Beginning Linear interpolator init')
+st = time.perf_counter()
+
+interp = RegularGridInterpolator(xdata_in, y_in, method='linear',bounds_error=False)
+nn_interp = RegularGridInterpolator(xdata_in, y_in, method='nearest',bounds_error=False)
+
+dat1 = interp(xdata_in2[:1000])
+print('Consistency check:', np.sum(dat1 - ygrid.flatten()[:1000]))
+
+mt = time.perf_counter()
+print('Starting interpolation')
+
+lin_values = interp(xdata)
+nearest_values = nn_interp(xdata)
+et = time.perf_counter()
+RMSE = np.sqrt(np.nanmean((lin_values- ydata)**2))
+
+plt.hist(np.log10(abs(lin_values-ydata)), bins='auto', alpha=0.5, label='Linear Interp. - 1e6 points',histtype='step',density=True)
+plt.hist(np.log10(abs(nearest_values-ydata)), bins='auto', alpha=0.5, label='Nearest Neighbour - 1e6 points',histtype='step',density=True)
+
+print('Linear grid interpolator RMSE: ', RMSE)
+print(f'Time for [interpolator init, interpolation] | [{mt-st:.3e}, {et-mt:.3e}]')
+
+Mvals = np.linspace(np.log10(8e4), np.log10(5e7), nM)
+qvals = np.linspace(np.log10(5e4), np.log10(5e7), nq)
+avals = np.linspace(0.01, 0.99, na)
+evals = np.linspace(0.01, 0.5, ne)
+Yvals = np.linspace(0.5, 0.99, nY)
+qSvals = np.linspace(0, np.pi, nTh)
+qKvals = np.linspace(0, np.pi, nTh)
+phiSvals = np.linspace(0, 2 * np.pi, nPh)
+tVals = np.linspace(0, 4, nt + 1)[1:]
+
+ydata_reshaped = ygrid.reshape((nM,nq,na,ne,nY,nTh,nPh,nTh,nt))
+
+xtest_in = []
+
+lists = [Mvals, qvals, avals, evals, Yvals, qSvals, phiSvals, qKvals, tVals]
+norms = [listy[-1] - listy[0] for listy in lists]
+nums = [nM, nq, na, ne, nY, nTh, nPh, nTh, nt]
+
+for i in range(ydata.size):
+    sub = []
+    for j in range(9):
+        sub.append((nums[j]-1)*(xdata[i, j] - lists[j][0])/norms[j])
+    xtest_in.append(sub)
+
+xtest_in = np.asarray(xtest_in).T
+mt = time.perf_counter()
+print('Starting interpolation')
+
+map_values = map_coordinates(cp.array(ydata_reshaped), cp.array(xtest_in), prefilter=True).get()
+et = time.perf_counter()
+
+RMSE = np.sqrt(np.nanmean((map_values - ydata)**2))
+
+print('map_coordinates grid interpolator RMSE: ', RMSE)
+print(f'Time for interpolation | {et-mt:.3e}')
+
+plt.hist(np.log10(abs(map_values-ydata)), bins='auto', alpha=0.5, label='n-Cubic Spine interp. - 1e6 points',histtype='step',density=True)
+
+ydata = np.log(ydata)  # network is trained on the log of the SNR
+
+device = 'cuda:0'
+model_name = 'm1'
+mlp = load_mlp(model_name, get_state_dict=True).to(device)
+mlp.eval()
+
+xmeanstd = np.load(f'../models/{model_name}_xdata_mean_std.npy')
+ymeanstd = np.load(f'../models/{model_name}_ydata_mean_std.npy')
+
+test_input = torch.Tensor(xdata)
+normed_input = norm_inputs(test_input, xmeanstd[0], xmeanstd[1]).float().to(device)
+st = time.perf_counter()
+with torch.no_grad():
+    out = []
+    for i in range(20):
+        output = mlp(normed_input[i * ydata.size // 20 : (i+1)*ydata.size // 20])
+        out.append(output.detach().cpu().numpy())
+
+et = time.perf_counter()
+print(f'Evaluation time: {et-st:.3e}s for {ydata.size:.1e} values, {(et-st)/ydata.size:.3e}s per data point.')
+output = np.concatenate(out)
+
+out_unnorm = np.exp(unnorm(output, ymeanstd[0], ymeanstd[1]))
+
+network_out = out_unnorm.flatten()
+
+plt.hist(np.log10(abs(network_out-ydata)), bins='auto', alpha=0.5,label='NN - 1e6 points', histtype='step', density=True)
+
+RMSE = np.sqrt(np.nanmean((network_out - ydata)**2))
+
+print('Neural network interpolator RMSE: ', RMSE)
+
+plt.legend()
+plt.xlabel('log10(|diff|) between truth and interpolated')
+plt.title('Neural Network (NN) performance vs. selected interpolation methods')
+plt.show()

mock_data/mock_data_function_1/scripts/model_test.py

+import torch
+import numpy as np
+import matplotlib.pyplot as plt
+import time
+from EMRI_DET.utilities import norm, norm_inputs, unnorm
+from EMRI_DET.nn.model_creation import load_mlp
+
+device = 'cuda:0'
+model_name = 'm1'
+mlp = load_mlp(model_name, get_state_dict=True).to(device)
+mlp.eval()
+
+xdata = np.load('../data/xtest.npy')
+ydata = np.log(np.load('../data/ytest.npy'))
+
+xmeanstd = np.load(f'../models/{model_name}_xdata_mean_std.npy')
+ymeanstd = np.load(f'../models/{model_name}_ydata_mean_std.npy')
+
+test_input = torch.Tensor(xdata)
+normed_input = norm_inputs(test_input, xmeanstd[0], xmeanstd[1]).float().to(device)
+normed_input = normed_input.float().to(device)
+ynorm = norm(ydata, ymeanstd[0], ymeanstd[1])
+
+st = time.perf_counter()
+nb = 50
+
+with torch.no_grad():
+    out = []
+    for i in range(nb):
+        output = mlp(normed_input[i * ydata.size // nb : (i+1)*ydata.size // nb])
+        out.append(output.detach().cpu().numpy())
+
+et = time.perf_counter()
+print(f'Evaluation time: {et-st:.3e}s for {ydata.size:.1e} values, {(et-st)/ydata.size:.3e}s per data point.')
+
+
+output = np.concatenate(out)
+out_unnorm = np.exp(unnorm(output, ymeanstd[0], ymeanstd[1]))
+
+threshold = np.median(np.exp(ydata))
+
+out_classified = np.zeros(shape=ydata.size)
+out_classified[out_unnorm.flatten() >= threshold] = 1
+
+truth_classified = np.zeros(shape=ydata.size)
+truth_classified[np.exp(ydata) >= threshold] = 1
+
+plt.hist(out_classified,bins=2,histtype='step')
+plt.hist(truth_classified, bins=2,histtype='step')
+plt.show()
+
+print('Root mean square error (unnorm): ',np.sqrt(np.mean((out_unnorm.flatten() - np.exp(ydata))**2)))
+print('Root mean square error (norm): ',np.sqrt(np.mean((output.flatten() - ynorm)**2)))
+print('Accuracy: ', 1 - np.mean(np.abs(out_classified - truth_classified)))
+
+
+# out_plot = out_unnorm.detach().numpy().flatten()
+# diffs = abs(np.log10(abs(out_plot/ydata)))
+#
+# plt.hist(diffs, bins='auto')
+# plt.xlabel('|log10(ratio)| between prediction and truth')
+# plt.show()
+#
+# diffs2 = np.log10(abs(out_plot-ydata))
+#
+# plt.hist(diffs2, bins='auto')
+# plt.xlabel('log10(|difference|) between prediction and truth')
+# plt.show()
+
+# plt.hist(np.exp(out_plot), bins='auto', label='NN', alpha=0.5, density=True)
+# plt.hist(np.exp(ydata), bins='auto', label='Truth', alpha=0.5, density=True)
+# plt.legend()
+# plt.xlabel('Mock SNR')
+# plt.show()
\ No newline at end of file