Only plot head-specific test set in results for multihead models

[joehart2001]

Issue:
When training a multi-head model, all heads' test sets are plotted on every head's results panel. For a 3-head
model, each panel shows test data from all 3 heads (each evaluated with its own readout), all labelled "Test".
This has two consequences:

• The relevant per-head test scatter is obscured by test data from other heads
• When heads span different energy/force ranges, the axes are forced to accommodate all heads' ranges, reducing resolution on the head of interest

Solution:

• run_train.py: Label each test set with its head name before passing to the plotter
• visualise_train.py: only plot the test data for the current head + add the head name to the legend label

Testing

• 2 epochs with: single head, multihead, naive FT, mutlihead FT

• Multihead before (3 heads) (cropped out training plots as not relevant here):

• Multihead after: