Triton neighbor list implementation

.github/workflows/publish.yml

@@ -6,166 +6,51 @@ on:
 
 jobs:
   build:
-    name: Build wheels on ${{ matrix.os }}-${{ matrix.accelerator }}
-    runs-on: ${{ matrix.os }}
-    strategy:
-      fail-fast: false
-      matrix:
-        os: [ubuntu-latest, ubuntu-24.04-arm, windows-2022, macos-latest]
-        accelerator: [cpu, cu118, cu126] #, cu128]
-        exclude:
-          - os: ubuntu-24.04-arm
-            accelerator: cu118
-          - os: ubuntu-24.04-arm
-            accelerator: cu126
-          # - os: ubuntu-24.04-arm
-          #   accelerator: cu128
-          - os: macos-latest
-            accelerator: cu118
-          - os: macos-latest
-            accelerator: cu126
-          # - os: macos-latest
-          #   accelerator: cu128
+    name: Create source distribution
+    runs-on: ubuntu-latest
 
     steps:
-      - name: Free space of Github Runner (otherwise it will fail by running out of disk)
-        if: matrix.os == 'ubuntu-latest'
-        run: |
-          sudo rm -rf /usr/share/dotnet
-          sudo rm -rf /opt/ghc
-          sudo rm -rf "/usr/local/share/boost"
-          sudo rm -rf "/usr/local/.ghcup"
-          sudo rm -rf "/usr/local/julia1.9.2"
-          sudo rm -rf "/usr/local/lib/android"
-          sudo rm -rf "$AGENT_TOOLSDIRECTORY"
-
-      - uses: actions/checkout@v4
+      - uses: actions/checkout@v5
 
-      - uses: actions/setup-python@v5
+      - uses: actions/setup-python@v6
         with:
           python-version: "3.13"
 
       - name: Install cibuildwheel
-        run: python -m pip install cibuildwheel==3.1.3
-
-      - name: Activate MSVC
-        uses: ilammy/msvc-dev-cmd@v1
-        with:
-          toolset: 14.29
-        if: matrix.os == 'windows-2022'
+        run: pip install build
 
-      - name: Build wheels
-        if: matrix.os != 'windows-2022'
-        shell: bash
-        run: python -m cibuildwheel --output-dir wheelhouse
-        env:
-          ACCELERATOR: ${{ matrix.accelerator }}
-          CPU_TRAIN: ${{ runner.os == 'macOS' && 'true' || 'false' }}
-
-      - name: Build wheels
-        if: matrix.os == 'windows-2022'
-        shell: cmd # Use cmd on Windows to avoid bash environment taking priority over MSVC variables
-        run: python -m cibuildwheel --output-dir wheelhouse
-        env:
-          ACCELERATOR: ${{ matrix.accelerator }}
-          DISTUTILS_USE_SDK: "1" # Windows requires this to use vc for building
-          SKIP_TORCH_COMPILE: "true"
+      - name: Build pypi package
+        run: python -m build --sdist
 
       - uses: actions/upload-artifact@v4
         with:
