correct pbc extension

mace/data/neighborhood.py

@@ -27,17 +27,18 @@ def get_neighborhood(
     max_positions = np.max(np.absolute(positions)) + 1
     # Extend cell in non-periodic directions
     # For models with more than 5 layers, the multiplicative constant needs to be increased.
+    temp_cell = np.copy(cell)
     if not pbc_x:
-        cell[:, 0] = max_positions * 5 * cutoff * identity[:, 0]
+        temp_cell[0, :] = max_positions * 5 * cutoff * identity[0, :]
     if not pbc_y:
-        cell[:, 1] = max_positions * 5 * cutoff * identity[:, 1]
+        temp_cell[1, :] = max_positions * 5 * cutoff * identity[1, :]
     if not pbc_z:
-        cell[:, 2] = max_positions * 5 * cutoff * identity[:, 2]
+        temp_cell[2, :] = max_positions * 5 * cutoff * identity[2, :]
 
     sender, receiver, unit_shifts = neighbour_list(
         quantities="ijS",
         pbc=pbc,
-        cell=cell,
+        cell=temp_cell,
         positions=positions,
         cutoff=cutoff,
         # self_interaction=True,  # we want edges from atom to itself in different periodic images

mace/data/utils.py

@@ -242,20 +242,20 @@ def load_from_xyz(
         atoms_without_iso_atoms = []
 
         for idx, atoms in enumerate(atoms_list):
-            isolated_atom_config = len(atoms) == 1 and atoms.info.get("config_type") == "IsolatedAtom"
+            isolated_atom_config = (
+                len(atoms) == 1 and atoms.info.get("config_type") == "IsolatedAtom"
+            )
             if isolated_atom_config:
                 if energy_key in atoms.info.keys():
-                    atomic_energies_dict[atoms.get_atomic_numbers()[0]] = (
-                        atoms.info[energy_key]
-                    )
+                    atomic_energies_dict[atoms.get_atomic_numbers()[0]] = atoms.info[
+                        energy_key
+                    ]
                 else:
                     logging.warning(
                         f"Configuration '{idx}' is marked as 'IsolatedAtom' "
                         "but does not contain an energy. Zero energy will be used."
                     )
-                    atomic_energies_dict[atoms.get_atomic_numbers()[0]] = np.zeros(
-                        1
-                    )
+                    atomic_energies_dict[atoms.get_atomic_numbers()[0]] = np.zeros(1)
             else:
                 atoms_without_iso_atoms.append(atoms)
 

mace/modules/utils.py

@@ -59,12 +59,9 @@ def compute_forces_virials(
     stress = torch.zeros_like(displacement)
     if compute_stress and virials is not None:
         cell = cell.view(-1, 3, 3)
-        volume = torch.einsum(
-            "zi,zi->z",
-            cell[:, 0, :],
-            torch.cross(cell[:, 1, :], cell[:, 2, :], dim=1),
-        ).unsqueeze(-1)
+        volume = torch.linalg.det(cell).abs().unsqueeze(-1)
         stress = virials / volume.view(-1, 1, 1)
+        stress = torch.where(torch.abs(stress) < 1e10, stress, torch.zeros_like(stress))
     if forces is None:
         forces = torch.zeros_like(positions)
     if virials is None: