Merge branch 'main' into update_csv_logger

Author: AntonioMirarchi

.github/workflows/CI.yml

@@ -12,29 +12,66 @@ on:
 
 jobs:
   test:
-    runs-on: "ubuntu-latest"
+    runs-on: ${{ matrix.os }}
+
+    strategy:
+      fail-fast: false
+      matrix:
+        os: ["ubuntu-latest", "macos-latest", "macos-13", "windows-2022"]
+        python-version: ["3.10"]
+
     defaults: # Needed for conda
       run:
         shell: bash -l {0}
 
     steps:
       - name: Check out
-        uses: actions/checkout@v2
+        uses: actions/checkout@v4
+
+      - uses: conda-incubator/setup-miniconda@v3
+        with:
+          python-version: ${{ matrix.python-version }}
+          channels: conda-forge
+          conda-remove-defaults: "true"
+        if: matrix.os != 'macos-13'
 
-      - name: Create a conda environment
-        uses: conda-incubator/setup-miniconda@v2
+      - uses: conda-incubator/setup-miniconda@v3
         with:
-          activate-environment: torchmd-net
-          environment-file: environment.yml
-          miniforge-variant: Mambaforge
-          use-mamba: true
+          python-version: ${{ matrix.python-version }}
+          channels: conda-forge
+          mamba-version: "*"
+          conda-remove-defaults: "true"
+        if: matrix.os == 'macos-13'
 
-      - name: Install the package
-        run: pip -vv install .
+      - name: Install OS-specific conda dependencies
+        run: |
+          if [[ "${{ runner.os }}" == "Linux" ]]; then
+            conda install --file conda_deps_linux.txt --channel conda-forge --override-channels
+          elif [[ "${{ runner.os }}" == "macos-13" ]]; then
+            mamba install --file conda_deps_osx.txt --channel conda-forge --override-channels
+          elif [[ "${{ runner.os }}" == "macOS" ]]; then
+            conda install --file conda_deps_osx.txt --channel conda-forge --override-channels
+          elif [[ "${{ runner.os }}" == "Windows" ]]; then
+            conda install --file conda_deps_win.txt --channel conda-forge --override-channels
+          fi
+
+      - name: Install testing packages
+        run: conda install -y -c conda-forge flake8 pytest psutil
 
       - name: List the conda environment
         run: conda list
 
+      - name: Build and install the package
+        run: |
+          if [[ "${{ runner.os }}" == "Windows" ]]; then
+            export LIB="C:/Miniconda/envs/test/Library/lib"
+            pip -vv install .
+          else
+            pip -vv install .
+          fi
+        env:
+          CPU_ONLY: 1
+
       - name: Lint with flake8
         run: |
           # stop the build if there are Python syntax errors or undefined names
@@ -43,8 +80,13 @@ jobs:
           flake8 . --count --exit-zero --max-complexity=10 --max-line-length=127 --statistics
 
       - name: Run tests
-        run: pytest -v -s
+        run: pytest -v -s --durations=10
+        env:
+          CPU_ONLY: 1
+          SKIP_TORCH_COMPILE: ${{ runner.os == 'Windows' && 'true' || 'false' }}
+          OMP_PREFIX: ${{ runner.os == 'macOS' && '/Users/runner/miniconda3/envs/test' || '' }}
+          CPU_TRAIN: ${{ runner.os == 'macOS' && 'true' || 'false' }}
+          LONG_TRAIN: "true"
 
       - name: Test torchmd-train utility
         run: torchmd-train --help
-          

conda_deps_linux.txt

@@ -0,0 +1,10 @@
+h5py
+nnpops
+pip
+libtorch >=2.5.1
+pytorch-cpu >=2.5.1
+pytorch_geometric
+lightning
+torchmetrics
+tqdm
+gxx

conda_deps_osx.txt

@@ -0,0 +1,12 @@
+h5py
+nnpops
+pip
+libtorch>=2.5.1
+pytorch-cpu>=2.5.1
+pytorch_geometric
+lightning
+torchmetrics
+tqdm
+clangxx
+llvm-openmp
+pybind11

conda_deps_win.txt

@@ -0,0 +1,13 @@
+h5py
+pip
+libtorch >=2.5.1
+pytorch-cpu >=2.5.1
+pytorch_geometric
+lightning
+torchmetrics
+tqdm
+vc
+vc14_runtime
+vs2015_runtime
+vs2019_win-64
+sleef

environment.yml

@@ -3,13 +3,11 @@ channels:
   - conda-forge
 dependencies:
   - h5py
-  - matplotlib-base
   - nnpops
   - pip
   - pytorch
   - pytorch_geometric
   - lightning
