Simulators

mythos supports multiple simulation backends. Each simulator exposes one or more observables (trajectories, scalars, etc.) that can be consumed by objectives during optimization.

For a walkthrough of running your first simulation, see Getting Started.

JaxMDSimulator

The JaxMDSimulator runs molecular dynamics entirely in Python using JAX-MD. Because the simulation is implemented in JAX, gradients can be computed via automatic differentiation, enabling direct optimization of energy function parameters without surrogate methods.

Key features:

• Full JAX autodiff through the simulation trajectory

• Pure-Python energy functions composed with the + operator

• Dynamic neighbor list updates each step

• Gradient checkpointing support to manage memory

• NVT Langevin dynamics with rigid-body support

Supported energy models: dna1, dna2, rna2, na1

See mythos.simulators.jax_md for the full API reference.

Examples: simulation notebook · optimization notebook

oxDNASimulator

The oxDNASimulator wraps the oxDNA C++/CUDA simulation code. Parameters are injected by modifying the oxDNA source (model.h) and recompiling the binary. Gradients are obtained via the DiffTRe algorithm (Boltzmann reweighting of reference trajectories).

Key features:

• Supports DNA1 and DNA2 interaction models

• Umbrella sampling via oxDNAUmbrellaSampler for melting temperature calculations

• Parameters baked into compiled binary (recompilation on update)

Supported energy models: dna1, dna2

Note

You must install oxDNA yourself. Pass the path to the oxDNA binary (simulation mode) or source directory (optimization mode) as parameters when constructing the simulator. See the oxDNA documentation for installation instructions.

See mythos.simulators.oxdna for the full API reference.

Examples: simulation notebook · advanced optimization notebook

GromacsSimulator

The GromacsSimulator runs GROMACS simulations for coarse-grained systems, particularly MARTINI lipid membranes. Energy function parameters are injected into the GROMACS topology (.top) file before each simulation via replace_params_in_topology(). Gradients are obtained via DiffTRe.

Key features:

• Two-phase simulation: equilibration followed by production

• MDP parameter overrides for runtime configuration

• Topology-file parameter injection for optimization

• Reads .trr trajectories via MDAnalysis

Supported energy models: martini/m2, martini/m3

Note

Requires a working GROMACS installation. The gmx binary must be available on your PATH. See the GROMACS installation guide for instructions.

See mythos.simulators.gromacs for the full API reference.

Examples: bottom-up optimization notebook · MARTINI optimization notebook

LAMMPSoxDNASimulator

The LAMMPSoxDNASimulator uses LAMMPS as a backend for running oxDNA-style simulations. Parameters are mapped from mythos energy function names to LAMMPS pair_coeff / bond_coeff directives via a REPLACEMENT_MAP and substituted into the LAMMPS input file.

Key features:

• LAMMPS-native implementation of oxDNA potentials

• Text-file parameter substitution (no recompilation)

• Temperature variable support

Supported energy models: dna1

Warning

LAMMPS does not implement BondedExcludedVolume. This energy term should be excluded from the energy function when using this simulator.

Note

Requires a working LAMMPS installation with the ASPHERE and CG-DNA packages enabled. The lmp binary must be available on your PATH. See the LAMMPS installation guide for instructions.

See mythos.simulators.lammps for the full API reference.

Creating a Custom Simulator

You can add your own simulation backend by subclassing the base simulator classes in mythos.simulators.base.

There are two base classes to choose from:

• ``Simulator`` — for in-process simulators (e.g. JAX-based). Override the run() method directly.

• ``InputDirSimulator`` — for file-based simulators that shell out to an external binary (e.g. GROMACS, LAMMPS). Implement the run_simulation(input_dir, ...) method; the base class handles copying input files to a temporary directory automatically based on its configured options.

Both are frozen chex dataclasses and must return a SimulatorOutput.

Minimal example using InputDirSimulator:

from pathlib import Path
from typing import Any

import chex
from typing_extensions import override

from mythos.simulators.base import InputDirSimulator, SimulatorOutput


@chex.dataclass(frozen=True, kw_only=True)
class MySimulator(InputDirSimulator):
    """A custom file-based simulator."""

    # Optional: override the default exposed observables
    # exposed_observables = ["trajectory", "energy"]

    @override
    def run_simulation(
        self, input_dir: Path, *args, opt_params: dict[str, Any], **kwargs
    ) -> SimulatorOutput:
        # 1. Write opt_params into input files in input_dir
        # 2. Run the external binary (subprocess, etc.)
        # 3. Read back results
        trajectory = ...  # load from output files
        return SimulatorOutput(observables=[trajectory])

Key points:

• Set exposed_observables (class variable) to declare which observables your simulator produces. Defaults to ["trajectory"].

• The exposes() method returns fully qualified names of the form "{observable}.{ClassName}.{name}" — these are the strings you reference in DiffTReObjective.required_observables.

• Use Simulator.create_n(n, name, ...) to create multiple instances with unique names for parallel optimization with RayOptimizer.

• The simulator should be serializable for compatibility with distributed optimizers like RayOptimizer. Avoid including non-serializable objects within the parameter definition of your simulator class.

See mythos.simulators.base for the full base class API.