test for grid search

Author: skfegan

CodeEntropy/levels/search.py

@@ -186,7 +186,7 @@ def get_distance(self, j_position, i_position, dimensions):
 
         return distance
 
-    def get_grid_neighbors(self, universe, search_object, mol_id, highest_level):
+    def get_grid_neighbors(self, universe, mol_id, highest_level):
         """
         Use MDAnalysis neighbor search to find neighbors.
 
@@ -237,4 +237,4 @@ def get_grid_neighbors(self, universe, search_object, mol_id, highest_level):
             # residues from the central molecule
             neighbors = search - fragment.residues
 
-        return neighbors
+        return neighbors.fragindices

tests/unit/CodeEntropy/levels/test_search.py

@@ -71,6 +71,59 @@ def test_get_RAD_neighbors(tmp_path: Path):
     assert neighbors == [151, 3, 75, 219, 229, 488, 460, 118, 230, 326]
 
 
+def test_get_grid_neighbors(tmp_path: Path):
+    """
+    Args:
+        tmp_path: Pytest provided temporatry directory
+    """
+    args = {}
+    search = Search()
+    system = "methane"
+    repo_root = Path(__file__).resolve().parents[4]
+    config_path = (
+        repo_root / "tests" / "regression" / "configs" / system / "config.yaml"
+    )
+
+    tmp_path.mkdir(parents=True, exist_ok=True)
+
+    raw = yaml.safe_load(config_path.read_text())
+    if not isinstance(raw, dict):
+        raise ValueError(
+            f"Config must parse to a dict. Got {type(raw)} from {config_path}"
+        )
+
+    cooked = Helpers._abspathify_config_paths(raw, base_dir=config_path.parent)
+    required: list[Path] = []
+    run1 = cooked.get("run1")
+    if isinstance(run1, dict):
+        ff = run1.get("force_file")
+        if isinstance(ff, str) and ff:
+            required.append(Path(ff))
+        for p in run1.get("top_traj_file") or []:
+            if isinstance(p, str) and p:
+                required.append(Path(p))
+
+    if required:
+        Helpers.ensure_testdata_for_system(system, required_paths=required)
+
+    runner = CodeEntropyRunner(tmp_path)
+    parser = runner._config_manager.build_parser()
+    args, _ = parser.parse_known_args()
+    args.end = run1.get("end")
+    args.top_traj_file = run1.get("top_traj_file")
+    args.file_format = run1.get("file_format")
+    assert args.end == 1
+
+    universe_operations = UniverseOperations()
+    universe = CodeEntropyRunner._build_universe(args, universe_operations)
+
+    neighbors = search.get_grid_neighbors(
+        universe=universe, mol_id=0, highest_level="united_atom"
+    )
+
+    assert (neighbors == [151, 3, 75, 219]).all
+
+
 def test_get_angle():
     search = Search()
     result1 = search.get_angle(a, b, c, dimensions)