run works

adaptivemd/engine/acemd/acemd.py

@@ -40,7 +40,7 @@ def __init__(self, conf_file, pdb_file, args=None):
     def call_format_str(self):
         return 'acemd %s {0}' % self.args
 
-    def task_run_trajectory(self, target):
+    def run(self, target):
         t = Task()
 
         initial_pdb = t.link(self['pdb_file_stage'], Location('initial.pdb'))

adaptivemd/engine/engine.py

@@ -2,9 +2,9 @@
 import random
 import os
 
-from adaptivemd.file import File, Location
+from adaptivemd.file import File
 from adaptivemd.generator import TaskGenerator
-from adaptivemd.mongodb import StorableMixin, SyncVariable
+from adaptivemd.mongodb import StorableMixin, ObjectSyncVariable
 from adaptivemd.task import Task
 
 
@@ -14,7 +14,27 @@ class Engine(TaskGenerator):
 
     """
 
-    def task_run_trajectory(self, target):
+    def __init__(self):
+        super(Engine, self).__init__()
+
+        self.types = {}
+
+        # set default output type if nothing is specified
+        self.add_output_type('master', 'output.dcd', 1)
+
+    @classmethod
+    def from_dict(cls, dct):
+        obj = super(Engine, cls).from_dict(dct)
+        obj.types = dct['types']
+        return obj
+
+    def to_dict(self):
+        dct = super(Engine, self).to_dict()
+        dct.update({
+            'types': self.types})
+        return dct
+
+    def run(self, target):
         """
         Create a task that returns a trajectory given in the input
 
@@ -31,6 +51,25 @@ def task_run_trajectory(self, target):
         """
         return None
 
+    def extend(self, target, length):
+        """
+        Create a task that extends a trajectory given in the input
+
+        Parameters
+        ----------
+        target : `Trajectory`
+            location of the target trajectory to be extended
+        length : int
+            number of additional frames to be computed
+
+        Returns
+        -------
+        `Task`
+            the task object containing the job description
+
+        """
+        return None
+
     def file_generators(self):
         """
         Return a list of function to be run with certain classes
@@ -47,9 +86,12 @@ def file_generators(self):
 
         """
         return {
-            Trajectory: self.task_run_trajectory
+            Trajectory: self.run
         }
 
+    def add_output_type(self, name, filename=None, stride=1):
+        self.types[name] = OutputTypeDescription(filename, stride)
+
 
 # ------------------------------------------------------------------------------
 # FILE TYPES
@@ -71,7 +113,9 @@ class Trajectory(File):
         the engine used to create the trajectory
     """
 
-    engine = SyncVariable('engine', lambda x: not bool(x))
+    _find_by = ['created', 'state', 'task', 'engine']
+
+    engine = ObjectSyncVariable('engine', 'generators', lambda x: not bool(x))
 
     def __init__(self, location, frame, length, engine=None):
         super(Trajectory, self).__init__(location)
@@ -104,12 +148,79 @@ def is_folder(self):
         return True
 
     def file(self, f):
-        return File(os.path.join(self.location, f))
+        if isinstance(f, basestring):
+            return File(os.path.join(self.location, f))
+        elif isinstance(f, OutputTypeDescription):
+            return self.file(f.filename)
 
     @property
     def restartable(self):
         return True
 
