openpathsampling · dwhswenson · Jul 27, 2021 · Feb 18, 2021 · Feb 19, 2021 · Feb 27, 2021
diff --git a/.coveragerc b/.coveragerc
@@ -7,3 +7,4 @@ exclude_lines =
     no-cov
     def __repr__
     raise NotImplementedError
+    __name__ == "__main__":
diff --git a/.github/workflows/test-suite.yml b/.github/workflows/test-suite.yml
@@ -28,16 +28,25 @@ jobs:
           - 3.8
           - 3.7
           - 3.6
+        INTEGRATIONS: [""]
+        include:
+          - CONDA_PY: 3.9
+            INTEGRATIONS: 'all-optionals'
 
     steps:
       - uses: actions/checkout@v2
+        with:
+          fetch-depth: 2
       - uses: actions/setup-python@v2
       - uses: conda-incubator/setup-miniconda@v2
         with:
           auto-update-conda: true
           python-version: ${{ matrix.CONDA_PY }}
       - name: "Install testing tools"
         run: python -m pip install -r ./devtools/tests_require.txt
+      - name: "Install integrations"
+        if: matrix.INTEGRATIONS == 'all-optionals'
+        run: conda install -c conda-forge -y openmm openmmtools mdtraj
       - name: "Install"
         run: |
           conda install pip
@@ -57,6 +66,8 @@ jobs:
           fi
           python autorelease_check.py --branch $BRANCH --even ${EVENT}
       - name: "Unit tests"
+        env:
+          PY_COLORS: "1"
         run: |
           python -c "import paths_cli"
           py.test -vv --cov --cov-report xml:cov.xml

diff --git a/docs/index.rst b/docs/index.rst
@@ -60,6 +60,7 @@ wrappers around well-tested OPS code.
    plugins
    parameters
    workflows
+   wizard
    full_cli
    api/index
 
diff --git a/docs/wizard.rst b/docs/wizard.rst
@@ -0,0 +1,8 @@
+.. _wizard:
+
+Writing tools for the Wizard
+============================
+
+The Wizard API is still rapidly in flux, and we don't recommend developing
+custom Wizard tools at this time. However, once its API is more stable, the
+Wizard will be extendable by outside developers.
diff --git a/paths_cli/commands/wizard.py b/paths_cli/commands/wizard.py
@@ -0,0 +1,12 @@
+import click
+from paths_cli.wizard.two_state_tps import TWO_STATE_TPS_WIZARD
+
+@click.command(
+    'wizard',
+    short_help="run wizard for setting up simulations",
+)
+def wizard():  # no-cov
+    TWO_STATE_TPS_WIZARD.run_wizard()
+
+CLI = wizard
+SECTION = "Simulation setup"
diff --git a/paths_cli/parsing/tools.py b/paths_cli/parsing/tools.py
@@ -0,0 +1,45 @@
+import numpy as np
+
+def custom_eval(obj, named_objs=None):
+    """Parse user input to allow simple math.
+
+    This allows certain user input to be treated as a simplified subset of
+    Python. In particular, this is intended to allow simple arithmetic. It
+    allows use of the modules numpy (as ``np``) and math, which provide
+    potentially useful functions (e.g., ``cos``) as well as constants (e.g.,
+    ``pi``).
+    """
+    string = str(obj)
+    # TODO: check that the only attribute access comes from a whitelist
+    # (parse the AST for that)
+    namespace = {
+        'np': __import__('numpy'),
+        'math': __import__('math'),
+    }
+    return eval(string, namespace)
+
+
+class UnknownAtomsError(RuntimeError):
+    pass
+
+def mdtraj_parse_atomlist(inp_str, n_atoms, topology=None):
+    """
+    n_atoms: int
+        number of atoms expected
+    """
+    # TODO: change n_atoms to the shape desired?
+    # TODO: enable the parsing of either string-like atom labels or numeric
+    indices = custom_eval(inp_str)
+
+    arr = np.array(indices)
+    if arr.dtype != int:
+        raise TypeError("Input is not integers")
+    if arr.shape != (1, n_atoms):
+        # try to clean it up
+        if len(arr.shape) == 1 and arr.shape[0] ==  n_atoms:
+            arr.shape = (1, n_atoms)
+        else:
+            raise TypeError(f"Invalid input. Requires {n_atoms} "
+                            "atoms.")
+
+    return arr
diff --git a/paths_cli/tests/parsing/test_tools.py b/paths_cli/tests/parsing/test_tools.py
@@ -0,0 +1,18 @@
+import pytest
+import numpy.testing as npt
+import numpy as np
+import math
+
+from paths_cli.parsing.tools import *
+
+@pytest.mark.parametrize('expr,expected', [
+    ('1+1', 2),
+    ('np.pi / 2', np.pi / 2),
+    ('math.cos(1.5)', math.cos(1.5)),
+])
+def test_custom_eval(expr, expected):
+    npt.assert_allclose(custom_eval(expr), expected)
+
+def test_mdtraj_parse_atomlist_bad_input():
+    with pytest.raises(TypeError, match="not integers"):
+        mdtraj_parse_atomlist("['a', 'b']", n_atoms=2)
diff --git a/paths_cli/tests/test_parameters.py b/paths_cli/tests/test_parameters.py
@@ -42,7 +42,6 @@ def undo_monkey_patch(stored_functions):
     importlib.reload(paths)
 
 
