Draft: Merge GNN paper progress into main

Adds all the progress towards the GNN paper. We added some new variables to the data_loading.py and some functionality to only look at certain features / disable the lower layer input for the GNNs here.

Once new models for D8 and D9 for energy and g/h separation are trained and tested we should update this branch and prepare to merge it into main