Install ci

.github/workflows/install-ci.yml

@@ -0,0 +1,67 @@
+name: Install CI
+
+on:
+  push:
+    branches: [main]
+  pull_request:
+    branches: [main, v2.0-refactor]
+  schedule:
+    - cron: '0 6 * * *'  # 6 AM UTC daily
+  workflow_dispatch:  # Allows manual triggering from GitHub UI
+
+jobs:
+  test-pip:
+    name: Test with pip (Python ${{ matrix.python-version }}, OS ${{ matrix.os }})
+    runs-on: ${{ matrix.os }}
+    strategy:
+      fail-fast: false
+      matrix:
+        os: [ubuntu-latest, macos-latest]
+        python-version: ["3.11","3.12","3.13"]
+
+    steps:
+    - uses: actions/checkout@v4
+
+    - uses: actions/setup-python@v5
+      with:
+        python-version: ${{ matrix.python-version }}
+
+    - name: Install package with pip
+      run: |
+        python -m pip install --upgrade pip
+        # pip 25.1+ supports dependency groups natively via --group
+        pip install -e . --group dev
+
+    - name: Run tests
+      run: |
+        pytest test/
+
+  test-uv:
+    name: Test with uv (Python ${{ matrix.python-version }}, OS ${{ matrix.os }})
+    runs-on: ${{ matrix.os }}
+    strategy:
+      fail-fast: false
+      matrix:
+        os: [ubuntu-latest, macos-latest]
+        python-version: ["3.11","3.12","3.13"]
+
+    steps:
+    - uses: actions/checkout@v4
+
+    - uses: actions/setup-python@v5
+      with:
+        python-version: ${{ matrix.python-version }}
+
+    - name: Install uv
+      run: |
+        curl -LsSf https://astral.sh/uv/install.sh | sh
+        echo "$HOME/.cargo/bin" >> $GITHUB_PATH
+
+    - name: Install package with uv
+      run: |
+        # Install core dependencies and development group
+        UV_PREVIEW=1 uv sync --group dev --preview-features extra-build-dependencies
+
+    - name: Run tests
+      run: |
+        uv run pytest test/

README.md

