Lens fix and new tests

src/dynsight/_internal/lens/lens.py

@@ -78,15 +78,15 @@ def neighbor_list_celllist_centers(  # noqa: C901, PLR0912
     n_cent = positions_cent.shape[0]
     r_cut2 = (r_cut - 1e-6) ** 2
 
-    cell_ids, head, next_, n_cell = build_cell_list(positions_env, box, r_cut)
+    _, head, next_, n_cell = build_cell_list(positions_env, box, r_cut)
     nx, ny, nz = n_cell
 
     n_neigh = np.zeros(n_cent, dtype=np.int32)
     # ---- count the neighbors for each center ----
     for i in prange(n_cent):
-        cx = cell_ids[i, 0]  # int(positions_cent[i, 0] / box[0] * nx) % nx
-        cy = cell_ids[i, 1]  # int(positions_cent[i, 1] / box[1] * ny) % ny
-        cz = cell_ids[i, 2]  # int(positions_cent[i, 2] / box[2] * nz) % nz
+        cx = int(positions_cent[i, 0] / box[0] * nx) % nx
+        cy = int(positions_cent[i, 1] / box[1] * ny) % ny
+        cz = int(positions_cent[i, 2] / box[2] * nz) % nz
 
         for dx in (-1, 0, 1):
             for dy in (-1, 0, 1):

tests/lens/case_data.py

@@ -0,0 +1,12 @@
+from dataclasses import dataclass
+
+
+@dataclass(frozen=True, slots=True)
+class LENSCaseData:
+    expected_lens: str
+    r_cut: float
+    delay: int
+    centers: str
+    selection: str
+    n_jobs: int
+    name: str

tests/lens/conftest.py

@@ -0,0 +1,74 @@
+import pytest
+
+from tests.lens.case_data import LENSCaseData
+
+
+@pytest.fixture(
+    scope="session",
+    params=(
+        # Case 0: Default case
+        lambda name: LENSCaseData(
+            expected_lens="c0_lens_rc3_d1_all_all_1.npy",
+            r_cut=3,
+            delay=1,
+            centers="all",
+            selection="all",
+            n_jobs=1,
+            name=name,
+        ),
+        # Case 1: changing cutoff
+        lambda name: LENSCaseData(
+            expected_lens="c1_lens_rc4_d1_all_all_1.npy",
+            r_cut=4,
+            delay=1,
+            centers="all",
+            selection="all",
+            n_jobs=1,
+            name=name,
+        ),
+        # Case 2: single center
+        lambda name: LENSCaseData(
+            expected_lens="c2_lens_rc4_d1_1_all_1.npy",
+            r_cut=4,
+            delay=1,
+            centers="id 1",
+            selection="all",
+            n_jobs=1,
+            name=name,
+        ),
+        # Case 3: single selection
+        lambda name: LENSCaseData(
+            expected_lens="c3_lens_rc4_d1_all_1_1.npy",
+            r_cut=4,
+            delay=1,
+            centers="all",
+            selection="id 1",
+            n_jobs=1,
+            name=name,
+        ),
+        # Case 4: parallel
+        lambda name: LENSCaseData(
+            expected_lens="c4_lens_rc4_d1_all_all_2.npy",
+            r_cut=4,
+            delay=1,
+            centers="all",
+            selection="all",
+            n_jobs=2,
+            name=name,
+        ),
+        # Case 5: changing delay
+        lambda name: LENSCaseData(
+            expected_lens="c5_lens_rc4_d2_all_all_1.npy",
+            r_cut=4,
+            delay=2,
+            centers="all",
+            selection="all",
+            n_jobs=1,
+            name=name,
+        ),
+    ),
+)
+def case_data(request: pytest.FixtureRequest) -> LENSCaseData:
+    return request.param(
+        f"{request.fixturename}{request.param_index}",  # type: ignore [attr-defined]
+    )

tests/lens/test_lens.py

@@ -1,94 +1,37 @@
-"""Test the consistency of LENS calculations with a control calculation.
-
-* Original author: Martina Crippa
-
-This test verifies that the LENS calculation (LENS and nn) yields the same
-values as a control calculation at different r_cut.
-
-Control file path:
-    - tests/systems/2_particles.xyz
-
-Dynsight function tested:
-    - dynsight.lens.list_neighbours_along_trajectory()
-    - dynsight.lens.compute_lens_over_trj()
-
-r_cuts checked:
-    - [2.5, 3.0, 3.5, 4.0, 4.5, 5.0, 5.5]
-"""
-
-from __future__ import annotations
+"""Pytest for dynsight.lens.compute_lens."""
 
 from pathlib import Path
 
+import MDAnalysis
 import numpy as np
 import pytest
 
 from dynsight.trajectory import Trj
 
-LENS_CUTOFF = [2.5, 3.0, 3.5, 4.0, 4.5, 5.0, 5.5]
-
-# ---------------- Fixtures ----------------
-
-
-@pytest.fixture(scope="module")
-def here() -> Path:
-    return Path(__file__).parent
-
+from .case_data import LENSCaseData
 
-@pytest.fixture(scope="module")
-def file_paths(here: Path) -> dict[str, Path]:
-    return {
-        "xyz": here / "../systems/2_particles.xyz",
-        "check_file": here / "../systems/LENS.npz",
-    }
 
+def test_lens(case_data: LENSCaseData) -> None:
+    original_dir = Path(__file__).resolve().parent
+    topology_file = original_dir / "../systems/balls_7_nvt.gro"
+    trajectory_file = original_dir / "../systems/balls_7_nvt.xtc"
+    expected_lens = original_dir / "test_lens" / case_data.expected_lens
+    universe = MDAnalysis.Universe(topology_file, trajectory_file)
 
-@pytest.fixture(scope="module")
-def trajectory(file_paths: dict[str, Path]) -> Trj:
-    return Trj.init_from_xyz(file_paths["xyz"], dt=1)
+    example_trj = Trj(universe)
 
+    test_lens = example_trj.get_lens(
+        r_cut=case_data.r_cut,
+        delay=case_data.delay,
+        centers=case_data.centers,
+        selection=case_data.selection,
+        n_jobs=case_data.n_jobs,
+    )
 
-# ---------------- Tests ----------------
-def test_lens(
-    trajectory: Trj,
-    file_paths: dict[str, Path],
-) -> None:
-    """Test the consistency of LENS calculations."""
-    check_file = np.load(file_paths["check_file"])
-
-    # Run LENS (and NN) calculation for different r_cuts
-    for i, r_cut in enumerate(LENS_CUTOFF):
-        reference_array = check_file[f"LENS_{i}"]
-
-        test_lens = trajectory.get_lens(r_cut=r_cut, respect_pbc=False)
-        test_lens_ds = np.array(
-            [np.concatenate(([0.0], tmp)) for tmp in test_lens.dataset]
-        )  # the original LENS function gave always 0.0 as first frame
-
-        assert np.allclose(reference_array[0], test_lens_ds), (
-            "LENS analyses provided different values "
-            f"compared to the control system for r_cut: {r_cut}."
-        )
-
-
-def test_nn(
-    trajectory: Trj,
-    file_paths: dict[str, Path],
-) -> None:
-    """Test the consistency of NN calculations."""
-    check_file = np.load(file_paths["check_file"])
-
-    # Run LENS (and NN) calculation for different r_cuts
-    for i, r_cut in enumerate(LENS_CUTOFF):
-        reference_array = check_file[f"LENS_{i}"]
-
-        _, test_nn = trajectory.get_coord_number(
-            r_cut=r_cut,
-            respect_pbc=False,
-        )
-        test_nn_ds = test_nn.dataset
-
-        assert np.allclose(reference_array[1], test_nn_ds), (
-            "NN analyses provided different values "
-            f"compared to the control system for r_cut: {r_cut}."
+    if not expected_lens.exists():
+        np.save(expected_lens, test_lens.dataset)
+        pytest.fail(
+            "LENS test files were not present. They have been created."
         )
+    exp_lens = np.load(expected_lens)
+    assert np.allclose(exp_lens, test_lens.dataset, atol=1e-6)

tests/neighbors/case_data.py

@@ -0,0 +1,11 @@
+from dataclasses import dataclass
+
+
+@dataclass(frozen=True, slots=True)
+class NNCaseData:
+    expected_nn: str
+    r_cut: float
+    centers: str
+    selection: str
+    n_jobs: int
+    name: str

tests/neighbors/conftest.py

@@ -0,0 +1,59 @@
+import pytest
+
+from tests.neighbors.case_data import NNCaseData
+
+
+@pytest.fixture(
+    scope="session",
+    params=(
+        # Case 0: Default case
+        lambda name: NNCaseData(
+            expected_nn="c0_nn_rc3_all_all_1.npy",
+            r_cut=3,
+            centers="all",
+            selection="all",
+            n_jobs=1,
+            name=name,
+        ),
+        # Case 1: changing cutoff
+        lambda name: NNCaseData(
+            expected_nn="c1_nn_rc4_all_all_1.npy",
+            r_cut=4,
+            centers="all",
+            selection="all",
+            n_jobs=1,
+            name=name,
+        ),
+        # Case 2: single center
+        lambda name: NNCaseData(
+            expected_nn="c2_nn_rc4_1_all_1.npy",
+            r_cut=4,
+            centers="id 1",
+            selection="all",
+            n_jobs=1,
+            name=name,
+        ),
+        # Case 3: single selection
+        lambda name: NNCaseData(
+            expected_nn="c3_nn_rc4_all_1_1.npy",
+            r_cut=4,
+            centers="all",
+            selection="id 1",
+            n_jobs=1,
+            name=name,
+        ),
+        # Case 4: parallel
+        lambda name: NNCaseData(
+            expected_nn="c4_nn_rc4_all_all_2.npy",
+            r_cut=4,
+            centers="all",
+            selection="all",
+            n_jobs=2,
+            name=name,
+        ),
+    ),
+)
+def case_data(request: pytest.FixtureRequest) -> NNCaseData:
+    return request.param(
+        f"{request.fixturename}{request.param_index}",  # type: ignore [attr-defined]
+    )

tests/neighbors/test_nn.py

@@ -0,0 +1,34 @@
+"""Pytest for dynsight.Trj.get_coord_number."""
+
+from pathlib import Path
+
+import MDAnalysis
+import numpy as np
+import pytest
+
+from dynsight.trajectory import Trj
+
+from .case_data import NNCaseData
+
+
+def test_nn(case_data: NNCaseData) -> None:
+    original_dir = Path(__file__).resolve().parent
+    topology_file = original_dir / "../systems/balls_7_nvt.gro"
+    trajectory_file = original_dir / "../systems/balls_7_nvt.xtc"
+    expected_nn = original_dir / "test_nn" / case_data.expected_nn
+    universe = MDAnalysis.Universe(topology_file, trajectory_file)
+
+    example_trj = Trj(universe)
+
+    _, test_nn = example_trj.get_coord_number(
+        r_cut=case_data.r_cut,
+        centers=case_data.centers,
+        selection=case_data.selection,
+        n_jobs=case_data.n_jobs,
+    )
+
+    if not expected_nn.exists():
+        np.save(expected_nn, test_nn.dataset)
+        pytest.fail("NN test files were not present. They have been created.")
+    exp_nn = np.load(expected_nn)
+    assert np.allclose(exp_nn, test_nn.dataset, atol=1e-6)