Merge pull request #14 from dwhswenson/restructure_file_functions

Restructure file functions

Author: dwhswenson

README.md

@@ -28,7 +28,6 @@ miscellaneous operations on OPS output files.
 **Miscellaneous Commands:**
 
 * `contents`:         List named objects from an OPS .nc file
-* `strip-snapshots`:  Remove coordinates/velocities from an OPS storage
 * `append`:           add objects from INPUT_FILE  to another file
 
 Full documentation is at https://openpathsampling-cli.readthedocs.io/; a brief

docs/for_core/cli.rst

@@ -79,10 +79,13 @@ foregoing the CLI tools to run simulations, some of the "miscellaneous"
 commands are likely to be quite useful. Here are some that are available in
 the CLI:
 
-* ``nclist``: list all the named objects in an OPS storage, organized by
+* ``contents``: list all the named objects in an OPS storage, organized by
   store (type); this is extremely useful to get the name of an object to use
   as command-line input to one of the simulation scripts
-* ``strip-snapshots``: create a copy of the input storage file with the
+.. * ``strip-snapshots``: create a copy of the input storage file with the
   details (coordinates/velocities) of all snapshots removed; this allows you
   to make a much smaller copy (with results of CVs) to copy back to a local
   computer for analysis
+* ``append`` : add an object from once OPS storage into another one; this is
+  useful for getting everything into a single file before running a
+  simulation

paths_cli/cli.py

@@ -128,10 +128,9 @@ def format_commands(self, ctx, formatter):
 OpenPathSampling is a Python library for path sampling simulations. This
 command line tool facilitates common tasks when working with
 OpenPathSampling. To use it, use one of the subcommands below. For example,
-you can get more information about the strip-snapshots (filesize reduction)
-tool with:
+you can get more information about the pathsampling tool with:
 
