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:

[ilyes319]

@vue1999 Could you check that?

[vue1999]

Thanks for this, I have checked the code and tested it, looks good to me. Could you please sync the branch with main and ensure the unit tests pass?

[joehart2001]

Thanks for this, I have checked the code and tested it, looks good to me. Could you please sync the branch with main and ensure the unit tests pass?

Thanks @vue1999 for testing. I've rebased now. lt seems the errors are coming from files i have not touched. any advice?

[vue1999]

Looks good, tested and unit tests pass✅

[vue1999]

Thanks for this, I have checked the code and tested it, looks good to me. Could you please sync the branch with main and ensure the unit tests pass?

Thanks @vue1999 for testing. I've rebased now. lt seems the errors are coming from files i have not touched. any advice?

Thanks, the unit tests were fixed in the main branch changes, so after rebasing they have succeeded as expected

[joehart2001]

Thanks for this, I have checked the code and tested it, looks good to me. Could you please sync the branch with main and ensure the unit tests pass?

Thanks @vue1999 for testing. I've rebased now. lt seems the errors are coming from files i have not touched. any advice?

Thanks, the unit tests were fixed in the main branch changes, so after rebasing they have succeeded as expected

Amazing, thanks!