Merge pull request #28 from steinmig/unittests

Add CI pipeline

.github/workflows/python-app.yml

@@ -0,0 +1,47 @@
+name: Test NeuralForceField package
+
+on: [push]
+
+jobs:
+  build:
+
+    runs-on: ubuntu-latest
+    strategy:
+      matrix:
+        # python-version: ["pypy3.10", "3.8", "3.9", "3.10", "3.11", "3.12", "3.13"]
+        python-version: ["3.10"]
+
+    steps:
+      - uses: actions/checkout@v4
+      - name: Set up Python ${{ matrix.python-version }}
+        uses: actions/setup-python@v5
+        with:
+          python-version: ${{ matrix.python-version }}
+      - name: Display Python version
+        run: python -c "import sys; print(sys.version)"
+      - name: Install basics
+        run: python -m pip install --upgrade pip setuptools wheel
+      - name: Install package
+        run: python -m pip install .
+      # - name: Install linters
+      #   run: python -m pip install flake8 mypy pylint
+      # - name: Install documentation requirements
+      #   run: python -m pip install -r docs/requirements.txt
+      # - name: Test with flake8
+      #   run: flake8 polymethod
+      # - name: Test with mypy
+      #   run: mypy polymethod
+      # - name: Test with pylint
+      #   run: pylint polymethod
+      - name: Test with pytest
+        run: |
+         pip install pytest pytest-cov
+         pytest nff/tests --doctest-modules --junitxml=junit/test-results-${{ matrix.python-version }}.xml --cov=nff --cov-report=xml --cov-report=html
+      - name: Upload pytest test results
+        uses: actions/upload-artifact@v4
+        with:
+          name: pytest-results-${{ matrix.python-version }}
+          path: junit/test-results-${{ matrix.python-version }}.xml
+        if: ${{ always() }}
+      # - name: Test documentation
+      #   run: sphinx-build docs/source docs/build

.gitignore

@@ -66,5 +66,17 @@ dist/
 sandbox_excited/
 build/
 
+# Editor files
+# vim
+*.swp
+*.swo
+
+# pycharm
+.idea/
+
+# coverage and tests
+junit
+.coverage
+
 # required exceptions
 !tutorials/models/ammonia/Ammonia.xyz

nff/data/dataset.py

@@ -86,6 +86,7 @@ def __init__(
         units: str = "kcal/mol",
         check_props: bool = True,
         do_copy: bool = True,
+        device: str = "cuda"
     ) -> None:
         """Constructor for Dataset class.
 
@@ -108,6 +109,7 @@ def __init__(
             self.props = props
         self.units = units
         self.to_units(units)
+        self.device = device
 
     def __len__(self) -> int:
         """Length of the dataset.
@@ -289,6 +291,7 @@ def _get_periodic_neighbor_list(
                 pbc=True,
                 cutoff=cutoff,
                 directed=(not undirected),
+                device=self.device,
             )
             nbrs, offs = atoms.update_nbr_list()
             nbrlist.append(nbrs)
@@ -444,6 +447,7 @@ def unwrap_xyz(self, mol_dic: dict) -> None:
                 numbers=self.props["nxyz"][i][:, 0],
                 cell=self.props["cell"][i],
                 pbc=True,
+                device=self.device
             )
 
             # recontruct coordinates based on subgraphs index
@@ -577,6 +581,7 @@ def gen_bond_prior(self, cutoff: float, bond_len_dict: dict | None = None) -> No
                 "cutoff": cutoff,
                 "cell": cell,
                 "nbr_torch": False,
+                "device": self.device
             }
 
             # the coordinates have been unwrapped and try to results offsets

nff/tests/conftest.py

@@ -0,0 +1,15 @@
+
+import os
+import pytest
+import torch
+
+torch.set_num_threads(int(os.getenv("OMP_NUM_THREADS", 1)))
+
+
+def pytest_addoption(parser):
+    parser.addoption("--device", action="store", default="cpu", help="Whether to use the CPU or GPU for the tests")
+
+
+@pytest.fixture
+def device(request):
+    return request.config.getoption("--device")

