aiqm / torchani

Accurate Neural Network Potential on PyTorch
https://aiqm.github.io/torchani/
MIT License
446 stars 125 forks source link

Training with a batch of crystals #613

Closed rschireman closed 2 years ago

rschireman commented 2 years ago

Hi all,

I see that from #605, the PBC implementation in torchani is not batched, or only supports a batch size of 1. I have modified my database to include the 3x3 cell matrix for each structure and can successfully load species, coordinates, energies, forces, and cells with torchani.data.load.

But, now the dimensions when the batch size is not 1 are not compatible with functions in aev.py. I'm more than happy to contribute and submit a pull request, but I wanted to reach out to make sure this wasn't already being developed and to also ask for some guidance. I'm mostly stuck on how to make the returned quantity from compute_shifts(cell: Tensor, pbc: Tensor, cutoff: float) -> Tensor: compatible with this line:

aev = compute_aev(species, coordinates, self.triu_index, self.constants(), self.sizes, (cell, shifts))

Thank you, Ray