From fa3a5b564234d229b8b58e370f37a9c774122582 Mon Sep 17 00:00:00 2001
From: Jinzhe Zeng <jinzhe.zeng@rutgers.edu>
Date: Mon, 13 Jun 2022 16:01:52 -0400
Subject: [PATCH] support reading xyz files

---
 dpdata/plugins/xyz.py | 18 ++++++++++++++++--
 dpdata/xyz/xyz.py     | 31 +++++++++++++++++++++++++++++++
 tests/test_xyz.py     | 21 ++++++++++++++++++++-
 3 files changed, 67 insertions(+), 3 deletions(-)

diff --git a/dpdata/plugins/xyz.py b/dpdata/plugins/xyz.py
index 5e2534a87..7ec9d2665 100644
--- a/dpdata/plugins/xyz.py
+++ b/dpdata/plugins/xyz.py
@@ -1,7 +1,7 @@
 import numpy as np
 
 from dpdata.xyz.quip_gap_xyz import QuipGapxyzSystems
-from dpdata.xyz.xyz import coord_to_xyz
+from dpdata.xyz.xyz import coord_to_xyz, xyz_to_coord
 from dpdata.format import Format
 
 @Format.register("xyz")
@@ -20,6 +20,20 @@ def to_system(self, data, file_name, **kwargs):
         with open(file_name, 'w') as fp:
             fp.write("\n".join(buff))
 
+    def from_system(self, file_name, **kwargs):
+        with open(file_name, 'r') as fp:
+            coords, types = xyz_to_coord(fp.read())
+        atom_names, atom_types, atom_numbs = np.unique(types,  return_inverse=True, return_counts=True)
+        return {
+            'atom_names': list(atom_names),
+            'atom_numbs': list(atom_numbs),
+            'atom_types': atom_types,
+            'coords': coords.reshape((1, *coords.shape)),
+            'cells': np.eye(3).reshape((1, 3, 3)) * 100,
+            'nopbc': True,
+            'orig': np.zeros(3),
+        }
+
 
 @Format.register("quip/gap/xyz")
 @Format.register("quip/gap/xyz_file")
@@ -29,4 +43,4 @@ def from_labeled_system(self, data, **kwargs):
 
     def from_multi_systems(self, file_name, **kwargs):
         # here directory is the file_name
-        return QuipGapxyzSystems(file_name)
\ No newline at end of file
+        return QuipGapxyzSystems(file_name)
diff --git a/dpdata/xyz/xyz.py b/dpdata/xyz/xyz.py
index 0bf1b6141..0ca5ac311 100644
--- a/dpdata/xyz/xyz.py
+++ b/dpdata/xyz/xyz.py
@@ -1,3 +1,5 @@
+from typing import Tuple
+
 import numpy as np
 
 def coord_to_xyz(coord: np.ndarray, types: list)->str:
@@ -26,3 +28,32 @@ def coord_to_xyz(coord: np.ndarray, types: list)->str:
     for at, cc in zip(types, coord):
         buff.append("{} {:.6f} {:.6f} {:.6f}".format(at, *cc))
     return "\n".join(buff)
+
+
+def xyz_to_coord(xyz: str) -> Tuple[np.ndarray, list]:
+    """Convert xyz format to coordinates and types.
+
+    Parameters
+    ----------
+    xyz : str
+        xyz format string
+
+    Returns
+    -------
+    coords : np.ndarray
+        coordinates, Nx3 array
+    types : list
+        list of types
+    """
+    symbols = []
+    coords = []
+    for ii, line in enumerate(xyz.split('\n')):
+        if ii == 0:
+            natoms = int(line.strip())
+        elif 2 <= ii <= 1 + natoms:
+            # symbol x y z
+            symbol, x, y, z = line.split()
+            coords.append((float(x), float(y), float(z)))
+            symbols.append(symbol)
+    return np.array(coords), symbols
+
diff --git a/tests/test_xyz.py b/tests/test_xyz.py
index 23bb3a81d..7d9649f94 100644
--- a/tests/test_xyz.py
+++ b/tests/test_xyz.py
@@ -2,8 +2,9 @@
 import numpy as np
 import tempfile
 from context import dpdata
+from comp_sys import CompSys, IsNoPBC
 
-class TestXYZ(unittest.TestCase):
+class TestToXYZ(unittest.TestCase):
     def test_to_xyz(self):
         with tempfile.NamedTemporaryFile('r') as f_xyz:
             dpdata.System(data={
@@ -17,3 +18,21 @@ def test_to_xyz(self):
             xyz0 = f_xyz.read().strip()
         xyz1 = "2\n\nC 0.000000 1.000000 2.000000\nO 3.000000 4.000000 5.000000"
         self.assertEqual(xyz0, xyz1)
+
+
+class TestFromXYZ(unittest.TestCase, CompSys, IsNoPBC):
+    def setUp(self):
+        self.places = 6
+        # considering to_xyz has been tested..
+        self.system_1 = dpdata.System(data={
+                "atom_names": ["C", "O"],
+                "atom_numbs": [1, 1],
+                "atom_types": np.array([0, 1]),
+                "coords": np.arange(6).reshape((1,2,3)),
+                "cells": np.zeros((1,3,3)),
+                "orig": np.zeros(3),
+                "nopbc": True,
+            })
+        with tempfile.NamedTemporaryFile('r') as f_xyz:
+            self.system_1.to("xyz", f_xyz.name)
+            self.system_2 = dpdata.System(f_xyz.name, fmt="xyz")