Utilities

Lower-level helper subpackages used throughout AIMMD. Most users never call these directly, but they are part of the public surface.

Resources

CPU/GPU introspection and process resource binding.

CPU visibility utilities.

This module provides small helpers to query which CPUs are available to the current process.

The returned CPU set is meant to reflect process-level restrictions imposed by:

• Linux CPU affinity (e.g., Slurm cgroups, taskset, sched_setaffinity),

• container/cgroup cpuset limitations.

If affinity information is not available, the implementation falls back to the number of physical cores reported by psutil.

Notes

• On Linux, os.sched_getaffinity() is the preferred source because it reflects the scheduler-visible CPU mask for the current PID.

• The fallback uses psutil.cpu_count(logical=False) (physical cores).

aimmd.resources.cpu.get_available_cpus()

Return the CPU ids currently available to the process.

Returns:

Sorted list of CPU ids. On Linux, this is derived from

os.sched_getaffinity(0) and therefore respects cpuset/affinity restrictions. If that fails, returns range(psutil.cpu_count(logical=False)). :rtype: list[int]

Notes

This function is intentionally conservative: it reports the CPUs the process can actually run on, not necessarily all CPUs installed on the machine.

aimmd.resources.cpu.get_num_cpus()

Return the number of CPUs currently available to the process.

Returns:

len(get_available_cpus()).

Return type:

int

GPU visibility utilities.

This module provides helpers to query which GPUs are visible to the current process.

Two notions are used:

• total GPUs (as reported by torch)

• available GPUs (as restricted by CUDA_VISIBLE_DEVICES)

The functions are designed for lightweight runtime checks in worker code and avoid introducing extra dependencies beyond torch.

Notes

• get_num_gpus() uses torch.cuda.device_count().

• get_available_gpus() inspects CUDA_VISIBLE_DEVICES:

• if unset, GPUs 0..N-1 are considered available

• if set to a comma-separated list, that list is interpreted as the

visible device ids

• When CUDA_VISIBLE_DEVICES is set to an empty string, no GPUs are available.

aimmd.resources.gpu.get_num_gpus()

Return the number of CUDA devices detected by torch.

Returns:

Number of GPUs returned by torch.cuda.device_count().

Returns 0 if torch fails to query CUDA devices. :rtype: int

Notes

Torch can emit warnings when probing device availability (driver mismatch, missing runtime, etc.). Warnings are suppressed within this call to keep worker logs clean.

aimmd.resources.gpu.get_available_gpus()

Return the GPU ids currently available to the process.

Returns:

If CUDA_VISIBLE_DEVICES is unset, returns list(range(get_num_gpus())).

If it is set, parses the comma-separated list and returns the integer ids. Empty entries are ignored. :rtype: list[int]

Notes

This function reports the ids visible inside the current process environment. If a scheduler (e.g., Slurm) sets CUDA_VISIBLE_DEVICES, this reflects the scheduler allocation.

Worker resource binding.

This module defines bind_resources(), a utility used by AIMMD workers to bind a process to a subset of CPUs and GPUs based on a per-node local index.

The goal is to make local multi-worker execution predictable:

• on workstations: map worker 0..K-1 to disjoint slices of CPU cores and GPUs

• on HPC with Slurm: respect allocations that are already restricted by cpusets and CUDA_VISIBLE_DEVICES, while still allowing explicit binding if requested. However, the default mode in this case is ‘skip’.

Binding model bind_resources operates on the resources visible to the current process:

• CPUs are obtained from aimmd.resources.cpu.get_available_cpus()

• GPUs are obtained from aimmd.resources.gpu.get_available_gpus()

The selection rule is:

• if resources_per_task == ‘all’:

use all visible resources

• if resources_per_task == ‘skip’:

do not bind, only report visible resources

• if resources_per_task is an integer:

assign a contiguous slice starting at localid * resources_per_task; if this exceeds the visible set, wrap around modulo the number available (and print a note about possible oversubscription)

Actual binding side effects CPU binding (when enabled) sets:

• OMP_NUM_THREADS, MKL_NUM_THREADS, OPENBLAS_NUM_THREADS

