Fix order of wrapping and strain (#221)

torchani/ase.py

@@ -77,6 +77,12 @@ class Calculator(ase.calculators.calculator.Calculator):
         coordinates = torch.tensor(self.atoms.get_positions())
         coordinates = coordinates.unsqueeze(0).to(self.device).to(self.dtype) \
                                  .requires_grad_('forces' in properties)
+
+        if pbc_enabled:
+            coordinates = utils.map2central(cell, coordinates, pbc)
+            if self.overwrite and atoms is not None:
+                atoms.set_positions(coordinates.detach().cpu().reshape(-1, 3).numpy())
+
         if 'stress' in properties:
             displacements = torch.zeros(3, 3, requires_grad=True,
                                         dtype=self.dtype, device=self.device)
@@ -87,9 +93,6 @@ class Calculator(ase.calculators.calculator.Calculator):
             coordinates = coordinates + strain_x + strain_y + strain_z
 
         if pbc_enabled:
-            coordinates = utils.map2central(cell, coordinates, pbc)
-            if self.overwrite and atoms is not None:
-                atoms.set_positions(coordinates.detach().cpu().reshape(-1, 3).numpy())
             if 'stress' in properties:
                 strain_x = self.strain(cell, displacement_x, 0)
                 strain_y = self.strain(cell, displacement_y, 1)