#!/usr/bin/env python3
# --------------------------------------------------------------------------------------#
# Copyright (c) 2026 MaxwellLink #
# This file is part of MaxwellLink. Repository: https://github.com/TaoELi/MaxwellLink #
# If you use this code, always credit and cite arXiv:2512.06173. #
# See AGENTS.md and README.md for details. #
# --------------------------------------------------------------------------------------#
from __future__ import annotations
import argparse
import subprocess
import shlex
import time
import shutil
import os
import socket, json
import numpy as np
try:
from .models import __drivers__
from .models.dummy_model import DummyModel
except ImportError:
from models import __drivers__
from models.dummy_model import DummyModel
description = """
A Python driver connecting to MaxwellLink, receiving E-field data and returning
the source amplitude vector for a quantum dynamics model.
"""
# The wire protocol (headers, framed ints/arrays/bytes, POSDATA/FORCEREADY
# packing, and the TCP connect-with-retry helper) lives in
# maxwelllink.sockets.protocol, which imports nothing from the rest of the
# package, so the driver shares one definition with the hubs without import
# cycles. Names are bound here so existing code that uses the module-level
# helpers keeps working unchanged.
from maxwelllink.sockets.protocol import ( # noqa: F401
BYE,
DT_FLOAT,
DT_INT,
FIELDDATA,
FORCEREADY,
GETFORCE,
GETSOURCE,
HAVEDATA,
HEADER_LEN,
INIT,
NEEDINIT,
POSDATA,
READY,
SOURCEREADY,
STATUS,
STOP,
_SocketClosed,
_connect_tcp_with_retry,
_pad12,
_recv_array,
_recv_bytes,
_recv_int,
_recv_msg,
_recv_posdata,
_recvall,
_send_array,
_send_bytes,
_send_force_ready,
_send_int,
_send_msg,
)
# helper function to determine whether this processor is the MPI master using mpi4py
def _am_master():
"""
Return True if this process is the MPI master rank (rank 0), otherwise False.
Notes
-----
Attempts to import ``mpi4py`` and query ``COMM_WORLD``. If ``mpi4py`` is not
available, returns ``True`` by treating the single process as rank 0.
"""
try:
from mpi4py import MPI as _MPI
_COMM = _MPI.COMM_WORLD
_RANK = _COMM.Get_rank()
except Exception:
_COMM = None
_RANK = 0
return _RANK == 0
def _read_value(s):
"""
Attempt to parse a string as ``int`` or ``float``; fall back to string/boolean.
Parameters
----------
s : str
Input token.
Returns
-------
int or float or bool or str
Parsed value.
"""
s = s.strip()
for cast in (int, float):
try:
return cast(s)
except ValueError:
continue
if s.lower() == "false":
return False
if s.lower() == "true":
return True
return s
def _read_args_kwargs(input_str):
"""
Parse a comma-separated string into positional and keyword arguments.
Parameters
----------
input_str : str
Comma-separated tokens. Positional values are bare; keyword values use
``key=value``. Booleans accept ``true``/``false`` (case-insensitive).
Returns
-------
tuple
``(args, kwargs)`` where ``args`` is a list and ``kwargs`` is a dict.
"""
args = []
kwargs = {}
tokens = input_str.split(",")
for token in tokens:
token = token.strip()
if "=" in token:
key, value = token.split("=", 1)
kwargs[key.strip()] = _read_value(value)
elif len(token) > 0:
args.append(_read_value(token))
return args, kwargs
[docs]
def run_driver(
unix=False,
address="localhost",
port: int = 31415,
timeout: float = 6000.0,
driver=DummyModel(),
sockets_prefix="/tmp/socketmxl_",
):
"""
Run the socket driver loop to communicate with MaxwellLink.
Parameters
----------
unix : bool, default: False
Use a UNIX domain socket if ``True``; otherwise use TCP/IP.
address : str, default: "localhost"
Hostname (TCP/IP) or UNIX socket name (when ``unix=True``).
port : int, default: 31415
TCP/IP port (ignored for UNIX sockets).
timeout : float, default: 6000.0
Socket timeout in seconds.
driver : DummyModel, default: DummyModel()
Quantum dynamics model implementing the driver interface.
sockets_prefix : str, default: ``"/tmp/socketmxl_"``
Prefix for UNIX domain socket paths (ignored for TCP/IP).
Notes
-----
Implements a simple message protocol with headers such as ``STATUS``,
``INIT``, ``POSDATA``/``FIELDDATA``, ``GETFORCE``/``GETSOURCE``, and ``STOP``.
"""
if unix:
sock = socket.socket(socket.AF_UNIX)
sock.connect(sockets_prefix + address)
else:
sock = _connect_tcp_with_retry(address, port, timeout)
initialized = False
have_result = False
pending_amp = None
additional_data = {}
dt_au = 0.0
molid = None
while True:
try:
msg = _recv_msg(sock)
except Exception:
# Treat EOF/timeouts during normal shutdown as clean exit
break
if msg == STATUS:
# Server is polling; we must reply with our state.
if not initialized:
_send_msg(sock, NEEDINIT)
elif have_result:
_send_msg(sock, HAVEDATA)
else:
_send_msg(sock, READY)
elif msg == INIT:
# Server sends INIT after we replied NEEDINIT
molid = _recv_int(sock)
init_json = json.loads(_recv_bytes(sock).decode("utf-8") or "{}")
dt_au = float(init_json.get("dt_au", 0.0))
print("[initialization] Time step in atomic units:", dt_au)
print("[initialization] Assigned a molecular ID:", molid)
driver.initialize(dt_au, molid)
initialized = True
print("[initialization] Finished initialization for molecular ID:", molid)
elif msg == FIELDDATA or msg == b"POSDATA":
# One step of data from server: treat "positions" as the E-field vector in a.u.
# This is to mirror i-pi's existing socket interface.
cell, icell, xyz = _recv_posdata(sock)
# effective [Ex, Ey, Ez] (a.u.) for this molecule
E = xyz[0]
# Stage the step (no commit)
driver.stage_step(E)
have_result = True
elif msg == GETSOURCE or msg == b"GETFORCE":
# Server asks us to return the result for this step
if not driver.have_result():
# it means the driver code was terminated during driver.propagate() and driver.calc_amp_vector()
# one way is to be defensive: return zero if we somehow got here without a computed result
pending_amp = np.zeros(3, float)
else:
pending_amp = driver.commit_step()
additional_data = driver.append_additional_data()
_send_force_ready(
sock,
energy_ha=0.0,
forces_Nx3_ha_per_bohr=pending_amp.reshape(1, 3),
virial_3x3_ha=np.zeros((3, 3)),
more=json.dumps(
additional_data,
ensure_ascii=False,
separators=(",", ":"),
sort_keys=True,
).encode("utf-8"),
)
have_result = False
pending_amp = None
elif msg == STOP:
# Acknowledge and leave gracefully
try:
_send_msg(sock, BYE)
finally:
print("Received STOP, exiting")
break
else:
raise RuntimeError(f"Unexpected header: {msg!r}")
[docs]
def mxl_driver_main():
"""
Parse CLI arguments and start the MaxwellLink socket driver.
Notes
-----
Constructs the selected model via ``__drivers__`` using the ``--model`` and
``--param`` options, then calls ``run_driver(...)``.
"""
parser = argparse.ArgumentParser(description=description)
parser.add_argument(
"-u",
"--unix",
action="store_true",
default=False,
help="Use a UNIX domain socket.",
)
parser.add_argument(
"-a",
"--address",
type=str,
default="localhost",
help="Host name (for INET sockets) or name of the UNIX domain socket to connect to.",
)
parser.add_argument(
"-S",
"--sockets_prefix",
type=str,
default="/tmp/socketmxl_",
help="Prefix used for the unix domain sockets. Ignored when using TCP/IP sockets.",
)
parser.add_argument(
"-p",
"--port",
type=int,
default=31415,
help="TCP/IP port number. Ignored when using UNIX domain sockets.",
)
parser.add_argument(
"-m",
"--model",
type=str,
default="dummy",
choices=list(__drivers__.keys()),
help="""Type of molecular / material model for computing dipole moments under EM field.
""",
)
parser.add_argument(
"-o",
"--param",
type=str,
default="",
help="""Parameters required to run the driver. Comma-separated list of values
""",
)
parser.add_argument(
"-v",
"--verbose",
action="store_true",
default=False,
help="Verbose output.",
)
args = parser.parse_args()
driver_args, driver_kwargs = _read_args_kwargs(args.param)
if args.model in __drivers__:
try:
d_f = __drivers__[args.model](
*driver_args, verbose=args.verbose, **driver_kwargs
)
except ImportError:
# specific errors have already been triggered
raise
except Exception as err:
print(f"Error setting up molecular dynamics model {args.model}")
print(__drivers__[args.model].__doc__)
print("Error trace: ")
raise err
elif args.model == "dummy":
d_f = DummyModel(verbose=args.verbose)
else:
raise ValueError("Unsupported driver model ", args.model)
run_driver(
unix=args.unix,
address=args.address,
port=args.port,
driver=d_f,
sockets_prefix=args.sockets_prefix,
)
def _clean_env_for_subprocess():
"""
Return a copy of the environment with MPI-related variables removed.
Returns
-------
dict
Sanitized environment dictionary suitable for launching child processes.
"""
env = os.environ.copy()
# Nuke anything that makes a child think it's an MPI rank
prefixes = (
"PMI_",
"PMIX_",
"OMPI_",
"MPI_",
"MPICH_",
"I_MPI_",
"HYDRA_",
"SLURM_",
"FI_",
"UCX_",
"PSM2_",
"PMI",
)
for k in list(env.keys()):
for p in prefixes:
if k.startswith(p):
env.pop(k, None)
break
# Some MPIs set these exact names without a prefix
for k in ("PMI_FD", "PMI_PORT", "PMI_ID", "PMI_RANK", "PMI_SIZE"):
env.pop(k, None)
return env
[docs]
def launch_driver(
command='--model tls --port 31415 --param "omega=0.242, mu12=187, orientation=2, pe_initial=1e-4" --verbose',
sleep_time=0.5,
):
"""
Launch the driver as a background subprocess for local testing.
Parameters
----------
command : str, default: '--model tls --port 31415 --param "omega=0.242, mu12=187, orientation=2, pe_initial=1e-4" --verbose'
Command-line arguments passed to ``mxl_driver.py``.
sleep_time : float, default: 0.5
Time to sleep (seconds) after launch to allow initialization.
Returns
-------
subprocess.Popen or None
The process handle on the master rank, otherwise ``None``.
"""
if _am_master():
print(f"Launching driver with command: mxl_driver.py {command}")
# launch the external driver (client)
driver_argv = shlex.split(shutil.which("mxl_driver.py") + " " + command)
# Use a fresh, non-blocking subprocess; inherit env/stdio for easy debugging
proc = subprocess.Popen(driver_argv, env=_clean_env_for_subprocess())
time.sleep(sleep_time)
return proc
else:
return None
[docs]
def terminate_driver(proc, timeout=2.0):
"""
Terminate a driver process launched by ``launch_driver``.
Parameters
----------
proc : subprocess.Popen or None
Process handle to terminate.
timeout : float, default: 2.0
Seconds to wait for graceful shutdown before escalating.
"""
if proc is not None and _am_master():
# Give it a moment to shut down naturally after the sim closes the socket
try:
proc.wait(timeout=timeout)
except subprocess.TimeoutExpired:
proc.terminate()
print("Driver did not exit cleanly, sent terminate signal")
try:
proc.wait(timeout=timeout)
except subprocess.TimeoutExpired:
proc.kill()
print("Driver did not terminate, sent kill signal")
if __name__ == "__main__":
mxl_driver_main()