• torch.set_num_threads(...)

• attempts to set process CPU affinity via psutil.Process().cpu_affinity(...)

GPU binding (when enabled) sets:

• CUDA_VISIBLE_DEVICES

• GPU_DEVICE_ORDINAL (used by some runtimes, including ROCm paths in torch)

No other resource managers are touched.

Notes

This function is intentionally pragmatic:

• it tries to bind resources when asked,

• but if affinity operations fail it continues and only emits a warning.

aimmd.resources.binding.bind_resources(localid, cpus_per_task='skip', gpus_per_task='skip')

Bind the current process to CPU/GPU resources for the given worker local id.

Parameters:

localid – Index of the worker on the current node. This is used to choose which

subset of resources the worker should take when cpus_per_task and/or gpus_per_task are integers. :type localid: int :param cpus_per_task: CPU allocation policy for this worker.

• ‘all’ : bind to all visible CPUs

• ‘skip’: do not bind; only report visible CPUs

• int : number of CPUs to allocate to this worker

Parameters:

gpus_per_task – GPU allocation policy for this worker.

• ‘all’ : bind to all visible GPUs

• ‘skip’: do not bind; only report visible GPUs

• 0 : bind to no GPU (unset visibility)

• int : number of GPUs to allocate to this worker

Notes

The selected resource lists are computed from the visible resources of the current process. This means scheduler restrictions (cpusets, CUDA_VISIBLE_DEVICES) are honored automatically.

Execution

Streaming subprocess execution and simple thread/process executors.

Utility functions for robust execution and task wrapping.

This module provides:

• execute_command: run a shell command via subprocess.Popen while: - streaming stdout in (near) real time, - periodically checking a Python-side stop condition, - enforcing a walltime limit, - handling graceful termination and forced kill, - optionally raising on non-zero exit codes.

• target_wrapper: wrap arbitrary Python callables to: - print a consistent start/exit banner including PID and thread ID, - surface error information in logs, - standardize process/thread task output for debugging.

Design notes

• execute_command is designed for long-running external commands (e.g., GROMACS) where you want continuous logging and cooperative stopping.

• Output streaming is implemented with select.select on the subprocess stdout file descriptor (POSIX-friendly approach).

• Process termination targets the process group via os.killpg, enabled by creating the subprocess in a new session (preexec_fn=os.setsid).

aimmd.execute.utils.execute_command(command, stop_condition=<function <lambda>>, walltime=inf, termination_timeout=60.0, raise_if_failure=True, log_file='stdout')

Execute a shell command while streaming output and polling a stop condition.

The command is executed in a dedicated subprocess (via subprocess.Popen) and its stdout/stderr stream is captured and forwarded in near real time to either:

• sys.stdout (default), or

• a file-like object supporting .write(…).

The loop repeatedly:

• reads and prints available output (non-blocking),

• checks whether the process has exited,

• evaluates stop_condition(),

• enforces a maximum walltime.

Parameters:

• command (str) – Command string executed with shell=True.

• stop_condition – Zero-argument callable returning bool.

If True, the subprocess is terminated and the function returns 0. :type stop_condition: callable, optional :param walltime: Maximum allowed runtime in seconds. If exceeded, the subprocess is terminated and the function returns 0. Default is math.inf (no walltime). :type walltime: float, optional :param termination_timeout: Seconds to wait after sending SIGTERM before escalating to SIGKILL. :type termination_timeout: float, optional :param raise_if_failure: If True, a non-zero subprocess exit code raises RuntimeError. If False, the exit code is returned. :type raise_if_failure: bool, optional :param log_file: Where to write logs/output.

• If ‘stdout’, output is written to sys.stdout.

• If a file-like object, .write(text) is used.

Returns:

• The subprocess exit code if it completes normally.

• 0 if stopped by stop_condition, walltime, or KeyboardInterrupt.

• 1 if an exception occurs and raise_if_failure is False.

Return type:

int

Raises:

RuntimeError – If the subprocess exits with non-zero status and raise_if_failure is True. Also raised if an internal exception occurs and raise_if_failure is True.

Notes

• The subprocess is launched in a new process group (os.setsid), allowing group termination via os.killpg.

