🌟 [FEATURE] Stress tensor

[keano130]

Is your feature request related to a problem? Please describe.
In order to execute npt-simulations, the stress of the system is needed, I wonder if it is possible to implement the option to calculate the stress. To calculate this stress, the derivative of the energy wrt the cell is needed, In evaluation mode however, the model does not store the gradients and therefore, this value can not be computed. In training mode everything is rescaled, and using training mode in inference is not the purpose of training mode I think. Adding a cell and stress key in GradientOutput causes an error that there is no key 'cell' in irreps_in[wrt] in line 62 in _grad_output.py. As I figure that the cell should not be included in the irreps_in, this options seems bad too.

Describe the solution you'd like
Method on the final_model (RescaleOutput) or GradientOutput to calculate stress, if a stress parameter in the config file is True

Additional context
In schnetpack this is implemented in https://github.com/atomistic-machine-learning/schnetpack/blob/master/src/schnetpack/atomistic/output_modules.py and https://github.com/atomistic-machine-learning/schnetpack/blob//src/schnetpack/atomistic/model.py

[Linux-cpp-lisp]

Hi @keano130 ,

This should not be too hard to implement; the main thing is adding the displacement. (Since our edge vectors are calculated from positions + cell @ periodic_image_number, they should be connected correctly in the autograd graph to displacements if it is added to both positions and cell.)
This, and indicating to GradientOutput that a cell must be present, can be done in a short StressOutput module. I can write this up quickly if you are willing to help test it 😄

[keano130]

Thank you for the quick response, I am definitely willing to help test it.

[Linux-cpp-lisp]

@keano130 check out the https://github.com/mir-group/nequip/tree/stress branch, which contains a first stab at an implementation (untested so far)

[Linux-cpp-lisp]

Added some basic tests, the model does run and produce numbers for the stress when PBC are used. (See minimal-stress.yaml.) Whether they're reasonable I can't really say.

One thing I'm not sure about is the multiplication conventions at https://github.com/mir-group/nequip/blob/stress/nequip/nn/_grad_output.py#L187. I had to modify them from the schnetpack version, particularly for the positions, so I'm not sure if they make sense. Let me know what you think, and if the code works for you.

For running simple NPT simulations, the easiest thing to do would be to add the stress to the ASE calculator (dynamics/nequip_calculator.py). This should be a pretty simple addition just involving some units stuff, and I'd happily accept a PR onto the stress branch adding it.

[keano130]

The stress_key should to be added to the rescaleoutput layer , I tested the stresses using finite difference and this test passes correctly.

Stress_dir.zip

[keano130]

Furthermore, I compared the stress predictions of the network trained on ab initio (dft) generated data with the stress predictions of dft and they seemed to be relatively close (relative mean absolute error of 30%) which is very possible as the network was not yet fully trained.

[Linux-cpp-lisp]

Great! That's partially trained on forces, predicting stresses?

Looking at stress_dir.zip I'm not sure I see the finite difference test you're talking about...

EDIT: I've added stress to the rescaling

[keano130]

Oh that's my bad, somehow i uploaded everything but the python file (I'm don't have my laptop with me, so I will upload it tomorrow, sorry for the delay)

[keano130]

Stress_test.zip
The 30% error was indeed trained on forces, predicting stress

[Linux-cpp-lisp]

Hi @keano130 ,

Wanted to let you know that stress is now being implemented in the stress-develop branch off of the latest develop — see PR #119. We hope to merge this soon after some more testing. The interface is similar to the previous stress branch, but is cleaned up and updated to be based on the latest code.

[Linux-cpp-lisp]

Hi @keano130,

Just a heads up that I've gone ahead and merged the (now much more complete) stress-develop branch into develop to end the merge hell that's starting to accumulate on that branch. Hopefully this makes it easier to use and test.

That said, stress support is still very much in beta and, as the warning I added says, please sanity check your results and let us know if you find things that seem strange or wrong.

[davidleocadio]

Hi @Linux-cpp-lisp . I've checked out the develop branch and tried calculating stress with two different methods, both giving me a Key Error: 'stress' suggesting stress isn't an available option. Could you point me on how to calculate stress with a NequIP calculator? The most straightforward way I've tried it, which I thought should work is

`NN=nequip_calculator.NequIPCalculator
NN=NN.from_deployed_model(model_path="./deployedstress.pth", \
species_to_type_name = {
    "Cu": "Cu",
})

CU=Atoms('Cu',positions=[(0, 0, 0),\
                                    (0.0*a, 0.5*a, 0.5*a),\
                                    (0.5*a, 0*a, 0.5*a),\
                                    (0.5*a, 0.5*a, 0*a)],\
                                        cell=[1*a,1*a,1*a],pbc=True)

NN.calculate(atoms=CU,properties=["energy","forces","stress"])

print(NN.results)

Giving an error such as

    stress = out[AtomicDataDict.STRESS_KEY].detach().cpu().numpy()
KeyError: 'stress'

Similarily adding the calculator to an Atoms object and running get_stress() doesn't work with the same error.

[Linux-cpp-lisp]

@davidleocadio does your model include a StressForceOutput? The model must be built to calculate stresses for the calculator to be able to get those outputs. See minimal_stress.yaml for an example of this in model_builders.

The good news is that if it doesn't, you can still reuse your trained weights an build a new model with a StressForceOutput that just loads the weights from your trained model without it. If this is the case let me know and I will show you how.

[davidleocadio]

@Linux-cpp-lisp great thanks. I wasn't aware of minimal_stress.yaml. However, I think the second option you mention where one uses the weights from a previously trained model would be very useful. Would you mind showing how one should do this?

[Linux-cpp-lisp]

@davidleocadio it'll be something like this:

1. pull the latest develop--- I renamed something in StressOutput to make this more consistant going forward. (Warning: this will break loading models trained with StressOutput from before the change, but it's easy to fix if you need to--- let me know if you do.)

2. make a my-model-with-stress.yaml that is a copy of your original config file

3. add the model_builders from minimal_stress.yaml and the load_model_state builder:

    model_builders:
      - SimpleIrrepsConfig
      - EnergyModel
      - PerSpeciesRescale
      - StressForceOutput
      - RescaleEnergyEtc
      - load_model_state
    
    load_model_state: /path/to/my/training_session/best_model.pth
   

   the load_model_state builder takes the previously constructed model---which is the same as the one you trained, except substituting StressForceOutput for ForceOutput---and loads the weights from the state_dict found at the path given by the load_model_state option.

4. Deploy my-model-with-stress.yaml:

   $ nequip-deploy build --model my-model-with-stress.yaml  deployed-stress.pth