Source code for pysisyphus.calculators.IPIClient

import struct
import socket

import numpy as np

from pysisyphus.socket_helper import send_closure, recv_closure, get_fmts

[docs] def ipi_client( addr, atoms, energy_getter, forces_getter, hessian_getter=None, hdrlen=12 ): atom_num = len(atoms) # Number of entries in a Caretsian forces/coords vector cartesians = 3 * atom_num # Formats needed for struct.pack, to cast variables to bytes. # Bytes needed, to store a forces/coords vector floats_bytes = 8 * cartesians # Virial is hardcoded to the zero vector. VIRIAL = struct.pack("d" * 9, *np.zeros(9)) ZERO = struct.pack("i", 0) fmts = get_fmts(cartesians) # Unix socket is hardcoded right now, but may also be inet-socket sock = socket.socket(socket.AF_UNIX, socket.SOCK_STREAM) addr = str(addr) sock.connect(addr) send_msg = send_closure(sock, hdrlen, fmts) recv_msg = recv_closure(sock, hdrlen, fmts) counter = 0 while True: try: # Lets start talking recv_msg(expect="STATUS") # The server initiates with a STATUS send_msg("READY") status = recv_msg(expect="STATUS") if status == "NEEDPOS": send_msg("HAVEPOS") _ = recv_msg(4, fmt="int")[0] # Recive atom num from IPI coords = np.array( recv_msg(floats_bytes, "floats") ) # Receive current coords assert coords.size % 3 == 0 # Assert Cartesian coordinates send_msg(atom_num, "int") # Just send back the current coordinates or translate all atoms in +X coords.reshape(-1, 3)[:, 0] += 1 send_msg(coords, "floats") continue # When the optimization converged EXIT will be returned .. not documented! if status == "EXIT": print("Exited!") break # It seems we have to send READY two times ... not documented! send_msg("READY") # The server then returns POSDATA. recv_msg(expect="POSDATA") # Receive cell vectors, inverse cell vectors and number of atoms ... # but we don't use them here, so we don't even try to convert them to something. sock.recv(72) # cell sock.recv(72) # icell ipi_atom_num = recv_msg(4, fmt="int")[0] assert ipi_atom_num == atom_num # ... and the current coordinates. coords = np.array(recv_msg(floats_bytes, "floats")) recv_msg(expect="STATUS") # Indicate, that calculation is possible send_msg("HAVEDATA") get_what = recv_msg() # Acutal QC calculations if get_what == "GETENERGY": energy = energy_getter(coords) print(f"Calculated energy: {energy:.6f}, counter={counter}") send_msg("ENERGYREADY") send_msg(energy, "float") if get_what == "GETFORCE": forces, energy = forces_getter(coords) print(f"Calculated energy & forces: {energy:.6f}, counter={counter}") send_msg("FORCEREADY") # Send everything to the server send_msg(energy, "float") send_msg(atom_num, "int") send_msg(forces, "floats") send_msg(VIRIAL, packed=True) # We don't want to send additional information, so just send 0. send_msg(ZERO, packed=True) elif get_what == "GETHESSIAN": hessian, energy = hessian_getter(coords) hessian_packed = struct.pack("d" * cartesians ** 2, *hessian.flatten()) print(f"Calculated energy & Hessian: {energy:.6f}, counter={counter}") send_msg("HESSIANREADY") # Send everything to the server send_msg(energy, "float") send_msg(atom_num, "int") send_msg(hessian_packed, packed=True) counter += 1 except Exception as err: raise err
[docs] def calc_ipi_client(addr, atoms, calc, queue=None, **kwargs): assert calc is not None, "Supplied calculator must not be None!" def energy_getter(coords): if queue is not None: queue.put(("coords", coords)) results = calc.get_energy(atoms, coords) energy = results["energy"] return energy def forces_getter(coords): if queue is not None: queue.put(("coords", coords)) results = calc.get_forces(atoms, coords) forces = results["forces"] energy = results["energy"] return forces, energy def hessian_getter(coords): if queue is not None: queue.put(("coords", coords)) results = calc.get_hessian(atoms, coords) hessian = results["hessian"] energy = results["energy"] return hessian, energy ipi_client(addr, atoms, energy_getter, forces_getter, hessian_getter, **kwargs)