Wizard rewrite for engines

dwhswenson · Oct 28, 2021 · Jul 24, 2021 · Sep 27, 2021 · Sep 30, 2021 · Oct 2, 2021
commit f1df73f1a2c4e18eda84e3657d72c02f3be4056f
diff --git a/paths_cli/wizard/engines.py b/paths_cli/wizard/engines.py
@@ -4,6 +4,8 @@
 )
 from functools import partial
 
+from paths_cli.wizard.wrap_compilers import WrapCategory
+
 SUPPORTED_ENGINES = {}
 for module in [openmm]:
     SUPPORTED_ENGINES.update(module.SUPPORTED)
@@ -24,8 +26,27 @@ def engines(wizard):
     engine = SUPPORTED_ENGINES[eng_name](wizard)
     return engine
 
+_ENGINE_HELP = "An engine describes how you'll do the actual dynamics."
+ENGINE_PLUGIN = WrapCategory(
+    name='engines',
+    ask="What will you use for the underlying engine?",
+    intro=("Let's make an engine. " + _ENGINE_HELP + " Most of the "
+           "details are given in files that depend on the specific "
+           "type of engine."),
+    helper=_ENGINE_HELP
+)
+
+
 if __name__ == "__main__":
     from paths_cli.wizard import wizard
-    wiz = wizard.Wizard({'engines': ('engines', 1, '=')})
-    wiz.run_wizard()
+    wiz = wizard.Wizard([])
+    choices = {
+        "OpenMM": openmm.OPENMM_PLUGIN,
+        load_label: partial(load_from_ops,
+                            store_name='engines',
+                            obj_name='engine')
+    }
+    ENGINE_PLUGIN.choices = choices
+    engine = ENGINE_PLUGIN(wiz)
+    print(engine)
 
diff --git a/paths_cli/wizard/helper.py b/paths_cli/wizard/helper.py
@@ -15,11 +15,29 @@ def force_exit(cmd, ctx):
     'q': raise_quit,
     'quit': raise_quit,
     '!q': force_exit,
+    '!quit': force_exit,
     'restart': raise_restart,
     'fuck': raise_restart,  # easter egg ;-)
     # TODO: add ls, perhaps others?
 }
 
+_QUIT_HELP = ("The !quit command (or !q) will quit the wizard, but only "
+              "after it checks whether you want to save the objects you've "
+              "created.")
+_FORCE_QUIT_HELP = ("The !!quit command (or !!q) will immediately force "
+                    "quit the wizard.")
+_RESTART_HELP = ("The !restart command will restart the creation of the "
+                 "current object in the wizard. Note that for objects that "
+                 "contain other objects (such as state definitions, which "
+                 "may be made of multiple volumes) this will restart from "
+                 "the outermost object (e.g., the state).")
+COMMAND_HELP_STR = {
+    'q': _QUIT_HELP,
+    'quit': _QUIT_HELP,
+    '!q': _FORCE_QUIT_HELP,
+    '!quit': _FORCE_QUIT_HELP,
+    'restart': _RESTART_HELP,
+}
 
 class Helper:
     def __init__(self, help_func):
@@ -29,14 +47,42 @@ def __init__(self, help_func):
             help_func = lambda args, ctx: text
         self.helper = help_func
         self.commands = HELPER_COMMANDS.copy()  # allows per-instance custom
+        self.command_help_str = COMMAND_HELP_STR.copy()
+        self.listed_commands = ['quit', '!quit', 'restart']
+
+    def command_help(self, help_args, context):
+        if help_args == "":
+            result = "The following commands can be used:\n"
+            result += "\n".join([f"* !{cmd}"
+                                 for cmd in self.listed_commands])
+        else:
+            try:
+                result = self.command_help_str[help_args]
+            except KeyError:
+                result = f"No help for !{help_args}."
+
+        return result
+
 
     def run_command(self, command, context):
         cmd_split = command.split()
         key = cmd_split[0]
         args = " ".join(cmd_split[1:])
-        return self.commands[key](args, context)
+        try:
+            cmd = self.commands[key]
+        except KeyError:
+            return f"Unknown command: {key}"
+
+        return cmd(args, context)
 
     def get_help(self, help_args, context):
+        # TODO: add default help (for ?!, etc)
+        if help_args != "" and help_args[0] == '!':
+            return self.command_help(help_args[1:], context)
+
+        if self.helper is None:
+            return "Sorry, no help available here."
+
         return self.helper(help_args, context)
 
     def __call__(self, user_input, context=None):
@@ -45,5 +91,3 @@ def __call__(self, user_input, context=None):
         func = {'?': self.get_help,
                 '!': self.run_command}[starter]
         return func(args, context)
-
-
diff --git a/paths_cli/wizard/load_from_ops.py b/paths_cli/wizard/load_from_ops.py
@@ -45,3 +45,5 @@ def load_from_ops(wizard, store_name, obj_name):
     storage = _get_ops_storage(wizard)
     obj = _get_ops_object(wizard, storage, store_name, obj_name)
     return obj
