From 6f001f197d98826b92b34b96c9cd51e375cae5a1 Mon Sep 17 00:00:00 2001
From: Davis Rens <166565828+davrens101@users.noreply.github.com>
Date: Sat, 7 Nov 2026 23:52:31 -0800
Subject: [PATCH] Preserve nonperiodic coordinates in graph conversion

---
 README.md                                 |  2 +-
 src/mattersim/datasets/utils/converter.py |  4 +-
 tests/datasets/utils/test_converter.py    | 45 +++++++++++++++++++++++
 3 files changed, 47 insertions(+), 4 deletions(-)

diff --git a/README.md b/README.md
index 28f9246..312369c 100644
--- a/README.md
+++ b/README.md
@@ -153,7 +153,7 @@ We kindly request that users of MatterSim version 1.0.0 cite our preprint availa
 
 ## Limitations
 
-**MatterSim-v1** is designed specifically for atomistic simulations of bulk materials. Applications or interpretations beyond this scope should be approached with caution. For instance, when using the model for simulations involving surfaces, interfaces, or properties influenced by long-range interactions, the results may be qualitatively accurate but are not suitable for quantitative analysis. In such cases, we recommend fine-tuning the model to better align with the specific application.
+**MatterSim-v1** is designed specifically for atomistic simulations of bulk materials. Applications or interpretations beyond this scope should be approached with caution. For instance, when using the model for simulations involving surfaces, interfaces, or properties influenced by long-range interactions, the results may be qualitatively accurate but are not suitable for quantitative analysis. For mixed-periodicity inputs, graph construction wraps coordinates only along periodic cell directions; this does not establish model accuracy for non-bulk systems. In such cases, we recommend fine-tuning the model to better align with the specific application.
 
 ## Trademarks
 
diff --git a/src/mattersim/datasets/utils/converter.py b/src/mattersim/datasets/utils/converter.py
index 80a75e4..4ff0656 100644
--- a/src/mattersim/datasets/utils/converter.py
+++ b/src/mattersim/datasets/utils/converter.py
@@ -213,9 +213,7 @@ def convert(
         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)
+        atoms.wrap()
         args = {}
         if self.model_type == "m3gnet":
             args["num_atoms"] = len(atoms)
diff --git a/tests/datasets/utils/test_converter.py b/tests/datasets/utils/test_converter.py
index 8d31ec3..635c1de 100644
--- a/tests/datasets/utils/test_converter.py
+++ b/tests/datasets/utils/test_converter.py
@@ -4,7 +4,9 @@
 """
 
 import numpy as np
+import pytest
 import torch
+from ase import Atoms
 from torch_geometric.data import Batch
 
 from mattersim.datasets.utils.converter import (
@@ -83,3 +85,46 @@ def test_pbc_fallback_uses_int64(self, water_molecule):
         converter = GraphConverter(model_type="m3gnet")
         graph = converter.convert(water_molecule)
         assert graph is not None
+
+    @pytest.mark.parametrize(
+        "pbc",
+        [
+            [True, True, False],
+            [True, False, False],
+        ],
+        ids=["slab_TTF", "wire_TFF"],
+    )
+    def test_mixed_pbc_wrapping_preserves_nonperiodic_coordinates(self, pbc):
+        """Wrapping must not create bonds across nonperiodic cell boundaries."""
+        nonperiodic_axis = pbc.index(False)
+        positions = np.zeros((3, 3), dtype=np.float64)
+        positions[:, 0] = 12.25
+        positions[:, nonperiodic_axis] = [1.0, 1.5, 11.2]
+
+        atoms = Atoms("Si3", positions=positions, cell=[10.0, 10.0, 10.0], pbc=pbc)
+        converter = GraphConverter(
+            model_type="m3gnet", twobody_cutoff=2.0, threebody_cutoff=2.0
+        )
+        graph = converter.convert(atoms)
+
+        np.testing.assert_allclose(
+            graph.atom_pos[:, nonperiodic_axis].detach().cpu().numpy(),
+            [1.0, 1.5, 11.2],
+            rtol=1e-5,
+            atol=1e-5,
+        )
+
+        np.testing.assert_allclose(
+            graph.atom_pos[:, 0].detach().cpu().numpy(),
+            2.25,
+            rtol=1e-5,
+            atol=1e-5,
+        )
+
+        # Exact edge pairs: only (0, 1) and (1, 0) within 2 A cutoff
+        edge_pairs = set(
+            zip(graph.edge_index[0].tolist(), graph.edge_index[1].tolist())
+        )
+        assert edge_pairs == {(0, 1), (1, 0)}
+
+        assert graph.num_three_body == 0