runner

Run calculations defined by a config.

Modules:

• absolv –

  Absolute solvation free energy calculations using OpenMM

Classes:

• PreparedSystem –

  A container for the prepared inputs for a particular solvent phase.

Functions:

• setup –

  Prepare each system to be simulated, namely the ligand in each solvent.

• run_eq –

  Perform a simulation at each lambda window and for each solvent.

• run_neq –

  Performs the simulations required to estimate the free energy using a

PreparedSystem

Bases: NamedTuple

A container for the prepared inputs for a particular solvent phase.

Attributes:

• system (System) –

  The alchemically modified OpenMM system.

• topology (Topology) –

  The OpenMM topology with any box vectors set.

• coords (Quantity) –

  The coordinates of the system.

system instance-attribute

system: System

The alchemically modified OpenMM system.

topology instance-attribute

topology: Topology

The OpenMM topology with any box vectors set.

coords instance-attribute

coords: Quantity

The coordinates of the system.

setup

setup(
    system: System,
    config: Config,
    force_field: ForceField | SystemGenerator,
    custom_alchemical_potential: str | None = None,
) -> tuple[PreparedSystem, PreparedSystem]

Prepare each system to be simulated, namely the ligand in each solvent.

Parameters:

• system (System) –

  The system to prepare.

• config (Config) –

  The config defining the calculation to perform.

• force_field (ForceField | SystemGenerator) –

  The force field, or a callable that transforms an OpenFF topology into an OpenMM system, without any alchemical modifications to run the calculations using.

  If a callable is specified, it should take arguments of an OpenFF topology, a unit wrapped numpy array of atom coords, and a string literal with a value of either "solvent-a" or "solvent-b".

• custom_alchemical_potential (str | None, default: None ) –

  A custom expression to use for the potential energy function that describes the chemical-alchemical intermolecular interactions.

  See the absolv.fep.apply_fep function for more details.

Returns:

• tuple[PreparedSystem, PreparedSystem] –

  The two prepared systems, corresponding to solvent-a and solvent-b respectively.

Source code in absolv/runner.py

def setup(
    system: absolv.config.System,
    config: absolv.config.Config,
    force_field: openff.toolkit.ForceField | absolv.utils.openmm.SystemGenerator,
    custom_alchemical_potential: str | None = None,
) -> tuple[PreparedSystem, PreparedSystem]:
    """Prepare each system to be simulated, namely the ligand in each solvent.

    Args:
        system: The system to prepare.
        config: The config defining the calculation to perform.
        force_field: The force field, or a callable that transforms an OpenFF
            topology into an OpenMM system, **without** any alchemical modifications
            to run the calculations using.

            If a callable is specified, it should take arguments of an OpenFF
            topology, a unit wrapped numpy array of atom coords, and a string
            literal with a value of either ``"solvent-a"`` or ``"solvent-b"``.
        custom_alchemical_potential: A custom expression to use for the potential
            energy function that describes the chemical-alchemical intermolecular
            interactions.

            See the ``absolv.fep.apply_fep`` function for more details.

    Returns:
        The two prepared systems, corresponding to solvent-a and solvent-b respectively.
    """

    solvated_a = _setup_solvent(
        "solvent-a",
        system.components_a,
        force_field,
        system.n_solute_molecules,
        system.n_solvent_molecules_a,
        custom_alchemical_potential,
    )
    solvated_b = _setup_solvent(
        "solvent-b",
        system.components_b,
        force_field,
        system.n_solute_molecules,
        system.n_solvent_molecules_b,
        custom_alchemical_potential,
    )

    if system.solvent_a is not None and config.pressure is not None:
        absolv.utils.openmm.add_barostat(
            solvated_a.system, config.temperature, config.pressure
        )
    if system.solvent_b is not None and config.pressure is not None:
        absolv.utils.openmm.add_barostat(
            solvated_b.system, config.temperature, config.pressure
        )

    return solvated_a, solvated_b

run_eq

run_eq(
    config: Config,
    prepared_system_a: PreparedSystem,
    prepared_system_b: PreparedSystem,
    platform: OpenMMPlatform = CUDA,
    output_dir: Path | None = None,
    parallel: bool = False,
) -> Result

Perform a simulation at each lambda window and for each solvent.

Parameters:

• config (Config) –

  The config defining the calculation to perform.

• prepared_system_a (PreparedSystem) –

  The prepared system a. See setup for more details.

• prepared_system_b (PreparedSystem) –

  The prepared system b. See setup for more details.

• platform (OpenMMPlatform, default: CUDA ) –

  The OpenMM platform to run using.

• output_dir (Path | None, default: None ) –

  The (optional) directory to save HREMD samples to.

• parallel (bool, default: False ) –

  Whether to run calculations for solvent A and solvent B in parallel. This is mostly useful when running HFE calculations where the vacuum phase will typically run on the CPU while the solvent phase will run on the GPU.

Source code in absolv/runner.py

