Energy forces

[RylieWeaver]

@allaffa Let me know what you think of this PR for energy and force computation. It should follow exactly what we discussed before. Some notable things:

(1) It requires data.energy and data.forces to exist and be named this way
(2) It balances force and energy prediction weighting implicitly, since we aren't able to pass multiple weights in single-task learning
(3) Compute forces is an argument to be specified in the json within "Training".
- If specified, it throws errors for anything other than nodal single task learning, as well as any model that doesn't
use positional information for prediction

[allaffa]

@RylieWeaver
This is exactly what we discussed and looks awesome.
Have you tried it with Lenndard Jones to make sure that it is compatible?

[RylieWeaver]

@allaffa

Awesome, thanks.

I was able to get it to work with Lennard Jones after some small changes to LJ's dataset creation and inference. For example, we now have data.energy and'data.forces instead of data.y and data.forces_pre_scaled. A next step could be updating the LennardJones example to that.

Do not merge it yet, because I am resolving any issues with the GitHub tests, but will let you know when they are passed.

[allaffa]

@allaffa

Awesome, thanks.

I was able to get it to work with Lennard Jones after some small changes to LJ's dataset creation and inference. For example, we now have data.energy and'data.forces instead of data.y and data.forces_pre_scaled. A next step could be updating the LennardJones example to that.

Do not merge it yet, because I am resolving any issues with the GitHub tests, but will let you know when they are passed.

@RylieWeaver
As long as the CI tests do not pass, I would not be allowed to merge this PR anyway. We have explicitly forbidden this option in the settings of the GitHub repository.

[allaffa]

@RylieWeaver

I think it would be good to also include a mew test that the CI workflow must execute where we check this capability to make sure that future changes in the code do not break it

[RylieWeaver]

@RylieWeaver As long as the CI tests do not pass, I would not be allowed to merge this PR anyway. We have explicitly forbidden this option in the settings of the GitHub repository.

Ok, that's good.

[RylieWeaver]

@RylieWeaver

I think it would be good to also include a mew test that the CI workflow must execute where we check this capability to make sure that future changes in the code do not break it

Ok, I can include this. Theoretically, any model that includes positional information should be able to predict forces in this method. Should we use that list for the testing?

[RylieWeaver]

@allaffa I believe I'll need to create a new dataset inside hydragnn/serialized_dataset for this testing, so that data.forces and data.energy are available. It seems like most of the other data is generated randomly from [0.0, 0.5, 1.0].

Is this correct and is random generation is fine for the forces/energy as well?

[allaffa]

@allaffa I believe I'll need to create a new dataset inside hydragnn/serialized_dataset for this testing, so that data.forces and data.energy are available. It seems like most of the other data is generated randomly from [0.0, 0.5, 1.0].

Is this correct and is random generation is fine for the forces/energy as well?

@RylieWeaver
I have an idea that I would like to run by you.
What about adding a frozen version of the LJ code as an additional example, where the example first generates the data, saves it into files, and then uses it for training? We don't need to generate many data samples and traing it fully. We just need to make sure that the training kicks off correctly without errors.

[RylieWeaver]

@allaffa I believe I'll need to create a new dataset inside hydragnn/serialized_dataset for this testing, so that data.forces and data.energy are available. It seems like most of the other data is generated randomly from [0.0, 0.5, 1.0].
Is this correct and is random generation is fine for the forces/energy as well?

@RylieWeaver I have an idea that I would like to run by you. What about adding a frozen version of the LJ code as an additional example, where the example first generates the data, saves it into files, and then uses it for training? We don't need to generate many data samples and traing it fully. We just need to make sure that the training kicks off correctly without errors.

@allaffa Yeah, I think this will be good to add LJ to the main branch. Would you be free to hammer out the specifics? (your calendar seems booked) If not, we can just hash out here what this will look like:

(1) So, what I think you're saying is to transfer the whole LJ example to the hydragnn/examples directory, including dataset creation, train.py, and json files. This will pull from both your fork and PyTorch_interatomic_potentials. It should be able to run similar to the qm9 example, where ____.py will detect if there is data and generate it if not.

(2) Then, would the examples/LJ dataset be simply copy-pasted over to hydragnn/serialized_dataset to use for examples, or is that unrelated? It would be 500 data points in line with the other examples

[allaffa]

@RylieWeaver
Additional changes that you forgot to address in the preview review of the PR

[allaffa]

@RylieWeaver
I committed just one tiny change on the year of the Copyright.

[RylieWeaver]

@RylieWeaver I committed just one tiny change on the year of the Copyright.

@allaffa Ok, cool, I just did a couple of cleanup things as well but it should pass all tests again.

[allaffa]

@pzhanggit @jychoi-hpc
I have been working with Rylie on this PR for quite some time.
Even if nominally all commits are authored by him, some portions of the code pushed in this PR were written by me in local scripts that I passed on to him.

Because of this, since I have been heavily involved in the code development of this PR myself, I would appreciate if you two could review this PR with fresh eyes. You may catch issues or concerns in terms of software design that I have not caught myself.

