deepmd-python-inference

# DeePMD-kit Python Inference Load a trained DeePMD-kit model in Python and predict energy, forces, and virial for atomic configurations. Also covers CLI-based testing with `dp test`. ## Quick Start ```python from deepmd.infer import DeepPot import numpy as np dp = DeepPot("model.pth") coord = np.array([[1, 0, 0], [0, 0, 1.5], [1, 0, 3]]).reshape([1, -1]) cell = np.diag(10 * np.ones(3)).reshape([1, -1]) atype = [1, 0, 1] e, f, v = dp.eval(coord, cell, atype) ``` ## Agent Responsibilities 1. Determine the model source: - Frozen model file (`.pth` for conventional PyTorch, `.pb` for TensorFlow, or `.pt2` for DPA4/SeZM) - Built-in pretrained model name (e.g., `DPA-3.2-5M`) - PyTorch checkpoint (`.pt`), whose stored model configuration must be inspected before choosing an inference or export path 1. Read `references/model-artifacts.md` for `.pt`/`.pt2` models or whenever the artifact route is unclear. 1. Read `references/held-out-evaluation.md` for complete labeled evaluation used for checkpoint selection or production admission. 1. Determine the inference task: - Single-frame prediction (energy, force, virial) - Batch prediction over multiple frames - Descriptor evaluation - Model deviation calculation - CLI-based testing against labeled data 1. Before using descriptor or embedding hooks, confirm that the selected artifact supports them; a loadable `.pt2` does not necessarily contain the serialized model required by those hooks. 1. Help the user prepare input arrays in the correct format. 1. Run inference and report results. ## Python API: DeepPot ### Load a Model ```python from deepmd.infer import DeepPot # From a frozen PyTorch model dp = DeepPot("model.pth") # From a frozen TensorFlow model dp = DeepPot("graph.pb") # From a frozen DPA4/SeZM model dp = DeepPot("model.pt2") # From a built-in pretrained model (auto-downloads if not cached) dp = DeepPot("DPA-3.2-5M") ``` Built-in pretrained model names include `DPA-3.3-1M`, `DPA-3.2-5M`, `DPA-3.1-3M`, `DPA3-Omol-Large`, etc. DeePMD-kit will automatically download and cache the model on first use. ### Predict Energy, Forces, and Virial ```python import numpy as np from deepmd.infer import DeepPot dp = DeepPot("model.pth") # Prepare inputs # coord: (nframes, natoms * 3) in Angstrom # cell: (nframes, 9) cell vectors in Angstrom, row-major # atype: list of atom type indices (length natoms) coord = np.array( [ [ 0.0, 0.0, 0.0, # atom 0 (O) 0.0, 0.0, 1.0, # atom 1 (H) 0.0, 1.0, 0.0, ] # atom 2 (H) ] ).reshape([1, -1]) cell = np.diag([10.0, 10.0, 10.0]).reshape([1, -1]) # atype indices correspond to type_map order in the model # e.g., if type_map = ["O", "H"], then O=0, H=1 atype = [0, 1, 1] e, f, v = dp.eval(coord, cell, atype) print(f"Energy (eV): {e}") # shape: (nframes, 1) print(f"Forces (eV/A): {f}") # shape: (nframes, natoms, 3) print(f"Virial (eV): {v}") # shape: (nframes, 9) ``` ### Non-periodic Systems For non-periodic (isolated) systems, pass `cell=None`: ```python e, f, v = dp.eval(coord, None, atype) ``` ### Batch Prediction Process multiple frames at once: ```python nframes = 10 natoms = 3 coords = np.random.rand(nframes, natoms * 3) cells = np.tile(np.diag([10.0, 10.0, 10.0]).reshape([1, -1]), (nframes, 1)) atype = [0, 1, 1] e, f, v = dp.eval(coords, cells, atype) # e: (nframes, 1) # f: (nframes, natoms, 3) # v: (nframes, 9) ``` ### Evaluate Descriptors Extract the descriptor (atomic environment representation) from the model: ```python descriptors = dp.eval_descriptor(coord, cell, atype) # shape: (nframes, natoms, ndesc) ``` This can also be done via CLI: ```bash dp eval-desc -m model.pth -s /path/to/system -o desc_output ``` ### Calculate Model Deviation Compare predictions from multiple models to estimate uncertainty: ```python from deepmd.infer import calc_model_devi, DeepPot coord = np.array([[1, 0, 0], [0, 0, 1.5], [1, 0, 3]]).reshape([1, -1]) cell = np.diag(10 * np.ones(3)).reshape([1, -1]) atype = [1, 0, 1] graphs = [DeepPot("model_0.pth"), DeepPot("model_1.pth")] model_devi = calc_model_devi(coord, cell, atype, graphs) ``` Important: avoid loading the same model multiple times in a loop, as this can cause memory leaks. ## CLI Testing: dp test Test a frozen model against labeled data: ```bash # Basic test dp --pt test -m model.pth -s /path/to/test_system -n 30 # Test with detailed output dp --pt test -m model.pth -s /path/to/test_system -n 30 -d test_detail ``` ### dp test Options | Option | Description | | ---------------- | ---------------------------------- | | `-m MODEL` | Path to the frozen model file | | `-s SYSTEM` | Path to the test data system | | `-n NUMB` | Number of test frames | | `-d DETAIL` | Output prefix for detailed results | | `--shuffle-test` | Shuffle test frames | ### Output `dp test` prints RMSE values for energy, force, and virial: ``` Energy RMSE : 1.234e-03 eV Energy RMSE/Natoms : 6.427e-06 eV Force RMSE : 2.345e-02 eV/A Virial RMSE : 5.678e-02 eV Virial RMSE/Natoms : 2.957e-04 eV ``` With `-d test_detail`, per-frame predictions are saved to files for further analysis. The 30-frame commands above are bounded examples, not complete held-out evaluation. Use `references/held-out-evaluation.md` when the result admits a model for production. ## Complete Example: Train, Freeze, and Inference ```python import subprocess import numpy as np from deepmd.infer import DeepPot # Step 1: Train (run in shell) # dp --pt train input.json # Step 2: Freeze (run in shell) # dp --pt freeze -o model.pth # Step 3: Python inference dp = DeepPot("model.pth") # Load test data from deepmd format coord = np.load("test_system/set.000/coord.npy") # (nframes, natoms*3) cell = np.load("test_system/set.000/box.npy") # (nframes, 9) atype_raw = np.loadtxt("test_system/type.raw", dtype=int).tolist() # Predict e, f, v = dp.eval(coord, cell, atype_raw) # Compare with reference ref_energy = np.load("test_system/set.000/energy.npy") ref_force = np.load("test_system/set.000/force.npy") natoms = len(atype_raw) energy_rmse = np.sqrt(np.mean((e.flatten() - ref_energy.flatten()) ** 2)) / natoms force_rmse = np.sqrt(np.mean((f.reshape(-1) - ref_force.reshape(-1)) ** 2)) print(f"Energy RMSE/atom: {energy_rmse:.6f} eV") print(f"Force RMSE: {force_rmse:.6f} eV/A") ``` ## Using Pretrained Models Directly Built-in pretrained models can be used without any training: ```python from deepmd.infer import DeepPot import numpy as np # Auto-downloads DPA-3.2-5M on first use dp = DeepPot("DPA-3.2-5M") # Water molecule example coord = np.array( [ [0.000, 0.000, 0.117], # O [0.000, 0.757, -0.469], # H [0.000, -0.757, -0.469], # H ] ).reshape([1, -1]) cell = np.diag([10.0, 10.0, 10.0]).reshape([1, -1]) atype = [0, 1, 1] # Check model's type_map for correct indices e, f, v = dp.eval(coord, cell, atype) print(f"Energy: {e[0][0]:.6f} eV") print(f"Forces:\n{f[0]}") ``` To download pretrained models explicitly: ```bash dp pretrained download DPA-3.3-1M dp pretrained download DPA-3.2-5M dp pretrained download DPA-3.1-3M dp pretrained download DPA-3.2-5M --cache-dir ./models ``` ## Input Array Format Reference | Array | Shape | Unit | Description | | ------- | -------------------- | -------- | --------------------------------------------- | | `coord` | (nframes, natoms\*3) | Angstrom | Atomic coordinates, flattened | | `cell` | (nframes, 9) | Angstrom | Cell vectors, row-major (a1x,a1y,a1z,a2x,...) | | `atype` | (natoms,) | - | Atom type indices matching model's type_map | | Output | Shape | Unit | Description | | ------ | -------------------- | ---- | ----------------------- | | `e` | (nframes, 1) | eV | Total energy per frame | | `f` | (nframes, natoms, 3) | eV/A | Forces on each atom | | `v` | (nframes, 9) | eV | Virial tensor per frame | ## Agent Checklist - [ ] Model file exists and is accessible (`.pb`, `.pth`, `.pt`, `.pt2`, or valid pretrained name) - [ ] An ambiguous `.pt` checkpoint was classified from its stored configuration, not its filename - [ ] The requested descriptor or embedding operation is supported by the specific artifact, not inferred from its suffix - [ ] `coord` array is shaped (nframes, natoms\*3) and in Angstrom - [ ] `cell` array is shaped (nframes, 9) or `None` for non-periodic systems - [ ] `atype` indices match the model's `type_map` ordering - [ ] For model deviation, multiple models are loaded only once (not in a loop) - [ ] Results are reported with correct units (eV, eV/A) ## References - [Python inference documentation](https://docs.deepmodeling.com/projects/deepmd/en/latest/inference/python.html) - [dp test documentation](https://docs.deepmodeling.com/projects/deepmd/en/latest/test/test.html) - [Pretrained model download](https://docs.deepmodeling.com/projects/deepmd/en/latest/model/pretrained.html) - [DeepPot API reference](https://docs.deepmodeling.com/projects/deepmd/en/latest/api_py/deepmd.infer.html) - [DeePMD-kit GitHub](https://github.com/deepmodeling/deepmd-kit)