nff/md/zhu_nakamura/dynamics_test.py → nff/tests/dynamics_test.py

@@ -15,8 +15,8 @@
 from ase.io.trajectory import Trajectory
 from ase import Atoms
 
-from nff.md.utils import mol_dot, mol_norm, ZhuNakamuraLogger, atoms_to_nxyz
-from nff.md.nvt_test import NoseHoover, NoseHooverChain
+from nff.md.utils_ax import mol_dot, mol_norm, ZhuNakamuraLogger, atoms_to_nxyz
+from nff.md.nvt_ax import NoseHoover, NoseHooverChain
 from nff.utils.constants import BOHR_RADIUS, FS_TO_AU, AMU_TO_AU, FS_TO_ASE, ASE_TO_FS, EV_TO_AU
 from nff.data import Dataset, collate_dicts
 from nff.utils.cuda import batch_to

nff/io/tests/test_ase.py → nff/tests/test_ase.py

@@ -5,6 +5,8 @@
 import numpy as np
 from ase import Atoms
 
+import pytest
+
 from nff.io.ase import AtomsBatch
 
 
@@ -19,6 +21,8 @@ def compare_dicts(d1: dict, d2: dict):
     for key, value in d1.items():
         if isinstance(value, dict):
             compare_dicts(value, d2[key])
+        elif isinstance(value, str):
+            assert value == d2[key]
         elif isinstance(value, Iterable):
             assert np.allclose(value, d2[key])
         else:
@@ -47,10 +51,17 @@ def get_ethanol():
     return Atoms(nxyz[:, 0].astype(int), positions=nxyz[:, 1:])
 
 
-# @ut.skip("skip this for now")
+@pytest.mark.usefixtures("device")  # Ensure the fixture is accessible
 class TestAtomsBatch(ut.TestCase):
     def setUp(self):
         self.ethanol = get_ethanol()
+        # Access the device value from the pytest fixture
+        self.device = self._test_fixture_device
+
+    @pytest.fixture(autouse=True)
+    def inject_device(self, device):
+        # Automatically set the fixture value to an attribute
+        self._test_fixture_device = device
 
     @ut.skip("skip this for now")
     def test_AtomsBatch(self):
@@ -111,7 +122,7 @@ def test_AtomsBatch(self):
             ]
         )
 
-        atoms_batch = AtomsBatch(self.ethanol, cutoff=2.5)
+        atoms_batch = AtomsBatch(self.ethanol, cutoff=2.5, device=self.device)
         atoms_batch.update_nbr_list()
 
         G1 = nx.from_edgelist(expected_nbrlist_cutoff_2dot5)
@@ -120,21 +131,21 @@ def test_AtomsBatch(self):
         assert nx.is_isomorphic(G1, G2)
 
     def test_get_batch(self):
-        atoms_batch = AtomsBatch(self.ethanol, cutoff=5)
+        atoms_batch = AtomsBatch(self.ethanol, cutoff=5, device=self.device)
         batch = atoms_batch.get_batch()
 
         assert "nxyz" in batch
 
     def test_from_atoms(self):
-        atoms_batch = AtomsBatch.from_atoms(self.ethanol, cutoff=2.5)
+        atoms_batch = AtomsBatch.from_atoms(self.ethanol, cutoff=2.5, device=self.device)
 
         # ensure atomic numbers, positions, and cell are the same
         assert np.allclose(atoms_batch.get_atomic_numbers(), self.ethanol.get_atomic_numbers())
         assert np.allclose(atoms_batch.get_positions(), self.ethanol.get_positions())
         assert np.allclose(atoms_batch.get_cell(), self.ethanol.get_cell())
 
     def test_copy(self):
