Communication with calculators over sockets

ASE can use sockets to communicate efficiently with certain external codes using the protocol of i-PI. This may significantly speed up geometry optimizations, dynamics and other algorithms in which ASE moves the atoms while the external code calculates energies, forces, and stress. Note that ASE does not require i-PI, but simply uses the same protocol.

The reference article for i-PI is Ceriotti, More, Manolopoulos, Comp. Phys. Comm. 185, 1019-1026 (2014).

Introduction

Normally, file-IO calculators in ASE launch a new process to calculate every atomic step. This is inefficient since the codes will need to either start from scratch or perform significant IO between steps.

Some codes can run in “driver mode” where a server provides atomic coordinates through a socket connection, and the code returns energies, forces, and stress to the server. That way the startup overhead is eliminated, and the codes can reuse and extrapolate wavefunctions and other quantities for increased efficiency.

ASE provides such a server in the form of a calculator.

Which codes can be used with socket I/O calculators?

Below is a list of codes that can run as clients, and whether ASE provides a calculator that supports doing so.

Program name Supported by ASE calculator
Quantum Espresso Yes
FHI-aims Yes
Siesta Yes
DFTB+ Yes, presumably (untested)
Yaff No; there is no ASE calculator for Yaff
cp2k No; ASE uses cp2k shell instead
Lammps No; ASE uses lammpsrun/lammpslib instead
ASE Yes - ASE provides a client as well
GPAW Yes, using the ASE client

The codes that are “not supported” by ASE can still be used as clients, but you will need to generate the input files and launch the client programs yourself.

Codes may require different commands, keywords, or compilation options in order to run in driver mode. See the code’s documentation for details. The i-PI documentation may also be useful.

How to use the ASE socket I/O interface

Example using Quantum Espresso

import sys

from ase.build import molecule
from ase.optimize import BFGS
from ase.calculators.espresso import Espresso
from ase.calculators.socketio import SocketIOCalculator

atoms = molecule('H2O', vacuum=3.0)
atoms.rattle(stdev=0.1)

# Environment-dependent parameters (please configure before running):
pseudopotentials = {'H': 'H.pbe-rrkjus.UPF',
                    'O': 'O.pbe-rrkjus.UPF'}
pseudo_dir='.'

# In this example we use a UNIX socket.  See other examples for INET socket.
# UNIX sockets are faster then INET sockets, but cannot run over a network.
# UNIX sockets are files.  The actual path will become /tmp/ipi_ase_espresso.
unixsocket = 'ase_espresso'

# Configure pw.x command for UNIX or INET.
#
# UNIX: --ipi {unixsocket}:UNIX
# INET: --ipi {host}:{port}
#
# See also QE documentation, e.g.:
#
#    https://www.quantum-espresso.org/Doc/pw_user_guide/node13.html
#
command = ('pw.x < PREFIX.pwi --ipi {unixsocket}:UNIX > PREFIX.pwo'
           .format(unixsocket=unixsocket))

espresso = Espresso(command=command,
                    ecutwfc=30.0,
                    pseudopotentials=pseudopotentials,
                    pseudo_dir=pseudo_dir)

opt = BFGS(atoms, trajectory='opt.traj',
           logfile='opt.log')

with SocketIOCalculator(espresso, log=sys.stdout,
                        unixsocket=unixsocket) as calc:
    atoms.calc = calc
    opt.run(fmax=0.05)

# Note: QE does not generally quit cleanly - expect nonzero exit codes.

Note

It is wise to ensure smooth termination of the connection. This can be done by calling calc.close() at the end or, more elegantly, by enclosing using the with statement as done in all examples here.

Example using FHI-aims

import sys

from ase.build import molecule
from ase.optimize import BFGS
from ase.calculators.aims import Aims
from ase.calculators.socketio import SocketIOCalculator

# Environment-dependent parameters -- please configure according to machine
# Note that FHI-aim support for the i-PI protocol must be specifically
# enabled at compile time, e.g.: make -f Makefile.ipi ipi.mpi
species_dir = '/home/aimsuser/src/fhi-aims.171221_1/species_defaults/light'
command = 'ipi.aims.171221_1.mpi.x'

# This example uses INET; see other examples for how to use UNIX sockets.
port = 31415

atoms = molecule('H2O', vacuum=3.0)
atoms.rattle(stdev=0.1)

aims = Aims(command=command,
            use_pimd_wrapper=('localhost', port),
            compute_forces=True,
            xc='LDA',
            species_dir=species_dir)

opt = BFGS(atoms, trajectory='opt.aims.traj', logfile='opt.aims.log')

with SocketIOCalculator(aims, log=sys.stdout, port=port) as calc:
    atoms.calc = calc
    opt.run(fmax=0.05)

Example using Siesta

import sys
from ase.build import molecule
from ase.calculators.siesta import Siesta
from ase.optimize import BFGS
from ase.calculators.socketio import SocketIOCalculator

unixsocket = 'siesta'

fdf_arguments = {'MD.TypeOfRun': 'Master',
                 'Master.code': 'i-pi',
                 'Master.interface': 'socket',
                 'Master.address': unixsocket,
                 'Master.socketType': 'unix'}

# To connect through INET socket instead, use:
#   fdf_arguments['Master.port'] = port
#   fdf_arguments['Master.socketType'] = 'inet'
# Optional, for networking:
#   fdf_arguments['Master.address'] = <hostname or IP address>

