Merge branch 'devel' into devel_thang

.github/workflows/benchmark.yml

@@ -0,0 +1,23 @@
+name: Python package
+
+on:
+  - push
+  - pull_request
+
+jobs:
+  build:
+    runs-on: ubuntu-latest
+    steps:
+    - uses: actions/checkout@v4
+    - name: Set up Python
+      uses: actions/setup-python@v5
+      with:
+        python-version: 3.12
+    - run: curl -LsSf https://astral.sh/uv/install.sh | sh
+    - name: Install dependencies
+      run: uv pip install --system .[amber,ase,pymatgen,benchmark] rdkit openbabel-wheel
+    - name: Run benchmarks
+      uses: CodSpeedHQ/action@v2
+      with:
+        token: ${{ secrets.CODSPEED_TOKEN }}
+        run: pytest benchmark/ --codspeed

.github/workflows/pyright.yml

@@ -0,0 +1,19 @@
+on:
+  - push
+  - pull_request
+
+name: Type checker
+jobs:
+  pyright:
+    name: pyright
+    runs-on: ubuntu-latest
+    steps:
+    - uses: actions/checkout@master
+    - uses: actions/setup-python@v5
+      with:
+        python-version: '3.12'
+    - run: pip install uv
+    - run: uv pip install --system -e .[amber,ase,pymatgen] rdkit openbabel-wheel
+    - uses: jakebailey/pyright-action@v2
+      with:
+        version: 1.1.363

.pre-commit-config.yaml

@@ -19,7 +19,7 @@ repos:
 # Python
 -   repo: https://github.com/astral-sh/ruff-pre-commit
     # Ruff version.
-    rev: v0.4.3
+    rev: v0.4.4
     hooks:
     - id: ruff
       args: ["--fix"]

benchmark/test_import.py

@@ -0,0 +1,20 @@
+from __future__ import annotations
+
+import subprocess
+import sys
+
+import pytest
+
+
+@pytest.mark.benchmark
+def test_import():
+    """Test import dpdata."""
+    subprocess.check_output(
+        [sys.executable, "-c", "'from dpdata import LabeledSystem'"]
+    ).decode("ascii")
+
+
+@pytest.mark.benchmark
+def test_cli():
+    """Test dpdata command."""
+    subprocess.check_output([sys.executable, "-m", "dpdata", "-h"]).decode("ascii")

docs/conf.py

@@ -11,6 +11,8 @@
 # add these directories to sys.path here. If the directory is relative to the
 # documentation root, use os.path.abspath to make it absolute, like shown here.
 #
+from __future__ import annotations
+
 import os
 import subprocess as sp
 import sys
@@ -207,3 +209,4 @@ def setup(app):
 # jupyterlite
 jupyterlite_contents = "./nb"
 jupyterlite_bind_ipynb_suffix = False
+jupyterlite_silence = False

docs/environment.yml

@@ -1,8 +1,9 @@
 name: dpdata
 channels:
   - https://repo.mamba.pm/emscripten-forge
-  - https://repo.mamba.pm/conda-forge
+  - conda-forge
 dependencies:
+  - xeus-python
   - numpy
   - scipy
   - monty

docs/make_format.py

@@ -1,8 +1,15 @@
+from __future__ import annotations
+
 import csv
 import os
+import sys
 from collections import defaultdict
 from inspect import Parameter, Signature, cleandoc, signature
-from typing import Literal
+
+if sys.version_info >= (3, 8):
+    from typing import Literal
+else:
+    from typing_extensions import Literal
 
 from numpydoc.docscrape import Parameter as numpydoc_Parameter
 from numpydoc.docscrape_sphinx import SphinxDocString

docs/nb/try_dpdata.ipynb

@@ -13,6 +13,8 @@
    "metadata": {},
    "outputs": [],
    "source": [
+    "from __future__ import annotations\n",
+    "\n",
     "import dpdata"
    ]
   },

docs/rtd_environment.yml

@@ -1,6 +1,7 @@
 name: dpdata
 channels:
-  - https://repo.mamba.pm/conda-forge
+  - conda-forge
 dependencies:
+  - mamba
   - pip:
     - ..[docs]

dpdata/__about__.py

@@ -1 +1,3 @@
+from __future__ import annotations
+
 __version__ = "unknown"

dpdata/__init__.py

@@ -1,36 +1,14 @@
-# monty needs lzma
-# See https://github.com/pandas-dev/pandas/pull/27882
-try:
-    import lzma  # noqa: F401
-except ImportError:
-
-    class fakemodule:
-        pass
-
-    import sys
-
-    sys.modules["lzma"] = fakemodule
+from __future__ import annotations
 
 from . import lammps, md, vasp
+from .bond_order_system import BondOrderSystem
 from .system import LabeledSystem, MultiSystems, System
 
 try:
     from ._version import version as __version__
 except ImportError:
     from .__about__ import __version__
 
