Skip to content

simulate #

Run OpenMM simulations.

Modules:

  • femto

    A comprehensive toolkit for predicting free energies

Functions:

  • simulate_state

    Simulate a system following the specified stages, at a given 'state' (i.e.

  • simulate_states

    Simulate the system at each 'state' using simulate_state.

simulate_state #

simulate_state(
    system: System,
    topology: Structure,
    state: dict[str, float],
    stages: list[SimulationStage],
    platform: OpenMMPlatform,
    reporter: OpenMMStateReporter | None = None,
    enforce_pbc: bool = False,
) -> State

Simulate a system following the specified stages, at a given 'state' (i.e. a set of context parameters, such as free energy lambda values)

Parameters:

  • system (System) –

    The system to simulate.

  • topology (Structure) –

    The topology to simulate.

  • state (dict[str, float]) –

    The state to simulate at.

  • stages (list[SimulationStage]) –

    The stages to run.

  • platform (OpenMMPlatform) –

    The accelerator to use.

  • reporter (OpenMMStateReporter | None, default: None ) –

    The reporter to use to record system statistics such as volume and energy.

  • enforce_pbc (bool, default: False ) –

    Whether to enforce periodic boundary conditions when retrieving the final coordinates.

Returns:

  • State

    The final coordinates and box vectors.

Source code in femto/md/simulate.py
def simulate_state(
    system: openmm.System,
    topology: parmed.Structure,
    state: dict[str, float],
    stages: list[femto.md.config.SimulationStage],
    platform: femto.md.constants.OpenMMPlatform,
    reporter: femto.md.reporting.openmm.OpenMMStateReporter | None = None,
    enforce_pbc: bool = False,
) -> openmm.State:
    """Simulate a system following the specified ``stages``, at a given 'state' (i.e.
    a set of context parameters, such as free energy lambda values)

    Args:
        system: The system to simulate.
        topology: The topology to simulate.
        state: The state to simulate at.
        stages: The stages to run.
        platform: The accelerator to use.
        reporter: The reporter to use to record system statistics such as volume and
            energy.
        enforce_pbc: Whether to enforce periodic boundary conditions when retrieving
            the final coordinates.

    Returns:
        The final coordinates and box vectors.
    """

    reporter = (
        reporter
        if reporter is not None
        else femto.md.reporting.openmm.OpenMMStateReporter(
            femto.md.reporting.NullReporter(), "", 999999999
        )
    )

    stage_counter = collections.defaultdict(int)

    coords = None

    for stage in stages:
        stage_name = f"{stage.type}-{stage_counter[stage.type]}"
        stage_counter[stage.type] += 1

        reporter.tag = f"equilibration/{stage_name}"

        simulation = _prepare_simulation(
            system, topology, state, coords, stage, platform
        )
        simulation.reporters.append(reporter)

        if isinstance(stage, femto.md.config.Minimization):
            reporter.report(simulation, simulation.context.getState(getEnergy=True))
            simulation.minimizeEnergy(
                stage.tolerance.value_in_unit(_KJ_PER_NM), stage.max_iterations
            )
            reporter.report(simulation, simulation.context.getState(getEnergy=True))
        elif isinstance(stage, femto.md.config.Anneal):
            femto.md.anneal.anneal_temperature(
                simulation,
                stage.temperature_initial,
                stage.temperature_final,
                stage.n_steps,
                stage.frequency,
            )
        elif isinstance(stage, femto.md.config.Simulation):
            simulation.step(stage.n_steps)
        else:
            raise NotImplementedError(f"unknown stage type {type(stage)}")

        _LOGGER.info(
            f"after {stage_name} "
            f"{femto.md.utils.openmm.get_simulation_summary(simulation)}"
        )
        coords = simulation.context.getState(
            getPositions=True,
            getVelocities=True,
            getForces=True,
            getEnergy=True,
            enforcePeriodicBox=enforce_pbc,
        )

    return coords

simulate_states #

simulate_states(
    system: System,
    topology: Structure,
    states: list[dict[str, float]],
    stages: list[SimulationStage],
    platform: OpenMMPlatform,
    reporter: OpenMMStateReporter | None = None,
    enforce_pbc: bool = False,
) -> list[State]

Simulate the system at each 'state' using simulate_state.

If running using MPI, each process will be responsible for simulating at a subset of the states.

Parameters:

  • system (System) –

    The system to simulate.

  • topology (Structure) –

    The topology of the system to simulate.

  • states (list[dict[str, float]]) –

    The states of the system to simulate.

  • stages (list[SimulationStage]) –

    The stages to run.

  • platform (OpenMMPlatform) –

    The accelerator to use.

  • reporter (OpenMMStateReporter | None, default: None ) –

    The reporter to use to record system statistics such as volume and energy.

  • enforce_pbc (bool, default: False ) –

    Whether to enforce periodic boundary conditions when retrieving the final coordinates.

Returns:

  • list[State]

    The final coordinates at each state.

Source code in femto/md/simulate.py
def simulate_states(
    system: openmm.System,
    topology: parmed.Structure,
    states: list[dict[str, float]],
    stages: list[femto.md.config.SimulationStage],
    platform: femto.md.constants.OpenMMPlatform,
    reporter: femto.md.reporting.openmm.OpenMMStateReporter | None = None,
    enforce_pbc: bool = False,
) -> list[openmm.State]:
    """Simulate the system at each 'state' using ``simulate_state``.

    If running using MPI, each process will be responsible for simulating at a subset
    of the states.

    Args:
        system: The system to simulate.
        topology: The topology of the system to simulate.
        states: The states of the system to simulate.
        stages: The stages to run.
        platform: The accelerator to use.
        reporter: The reporter to use to record system statistics such as volume and
            energy.
        enforce_pbc: Whether to enforce periodic boundary conditions when retrieving
            the final coordinates.

    Returns:
        The final coordinates at each state.
    """

    with femto.md.utils.mpi.get_mpi_comm() as mpi_comm:
        # figure out how many states will be run in this MPI process
        n_local_states, state_offset = femto.md.utils.mpi.divide_tasks(
            mpi_comm, len(states)
        )

        final_coords: dict[int, openmm.State] = {}

        for i in range(n_local_states):
            state_idx = i + state_offset

            final_coords[state_idx] = simulate_state(
                system,
                topology,
                states[i + state_offset],
                stages,
                platform,
                reporter if state_idx == 0 else None,
                enforce_pbc=enforce_pbc,
            )

        final_coords = femto.md.utils.mpi.reduce_dict(final_coords, mpi_comm)

    return [final_coords[i] for i in range(len(states))]