+
+
diff --git a/paths_cli/wizard/openmm.py b/paths_cli/wizard/openmm.py
@@ -13,6 +13,69 @@
     SimpleParameter, InstanceBuilder, load_custom_eval, CUSTOM_EVAL_ERROR
 )
 
+from paths_cli.wizard.wrap_compilers import (
+    WizardProxyParameter, WrapCompilerWizardPlugin
+)
+
+from paths_cli.compiling.engines import OPENMM_PLUGIN as OPENMM_COMPILING
+
+_where_is_xml = "Where is the XML file for your OpenMM {obj_type}?"
+_xml_help = (
+    "You can write OpenMM objects like systems and integrators to XML "
+    "files using the XMLSerializer class. Learn more here:\n"
+    + OPENMM_SERIALIZATION_URL
+)
+
+_OPENMM_REQS = ("To use OpenMM in the OPS wizard, you'll need to provide a "
+                "file with topology information (usually a PDB), as well "
+                "as XML versions of your OpenMM integrator and system "
+                "objects.")
+
+OPENMM_PLUGIN = WrapCompilerWizardPlugin(
+    name="OpenMM",
+    category="engines",
+    description=("OpenMM is an GPU-accelerated library for molecular "
+                 "dynamics. " + _OPENMM_REQS),
+    intro="Great! OpenMM gives you a lots of flexibility" + _OPENMM_REQS,
+    parameters=[
+        WizardProxyParameter(
+            name='topology',
+            ask="Where is a PDB file describing your system?",
+            helper=("We use a PDB file to set up the OpenMM simulation. "
+                    "Please provide the path to your PDB file."),
+            error=FILE_LOADING_ERROR_MSG,
+        ),
+        WizardProxyParameter(
+            name='integrator',
+            ask=_where_is_xml.format(obj_type='integrator'),
+            helper=_xml_help,
+            error=FILE_LOADING_ERROR_MSG,
+        ),
+        WizardProxyParameter(
+            name='system',
+            ask=_where_is_xml.format(obj_type="system"),
+            helper=_xml_help,
+            error=FILE_LOADING_ERROR_MSG,
+        ),
+        WizardProxyParameter(
+            name='n_steps_per_frame',
+            ask="How many MD steps per saved frame?",
+            helper=("Your integrator has a time step associated with it. "
+                    "We need to know how many time steps between the "
+                    "frames we actually save during the run."),
+            error=CUSTOM_EVAL_ERROR,
+        ),
+        WizardProxyParameter(
+            name='n_frames_max',
+            ask="How many frames before aborting a trajectory?",
+            helper=None,
+            error=CUSTOM_EVAL_ERROR,
+        ),
+    ],
+    compiler_plugin=OPENMM_COMPILING,
+)
+
+
 ### TOPOLOGY
 
 def _topology_loader(filename):
@@ -33,8 +96,7 @@ def _topology_loader(filename):
 _xml_help = (
     "You can write OpenMM objects like systems and integrators to XML "
     "files using the XMLSerializer class. Learn more here:\n"
-    "http://docs.openmm.org/latest/api-python/generated/"
-    "simtk.openmm.openmm.XmlSerializer.html"
+    + OPENMM_SERIALIZATION_URL
 )
 # TODO: switch to using load_openmm_xml from input file setup
 def _openmm_xml_loader(xml):
@@ -174,5 +236,6 @@ def openmm(wizard):
 if __name__ == "__main__":
     from paths_cli.wizard.wizard import Wizard
     wizard = Wizard([])
-    engine = openmm_builder(wizard)
+    # engine = openmm_builder(wizard)
+    engine = OPENMM_PLUGIN(wizard)
     print(engine)
diff --git a/paths_cli/wizard/parameters.py b/paths_cli/wizard/parameters.py
@@ -1,4 +1,4 @@
-from paths_cli.parsing.tools import custom_eval
+from paths_cli.compiling.tools import custom_eval
 import importlib
 
 from paths_cli.wizard.helper import Helper

diff --git a/paths_cli/wizard/wizard.py b/paths_cli/wizard/wizard.py
@@ -72,8 +72,8 @@ def _speak(self, content, preface):
     def ask(self, question, options=None, default=None, helper=None):
         result = self.console.input("🧙 " + question + " ")
         self.console.print()
-        if helper and result.startswith("?"):
-            helper(self, result)
+        if helper and result[0] in ["?", "!"]:
+            self.say(helper(result))
             return
         return result
 
@@ -90,6 +90,32 @@ def bad_input(self, content, preface="👺 "):
         # just changes the default preface; maybe print 1st line red?
         self.say(content, preface)
 