-# BondOrder System has dependency on rdkit
-try:
-    # prevent conflict with dpdata.rdkit
-    import rdkit as _  # noqa: F401
-
-    USE_RDKIT = True
-except ModuleNotFoundError:
-    USE_RDKIT = False
-
-if USE_RDKIT:
-    from .bond_order_system import BondOrderSystem
-
 __all__ = [
     "__version__",
     "lammps",

dpdata/__main__.py

@@ -1,3 +1,5 @@
+from __future__ import annotations
+
 from dpdata.cli import dpdata_cli
 
 if __name__ == "__main__":

dpdata/abacus/md.py

@@ -1,3 +1,5 @@
+from __future__ import annotations
+
 import os
 import warnings
 

dpdata/abacus/relax.py

@@ -1,3 +1,5 @@
+from __future__ import annotations
+
 import os
 
 import numpy as np

dpdata/abacus/scf.py

@@ -1,3 +1,5 @@
+from __future__ import annotations
+
 import os
 import re
 import warnings

dpdata/amber/mask.py

@@ -1,5 +1,7 @@
 """Amber mask."""
 
+from __future__ import annotations
+
 try:
     import parmed
 except ImportError:

dpdata/amber/md.py

@@ -1,8 +1,9 @@
+from __future__ import annotations
+
 import os
 import re
 
 import numpy as np
-from scipy.io import netcdf_file
 
 from dpdata.amber.mask import pick_by_amber_mask
 from dpdata.unit import EnergyConversion
@@ -44,6 +45,8 @@ def read_amber_traj(
     labeled : bool
         Whether to return labeled data
     """
+    from scipy.io import netcdf_file
+
     flag_atom_type = False
     flag_atom_numb = False
     amber_types = []

dpdata/amber/sqm.py

@@ -1,3 +1,5 @@
+from __future__ import annotations
+
 import numpy as np
 
 from dpdata.periodic_table import ELEMENTS

dpdata/ase_calculator.py

@@ -1,6 +1,8 @@
-from typing import TYPE_CHECKING, List, Optional
+from __future__ import annotations
 
-from ase.calculators.calculator import (
+from typing import TYPE_CHECKING
+
+from ase.calculators.calculator import (  # noqa: TID253
     Calculator,
     PropertyNotImplementedError,
     all_changes,
@@ -23,7 +25,10 @@ class DPDataCalculator(Calculator):
         dpdata driver
     """
 
-    name = "dpdata"
+    @property
+    def name(self) -> str:
+        return "dpdata"
+
     implemented_properties = ["energy", "free_energy", "forces", "virial", "stress"]
 
     def __init__(self, driver: Driver, **kwargs) -> None:
@@ -32,9 +37,9 @@ def __init__(self, driver: Driver, **kwargs) -> None:
 
     def calculate(
         self,
-        atoms: Optional["Atoms"] = None,
-        properties: List[str] = ["energy", "forces"],
-        system_changes: List[str] = all_changes,
+        atoms: Atoms | None = None,
+        properties: list[str] = ["energy", "forces"],
+        system_changes: list[str] = all_changes,
     ):
         """Run calculation with a driver.
 
@@ -48,10 +53,10 @@ def calculate(
         system_changes : List[str], optional
             unused, only for function signature compatibility, by default all_changes
         """
-        if atoms is not None:
-            self.atoms = atoms.copy()
+        assert atoms is not None
+        atoms = atoms.copy()
 
-        system = dpdata.System(self.atoms, fmt="ase/structure")
+        system = dpdata.System(atoms, fmt="ase/structure")
         data = system.predict(driver=self.driver).data
 
         self.results["energy"] = data["energies"][0]

dpdata/bond_order_system.py

@@ -1,9 +1,10 @@
 # %%
 # Bond Order System
+from __future__ import annotations
+
 from copy import deepcopy
 
 import numpy as np
-from rdkit.Chem import Conformer
 
 import dpdata.rdkit.utils
 from dpdata.rdkit.sanitize import Sanitizer
@@ -97,11 +98,14 @@ def from_fmt_obj(self, fmtobj, file_name, **kwargs):
         mol = fmtobj.from_bond_order_system(file_name, **kwargs)
         self.from_rdkit_mol(mol)
         if hasattr(fmtobj.from_bond_order_system, "post_func"):
-            for post_f in fmtobj.from_bond_order_system.post_func:
+            for post_f in fmtobj.from_bond_order_system.post_func:  # type: ignore
                 self.post_funcs.get_plugin(post_f)(self)
         return self
 
     def to_fmt_obj(self, fmtobj, *args, **kwargs):
+        from rdkit.Chem import Conformer
+
+        assert self.rdkit_mol is not None
         self.rdkit_mol.RemoveAllConformers()
         for ii in range(self.get_nframes()):
             conf = Conformer()
@@ -144,9 +148,9 @@ def get_formal_charges(self):
         """Return the formal charges on each atom."""
         return self.data["formal_charges"]
 
-    def copy(self):
+    def copy(self):  # type: ignore
         new_mol = deepcopy(self.rdkit_mol)
-        self.__class__(data=deepcopy(self.data), rdkit_mol=new_mol)
+        return self.__class__(data=deepcopy(self.data), rdkit_mol=new_mol)
 
     def __add__(self, other):
         raise NotImplementedError(

dpdata/cli.py

@@ -1,7 +1,8 @@
 """Command line interface for dpdata."""
 
+from __future__ import annotations
+
 import argparse
-from typing import Optional
 
 from . import __version__
 from .system import LabeledSystem, MultiSystems, System
@@ -59,11 +60,11 @@ def convert(
     *,
     from_file: str,
     from_format: str = "auto",
-    to_file: Optional[str] = None,
-    to_format: Optional[str] = None,
+    to_file: str | None = None,
+    to_format: str | None = None,
     no_labeled: bool = False,
     multi: bool = False,
-    type_map: Optional[list] = None,
+    type_map: list | None = None,
     **kwargs,
 ):
     """Convert files from one format to another one.

dpdata/cp2k/cell.py

@@ -1,4 +1,5 @@
 # %%
+from __future__ import annotations
 
 import numpy as np
 

dpdata/cp2k/output.py

@@ -1,4 +1,6 @@
 # %%
+from __future__ import annotations
+
 import math
 import re
 from collections import OrderedDict