-  - pydantic
   - torchmetrics
   - tqdm
   # Dev tools

setup.py

@@ -7,6 +7,9 @@
 import torch
 from torch.utils.cpp_extension import BuildExtension, CUDAExtension, include_paths, CppExtension
 import os
+import sys
+
+is_windows = sys.platform == 'win32'
 
 try:
     version = (
@@ -43,7 +46,7 @@ def set_torch_cuda_arch_list():
 ExtensionType = CppExtension if not use_cuda else CUDAExtension
 extensions = ExtensionType(
     name='torchmdnet.extensions.torchmdnet_extensions',
-    sources=[os.path.join(extension_root, "extensions.cpp")] + neighbor_sources,
+    sources=[os.path.join(extension_root, "torchmdnet_extensions.cpp")] + neighbor_sources,
     include_dirs=include_paths(),
     define_macros=[('WITH_CUDA', 1)] if use_cuda else [],
 )
@@ -58,5 +61,5 @@ def set_torch_cuda_arch_list():
             'build_ext': BuildExtension.with_options(no_python_abi_suffix=True, use_ninja=False)},
         include_package_data=True,
         entry_points={"console_scripts": ["torchmd-train = torchmdnet.scripts.train:main"]},
-        package_data={"torchmdnet": ["extensions/torchmdnet_extensions.so"]},
+        package_data={"torchmdnet": ["extensions/torchmdnet_extensions.so"] if not is_windows else ["extensions/torchmdnet_extensions.dll"]},
     )

tests/test_cfconv.py

@@ -4,21 +4,32 @@
 
 import pytest
 from pytest import mark
-import torch as pt
-from torchmdnet.models.torchmd_gn import CFConv as RefCFConv
-from torchmdnet.models.utils import OptimizedDistance, GaussianSmearing, ShiftedSoftplus
 
-from NNPOps.CFConv import CFConv
-from NNPOps.CFConvNeighbors import CFConvNeighbors
+try:
+    import NNPOps
 
+    nnpops_available = True
+except ImportError:
+    nnpops_available = False
 
+
+@pytest.mark.skipif(not nnpops_available, reason="NNPOps not available")
 @mark.parametrize("device", ["cpu", "cuda"])
 @mark.parametrize(
     ["num_atoms", "num_filters", "num_rbfs"],
     [(3, 5, 7), (3, 7, 5), (5, 3, 7), (5, 7, 3), (7, 3, 5), (7, 5, 3)],
 )
 @mark.parametrize("cutoff_upper", [5.0, 10.0])
 def test_cfconv(device, num_atoms, num_filters, num_rbfs, cutoff_upper):
+    import torch as pt
+    from torchmdnet.models.torchmd_gn import CFConv as RefCFConv
+    from torchmdnet.models.utils import (
+        OptimizedDistance,
+        GaussianSmearing,
+        ShiftedSoftplus,
+    )
+    from NNPOps.CFConv import CFConv
+    from NNPOps.CFConvNeighbors import CFConvNeighbors
 
     if not pt.cuda.is_available() and device == "cuda":
         pytest.skip("No GPU")

tests/test_datamodule.py

@@ -28,6 +28,7 @@ def test_datamodule_create(tmpdir):
     dl2 = data._get_dataloader(data.train_dataset, "train", store_dataloader=False)
     assert dl1 is not dl2
 
+
 def test_dataloader_get(tmpdir):
     args = load_example_args("graph-network")
     args["train_size"] = 800
@@ -50,17 +51,23 @@ def test_dataloader_get(tmpdir):
     # Assert that the dataloader is not empty
     assert len(data.train_dataloader()) > 0
 
+
 @mark.parametrize("energy,forces", [(True, True), (True, False), (False, True)])
 @mark.parametrize("has_atomref", [True, False])
 def test_datamodule_standardize(energy, forces, has_atomref, tmpdir):
     args = load_example_args("graph-network")
     args["standardize"] = True
-    args["train_size"] = 800
-    args["val_size"] = 100
-    args["test_size"] = 100
+    args["train_size"] = 80
+    args["val_size"] = 10
+    args["test_size"] = 10
     args["log_dir"] = tmpdir
+    args["batch_size"] = 2
+    args["inference_batch_size"] = 2
+    args["num_workers"] = 0
 
-    dataset = DummyDataset(energy=energy, forces=forces, has_atomref=has_atomref)
+    dataset = DummyDataset(
+        num_samples=100, energy=energy, forces=forces, has_atomref=has_atomref
+    )
     data = DataModule(args, dataset=dataset)
     data.prepare_data()
     data.setup("fit")