• Output is read via non-blocking polling (select.select).

• On stop/interrupt/exception, remaining buffered output is read and printed.

aimmd.execute.utils.target_wrapper(target, name, *args, **kwargs)

Wrap a callable to provide standardized logging and error reporting.

This wrapper is used by aimmd.execute.base.TaskExecutor when launching targets in threads/processes (especially important under multiprocessing “spawn”, where child initialization can be opaque).

Parameters:

• target (callable) – The function to call.

• name (str) – Human-readable identifier for log messages.

• *args – Positional arguments forwarded to target.

• **kwargs – Keyword arguments forwarded to target.

Return type:

None

Raises:

Exception – Re-raises any exception raised by target after recording a message.

Notes

• Prints PID and thread ID to help debugging concurrency issues.

• If an exception occurs, the exception message is included in the final line.

Process-based task executor implementation.

This module defines ProcessExecutor, a concrete implementation of aimmd.execute.base.TaskExecutor backed by multiprocessing.Process objects (spawn context).

The executor provides:

• Process construction via the global ctx spawn context

• Graceful termination (terminate)

• Forced termination (also terminate; Python processes do not have a distinct “kill” API in the standard library)

• Resource cleanup via join

Notes

• This executor relies on multiprocessing.get_context(‘spawn’) from aimmd.execute.base (imported as ctx).

• The _closed state is backend-specific and uses a private attribute of multiprocessing.Process. This is practical but not part of the public API; behavior may vary across Python versions.

class aimmd.execute.processes.ProcessExecutor

Bases: TaskExecutor

Task executor using multiprocessing.Process.

Instances of this executor manage a list of processes created with a spawn start method (safer across platforms and with CUDA/GPU contexts).

__task_name__

Human-readable label for the backend task type.

Type:

str

Initialize an empty TaskExecutor.

Thread-based task executor implementation.

This module defines ThreadExecutor, a concrete implementation of aimmd.execute.base.TaskExecutor backed by threading.Thread.

Threads are created as daemon threads, meaning:

• they will not prevent Python from exiting if only daemon threads remain,

• they are appropriate for background helper tasks,

• but they are not suitable when you require guaranteed cleanup/finalization.

Notes

• Python does not provide a safe, general mechanism to “terminate” threads. Therefore, TaskExecutor.stop() relies on backend hooks; for threads, _terminate/_kill are intentionally not implemented here. Consumers should design thread targets to exit cooperatively (e.g., by checking a shared stop condition/event).

class aimmd.execute.threads.ThreadExecutor

Bases: TaskExecutor

Task executor using threading.Thread.

__task_name__

Human-readable label for the backend task type.

Type:

str

Initialize an empty TaskExecutor.

Caching

Robust, lock-protected caches for trajectory readers and NumPy arrays.

Safe .npy file utilities and a lightweight .npy reader cache.

This module has two responsibilities:

1. Safe on-disk storage of NumPy arrays - save_npy: write an array atomically (temp file + replace),

protected by a lock.

• load_npy: read an array under the same lock.

• update_npy: update selected rows of an existing .npy file in place,

also protected by a lock.

2. Faster repeated reads of .npy files - NpyReaderCache: a small in-memory cache for arrays loaded from .npy

files. It avoids repeated np.load on the same file.

Locking model All operations use a per-file lock located next to the .npy file:

• Data: <folder>/<name>.npy

• Lock: <folder>/.<name>.lock

The lock prevents readers from loading a file while it is being written or updated.

Atomic full writes (save_npy) save_npy writes to a temporary file in the same directory and then replaces the target file using os.replace. On typical POSIX filesystems, this makes the write appear atomic to readers.

In-place updates (update_npy) update_npy is optimized for the case where you want to overwrite only some rows (axis 0) without loading the full array into memory.

Important constraints (by design) update_npy is intentionally narrow and assumes:

• The .npy header is treated as a fixed 128-byte region.

• Array data are assumed to start at byte offset 128.

• Only simple dtypes are supported (no structured dtypes).

• Trailing dimensions (shape[1:]) must match between file and update data.