@@ -7,6 +7,8 @@
 [![Project Status: Active - The project has reached a stable, usable state and is being actively developed.](https://www.repostatus.org/badges/latest/active.svg)](https://www.repostatus.org/#active)
 [![GitHub](https://img.shields.io/github/license/NVIDIA/physicsnemo)](https://github.com/NVIDIA/physicsnemo/blob/master/LICENSE.txt)
 [![Code style: black](https://img.shields.io/badge/code%20style-black-000000.svg)](https://github.com/psf/black)
+[![Install CI](https://github.com/coreyjadams/physicsnemo/actions/workflows/install-ci.yml/badge.svg)](https://github.com/coreyjadams/physicsnemo/actions/workflows/install-ci.yml)
+
 <!-- markdownlint-enable -->
 [**NVIDIA PhysicsNeMo**](#what-is-physicsnemo)
 | [**Documentation**](https://docs.nvidia.com/deeplearning/physicsnemo/physicsnemo-core/index.html)

physicsnemo/nn/neighbors/_radius_search/_warp_impl.py

@@ -100,7 +100,8 @@ def count_neighbors(
         torch.cumsum(torch_result_count, dim=0, out=torch_offset[1:])
 
         # Create a pinned buffer on CPU to receive the count
-        pinned_buffer = torch.zeros(1, dtype=torch.int32, pin_memory=True)
+        pin_memory = torch.cuda.is_available()
+        pinned_buffer = torch.zeros(1, dtype=torch.int32, pin_memory=pin_memory)
         # Copy the last element to pinned memory
         pinned_buffer.copy_(torch_offset[-1:])
         total_count = pinned_buffer.item()

pyproject.toml

@@ -9,10 +9,8 @@ authors = [
 description = "A deep learning framework for AI-driven multi-physics systems"
 readme = "README.md"
 
-# Python 3.10 is EOL in 2026.  Migrating to 3.11+
-# Pytorch support for torch.compile in 3.14 is not ready as of Dec 2025.
-# Remove constraint against 3.14 when it arrives.
-requires-python = ">=3.10,<3.14"
+
+requires-python = ">=3.11,<3.14"
 
 license = "Apache-2.0"
 
@@ -59,21 +57,25 @@ Changelog = "https://github.com/NVIDIA/physicsnemo/blob/main/CHANGELOG.md"
 #####################################################################
 
 [tool.uv]
-no-build-isolation-package = [
-    "torch_scatter",
-    "torch_cluster",
-    "earth2grid",
-]
 managed = true
 
-# torch-sparse requires us to tell uv about it's dependency on torch.
 [tool.uv.extra-build-dependencies]
 torch-sparse = ["torch"]
+torch-cluster = ["torch"]
+torch-scatter = ["torch"]
 earth2grid = ["setuptools", "torch"] 
 
-# Earth2 grid is not on pypi ....
+[[tool.uv.index]]
+name = "nvidia"
+url = "https://pypi.nvidia.com"
+
+
 [tool.uv.sources]
 earth2grid = { url = "https://github.com/NVlabs/earth2grid/archive/main.tar.gz" }
+cuml-cu13 = { index = "nvidia" }
+pylibraft-cu13 = { index = "nvidia" }
+cuml-cu12 = { index = "nvidia" } # TODO setup optional install groups for CUDA 12 installations
+pylibraft-cu12 = { index = "nvidia" }
 
 #####################################################################
 # Flags Controlling the local build of physicsnemo
@@ -132,11 +134,11 @@ utils-extras = [
     "line_profiler",
     "vtk",
     "stl",
-    "rmm",
-    "cupy",
+    "pylibraft-cu13", # Provides rmm
+    "cupy-cuda13x",
 ]
 nn-extras = [
-    "cuml",
+    "cuml-cu13",
     "transformer_engine[pytorch]",
     "scipy",
     "natten",
@@ -210,4 +212,4 @@ Pangu = "physicsnemo.models.pangu:Pangu"
 Fengwu = "physicsnemo.models.fengwu:Fengwu"
 SwinRNN = "physicsnemo.models.swinvrnn:SwinRNN"
 EDMPrecondSR = "physicsnemo.models.diffusion:EDMPrecondSR"
-UNet = "physicsnemo.models.diffusion:UNet"
+UNet = "physicsnemo.models.diffusion:UNet"

test/metrics/test_metrics_general.py

@@ -608,6 +608,10 @@ def test_wasserstein(device, mean, variance, rtol: float = 1e-3, atol: float = 1
 
 
 def test_means_var(device, rtol: float = 1e-3, atol: float = 1e-3):
+
+    if not torch.cuda.is_available():
+        pytest.skip("CUDA required for this test.")
+
     DistributedManager._shared_state = {}
     if (device == "cuda:0") and (not DistributedManager.is_initialized()):
         os.environ["MASTER_ADDR"] = "localhost"

test/nn/neighbors/test_radius_search.py

@@ -14,8 +14,6 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
-import sys
-
 import pytest
 import torch
 
@@ -25,7 +23,6 @@
 )
 
 
-@pytest.mark.skipif(sys.platform == "darwin", reason="Torch doesn't support pinned memory on mac")
 @pytest.mark.parametrize("return_dists", [True, False])
 @pytest.mark.parametrize("return_points", [True, False])
 @pytest.mark.parametrize("max_points", [5, None])
@@ -200,7 +197,6 @@ def test_radius_search(
 
 
 
-@pytest.mark.skipif(sys.platform == "darwin", reason="Torch doesn't support pinned memory on mac")
 def test_radius_search_torch_compile_no_graph_break(device):
     # Cuda curnently disabled in this test, but it does work.
 
@@ -239,7 +235,6 @@ def search_fn(points, queries):
         assert torch.allclose(eager, compiled, atol=1e-6)
 
 
-@pytest.mark.skipif(sys.platform == "darwin", reason="Torch doesn't support pinned memory on mac")
 def test_opcheck(device):
 
     if device == "cpu":
@@ -255,7 +250,6 @@ def test_opcheck(device):
     )
 
 
-@pytest.mark.skipif(sys.platform == "darwin", reason="Torch doesn't support pinned memory on mac")
 @pytest.mark.parametrize("max_points", [22, None])
 def test_radius_search_comparison(device, max_points):
     torch.manual_seed(42)
@@ -305,7 +299,6 @@ def test_radius_search_comparison(device, max_points):
         assert torch.allclose(distance_warp.sum(), distance_torch.sum())
 
 
-@pytest.mark.skipif(sys.platform == "darwin", reason="Torch doesn't support pinned memory on mac")
 @pytest.mark.parametrize("max_points", [8, None])
 def test_radius_search_gradients(device, max_points):
     # Gradients are only supported to flow through the output points.
@@ -355,7 +348,6 @@ def test_radius_search_gradients(device, max_points):
     # assert torch.allclose(qrs_grad_warp, qrs_grad_torch, atol=1e-5), "Query gradients do not match"
 
 
-@pytest.mark.skipif(sys.platform == "darwin", reason="Torch doesn't support pinned memory on mac")
 @pytest.mark.parametrize("precision", [torch.bfloat16, torch.float16, torch.float32])
 @pytest.mark.parametrize("max_points", [8, None])
 def test_radius_search_reduced_precision(device, precision, max_points):