atst-tools 2.0.0__py3-none-any.whl

atst_tools/__init__.py

@@ -0,0 +1,51 @@
+"""ATST-Tools package metadata helpers."""
+
+from importlib.metadata import PackageNotFoundError, version as _metadata_version
+from pathlib import Path
+
+
+def _source_tree_version() -> str | None:
+    """Return the project version from a local source-tree ``pyproject.toml``."""
+    for parent in Path(__file__).resolve().parents:
+        pyproject = parent / "pyproject.toml"
+        if not pyproject.exists():
+            continue
+
+        in_project = False
+        for raw_line in pyproject.read_text(encoding="utf-8").splitlines():
+            line = raw_line.strip()
+            if line == "[project]":
+                in_project = True
+                continue
+            if in_project and line.startswith("["):
+                return None
+            if in_project and line.startswith("version"):
+                _, value = line.split("=", 1)
+                return value.strip().strip('"').strip("'")
+    return None
+
+
+def package_version() -> str:
+    """Return the ATST-Tools package version.
+
+    The package version is governed by ``pyproject.toml``. Source-tree runs
+    read that file directly, while installed-package runs use distribution
+    metadata generated from the same project version.
+
+    Returns:
+        Project package version, or ``"unknown"`` when neither source-tree nor
+        installed-package metadata is available.
+    """
+    return _source_tree_version() or _installed_version()
+
+
+def _installed_version() -> str:
+    try:
+        return _metadata_version("atst-tools")
+    except PackageNotFoundError:
+        return "unknown"
+
+
+__version__ = package_version()
+
+__all__ = ["__version__", "package_version"]

atst_tools/calculators/abacuslite_backend.py

@@ -0,0 +1,81 @@
+"""Resolve the ABACUS ASE backend used by ATST-Tools."""
+
+from __future__ import annotations
+
+import re
+import shlex
+from typing import Literal
+
+from ase.calculators.genericfileio import read_stdout
+
+
+BackendSource = Literal["external", "vendored"]
+
+
+def _load_abacuslite_backend():
+    try:
+        from abacuslite import Abacus, AbacusProfile
+
+        return Abacus, AbacusProfile, "external"
+    except ImportError:
+        from atst_tools.external.ASE_interface.abacuslite import Abacus, AbacusProfile
+
+        return Abacus, AbacusProfile, "vendored"
+
+
+Abacus, AbacusProfile, BACKEND_SOURCE = _load_abacuslite_backend()
+
+
+def _default_version_command(command: str) -> list[str]:
+    """Return the default bare ABACUS version-probe command."""
+    parts = shlex.split(command)
+    if not parts:
+        return ["abacus", "--version"]
+    executable = parts[-1] if parts[0] in {"mpirun", "mpiexec", "srun"} else parts[0]
+    return [executable, "--version"]
+
+
+class ATSTAbacusProfile(AbacusProfile):
+    """ABACUS profile with ATST-managed version probing.
+
+    The run command may be MPI-wrapped for calculations, but version probing is
+    a lightweight environment check and defaults to a bare executable call.
+    """
+
+    def __init__(self, *args, version_command: str | None = None, **kwargs):
+        """Initialize an ATST ABACUS profile.
+
+        Args:
+            *args: Positional arguments passed to abacuslite's profile.
+            version_command: Optional full command used for version probing.
+            **kwargs: Keyword arguments passed to abacuslite's profile.
+        """
+        super().__init__(*args, **kwargs)
+        self.version_command = version_command
+
+    @staticmethod
+    def parse_version(stdout: str) -> str:
+        """Parse ABACUS version output from legacy and LTS banner formats."""
+        for pattern in (r"ABACUS version (\S+)", r"ABACUS\s+(v\S+)"):
+            match = re.search(pattern, stdout)
+            if match is not None:
+                return match.group(1)
+        raise RuntimeError(f"Could not parse ABACUS version from output: {stdout!r}")
+
+    def version(self) -> str:
+        """Return the ABACUS version using ATST's version-probe command."""
+        command = (
+            shlex.split(self.version_command)
+            if self.version_command
+            else _default_version_command(self.command)
+        )
+        return self.parse_version(read_stdout(command))
+
+
+__all__ = [
+    "Abacus",
+    "AbacusProfile",
+    "ATSTAbacusProfile",
+    "BACKEND_SOURCE",
+    "BackendSource",
+]

atst_tools/calculators/base.py

@@ -0,0 +1,10 @@
+from abc import ABC, abstractmethod
+from typing import Any, Dict
+
+class BaseCalculator(ABC):
+    """Abstract base class for ATST calculators."""
+
+    @abstractmethod
+    def get_calculator(self, **kwargs) -> Any:
+        """Return a configured ASE calculator instance."""
+        pass

atst_tools/calculators/dp.py