+    def run(self):
+        if self.engine:
+            return self.engine.run(self)
+        else:
+            return None
+
+    def extend(self, length):
+        """
+        Get a task to extend this trajectory if the engine is set
+
+        Parameters
+        ----------
+        length : int
+            the length to extend by
+
+        Returns
+        -------
+        `Task`
+            the task object
+        """
+        if self.engine:
+            return self.engine.extend(self, length)
+        else:
+            return None
+
+    def outputs(self, outtype):
+        """
+        Get a location to the file containing the output by given name
+
+        Parameters
+        ----------
+        outtype : str ot `OutputTypeDescription`
+
+        Returns
+        -------
+        `File`
+            a file location that points to the concrete file that contains
+            the data for a particular output type
+
+        """
+        if self.engine:
+            if isinstance(outtype, basestring):
+                if outtype in self.engine.types:
+                    return self.file(self.engine.types[outtype])
+            elif isinstance(outtype, OutputTypeDescription):
+                return self.file(outtype)
+
+        return None
+
+    @property
+    def types(self):
+        """
+        Return the OutputTypeDescriptions for this trajectory
+        Returns
+        -------
+        dict str: `OutputTypeDescription`
+            the output description dict of the engine
+
+        """
+        if self.engine:
+            return self.engine.types
+
+        return None
+
 
 class Frame(StorableMixin):
     """
@@ -156,12 +267,24 @@ class TrajectoryGenerationTask(Task):
             'trajectory'
         ]
 
+    def _default_success(self, scheduler):
+        super(TrajectoryGenerationTask, self)._default_success(scheduler)
+
+        # # give the used engine the credit for making the trajectory
+        # for t in self.targets:
+        #     if isinstance(t, Trajectory):
+        #         t.engine = self.generator
+
     def __init__(self, generator=None, trajectory=None):
         super(TrajectoryGenerationTask, self).__init__(generator)
+
+        # set this engine to be run by this
         self.trajectory = trajectory
+        if trajectory:
+            trajectory.engine = self.generator
 
     def extend(self, length):
-        t = self.generator.task_extend_trajectory(self.trajectory, length)
+        t = self.generator.extend(self.trajectory, length)
 
         # this is not really necessary since we require internally that the source exists
         # but this will cause all dependencies to be submitted, too
@@ -193,3 +316,14 @@ def ready(self):
             return False
 
         return True
+
+
+class OutputTypeDescription(StorableMixin):
+    def __init__(self, filename=None, stride=1):
+        super(OutputTypeDescription, self).__init__()
+
+        if filename is None:
+            filename = 'stride-%d.dcd' % stride
+
+        self.filename = filename
+        self.stride = stride

adaptivemd/engine/openmm/openmm.py

@@ -1,6 +1,6 @@
 import os
+import ujson
 
-from adaptivemd.task import PythonTask
 from adaptivemd.file import Location, File
 from adaptivemd.engine import Engine, Frame, Trajectory, \
     TrajectoryGenerationTask, TrajectoryExtensionTask
@@ -25,7 +25,7 @@ class OpenMMEngine(Engine):
         a list of arguments passed to the `openmmrun.py` script
     """
 
-    def __init__(self, system_file, integrator_file, pdb_file, args=None, restartable=True):
+    def __init__(self, system_file, integrator_file, pdb_file, args=None, restartable=False):
         super(OpenMMEngine, self).__init__()
 
         self._items = dict()
@@ -41,14 +41,24 @@ def __init__(self, system_file, integrator_file, pdb_file, args=None, restartabl
             self.initial_staging.append(stage)
 
         if args is None:
-            args = '-p CPU --store-interval 1'
+            args = '-p CPU'
 
         self.args = args
         self.restartable = restartable
 
-    @property
-    def call_format_str(self):
-        return 'python openmmrun.py %s {3} -t {0} --length {1} {2}' % self.args
+    @classmethod
+    def from_dict(cls, dct):
+        obj = super(OpenMMEngine, cls).from_dict(dct)
+        obj.args = dct['args']
+        obj.restartable = dct['restartable']
+        return obj
+
+    def to_dict(self):
+        dct = super(OpenMMEngine, self).to_dict()
+        dct.update({
+            'restartable': self.restartable,
+            'args': self.args})
+        return dct
 
     @staticmethod
     def then_func_import(project, task, data, inputs):
@@ -57,7 +67,14 @@ def then_func_import(project, task, data, inputs):
             if f not in project.files:
                 project.files.update(f)
 
-    def task_run_trajectory(self, target):
+    def _create_output_str(self):
+        d = dict()
+        for name, opt in self.types.iteritems():
+            d[name] = opt.to_dict()
+
+        return '--types="%s"' % ujson.dumps(d).replace('"', "'")
+
+    def run(self, target):
         t = TrajectoryGenerationTask(self, target)
 
         initial_pdb = t.link(self['pdb_file_stage'], Location('initial.pdb'))
@@ -91,8 +108,9 @@ def task_run_trajectory(self, target):
         # create the directory
         t.touch(output)
 
-        cmd = 'python openmmrun.py {args} -t {pdb} --length {length} {output}'.format(
+        cmd = 'python openmmrun.py {args} {types} -t {pdb} --length {length} {output}'.format(
             pdb=input_pdb,
+            types=self._create_output_str(),
             length=target.length,
             output=output,
             args=self.args,
@@ -103,7 +121,7 @@ def task_run_trajectory(self, target):
 
         return t
 
-    def task_extend_trajectory(self, source, length):
+    def extend(self, source, length):
         if length < 0:
             return []
 
@@ -157,16 +175,16 @@ def task_extend_trajectory(self, source, length):
 
         return t
 
-    def task_import_trajectory_folder(self, source):
-        t = PythonTask(self)
-
-        t.link(self['pdb_file_stage'], Location('initial.pdb'))
-        t.call(scan_trajectories, source)
-
-        # call `then_func_import` after success
-        t.then('then_func_import')
-
-        return t
+    # def task_import_trajectory_folder(self, source):
+    #     t = PythonTask(self)
+    #
+    #     t.link(self['pdb_file_stage'], Location('initial.pdb'))
+    #     t.call(scan_trajectories, source=source)
+    #
+    #     # call `then_func_import` after success
+    #     t.then('then_func_import')
+    #
+    #     return t
 
 
 def scan_trajectories(source):