From 1c0c0906447fc695a2dfb196f938465411f7f4cd Mon Sep 17 00:00:00 2001
From: Jinzhe Zeng <jinzhe.zeng@rutgers.edu>
Date: Thu, 20 Jan 2022 18:26:58 -0500
Subject: [PATCH] add xyz format (to)

This commit adds support for converting data to pure XYZ format. It's a quick way to visualize data using VMD or PyMol - these two software do support XYZ format file.
---
 dpdata/plugins/xyz.py | 19 +++++++++++++++++++
 dpdata/xyz/xyz.py     | 28 ++++++++++++++++++++++++++++
 tests/test_xyz.py     | 19 +++++++++++++++++++
 3 files changed, 66 insertions(+)
 create mode 100644 dpdata/xyz/xyz.py
 create mode 100644 tests/test_xyz.py

diff --git a/dpdata/plugins/xyz.py b/dpdata/plugins/xyz.py
index 71bf92683..5e2534a87 100644
--- a/dpdata/plugins/xyz.py
+++ b/dpdata/plugins/xyz.py
@@ -1,6 +1,25 @@
+import numpy as np
+
 from dpdata.xyz.quip_gap_xyz import QuipGapxyzSystems
+from dpdata.xyz.xyz import coord_to_xyz
 from dpdata.format import Format
 
+@Format.register("xyz")
+class XYZFormat(Format):
+    """XYZ foramt.
+
+    Examples
+    --------
+    >>> s.to("xyz", "a.xyz")
+    """
+    def to_system(self, data, file_name, **kwargs):
+        buff = []
+        types = np.array(data['atom_names'])[data['atom_types']]
+        for cc in data['coords']:
+            buff.append(coord_to_xyz(cc, types))
+        with open(file_name, 'w') as fp:
+            fp.write("\n".join(buff))
+
 
 @Format.register("quip/gap/xyz")
 @Format.register("quip/gap/xyz_file")
diff --git a/dpdata/xyz/xyz.py b/dpdata/xyz/xyz.py
new file mode 100644
index 000000000..0bf1b6141
--- /dev/null
+++ b/dpdata/xyz/xyz.py
@@ -0,0 +1,28 @@
+import numpy as np
+
+def coord_to_xyz(coord: np.ndarray, types: list)->str:
+    """Convert coordinates and types to xyz format.
+    
+    Parameters
+    ----------
+    coord: np.ndarray
+        coordinates, Nx3 array
+    types: list
+        list of types
+    
+    Returns
+    -------
+    str
+        xyz format string
+    
+    Examples
+    --------
+    >>> coord_to_xyz(np.ones((1,3)), ["C"])
+    1
+
+    C 1.000000 1.000000 1.000000
+    """
+    buff = [str(len(types)), '']
+    for at, cc in zip(types, coord):
+        buff.append("{} {:.6f} {:.6f} {:.6f}".format(at, *cc))
+    return "\n".join(buff)
diff --git a/tests/test_xyz.py b/tests/test_xyz.py
new file mode 100644
index 000000000..23bb3a81d
--- /dev/null
+++ b/tests/test_xyz.py
@@ -0,0 +1,19 @@
+import unittest
+import numpy as np
+import tempfile
+from context import dpdata
+
+class TestXYZ(unittest.TestCase):
+    def test_to_xyz(self):
+        with tempfile.NamedTemporaryFile('r') as f_xyz:
+            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),
+            }).to("xyz", f_xyz.name)
+            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)