Worker¶

aimmd.Worker is the atomic execution unit. A worker runs exactly one task at a time — shoot (committor-guided two-way shooting), free (unbiased simulation from a state), train (fit the committor model and refresh the adaptive sampling state), or kinetics_convergence — handling signal-driven cooperative termination, resource binding, and log redirection.

class aimmd.Worker(params, directory='.', localid=0, cpus_per_task='skip', gpus_per_task='skip', log_file='stdout', walltime=inf, nsteps=inf, nframes=inf, termination_timeout=60.0)[source]¶

Bases: WorkerHelpers, WorkerMagic, WorkerProperties, WorkerRun, WorkerSimulate, WorkerFree, WorkerShoot, WorkerTrain

Execute AIMMD tasks in an isolated worker process.

A Worker is the atomic execution unit used by AIMMD on both local machines and clusters. It is designed to run as a dedicated process that:

installs SIGINT/SIGTERM handlers and reacts cooperatively to termination,
optionally binds CPU/GPU resources based on a localid and policy,
redirects stdout/stderr to a per-worker log file when requested,
executes exactly one task at a time via run().

Tasks The canonical task names dispatched by run() are:

'shoot': committor-guided path sampling (core enhanced sampling),
'free': free simulations around a chosen state,
'train': update the committor model and adaptive sampling state,
'kinetics_convergence': retrain on growing sub-samples of the path ensemble and report how rate estimates converge with training-set size.

:param The constructor is aliased to WorkerHelpers._init(). See that method: :param for the full signature and meaning of worker initialization arguments: :param (params/directory/localid: :param resource policies: :param stop conditions: :param logging).:

Notes

Workers intentionally mutate process-global state (signal handlers, stdout/stderr redirection). This is safe because workers are meant to run as isolated subprocesses.
Cache state is cleared per task (via WorkerRun) to avoid leaking readers/arrays across tasks or directories.

Initialize a worker process.

The worker process runs one AIMMD task at a time (free simulation, shooting simulation, NN training), typically under a batch scheduler where CPU/GPU resources are allocated per worker.

This initializer:

stores core worker attributes (parameters, directory, local ID, stop-condition thresholds),
installs SIGTERM/SIGINT handlers that record a pending termination request in termination_signal,
wires logging through log_file (usually a property on the concrete worker that redirects stdout/stderr and/or opens a file).

Parameters:: params – Either an already-instantiated Params, or a

path to a parameters file that can be loaded into Params. :type params: str or aimmd.params.Params :param directory: Working directory for the worker. Stored in both directory and _directory. Default is '.'. :type directory: str, optional :param localid: Local worker index, used for deterministic resource binding. Default is 0. :type localid: int, optional :param cpus_per_task: CPU allocation policy for this worker:

'skip': do not explicitly bind; only report availability.
'share': divide available CPUs among workers.
'all': bind all available CPUs to this worker.
int: bind exactly this many CPUs (policy interpreted by

bind_resources()).

Default is 'skip'. :type cpus_per_task: {‘skip’, ‘share’, ‘all’} or int, optional :param gpus_per_task: GPU allocation policy, analogous to cpus_per_task. Default is 'skip'. :type gpus_per_task: {‘skip’, ‘share’, ‘all’} or int, optional :param log_file: Logging target. Assigned via log_file, which is expected to be implemented by the concrete worker class. If 'stdout', output is left on the original stdout/stderr. Default is 'stdout'. :type log_file: {‘stdout’} or str or file-like, optional :param walltime: Maximum allowed walltime in seconds for the current task before the worker should stop. Default is inf (no limit). :type walltime: float, optional :param nsteps: Maximum number of steps before stopping. Default is inf. :type nsteps: float, optional :param nframes: Maximum number of frames before stopping. Default is inf. :type nframes: float, optional :param termination_timeout: Grace period (seconds) used by higher-level logic to allow a task to terminate cleanly after a stop/termination request. This mixin stores the value but does not enforce it directly. Default is 60. seconds. :type termination_timeout: float, optional

