[DRAFT] Add partial charge prediction to EnergyDipoleMACE

[jan-vavrin]

This change would allow the EnergyDipoleMACE model to predict (and be trained on) partial charges from DFT calculations. It is based on this discussion.

The model already predicts charges internally, so the only change is that we add a loss function for the charges in addition to energy, forces and dipole.

We also normalise the predicted charges to sum to the total charge -- this is done in AtomicDielectricMACE, but has not been done in EnergyDipoleMACE so far.

The way it is implemented now means that partial charge prediction is always turned on. However, we would likely want to make this feature optional (via a command line argument like "--predict_charges") and then check that charges are provided in the training set if this argument is true. However, I am not sure where is the best place to do this in the code.

To train the model, something like this ought to work:
mace_run_train --model EnergyDipolesMACE --loss energy_forces_dipole_charges --charges_weight 10.0 --error_table EnergyDipoleChargeRMSE --embedding_specs="{'charge': {'embed_type': 'continuous', 'min': -2, 'max': 2}}" --use_embedding_readout True [other usual args]

Ideally, the "--predict_charges" argument would automatically set "--loss energy_forces_dipole_charges" and "--error_table EnergyDipoleChargeRMSE" and maybe also the charge embedding.

The training set needs to have energies, forces, partial charges, dipole and the total charge (though in principle, the total charge is just the sum of partial charges). An example training set I have in mind is train.extxyz.txt, which is just a bunch of clusters including H3O+, H2O and OH- (at a very low level of theory, but should be fine as an example).

[ilyes319]

Hey @jan-vavrin, thank you for that PR could you add some tests for it? Also make sure all the CLI tests are passing.