@@ -0,0 +1,112 @@
+"""DeepMD-kit ASE calculator adapter."""
+
+from __future__ import annotations
+
+import json
+import os
+from typing import Any, Dict, Hashable
+
+from ase.calculators.calculator import Calculator
+
+
+def is_dp_calculator(name: str) -> bool:
+    """Return whether a calculator name refers to the DeepMD-kit adapter."""
+    return name.lower() in {"dp", "deepmd"}
+
+
+def dp_section(config: Dict[str, Any]) -> Dict[str, Any]:
+    """Return DeepMD-kit settings from supported config layouts."""
+    calculator = config.get("calculator", {})
+    if isinstance(calculator, dict):
+        if "dp" in calculator:
+            return dict(calculator.get("dp") or {})
+        if calculator.get("name") == "deepmd":
+            return dict(calculator.get("deepmd") or {})
+    if "parameters" in config:
+        return dict(config["parameters"])
+    return {}
+
+
+def dp_share_calculator(config: Dict[str, Any], default: bool = True) -> bool:
+    """Return the configured DP calculator sharing policy."""
+    return bool(dp_section(config).get("share_calculator", default))
+
+
+def should_share_calculator(name: str, config: Dict[str, Any], parallel: bool = False) -> bool:
+    """Return whether workflow images should share a single calculator instance."""
+    return is_dp_calculator(name) and not parallel and dp_share_calculator(config)
+
+
+def _normalize_type_dict(dp_params: Dict[str, Any]) -> dict[str, int] | None:
+    type_map = dp_params.pop("type_map", None)
+    type_dict = dp_params.pop("type_dict", None)
+    if type_map is not None and type_dict is not None:
+        raise ValueError("calculator.dp.type_map and calculator.dp.type_dict are mutually exclusive")
+    if type_dict is not None:
+        return {str(symbol): int(index) for symbol, index in dict(type_dict).items()}
+    if type_map is None:
+        return None
+    return {str(symbol): index for index, symbol in enumerate(type_map)}
+
+
+def _cache_key(model: str, params: Dict[str, Any]) -> tuple[Hashable, ...]:
+    serializable = json.dumps(params, sort_keys=True, default=str)
+    return (os.path.abspath(os.path.expanduser(model)), serializable)
+
+
+class DeepPotentialFactory:
+    """Factory for creating DeepMD-kit ASE calculators with optional sharing."""
+
+    _instances: Dict[tuple[Hashable, ...], Calculator] = {}
+
+    @staticmethod
+    def get_calculator(
+        config: Dict[str, Any],
+        shared: bool | None = None,
+        **kwargs: Any,
+    ) -> Calculator:
+        """Create a DeepMD-kit ASE calculator.
+
+        Args:
+            config: ATST-Tools configuration dictionary.
+            shared: Override the configured calculator sharing policy.
+            **kwargs: Workflow-local calculator construction hints.
+
+        Returns:
+            Configured ``deepmd.calculator.DP`` instance.
+        """
+        try:
+            from deepmd.calculator import DP
+        except ImportError as exc:
+            raise ImportError(
+                "deepmd-kit is not installed. Install it to use the DP calculator."
+            ) from exc
+
+        dp_params = dp_section(config)
+        dp_params.update(kwargs)
+
+        model_file = dp_params.pop("model", None)
+        if not model_file:
+            raise ValueError("Missing required field calculator.dp.model")
+
+        omp = dp_params.pop("omp", None)
+        if omp is not None:
+            os.environ["OMP_NUM_THREADS"] = str(int(omp))
+
+        share = dp_share_calculator(config) if shared is None else bool(shared)
+        dp_params.pop("share_calculator", None)
+        dp_params.pop("directory", None)
+
+        type_dict = _normalize_type_dict(dp_params)
+        constructor_params: Dict[str, Any] = dict(dp_params)
+        if type_dict is not None:
+            constructor_params["type_dict"] = type_dict
+
+        key = _cache_key(model_file, constructor_params)
+        if share and key in DeepPotentialFactory._instances:
+            return DeepPotentialFactory._instances[key]
+
+        calc = DP(model=model_file, **constructor_params)
+        if share:
+            DeepPotentialFactory._instances[key] = calc
+        return calc

atst_tools/calculators/factory.py

