Enables path-to-file writing (#29)

* Enables path-to-file writing

* enable path for plotting

* Adds path to example script

MDRestraintsGenerator/datatypes.py

@@ -14,6 +14,7 @@
 """
 
 import MDRestraintsGenerator.writers as writers
+from pathlib import Path
 
 import MDAnalysis as mda
 import numpy as np
@@ -56,7 +57,7 @@ def mean_squared(self):
         """Returns (value-mean)**2 """
         self.ms_values = (self.values - self.mean)**2
 
-    def plot(self, picked_frame=None):
+    def plot(self, picked_frame=None, path='./'):
         """Plots the data
 
         Input
@@ -133,7 +134,7 @@ def get_hist_binrange(atype, in_values):
         xstring = f"{self.atype} [{self.units}]"
         plt.xlabel(xstring)
         plt.ylabel("Number of frames")
-        filename = f"{self.filename}.png"
+        filename = f"{path}/{self.filename}.png"
         plt.legend(loc="best")
         plt.savefig(filename, format="png")
         plt.close()
@@ -382,12 +383,18 @@ def plot(self, frame=None, path=None):
             except AttributeError:
                 pass
 
-        self.bond.plot(picked_frame=frame)
-        self.angles[0].plot(picked_frame=frame)
-        self.angles[1].plot(picked_frame=frame)
-        self.dihedrals[0].plot(picked_frame=frame)
-        self.dihedrals[1].plot(picked_frame=frame)
-        self.dihedrals[2].plot(picked_frame=frame)
+        if path is not None:
+            Path(path).mkdir(parents=True, exist_ok=True)
+            dirpath = path
+        else:
+            dirpath = './'
+
+        self.bond.plot(picked_frame=frame, path=dirpath)
+        self.angles[0].plot(picked_frame=frame, path=dirpath)
+        self.angles[1].plot(picked_frame=frame, path=dirpath)
+        self.dihedrals[0].plot(picked_frame=frame, path=dirpath)
+        self.dihedrals[1].plot(picked_frame=frame, path=dirpath)
+        self.dihedrals[2].plot(picked_frame=frame, path=dirpath)
 
     def write(self, frame=None, path=None, force_constant=10.0, outtype="GMX"):
         """Writes out boresch restraint
@@ -422,15 +429,22 @@ def write(self, frame=None, path=None, force_constant=10.0, outtype="GMX"):
                           f"restraint atoms.")
                 raise RuntimeError(errmsg)
 
-        self._write_gmx(index=frame, force_constant=force_constant)
+        if path is not None:
+            Path(path).mkdir(parents=True, exist_ok=True)
+            dirpath = path
+        else:
+            dirpath = '.'
+
+        self._write_gmx(index=frame, path=dirpath,
+                        force_constant=force_constant)
 
-    def _write_gmx(self, index, force_constant):
+    def _write_gmx(self, index, path, force_constant):
         # seek corre
         self.atomgroup.universe.trajectory[index]
-        self.atomgroup.write('ClosestRestraintFrame.gro')
+        self.atomgroup.write(f'{path}/ClosestRestraintFrame.gro')
 
         # write out top file
-        with open('BoreschRestraint.top', 'w') as rfile:
+        with open(f'{path}/BoreschRestraint.top', 'w') as rfile:
             rfile.write('\n')
             rfile.write('; restraints\n')
             rfile.write('[ intermolecular_interactions ]\n')

MDRestraintsGenerator/tests/test_datatypes.py

@@ -7,13 +7,20 @@
 from .datafiles import T4_TPR, T4_XTC, T4_OGRO
 import pytest
 import os
+from pathlib import Path
+from numpy.testing import assert_almost_equal
 
 
 @pytest.fixture(scope='module')
 def u():
     return mda.Universe(T4_TPR, T4_XTC)
 
 
+@pytest.fixture(scope='module')
+def u_gro():
+    return mda.Universe(T4_OGRO)
+
+
 def test_vector_store():
     n_frames = 100
     vector = dtypes.VectorData(n_frames)
@@ -62,6 +69,22 @@ def test_Boresch_plotting(tmpdir, u, frame):
             assert os.path.isfile(name)
 
 
+def test_Boresch_plotting_path(tmpdir, u):
+    l_atoms = [0, 1, 2]
+    p_atoms = [4, 5, 6]
+
+    boresch = dtypes.BoreschRestraint(u, l_atoms, p_atoms)
+
+    boresch.analyze()
+
+    with tmpdir.as_cwd():
+        boresch.plot(frame=None, path='testdir')
+        for name in ['./testdir/bond_1.png', './testdir/angle_1.png',
+                     './testdir/angle_2.png', './testdir/dihedral_1.png', 
+                     './testdir/dihedral_2.png', './testdir/dihedral_3.png']:
+            assert os.path.isfile(name)
+
+
 def test_Boresch_plotting_notanalysed(tmpdir, u):
     l_atoms = [0, 1, 2]
     p_atoms = [4, 5, 6]
@@ -74,13 +97,28 @@ def test_Boresch_plotting_notanalysed(tmpdir, u):
         boresch.plot()
 
 
-def test_Boresch_write_colinear():
-    u = mda.Universe(T4_OGRO)
+def test_Boresch_write_path(tmpdir, u_gro):
+    l_atoms = [0, 1, 2]
+    p_atoms = [4, 5, 6]
+
+    boresch = dtypes.BoreschRestraint(u_gro, l_atoms, p_atoms)
+
+    boresch.store_frame(0)
+
+    with tmpdir.as_cwd():
+        Path('./testdir').mkdir()
+        boresch.write(frame=0, path='./testdir')
+
+        u2_gro = mda.Universe('./testdir/ClosestRestraintFrame.gro')
 
+        assert_almost_equal(u_gro.atoms.positions, u2_gro.atoms.positions)
+
+
+def test_Boresch_write_colinear(u_gro):
     l_atoms = [2606, 2610, 2611]
     p_atoms = [1218, 1216, 1214]
 
-    boresch = dtypes.BoreschRestraint(u, l_atoms, p_atoms)
+    boresch = dtypes.BoreschRestraint(u_gro, l_atoms, p_atoms)
 
     boresch.store_frame(0)
 

scripts/BoreschRestraintGMX.py

@@ -40,6 +40,8 @@ def parse_args():
                             help='simulation temperature')
         parser.add_argument('--force_constant', type=float, default=10.0,
                             help='restraint force constant')
+        parser.add_argument('--outpath', default='./',
+                            help='output path for writing files')
         args = parser.parse_args()
         return args
 
@@ -73,10 +75,11 @@ def parse_args():
     boresch.run()
 
     # Plot out the statistics
-    boresch.restraint.plot()
+    boresch.restraint.plot(path=args.outpath)
 
     # Write out the intermolecular section to a topology
-    boresch.restraint.write(force_constant=args.force_constant)
+    boresch.restraint.write(path=args.outpath,
+                            force_constant=args.force_constant)
 
     dG_off = boresch.restraint.standard_state(force_constant=args.force_constant,
                                               temperature=args.temperature)