Address socket i-PI review feedback

Author: jianrui geng

docs/advanced/input_files/input-main.md

@@ -9,6 +9,7 @@
     - [suffix](#suffix)
     - [ntype](#ntype)
     - [calculation](#calculation)
+    - [socket\_driver](#socket_driver)
     - [esolver\_type](#esolver_type)
     - [symmetry](#symmetry)
     - [symmetry\_prec](#symmetry_prec)
@@ -583,7 +584,6 @@
   - relax: perform structure relaxation calculations, the relax_nmax parameter depicts the maximal number of ionic iterations
   - cell-relax: perform cell relaxation calculations
   - md: perform molecular dynamics simulations
-  - socket: run as a socket client for external drivers using the i-PI protocol
   - get_pchg: obtain partial (band-decomposed) charge densities (for LCAO basis only). See out_pchg for more information
   - get_wf: obtain real space wave functions (for LCAO basis only). See out_wfc_norm and out_wfc_re_im for more information
   - get_s: obtain the overlap matrix formed by localized orbitals (for LCAO basis with multiple k points). the file name is SR.csr with file format being the same as that generated by out_mat_hs2
@@ -593,6 +593,15 @@
   - test_neighbour: obtain information of neighboring atoms (for LCAO basis only), please specify a positive search_radius manually
 - **Default**: scf
 
+### socket_driver
+
+- **Type**: Boolean
+- **Availability**: *calculation==scf*
+- **Description**: If set to True, ABACUS keeps the calculation type as scf and receives atomic positions from an external driver through the i-PI socket protocol.
+
+  > Note: Use calculation = scf with socket_driver = True. Set ABACUS_SOCKET_ADDRESS to host:port or path:UNIX to choose the socket endpoint; if unset, localhost:31415 is used.
+- **Default**: False
+
 ### esolver_type
 
 - **Type**: String
@@ -844,7 +853,7 @@
 - **Description**: Charge extrapolation method for MD, relaxation, and socket-driven calculations.
 
   When set to default, ABACUS chooses second-order for md, first-order for
-  relax/cell-relax/socket, and atomic for other calculations. Socket-driven
+  relax/cell-relax and socket_driver calculations, and atomic for other calculations. Socket-driven
   molecular dynamics can explicitly set second-order if the external driver
   updates structures smoothly enough for second-order extrapolation.
 - **Default**: default

docs/advanced/interface/ase.md

@@ -103,6 +103,64 @@ In the new implementation, we limit the range of functionalties supported to mai
 
 Please read the examples in `interfaces/ASE_interface/examples/` for more details.
 
+### Socket I/O with ASE
+
+For socket-driven ASE workflows, use the `AbacusSocketIO` calculator. ASE runs the i-PI socket server, while ABACUS keeps `calculation=scf` and is launched with `socket_driver=1` as the client. The protocol is simple: ASE sends atomic positions and cell data to ABACUS; ABACUS evaluates one SCF step for that structure and returns energy, forces, and virial. See the [ASE socket I/O documentation](https://ase-lib.org/ase/calculators/socketio/socketio.html) and the i-PI reference paper, [Ceriotti et al., Comput. Phys. Commun. 185, 1019-1026 (2014)](https://doi.org/10.1016/j.cpc.2013.10.027), for the protocol background.
+
+Build ABACUS as usual before using this interface. PW-only builds work with `basis_type=pw`; LCAO socket calculations require an LCAO-enabled executable. No extra socket library is required.
+
+With CMake, choose the executable according to the basis:
+
+```bash
+cmake -S . -B build-pw -DENABLE_MPI=ON -DENABLE_LCAO=OFF
+cmake --build build-pw --target abacus_pw_para -j
+
+cmake -S . -B build-lcao -DENABLE_MPI=ON -DENABLE_LCAO=ON
+cmake --build build-lcao --target abacus_basic_para -j
+```
+
+With the ABACUS toolchain workflow, build the normal ABACUS executable with LCAO support when `basis_type=lcao` is needed, then pass that executable to `AbacusProfile(command=...)`. The ASE interface can be installed from this repository with:
+
+```bash
+cd interfaces/ASE_interface
+pip install .
+```
+
+A minimal socket calculator setup is:
+
+```python
+from ase.optimize import BFGS
+from abacuslite import AbacusProfile, AbacusSocketIO
+
+aprof = AbacusProfile(
+    command="mpirun -np 4 /path/to/abacus",
+    pseudo_dir="/path/to/pseudopotentials",
+    orbital_dir="/path/to/orbitals",
+    omp_num_threads=1,
+)
+
+abacus = AbacusSocketIO(
+    profile=aprof,
+    directory="socketio",
+    unixsocket="abacus_si",
+    pseudopotentials={"Si": "Si_ONCV_PBE-1.0.upf"},
+    basissets={"Si": "Si_gga_8au_100Ry_2s2p1d.orb"},
+    inp={"calculation": "scf", "basis_type": "lcao", "kspacing": 0.1},
+)
+
+with abacus as calc:
+    atoms.calc = calc
+    BFGS(atoms).run(fmax=0.05)
+```
+
+`AbacusSocketIO` sets `socket_driver=1` and `cal_force=1` automatically. The socket endpoint is selected by the calculator arguments and passed to ABACUS through `ABACUS_SOCKET_ADDRESS`, for example `localhost:31415` or `/tmp/ipi_abacus_si:UNIX`. Calling `atoms.get_potential_energy()` is supported, but the ABACUS client still computes forces because the i-PI `GETFORCE` exchange returns energy, forces, and virial as one response.
+
+A socket calculator owns one ABACUS process initialized from one fixed `INPUT`/`STRU` setup. Reuse the same `AbacusSocketIO` instance only for position updates under the same electronic-structure settings and the same cell. Do not change `kpts`, `kspacing`, `nspin`, `basis_type`, `basissets`, pseudopotentials, species, atom count, cell, or other core `INPUT`/`STRU` parameters through an existing socket calculator; create a new `AbacusSocketIO` instance and a new ABACUS client process for those changes. `AbacusSocketIO` rejects cell changes before sending them to ABACUS, and the ABACUS socket driver also checks incoming POSDATA cells against the initial `STRU` cell and exits if they differ.
+
+The i-PI protocol does not transmit element symbols. `AbacusSocketIO` therefore sorts the internal socket atoms with the same first-occurrence species grouping used when writing `STRU`, and maps returned forces back to the original ASE `Atoms` order. This avoids silent force/atom mismatches when structures are read from CIF, extxyz, POSCAR, or other formats whose atom order is not already grouped for ABACUS. Users should not manually reorder atoms for socket I/O; pass the physical ASE `Atoms` object directly to the calculator.
+
+A complete fixed-cell validation and benchmark example is available in `interfaces/ASE_interface/examples/socketio.py`.
+
 ## SPAP Analysis
 
 [SPAP](https://github.com/chuanxun/StructurePrototypeAnalysisPackage) (Structure Prototype Analysis Package) is written by Dr. Chuanxun Su to analyze symmetry and compare similarity of large amount of atomic structures. The coordination characterization function (CCF) is used to 

docs/parameters.yaml

@@ -28,7 +28,6 @@ parameters:
       * relax: perform structure relaxation calculations, the relax_nmax parameter depicts the maximal number of ionic iterations
       * cell-relax: perform cell relaxation calculations
       * md: perform molecular dynamics simulations
-      * socket: run as a socket client for external drivers using the i-PI protocol
       * get_pchg: obtain partial (band-decomposed) charge densities (for LCAO basis only). See out_pchg for more information
       * get_wf: obtain real space wave functions (for LCAO basis only). See out_wfc_norm and out_wfc_re_im for more information
       * get_s: obtain the overlap matrix formed by localized orbitals (for LCAO basis with multiple k points). the file name is SR.csr with file format being the same as that generated by out_mat_hs2
@@ -39,6 +38,16 @@ parameters:
     default_value: scf
     unit: ""
     availability: ""
+  - name: socket_driver
+    category: System variables
+    type: Boolean
+    description: |
+      If set to True, ABACUS keeps the calculation type as scf and receives atomic positions from an external driver through the i-PI socket protocol.
+
+      [NOTE] Use calculation = scf with socket_driver = True. Set ABACUS_SOCKET_ADDRESS to host:port or path:UNIX to choose the socket endpoint; if unset, localhost:31415 is used.
+    default_value: "False"
+    unit: ""
+    availability: calculation==scf
   - name: esolver_type
     category: System variables
     type: String
@@ -333,7 +342,7 @@ parameters:
       Charge extrapolation method for MD, relaxation, and socket-driven calculations.
 
       When set to default, ABACUS chooses second-order for md, first-order for
-      relax/cell-relax/socket, and atomic for other calculations. Socket-driven
+      relax/cell-relax and socket_driver calculations, and atomic for other calculations. Socket-driven
       molecular dynamics can explicitly set second-order if the external driver
       updates structures smoothly enough for second-order extrapolation.
     default_value: default

interfaces/ASE_interface/README.md

@@ -32,6 +32,7 @@ Please refer to the example scripts in the `examples` folder. Recommended learni
 7. **constraintmd.py** - Constrained molecular dynamics simulation
 8. **metadynamics.py** - Metadynamics simulation
 9. **neb.py** - Nudged Elastic Band (NEB) calculation
+10. **socketio.py** - ASE optimization with the AbacusSocketIO calculator running ABACUS as an i-PI socket client
 
 More usage examples will be provided in future versions.
 

interfaces/ASE_interface/abacuslite/core.py

@@ -45,6 +45,7 @@
     GenericFileIOCalculator,
     read_stdout
 )
+from ase.calculators.socketio import SocketIOCalculator
 from ase.atoms import Atoms
 from ase.dft.kpoints import BandPath
 from ase.io import read
@@ -135,6 +136,21 @@ def get_calculator_command(self, inputfile) -> List[str]:
         # additional inputfile argument is not used.
         return []
 
+    def socketio_argv_inet(self, port: Optional[int] = None) -> List[str]:
+        port = 31415 if port is None else port
+        return [
+            'env',
+            f'ABACUS_SOCKET_ADDRESS=localhost:{port}',
+            *self._split_command,
+        ]
+
+    def socketio_argv_unix(self, socket: str) -> List[str]:
+        return [
+            'env',
+            f'ABACUS_SOCKET_ADDRESS=/tmp/ipi_{socket}:UNIX',
+            *self._split_command,
+        ]
+
     def version(self) -> str:
         '''get the abacus version information'''
         cmd_ = [*self._split_command, '--version']
@@ -443,6 +459,17 @@ def __init__(self,
             directory=directory,
         )
 
+    def write_input(self, atoms, properties=None, system_changes=None):
+        if properties is None:
+            properties = self.template.implemented_properties
+        self.template.write_input(
+            profile=self.profile,
+            directory=Path(self.directory),
+            atoms=atoms,
+            parameters=self.parameters,
+            properties=properties,
+        )
+
     @classmethod
     def restart(cls, profile=None, directory='.', **kwargs):
         '''instantiate one ABACUS calculator from an existing job directory,
@@ -557,11 +584,200 @@ def band_structure(self, efermi=None):
         from ase.spectrum.band_structure import get_band_structure
         return get_band_structure(calc=self, reference=efermi)
 
+class AbacusSocketIO(SocketIOCalculator):
+    """ASE socket I/O calculator that launches ABACUS as an i-PI client.
+
+    A socket calculator owns one ABACUS process with one fixed INPUT/STRU
+    setup. The i-PI protocol can update positions, but electronic-structure
+    parameters such as k-points, spin, basis, pseudopotentials, and species
+    require a new calculator instance. Energy-only ASE calls are accepted, but
+    ABACUS still computes forces because i-PI GETFORCE returns energy, forces,
+    and virial together.
+    """
+
+    def __init__(self,
+                 profile=None,
+                 directory='.',
+                 port=None,
+                 unixsocket=None,
+                 timeout=None,
+                 log=None,
+                 **kwargs):
+        inp = self._socket_inp(kwargs.pop('inp', {}))
+        self.abacus = Abacus(
+            profile=profile,
+            directory=directory,
+            inp=inp,
+            **kwargs,
+        )
+        self._reference_cell = None
+        super().__init__(
+            port=port,
+            unixsocket=unixsocket,
+            timeout=timeout,
+            log=log,
+            launch_client=self._launch_client,
+        )
+
+    def calculate(self, atoms=None, properties=['energy'], system_changes=None):
+        from ase.calculators.calculator import (
+            PropertyNotImplementedError,
+            all_changes,
+        )
+        from ase.stress import full_3x3_to_voigt_6_stress
+
+        if system_changes is None:
+            system_changes = all_changes
+        if atoms is None:
+            atoms = self.atoms
+        if atoms is None:
+            raise ValueError('AbacusSocketIO.calculate requires atoms')
+
+        bad = [change for change in system_changes
+               if change not in self.supported_changes]
+        if self.atoms is not None and any(bad):
+            raise PropertyNotImplementedError(
+                'Cannot change {} through IPI protocol. '
+                'Please create new socket calculator.'
+                .format(bad if len(bad) > 1 else bad[0]))
+
+        self._check_fixed_cell(atoms)
+        order = self._socket_sort_indices(atoms)
+        socket_atoms = atoms[order]
+        self.atoms = atoms.copy()
+
+        if self.server is None:
+            self.server = self.launch_server()
+            proc = self.launch_client(socket_atoms, properties,
+                                      port=self._port,
+                                      unixsocket=self._unixsocket)
+            self.server.proc = proc
+
+        results = self.server.calculate(socket_atoms)
+        results['free_energy'] = results['energy']
+        virial = results.pop('virial')
+        if self.atoms.cell.rank == 3 and any(self.atoms.pbc):
+            vol = atoms.get_volume()
+            results['stress'] = -full_3x3_to_voigt_6_stress(virial) / vol
+        if 'forces' in results:
+            results['forces'] = self._forces_to_input_order(
+                results['forces'], order)
+        self.results.update(results)
+
+    def _check_fixed_cell(self, atoms):
+        from ase.calculators.calculator import PropertyNotImplementedError
+
+        cell = atoms.cell.array.copy()
+        if self._reference_cell is None:
+            self._reference_cell = cell
+            return
+        max_delta = np.max(np.abs(cell - self._reference_cell))
+        if max_delta > 1.0e-10:
+            raise PropertyNotImplementedError(
+                'AbacusSocketIO is fixed-cell only; create a new socket '
+                'calculator for a changed cell, or use the normal Abacus '
+                'FileIO calculator for variable-cell workflows.'
+            )
+
+    def set(self, **kwargs):
+        if kwargs:
+            raise ValueError(
+                'AbacusSocketIO input parameters are fixed after construction; '
+                'create a new AbacusSocketIO calculator to change k-points, '
+                'spin, basis, pseudopotentials, species, or other INPUT/STRU '
+                'settings.'
+            )
+        return super().set(**kwargs)
+
+    def _launch_client(self, atoms, properties=None, port=None, unixsocket=None):
+        from subprocess import Popen
+
+        if properties is None:
+            properties = self.abacus.template.implemented_properties
+
+        directory = Path(self.abacus.directory)
+        directory.mkdir(exist_ok=True, parents=True)
+
+        if hasattr(self.abacus, 'write_inputfiles'):
+            self.abacus.write_inputfiles(atoms, properties)
+        else:
+            self.abacus.write_input(atoms, properties=properties)
+
+        if unixsocket is not None:
+            argv = self.abacus.profile.socketio_argv_unix(socket=unixsocket)
+        else:
+            argv = self.abacus.profile.socketio_argv_inet(port=port)
+
+        stdout = open(directory / self.abacus.template.outputname, 'w')
+        stderr = open(directory / self.abacus.template.errorname, 'w')
+        try:
+            return Popen(argv, cwd=directory, env=os.environ,
+                         stdout=stdout, stderr=stderr)
+        finally:
+            stdout.close()
+            stderr.close()
+
+    @staticmethod
+    def _socket_inp(inp):
+        inp = dict(inp)
+        calculation = inp.get('calculation', 'scf')
+        if calculation != 'scf':
+            raise ValueError('ABACUS socket I/O requires calculation="scf"')
+        inp.update({
+            'calculation': 'scf',
+            'socket_driver': 1,
+            'cal_force': 1,
+        })
+        return inp
+
+    @staticmethod
+    def _socket_sort_indices(atoms):
+        return species_group_indices(atoms.get_chemical_symbols())
+
+    @staticmethod
+    def _forces_to_input_order(forces, order):
+        reordered = np.empty_like(forces)
+        for sorted_index, original_index in enumerate(order):
+            reordered[original_index] = forces[sorted_index]
+        return reordered
+
+
 class TestAbacusCalculator(unittest.TestCase):
 
     here = Path(__file__).parent
     pporb = here.parent.parent.parent / 'tests' / 'PP_ORB'
 
+    def test_socketio_species_order_mapping(self):
+        atoms = Atoms(symbols=['Si', 'O', 'C', 'Si', 'O', 'C'])
+        order = AbacusSocketIO._socket_sort_indices(atoms)
+        self.assertEqual(order, [0, 3, 1, 4, 2, 5])
+
+        socket_forces = np.arange(18).reshape(6, 3)
+        input_forces = AbacusSocketIO._forces_to_input_order(
+            socket_forces, order)
+
+        expected = np.empty_like(socket_forces)
+        for sorted_index, original_index in enumerate(order):
+            expected[original_index] = socket_forces[sorted_index]
+        np.testing.assert_array_equal(input_forces, expected)
+
+    def test_socketio_rejects_parameter_changes(self):
+        calc = object.__new__(AbacusSocketIO)
+        with self.assertRaisesRegex(ValueError, 'fixed after construction'):
+            calc.set(kpts={'mode': 'mp-sampling', 'nk': [2, 2, 2]})
+
+    def test_socketio_rejects_cell_changes(self):
+        from ase.calculators.calculator import PropertyNotImplementedError
+
+        calc = object.__new__(AbacusSocketIO)
+        calc.atoms = Atoms('Si', cell=[5.0, 5.0, 5.0], pbc=True)
+        calc._reference_cell = calc.atoms.cell.array.copy()
+
+        changed = calc.atoms.copy()
+        changed.cell[0, 0] = 5.1
+        with self.assertRaisesRegex(PropertyNotImplementedError, 'fixed-cell'):
+            calc._check_fixed_cell(changed)
+
     def test_calculator_results(self):
         from ase.build.bulk import bulk
         silicon = bulk('Si', crystalstructure='diamond', a=5.43)