Skip to content

Commit

Permalink
add xyz format (to) (#240)
Browse files Browse the repository at this point in the history
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.
  • Loading branch information
njzjz authored Jan 21, 2022
1 parent 25acc84 commit 3edbc24
Show file tree
Hide file tree
Showing 3 changed files with 66 additions and 0 deletions.
19 changes: 19 additions & 0 deletions dpdata/plugins/xyz.py
Original file line number Diff line number Diff line change
@@ -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")
Expand Down
28 changes: 28 additions & 0 deletions dpdata/xyz/xyz.py
Original file line number Diff line number Diff line change
@@ -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)
19 changes: 19 additions & 0 deletions tests/test_xyz.py
Original file line number Diff line number Diff line change
@@ -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)

0 comments on commit 3edbc24

Please sign in to comment.