Save unitcell when writing trr

package/MDAnalysis/coordinates/TRR.py

@@ -4,7 +4,7 @@
 from . import base
 from ..lib.formats.trr import TRRFile
 from ..core.AtomGroup import Universe, AtomGroup
-from ..lib.mdamath import triclinic_box
+from ..lib.mdamath import triclinic_box, triclinic_vectors
 
 
 class TRRReader(base.Reader):
@@ -43,12 +43,10 @@ def __init__(self, filename, convert_units=True, sub=None, **kwargs):
         self.ts.time = time_1
         self.ts.dt = time_2 - time_1
         self.ts.dimensions = triclinic_box(*box)
-        self.box = box
-        self.ts._unitcell = self.ts.dimensions
 
         if self.convert_units:
             self.convert_pos_from_native(self.ts._pos)
-            self.convert_pos_from_native(self.ts._unitcell[:3])
+            self.convert_pos_from_native(self.ts.dimensions[:3])
             self.convert_velocities_from_native(self.ts.velocities)
             self.convert_forces_from_native(self.ts.forces)
 
@@ -113,8 +111,7 @@ class TRRWriter(base.Writer):
     units = {'time': 'ps', 'length': 'nm', 'velocity': 'nm/ps',
              'force': 'kJ/(mol*nm)'}
 
-    def __init__(self, filename, convert_units=True, n_atoms=0,
-                 **kwargs):
+    def __init__(self, filename, convert_units=True, n_atoms=0, **kwargs):
         self._convert_units = convert_units
         self.trr = TRRFile(filename, 'w')
         self.n_atoms = n_atoms
@@ -152,8 +149,11 @@ def write_next_timestep(self, ts):
             forces = np.empty((self.n_atoms, 3), dtype=np.float32)
         time = ts.time
         step = ts.frame
-        box = np.eye(3)
-        print(velo.shape)
+
+        if self._convert_units:
+            dimensions = self.convert_dimensions_to_unitcell(ts)
+
+        box = triclinic_vectors(dimensions)
 
         # TODO: write dummy for
         self.trr.write(xyz, velo, forces, box, step, time, 1)