Add nequip to train.

[harveydevereux]

This will add nequip as a janus train [architecture] option, where the [architecture] is now a required argument to train. All Architectures other than "mace", "mace_mp", "mace_off", "mace_omol", and "nequip" raise an error.

It does work with these inputs inputs.zip

Todo:

• Testcase with janus train nequip.
• Documentation.

One thing I noticed is that janus train mace ... will currently output outside janus_results, but the way I have it for nequip it goes into janus_results. What's better? E.g. for mace,

checkpoints  config.yml  janus_results  logs  results  test_compiled.model  test.model

Or is that something that is control by the mace config.yml? For nequip I've had to set it manually so I chose ./janus_results.

---

Details

Similar to mace.cli.run_train.run, nequip has a nequip.scripts.train.main entrypoint - but it uses Hydra to manage configuration here.

I've done similar to the current code and used the main function rather than a subprocess. This means using the Compose API which at least one downsides - the hydra.runtime.output_dir is not set, but is required, and we have to manually set the HydraConfig singleton (as far as I've found...).

This does mean that we have a match that obtains the runner function and its arguments, seperating out modules to avoid conflicts (mace/nequip in fact is one).

[ElliottKasoar]

Thank you so much for this!

How does nequip fine-tuning work?

One thing I noticed is that janus train mace ... will currently output outside janus_results, but the way I have it for nequip it goes into janus_results. What's better? E.g. for mace,

Perhaps we should set the default value for mace's --work_dir (https://github.com/ACEsuit/mace/blob/main/mace/tools/arg_parser.py#L38) to janus-core, and then do the equivalent for any new architecture?

This does mean that we have a match that obtains the runner function and its arguments, seperating out modules to avoid conflicts (mace/nequip in fact is one).

I think this was probably unavoidable given how different training entry points will be, so I think this is fine?