Stress tensor

[jduerholt]

Dear schnetpack developers,

is it planned to also implement the calculation of the stress tensor in schnetpack? This would enable to perform NPT simulations.

Best

Johannes

[mgastegger]

Hi Johannes, we plan to implement periodic boundary conditions and NPT simulations at one point in the future, so there will be stress tensors. This will take some time, but the standard static stress tensor should be straightforward. I will have a look.

[mgastegger]

Hello Johannes, we finally added the stress tensor functionality to the atomistic layer. The TorchEnvironmentProvider should be used in order to obtain the right behavior for periodic boundary conditions. It is also possible to use the stress tensor for training with spk_run.py by specifying the --stress flag.

[ruoitrau86]

Hi, Is there any update soon for this implementation for the ASE interface? as I see that Implemented_properties = [energy, forces] only at this time.

[jduerholt]

Thank you very much for the implementation and the possibility to use it also for training. I will start using it soon!

[mgastegger]

I will have a look to add the stress to the ASE calculator.

[jduerholt]

One further question: You mention that the torch environment provider should be used to obtain the right behavior for PBCs. From my understanding the same should also be achieved by using the ASE environment provider, or?

[mgastegger]

The ASE calculator should now be able to handle the stress tensor.

[jduerholt]

I had a look on the implementation of the stress tensor in the ASE. The following units are specified

stress_units="eV/Angstrom/Angstrom"

I am now a bit puzzled since in general a stress tensor has units of energy/volume like a pressure. Also ASE is using this definition from my understanding. Can you help me here?

[mgastegger]

First of all, you are right, the stress tensor units should be Energy/Length^3. I fixed this in the branch. The argument stress_units specifies the units used by the ML model. These are then used to convert the predictions to the (hopefully) correct ASE units. One could also e.g. have kcal/mol as energy units.

[jduerholt]

Thank you very much. Could you perhaps give me a brief explanation how you calculate the tensor? I see, that you differentiate with respect to displacements. What does this mean?

[mgastegger]

We model the stress tensor as the derivative of the potential energy with respect to the strain (displacement) on the unit cell, which is then divided by the volume (see e.g. Knuth et. al. Comput. Phys. Commun 190, 33-50, 2015): a and b are the Cartesian components e.g. xx, xy, xz, ...

To do so, we define the displacements as 3 by 3 tensor initialized to 0 (displacement) and apply it to the coordinates as well as the cell. The whole model then builds on the distances computed in this manner. When evaluating the model, we take the derivative of the predicted energy with respect to the displacements and divide it by the volume of the cell, which should give the stress tensor.

[Dankomaister]

This is a very interesting feature, any update on the progress to including this in the master branch?

[mgastegger]

The updated ASE calculator is now in the master branch.

[Dankomaister]

Great!

I tried evaluating the stress tensor on my dataset of periodic structures (diamond and lonsdaleite) and after only a few epochs (training on energy and forces) I was able to reach a MAE in the stress tensor below 0.1 eV/Å^3 (interestingly the MAE in the stress tensor show sign of overfitting while the energy and force MAE are continuously decreasing).

I did however run into a problem. I’m using ASE to create my dataset by reading results from VASP (vasprun.xml). The issue is that ASE only saves the symmetric stress tensor in Voigt order (xx, yy, zz, yz, xz, xy) not the full stress tensor (which I assume schnetpack is modelling?).

I solved that problem by simply replacing the atoms._calc.results['stress'] array with the full 3x3 stress tensor read from vasprun.xml but I ran into another issue. When ASE saves the atoms._calc.results['stress'] array to the .db file it flattens the array. This is a problem since the dimensionality of the stress tensor from schnetpack (3x3) is then different from what is read from the .db file (9x1) which results in errors when calculating metrics (MeanAbsoluteError).

This can be solved by simply reshaping the stress array in the add_batch function in the MeanAbsoluteError metric class. However, I feel that this is not the best solution, does anyone have a better solution for getting the full stress tensor correctly written to the .db file using ASE?

[mgastegger]

Glad to hear the stress tensor seems to work. We use the full (symmetric) 3x3 tensor in SchNetPack.

Regarding your problem, if you just want to generate the database for training, you could also create a custom AtomsData dataset, where you add the stress tensor as an extra property. E.g. by adding to the properties dict in add_system(atoms, **properties). In general, we store all the properties into the additional data entry of the ASE db, since we encountered similar problems to yours when trying to use the ASE Atoms functionality.

I don't know how helpful this is, but it seems like one can disable the Voigt order in atoms.get_stress via voigt=False.

And it also seems that ASE expects the Voigt stress convention in the calculator output, so I will have to change this in the calculator.

[Dankomaister]

Thanks for the suggestion!

I will think about what the best solution is. Since SchNetPack returns the 3x3 symmetric stress tensor I can just convert it to an 1D array in Voigt order. Then it is comparable with the stress as read from VASP by ASE if I just convert the unit also… Or maybe I will write a custom function to read the properties from VASP and store them into the additional data entry of the ASE db like you suggested. Then I don't have to do any conversion or change the SchNetPack code.

[Dankomaister]

Hi, I have encountered a problem when calculating stress using different batch sizes.

When I use a batch size of 1 then the results are close to the true values (5-25% error) which is expected since the training data only contains structures with small stress. However, when I increase the batch size the stress for certain structures becomes infinite. This is especially bad when a batch size of 3 is used, then the calculated stress for all structures becomes +-inf.

My guess is that this has something to do with the padding causing a zero volume? But I don't really see how the stress for all structures can affected, my thinking is that only structures in the last batch should be affected if it is a padding problem.

Any ideas what could cause this?

[mgastegger]

Hm. Are you sure this is a problem of the batch size? Issues with the padding should in principle only happen, when the molecule size changes withing a batch. And the volume is computed based on the cell, which is never padded (since it is always 3x3). There could be a problem with the broadcasting somewhere, though. I will have a look. Otherwise I would check whether there are data points with all zeros unit cells. This could also cause the zero volume division.

[Dankomaister]

Double checked, and yes there is definitely something strange going on when I change the batch size. What I mean with padding is that the number of total images is not even divisible with the batch size. But also for the testing data I'm using the number of atoms in the images also changes.

Let me see if i can make an MWE and upload here.

[Dankomaister]

MWE_stress_bug.zip Ok, check if this MWE works!

Now when I think about it, I also had some problem when training, where the loss function (which includes the stress) became inf for certain batch sizes. This I solved by setting drop_last=True in the AtomsLoader however this is not an ideal solution when evaluating against testing data :)

[mgastegger]

Thank you very much for the example, it helped nailing down the issue. There was a problem with the computation of the volume when the stress tensor was evaluated. It should work now.

[Dankomaister]

Great happy to help! schnetpack keeps getting better and better.