Return type:: None

See also

aimmd.resources.bind_resources

Implements the CPU/GPU binding policy based on localid and the per-task configuration.

free(target_state=0, k=0, total=1, wait=False)¶

Public convenience wrapper for the free-simulation task.

Parameters:: target_state – Target state for free simulations.

If int, it is interpreted as an index into params.states.
If str, it is used directly as the state label.

Default is 0. :type target_state: int or str, optional :param k: Worker slot index within a group of total free simulations. Used to compute the starting trajectory number and to select initial paths deterministically when needed. Default is 0. :type k: int, optional :param total: Total number of concurrent free simulations across workers. The current worker advances trajectory indices by total so that each worker writes a disjoint subsequence of trajectory names. Default is 1. :type total: int, optional :param wait: If True and params.nbins > 1, wait until the network and current bins/densities can be loaded before starting. Default is False. :type wait: bool, optional

Returns:: Whatever Worker.run() returns for the 'free' task.
Return type:: object

property initial_paths¶

Initial paths available to the worker.

The worker expects initial paths to live under a folder named according to the sorted end states, e.g. initial('A','B'). This property first attempts to load existing paths from {directory}/initial{states}/*. If none are found, it lazily initializes them via Launcher in _directory and tries again.

Returns:: Ensemble containing all initial paths located (or created) for the

current end-state set. :rtype: aimmd.pathensemble.PathEnsemble

Notes

The state tuple is obtained from params.sorted_states and is embedded directly into the folder name.
The lazy creation branch imports Launcher locally to reduce import-time coupling and avoid circular imports.

kinetics_convergence(fractions=None, save_file='kinetics_convergence.npy', network_save_pattern='{directory}/network{states}.kcv{fraction_pct:03d}.h5', **kwargs)¶

Assess how kinetics estimates converge with growing training-set size.

For each requested fraction the path ensemble is sub-sampled (see Fraction sampling below), the network is retrained from scratch on the sub-sample, and the AIMMD reweighting is used to estimate the rate constants k12 and k21. The results are collected, saved, and returned as a structured NumPy array so they can immediately be plotted or stored for later analysis.

This is a convenience wrapper around Worker.run() with task 'kinetics_convergence'; it therefore inherits all of the standard run-wrapper behaviour (directory switching, cache clearing, resource binding, etc.).

Parameters:: fractions – Fractions of the full path ensemble to include in each training

run. Values must lie in (0, 1]. The list is de-duplicated and sorted in ascending order before use. Default is [0.2, 0.4, 0.6, 0.8, 1.0] (five 20 %-increment steps). :type fractions: list of float, optional :param save_file: Path (relative to the worker directory) where the result array is written as a .npy file. Default is 'kinetics_convergence.npy'. :type save_file: str, optional :param network_save_pattern: Format string used to derive the filename for the network checkpoint saved after each fraction’s training run. The following placeholders are available:

{directory} — the worker directory (same as

Worker.directory).

{states} — the sorted state-label string, e.g. 'ARB'

(same as Params.sorted_states).

{fraction} — the fraction as a float, e.g. 0.2.
{fraction_pct} — the fraction as an integer percentage,

e.g. 20.

The default pattern '{directory}/network{states}.kcv{fraction_pct:03d}.h5' produces filenames such as run_example/networkARB.kcv020.h5.

Pass None to skip saving networks entirely. :type network_save_pattern: str or None, optional :param **kwargs: Additional keyword arguments forwarded to the fit routine (e.g. epochs, lr, batch_size). Stop-condition keys (walltime, nsteps, nframes) are consumed by the run wrapper and do not reach the fit routine.

Returns:

numpy.ndarray – Structured array with dtype

[('fraction', float), ('n_frames', float), ('k12', float), ('k21', float), ('k12_rw', float), ('k21_rw', float)] and one row per requested fraction. n_frames is the total frame count of the sub-sampled ensemble at that fraction (useful for converting fraction-of-training to physical sampling time). k12/k21 hold the uncorrected AIMMD rate estimates in units of [1/dt]; k12_rw/k21_rw hold the Tiwary-Parrinello bias-reweighted estimates (filled only when params.record_bias is True, left as nan otherwise). Entries for fractions where training failed are also left as nan.

Fraction sampling
—————–
Paths are sub-sampled *per source so that every shooting chain and*
every free-simulation trajectory contributes the same fraction to the
training set. Concretely
* For each shooting chain of length N, the first – max(1, round(N * fraction)) paths are used.
* For each free trajectory of length N (frames), the first – max(1, round(N * fraction)) frames are used.
Using the *first paths/frames preserves the temporal ordering of the*
Markov chains, which is the most natural sub-sample for a convergence
analysis.

Notes

The fit routine (params.fit) always resets the network parameters before training (via network.reset_parameters()), so each fraction is trained independently from a random initialisation.
The current network state is saved before the analysis and restored afterwards, leaving the worker in the same state it was in before the call.
Temporary *.kcv.npy cache files are written to disk during value computation and removed again before the method returns.
When params.record_bias is True, the per-frame bias cache is populated for the sub-sampled ensemble before reweighting, the reactive-region bias check (check_reactive_bias()) is applied, and the bias-reweighted rates are computed via compute_bias_corrections() and stored in the k12_rw / k21_rw fields of the result.

Examples

Basic usage after a completed run:

worker = aimmd.Worker(params, ‘run_example’) results = worker.kinetics_convergence() # saves run_example/networkARB.kcv020.h5, .kcv040.h5, … .kcv100.h5 print(results[‘fraction’], results[‘k12’], results[‘k21’])

Custom fractions, fewer training epochs, and a custom network naming:

results = worker.kinetics_convergence( fractions=[0.25, 0.5, 0.75, 1.0], network_save_pattern=( ‘checkpoints/network{states}_f{fraction:.2f}.h5’), epochs=100, )

Disable per-fraction network saving:

results = worker.kinetics_convergence(network_save_pattern=None)

See also

Worker.train

The standard training task.

property log_file¶

Current logging target.

The worker uses this property to route all print output by redirecting sys.stdout and sys.stderr:

If a file-like object is provided, it becomes the new stdout/stderr.
If a string is provided, it is treated as a filename relative to directory and opened in append mode (line-buffered).
If 'stdout' or None is provided, stdout/stderr are restored to their original streams captured at initialization.

Returns:: The current log target stored in _log_file. This is either

a file-like object or the original stdout stream. :rtype: object

property must_stop¶

Whether the current task should stop.

The worker stops if:

a termination signal was received (SIGTERM or SIGINT),
walltime exceeded: time.time() - _t0 >= walltime,
frame limit exceeded: _total_frames >= nframes,
step limit exceeded: _total_steps >= nsteps.

If a stop-condition threshold is exceeded, this property sets termination_signal to 2 (used as a SIGINT-like marker) and returns True.

Returns:

bool – True if the worker should stop, otherwise False.

Side Effects
————
termination_signal – Set to 2 when a stop-condition threshold triggers and no

termination signal was previously recorded.

Notes

This property assumes that the concrete worker maintains the counters _total_frames, _total_steps and the task start time _t0 (seconds since epoch).

run(task, *args, **kwargs)¶

Run a single worker task.

Parameters:: task – Task identifier. Supported values are:

'shoot': dispatches to _shoot()
'free': dispatches to _free()
'train': dispatches to _train()
'kinetics_convergence': dispatches to _kinetics_convergence()

Parameters:

*args – Positional arguments forwarded to the selected task method.
**kwargs – Keyword arguments forwarded to the selected task method. This method

also consumes stop-condition keys via _update_stop_condition() (typically nsteps, nframes, walltime), removing them from kwargs before dispatch.

Returns:

The return value of the selected task method.

Return type:

object

Raises:

TypeError – If task does not match a supported task name.
Exception – Re-raises any exception thrown by task execution after optional traceback reporting.
Side Effects –
------------ –

:raises - Changes the current working directory to params.parent for the: duration of the task. :raises - Clears global caches MDA_CACHE and: NPY_CACHE. :raises - Resets per-task counters (_t0, _total_steps,: _total_frames). :raises - Ensures cleanup in a finally block even on errors.:

shoot(target_state=1, k=0, sweep=False, sweep_target=inf, nchains_per_worker=1)¶

Public convenience wrapper for the shooting task.

Do (AI-enhanced) path sampling by two-way shooting in the chosen state.

Parameters:: target_state – Target state used to select shooting points and to name output

folders.

If int, interpreted as an index into params.states (a, r, b).
If str, interpreted as the state label directly.

The common convention is: a (0), r (1), b (2), but any label supported by params.states is accepted. Default is 1. :type target_state: int or str, optional :param k: Worker index used to disambiguate output folders (e.g., chainR{k}) and to offset cycling of initial paths. Default is 0. If nchains_per_worker > 1, it is just the first shooting chain associated with the worker, with the others following sequentially. :type k: int, optional :param sweep: If False, run committor-guided shooting with a selection pool (enhanced sampling in the reactive region).

If True, run sweep-mode shooting intended for committor validation: repeatedly shoot from a predefined set of frames (from the merged initial ensemble) to empirically estimate outcome probabilities (e.g. fraction reaching state 1 vs state 0). Frames are picked round-robin by global least-covered coverage shared across all sweep workers (not a fixed per-worker cycle), and the campaign stops on a global shot target (see sweep_target). Default is False. In sweep mode nchains_per_worker is ignored (each sweep worker manages a single sweep{t}{k} chain). :type sweep: bool, optional :param sweep_target: Sweep mode only: total number of committed shots, summed across all sweep{t}* worker folders, at which the campaign is complete. Each worker stops once the global committed count reaches this value (so an uneven set of workers still hits the intended total exactly, and a finished campaign can be extended just by raising the target and resubmitting). Default is inf (run until walltime). Ignored when sweep is False. :type sweep_target: float, optional :param nchains_per_worker: If > 1, the worker will manage more than one chain. A higher value of nchains_per_worker tends to regularize the training set and thus improve performance. If running only one shooting worker and selection_pool_size=1, nchains_per_worker=10 is recommended. Ignored when sweep is True. :type nchains_per_worker: int, optional, default = 1

Returns:: Whatever Worker.run() returns for the 'shoot' task.
Return type:: object

property total_frames¶

Total number of frames ingested/produced so far.

This wraps the internal tqdm counter _total_frames.

Returns:: Current frame counter. Returns 0 if the progress bar is not

initialized. :rtype: int

property total_steps¶

Total number of completed trajectory steps reported by this worker.

This is a convenience accessor around the internal tqdm counter _total_steps created/updated by _set_progress_bar().

Returns:: Current progress value. Returns 0 if the progress bar is not

initialized. :rtype: int

train(nrounds=1, keep_running=False, **kwargs)¶

Public convenience wrapper for the training task.

This method forwards to Worker.run() with task='train' and passes nrounds and keep_running as positional arguments.

Parameters:: nrounds – Number of training rounds to perform. A “round” corresponds to one

successful call to the fit routine (i.e., a round only counts if training returns at least one loss value). Default is 1. :type nrounds: int or float, optional :param keep_running: If False, the worker stops once nrounds successful rounds have completed. If True, the worker keeps monitoring the path ensemble and continues to refresh bins/densities when new frames appear, even after completing nrounds. Default is False. :type keep_running: bool, optional :param **kwargs: Additional keyword arguments forwarded to the fit routine via Worker.run() / _train(). Stop-condition keywords (typically walltime, nsteps, nframes) may be consumed by the run wrapper, depending on the worker implementation.

Returns:: Whatever Worker.run() returns for the 'train' task.
Return type:: object