atst-tools 2.3.2__tar.gz → 2.3.3__tar.gz

{atst_tools-2.3.2/src/atst_tools.egg-info → atst_tools-2.3.3}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: atst-tools
-Version: 2.3.2
+Version: 2.3.3
 Summary: Advanced ASE Transition State Tools for ABACUS and Deep-Potential
 License-Expression: LGPL-3.0-or-later
 Requires-Python: >=3.10
@@ -24,6 +24,7 @@ Requires-Dist: mpi4py>=4.1.2; extra == "parallel"
 Provides-Extra: test
 Requires-Dist: pytest<10,>=8.4; extra == "test"
 Requires-Dist: tomli<3,>=2.0; extra == "test"
+Requires-Dist: pyyaml<7,>=6; extra == "test"
 Provides-Extra: release
 Requires-Dist: build<2,>=1.5; extra == "release"
 Requires-Dist: twine<7,>=6.2; extra == "release"
@@ -33,11 +34,12 @@ Requires-Dist: pytest<10,>=8.4; extra == "dev"
 Requires-Dist: build<2,>=1.5; extra == "dev"
 Requires-Dist: twine<7,>=6.2; extra == "dev"
 Requires-Dist: tomli<3,>=2.0; extra == "dev"
+Requires-Dist: pyyaml<7,>=6; extra == "dev"
 Dynamic: license-file
 
 # ATST-Tools
 
