(feat) XYZWriter: support user-defined precision

MDAnalysis · Oct 29, 2024 · de76d90 · de76d90
1 parent e776f12
commit de76d90
Show file tree

Hide file tree

Showing 4 changed files with 30 additions and 5 deletions.
diff --git a/package/AUTHORS b/package/AUTHORS
@@ -244,6 +244,7 @@ Chronological list of authors
   - Fabian Zills
   - Laksh Krishna Sharma
   - Matthew Davies
+  - Jia-Xin Zhu
 
 
 External code

diff --git a/package/CHANGELOG b/package/CHANGELOG
@@ -18,7 +18,7 @@ The rules for this file:
          ljwoods2, aditya292002, pstaerk, PicoCentauri, BFedder,
          tyler.je.reddy, SampurnaM, leonwehrhan, kainszs, orionarcher, 
          yuxuanzhuang, PythonFZ, laksh-krishna-sharma, orbeckst, MattTDavies,
-         talagayev, aya9aladdin
+         talagayev, aya9aladdin, ChiahsinChu
 
  * 2.8.0
 
@@ -63,6 +63,7 @@ Fixes
  * Fix groups.py doctests using sphinx directives (Issue #3925, PR #4374)
 
 Enhancements
+ * Added `precision` for XYZWriter (PR #4771)
  * Removed type and mass guessing from all topology parsers (PR #3753)
  * Added guess_TopologyAttrs() API to the Universe to handle attribute
    guessing (PR #3753)

diff --git a/package/MDAnalysis/coordinates/XYZ.py b/package/MDAnalysis/coordinates/XYZ.py
@@ -140,8 +140,15 @@ class XYZWriter(base.WriterBase):
     # these are assumed!
     units = {'time': 'ps', 'length': 'Angstrom'}
 
-    def __init__(self, filename, n_atoms=None, convert_units=True,
-                 remark=None, **kwargs):
+    def __init__(
+        self,
+        filename,
+        n_atoms=None,
+        convert_units=True,
+        remark=None,
+        precision=5,
+        **kwargs,
+    ):
         """Initialize the XYZ trajectory writer
 
         Parameters
@@ -161,6 +168,8 @@ def __init__(self, filename, n_atoms=None, convert_units=True,
         remark: str (optional)
             single line of text ("molecule name"). By default writes MDAnalysis
             version and frame
+        precision: int (optional)
+            set precision of saved trjactory to this number of decimal places.
 
 
         .. versionchanged:: 1.0.0
@@ -175,6 +184,7 @@ def __init__(self, filename, n_atoms=None, convert_units=True,
         self.remark = remark
         self.n_atoms = n_atoms
         self.convert_units = convert_units
+        self.precision = precision
 
         # can also be gz, bz2
         self._xyz = util.anyopen(self.filename, 'wt')
@@ -296,8 +306,10 @@ def _write_next_frame(self, ts=None):
 
         # Write content
         for atom, (x, y, z) in zip(self.atomnames, coordinates):
-            self._xyz.write("{0!s:>8}  {1:10.5f} {2:10.5f} {3:10.5f}\n"
-                            "".format(atom, x, y, z))
+            self._xyz.write(
+                "{0!s:>8}  {1:10.{p}f} {2:10.{p}f} {3:10.{p}f}\n"
+                "".format(atom, x, y, z, p=self.precision)
+            )
 
 
 class XYZReader(base.ReaderBase):

diff --git a/testsuite/MDAnalysisTests/coordinates/test_xyz.py b/testsuite/MDAnalysisTests/coordinates/test_xyz.py
@@ -122,6 +122,17 @@ def test_remark(self, universe, remarkout, remarkin, ref, tmpdir):
 
                 assert lines[1].strip() == remarkin
 
+    def test_precision(self, universe, tmpdir):
+        outfile = "write-precision.xyz"
+        precision = 10
+
+        with tmpdir.as_cwd():
+            universe.atoms.write(outfile, precision=precision)
+            with open(outfile, "r") as xyzout:
+                lines = xyzout.readlines()
+                # check that the precision is set correctly
+                assert len(lines[2].split()[1].split(".")[1]) == precision
+
 
 class XYZ_BZ_Reference(XYZReference):
     def __init__(self):