tests/test_dataset_comp6.py

@@ -39,7 +39,7 @@ def test_dataset_s66x8():
             ),
             atol=1e-4,
         )
-        assert pt.allclose(sample.y, pt.tensor([[-5755.7288331]],dtype=pt.float64))
+        assert pt.allclose(sample.y, pt.tensor([[-5755.7288331]], dtype=pt.float64))
         assert pt.allclose(
             sample.neg_dy,
             -pt.tensor(
@@ -65,3 +65,8 @@ def test_dataset_s66x8():
         data_loader = DataLoader(dataset=data_set, batch_size=32, num_workers=2)
         for batch in data_loader:
             pass
+
+        # Extra cleanup for Windows tests which are sensitive to open mmap files
+        del sample
+        del data_loader
+        del data_set

tests/test_module.py

@@ -12,6 +12,7 @@
 from torchmdnet.module import LNNP
 from torchmdnet.data import DataModule
 from torchmdnet import priors
+import os
 
 from utils import load_example_args, DummyDataset
 
@@ -29,7 +30,20 @@ def test_load_model():
 @mark.parametrize("model_name", models.__all_models__)
 @mark.parametrize("use_atomref", [True, False])
 @mark.parametrize("precision", [32, 64])
+@mark.skipif(
+    os.getenv("LONG_TRAIN", "false") == "false", reason="Skipping long train test"
+)
 def test_train(model_name, use_atomref, precision, tmpdir):
+    import torch
+
+    torch.set_num_threads(1)
+
+    accelerator = "auto"
+    if os.getenv("CPU_TRAIN", "false") == "true":
+        # OSX MPS backend runs out of memory on Github Actions
+        torch.set_default_device("cpu")
+        accelerator = "cpu"
+
     args = load_example_args(
         model_name,
         remove_prior=not use_atomref,
@@ -43,6 +57,64 @@ def test_train(model_name, use_atomref, precision, tmpdir):
         num_rbf=16,
         batch_size=8,
         precision=precision,
+        num_workers=0,
+    )
+    datamodule = DataModule(args, DummyDataset(has_atomref=use_atomref))
+
+    prior = None
+    if use_atomref:
+        prior = getattr(priors, args["prior_model"])(dataset=datamodule.dataset)
+        args["prior_args"] = prior.get_init_args()
+
+    module = LNNP(args, prior_model=prior)
+
+    trainer = pl.Trainer(
+        max_steps=10,
+        default_root_dir=tmpdir,
+        precision=args["precision"],
+        inference_mode=False,
+        accelerator=accelerator,
+        num_nodes=1,
+        devices=1,
+        use_distributed_sampler=False,
+    )
+    trainer.fit(module, datamodule)
+    trainer.test(module, datamodule)
+
+
+@mark.parametrize("model_name", models.__all_models__)
+@mark.parametrize("use_atomref", [True, False])
+@mark.parametrize("precision", [32, 64])
+def test_dummy_train(model_name, use_atomref, precision, tmpdir):
+    import torch
+
+    torch.set_num_threads(1)
+
+    accelerator = "auto"
+    if os.getenv("CPU_TRAIN", "false") == "true":
+        # OSX MPS backend runs out of memory on Github Actions
+        torch.set_default_device("cpu")
+        accelerator = "cpu"
+
+    extra_args = {}
+    if model_name != "tensornet":
+        extra_args["num_heads"] = 2
+
+    args = load_example_args(
+        model_name,
+        remove_prior=not use_atomref,
+        train_size=0.05,
+        val_size=0.01,
+        test_size=0.01,
+        log_dir=tmpdir,
+        derivative=True,
+        embedding_dimension=2,
+        num_layers=1,
+        num_rbf=4,
+        batch_size=2,
+        precision=precision,
+        num_workers=0,
+        **extra_args,
     )
     datamodule = DataModule(args, DummyDataset(has_atomref=use_atomref))
 
@@ -53,6 +125,15 @@ def test_train(model_name, use_atomref, precision, tmpdir):
 
     module = LNNP(args, prior_model=prior)
 
-    trainer = pl.Trainer(max_steps=10, default_root_dir=tmpdir, precision=args["precision"],inference_mode=False)
+    trainer = pl.Trainer(
+        max_steps=10,
+        default_root_dir=tmpdir,
+        precision=args["precision"],
+        inference_mode=False,
+        accelerator=accelerator,
+        num_nodes=1,
+        devices=1,
+        use_distributed_sampler=False,
+    )
     trainer.fit(module, datamodule)
     trainer.test(module, datamodule)