AIMMD multi-system (multi-ligand) tutorial¶
This notebook shows how one AIMMD params file can drive several chemical
systems at once and train a single shared committor model that works for
either system. It is the multi-system counterpart of 1_toy_1d.ipynb and, like
it, runs end-to-end on the toy engine (no GROMACS / no GPU) in a couple of
minutes.
We set up two toy “ligands” with different atom counts (1 and 2 atoms) to
demonstrate the part that makes a shared model non-trivial: a fixed
descriptor_transform / atom_types encoding maps both systems into the same
network input space.
What we cover:
Two toy systems and their initial transition paths.
A single multi-system
params.py(multi_system=True).A shared-network campaign (
multi_system_share_network=True): one trainer, one network, balanced pooling of both systems.Inspecting the per-system run layout and the one shared network.
The separate-networks mode for contrast.
Per-system kinetics convergence with the shared model.
An in-state bias (OPES-style) run with per-system Tiwary-Parrinello bias-reweighted rates.
0. Imports and working directory¶
Everything is created in a fresh temporary directory so the notebook is self-contained and repeatable.
import os
import json
import tempfile
import numpy as np
import MDAnalysis as mda
import matplotlib.pyplot as plt
import aimmd
work = tempfile.mkdtemp(prefix="aimmd_multi_")
os.chdir(work)
print("working directory:", work)
/home/lichtinger/anaconda3/envs/new_aimmd/lib/python3.13/site-packages/tqdm/auto.py:21: TqdmWarning: IProgress not found. Please update jupyter and ipywidgets. See https://ipywidgets.readthedocs.io/en/stable/user_install.html
from .autonotebook import tqdm as notebook_tqdm
working directory: /tmp/aimmd_multi_2aieb_61
1. Two toy systems with different atom counts¶
Each “ligand” is a toy particle whose first coordinate sweeps 0 -> 10, which
is a clean A -> R -> B transition for the state function below. System s1
has 1 atom; system s2 has 2 atoms (the extra atom is inert and ignored
by the analysis functions). Writing different atom counts is the whole point:
the shared network must consume both.
def write_initial(fname, n_atoms, n_frames=50):
u = mda.Universe.empty(n_atoms, trajectory=True)
sweep = np.linspace(0.0, 10.0, n_frames)
with mda.Writer(fname, n_atoms) as w:
for x in sweep:
u.atoms.positions = np.column_stack([
np.full(n_atoms, x), np.full(n_atoms, 5.0),
np.full(n_atoms, 5.0)]).astype(np.float32)
w.write(u.atoms)
write_initial("s1.xtc", n_atoms=1)
write_initial("s2.xtc", n_atoms=2)
print("s1 atoms:", mda.Universe("s1.xtc").atoms.n_atoms)
print("s2 atoms:", mda.Universe("s2.xtc").atoms.n_atoms)
s1 atoms: 1
s2 atoms: 2
/home/lichtinger/anaconda3/envs/new_aimmd/lib/python3.13/site-packages/MDAnalysis/core/universe.py:1916: UserWarning: there is no reference attributes in this universe to guess types from
warnings.warn(str(e))
/home/lichtinger/anaconda3/envs/new_aimmd/lib/python3.13/site-packages/MDAnalysis/core/universe.py:1916: UserWarning: there is no reference attributes (elements, types, or names) in this universe to guess mass from
warnings.warn(str(e))
2. One multi-system params file¶
The key fields:
multi_system = Trueturns on multi-system mode.multi_system_share_network = Truetrains one network for both systems.system_ids, list-valuedtopology, and groupedinitial_pathsgive one entry per system.states_function(traj, system_id=...)uses thesystem_idkeyword to apply a per-ligand state boundary.descriptor_transform(coords, system_id=...)maps any atom count to a fixed width (here the first atom’s x), so oneLinear(1, ...)network consumes both systems. (For real graph networks you would instead pass a fixedatom_typestable; the mechanism is identical.)fitsimply forwards to the default trainer — its only change for multi-system is that it now accepts a list of PathEnsembles.
PARAMS = """
import numpy as np
import torch
from aimmd.network import fit as _fit
from aimmd.network.rescalable import Rescalable
engine = 'toy'
multi_system = True
multi_system_share_network = True
system_ids = ['s1', 's2']
topology = ['s1.xtc', 's2.xtc']
initial_paths = [['s1.xtc'], ['s2.xtc']]
free_overriding_states = 'all'
def toy_mdrun(ts):
for _ in range(50):
ts.positions = (ts.positions + .05 * np.random.normal()) % 10
def states_function(trajectory, system_id=None):
cut_a = 2.0 if system_id == 's1' else 2.5 # per-ligand A boundary
out = []
for frame in trajectory:
x = frame.positions[0, 0]
out.append('A' if x < cut_a else ('B' if x > 8.0 else 'R'))
return np.array(out, dtype='<U1')
def descriptor_transform(coordinates, system_id=None):
arr = np.asarray(coordinates).reshape(np.asarray(coordinates).shape[0], -1, 3)
return arr[:, :1, 0] # fixed width for any system
class Network(Rescalable):
def __init__(self):
super().__init__()
self.input = torch.nn.Linear(1, 16)
self.act = torch.nn.ReLU()
self.output = torch.nn.Linear(16, 1)
self.reset_parameters()
def forward(self, x):
return self.output(self.act(self.input(x[:, :1])))
def reset_parameters(self):
self.input.reset_parameters(); self.output.reset_parameters()
network = Network()
def fit(params, pathensemble, verbose=False, worker=None):
return _fit(params, pathensemble, nbins=0, state_bins='all', augment='no',
lr=1e-3, epochs=40, batch_size=128, stop=1e9,
in_memory=True, graphs=False, verbose=verbose, worker=worker)
"""
with open("params.py", "w") as fh:
fh.write(PARAMS)
print("wrote params.py")
wrote params.py
3. Load the params and inspect the multi-system plumbing¶
Params builds one MDAnalysis universe per system and keeps the initial paths
grouped per system.
params = aimmd.Params.load("params.py")
print("multi_system :", params.multi_system)
print("share network :", params.multi_system_share_network)
print("system_ids :", params.system_ids)
print("topology (per system) :", params.topology)
print("universe atom counts :",
{sid: params.universe_of(sid).atoms.n_atoms for sid in params.system_ids})
print("initial-path groups :", len(params.initial_paths))
Written full params with descriptions to 'params1.py'
multi_system : True
share network : True
system_ids : ['s1', 's2']
topology (per system) : ['s1.xtc', 's2.xtc']
universe atom counts : {'s1': 1, 's2': 2}
initial-path groups : 2
6. The shared model predicts the committor for both systems¶
We load the one shared network and evaluate the committor along the reaction
coordinate for each system. Because the systems have different atom counts,
this only works thanks to the shared descriptor_transform encoding.
import torch
from scipy.special import expit
params.update_network("run1") # loads run1/networkARB.h5 (shared)
xs = np.linspace(0, 10, 60)
plt.figure(figsize=(6, 4))
for sid, n_atoms in (("s1", 1), ("s2", 2)):
coords = np.zeros((len(xs), n_atoms * 3), dtype=np.float32)
coords[:, 0] = xs # atom-0 x = reaction coordinate
logit = params.values_function(coords, system_id=sid)
plt.plot(xs, expit(logit), "o-", ms=3, label=f"system {sid}")
plt.axhline(0.5, color="k", lw=0.7, ls="--")
plt.xlabel("reaction coordinate (atom-0 x)")
plt.ylabel("committor $p_B$")
plt.title("One shared network, two systems")
plt.legend(); plt.tight_layout(); plt.show()
7. Separate networks (for contrast)¶
With multi_system_share_network = False, each ligand trains its own network
in its own subfolder (one trainer per system; trainers_share_gpu decides
whether they share a GPU). Everything else — the single params file, the
subfolder layout, the per-system worker counts — is identical.
sep = PARAMS.replace("multi_system_share_network = True",
"multi_system_share_network = False")
os.makedirs("sep", exist_ok=True)
write_initial("sep/s1.xtc", 1); write_initial("sep/s2.xtc", 2)
with open("sep/params.py", "w") as fh:
fh.write(sep)
here = os.getcwd(); os.chdir("sep")
try:
aimmd.Launcher("params.py", "run1").run(n=1, n1=1, n2=1, nsteps=10, walltime=300)
print("per-system networks :",
{sid: os.path.isfile(f"run1/{sid}/networkARB.h5") for sid in ("s1", "s2")})
print("shared network at root :", os.path.isfile("run1/networkARB.h5"))
finally:
os.chdir(here)
Written full params with descriptions to 'params1.py'
Assigned parameters file 'params1.py'
+++ created 'run1'
+++ created 'run1/s1'
+++ created 'run1/s1/initialARB'
+++ saved 'run1/s1/initialARB/s1.xtc.trr' (from: 's1.xtc')
+++ saved 'run1/s1/initialARB/s1.xtc.trr.states.npy'
+++ created run1/s1/freeA
+++ created run1/s1/freeB
+++ created run1/s1/chainR0
+++ created 'run1/s2'
+++ created 'run1/s2/initialARB'
+++ saved 'run1/s2/initialARB/s2.xtc.trr' (from: 's2.xtc')
+++ saved 'run1/s2/initialARB/s2.xtc.trr.states.npy'
+++ created run1/s2/freeA
+++ created run1/s2/freeB
+++ created run1/s2/chainR0
/home/lichtinger/anaconda3/envs/new_aimmd/lib/python3.13/site-packages/MDAnalysis/core/universe.py:1916: UserWarning: there is no reference attributes in this universe to guess types from
warnings.warn(str(e))
/home/lichtinger/anaconda3/envs/new_aimmd/lib/python3.13/site-packages/MDAnalysis/core/universe.py:1916: UserWarning: there is no reference attributes (elements, types, or names) in this universe to guess mass from
warnings.warn(str(e))
/home/lichtinger/anaconda3/envs/new_aimmd/lib/python3.13/site-packages/MDAnalysis/core/universe.py:1916: UserWarning: there is no reference attributes in this universe to guess types from
warnings.warn(str(e))
/home/lichtinger/anaconda3/envs/new_aimmd/lib/python3.13/site-packages/MDAnalysis/core/universe.py:1916: UserWarning: there is no reference attributes (elements, types, or names) in this universe to guess mass from
warnings.warn(str(e))
/home/lichtinger/anaconda3/envs/new_aimmd/lib/python3.13/site-packages/MDAnalysis/core/universe.py:1916: UserWarning: there is no reference attributes in this universe to guess types from
warnings.warn(str(e))
/home/lichtinger/anaconda3/envs/new_aimmd/lib/python3.13/site-packages/MDAnalysis/core/universe.py:1916: UserWarning: there is no reference attributes (elements, types, or names) in this universe to guess mass from
warnings.warn(str(e))
/home/lichtinger/anaconda3/envs/new_aimmd/lib/python3.13/site-packages/MDAnalysis/core/universe.py:1916: UserWarning: there is no reference attributes in this universe to guess types from
warnings.warn(str(e))
/home/lichtinger/anaconda3/envs/new_aimmd/lib/python3.13/site-packages/MDAnalysis/core/universe.py:1916: UserWarning: there is no reference attributes (elements, types, or names) in this universe to guess mass from
warnings.warn(str(e))
/home/lichtinger/anaconda3/envs/new_aimmd/lib/python3.13/site-packages/MDAnalysis/core/universe.py:1916: UserWarning: there is no reference attributes in this universe to guess types from
warnings.warn(str(e))
/home/lichtinger/anaconda3/envs/new_aimmd/lib/python3.13/site-packages/MDAnalysis/core/universe.py:1916: UserWarning: there is no reference attributes (elements, types, or names) in this universe to guess mass from
warnings.warn(str(e))
/home/lichtinger/anaconda3/envs/new_aimmd/lib/python3.13/site-packages/MDAnalysis/core/universe.py:1916: UserWarning: there is no reference attributes in this universe to guess types from
warnings.warn(str(e))
/home/lichtinger/anaconda3/envs/new_aimmd/lib/python3.13/site-packages/MDAnalysis/core/universe.py:1916: UserWarning: there is no reference attributes (elements, types, or names) in this universe to guess mass from
warnings.warn(str(e))
/home/lichtinger/anaconda3/envs/new_aimmd/lib/python3.13/site-packages/MDAnalysis/core/universe.py:1916: UserWarning: there is no reference attributes in this universe to guess types from
warnings.warn(str(e))
/home/lichtinger/anaconda3/envs/new_aimmd/lib/python3.13/site-packages/MDAnalysis/core/universe.py:1916: UserWarning: there is no reference attributes (elements, types, or names) in this universe to guess mass from
warnings.warn(str(e))
/home/lichtinger/anaconda3/envs/new_aimmd/lib/python3.13/site-packages/MDAnalysis/core/universe.py:1916: UserWarning: there is no reference attributes in this universe to guess types from
warnings.warn(str(e))
/home/lichtinger/anaconda3/envs/new_aimmd/lib/python3.13/site-packages/MDAnalysis/core/universe.py:1916: UserWarning: there is no reference attributes (elements, types, or names) in this universe to guess mass from
warnings.warn(str(e))
/home/lichtinger/anaconda3/envs/new_aimmd/lib/python3.13/site-packages/MDAnalysis/core/universe.py:1916: UserWarning: there is no reference attributes in this universe to guess types from
warnings.warn(str(e))
/home/lichtinger/anaconda3/envs/new_aimmd/lib/python3.13/site-packages/MDAnalysis/core/universe.py:1916: UserWarning: there is no reference attributes (elements, types, or names) in this universe to guess mass from
warnings.warn(str(e))
[PID 687250, TID 140155376600896: "run1/s2" ARB trainer] starting (Fri Jun 19 11:42:09 2026)
... args: ('params1.py', 'run1/s2', np.int64(3), np.int64(3), 1, 'trainARB.log', 300.0, np.float64(10.0), np.float64(inf), 59.0, 'train', np.float64(inf), True)
[PID 687250, TID 140155376600896: "run1/s2" ARB trainer] exited correctly
[PID 687242, TID 139653970908992: "run1/s1" freeA (worker0)] starting (Fri Jun 19 11:42:09 2026)
... args: ('params1.py', 'run1/s1', np.int64(0), np.int64(3), 1, 'freeA/worker0.log', inf, inf, inf, 59.0, 'free', 'A', 0, 1, np.True_)
[PID 687242, TID 139653970908992: "run1/s1" freeA (worker0)] exited correctly
[PID 687246, TID 140131098457920: "run1/s2" freeA (worker0)] starting (Fri Jun 19 11:42:09 2026)
... args: ('params1.py', 'run1/s2', np.int64(4), np.int64(3), 1, 'freeA/worker0.log', inf, inf, inf, 59.0, 'free', 'A', 0, 1, np.True_)
[PID 687246, TID 140131098457920: "run1/s2" freeA (worker0)] exited correctly
[PID 687245, TID 140556642162496: "run1/s1" ARB trainer] starting (Fri Jun 19 11:42:09 2026)
... args: ('params1.py', 'run1/s1', np.int64(3), np.int64(3), 1, 'trainARB.log', 300.0, np.float64(10.0), np.float64(inf), 59.0, 'train', np.float64(inf), True)
[PID 687245, TID 140556642162496: "run1/s1" ARB trainer] exited correctly
[PID 687244, TID 139696671479616: "run1/s1" chainR0] starting (Fri Jun 19 11:42:09 2026)
... args: ('params1.py', 'run1/s1', np.int64(2), np.int64(3), 1, 'chainR0/worker.log', inf, inf, inf, 59.0, 'shoot', 'R', 0, False)
[PID 687244, TID 139696671479616: "run1/s1" chainR0] exited correctly
[PID 687248, TID 139690070644544: "run1/s2" freeB (worker0)] starting (Fri Jun 19 11:42:09 2026)
... args: ('params1.py', 'run1/s2', np.int64(5), np.int64(3), 1, 'freeB/worker0.log', inf, inf, inf, 59.0, 'free', 'B', 0, 1, np.True_)
[PID 687248, TID 139690070644544: "run1/s2" freeB (worker0)] exited correctly
[PID 687243, TID 139869902411584: "run1/s1" freeB (worker0)] starting (Fri Jun 19 11:42:09 2026)
... args: ('params1.py', 'run1/s1', np.int64(1), np.int64(3), 1, 'freeB/worker0.log', inf, inf, inf, 59.0, 'free', 'B', 0, 1, np.True_)
[PID 687243, TID 139869902411584: "run1/s1" freeB (worker0)] exited correctly
[PID 687249, TID 139816219973440: "run1/s2" chainR0] starting (Fri Jun 19 11:42:09 2026)
... args: ('params1.py', 'run1/s2', np.int64(6), np.int64(3), 1, 'chainR0/worker.log', inf, inf, inf, 59.0, 'shoot', 'R', 0, False)
[PID 687249, TID 139816219973440: "run1/s2" chainR0] exited correctly
per-system networks : {'s1': True, 's2': True}
shared network at root : False
Summary¶
With a single params file and the multi_system switches we trained one
committor network that serves two systems with different atom counts, balanced
so neither dominates — contrasted it with the separate-network mode, and ran an
in-state-biased variant whose per-system Tiwary-Parrinello reweighting
recovers unbiased rates. The same configuration scales to real systems (e.g. a
graph/PaiNN network with a fixed atom_types table over many ligands), to OPES
biases via the per-system gmx_mdrun -plumed string, and to live HPC campaigns
via Launcher.create_job.