If these constraints are violated, update_npy raises a RuntimeError.

Rationale AIMMD commonly produces large arrays that are:

• appended/updated incrementally by one process,

• read frequently by others for monitoring or downstream computation.

This module provides robust semantics for that workflow.

aimmd.cache.npy.save_npy(fname, array, timeout=10.0)

Save an array to a .npy file safely (lock + atomic replace).

Parameters:

• fname (str) – Target .npy filename.

• array (numpy.ndarray) – Array to write.

• timeout (float) – FileLock timeout.

• Effects (Side)

• ------------

• fname. (- Writes a temporary hidden file next to)

• fname.

• os.replace. (- Replaces fname atomically via)

Notes

This function is intended for complete rewrites of a file. For sparse row updates, use update_npy().

To avoid I/O error, it tries saving at most ten time, waiting 0.1 seconds before each successive attempt.

aimmd.cache.npy.load_npy(fname, timeout=10.0)

Load an array from a .npy file safely (under a lock).

Parameters:

• fname (str) – .npy file to read.

• timeout (float, default 5.0) – Maximum time (seconds) to wait for the file lock.

Returns:

Loaded array, or None if:

• the file does not exist,

• the lock cannot be acquired,

• the file cannot be loaded for any reason.

Return type:

numpy.ndarray or None

Notes

This function is intentionally permissive and returns None on failure.

aimmd.cache.npy.update_npy(fname, data, indices, timeout=10.0)

Update selected rows (axis 0) of a .npy file in place.

This function is optimized for high-performance computing environments where network filesystems (NFS/Lustre) may exhibit transient I/O latencies. It combines file locking, atomic binary writes, and a retry mechanism to prevent race conditions and ‘Errno 5’ (I/O) errors from crashing long-running simulations.

Parameters:

• fname (str) – Path to the target .npy file.

• data – The data to be written. Must match the trailing shape of the existing

file. :type data: array-like :param indices: The row index (or indices) along axis 0 where data should be written. :type indices: int or array-like of int :param timeout: FileLock timeout. :type timeout: float

Raises:

• RuntimeError –

  • If the array contains structured dtypes (not supported). - If the stored dtype in the file does not match the input data dtype. - If the trailing shape (axis 1+) does not match the file.

• OSError –

  • If the filesystem remains unresponsive after the maximum number of retries.

• Implementation Details –

• ---------------------- –

• 1. Header Assumption – Assumes a fixed 128-byte header for fast seeking.:

• 2. Atomic Growth – Uses file.truncate to grow the file if indices exceed: size.

• 3. Resilience – Wraps binary I/O in a 5-attempt retry loop with exponential: backoff.

• 4. Data Integrity – Uses os.fsync to ensure bits are physically committed: to disk.

class aimmd.cache.npy.NpyReaderCache

Bases: AbstractCache

In-memory cache for arrays loaded from .npy files.

Behavior

• Loads arrays using load_npy().

• Marks arrays as read-only before returning/caching them.

• Uses the base cache eviction mechanism when the heuristic size budget is exceeded.

max_size

Heuristic cache budget, set to half of the available system memory at import time. Leaving headroom for the network, MDA universes, NumPy temporaries, and the OS page cache is more important than maximizing cache occupancy: filling all available RAM with cached arrays pushes the process into swap or an OOM kill before eviction has a chance to run.

Type:

int

Notes

• This cache is process-local (not shared across processes).

• _size is overridden to charge cached arrays by their nbytes, so the budget is honored even for arrays returned by np.load (which have OWNDATA=False and would otherwise be undercounted by sys.getsizeof).

Initialize an empty cache instance.

Notes

• _cache preserves insertion order for FIFO eviction.

• total_size tracks a heuristic byte budget via sys.getsizeof.

max_size = 3223717888

MDAnalysis reader cache for robust trajectory opening.

AIMMD frequently needs to open trajectories repeatedly, potentially while they are being written by other processes. Some MDAnalysis formats maintain auxiliary offset and lock files (e.g., XTC/TRR offsets), which can become stale and lead to hangs or errors.

This module mitigates that by:

• removing MDAnalysis-generated offset/lock files before opening,