atoms = molecule('H2O', vacuum=3.0)
atoms.rattle(stdev=0.1)

siesta = Siesta(fdf_arguments=fdf_arguments)
opt = BFGS(atoms, trajectory='opt.siesta.traj', logfile='opt.siesta.log')

with SocketIOCalculator(siesta, log=sys.stdout,
                        unixsocket=unixsocket) as calc:
    atoms.calc = calc
    opt.run(fmax=0.05)

# Note: Siesta does not exit cleanly - expect nonzero exit codes.

For codes other than these, see the next section.

Run server and client manually

ASE can run as a client using the SocketClient class. This may be useful for controlling calculations remotely or using a serial process to control a parallel one.

This example will launch a server without (necessarily) launching any client:

import sys

from ase.build import molecule
from ase.io import write
from ase.optimize import BFGS
from ase.calculators.socketio import SocketIOCalculator

unixsocket = 'ase_server_socket'

atoms = molecule('H2O', vacuum=3.0)
atoms.rattle(stdev=0.1)
write('initial.traj', atoms)

opt = BFGS(atoms, trajectory='opt.driver.traj', logfile='opt.driver.log')

with SocketIOCalculator(log=sys.stdout,
                        unixsocket=unixsocket) as calc:
    # Server is now running and waiting for connections.
    # If you want to launch the client process here directly,
    # instead of manually in the terminal, uncomment these lines:
    #
    # from subprocess import Popen
    # proc = Popen([sys.executable, 'example_client_gpaw.py'])

    atoms.calc = calc
    opt.run(fmax=0.05)

Run it and then run the client:

from __future__ import print_function
from ase.io import read
from ase.calculators.socketio import SocketClient
from gpaw import GPAW, Mixer

# The atomic numbers are not transferred over the socket, so we have to
# read the file
atoms = read('initial.traj')
unixsocket = 'ase_server_socket'

atoms.calc = GPAW(mode='lcao',
                  basis='dzp',
                  txt='gpaw.client.txt',
                  mixer=Mixer(0.7, 7, 20.0))

client = SocketClient(unixsocket=unixsocket)

# Each step of the loop changes the atomic positions, but the generator
# yields None.
for i, _ in enumerate(client.irun(atoms, use_stress=False)):
    print('step:', i)

This also demonstrates how to use the interface with GPAW. Instead of running the client script, it is also possible to run any other program that acts as a client. This includes the codes listed in the compatibility table above.

Module documentation

class ase.calculators.socketio.SocketIOCalculator(calc=None, port=None, unixsocket=None, timeout=None, log=None)[source]

Initialize socket I/O calculator.

This calculator launches a server which passes atomic coordinates and unit cells to an external code via a socket, and receives energy, forces, and stress in return.

ASE integrates this with the Quantum Espresso, FHI-aims and Siesta calculators. This works with any external code that supports running as a client over the i-PI protocol.

Parameters:

calc: calculator or None

If calc is not None, a client process will be launched using calc.command, and the input file will be generated using calc.write_input(). Otherwise only the server will run, and it is up to the user to launch a compliant client process.

port: integer

port number for socket. Should normally be between 1025 and 65535. Typical ports for are 31415 (default) or 3141.

unixsocket: str or None

if not None, ignore host and port, creating instead a unix socket using this name prefixed with /tmp/ipi_. The socket is deleted when the calculator is closed.

timeout: float >= 0 or None

timeout for connection, by default infinite. See documentation of Python sockets. For longer jobs it is recommended to set a timeout in case of undetected client-side failure.

log: file object or None (default)

logfile for communication over socket. For debugging or the curious.

In order to correctly close the sockets, it is recommended to use this class within a with-block:

>>> with SocketIOCalculator(...) as calc:
...    atoms.calc = calc
...    atoms.get_forces()
...    atoms.rattle()
...    atoms.get_forces()

It is also possible to call calc.close() after use. This is best done in a finally-block.

class ase.calculators.socketio.SocketClient(host='localhost', port=None, unixsocket=None, timeout=None, log=None, comm=None)[source]

Create client and connect to server.

Parameters:

host: string
Hostname of server. Defaults to localhost
port: integer or None
Port to which to connect. By default 31415.
unixsocket: string or None
If specified, use corresponding UNIX socket. See documentation of unixsocket for SocketIOCalculator.
timeout: float or None
See documentation of timeout for SocketIOCalculator.
log: file object or None
Log events to this file
comm: communicator or None
MPI communicator object. Defaults to ase.parallel.world. When ASE runs in parallel, only the process with world.rank == 0 will communicate over the socket. The received information will then be broadcast on the communicator. The SocketClient must be created on all ranks of world, and will see the same Atoms objects.

The SocketServer allows launching a server without the need to create a calculator:

class ase.calculators.socketio.SocketServer(client_command=None, port=None, unixsocket=None, timeout=None, cwd=None, log=None)[source]

Create server and listen for connections.

Parameters:

client_command: Shell command to launch client process, or None
The process will be launched immediately, if given. Else the user is expected to launch a client whose connection the server will then accept at any time. One calculate() is called, the server will block to wait for the client.
port: integer or None
Port on which to listen for INET connections. Defaults to 31415 if neither this nor unixsocket is specified.
unixsocket: string or None
Filename for unix socket.
timeout: float or None
timeout in seconds, or unlimited by default. This parameter is passed to the Python socket object; see documentation therof
log: file object or None
useful debug messages are written to this.