Metadata-Version: 2.4
Name: deepmd-torchsim
Version: 0.1.1
Summary: DeePMD-kit (PyTorch backend) ModelInterface implementation for torch-sim
Author-email: rverma7@ncsu.edu
License-Expression: MIT
Requires-Python: >=3.12
Description-Content-Type: text/markdown
License-File: LICENSE
Requires-Dist: torch-sim-atomistic>=0.6.1
Requires-Dist: torch>=2
Requires-Dist: numpy
Requires-Dist: ase>=3.26
Provides-Extra: deepmd
Requires-Dist: deepmd-kit==3.1.3; extra == "deepmd"
Requires-Dist: torch==2.10.0; extra == "deepmd"
Requires-Dist: mpich; extra == "deepmd"
Requires-Dist: e3nn; extra == "deepmd"
Provides-Extra: test
Requires-Dist: pytest>=8; extra == "test"
Dynamic: license-file

# deepmd-torchsim

A [torch-sim](https://github.com/TorchSim/torch-sim) `ModelInterface` implementation
for [DeePMD-kit](https://github.com/deepmodeling/deepmd-kit)'s PyTorch backend.

## Install

The package (`DeepmdModel`, torch-sim `ModelInterface` wrapper) can be installed with either `pip` and `uv`:

```bash
# from PyPI (once published) or a local checkout
pip install deepmd-torchsim
uv pip install deepmd-torchsim
uv add deepmd-torchsim

# editable, from a local checkout 
pip install -e .
uv pip install -e .
```

### Getting a working `deepmd-kit` backend

```bash
pip install "deepmd-torchsim[deepmd]"
uv pip install "deepmd-torchsim[deepmd]"
```

The `deepmd` extra pins an working `deepmd-kit==3.1.3`, with `torch==2.10.0`.

## Usage

```python
import torch
from deepmd_torchsim import DeepmdModel
import torch_sim as ts
from ase.build import molecule

model = DeepmdModel(
    model_path="frozen_model.pth",
    device=torch.device("cuda" if torch.cuda.is_available() else "cpu"),
    compute_forces=True,
    compute_stress=True,
)

state = ts.io.atoms_to_state([molecule("H2O")], model.device, model.dtype)
results = model(state)
print(results["energy"])  # [n_systems]
print(results["forces"])  # [n_atoms, 3]
print(results["stress"])  # [n_systems, 3, 3]
```

For multitask/multi-domain foundation checkpoints (e.g. DPA-3), pass `head=`
to select which trained domain to evaluate with:

```python
model = DeepmdModel(model_path="DPA-3.1-3M.pt", head="Omat24")
```
## Tests

```bash
pytest tests/
```

## License

This project is licensed under the [MIT License](https://github.com/rahulumrao/deepmd_torchsim/blob/main/LICENSE).

## Author
Rahul Verma \
Email: rverma7@ncsu.edu