• opening with refresh_offsets=True,

• detecting how many frames are safely readable,

• returning only the safe prefix of the trajectory.

Integration This cache is constructed during aimmd._init.initialize() and stored as aimmd._config.MDA_CACHE for global reuse.

Notes

This cache is process-local. It does not coordinate between processes beyond filesystem-level effects (offset/lock file removal).

aimmd.cache.mda.Reader(*args, **kwargs)

aimmd.cache.mda.remove_offset_files(fname)

Remove MDAnalysis offset and lock files for a given trajectory file.

Parameters:

• fname (str) – Trajectory file path.

• Effects (Side)

• ------------

• present) (Deletes auxiliary files next to the trajectory (if)

• typically

• .<name>_offsets.npz (-)

• .<name>_offsets.lock (-)

Notes

The removal is performed in a loop until both files are absent. This preserves the current behavior under concurrent creation/deletion.

aimmd.cache.mda.count_safe_frames(reader)

Determine how many frames are safely readable from an MDAnalysis reader.

Parameters:

reader – An opened MDAnalysis reader supporting __len__ and random access

via reader[i]. :type reader: object

Returns:

The largest n such that frames [0 .. n-1] are readable.

Return type:

int

Raises:

RuntimeError – If no readable frame can be found.

Notes

Algorithm (preserved):

• Start from len(reader).

• Try to access the last frame.

• On (StopIteration, EOFError, OSError), decrement and retry.

class aimmd.cache.mda.MDAReaderCache

Bases: AbstractCache

Cache for MDAnalysis readers.

Behavior

• Validates file exists and extension is supported.

• Removes stale offset/lock files.

• Retries opening multiple times.

• Returns a reader sliced to only safe frames.

max_size

Size budget (currently a placeholder heuristic).

Type:

int

Notes

max_size is not a byte-accurate budget for readers; it is used only to limit how many objects remain in the cache (heuristically).

Initialize an empty cache instance.

Notes

• _cache preserves insertion order for FIFO eviction.

• total_size tracks a heuristic byte budget via sys.getsizeof.

max_size = 24000

Core utilities

General-purpose helper functions used across AIMMD.

This module intentionally collects small, dependency-light utilities that are reused across the package. These functions are designed to:

• stay independent of high-level AIMMD classes,

• be standalone and easy to test,

• avoid heavy imports of AIMMD internals,

• keep side effects explicit and localized.

The utilities here cover:

• indexing helpers for block-structured data,

• lightweight time formatting for logs,

• common list/array operations,

• boolean array inspection,

• filesystem convenience utilities,

• small simulation helpers (e.g., random velocity initialization),

• cache-related filename replacement,

• MDAnalysis trajectory helpers.

Design notes

• Functions are generally pure unless their purpose is I/O (e.g., remove, replace_in_cache) or they explicitly sample randomness (e.g., randomize_velocities).

• For performance and portability, implementations avoid complex dependencies.

• Some helpers assume conventions used throughout AIMMD (e.g., “offsets” as cumulative block boundaries).

Contents Indexing

• get_local_index: map a global index to (block_index, local_index) using offsets.

Time formatting

• now: current time string for logs.

List/array helpers

• cycle: rotate a list by n.

• concatenate: concatenate non-empty arrays.

• merge_ranges: merge overlapping integer ranges.

Boolean array inspection

• longest_true_segment: find the longest continuous portion which is true

Filesystem helpers

• remove: remove files matched by patterns (supports wildcards).

• unique_path: generate a non-existing filesystem path by incrementing a suffix.

Simulation helpers

• randomize_velocities: Maxwell-Boltzmann sampling + COM removal.

Cache utilities

• replace_in_cache: rename files and keep cache keys consistent.

Trajectory helpers (MDAnalysis)

• memory_reader_from_timesteps: pack timesteps into a MemoryReader (time info caveat).

• process_state: normalize a state label (e.g., ‘A’, ‘B’, ‘R’) or numeric indices.

• extend_array: pad an array to a minimum length (immutable result).

• extract_folder_and_name: split a path into (folder, basename).

• guess_masses: heuristic mass assignment for Martini beads.

MDAnalysis dependency note Some functions in the last section refer to Timestep and MemoryReader but do not import them in this module. This is intentional in the current code layout: the caller environment is expected to provide these names (typically via MDAnalysis imports). If you want this module to be self-contained, add explicit imports from MDAnalysis (but that would be a code change).

aimmd.core.utils.now()

Current local time formatted for log messages.

Returns:

Time formatted as "(Mon Feb 25 13:37:00 2026)" using time.ctime().

Return type:

str

Notes

• The output format is intentionally stable and human-readable, suitable for log files and HPC stdout.

• Timezone and locale are inherited from the running environment.

aimmd.core.utils.get_local_index(i, offsets, clip=False)

Convert a global index into (block_index, local_index) given cumulative offsets.

Parameters:

• i (int) – Global index to map.

• offsets – Cumulative offsets defining block boundaries.

Convention used by AIMMD:

• offsets[k] is the starting global index of block k,

• offsets[-1] is the first index after the last valid element

(i.e., a total length). :type offsets: array-like of int :param clip: If False (default), raise IndexError when i is outside bounds. If True, clip to the closest valid index. :type clip: bool, optional

Returns:

(block_index, local_index), where local_index is relative to the

start of the selected block. :rtype: tuple[int, int]

Raises:

IndexError – If clip is False and i is outside valid bounds.

Notes

This helper is used for block-structured trajectory storage, where a large logical index space is composed of multiple blocks (e.g., trajectory chunks). The function uses bisect_right to find the first block boundary strictly greater than i.

aimmd.core.utils.cycle(lst, n)

Rotate a list by n positions.

Parameters:

• lst (list) – Input list.

• n (int) – Rotation offset. Positive values rotate left, negative rotate right.

Returns:

Rotated list.

Return type:

list

Notes

• This is a convenience wrapper around slicing.

• n is reduced modulo len(lst).

aimmd.core.utils.concatenate(arrays, **kwargs)

Concatenate arrays along axis 0, skipping empty arrays.

Parameters:

• arrays (iterable of array-like) – Input arrays. Elements with len(array) == 0 are ignored.

• **kwargs – Passed to numpy.concatenate.

Returns:

Concatenated array. If all inputs are empty, returns np.array([]).

Return type:

np.ndarray

Notes

This function is useful when building datasets incrementally, where some components may be empty (e.g., no samples collected for a state).

aimmd.core.utils.extend_array(instance, min_length)

Extend an array to at least min_length along axis 0 by zero-padding.

Parameters:

• instance (np.ndarray) – Original array (assumed at least 1D).

• min_length (int) – Target minimum length.

Returns:

If instance is already long enough, it is returned unchanged.

Otherwise, a new zero-padded array is returned. :rtype: np.ndarray

Notes

• The returned array is marked read-only (flags.writeable = False). This supports cache-like usage where consumers must not mutate the array.

• Padding is applied only along axis 0; all trailing dimensions are preserved.

aimmd.core.utils.merge_ranges(ranges)

Merge overlapping integer ranges.

Parameters:

ranges – Input ranges as (begin, end) tuples. Ranges are merged when the next

begin is <= the current end. :type ranges: iterable of tuple(int, int)

Returns:

Sorted and merged ranges.

Return type:

list[tuple(int, int)]

Notes

• This is a generic interval merge helper; it does not enforce a specific inclusive/exclusive convention beyond the overlap test.

• The input is sorted, so the function is O(n log n) due to sorting.

aimmd.core.utils.longest_true_segment(mask, left=True)

Return the half-open interval [b, e) of the longest contiguous True segment.

Parameters:

• mask (array-like of bool) – Input boolean array.

• left – If multiple segments have the same maximal length:

• True -> return the first one

• False -> return the last one

Returns:

b, e – Start and end indices of the longest True segment, with e excluded.

Return type:

int, int

Notes

If no element is True, returns (0, 0).

aimmd.core.utils.remove(*patterns, except_for=[], verbose=True)

Remove files matched by one or more patterns.

Parameters:

*patterns – Filenames or glob-style patterns (wildcards supported).

Examples: "foo.txt", "*.npy", "temp/[0-9]*.xtc". :type *patterns: str :param except_for: Filenames to skip even if matched. :type except_for: str or list[str], optional :param verbose: If True, print status messages to stdout. :type verbose: bool, optional

Notes

• Directories are skipped (IsADirectoryError).

• Missing files are ignored (FileNotFoundError).

• Patterns are expanded using glob.glob when they contain wildcard tokens.

Side effects Deletes files from disk.

aimmd.core.utils.unique_path(path, extension=None)

Generate a unique filesystem path by incrementing a numeric suffix.

Parameters:

• path (str or pathlib.Path) – Candidate path.

• extension – If provided, force this suffix (e.g., ".xtc"). If path does not

already have this suffix, the extension is appended to the filename. :type extension: str, optional

Returns:

A path that does not exist on disk.

Return type:

pathlib.Path

Notes

• If path already exists, a numeric suffix is appended/updated: file -> file1 -> file2 -> ....

• If the stem already ends with digits, they are treated as the initial counter and incremented.

aimmd.core.utils.replace_in_cache(cache, old_name, new_name, prefixes=[''])

Rename file(s) on disk and update an in-memory cache key if present.

Parameters:

cache – Cache instance exposing a dict-like attribute cache._cache mapping

filenames to cached objects. :type cache: object :param old_name: Old filename (relative to each prefix). :type old_name: str :param new_name: New filename (relative to each prefix). :type new_name: str :param prefixes: Prefix strings prepended to the names. Typical usage is to provide directory prefixes like ["run1/", "run2/"] or [""] (default). :type prefixes: list[str], optional

Notes

• Each prefix is attempted independently.

• If os.replace fails for a prefix (missing file, permission, etc.), that prefix is silently skipped.

• If the old filename exists in the cache, the key is moved to the new name.

Side effects Renames files on disk via os.replace.

aimmd.core.utils.randomize_velocities(masses, T)

Sample velocities from a Maxwell–Boltzmann distribution and remove COM motion.

Parameters:

• masses (array-like) – Particle masses in atomic mass units (a.m.u.).

• T (float) – Temperature in Kelvin.

Returns:

Velocities of shape (n, 3), in Angstrom/ps.

Return type:

np.ndarray

Notes

• Uses \(k_B = 0.0083144621\) kJ/mol/K (gas constant in kJ/mol/K form).

• The per-particle standard deviation is:

  std_i = sqrt(kB * T / m_i) * 10

  where the factor 10 converts the internal unit scaling used by this code (A/ps) consistently with typical MD conventions in this project.

• After sampling, the mass-weighted center-of-mass velocity is subtracted.

Determinism This function uses NumPy’s global RNG. Control determinism by setting np.random.seed(…) in the caller.

aimmd.core.utils.memory_reader_from_timesteps(*list_of_timesteps)

Build an MDAnalysis MemoryReader from one or more timesteps (or iterables).

Parameters:

*list_of_timesteps – Any combination of:

• a single MDAnalysis Timestep,

• an iterable of timesteps,

• nested iterables of timesteps.

Returns:

An in-memory trajectory-like reader containing positions, velocities,

and dimensions. :rtype: MemoryReader

Notes

• Time metadata is not preserved (as stated by the original code).

• Missing velocities are filled with zeros for each timestep.

• dt is taken from the last visited timestep.

• This function expects Timestep and MemoryReader to be available in the module namespace (not imported here).

aimmd.core.utils.process_state(state, allowed_states='ARB')

Normalize a state label.

Parameters:

state – One of:

• a label in allowed_states (case-insensitive), e.g. 'A', 'B', 'R';

• a numeric index selecting from allowed_states, e.g. 0 -> 'A'.

• a numeric string (e.g. "1") is interpreted as an integer index.

Parameters:

allowed_states (str, optional) – Allowed state symbols in positional order.

Returns:

Normalized single-character state label.

Return type:

str

Raises:

TypeError – If the provided state is not valid.

Notes

This helper makes state handling uniform across:

• user-facing interfaces (strings),

• internal indexing (integers),

• serialized configurations (numeric strings).

aimmd.core.utils.extract_folder_and_name(fname)

Split a path string into (folder, name).

Parameters:

fname (str) – A filesystem path.

Returns:

(folder, name) where:

• name is the basename (final path component),

• folder is the parent directory, or "." if not present.

Return type:

tuple[str, str]

Notes

This helper exists because several cache utilities build auxiliary files next to the target file (locks, temp files, offsets), and need a robust way to form sibling paths.

aimmd.core.utils.guess_masses(atoms)

Heuristic mass assignment for an MDAnalysis AtomGroup.

Parameters:

atoms (AtomGroup) – MDAnalysis AtomGroup.

Returns:

Per-atom masses.

Return type:

np.ndarray

Notes

• Martini heuristic: If any of the first ~50 atoms has a name starting with BB, SC, or PO4, this function assumes a Martini-like coarse-grained system and returns a uniform mass of 72.0 for all beads.

• Otherwise, returns atoms.masses as provided by the topology.

Rationale In several workflows it is preferable to assign a reasonable default mass rather than propagate missing/zero masses (which can destabilize velocity initialization and COM removal).

aimmd.core.utils.accepts_system_id(function)

Whether function can be called with a system_id keyword argument.

Used to thread the multi-system system_id through the user data functions (states_function, descriptors_function, values_function) and descriptor_transform in a fully backward-compatible way: functions written for single-system runs (which take only the data argument) are detected here and keep being called without system_id.

Returns True if the function declares an explicit system_id parameter (positional-or-keyword or keyword-only) or accepts arbitrary **kwargs. Returns False for anything whose signature cannot be inspected.

Engines

The pure-Python toy engine used by the tests and tutorials.

Toy simulation engine.

This module implements ToyEngine, a minimal integrator that advances an MDAnalysis timestep object in pure Python and writes an output trajectory.

Intended use

• Method development.

• Code development and debugging.

• Unit tests / integration tests where GROMACS is not required.

• Demonstration runs with a controllable slowdown.

Core idea The engine:

1. loads an existing trajectory file (via MDA_CACHE),

2. creates a new writable trajectory file (either appending or creating a new part),

3. replays old frames (if appending),

4. generates new frames by repeatedly calling a user-supplied mdrun(ts) callback,

5. stops when walltime is reached or stop_condition() becomes true.

Concurrency model The engine runs two threads:

• one thread performs trajectory writing and timestep updates,

• one thread polls stop_condition() at short intervals.

This mirrors how “real” engines behave in AIMMD: a long-running simulation is interruptible by an external condition.

Notes

• This is a toy engine. It does not perform physical integration by itself. The mdrun callback is responsible for modifying ts.positions (and any other state) in a meaningful way.

class aimmd.engines.toy.ToyEngine(mdrun=None, slowdown=0.01, extension='.xtc')

Bases: object

Minimal engine that advances an MDAnalysis timestep via a callback.

Parameters:

mdrun – Function that advances the MDAnalysis timestep. It is called as

mdrun(ts) where ts is an MDAnalysis Timestep. If None, a no-op function is used (ts remains unchanged). :type mdrun: callable, optional :param slowdown: Seconds to sleep between generated frames. This is only for throttling the toy engine to avoid tight loops. Moreover, it avoid that the rest of the AIMMD code does not keep up with the simualtion speed. :type slowdown: float, default 0.01 :param extension: Trajectory extension used for input/output files. :type extension: str, default ‘.xtc’

Notes

ToyEngine is callable. Calling an instance runs the simulation loop and returns a status code (0/1) or raises, depending on raise_if_failure.

Construct a ToyEngine instance.

Parameters:

• mdrun (callable, optional) – Callback used to evolve the timestep.

• slowdown (float, default 0.01) – Sleep time between frames (seconds).

• extension (str, default '.xtc') – Output trajectory extension.

__init__(mdrun=None, slowdown=0.01, extension='.xtc')

Construct a ToyEngine instance.

Parameters:

• mdrun (callable, optional) – Callback used to evolve the timestep.

• slowdown (float, default 0.01) – Sleep time between frames (seconds).

• extension (str, default '.xtc') – Output trajectory extension.