🌟 [FEATURE] Stress tensor about nequip HOT 15 CLOSED

@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
   

from nequip.

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

from nequip.

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 😄

from nequip.

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

from nequip.

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

from nequip.

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.

from nequip.

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

from nequip.

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.

from nequip.

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)

from nequip.

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

from nequip.

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.

from nequip.

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.

from nequip.

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.

from nequip.

@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.

from nequip.

@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?

from nequip.