feat: add pip installation support with optional torch-scatter dependency

.github/workflows/test.yml

@@ -25,7 +25,7 @@ jobs:
 
     - name: Install dependencies
       run: |
-        uv sync --group dev
+        uv sync --group dev --extra scatter
 
     - name: Run tests
       run: |

README.md

@@ -39,13 +39,28 @@ uv sync
 uv sync --find-links https://data.pyg.org/whl/torch-2.5.0+cu121.html
 ```
 
-### Using pip (from PyPI - coming soon)
+### Using pip
 
 ```bash
+# Basic install (all core features except grid integration / LDOS)
 pip install dftio
+
+# Install with grid integration support (LDOS calculations)
+pip install "dftio[scatter]" -f https://data.pyg.org/whl/torch-2.5.0+cpu.html
+
+# Install with all optional dependencies (includes scatter + dev tools)
+pip install "dftio[full]" -f https://data.pyg.org/whl/torch-2.5.0+cpu.html
+
+# For GPU users, replace 'cpu' with your CUDA version (e.g., cu121)
+pip install "dftio[scatter]" -f https://data.pyg.org/whl/torch-2.5.0+cu121.html
+
+# Or install from requirements files
+pip install -r requirements.txt          # core only
+pip install -r requirements-full.txt     # core + scatter + dev
+pip install -r requirements-dev.txt      # core + dev tools
 ```
 
-**Note**: dftio depends on `torch-scatter` which requires special handling. The install script automatically manages this for you.
+**Note**: The `scatter` extra installs `torch-scatter`, which is only needed for LDOS (Local Density of States) calculations via `dftio.calc.ldos`. All other dftio functionality (parsing, data structures, CLI) works without it.
 
 ## Supports
 

dftio/op/grid_int.py

@@ -1,7 +1,6 @@
 from ..datastruct import PrimitiveFieldsNeighborList
 import torch
 import ase.data as data
-from torch_scatter import scatter_sum
 import numpy as np
 
 atomic_numbers_r = dict(zip(data.atomic_numbers.values(), data.atomic_numbers.keys()))
@@ -32,7 +31,18 @@ def __init__(self, atomic_numbers, pbc, cell, coordinates, grids, atomic_basis,
         self.cell_shift = torch.from_numpy(np.concatenate(self.cell_shift, axis=0, dtype=np.int32))
 
     def integrate(self, weights=None):
-
+        try:
+            from torch_scatter import scatter_sum
+        except ImportError as err:
+            raise ImportError(
+                "torch-scatter is required for grid integration (LDOS calculations). "
+                "Install dftio with the scatter extra:\n"
+                "  pip install \"dftio[scatter]\" -f https://data.pyg.org/whl/torch-<version>+<variant>.html\n"
+                "Or install all optional dependencies:\n"
+                "  pip install \"dftio[full]\" -f https://data.pyg.org/whl/torch-<version>+<variant>.html\n"
+                "For detailed instructions, see: https://deepmodeling.github.io/dftio/installation.html"
+            ) from err
+
         ngrid = len(self.grids)
         dtype = weights.dtype if weights is not None else self.dtype
         results = torch.zeros(ngrid, dtype=dtype)

docs/installation.md

@@ -38,11 +38,27 @@ If you prefer to manage the installation yourself, you can use `uv`.
     ```
     Including the `--group dev` flag will also install the packages required for testing and building documentation.
 
-## Using pip (from PyPI)
+## Using pip
 
