full ensemble fix, feature model fix #332
Conversation
apax/md/ase_calc.py
Outdated
| ensemble = {k + "_ensemble": v for k, v in results.items()} | ||
| results = {k: jnp.mean(v, axis=0) for k, v in results.items()} | ||
| results.update(uncertainty) | ||
| results.update(ensemble) |
There was a problem hiding this comment.
Does this mean, when using any kind of ensemble, we will have:
{energy: float, energy_ensemble: list[float], forces: ndarray[N, 3], forces_ensemble: ndarray[m, N, 3]}
and so on? I see where this can be very useful, but for the sake of performance it might be a good idea to make this toggleable? For e.g. saving to h5 the amount of data will increase noticeably?
There was a problem hiding this comment.
I also opened zincware/ZnH5MD#136 but I think this option should be available on both sides?
There was a problem hiding this comment.
Yes these are now available for all types of ensemble. passing arguments to the ASE calculator in IPS is a bit cumbersome at the moment. You always need a second node in addition to the training node (otherwise you have to retrain when changing the ASE options).
The inference speed should actually not be affected, but yes the storage size will increase significantly. Storing the ensemble predictions is relevant for some UQ metrics and reweighting.
I guess I can add an optional white list to the ASE calculator which filters the results by keys. same for jaxMD.
it seems like a cleaner solution to me to specify this in the h5 writer, but I can also implement it on the apax side.
There was a problem hiding this comment.
I guess it doesn't hurt to implement it here as well, but I agree that it is more useful on the znh5md side.
There was a problem hiding this comment.
Ah right, I forgot about ASE. Yeah I can do that for ASE as well, but it's not going to be very ergonomic for IPSuite.
There was a problem hiding this comment.
I have to say that this would be really unergonomical when using the ASECalculator in IPSuite. I think models should return what they return and which properties are logged should be part of the trajectory writer. If I implement this as part of the ASECalc, it would mean that we train a model with a specified list of outputs and changing these properties would require a retraining or using a separate Node.
So I think this was a good addition to jaxMD (as it is implemented in the trajectory writer), but for the ASECalculator it makes everything more complicated.
for more information, see https://pre-commit.ci
|
pre-commit.ci autofix |
for more information, see https://pre-commit.ci
|
pre-commit.ci autofix |
for more information, see https://pre-commit.ci
for more information, see https://pre-commit.ci
|
pre-commit.ci autofix |
added ensemble output for full ensembles to ASE calc
minor changes to feature model for enhanced sampling