-          name: ${{ matrix.accelerator }}-cibw-wheels-${{ matrix.os }}-${{ strategy.job-index }}
-          path: ./wheelhouse/*.whl
+          name: source_dist
+          path: dist/*.tar.gz
 
-  publish-to-public-pypi:
+  publish-to-pypi:
     name: >-
       Publish Python 🐍 distribution 📦 to PyPI
     needs:
       - build
     runs-on: ubuntu-latest
     environment:
       name: pypi
+      url: https://pypi.org/p/torchmd-net
     permissions:
       id-token: write # IMPORTANT: mandatory for trusted publishing
-    strategy:
-      fail-fast: false
-      matrix:
-        accelerator: [cpu, cu118, cu126] #, hip, cu124, cu126, cu128]
 
     steps:
       - name: Download all the dists
         uses: actions/download-artifact@v4
         with:
-          pattern: "${{ matrix.accelerator }}-cibw-wheels*"
           path: dist/
           merge-multiple: true
 
       - name: Publish distribution 📦 to PyPI
         uses: pypa/gh-action-pypi-publish@release/v1
         with:
           password: ${{ secrets.TMDNET_PYPI_API_TOKEN }}
-          skip-existing: true
-
-  # publish-to-accelera-pypi:
-  #   name: >-
-  #     Publish Python 🐍 distribution 📦 to Acellera PyPI
-  #   needs:
-  #     - build
-  #   runs-on: ubuntu-latest
-  #   permissions: # Needed for GCP authentication
-  #     contents: "read"
-  #     id-token: "write"
-  #   strategy:
-  #     fail-fast: false
-  #     matrix:
-  #       accelerator: [cpu, cu118, cu126, cu128]
-
-  #   steps:
-  #     - uses: actions/checkout@v4 # Needed for GCP authentication for some reason
-
-  #     - name: Set up Cloud SDK
-  #       uses: google-github-actions/auth@v2
-  #       with:
-  #         workload_identity_provider: ${{ secrets.GCP_WORKLOAD_IDENTITY_PROVIDER }}
-  #         service_account: ${{ secrets.GCP_PYPI_SERVICE_ACCOUNT }}
-
-  #     - name: Download all the dists
-  #       uses: actions/download-artifact@v4
-  #       with:
-  #         pattern: "${{ matrix.accelerator }}-cibw-wheels*"
-  #         path: dist/
-  #         merge-multiple: true
-
-  #     - name: Publish distribution 📦 to Acellera PyPI
-  #       run: |
-  #         pip install build twine keyring keyrings.google-artifactregistry-auth
-  #         pip install -U packaging
-  #         twine upload --repository-url https://us-central1-python.pkg.dev/pypi-packages-455608/${{ matrix.accelerator }} dist/* --verbose --skip-existing
-
-  # publish-to-official-pypi:
-  #   name: >-
-  #     Publish Python 🐍 distribution 📦 to PyPI
-  #   needs:
-  #     - build
-  #   runs-on: ubuntu-latest
-  #   environment:
-  #     name: pypi
-  #     url: https://pypi.org/p/torchmd-net
-  #   permissions:
-  #     id-token: write # IMPORTANT: mandatory for trusted publishing
-
-  #   steps:
-  #     - name: Download all the dists
-  #       uses: actions/download-artifact@v4
-  #       with:
-  #         pattern: "cu118-cibw-wheels*"
-  #         path: dist/
-  #         merge-multiple: true
-
-  #     - name: Publish distribution 📦 to PyPI
-  #       uses: pypa/gh-action-pypi-publish@release/v1
-  #       with:
-  #         password: ${{ secrets.TMDNET_PYPI_API_TOKEN }}
-  #         skip_existing: true
+          skip_existing: true
 
   github-release:
     name: >-

.github/workflows/test.yml

@@ -17,93 +17,28 @@ jobs:
           ["ubuntu-latest", "ubuntu-22.04-arm", "macos-latest", "windows-2022"]
         python-version: ["3.13"]
 
-    defaults: # Needed for conda
-      run:
-        shell: bash -l {0}
-
     steps:
       - name: Check out
-        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'
+        uses: actions/checkout@v5
 
-      - uses: conda-incubator/setup-miniconda@v3
+      - name: Install uv
+        uses: astral-sh/setup-uv@v7
         with:
           python-version: ${{ matrix.python-version }}
-          channels: conda-forge
-          mamba-version: "*"
-          conda-remove-defaults: "true"
-        if: matrix.os == 'macos-13'
-
-      - name: Install OS-specific compilers
-        run: |
-          if [[ "${{ matrix.os }}" == "ubuntu-22.04-arm" ]]; then
-            conda install gxx --channel conda-forge --override-channels
-          elif [[ "${{ runner.os }}" == "Linux" ]]; then
-            conda install gxx --channel conda-forge --override-channels
-          elif [[ "${{ runner.os }}" == "macOS" ]]; then
-            conda install llvm-openmp pybind11 --channel conda-forge --override-channels
-            echo "CC=clang" >> $GITHUB_ENV
-            echo "CXX=clang++" >> $GITHUB_ENV
-          elif [[ "${{ runner.os }}" == "Windows" ]]; then
-            conda install vc vc14_runtime vs2015_runtime --channel conda-forge --override-channels
-          fi
-
-      - name: List the conda environment
-        run: conda list
 
-      - name: Install testing packages
-        run: conda install -y -c conda-forge flake8 pytest psutil python-build
+      - name: Install the project
+        run: uv sync --all-extras --dev
 
       - name: Lint with flake8
         run: |
           # stop the build if there are Python syntax errors or undefined names
-          flake8 . --count --select=E9,F63,F7,F82 --show-source --statistics
+          uv run flake8 ./torchmdnet --count --select=E9,F63,F7,F82 --show-source --statistics
           # exit-zero treats all errors as warnings. The GitHub editor is 127 chars wide
-          flake8 . --count --exit-zero --max-complexity=10 --max-line-length=127 --statistics
-
-      - name: Set pytorch index
-        run: |
-          if [[ "${{ runner.os }}" == "Windows" ]]; then
-            mkdir -p "C:\ProgramData\pip"
-            echo "[global]
-            extra-index-url = https://download.pytorch.org/whl/cpu" > "C:\ProgramData\pip\pip.ini"
-          else
-            mkdir -p $HOME/.config/pip
-            echo "[global]
-            extra-index-url = https://download.pytorch.org/whl/cpu" > $HOME/.config/pip/pip.conf
-          fi
-
-      - name: Build and install the package
-        run: |
-          if [[ "${{ runner.os }}" == "Windows" ]]; then
-            export LIB="C:/Miniconda/envs/test/Library/lib"
-          fi
-          python -m build            
-          pip install dist/*.whl
-        env:
-          ACCELERATOR: "cpu"
-
-      # - name: Install nnpops
-      #   if: matrix.os == 'ubuntu-latest' || matrix.os == 'macos-latest'
-      #   run: conda install nnpops --channel conda-forge --override-channels
-
-      - name: List the conda environment
-        run: conda list
+          uv run flake8 ./torchmdnet --count --exit-zero --max-complexity=10 --max-line-length=127 --statistics
 
       - name: Run tests
-        run: pytest -v -s --durations=10
-        env:
-          ACCELERATOR: "cpu"
-          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"
+        # For example, using `pytest`
+        run: uv run pytest tests
 
       - name: Test torchmd-train utility
-        run: torchmd-train --help
+        run: uv run torchmd-train --help

README.md

@@ -21,19 +21,17 @@ Documentation is available at https://torchmd-net.readthedocs.io
 
 
 ## Installation  
-TorchMD-Net is available as a pip installable wheel as well as in [conda-forge](https://conda-forge.org/)
+TorchMD-Net is available as a pip package as well as in [conda-forge](https://conda-forge.org/)
 
-TorchMD-Net provides builds for CPU-only, CUDA 11 and CUDA 12. CPU versions are only provided as reference,
-as the performance will be extremely limited.
-Depending on which variant you wish to install, you can install it with one of the following commands:
+As TorchMD-Net depends on PyTorch we need to add additional index URLs to the installation command as per [pytorch](https://pytorch.org/get-started/locally/)
 
 ```sh
-# The following will install the CUDA 11.8 version
-pip install torchmd-net-cu11 --extra-index-url https://download.pytorch.org/whl/cu118
-# The following will install the CUDA 12.4 version
-pip install torchmd-net-cu12 --extra-index-url https://download.pytorch.org/whl/cu124
-# The following will install the CPU only version (not recommended)
-pip install torchmd-net-cpu --extra-index-url https://download.pytorch.org/whl/cpu
+# The following will install TorchMD-Net with PyTorch CUDA 11.8 version
+pip install torchmd-net --extra-index-url https://download.pytorch.org/whl/cu118
+# The following will install TorchMD-Net with PyTorch CUDA 12.4 version
+pip install torchmd-net --extra-index-url https://download.pytorch.org/whl/cu124
+# The following will install TorchMD-Net with PyTorch CPU only version (not recommended)
+pip install torchmd-net --extra-index-url https://download.pytorch.org/whl/cpu
 ```
 
 Alternatively it can be installed with conda or mamba with one of the following commands.
@@ -46,7 +44,7 @@ mamba install torchmd-net cuda-version=12.4
 
 ### Install from source  
 
-TorchMD-Net is installed using pip, but you will need to install some dependencies before. Check [this documentation page](https://torchmd-net.readthedocs.io/en/latest/installation.html#install-from-source).  
+TorchMD-Net is installed using pip with `pip install -e .` to create an editable install.  
 
 ## Usage
 Specifying training arguments can either be done via a configuration yaml file or through command line arguments directly. Several examples of architectural and training specifications for some models and datasets can be found in [examples/](https://github.com/torchmd/torchmd-net/tree/main/examples). Note that if a parameter is present both in the yaml file and the command line, the command line version takes precedence.