[RylieWeaver]

Double-clicking on this, let me know what you think @pzhanggit @jychoi-hpc

There are some more PRs in waiting that require this one first, so I'd like to get this one approved or revised soon if possible. Thanks.

[allaffa]

Double-clicking on this, let me know what you think @pzhanggit @jychoi-hpc

There are some more PRs in waiting that require this one first, so I'd like to get this one approved or revised soon if possible. Thanks.

@pzhanggit and @jychoi-hpc

For the interest of time, please focus on reviewing only the two files that change the core capabilities of hydragnn, namely:

• hydragnn/models/Base.py
• hydragnn/train/train_validate_test.py

All the other files are in the "examples" folder, therefore they do no affect the core of HydraGNN, and are the result of extensive work between me, Rylie, and Pranav anyway. The reason why I asked for you to review the PR is because of the two files above.

[jychoi-hpc]

Thank you for the update. I looked over the base and train script. It looks good to me.

[pzhanggit]

Why is grad_outputs used here? This would lead to sum of the gradients, right?

[RylieWeaver]

I don't think so.

Here's my understanding of it:
The grad_outputs argument essentially serves as a weight to multiply the gradients. This is necessary since graph_energy_pred is a higher dimension than just a scalar. The torch.ones_like() here will assign a weight of 1 to each gradient, which I believe is what we want. Otherwise, there would be a force prediction which is c*grad_E where c is a constant not equal to 1, which is nonphysical.

Specifically to your question:
No, I don't believe that it results in a sum. The output shape of forces_pred will be in the same shape as data.pos. The grad_outputs specifically is for weighting the gradients before a sum, but that extra step of summing is outside of autograd calculation I believe.

Does this answer your question?

@pzhanggit

[pzhanggit]

I don't think so.

Here's my understanding of it: The grad_outputs argument essentially serves as a weight to multiply the gradients. This is necessary since graph_energy_pred is a higher dimension than just a scalar. The torch.ones_like() here will assign a weight of 1 to each gradient, which I believe is what we want. Otherwise, there would be a force prediction which is c*grad_E where c is a constant not equal to 1, which is nonphysical.

Specifically to your question: No, I don't believe that it results in a sum. The output shape of forces_pred will be in the same shape as data.pos. The grad_outputs specifically is for weighting the gradients before a sum, but that extra step of summing is outside of autograd calculation I believe.

Does this answer your question?

@pzhanggit

Thanks. So just to verify if my understanding is correct: grad_outputs is needed because we're calculating gradients of a batch of samples here. And since grad_outputs should be a sequence of length matching output containing the “vector” in vector-Jacobian product, it provides a way to aggregate/sum the gradients across all the samples in the batch. This is equivalent to calculate the gradients iteratively for each sample, since the cross-gradients between samples would be zero.

[RylieWeaver]

@pzhanggit

grad_outputs is needed because we're calculating gradients of a batch of samples here

Yep, that's right.

And since grad_outputs should be a sequence of length matching output containing the “vector” in vector-Jacobian product

I think so. Yes, grad_outputs should be a sequence, and it should be the same shape as whatever you're predicting. It then multiplies the matrix of gradients (row-wise, not normal matrix multiplication), which I think is the vector-Jacobian product you're referring to.

it provides a way to aggregate/sum the gradients across all the samples in the batch.

I think yes. It's scaling those gradients, which would be relevant in an aggregation/sum. Although, it does not do that aggregation/sum.

This is equivalent to calculate the gradients iteratively for each sample, since the cross-gradients between samples would be zero.

Yep.

Are there any parts still unclear?

[pzhanggit]

@pzhanggit

grad_outputs is needed because we're calculating gradients of a batch of samples here

Yep, that's right.

And since grad_outputs should be a sequence of length matching output containing the “vector” in vector-Jacobian product

I think so. Yes, grad_outputs should be a sequence, and it should be the same shape as whatever you're predicting. It then multiplies the matrix of gradients (row-wise, not normal matrix multiplication), which I think is the vector-Jacobian product you're referring to.

it provides a way to aggregate/sum the gradients across all the samples in the batch.

I think yes. It's scaling those gradients, which would be relevant in an aggregation/sum. Although, it does not do that aggregation/sum.

This is equivalent to calculate the gradients iteratively for each sample, since the cross-gradients between samples would be zero.

Yep.

Are there any parts still unclear?

Scaling those gradients?

[allaffa]

@pzhanggit
The scaling is done with torch.ones_like(graph_energy_pred)
In our case, we use vectors of ones because we do NOT want to scale. However, you could apply any multiplying factor (or even provide a customized vector with different values for each entry).

[RylieWeaver]

@pzhanggit @allaffa ^ on what Max said. Also, feel free to @ me so In reply faster :)

[allaffa]

@pzhanggit @RylieWeaver

I personally found the explanation given by ChatGPT very useful. I attach it here hoping that you will find it useful too.

[RylieWeaver]

Yep. Wow, the o1-preview is great

[pzhanggit]

thanks.