Adding stress

[JonathanSchmidt1]

Thank you for the great code base already. For our application we would need to also calculate derivatives with respect to strain. As I am not that familiar with your code yet I wanted to ask whether you see any pitfalls when implementing derivatives with respect to strain.

the standard way to do it would be in this fashion here (this is from nequip):

displacement = torch.zeros(
            (num_batch, 3, 3),
            dtype=pos.dtype,
            device=pos.device,
        )
        displacement.requires_grad_(True)
symmetric_displacement = 0.5 * (displacement + displacement.transpose(-1, -2))
 pos = pos + torch.bmm(
            pos.unsqueeze(-2), symmetric_displacement[batch]
        ).squeeze(-2)
        # we only displace the cell if we have one:
        if has_cell:
            # bmm is num_batch in batch
            # here we apply the distortion to the cell as well
            # this is critical also for the correctness
            # if we didn't symmetrize the distortion, since without this
            # there would then be an infinitesimal rotation of the positions
            # but not cell, and it thus wouldn't be global and have
            # no effect due to equivariance/invariance.
            cell = cell + torch.bmm(
                cell, symmetric_displacement
            )

        # Call model and get gradients
        data = self.func(data)

        grads = torch.autograd.grad(
            [data[AtomicDataDict.TOTAL_ENERGY_KEY].sum()],
            [pos, data["_displacement"]],
            create_graph=self.training,  # needed to allow gradients of this output during training
        )

        # Put negative sign on forces
        forces = grads[0]
        if forces is None:
            # condition needed to unwrap optional for torchscript
            assert False, "failed to compute forces autograd"
        forces = torch.neg(forces)
        data[AtomicDataDict.FORCE_KEY] = forces

        # Store virial
        virial = grads[1]

A small add-on question. Did you ever try to get higher order derivatives for other responses?

[zhangylch]

Thank you for the great code base already. For our application we would need to also calculate derivatives with respect to strain. As I am not that familiar with your code yet I wanted to ask whether you see any pitfalls when implementing derivatives with respect to strain.

the standard way to do it would be in this fashion here (this is from nequip):

displacement = torch.zeros(
            (num_batch, 3, 3),
            dtype=pos.dtype,
            device=pos.device,
        )
        displacement.requires_grad_(True)
symmetric_displacement = 0.5 * (displacement + displacement.transpose(-1, -2))
 pos = pos + torch.bmm(
            pos.unsqueeze(-2), symmetric_displacement[batch]
        ).squeeze(-2)
        # we only displace the cell if we have one:
        if has_cell:
            # bmm is num_batch in batch
            # here we apply the distortion to the cell as well
            # this is critical also for the correctness
            # if we didn't symmetrize the distortion, since without this
            # there would then be an infinitesimal rotation of the positions
            # but not cell, and it thus wouldn't be global and have
            # no effect due to equivariance/invariance.
            cell = cell + torch.bmm(
                cell, symmetric_displacement
            )

        # Call model and get gradients
        data = self.func(data)

        grads = torch.autograd.grad(
            [data[AtomicDataDict.TOTAL_ENERGY_KEY].sum()],
            [pos, data["_displacement"]],
            create_graph=self.training,  # needed to allow gradients of this output during training
        )

        # Put negative sign on forces
        forces = grads[0]
        if forces is None:
            # condition needed to unwrap optional for torchscript
            assert False, "failed to compute forces autograd"
        forces = torch.neg(forces)
        data[AtomicDataDict.FORCE_KEY] = forces

        # Store virial
        virial = grads[1]

A small add-on question. Did you ever try to get higher order derivatives for other responses?

Thank you for your interest in our code. Are you trying to ask if we include the calculation of virial tensor?

Concerning the add question, we have ever got the three-order response (hyper-polarzability) for the NMA if I remember correctly.

[JonathanSchmidt1]

I could not find it in the code so I assumed it was not implemented? So I wanted to know if you think it's easy to add in the suggested fashion?

Concerning the add question, we have ever got the three-order response (hyper-polarzability) for the NMA if I remember correctly.

Nice

[zhangylch]

I could not find it in the code so I assumed it was not implemented? So I wanted to know if you think it's easy to add in the suggested fashion?

Yes, these codes do not support outputting virial tensors. Although we are considering adding that functionality to REANN and FIREANN, it will take some time to complete it. We need to find the suitable way to do it. How urgent is your need?

[JonathanSchmidt1]

The project is supposed to start middle of January start of February. Then we will have to decide on a framework for the force field. As we are looking for response functions there are not many options and FIREANN would definitely be the best if we have stress.

[zhangylch]

The project is supposed to start middle of January start of February. Then we will have to decide on a framework for the force field. As we are looking for response functions there are not many options and FIREANN would definitely be the best if we have stress.

Thank you for your interest. I will keep you informed when we have update on this implementation.

[zhangylch]

The project is supposed to start middle of January start of February. Then we will have to decide on a framework for the force field. As we are looking for response functions there are not many options and FIREANN would definitely be the best if we have stress.

Hello, I have added the calculation of stress tensor in the branch "Equi-FIREANN" and there is an example showing how to ouput the stress in the pes/example/test.py. Hopefully it would be helpful for your project! If there is any problem with the codes, please let me know.

[JonathanSchmidt1]

Thank you so much