+    @get_object
+    def ask_enumerate_dict(self, question, options, helper=None):
+        self.say(question)
+        opt_string = "\n".join([f" {(i+1):>3}. {opt}"
+                                for i, opt in enumerate(options)])
+        self.say(opt_string, preface=" "*3)
+        choice = self.ask("Please select an option:", helper=helper)
+
+        # indicates input was handled by helper -- so ask again
+        if choice is None:
+           return None
+
+        # select by string
+        if choice in options:
+            return options[choice]
+
+        # select by number
+        try:
+            num = int(choice) - 1
+            result = list(options.values())[num]
+        except Exception:
+            self.bad_input(f"Sorry, '{choice}' is not a valid option.")
+            result = None
+
+        return result
+
     def ask_enumerate(self, question, options):
         """Ask the user to select from a list of options"""
         self.say(question)
@@ -113,6 +139,17 @@ def ask_enumerate(self, question, options):
 
         return result
 
+    @get_object
+    def ask_load(self, question, loader, helper=None):
+        as_str = self.ask(question, helper=helper)
+        try:
+            result = loader(as_str)
+        except Exception as e:
+            self.exception(f"Sorry, I couldn't understand the input "
+                           f"'{as_str}'.", e)
+            return
+        return result
+
     # this should match the args for wizard.ask
     @get_object
     def ask_custom_eval(self, question, options=None, default=None,

diff --git a/paths_cli/wizard/wrap_compilers.py b/paths_cli/wizard/wrap_compilers.py
@@ -0,0 +1,106 @@
+NO_PARAMETER_LOADED = object()
+
+from .helper import Helper
+
+class WizardProxyParameter:
+    def __init__(self, name, ask, helper, error):
+        self.name = name
+        self.ask = ask
+        self.helper = Helper(helper)
+        self.error = error
+        self.loader = None
+
+    def register_loader(self, loader):
+        if self.loader is not None:
+            raise RuntimeError("Already have a loader for this parameter")
+        self.loader = loader
+
+    def _process_input(self, user_str):
+        if user_str[0] in ['?', '!']:
+            pass
+
+    def __call__(self, wizard):
+        obj = NO_PARAMETER_LOADED
+        while obj is NO_PARAMETER_LOADED:
+            obj = wizard.ask_load(self.ask, self.loader, self.helper)
+
+        return obj
+
+
+class WrapCompilerWizardPlugin:
+    def __init__(self, name, category, parameters, compiler_plugin,
+                 prerequisite=None, intro=None, description=None):
+        self.parameters = parameters
+        self.compiler_plugin = compiler_plugin
+        self.prerequisite = prerequisite
+        self.intro = intro
+        self.description = description
+        loaders = {p.name: p.loader for p in self.compiler_plugin.parameters}
+        for param in self.parameters:
+            param.register_loader(loaders[param.name])
+
+    def __call__(self, wizard):
+        if self.intro is not None:
+            wizard.say(self.intro)
+
+        if self.prerequisite is not None:
+            self.prerequisite(wizard)
+
+        dct = {param.name: param(wizard) for param in self.parameters}
+        result = self.compiler_plugin.builder(**dct)
+
+        return result
+
+class CategoryHelpFunc:
+    def __init__(self, category, basic_help=None):
+        self.category = category
+        if basic_help is None:
+            basic_help = f"Sorry, no help available for {category.name}."
+        self.basic_help = basic_help
+
+    def __call__(self, help_args, context):
+        if not help_args:
+            return self.basic_help
+        help_dict = {}
+        for num, (name, obj) in enumerate(self.category.choices.items()):
+            try:
+                help_str = obj.description
+            except Exception:
+                help_str = f"Sorry, no help available for '{name}'."
+            help_dict[str(num+1)] = help_str
+            help_dict[name] = help_str
+
+        try:
+            result = help_dict[help_args]
+        except KeyError:
+            result = None
+        return result
+
+
+class WrapCategory:
+    def __init__(self, name, ask, helper=None, intro=None):
+        self.name = name
+        if isinstance(intro, str):
+            intro = [intro]
+        self.intro = intro
+        self.ask = ask
+        if helper is None:
+            helper = Helper(CategoryHelpFunc(self))
+        if isinstance(helper, str):
+            helper = Helper(CategoryHelpFunc(self, helper))
+
+        self.helper = helper
+        self.choices = {}
+
+    def register_plugin(self, plugin):
+        self.choices[plugin.name] = plugin
+
+    def __call__(self, wizard):
+        for line in self.intro:
+            wizard.say(line)
+
+        selected = wizard.ask_enumerate_dict(self.ask, self.choices,
+                                             self.helper)
+        obj = selected(wizard)
+        return obj
+
Original file line number	Diff line number	Diff line change
Expand Up		@@ -45,3 +45,5 @@ def load_from_ops(wizard, store_name, obj_name):
		storage = _get_ops_storage(wizard)
		obj = _get_ops_object(wizard, storage, store_name, obj_name)
		return obj