Merge pull request #46 from mattwthompson/write_xyz

Add xyz writer

topology/core/site.py

@@ -20,16 +20,23 @@ def __init__(self,
             self.position = u.nm * np.zeros(3)
         else:
             self.position = _validate_position(position)
-        if element:
-            self.element = element
 
+        self._element = element
         self._atom_type = _validate_atom_type(atom_type)
         self._charge = _validate_charge(charge)
         self._connections = list()
 
     def add_connection(self, other_site):
         self._connections.append(other_site)
 
+    @property
+    def element(self):
+        return self._element
+
+    @element.setter
+    def element(self, element):
+        self._element = element
+
     @property
     def connections(self):
         return self._connections

topology/formats/xyz.py

@@ -1,3 +1,5 @@
+import datetime
+
 import numpy as np
 import unyt as u
 
@@ -33,3 +35,22 @@ def read_xyz(filename):
             raise ValueError(msg.format(n_atoms))
 
     return top
+
+def write_xyz(top, filename):
+    with open(filename, 'w') as out_file:
+        out_file.write('{:d}\n'.format(top.n_sites))
+        out_file.write('{} {} written by topology at {}\n'.format(
+            top.name,
+            filename,
+            str(datetime.datetime.now())))
+        for idx, site in enumerate(top.site_list):
+            # TODO: Better handling of element guessing and site naming
+            if site.element is not None:
+                tmp_name = site.element.symbol
+            else:
+                tmp_name = 'X'
+            out_file.write('{0} {1:8.3f} {2:8.3f} {3:8.3f}\n'.format(
+                tmp_name,
+                site.position[0].in_units(u.angstrom).value,
+                site.position[1].in_units(u.angstrom).value,
+                site.position[2].in_units(u.angstrom).value))

topology/tests/test_xyz.py

@@ -1,10 +1,10 @@
 import unyt as u
 import pytest
 
-from topology.formats.xyz import read_xyz
+from topology.formats.xyz import read_xyz, write_xyz
 from topology.tests.base_test import BaseTest
 from topology.utils.io import get_fn
-
+from topology.testing.utils import allclose
 
 class TestXYZ(BaseTest):
     def test_read_xyz(self):
@@ -26,3 +26,17 @@ def test_wrong_n_atoms(self):
             read_xyz(get_fn('too_few_atoms.xyz'))
         with pytest.raises(ValueError):
             read_xyz(get_fn('too_many_atoms.xyz'))
+
+    def test_write_xyz(self):
+        top = read_xyz(get_fn('ethane.xyz'))
+        write_xyz(top, 'tmp.xyz')
+
+    def test_full_io(self):
+        original_top = read_xyz(get_fn('ethane.xyz'))
+
+        write_xyz(original_top, 'full_conversion.xyz')
+        new_top = read_xyz('full_conversion.xyz')
+
+        assert original_top.n_sites == new_top.n_sites
+        assert original_top.n_connections == new_top.n_connections
+        assert allclose(original_top.positions(), new_top.positions())