-    openpathsampling strip-snapshots --help
+    openpathsampling pathsampling --help
 """
 
 @click.command(cls=OpenPathSamplingCLI, name="openpathsampling",

paths_cli/commands/strip_snapshots.py

@@ -0,0 +1,103 @@
+"""Tools to facilitate copying files.
+
+This is mainly aimed at cases where a file is being copied with some sort of
+modification, or where CVs need to be disk-cached.
+"""
+
+import click
+from tqdm.auto import tqdm
+from paths_cli.parameters import (
+    Option, Argument, HELP_MULTIPLE, StorageLoader, OPSStorageLoadNames
+)
+
+INPUT_APPEND_FILE = StorageLoader(
+    param=Argument('append_file',
+                   type=click.Path(writable=True, readable=True)),
+    mode='a'
+)
+
+class PrecomputeLoadNames(OPSStorageLoadNames):
+    def get(self, storage, name):
+        if len(name) == 0:
+            return list(getattr(storage, self.store))
+        elif len(name) == 1 and name[0] == '--':
+            return []
+
+        return super(PrecomputeLoadNames, self).get(storage, name)
+
+PRECOMPUTE_CVS = PrecomputeLoadNames(
+    param=Option('--cv', type=str, multiple=True,
+                 help=('name of CV to precompute; if not specified all will'
+                       + ' be used' + HELP_MULTIPLE
+                       + ' (use `--cv --` to disable precomputing)')),
+    store='cvs'
+)
+
+
+def make_blocks(listlike, blocksize):
+    """Make blocks out of a listlike object.
+
+    Parameters
+    ----------
+    listlike : Iterable
+        must be an iterable that supports slicing
+    blocksize : int
+        number of objects per block
+
+
+    Returns
+    -------
+    List[List[Any]] :
+        the input iterable chunked into blocks
+    """
+    n_objs = len(listlike)
+    partial_block = 1 if n_objs % blocksize else 0
+    n_blocks = (n_objs // blocksize) + partial_block
+    minval = lambda i: i * blocksize
+    maxval = lambda i: min((i + 1) * blocksize, n_objs)
+    blocks = [listlike[minval(i):maxval(i)] for i in range(n_blocks)]
+    return blocks
+
+
+def precompute_cvs(cvs, block):
+    """Calculate a CV for a a given block.
+
+    Parameters
+    ----------
+    cvs : List[:class:`openpathsampling.CollectiveVariable`]
+        CVs to precompute
+    block : List[Any]
+        b
+    """
+    for cv in cvs:
+        cv.enable_diskcache()
+        _ = cv(block)
+
+
+def precompute_cvs_func_and_inputs(input_storage, cvs, blocksize):
+    """
+    Parameters
+    ----------
+    input_storage : :class:`openpathsampling.Storage`
+        storage file to read from
+    cvs : List[:class:`openpathsampling.CollectiveVariable`]
+        list of CVs to precompute; if None, use all CVs in ``input_storage``
+    blocksize : int
+        number of snapshots per block to precompute
+    """
+    if cvs is None:
+        cvs = list(input_storage.cvs)
+
+    precompute_func = lambda inps: precompute_cvs(cvs, inps)
+    snapshot_proxies = input_storage.snapshots.all().as_proxies()
+    snapshot_blocks = make_blocks(snapshot_proxies, blocksize)
+    return precompute_func, snapshot_blocks
+
+
+def rewrite_file(stage_names, stage_mapping):
+    stages = tqdm(stage_names, desc="All stages")
+    for stage in stages:
+        store_func, inputs = stage_mapping[stage]
+        desc = "This stage: {}".format(stage)
+        for obj in tqdm(inputs, desc=desc, leave=False):
+            store_func(obj)

paths_cli/file_copying.py

@@ -88,7 +88,6 @@ def test_append_remove_tag(tps_network_and_traj):
         result = runner.invoke(append,
                                [in_file, '-a', "output.nc",
                                 "--tag", 'template', '--save-tag', ''])
-        print(result.output)
         assert result.exception is None
         assert result.exit_code == 0
 

paths_cli/tests/commands/test_append.py

@@ -0,0 +1,129 @@
+import collections
+import functools
+import os
+import tempfile
+from unittest.mock import MagicMock, patch
+import pytest
+
+import openpathsampling as paths
+from openpathsampling.tests.test_helpers import make_1d_traj
+
+from paths_cli.file_copying import *
+
+class Test_PRECOMPUTE_CVS(object):
+    def setup(self):
+        self.tmpdir = tempfile.mkdtemp()
+        self.storage_filename = os.path.join(self.tmpdir, "test.nc")
+        self.storage = paths.Storage(self.storage_filename, mode='w')
+        snap = make_1d_traj([1])[0]
+        self.storage.save(snap)
+        self.cv_x = paths.CoordinateFunctionCV("x", lambda s: s.xyz[0][0])
+        self.cv_y = paths.CoordinateFunctionCV("y", lambda s: s.xyz[0][1])
+        self.storage.save([self.cv_x, self.cv_y])
+
+    def teardown(self):
+        self.storage.close()
+
+        for filename in os.listdir(self.tmpdir):
+            os.remove(os.path.join(self.tmpdir, filename))
+        os.rmdir(self.tmpdir)
+
+    @pytest.mark.parametrize('getter', ['x', None, '--'])
+    def test_get(self, getter):
+        expected = {'x': [self.cv_x],
+                    None: [self.cv_x, self.cv_y],
+                    '--': []}[getter]
+        getter = [] if getter is None else [getter]  # CLI gives a list
+        cvs = PRECOMPUTE_CVS.get(self.storage, getter)
+        assert len(cvs) == len(expected)
+        assert set(cvs) == set(expected)
+
+
+@pytest.mark.parametrize('blocksize', [2, 3, 5, 10, 12])
+def test_make_blocks(blocksize):
+    expected_lengths = {2: [2, 2, 2, 2, 2],
+                        3: [3, 3, 3, 1],
+                        5: [5, 5],
+                        10: [10],
+                        12: [10]}[blocksize]
+    ll = list(range(10))
+    blocks = make_blocks(ll, blocksize)
+    assert [len(block) for block in blocks] == expected_lengths
+    assert sum(blocks, []) == ll
+
+
+class TestPrecompute(object):
+    def setup(self):
+        class RunOnceFunction(object):
+            def __init__(self):
+                self.previously_seen = set([])
+
+            def __call__(self, snap):
+                if snap in self.previously_seen:
+                    raise AssertionError("Second CV eval for " + str(snap))
+                self.previously_seen.update({snap})
+                return snap.xyz[0][0]
+
+        self.cv = paths.FunctionCV("test", RunOnceFunction())
+        traj = make_1d_traj([2, 1])
+        self.snap = traj[0]
+        self.other_snap = traj[1]
+
+    def test_precompute_cvs(self):
+        precompute_cvs([self.cv], [self.snap])
+        assert self.cv.f.previously_seen == {self.snap}
+        recalced = self.cv(self.snap)  # AssertionError if func called
+        assert recalced == 2
+        assert self.cv.diskcache_enabled is True
+
+    @pytest.mark.parametrize('cvs', [['test'], None])
+    def test_precompute_cvs_and_inputs(self, cvs):
+        with tempfile.TemporaryDirectory() as tmpdir:
+            storage = paths.Storage(os.path.join(tmpdir, "test.nc"),
+                                    mode='w')
+            traj = make_1d_traj(list(range(10)))
+            cv = paths.FunctionCV("test", lambda s: s.xyz[0][0])
+            storage.save(traj)
+            storage.save(cv)
+
+            if cvs is not None:
+                cvs = [storage.cvs[cv] for cv in cvs]
+
+            precompute_func, blocks = precompute_cvs_func_and_inputs(
+                input_storage=storage,
+                cvs=cvs,
+                blocksize=2
+            )
+            assert len(blocks) == 5
+            for block in blocks:
+                assert len(block) == 2
+
+            # smoke test: only effect should be caching results
+            precompute_func(blocks[0])
+
+
+def test_rewrite_file():
+    # making a mock for storage instead of actually testing integration
+    class FakeStore(object):
+        def __init__(self):
+            self._stores = collections.defaultdict(list)
+
+        def store(self, obj, store_name):
+            self._stores[store_name].append(obj)
+
+    stage_names = ['foo', 'bar']
+    storage = FakeStore()
+    store_funcs = {
+        name: functools.partial(storage.store, store_name=name)
+        for name in stage_names
+    }
+    stage_mapping = {
+        'foo': (store_funcs['foo'], [0, 1, 2]),
+        'bar': (store_funcs['bar'], [[3], [4], [5]])
+    }
+    silent_tqdm = lambda x, desc=None, leave=True: x
+    with patch('paths_cli.file_copying.tqdm', silent_tqdm):
+        rewrite_file(stage_names, stage_mapping)
+
+    assert storage._stores['foo'] == [0, 1, 2]
+    assert storage._stores['bar'] == [[3], [4], [5]]