-[![Version](https://img.shields.io/badge/version-2.3.2-blue)](pyproject.toml)
+[![Version](https://img.shields.io/badge/version-2.3.3-blue)](pyproject.toml)
 [![Python](https://img.shields.io/badge/python-%3E%3D3.10-blue)](pyproject.toml)
 [![License](https://img.shields.io/badge/license-LGPL--v3-blue)](#license)
 
@@ -67,11 +69,11 @@ Examples: [GitHub](https://github.com/QuantumMisaka/atst-tools/tree/main/example
 | Intrinsic reaction coordinates | `irc` |
 | ASE-driven or ABACUS-native molecular dynamics | `md` |
 | Experimental Direct MaxFlux path search | `dmf` |
-| Experimental ABACUS constant-potential single points and serial scans | `constant_potential` |
+| ABACUS constant-potential single points and serial scans | `constant_potential` |
 
-Experimental constant-potential fixed-cell relaxation, NEB and ASE-driven MD
+Constant-potential fixed-cell relaxation, serial NEB and ASE-driven MD
 use the explicit `compensated_gate` boundary. See the
-[configuration reference](docs/user/CONFIG_REFERENCE.md#211-constant-potential-evaluation-development-candidate)
+[configuration reference](docs/user/CONFIG_REFERENCE.md#211-constant-potential-evaluation)
 for setup and supported boundaries.
 
 Local commands prepare paths and ABACUS inputs, inspect trajectories, and
@@ -107,10 +109,10 @@ backend inputs.
 
 ## Quick Start
 
-Obtain examples at the tag matching your installed version (shown for 2.3.2):
+Obtain examples at the tag matching your installed version (shown for 2.3.3):
 
 ```bash
-git clone --branch v2.3.2 https://github.com/QuantumMisaka/atst-tools.git
+git clone --branch v2.3.3 https://github.com/QuantumMisaka/atst-tools.git
 cd atst-tools
 ```
 
@@ -149,7 +151,7 @@ methods or DP variants. A complete walkthrough is also available in the
 | Use CLI commands | [CLI reference](docs/user/CLI_REFERENCE.md) |
 | Configure ABACUS input/output and parallelism | [ABACUS wrapper guide](docs/user/ABACUSLITE_WRAPPER_GUIDE.md) |
 | Embed workflows or use the JSON process runner | [Python API](docs/user/PYTHON_API_REFERENCE.md) |
-| Check support and releases | [Feature matrix](docs/reports/FEATURE_STATUS_MATRIX.md) · [2.3.2 release notes](docs/releases/RELEASE_NOTES_2.3.2.md) |
+| Check support and releases | [Feature matrix](docs/reports/FEATURE_STATUS_MATRIX.md) · [2.3.3 release notes](docs/releases/RELEASE_NOTES_2.3.3.md) |
 | Contribute | [Developer handover](docs/developer/HANDOVER.md) · [Source architecture](src/atst_tools/README.md) |
 
 The primary documentation is in English, with a [Chinese user guide](docs/user/USER_GUIDE_CN.md).

{atst_tools-2.3.2 → atst_tools-2.3.3}/README.md

@@ -1,6 +1,6 @@
 # ATST-Tools
 
-[![Version](https://img.shields.io/badge/version-2.3.2-blue)](pyproject.toml)
+[![Version](https://img.shields.io/badge/version-2.3.3-blue)](pyproject.toml)
 [![Python](https://img.shields.io/badge/python-%3E%3D3.10-blue)](pyproject.toml)
 [![License](https://img.shields.io/badge/license-LGPL--v3-blue)](#license)
 
@@ -30,11 +30,11 @@ Examples: [GitHub](https://github.com/QuantumMisaka/atst-tools/tree/main/example
 | Intrinsic reaction coordinates | `irc` |
 | ASE-driven or ABACUS-native molecular dynamics | `md` |
 | Experimental Direct MaxFlux path search | `dmf` |
-| Experimental ABACUS constant-potential single points and serial scans | `constant_potential` |
+| ABACUS constant-potential single points and serial scans | `constant_potential` |
 
-Experimental constant-potential fixed-cell relaxation, NEB and ASE-driven MD
+Constant-potential fixed-cell relaxation, serial NEB and ASE-driven MD
 use the explicit `compensated_gate` boundary. See the
-[configuration reference](docs/user/CONFIG_REFERENCE.md#211-constant-potential-evaluation-development-candidate)
+[configuration reference](docs/user/CONFIG_REFERENCE.md#211-constant-potential-evaluation)
 for setup and supported boundaries.
 
 Local commands prepare paths and ABACUS inputs, inspect trajectories, and
@@ -70,10 +70,10 @@ backend inputs.
 
 ## Quick Start
 
-Obtain examples at the tag matching your installed version (shown for 2.3.2):
+Obtain examples at the tag matching your installed version (shown for 2.3.3):
 
 ```bash
-git clone --branch v2.3.2 https://github.com/QuantumMisaka/atst-tools.git
+git clone --branch v2.3.3 https://github.com/QuantumMisaka/atst-tools.git
 cd atst-tools
 ```
 
@@ -112,7 +112,7 @@ methods or DP variants. A complete walkthrough is also available in the
 | Use CLI commands | [CLI reference](docs/user/CLI_REFERENCE.md) |
 | Configure ABACUS input/output and parallelism | [ABACUS wrapper guide](docs/user/ABACUSLITE_WRAPPER_GUIDE.md) |
 | Embed workflows or use the JSON process runner | [Python API](docs/user/PYTHON_API_REFERENCE.md) |
-| Check support and releases | [Feature matrix](docs/reports/FEATURE_STATUS_MATRIX.md) · [2.3.2 release notes](docs/releases/RELEASE_NOTES_2.3.2.md) |
+| Check support and releases | [Feature matrix](docs/reports/FEATURE_STATUS_MATRIX.md) · [2.3.3 release notes](docs/releases/RELEASE_NOTES_2.3.3.md) |
 | Contribute | [Developer handover](docs/developer/HANDOVER.md) · [Source architecture](src/atst_tools/README.md) |
 
 The primary documentation is in English, with a [Chinese user guide](docs/user/USER_GUIDE_CN.md).

{atst_tools-2.3.2 → atst_tools-2.3.3}/pyproject.toml

@@ -4,7 +4,7 @@ build-backend = "setuptools.build_meta"
 
 [project]
 name = "atst-tools"
-version = "2.3.2"
+version = "2.3.3"
 description = "Advanced ASE Transition State Tools for ABACUS and Deep-Potential"
 readme = "README.md"
 license = "LGPL-3.0-or-later"
@@ -24,13 +24,14 @@ dependencies = [
 plot = ["matplotlib>=3.9,<4"]
 dp = ["deepmd-kit>=3.1.3"]
 parallel = ["mpi4py>=4.1.2"]
-test = ["pytest>=8.4,<10", "tomli>=2.0,<3"]
+test = ["pytest>=8.4,<10", "tomli>=2.0,<3", "pyyaml>=6,<7"]
 release = ["build>=1.5,<2", "twine>=6.2,<7", "tomli>=2.0,<3"]
 dev = [
     "pytest>=8.4,<10",
     "build>=1.5,<2",
     "twine>=6.2,<7",
     "tomli>=2.0,<3",
+    "pyyaml>=6,<7",
 ]
 
 [project.scripts]

{atst_tools-2.3.2 → atst_tools-2.3.3}/src/atst_tools/api/services.py

@@ -24,6 +24,7 @@ from atst_tools.api.models import (
     WorkflowResult,
 )
 from atst_tools.calculators.abacuslite_backend import BACKEND_SOURCE
+from atst_tools.calculators.factory import socket_session
 from atst_tools.utils.artifacts import (
     SCHEMA_VERSION,
     add_backend_identity_to_manifest,
@@ -915,7 +916,10 @@ def _run_workflow(
         failure = None
         dispatch_started = time.monotonic()
         try:
-            value = _dispatch_normalized(config, options)
+            # Persistent ABACUS socket clients built by this workflow must not
+            # outlive it; clients owned by other callers are left untouched.
+            with socket_session():
+                value = _dispatch_normalized(config, options)
         except ATSTAPIError as exc:
             failure = exc
         except Exception as exc:

{atst_tools-2.3.2 → atst_tools-2.3.3}/src/atst_tools/calculators/abacuslite_backend.py

@@ -20,12 +20,32 @@ def _load_abacuslite_backend():
             "abacuslite is required by ATST-Tools; install abacuslite>=1.2.1,<2",
             name="abacuslite",
         )
-    from abacuslite import Abacus, AbacusProfile
+    from abacuslite import Abacus, AbacusProfile, AbacusSocketIO
 
-    return Abacus, AbacusProfile, "external"
+    return Abacus, AbacusProfile, AbacusSocketIO, "external"
 
 
-Abacus, AbacusProfile, BACKEND_SOURCE = _load_abacuslite_backend()
+Abacus, AbacusProfile, AbacusSocketIO, BACKEND_SOURCE = _load_abacuslite_backend()
+
+
+def socket_input_parameters(parameters: dict) -> dict:
+    """Apply the backend's socket-mode INPUT rules to ABACUS parameters.
+
+    Delegates to ``AbacusSocketIO._socket_inp`` so prepared INPUT files match
+    what the socket calculator writes: ``socket_driver = 1``,
+    ``calculation = scf`` and normalized ``cal_force``/``cal_stress``.
+
+    Args:
+        parameters: ABACUS INPUT parameters.
+
+    Returns:
+        A new parameter mapping with socket-mode values applied.
+
+    Raises:
+        ValueError: If ``calculation`` is not ``scf`` or a ``cal_*`` flag is
+            not a boolean value.
+    """
+    return AbacusSocketIO._socket_inp(parameters)
 
 
 def import_generalio():
@@ -131,6 +151,7 @@ class ATSTAbacusProfile(AbacusProfile):
 __all__ = [
     "Abacus",
     "AbacusProfile",
+    "AbacusSocketIO",
     "ATSTAbacusProfile",
     "BACKEND_SOURCE",
     "BackendSource",

{atst_tools-2.3.2 → atst_tools-2.3.3}/src/atst_tools/calculators/constant_potential.py

@@ -20,6 +20,7 @@ from typing import Any
 import numpy as np
 from ase.calculators.calculator import Calculator, all_changes
 
+from atst_tools.utils.config_schema import ABACUS_CONTROL_KEYS
 from atst_tools.utils.electrochemistry import chebyshev_derivative, surface_area
 
 
@@ -61,7 +62,7 @@ _CURRENT_SCF_FAILURE_MARKER = re.compile(
 CP_FACTS_INFO_KEY = "atst_constant_potential_facts"
 CP_IDENTITY_INFO_KEY = "atst_constant_potential_identity"
 CP_FACTS_SCHEMA = "atst-constant-potential-facts-v1"
-_ABACUS_CONTROL_KEYS = {"command", "mpi", "omp", "directory", "parameters", "version_command"}
+_ABACUS_CONTROL_KEYS = ABACUS_CONTROL_KEYS
 _FIXED_BOUNDARY_KEYS = (
     "gate_flag",
     "imp_sol",
@@ -359,7 +360,8 @@ def _same_finite_float(left: Any, right: Any) -> bool:
 
 
 def _constant_potential_facts_error(
-    atoms: Any, payload: Any, *, require_duplicate_identity: bool = True
+    atoms: Any, payload: Any, *, require_duplicate_identity: bool = True,
+    check_geometry: bool = True,
 ) -> str | None:
     """Return a fail-closed validation error for one durable CP envelope."""
     if atoms is None or not hasattr(atoms, "info"):
@@ -370,12 +372,13 @@ def _constant_potential_facts_error(
         return "constant-potential facts schema or workflow is unsupported"
     if payload.get("status") != "complete" or payload.get("cp_converged") is not True:
         return "constant-potential facts are not a complete converged result"
-    try:
-        geometry_fingerprint = constant_potential_geometry_fingerprint(atoms)
-    except (AttributeError, TypeError, ValueError):
-        return "constant-potential facts geometry cannot be fingerprinted"
-    if payload.get("geometry_fingerprint") != geometry_fingerprint:
-        return "constant-potential facts geometry fingerprint does not match"
+    if check_geometry:
+        try:
+            geometry_fingerprint = constant_potential_geometry_fingerprint(atoms)
+        except (AttributeError, TypeError, ValueError):
+            return "constant-potential facts geometry cannot be fingerprinted"
+        if payload.get("geometry_fingerprint") != geometry_fingerprint:
+            return "constant-potential facts geometry fingerprint does not match"
 
     identity = payload.get("identity")
     facts = payload.get("facts")
@@ -612,12 +615,14 @@ def publish_constant_potential_facts(
     return copy.deepcopy(payload)
 
 
-def read_constant_potential_facts(atoms: Any) -> dict[str, Any] | None:
-    """Read and validate CP facts persisted on an ASE atoms object."""
+def read_constant_potential_facts(
+    atoms: Any, *, check_geometry: bool = True
+) -> dict[str, Any] | None:
+    """Read valid CP facts; restart checks geometry while NEB may use ASE state."""
     if atoms is None or not hasattr(atoms, "info"):
         return None
     payload = atoms.info.get(CP_FACTS_INFO_KEY)
-    if _constant_potential_facts_error(atoms, payload) is not None:
+    if _constant_potential_facts_error(atoms, payload, check_geometry=check_geometry) is not None:
         return None
     return copy.deepcopy(dict(payload))
 
@@ -751,7 +756,14 @@ def _read_vacuum_from_logs(directory: str | Path, suffix: str | None = None) ->
 
 
 def _backend_log_paths(calculator: Any) -> tuple[Path, ...]:
-    """Return only the current backend's SCF log candidates."""
+    """Return only the current backend's SCF log candidates.
+
+    An abacuslite ``AbacusSocketIO`` keeps its ABACUS working directory and
+    INPUT parameters on the wrapped FileIO calculator (``calculator.abacus``);
+    its own ``directory`` is ASE's default ``'.'``, which is not where ABACUS
+    writes.
+    """
+    calculator = getattr(calculator, "abacus", calculator)
     directory = getattr(calculator, "directory", None)
     if directory is None:
         return ()
@@ -765,7 +777,15 @@ def _backend_log_paths(calculator: Any) -> tuple[Path, ...]:
 
 
 def _backend_scf_converged(calculator: Any, results: Mapping[str, Any]) -> bool | None:
-    """Read an explicit SCF-validity fact from results or the current log."""
+    """Read an explicit SCF-validity fact from results, the backend, or its log.
+
+    Args:
+        calculator: Backend calculator that produced ``results``.
+        results: Result mapping of the evaluation being validated.
+
+    Returns:
+        ``True``/``False`` for an explicit fact, ``None`` when no evidence exists.
+    """
     for key in ("scf_converged", "scf_convergence", "scf_status", "convergence"):
         if key not in results or results[key] is None:
             continue
@@ -782,6 +802,11 @@ def _backend_scf_converged(calculator: Any, results: Mapping[str, Any]) -> bool
             return bool(int(value))
         except (TypeError, ValueError):
             return None
+    # abacuslite's socket client reports the per-evaluation SCF state from the
+    # ABACUS i-PI metadata as an attribute; it is reset before every request.
+    socket_converged = getattr(calculator, "last_scf_converged", None)
+    if isinstance(socket_converged, (bool, np.bool_)):
+        return bool(socket_converged)
     for path in _backend_log_paths(calculator):
         if not path.is_file():
             continue

{atst_tools-2.3.2 → atst_tools-2.3.3}/src/atst_tools/calculators/factory.py

@@ -2,18 +2,21 @@
 
 from __future__ import annotations
 
+import itertools
 import os
 import re
 import shlex
 import logging
 import math
+from contextlib import contextmanager
+from contextvars import ContextVar
 from copy import deepcopy
-from typing import Any, Dict
+from typing import Any, Dict, Iterator
 
 from ase.calculators.calculator import Calculator
 
 from atst_tools.calculators.abacuslite_backend import (
-    Abacus, ATSTAbacusProfile, BACKEND_SOURCE, _default_version_command,
+    Abacus, AbacusSocketIO, ATSTAbacusProfile, BACKEND_SOURCE, _default_version_command,
 )
 from atst_tools.calculators.constant_potential import (
     ConstantPotentialCalculator,
@@ -23,13 +26,18 @@ from atst_tools.calculators.dp import DeepPotentialFactory
 from atst_tools.runtime import counters as runtime_counters
 from atst_tools.runtime import launch as runtime_launch
 from atst_tools.utils.config_schema import (
+    ABACUS_CONTROL_KEYS,
+    SOCKET_FIXED_CELL_MD_ENSEMBLES,
+    SOCKET_SUPPORTED_WORKFLOWS,
     ConstantPotentialConfig,
     validate_constant_potential_md_scope,
 )
 from atst_tools.utils.mpi import mpi_launcher_detected
 
 
-_ABACUS_CONTROL_KEYS = {"command", "mpi", "omp", "directory", "parameters", "version_command"}
+_ABACUS_CONTROL_KEYS = ABACUS_CONTROL_KEYS
+# INPUT keys that AbacusSocketIO validates only when passed through ``inp``.
+_SOCKET_INP_KEYS = ("calculation", "cal_force", "cal_stress")
 _MPI_ENV_KEYS_TO_CLEAR = (
     "OMPI_COMM_WORLD_SIZE",
     "OMPI_COMM_WORLD_RANK",
@@ -46,6 +54,9 @@ LOGGER = logging.getLogger(__name__)
 _ABACUS_BACKEND_LOGGED = False
 _SHELL_ASSIGNMENT_RE = re.compile(r"^[A-Za-z_][A-Za-z0-9_]*=.*$")
 _MPI_LAUNCHERS = {"mpirun", "mpiexec", "srun"}
+_SOCKET_NAME_COUNTER = itertools.count()
+#: Socket clients owned by the innermost active :func:`socket_session`.
+_SOCKET_SESSION: ContextVar[list[Any] | None] = ContextVar("atst_abacus_socket_session", default=None)
 
 
 def _abacus_section(config: Dict[str, Any]) -> Dict[str, Any]:
@@ -133,6 +144,96 @@ def _validate_abacus_command(command: str) -> None:
         )
 
 
+def _socket_options(raw: Any) -> Dict[str, Any] | None:
+    """Normalize ``calculator.abacus.socket`` into AbacusSocketIO keywords.
+
+    Args:
+        raw: ``None``/``False`` (disabled), ``True`` (enabled with defaults), or
+            a mapping with optional ``enabled``, ``unixsocket``, ``port`` and
+            ``timeout`` keys.
+
+    Returns:
+        Keyword arguments for ``AbacusSocketIO`` or ``None`` when disabled.
+
+    Raises:
+        ValueError: If the value has an unsupported shape or both ``port`` and
+            ``unixsocket`` are given.
+    """
+    if raw is None or raw is False:
+        return None
+    if raw is True:
+        raw = {}
+    if not isinstance(raw, dict):
+        raise ValueError("calculator.abacus.socket must be a boolean or a mapping")
+    options = dict(raw)
+    if not options.pop("enabled", True):
+        return None
+    unknown = set(options) - {"unixsocket", "port", "timeout"}
+    if unknown:
+        raise ValueError(f"Unknown calculator.abacus.socket keys: {sorted(unknown)}")
+    if options.get("port") is not None and options.get("unixsocket") is not None:
+        raise ValueError("calculator.abacus.socket accepts either port or unixsocket, not both")
+    if options.get("port") is None and options.get("unixsocket") is None:
+        # A per-process, per-calculator name keeps concurrent runs on one node
+        # from colliding on /tmp/ipi_<name>.
+        options["unixsocket"] = f"atst_{os.getpid()}_{next(_SOCKET_NAME_COUNTER)}"
+    options.setdefault("timeout", None)
+    return options
+
+
+def _check_socket_workflow(config: Dict[str, Any]) -> None:
+    """Reject socket mode for workflows that are not single fixed-cell runs."""
+    calculation = config.get("calculation")
+    if not isinstance(calculation, dict):
+        return
+    workflow = calculation.get("type")
+    if workflow is not None and workflow not in SOCKET_SUPPORTED_WORKFLOWS:
+        raise ValueError(
+            "calculator.abacus.socket is supported only for calculation.type in "
+            f"{list(SOCKET_SUPPORTED_WORKFLOWS)}; got {workflow!r}"
+        )
+    if workflow == "md":
+        if calculation.get("driver", "ase") != "ase":
+            raise ValueError("calculator.abacus.socket requires calculation.driver=ase for md")
+        ensemble = calculation.get("ensemble", "nvt")
+        if ensemble not in SOCKET_FIXED_CELL_MD_ENSEMBLES:
+            raise ValueError(
+                "calculator.abacus.socket requires a fixed-cell md ensemble "
+                f"{list(SOCKET_FIXED_CELL_MD_ENSEMBLES)}; got {ensemble!r}"
+            )
+
+
+def _close_socket_calculator(calculator: Any) -> None:
+    """Close one socket client; failures are logged, never raised."""
+    try:
+        calculator.close()
+    except Exception as exc:  # pragma: no cover - best-effort cleanup
+        LOGGER.warning("Failed to close ABACUS socket calculator: %s", exc)
+
+
+@contextmanager
+def socket_session() -> Iterator[list[Any]]:
+    """Own the ABACUS socket clients built while this context is active.
+
+    Clients created by :class:`AbacusFactory` inside the context are closed
+    when it exits, including on error. Clients created outside any session
+    belong to the caller, who closes them. Sessions are context-local (thread
+    and asyncio safe) and may nest; each session closes only the clients it
+    created.
+
+    Yields:
+        The list of socket clients owned by this session.
+    """
+    owned: list[Any] = []
+    token = _SOCKET_SESSION.set(owned)
+    try:
+        yield owned
+    finally:
+        _SOCKET_SESSION.reset(token)
+        while owned:
+            _close_socket_calculator(owned.pop())
+
+
 def _resolve_directory(path: str | None) -> str | None:
     if path is None:
         return None
@@ -197,7 +298,13 @@ class AbacusFactory:
             omp_num_threads=omp,
             version_command=version_command,
         )
-        calculator = Abacus(directory=directory, profile=profile, **parameters)
+        socket = _socket_options(abacus_config.get("socket"))
+        if socket is None:
+            calculator = Abacus(directory=directory, profile=profile, **parameters)
+        else:
+            calculator = AbacusFactory._socket_calculator(
+                config, directory, profile, parameters, socket
+            )
         return runtime_counters.instrument_calculator(
             calculator,
             build_key="abacus.calculator_built",
@@ -205,6 +312,42 @@ class AbacusFactory:
         )
 
 
+    @staticmethod
+    def _socket_calculator(
+        config: Dict[str, Any],
+        directory: str,
+        profile: ATSTAbacusProfile,
+        parameters: Dict[str, Any],
+        socket: Dict[str, Any],
+    ) -> Calculator:
+        """Build an ``AbacusSocketIO`` calculator owned by the active socket session.
+
+        Endpoint conflicts are left to the backend: binding an occupied Unix
+        socket or TCP port fails when the first calculation starts the server.
+        """
+        _check_socket_workflow(config)
+        endpoint = (
+            f"unix:{socket['unixsocket']}" if socket.get("unixsocket") is not None else f"port:{socket['port']}"
+        )
+        parameters = dict(parameters)
+        inp = dict(parameters.pop("inp", {}) or {})
+        for key in _SOCKET_INP_KEYS:
+            if key in parameters:
+                inp[key] = parameters.pop(key)
+        calculator = AbacusSocketIO(
+            profile=profile,
+            directory=directory,
+            inp=inp,
+            **socket,
+            **parameters,
+        )
+        session = _SOCKET_SESSION.get()
+        if session is not None:
+            session.append(calculator)
+        LOGGER.info("Using ABACUS i-PI socket client (%s).", endpoint)
+        return calculator
+
+
 class CalculatorFactory:
     """Unified factory for calculator construction."""
 
@@ -218,6 +361,11 @@ class CalculatorFactory:
                 else None
             )
             if constant_potential is not None:
+                if _socket_options(_abacus_section(config).get("socket")) is not None:
+                    raise ValueError(
+                        "calculator.abacus.socket cannot be combined with "
+                        "calculator.constant_potential"
+                    )
                 return _constant_potential_calculator(config, constant_potential, kwargs)
             return AbacusFactory.get_calculator(config, **kwargs)
         if name in {"dp", "deepmd"}:

{atst_tools-2.3.2 → atst_tools-2.3.3}/src/atst_tools/mep/autoneb.py

@@ -22,6 +22,7 @@ from ase.optimize import FIRE, BFGS
 
 from atst_tools.mep.neb import AbacusNEB
 from atst_tools.calculators.factory import CalculatorFactory
+from atst_tools.calculators.constant_potential import read_constant_potential_facts
 from atst_tools.calculators.dp import is_dp_calculator, should_share_calculator
 from atst_tools.utils.artifacts import write_artifact_manifest
 from atst_tools.utils.geometry import prepare_fixed_cell_images, fixed_cell_metadata
@@ -32,8 +33,13 @@ from atst_tools.utils.convergence import (
     as_step_count,
     emit_unconverged_advisory,
 )
-from atst_tools.utils.neb_endpoints import endpoint_policy, ensure_neb_endpoint_results
+from atst_tools.utils.neb_endpoints import (
+    endpoint_policy,
+    ensure_neb_endpoint_results,
+    expected_cp_endpoint_identity,
+)
 from atst_tools.utils.neb_endpoints import freeze_current_results, freeze_results, get_endpoint_results
+from atst_tools.workflows.endpoint_relax import relax_neb_endpoints
 from atst_tools.utils.mpi import (
     get_ase_world,
     rank_owns_local_image,
@@ -813,6 +819,7 @@ class AutoNEBRunner:
             raise initial_chain_error
         self._image_index_by_id = {id(image): index for index, image in enumerate(self.init_chain)}
         self._active_autoneb = None
+        self._endpoint_stages = []
         self._apply_legacy_endpoint_conditions()
 
     def _apply_legacy_endpoint_conditions(self):
@@ -827,6 +834,7 @@ class AutoNEBRunner:
                     results[0],
                     results[1],
                     self.init_chain[index].info.get("atst_endpoint_result"),
+                    read_constant_potential_facts(self.init_chain[index], check_geometry=False),
                 )
         reference = self.init_chain[len(self.init_chain) // 2]
         magmoms = reference.get_initial_magnetic_moments()
@@ -835,8 +843,8 @@ class AutoNEBRunner:
             endpoint.set_initial_magnetic_moments(magmoms)
             if constraints:
                 endpoint.set_constraint(deepcopy(constraints))
-        for index, (energy, forces, status) in endpoint_results.items():
-            freeze_results(self.init_chain[index], energy, forces, status=status)
+        for index, (energy, forces, status, cp_facts) in endpoint_results.items():
+            freeze_results(self.init_chain[index], energy, forces, status=status, cp_facts=cp_facts)
 
     def _base_directory(self):
         base_dir = self.calc_config['directory']
@@ -910,17 +918,39 @@ class AutoNEBRunner:
         self.attach_calculators(images, synchronize=False)
 
     def _prepare_endpoint_results(self):
-        """Prepare endpoint single-point results once and synchronize images."""
+        """Relax missing-result endpoints, then synchronize their results."""
         base_dir = self._base_directory()
         policy = endpoint_policy(self.calc_config, default="auto")
+        cp_identity = expected_cp_endpoint_identity(self.config)
+        endpoint_config = self.calc_config["endpoint_optimization"]
+        self._endpoint_stages = []
 
         def prepare():
+            if endpoint_config["enabled"]:
+                if policy == "never":
+                    ensure_neb_endpoint_results(
+                        self.init_chain,
+                        lambda directory: None,
+                        policy="never",
+                        context="AutoNEB",
+                        constant_potential_identity=cp_identity,
+                    )
+                self._endpoint_stages = relax_neb_endpoints(
+                    self.init_chain,
+                    endpoint_config,
+                    lambda label, atoms: self._get_calculator(
+                        f"{base_dir}/endpoint_{label}_relax"
+                    ),
+                    self._get_optimizer(),
+                    constant_potential_identity=cp_identity,
+                )
             ensure_neb_endpoint_results(
                 self.init_chain,
                 lambda directory: self._get_calculator(f"{base_dir}/{directory}"),
                 policy=policy,
                 directories=("endpoint_initial", "endpoint_final"),
                 context="AutoNEB",
+                constant_potential_identity=cp_identity,
             )
 
         if not self.parallel:
@@ -1167,7 +1197,8 @@ class AutoNEBRunner:
                 workflow="autoneb",
                 world=self.world,
             )
-        stages = [record.to_manifest() for record in iteration_records]
+        stages = [record.to_manifest() for record in self._endpoint_stages]
+        stages.extend(record.to_manifest() for record in iteration_records)
         if final_record is not None:
             stages.append(final_record.to_manifest())