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.

Client program	Supported by ASE calculator
Abinit	Yes
ASE	Yes - ASE provides a client as well
cp2k	No; ASE uses cp2k shell instead
DFTB+	Yes
FHI-aims	Yes
GPAW	Yes, using the ASE client
Lammps	No; ASE uses lammpsrun/lammpslib instead
NWChem	Yes
Quantum Espresso	Yes
Siesta	Yes
Yaff	No; there is no ASE calculator for Yaff

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),
            # alternative: ('UNIX:mysocketname', 31415)
            # (numeric port must be given even with Unix socket)
            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:
    # For running with UNIX socket, put unixsocket='mysocketname'
    # instead of port cf. aims parameters above
    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.

Example using DFTB+

import sys
from ase.build import molecule
from ase.calculators.dftb import Dftb
from ase.calculators.socketio import SocketIOCalculator
from ase.optimize import BFGS


atoms = molecule('H2O')
dftb = Dftb(Hamiltonian_MaxAngularMomentum_='',
            Hamiltonian_MaxAngularMomentum_O='"p"',
            Hamiltonian_MaxAngularMomentum_H='"s"',
            Driver_='',
            Driver_Socket_='',
            Driver_Socket_File='Hello')
opt = BFGS(atoms, trajectory='test.traj')


with SocketIOCalculator(dftb, log=sys.stdout, unixsocket='Hello') as calc:
    atoms.calc = calc
    opt.run(fmax=0.01)

Note

The DFTB+ script did not work with INET sockets. This may have been a problem on the test machine. The relevant keyword is Driver_Socket_Port=<portnumber> in case someone wants to test.

Example using NWChem

import sys

from ase.build import molecule
from ase.optimize import BFGS
from ase.calculators.nwchem import NWChem
from ase.calculators.socketio import SocketIOCalculator

atoms = molecule('H2O')
atoms.rattle(stdev=0.1)

unixsocket = 'ase_nwchem'

nwchem = NWChem(theory='scf',
                task='optimize',
                driver={'socket': {'unix': unixsocket}})

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

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

Example using Abinit

from ase.build import bulk
from ase.optimize import BFGS
from ase.calculators.abinit import Abinit
from ase.calculators.socketio import SocketIOCalculator
from ase.constraints import ExpCellFilter


atoms = bulk('Si')
atoms.rattle(stdev=0.1, seed=42)

# Configuration parameters; please edit as appropriate
pps = '/path/to/pseudopotentials'
pseudopotentials = {'Si': '14-Si.LDA.fhi'}
exe = 'abinit'

unixsocket = 'ase_abinit'
command = f'{exe} PREFIX.in --ipi {unixsocket}:UNIX > PREFIX.log'
# (In the command, note that PREFIX.in must precede --ipi.)


configuration_kwargs = dict(
    command=command,
    pp_paths=[pps],
    v8_legacy_format=False,
)


# Implementation note: Socket-driven calculations in Abinit inherit several
# controls for from the ordinary cell optimization code.  We have to hack those
# variables in order for Abinit not to decide that the calculation converged:
boilerplate_kwargs = dict(
    ionmov=28,  # activate i-pi/socket mode
    expert_user=1,  # Ignore warnings (chksymbreak, chksymtnons, chkdilatmx)
    optcell=2,  # allow the cell to relax
    tolmxf=1e-300,  # Prevent Abinit from thinking we "converged"
    ntime=100_000,  # Allow "infinitely" many iterations in Abinit
    ecutsm=0.5,  # Smoothing PW cutoff energy (mandatory for cell optimization)
)


kwargs = dict(
    ecut=5 * 27.3,
    tolvrs=1e-8,
    kpts=[2, 2, 2],
    **boilerplate_kwargs,
    **configuration_kwargs,
)

abinit = Abinit(**kwargs)

opt = BFGS(ExpCellFilter(atoms),
           trajectory='opt.traj')

with SocketIOCalculator(abinit, unixsocket=unixsocket) as atoms.calc:
    opt.run(fmax=0.01)

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 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, *, launch_client=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(port=None, unixsocket=None, timeout=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.