-        atoms_batch = AtomsBatch(self.ethanol, cutoff=2.5)
+        atoms_batch = AtomsBatch(self.ethanol, cutoff=2.5, device=self.device)
         atoms_batch.get_batch()  # update props
         atoms_batch_copy = atoms_batch.copy()
 
@@ -154,7 +165,7 @@ def test_copy(self):
         assert atoms_batch.requires_large_offsets == atoms_batch_copy.requires_large_offsets
 
     def test_fromdict(self):
-        atoms_batch = AtomsBatch(self.ethanol, cutoff=2.5)
+        atoms_batch = AtomsBatch(self.ethanol, cutoff=2.5, device=self.device)
         ab_dict = atoms_batch.todict(update_props=True)
         ab_from_dict = AtomsBatch.fromdict(ab_dict)
 
@@ -183,6 +194,7 @@ def test_fromdict(self):
         compare_dicts(ab_dict_props, ab_dict_again_props)
 
 
+@pytest.mark.usefixtures("device")  # Ensure the fixture is loaded
 class TestPeriodic(ut.TestCase):
     def setUp(self):
         nxyz = np.array(
@@ -205,9 +217,15 @@ def setUp(self):
                 [0.0, 0.0, 5.51891759],
             ]
         )
-        self.quartz = AtomsBatch(nxyz[:, 0].astype(int), positions=nxyz[:, 1:], cell=lattice, pbc=True)
-
-    def test_ase(self):
+        self.quartz = AtomsBatch(nxyz[:, 0].astype(int), positions=nxyz[:, 1:], cell=lattice, pbc=True,
+                                 device=self._test_fixture_device)
+
+    @pytest.fixture(autouse=True)
+    def inject_device(self, device):
+        # Automatically set the fixture value to an attribute
+        self._test_fixture_device = device
+
+    def test_print(self):
         print(self.quartz)
 
     def test_nbrlist(self):
@@ -469,7 +487,6 @@ def test_nbrlist(self):
             ]
         )
         assert np.allclose(nbrlist, expected_nbrlist)
-        print(offsets)
 
 
 if __name__ == "__main__":

nff/data/tests/test_dataset.py → nff/tests/test_data/test_dataset.py

@@ -6,6 +6,8 @@
 import numpy as np
 import torch
 
+import pytest
+
 from nff.data.dataset import (
     Dataset,
     concatenate_dict,
@@ -14,8 +16,8 @@
 )
 
 current_path = Path(__file__).parent
-DATASET_PATH = current_path / "../../../tutorials/data/dataset.pth.tar"
-PEROVSKITE_DATA_PATH = current_path / "./data/SrIrO3_bulk_55_nff_all_dataset.pth.tar"
+DATASET_PATH = os.path.join(current_path, "..", "..", "..", "tutorials", "data", "dataset.pth.tar")
+PEROVSKITE_DATA_PATH = os.path.join(current_path, "data", "SrIrO3_bulk_55_nff_all_dataset.pth.tar")
 TARG_NAME = "formula"
 VAL_SIZE = 0.1
 TEST_SIZE = 0.1
@@ -223,6 +225,7 @@ def test_inexistent_list_lists(self):
         self.assertEqual(ab, expected)
 
 
+@pytest.mark.usefixtures("device")  # Ensure the fixture is accessible
 class TestPeriodicDataset(unittest.TestCase):
     def setUp(self):
         self.quartz = {
@@ -248,7 +251,12 @@ def setUp(self):
             ),
         }
 
-        self.qtz_dataset = Dataset(concatenate_dict(*[self.quartz] * 3))
+        self.qtz_dataset = Dataset(concatenate_dict(*[self.quartz] * 3), device=self._test_fixture_device)
+
+    @pytest.fixture(autouse=True)
+    def inject_device(self, device):
+        # Automatically set the fixture value to an attribute
+        self._test_fixture_device = device
 
     def test_neighbor_list(self):
         nbrs, offs = self.qtz_dataset.generate_neighbor_list(cutoff=5)