From eeb275c52c2b5d0a6652508450eda265d034ab0f Mon Sep 17 00:00:00 2001
From: Xixian Liu <104359500+ZeroKnighting@users.noreply.github.com>
Date: Wed, 4 Dec 2024 22:40:17 +0800
Subject: [PATCH] fix: non-pbc handle (#50)

* fix: non-pbc handle

* fix: zero cell detect

* fix bugs

---------

Co-authored-by: Xixian <v-xixianliu@microsoft.com>
---
 src/mattersim/datasets/utils/convertor.py | 53 +++++++++++++++--------
 1 file changed, 36 insertions(+), 17 deletions(-)

diff --git a/src/mattersim/datasets/utils/convertor.py b/src/mattersim/datasets/utils/convertor.py
index ef85821..dc36687 100644
--- a/src/mattersim/datasets/utils/convertor.py
+++ b/src/mattersim/datasets/utils/convertor.py
@@ -1,4 +1,5 @@
 # -*- coding: utf-8 -*-
+import warnings
 from typing import Optional, Tuple
 
 import ase
@@ -10,6 +11,9 @@
 
 from .threebody_indices import compute_threebody as _compute_threebody
 
+# Ensure the warning is only shown once
+warnings.filterwarnings("once", category=UserWarning)
+
 """
 Supported Properties:
     - "num_nodes"(set by default)  ## int
@@ -110,24 +114,13 @@ def get_fixed_radius_bonding(
     Returns:
         center_indices, neighbor_indices, images, distances
     """
-    if isinstance(structure, Atoms):
-        pbc_ = np.array(structure.pbc, dtype=int)
-        if np.all(pbc_ < 0.1) or not pbc:
-            lattice_matrix = np.array(
-                [[1000.0, 0.0, 0.0], [0.0, 1000.0, 0.0], [0.0, 0.0, 1000.0]],
-                dtype=float,
-            )
-            pbc_ = np.array([0, 0, 0], dtype=int)
-        else:
-            lattice_matrix = np.ascontiguousarray(
-                structure.cell[:], dtype=float
-            )  # noqa: E501
+    pbc_ = np.array(structure.pbc, dtype=int)
 
-        cart_coords = np.ascontiguousarray(
-            np.array(structure.positions), dtype=float
-        )  # noqa: E501
-    else:
-        raise ValueError("structure type not supported")
+    lattice_matrix = np.ascontiguousarray(structure.cell[:], dtype=float)  # noqa: E501
+
+    cart_coords = np.ascontiguousarray(
+        np.array(structure.positions), dtype=float
+    )  # noqa: E501
     r = float(cutoff)
 
     (
@@ -192,6 +185,32 @@ def convert(
             pbc: bool, whether to use periodic boundary condition, default True
         """
         # normalize the structure
+        if isinstance(atoms, Atoms):
+            pbc_ = np.array(atoms.pbc, dtype=int)
+            if np.all(pbc_ < 0.01) or not pbc:
+                min_x = np.min(atoms.positions[:, 0])
+                min_y = np.min(atoms.positions[:, 1])
+                min_z = np.min(atoms.positions[:, 2])
+                max_x = np.max(atoms.positions[:, 0])
+                max_y = np.max(atoms.positions[:, 1])
+                max_z = np.max(atoms.positions[:, 2])
+                x_len = max((max_x - min_x) * 10, 1000)
+                y_len = max((max_y - min_y) * 10, 1000)
+                z_len = max((max_z - min_z) * 10, 1000)
+                lattice_matrix = np.array(
+                    [[x_len, 0.0, 0.0], [0.0, y_len, 0.0], [0.0, 0.0, z_len]],
+                    dtype=float,
+                )
+                pbc_ = np.array([1, 1, 1], dtype=int)
+                warnings.warn("No PBC detected, using a large supercell", UserWarning)
+                atoms.set_cell(lattice_matrix)
+                atoms.set_pbc(pbc_)
+            else:
+                if np.all(atoms.cell < 1e-5):
+                    raise ValueError("Cell vectors are too small")
+        else:
+            raise ValueError("structure type not supported")
+
         scaled_pos = atoms.get_scaled_positions()
         scaled_pos = np.mod(scaled_pos, 1)
         atoms.set_scaled_positions(scaled_pos)