def run_eq(
    config: absolv.config.Config,
    prepared_system_a: PreparedSystem,
    prepared_system_b: PreparedSystem,
    platform: femto.md.constants.OpenMMPlatform = femto.md.constants.OpenMMPlatform.CUDA,
    output_dir: pathlib.Path | None = None,
    parallel: bool = False,
) -> absolv.config.Result:
    """Perform a simulation at each lambda window and for each solvent.

    Args:
        config: The config defining the calculation to perform.
        prepared_system_a: The prepared system a. See ``setup`` for more details.
        prepared_system_b: The prepared system b. See ``setup`` for more details.
        platform: The OpenMM platform to run using.
        output_dir: The (optional) directory to save HREMD samples to.
        parallel: Whether to run calculations for solvent A and solvent B in
            parallel. This is mostly useful when running HFE calculations where
            the vacuum phase will typically run on the CPU while the solvent phase
            will run on the GPU.
    """

    output_dir_a = None if output_dir is None else output_dir / "solvent-a"
    output_dir_b = None if output_dir is None else output_dir / "solvent-b"

    args = [
        (config.alchemical_protocol_a, prepared_system_a, output_dir_a),
        (config.alchemical_protocol_b, prepared_system_b, output_dir_b),
    ]
    run_fn = functools.partial(
        _run_eq_phase, temperature=config.temperature, platform=platform
    )

    if parallel:
        with multiprocessing.Pool(2) as pool:
            results = list(pool.starmap(run_fn, args))
    else:
        results = [run_fn(*args[0]), run_fn(*args[1])]

    results_a, overlap_a = results[0]
    results_b, overlap_b = results[1]

    dg_a, dg_a_std = results_a["ddG_kcal_mol"], results_a["ddG_error_kcal_mol"]
    # overlap_a = overlap_a["overlap_0"]
    dg_b, dg_b_std = results_b["ddG_kcal_mol"], results_b["ddG_error_kcal_mol"]
    # overlap_b = overlap_b["overlap_0"]

    return absolv.config.Result(
        dg_solvent_a=dg_a * openmm.unit.kilocalorie_per_mole,
        dg_std_solvent_a=dg_a_std * openmm.unit.kilocalorie_per_mole,
        dg_solvent_b=dg_b * openmm.unit.kilocalorie_per_mole,
        dg_std_solvent_b=dg_b_std * openmm.unit.kilocalorie_per_mole,
    )

run_neq

run_neq(
    config: Config,
    prepared_system_a: PreparedSystem,
    prepared_system_b: PreparedSystem,
    platform: OpenMMPlatform = CUDA,
) -> Result

Performs the simulations required to estimate the free energy using a non-equilibrium method.

These include equilibrium simulations at the two end states (i.e. fully interacting and fully de-coupled solute) for each solvent followed by non-equilibrium switching simulations between each end state to compute the forward and reverse work values.

Parameters:

• config (Config) –

  The config defining the calculation to perform.

• prepared_system_a (PreparedSystem) –

  The prepared system a. See setup for more details.

• prepared_system_b (PreparedSystem) –

  The prepared system b. See setup for more details.

• platform (OpenMMPlatform, default: CUDA ) –

  The OpenMM platform to run using.

Source code in absolv/runner.py

def run_neq(
    config: absolv.config.Config,
    prepared_system_a: PreparedSystem,
    prepared_system_b: PreparedSystem,
    platform: femto.md.constants.OpenMMPlatform = femto.md.constants.OpenMMPlatform.CUDA,
) -> absolv.config.Result:
    """Performs the simulations required to estimate the free energy using a
    non-equilibrium method.

    These include **equilibrium** simulations at the two end states (i.e. fully
    interacting and fully de-coupled solute) for each solvent followed by
    non-equilibrium switching simulations between each end state to compute the
    forward and reverse work values.

    Args:
        config: The config defining the calculation to perform.
        prepared_system_a: The prepared system a. See ``setup`` for more details.
        prepared_system_b: The prepared system b. See ``setup`` for more details.
        platform: The OpenMM platform to run using.
    """

    with tempfile.TemporaryDirectory() as tmp_dir:
        tmp_dir = pathlib.Path(tmp_dir)

        solvent_a_dir = tmp_dir / "solvent-a"
        solvent_b_dir = tmp_dir / "solvent-b"

        _run_phase_end_states(
            config.alchemical_protocol_a,
            config.temperature,
            prepared_system_a,
            solvent_a_dir,
            platform,
        )
        dg_a, dg_std_a = _run_switching(
            config.alchemical_protocol_a,
            config.temperature,
            prepared_system_a,
            solvent_a_dir,
            platform,
        )

        _run_phase_end_states(
            config.alchemical_protocol_b,
            config.temperature,
            prepared_system_b,
            solvent_b_dir,
            platform,
        )
        dg_b, dg_std_b = _run_switching(
            config.alchemical_protocol_b,
            config.temperature,
            prepared_system_b,
            solvent_b_dir,
            platform,
        )

    return absolv.config.Result(
        dg_solvent_a=dg_a.in_units_of(openmm.unit.kilocalories_per_mole),
        dg_std_solvent_a=dg_std_a.in_units_of(openmm.unit.kilocalories_per_mole),
        dg_solvent_b=dg_b.in_units_of(openmm.unit.kilocalories_per_mole),
        dg_std_solvent_b=dg_std_b.in_units_of(openmm.unit.kilocalories_per_mole),
    )