@@ -0,0 +1,126 @@
+"""Calculator factories for ATST-Tools workflows."""
+
+from __future__ import annotations
+
+import os
+import shlex
+import logging
+from typing import Any, Dict
+
+from ase.calculators.calculator import Calculator
+
+from atst_tools.calculators.abacuslite_backend import Abacus, ATSTAbacusProfile, BACKEND_SOURCE
+from atst_tools.calculators.dp import DeepPotentialFactory
+
+
+_ABACUS_CONTROL_KEYS = {"command", "mpi", "omp", "directory", "parameters", "version_command"}
+LOGGER = logging.getLogger(__name__)
+_ABACUS_BACKEND_LOGGED = False
+
+
+def _abacus_section(config: Dict[str, Any]) -> Dict[str, Any]:
+    """Return ABACUS calculator settings from the supported config layouts."""
+    if "calculator" in config:
+        return dict(config.get("calculator", {}).get("abacus", {}))
+    if "abacus" in config:
+        return dict(config.get("abacus", {}))
+    return dict(config)
+
+
+def _as_mp_kpts(kpts: Any) -> Any:
+    if isinstance(kpts, list) and len(kpts) == 3:
+        return {
+            "mode": "mp-sampling",
+            "nk": kpts,
+            "gamma-centered": True,
+            "kshift": [0, 0, 0],
+        }
+    return kpts
+
+
+def _build_abacus_command(command: str, mpi: int) -> str:
+    if "{mpi}" in command:
+        return command.format(mpi=mpi)
+
+    executable = shlex.split(command)[0] if command.strip() else "abacus"
+    if mpi > 1 and executable not in {"mpirun", "mpiexec", "srun"}:
+        return f"mpirun -np {mpi} {command}"
+    return command or "abacus"
+
+
+def _resolve_directory(path: str | None) -> str | None:
+    if path is None:
+        return None
+    return os.path.abspath(os.path.expanduser(path))
+
+
+class AbacusFactory:
+    """Factory for creating ABACUS ASE calculators through abacuslite."""
+
+    @staticmethod
+    def _log_backend_source_once() -> None:
+        global _ABACUS_BACKEND_LOGGED
+        if not _ABACUS_BACKEND_LOGGED:
+            LOGGER.info("Using %s abacuslite backend for ABACUS calculator.", BACKEND_SOURCE)
+            _ABACUS_BACKEND_LOGGED = True
+
+    @staticmethod
+    def get_calculator(
+        config: Dict[str, Any],
+        directory: str | None = None,
+        mpi: int | None = None,
+        omp: int | None = None,
+        **kwargs: Any,
+    ) -> Calculator:
+        AbacusFactory._log_backend_source_once()
+        abacus_config = _abacus_section(config)
+        raw_parameters = dict(abacus_config.get("parameters", {}))
+
+        parameters = {
+            key: value
+            for key, value in abacus_config.items()
+            if key not in _ABACUS_CONTROL_KEYS
+        }
+        parameters.update(raw_parameters)
+        parameters.update(kwargs)
+
+        if "pp" in parameters:
+            parameters["pseudopotentials"] = parameters.pop("pp")
+        if "basis" in parameters:
+            parameters["basissets"] = parameters.pop("basis")
+        if "basis_dir" in parameters:
+            parameters["orbital_dir"] = parameters.pop("basis_dir")
+        if "kpts" in parameters:
+            parameters["kpts"] = _as_mp_kpts(parameters["kpts"])
+
+        pseudo_dir = _resolve_directory(parameters.pop("pseudo_dir", None))
+        orbital_dir = _resolve_directory(parameters.pop("orbital_dir", None))
+
+        mpi = int(mpi if mpi is not None else abacus_config.get("mpi", 1))
+        omp = int(omp if omp is not None else abacus_config.get("omp", 1))
+        directory = directory or abacus_config.get("directory", ".")
+        command = _build_abacus_command(abacus_config.get("command", "abacus"), mpi)
+        version_command = abacus_config.get("version_command")
+
+        os.environ["OMP_NUM_THREADS"] = str(omp)
+        profile = ATSTAbacusProfile(
+            command=command,
+            pseudo_dir=pseudo_dir,
+            orbital_dir=orbital_dir,
+            omp_num_threads=omp,
+            version_command=version_command,
+        )
+        return Abacus(directory=directory, profile=profile, **parameters)
+
+
+class CalculatorFactory:
+    """Unified factory for calculator construction."""
+
+    @staticmethod
+    def get_calculator(name: str, config: Dict[str, Any], **kwargs: Any) -> Calculator:
+        name = name.lower()
+        if name == "abacus":
+            return AbacusFactory.get_calculator(config, **kwargs)
+        if name in {"dp", "deepmd"}:
+            return DeepPotentialFactory.get_calculator(config, **kwargs)
+        raise ValueError(f"Unsupported calculator: {name}. Supported: 'abacus', 'dp'")

atst_tools/external/ASE_interface/__init__.py

@@ -0,0 +1 @@
+"""Vendored ABACUS ASE interface snapshot."""

atst_tools/external/ASE_interface/abacuslite/__init__.py

@@ -0,0 +1,13 @@
+'''
+interfaces to Atomic-orbital Based Ab-initio Computation at UStc (ABACUS),
+for more information about this DFT calculator, 
+please refer to the Github official repository:
+https://github.com/deepmodeling/abacus-develop
+and online-manual:
+https://abacus.deepmodeling.com/en/latest/index.html
+
+For a more complete ABACUS pre-/post-processing workflow package,
+please refer to the ABACUSTest package:
+https://github.com/pxlxingliang/abacus-test
+'''
+from .core import *