Merge branch 'feat/set_default_to_none_for_atom_selection' into mikem…

…henry-patch-3
choderalab · Jun 9, 2023 · 1505004 · 1505004
2 parents 0474f9a + bd3a6e6
commit 1505004
Show file tree

Hide file tree

Showing 4 changed files with 49 additions and 33 deletions.
diff --git a/examples/new-cli/template.yaml b/examples/new-cli/template.yaml
@@ -67,7 +67,7 @@ trajectory_directory: lig0to1
 trajectory_prefix: out
 
 # Atoms to store in NetCDF files (MDTraj selection syntax)
-atom_selection: not water
+atom_selection: null
 
 # Calculation phases to run
 # Permitted phases: ['complex', 'solvent', 'vacuum']

diff --git a/perses/app/setup_relative_calculation.py b/perses/app/setup_relative_calculation.py
@@ -584,7 +584,7 @@ def run_setup(setup_options, serialize_systems=True, build_samplers=True):
         _logger.info(f"\tatom selection detected: {atom_selection}")
     else:
         _logger.info(f"\tno atom selection detected: default to all.")
-        atom_selection = 'all'
+        atom_selection = 'none'
 
     if setup_options['fe_type'] == 'neq':
         _logger.info(f"\tInstantiating nonequilibrium switching FEP")
@@ -660,7 +660,7 @@ def run_setup(setup_options, serialize_systems=True, build_samplers=True):
                 if atom_selection:
                     selection_indices = htf[phase].hybrid_topology.select(atom_selection)
                 else:
-                    selection_indices = None
+                    selection_indices = []
 
                 storage_name = AnyPath(trajectory_directory) / f"{trajectory_prefix}-{phase}.nc"
                 _logger.info(f'\tstorage_name: {storage_name}')

diff --git a/perses/tests/conftest.py b/perses/tests/conftest.py
@@ -1,4 +1,5 @@
 import pytest
+import yaml
 
 
 def pytest_addoption(parser):
@@ -25,3 +26,44 @@ def pytest_collection_modifyitems(config, items):
 def in_tmpdir(tmpdir):
     with tmpdir.as_cwd():
         yield
+
+@pytest.fixture
+def input_yaml_template():
+    input_yaml = """
+                atom_selection: null
+                checkpoint_interval: 5
+                fe_type: repex
+                forcefield_files:
+                - amber/ff14SB.xml
+                - amber/tip3p_standard.xml
+                - amber/tip3p_HFE_multivalent.xml
+                - amber/phosaa10.xml
+                n_cycles: 10
+                n_equilibration_iterations: 10
+                n_states: 3
+                n_steps_per_move_application: 50
+                new_ligand_index: 15
+                old_ligand_index: 14
+                phases:
+                - vacuum
+                pressure: 1.0
+                save_setup_pickle_as: fesetup_hbonds.pkl
+                small_molecule_forcefield: openff-2.0.0
+                solvent_padding: 9.0
+                temperature: 300.0
+                timestep: 4.0
+                trajectory_directory: cdk2_repex_hbonds
+                trajectory_prefix: cdk2
+                """
+    return input_yaml
+
+@pytest.fixture
+def input_template_obj_default_selection(input_yaml_template):
+    input_obj = yaml.safe_load(input_yaml_template)
+    return input_obj
+
+
+@pytest.fixture
+def input_template_not_water_selection(input_template_obj_default_selection):
+    input_template_obj_default_selection["atom_selection"] = "not water"
+    return input_template_obj_default_selection
diff --git a/perses/tests/test_resume.py b/perses/tests/test_resume.py
@@ -1,7 +1,6 @@
 import logging
 import os
 import subprocess
-import tempfile
 
 import pytest
 import simtk.openmm.app as app
@@ -20,7 +19,9 @@
 
 
 @pytest.mark.gpu_ci
-def test_cli_resume_repex():
+@pytest.mark.parametrize("input_params", ["input_template_obj_default_selection",
+                                          "input_template_not_water_selection"])
+def test_cli_resume_repex(input_params, request):
 
     with enter_temp_directory() as temp_dir:
         os.chdir(temp_dir)
@@ -33,34 +34,7 @@ def test_cli_resume_repex():
             os.path.join("data", "cdk2-example", "CDK2_ligands_shifted.sdf"),
         )
 
-        document = """
-        atom_selection: not water
-        checkpoint_interval: 5
-        fe_type: repex
-        forcefield_files:
-        - amber/ff14SB.xml
-        - amber/tip3p_standard.xml
-        - amber/tip3p_HFE_multivalent.xml
-        - amber/phosaa10.xml
-        n_cycles: 10
-        n_equilibration_iterations: 10
-        n_states: 3
-        n_steps_per_move_application: 50
-        new_ligand_index: 15
-        old_ligand_index: 14
-        phases:
-        - vacuum
-        pressure: 1.0
-        save_setup_pickle_as: fesetup_hbonds.pkl
-        small_molecule_forcefield: openff-2.0.0
-        solvent_padding: 9.0
-        temperature: 300.0
-        timestep: 4.0
-        trajectory_directory: cdk2_repex_hbonds
-        trajectory_prefix: cdk2
-        """
-
-        y_doc = yaml.load(document, Loader=yaml.UnsafeLoader)
+        y_doc = request.getfixturevalue(input_params)
         y_doc["protein_pdb"] = protein_pdb
         y_doc["ligand_file"] = ligand_file