I found that the current implementation has several performance issues regarding gradient wrt. cell. This PR modifies that. Since the changes are relatively much, I will put some comments.

Change summary:

• Use shift for gradient instead of cell.
• shift is now length scale instead cell unit.
• Calculate Voigt notation style stress directly

Also, this PR contains bugfix related to sked cell.

When the length of atoms is 1, the routine raises error.

from ase.build import molecule
from ase.calculators.dftd3 import DFTD3
from torch_dftd.torch_dftd3_calculator import TorchDFTD3Calculator

if __name__ == "__main__":
    atoms = molecule("H")
    # device="cuda:0" for fast GPU computation.
    calc = TorchDFTD3Calculator(atoms=atoms, device="cpu", damping="bj")

    energy = atoms.get_potential_energy()
    forces = atoms.get_forces()

    print(f"energy {energy} eV")
    print(f"forces {forces}")

Traceback (most recent call last):
  File "quick.py", line 12, in <module>
    energy = atoms.get_potential_energy()
  File "/home/ahayashi/envs/dftd/lib/python3.8/site-packages/ase/atoms.py", line 731, in get_potential_energy
    energy = self._calc.get_potential_energy(self)
  File "/home/ahayashi/envs/dftd/lib/python3.8/site-packages/ase/calculators/calculator.py", line 709, in get_potential_energy
    energy = self.get_property('energy', atoms)
  File "/home/ahayashi/torch-dftd/torch_dftd/torch_dftd3_calculator.py", line 141, in get_property
    dftd3_result = Calculator.get_property(self, name, atoms, allow_calculation)
  File "/home/ahayashi/envs/dftd/lib/python3.8/site-packages/ase/calculators/calculator.py", line 737, in get_property
    self.calculate(atoms, [name], system_changes)
  File "/home/ahayashi/torch-dftd/torch_dftd/torch_dftd3_calculator.py", line 119, in calculate
    results = self.dftd_module.calc_energy(**input_dicts, damping=self.damping)[0]
  File "/home/ahayashi/torch-dftd/torch_dftd/nn/base_dftd_module.py", line 75, in calc_energy
    E_disp = self.calc_energy_batch(
  File "/home/ahayashi/torch-dftd/torch_dftd/nn/dftd3_module.py", line 86, in calc_energy_batch
    E_disp = d3_autoev * edisp(
  File "/home/ahayashi/torch-dftd/torch_dftd/functions/dftd3.py", line 189, in edisp
    c6 = _getc6(Zi, Zj, nci, ncj, c6ab=c6ab, k3=k3)  # c6 coefficients
  File "/home/ahayashi/torch-dftd/torch_dftd/functions/dftd3.py", line 97, in _getc6
    k3_rnc = torch.where(cn0 > 0.0, k3 * r, -1.0e20 * torch.ones_like(r)).view(n_edges, -1)
RuntimeError: cannot reshape tensor of 0 elements into shape [0, -1] because the unspecified dimension size -1 can be any value and is ambiguous

I found that the current implementation has several performance issues regarding gradient wrt. cell. This PR modifies it. Since the changes are relatively much, I will put some comments.

Change summary:

• Use shift for gradient instead of cell.
• shift is now length scale instead cell unit.
• Calculate Voigt notation style stress directly

Also, this PR contains bugfix related to sked cell.

I found that test function test_calc_energy_force_stress_device_batch_abc unintentionally ignores abc argument.

This PR modified related implementation to work it.

In addition, corner case correspondence when the total number of atom is zero is also added. (n_graphs cannot be calculated from batch_edge when len(batch_edge) == 0.)

The _gettc6 routine now works correctly even when the number of adjacencies is 0. Instead of calling calc_neighbor_by_pymatgen when the number of atoms is 0 and the periodic boundary condition, it now return edge_index, S for adjacency 0. In my environment, using the result of torch.sum for the size of torch.zeros caused an error, so I changed it to cast the result of sum to int.

In test_torch_dftd3_calculator.py/_assert_energy_force_stress_equal, there is a code below.

    if np.all(atoms.pbc == np.array([True, True, True])):
        s1 = atoms.get_stress()
        s2 = atoms.get_stress()
        assert np.allclose(s1, s2, atol=1e-5, rtol=1e-5)

This code cannot compare the results of stresses of calc1 and calc2. Both s1 and s2 are the stress of calc2.