-
 class ParameterTest(object):
     def test_parameter(self):
         # this is just a smoke test to contrast with the ValueError case

diff --git a/paths_cli/tests/utils.py b/paths_cli/tests/utils.py
@@ -0,0 +1,18 @@
+import urllib.request
+
+try:
+    urllib.request.urlopen('https://www.google.com')
+except:
+    HAS_INTERNET = False
+else:
+    HAS_INTERNET = True
+
+def assert_url(url):
+    if not HAS_INTERNET:
+        pytest.skip("Internet connection seems faulty")
+
+    # TODO: On a 404 this will raise a urllib.error.HTTPError. It would be
+    # nice to give some better output to the user here.
+    resp = urllib.request.urlopen(url)
+    assert resp.status == 200
+
diff --git a/paths_cli/tests/wizard/__init__.py b/paths_cli/tests/wizard/__init__.py
diff --git a/paths_cli/tests/wizard/conftest.py b/paths_cli/tests/wizard/conftest.py
@@ -0,0 +1,79 @@
+import pytest
+
+import openpathsampling as paths
+import mdtraj as md
+
+@pytest.fixture
+def ad_openmm(tmpdir):
+    """
+    Provide directory with files to start alanine depeptide sim in OpenMM
+    """
+    mm = pytest.importorskip('simtk.openmm')
+    u = pytest.importorskip('simtk.unit')
+    openmmtools = pytest.importorskip('openmmtools')
+    md = pytest.importorskip('mdtraj')
+    testsystem = openmmtools.testsystems.AlanineDipeptideVacuum()
+    integrator = openmmtools.integrators.VVVRIntegrator(
+        300 * u.kelvin,
+        1.0 / u.picosecond,
+        2.0 * u.femtosecond
+    )
+    traj = md.Trajectory([testsystem.positions.value_in_unit(u.nanometer)],
+                         topology=testsystem.mdtraj_topology)
+    files = {'integrator.xml': integrator,
+             'system.xml': testsystem.system}
+    with tmpdir.as_cwd():
+        for fname, obj in files.items():
+            with open(fname, mode='w') as f:
+                f.write(mm.XmlSerializer.serialize(obj))
+
+        traj.save('ad.pdb')
+
+    return tmpdir
+
+@pytest.fixture
+def ad_engine(ad_openmm):
+    with ad_openmm.as_cwd():
+        pdb = md.load('ad.pdb')
+        topology = paths.engines.MDTrajTopology(
+            pdb.topology
+        )
+        engine = paths.engines.openmm.Engine(
+            system='system.xml',
+            integrator='integrator.xml',
+            topology=topology,
+            options={'n_steps_per_frame': 10,
+                     'n_frames_max': 10000}
+        ).named('ad_engine')
+    return engine
+
+@pytest.fixture
+def toy_engine():
+    pes = (paths.engines.toy.OuterWalls([1.0, 1.0], [1.0, 1.0])
+           + paths.engines.toy.Gaussian(-1.0, [12.0, 12.0], [-0.5, 0.0])
+           + paths.engines.toy.Gaussian(-1.0, [12.0, 12.0], [0.5, 0.0]))
+    topology = paths.engines.toy.Topology(n_spatial=2,
+                                          masses=[1.0],
+                                          pes=pes)
+    integ = paths.engines.toy.LangevinBAOABIntegrator(
+        dt=0.02,
+        temperature=0.1,
+        gamma=2.5
+    )
+    options = {'integ': integ,
+               'n_frames_max': 5000,
+               'n_steps_per_frame': 1}
+    engine = paths.engines.toy.Engine(
+        options=options,
+        topology=topology
+    ).named('toy-engine')
+    return engine
+
+
+@pytest.fixture
+def tps_network():
+    cv = paths.CoordinateFunctionCV('x', lambda s: s.xyz[0][0])
+    state_A = paths.CVDefinedVolume(cv, float("-inf"), 0).named("A")
+    state_B = paths.CVDefinedVolume(cv, 0, float("inf")).named("B")
+    network = paths.TPSNetwork(state_A, state_B).named('tps-network')
+    return network
diff --git a/paths_cli/tests/wizard/mock_wizard.py b/paths_cli/tests/wizard/mock_wizard.py
@@ -0,0 +1,50 @@
+from paths_cli.wizard.wizard import Wizard
+from unittest import mock
+
+def make_mock_wizard(inputs):
+    wizard = Wizard([])
+    wizard.console.input = mock.Mock(return_value=inputs)
+    return wizard
+
+def make_mock_retry_wizard(inputs):
+    wizard = Wizard([])
+    wizard.console.input = mock.Mock(side_effect=inputs)
+    return wizard
+
+class MockConsole:
+    def __init__(self, inputs=None):
+        if isinstance(inputs, str):
+            inputs = [inputs]
+        elif inputs is None:
+            inputs = []
+        self.inputs = inputs
+        self._input_iter = iter(inputs)
+        self.log = []
+        self.width = 80
+        self.input_call_count = 0
+
+    def print(self, content=""):
+        self.log.append(content)
+
+    def input(self, content):
+        self.input_call_count += 1
+        try:
+            user_input = next(self._input_iter)
+        except StopIteration as e:
+            print(self.log_text)
+            raise e
+
+        self.log.append(content + " " + user_input)
+        return user_input
+
+    @property
+    def log_text(self):
+        return "\n".join(self.log)
+
+def mock_wizard(inputs):
+    wizard = Wizard([])
+    console = MockConsole(inputs)
+    wizard.console = console
+    return wizard
+
+
diff --git a/paths_cli/tests/wizard/test_core.py b/paths_cli/tests/wizard/test_core.py
@@ -0,0 +1,30 @@
+import pytest
+from unittest import mock
+
+from openpathsampling.experimental.storage.collective_variables import \
+        CoordinateFunctionCV
+
+from paths_cli.wizard.core import *
+
+from paths_cli.tests.wizard.mock_wizard import (
+    make_mock_wizard, make_mock_retry_wizard
+)
+
+@pytest.mark.parametrize('req,expected', [
+    (('foo', 2, 2), '2'), (('foo', 2, float('inf')), 'at least 2'),
+    (('foo', 0, 2), 'at most 2'),
+    (('foo', 1, 3), 'at least 1 and at most 3')
+])
+def test_interpret_req(req, expected):
+    assert interpret_req(req) == expected
+
+@pytest.mark.parametrize('length,expected', [
+    (0, 'foo'), (1, 'baz'), (2, 'quux'),
+])
+def test_get_missing_object(length, expected):
+    dct = dict([('bar', 'baz'), ('qux', 'quux')][:length])
+    fallback = lambda x: 'foo'
+    wizard = make_mock_wizard('2')
+    result = get_missing_object(wizard, dct, display_name='string',
+                                fallback_func=fallback)
+    assert result == expected