Move engine specific functions into respective files

mdbenchmark/mdengines/gromacs.py

@@ -32,6 +32,19 @@
 NAME = 'gromacs'
 
 
+def prepare_benchmark(name, *args, **kwargs):
+    sim = kwargs['sim']
+
+    full_filename = name + '.tpr'
+    if name.endswith('.tpr'):
+        full_filename = name
+        name = name[:-4]
+
+    copyfile(full_filename, sim[full_filename].relpath)
+
+    return name
+
+
 def check_input_file_exists(name):
     """Check if the TPR file exists.
     """

mdbenchmark/mdengines/namd.py

@@ -30,6 +30,27 @@
 NAME = 'namd'
 
 
+def prepare_benchmark(name, *args, **kwargs):
+    sim = kwargs['sim']
+
+    if name.endswith('.namd'):
+        name = name[:-5]
+
+    namd = '{}.namd'.format(name)
+    psf = '{}.psf'.format(name)
+    pdb = '{}.pdb'.format(name)
+
+    with open(namd) as fh:
+        analyze_namd_file(fh)
+        fh.seek(0)
+
+    copyfile(namd, sim[namd].relpath)
+    copyfile(psf, sim[psf].relpath)
+    copyfile(pdb, sim[pdb].relpath)
+
+    return name
+
+
 def analyze_namd_file(fh):
     """ Check whether the NAMD config file has any relative imports or variables
     """

mdbenchmark/mdengines/utils.py

@@ -143,57 +143,14 @@ def cleanup_before_restart(engine, sim):
         os.remove(fn)
 
 
-def prepare_gromacs_benchmark(name, *args, **kwargs):
-    sim = kwargs['sim']
-
-    full_filename = name + '.tpr'
-    if name.endswith('.tpr'):
-        full_filename = name
-        name = name[:-4]
-
-    copyfile(full_filename, sim[full_filename].relpath)
-
-    return name
-
-
-def prepare_namd_benchmark(name, *args, **kwargs):
-    sim = kwargs['sim']
-
-    if name.endswith('.namd'):
-        name = name[:-5]
-
-    namd = '{}.namd'.format(name)
-    psf = '{}.psf'.format(name)
-    pdb = '{}.pdb'.format(name)
-
-    with open(namd) as fh:
-        analyze_namd_file(fh)
-        fh.seek(0)
-
-    copyfile(namd, sim[namd].relpath)
-    copyfile(psf, sim[psf].relpath)
-    copyfile(pdb, sim[pdb].relpath)
-
-    return name
-
-
-def prepare_benchmark(engine, name, *args, **kwargs):
-    if engine.NAME == 'gromacs':
-        name = prepare_gromacs_benchmark(name, *args, **kwargs)
-    elif engine.NAME == 'namd':
-        name = prepare_namd_benchmark(name, *args, **kwargs)
-
-    return name
-
-
 def write_benchmark(engine, base_directory, template, nodes, gpu, module, name,
                     host, time):
     """Generate a benchmark folder with the respective Sim object."""
     # Create the `mds.Sim` object
     sim = mds.Sim(base_directory['{}/'.format(nodes)])
 
     # Do MD engine specific things. Here we also format the name.
-    name = prepare_benchmark(engine=engine, name=name, sim=sim)
+    name = engine.prepare_benchmark(name=name, sim=sim)
 
     # Add categories to the `Sim` object
     sim.categories = {