-*Coming soon. Once `dftio` is published to the Python Package Index (PyPI), you will be able to install it directly with `pip`.*
+You can install `dftio` directly with pip:
 
 ```bash
-# This will be enabled in a future release
-# pip install dftio
-```
+# Basic install (all core features except grid integration / LDOS)
+pip install dftio
+
+# Install with grid integration support for LDOS calculations
+pip install "dftio[scatter]" -f https://data.pyg.org/whl/torch-2.5.0+cpu.html
+
+# Install with all optional dependencies (scatter + dev tools)
+pip install "dftio[full]" -f https://data.pyg.org/whl/torch-2.5.0+cpu.html
+
+# For GPU users, replace 'cpu' with your CUDA version (e.g., cu121, cu124)
+pip install "dftio[scatter]" -f https://data.pyg.org/whl/torch-2.5.0+cu121.html
+
+# Or install from requirements files
+pip install -r requirements.txt          # core only
+pip install -r requirements-full.txt     # core + scatter + dev
+pip install -r requirements-dev.txt      # core + dev tools
+```
+
+> **Note:** The `scatter` extra installs `torch-scatter`, which is only needed for LDOS (Local Density of States) calculations via `dftio.calc.ldos`. All other dftio functionality works without it.

install.sh

@@ -45,9 +45,9 @@ if ! command -v uv &> /dev/null; then
 fi
 
 # Sync dependencies with the specified find-links
-echo "Installing dftio with torch_scatter ($VARIANT version)..."
+echo "Installing dftio with torch_scatter ($VARIANT version) [scatter extra]..."
 echo "Using Python: $(which python)"
-uv sync --python $(which python) --find-links "$FIND_LINKS_URL"
+uv sync --python "$(which python)" --extra scatter --find-links "$FIND_LINKS_URL"
 
 echo ""
 echo "✅ Installation complete!"

pyproject.toml

@@ -22,7 +22,6 @@ dependencies = [
     "lmdb==1.4.1",
     "sisl>=0.14.3",
     "dpdata>=0.2.20",
-    "torch-scatter==2.1.2",
     "tqdm",
 ]
 
@@ -35,6 +34,17 @@ dev = [
     "sphinx-autodoc-typehints",
 ]
 
+[project.optional-dependencies]
+scatter = ["torch-scatter==2.1.2"]
+full = ["dftio[scatter]", "dftio[dev]"]
+dev = [
+    "pytest>=7.2.0",
+    "pytest-order==1.2.0",
+    "pytest-cov",
+    "jupyter-book>=0.15,<1.0",
+    "sphinx-autodoc-typehints",
+]
+
 [project.scripts]
 dftio = "dftio.__main__:main"
 

requirements-dev.txt

@@ -0,0 +1,6 @@
+-r requirements.txt
+pytest>=7.2.0
+pytest-order==1.2.0
+pytest-cov
+jupyter-book>=0.15,<1.0
+sphinx-autodoc-typehints

requirements-full.txt

@@ -0,0 +1,4 @@
+--find-links https://data.pyg.org/whl/torch-2.5.0+cpu.html
+-r requirements.txt
+-r requirements-dev.txt
+torch-scatter==2.1.2

requirements.txt

@@ -0,0 +1,14 @@
+numpy
+scipy>=1.11,<1.12
+matplotlib
+torch>=2.0.0,<2.5.1
+ase
+pyyaml
+future
+dargs==0.4.4
+e3nn>=0.5.1
+h5py>=3.7.0,<3.11.0,!=3.10.0
+lmdb==1.4.1
+sisl>=0.14.3
+dpdata>=0.2.20
+tqdm

test/test_grid_int.py

@@ -5,6 +5,14 @@
 from dftio.datastruct import AtomicBasis
 import ase.data as data
 
+try:
+    import torch_scatter  # noqa: F401
+    _SCATTER_AVAILABLE = True
+except ImportError:
+    _SCATTER_AVAILABLE = False
+
+needs_scatter = pytest.mark.skipif(not _SCATTER_AVAILABLE, reason="torch-scatter not installed")
+
 class MockAtomicBasis:
     def __init__(self, atomic_numbers):
         self.atomic_numbers = atomic_numbers
@@ -44,6 +52,7 @@ def test_single_grid_integrator_init(mock_atomic_basis):
     assert sgi.coordinates.shape == (1, 3)
     assert sgi.grids.shape == (1, 3)
 
+@needs_scatter
 def test_integrate(mock_atomic_basis):
     """Test integrate method."""
     atomic_numbers = [1]
@@ -69,3 +78,41 @@ def test_integrate(mock_atomic_basis):
     weights = torch.tensor([1.0])
     result_w = sgi.integrate(weights=weights)
     assert result_w.shape == (1,)
+
+
+def test_integrate_without_scatter_raises_error(mock_atomic_basis, monkeypatch):
+    """Test that integrate() raises helpful ImportError when torch-scatter is not installed."""
+    import builtins
+    import sys
+
+    # Remove torch_scatter from sys.modules so cached import is bypassed
+    sys.modules.pop("torch_scatter", None)
+    sys.modules.pop("torch_scatter.scatter_sum", None)
+
+    # Store the original import function
+    _original_import = builtins.__import__
+
+    def _mock_import(name, globals=None, locals=None, fromlist=(), level=0):
+        if name == "torch_scatter":
+            raise ModuleNotFoundError("No module named 'torch_scatter'")
+        return _original_import(name, globals, locals, fromlist, level)
+
+    monkeypatch.setattr(builtins, "__import__", _mock_import)
+
+    atomic_numbers = [1]
+    pbc = [True, True, True]
+    cell = np.eye(3)
+    coordinates = np.array([[0.0, 0.0, 0.0]])
+    grids = np.array([[0.1, 0.1, 0.1]])
+
+    sgi = SingleGridIntegrator(
+        atomic_numbers=atomic_numbers,
+        pbc=pbc,
+        cell=cell,
+        coordinates=coordinates,
+        grids=grids,
+        atomic_basis={"H": mock_atomic_basis},
+    )
+
+    with pytest.raises(ImportError, match="torch-scatter is required"):
+        sgi.integrate()

test/test_ldos.py

@@ -4,6 +4,14 @@
 from dftio.calc.ldos import LDOS
 from dftio.datastruct import AtomicBasis
 
+try:
+    import torch_scatter  # noqa: F401
+    _SCATTER_AVAILABLE = True
+except ImportError:
+    _SCATTER_AVAILABLE = False
+
+needs_scatter = pytest.mark.skipif(not _SCATTER_AVAILABLE, reason="torch-scatter not installed")
+
 class MockAtomicBasis:
     def __init__(self, atomic_numbers):
         self.atomic_numbers = atomic_numbers
@@ -40,6 +48,7 @@ def test_ldos_init(mock_atomic_basis):
     assert ldos.natoms == 1
     assert ldos.nspin == 2
 
+@needs_scatter
 def test_ldos_get(mock_atomic_basis):
     """Test LDOS get method."""
     atomic_numbers = [1]