plot data

examples/analysis_example.py

@@ -3,11 +3,13 @@
 from molearn.models.foldingnet import AutoEncoder
 from molearn.analysis import MolearnAnalysis
 from molearn.data import PDBData
+import matplotlib.pyplot as plt
 
 
 print("> Loading network parameters...")
 
 fname = f'xbb_foldingnet_checkpoints{os.sep}checkpoint_no_optimizer_state_dict_epoch167_loss0.003259085263643.ckpt'
+# change 'cpu' to 'cuda' if you have a suitable cuda enabled device
 checkpoint = torch.load(fname, map_location=torch.device('cpu'))
 net = AutoEncoder(**checkpoint['network_kwargs'])
 net.load_state_dict(checkpoint['model_state_dict'])
@@ -47,12 +49,26 @@
 err_train = MA.get_error('training')
 err_test = MA.get_error('test')
 
+print(f'Mean RMSD is {err_train.mean()} for training set and {err_test.mean()} for test set')
+fig, ax = plt.subplots()
+violin = ax.violinplot([err_train, err_test], showmeans = True, )
+ax.set_xticks([1,2])
+ax.set_title('RMSD of training and test set')
+ax.set_xticklabels(['Training', 'Test'])
+plt.savefig('RMSD_plot.png')
+
+
 print("> generating error landscape")
 # build a 50x50 grid. By default, it will be 10% larger than the region occupied
 # by all loaded datasets
 MA.setup_grid(50)
 landscape_err_latent, landscape_err_3d, xaxis, yaxis = MA.scan_error()
 
+fig, ax = plt.subplots()
+c = ax.pcolormesh(xaxis, yaxis, landscape_err_latent)
+plt.savefig('Error_grid.png')
+
+
 ## to visualise the GUI, execute the code above in a Jupyter notebook, then call:
 # from molearn.analysis import MolearnGUI
 # MolearnGUI(MA)