Skip to content

Commit

Permalink
Decouple config and system
Browse files Browse the repository at this point in the history
  • Loading branch information
@SimonBoothroyd
SimonBoothroyd committed Jan 15, 2024
1 parent 245b046 commit 5dfce5b
Show file tree
Hide file tree
Showing 5 changed files with 79 additions and 88 deletions.
37 changes: 19 additions & 18 deletions absolv/config.py
View file
Original file line number Diff line number Diff line change
Expand Up @@ -86,24 +86,32 @@ def _validate_solvent_b(cls, value):
), "at least one solvent must be specified when `solvent_b` is not none"
return value

def to_components(self) -> tuple[list[tuple[str, int]], list[tuple[str, int]]]:
"""Converts this object into two lists: one containing the identities and
counts of the molecules present in the first solution, and one containing the
same for the second solution.
@property
def components_a(self) -> list[tuple[str, int]]:
"""Returns the identities and counts of the molecules present in the first
system.
The identity and amount are stored in a tuple as a SMILES pattern and integer
count.
Returns:
The SMILES representation and count of each component.
"""

components_a = [*self.solutes.items()] + (
return [*self.solutes.items()] + (
[] if self.solvent_a is None else [*self.solvent_a.items()]
)
components_b = [*self.solutes.items()] + (

@property
def components_b(self) -> list[tuple[str, int]]:
"""Returns the identities and counts of the molecules present in the second
system.
Returns:
The SMILES representation and count of each component.
"""

return [*self.solutes.items()] + (
[] if self.solvent_b is None else [*self.solvent_b.items()]
)

return components_a, components_b


class MinimizationProtocol(pydantic.BaseModel):
"""Configure how a system should be energy minimized."""
Expand Down Expand Up @@ -308,14 +316,7 @@ class NonEquilibriumProtocol(pydantic.BaseModel):


class Config(pydantic.BaseModel):
"""A schema that fully defines the inputs needed to compute the transfer free energy
of a solvent between to solvents, or between a solvent and vacuum."""

system: System = pydantic.Field(
...,
description="A description of the system under investigation, including the "
"identity of the solute and the two solvent phases.",
)
"""Configure a transfer free energy calculation."""

temperature: OpenMMQuantity[openmm.unit.kelvin] = pydantic.Field(
..., description="The temperature to calculate at [K]."
Expand Down
22 changes: 10 additions & 12 deletions absolv/runner.py
View file
Original file line number Diff line number Diff line change
Expand Up @@ -87,13 +87,15 @@ def _setup_solvent(


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
Expand All @@ -112,32 +114,28 @@ def setup(
The two prepared systems, corresponding to solvent-a and solvent-b respectively.
"""

n_solute_molecules = config.system.n_solute_molecules

components_a, components_b = config.system.to_components()

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

if config.system.solvent_a is not None and config.pressure is not None:
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 config.system.solvent_b is not None and config.pressure is not None:
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
)
Expand Down
6 changes: 2 additions & 4 deletions absolv/tests/test_config.py
View file
Original file line number Diff line number Diff line change
Expand Up @@ -54,10 +54,8 @@ def test_to_components(self):
solutes={"CO": 1, "CCO": 2}, solvent_a={"O": 3}, solvent_b={"OCO": 4}
)

components_a, components_b = system.to_components()

assert components_a == [("CO", 1), ("CCO", 2), ("O", 3)]
assert components_b == [("CO", 1), ("CCO", 2), ("OCO", 4)]
assert system.components_a == [("CO", 1), ("CCO", 2), ("O", 3)]
assert system.components_b == [("CO", 1), ("CCO", 2), ("OCO", 4)]


class TestEquilibriumProtocol:
Expand Down
3 changes: 1 addition & 2 deletions docs/user-guide/overview.md
View file
Original file line number Diff line number Diff line change
Expand Up @@ -107,7 +107,6 @@ These individual components are then combined into a single configuration object
import absolv.config

config = absolv.config.Config(
system=system,
temperature=temperature,
pressure=pressure,
alchemical_protocol_a=alchemical_protocol_a,
Expand All @@ -122,7 +121,7 @@ import openff.toolkit
force_field = openff.toolkit.ForceField("openff-2.1.0.offxml")

import absolv.runner
prepared_system_a, prepared_system_b = absolv.runner.setup(config, force_field)
prepared_system_a, prepared_system_b = absolv.runner.setup(system, config, force_field)
```

run the calculation:
Expand Down
99 changes: 47 additions & 52 deletions regression/run.py
View file
Original file line number Diff line number Diff line change
Expand Up @@ -51,58 +51,52 @@
)


def default_config_neq(system: absolv.config.System) -> absolv.config.Config:
return absolv.config.Config(
system=system,
temperature=DEFAULT_TEMPERATURE,
pressure=DEFAULT_PRESSURE,
alchemical_protocol_a=absolv.config.NonEquilibriumProtocol(
switching_protocol=absolv.config.SwitchingProtocol(
n_electrostatic_steps=60,
n_steps_per_electrostatic_step=100, # 12 ps
n_steric_steps=0,
n_steps_per_steric_step=0,
)
),
alchemical_protocol_b=absolv.config.NonEquilibriumProtocol(
switching_protocol=absolv.config.SwitchingProtocol(
n_electrostatic_steps=60,
n_steps_per_electrostatic_step=100, # 12 ps
n_steric_steps=190,
n_steps_per_steric_step=100, # 38 ps
)
),
)


def default_config_eq(system: absolv.config.System) -> absolv.config.Config:
return absolv.config.Config(
system=system,
temperature=DEFAULT_TEMPERATURE,
pressure=DEFAULT_PRESSURE,
alchemical_protocol_a=absolv.config.EquilibriumProtocol(
production_protocol=absolv.config.HREMDProtocol(
n_steps_per_cycle=500,
n_cycles=2000,
integrator=femto.md.config.LangevinIntegrator(
timestep=1.0 * openmm.unit.femtosecond
),
DEFAULT_CONFIG_NEQ = absolv.config.Config(
temperature=DEFAULT_TEMPERATURE,
pressure=DEFAULT_PRESSURE,
alchemical_protocol_a=absolv.config.NonEquilibriumProtocol(
switching_protocol=absolv.config.SwitchingProtocol(
n_electrostatic_steps=60,
n_steps_per_electrostatic_step=100, # 12 ps
n_steric_steps=0,
n_steps_per_steric_step=0,
)
),
alchemical_protocol_b=absolv.config.NonEquilibriumProtocol(
switching_protocol=absolv.config.SwitchingProtocol(
n_electrostatic_steps=60,
n_steps_per_electrostatic_step=100, # 12 ps
n_steric_steps=190,
n_steps_per_steric_step=100, # 38 ps
)
),
)
DEFAULT_CONFIG_EQ = absolv.config.Config(
temperature=DEFAULT_TEMPERATURE,
pressure=DEFAULT_PRESSURE,
alchemical_protocol_a=absolv.config.EquilibriumProtocol(
production_protocol=absolv.config.HREMDProtocol(
n_steps_per_cycle=500,
n_cycles=2000,
integrator=femto.md.config.LangevinIntegrator(
timestep=1.0 * openmm.unit.femtosecond
),
lambda_sterics=absolv.config.DEFAULT_LAMBDA_STERICS_VACUUM,
lambda_electrostatics=absolv.config.DEFAULT_LAMBDA_ELECTROSTATICS_VACUUM,
),
alchemical_protocol_b=absolv.config.EquilibriumProtocol(
production_protocol=absolv.config.HREMDProtocol(
n_steps_per_cycle=500,
n_cycles=1000,
integrator=femto.md.config.LangevinIntegrator(
timestep=4.0 * openmm.unit.femtosecond
),
lambda_sterics=absolv.config.DEFAULT_LAMBDA_STERICS_VACUUM,
lambda_electrostatics=absolv.config.DEFAULT_LAMBDA_ELECTROSTATICS_VACUUM,
),
alchemical_protocol_b=absolv.config.EquilibriumProtocol(
production_protocol=absolv.config.HREMDProtocol(
n_steps_per_cycle=500,
n_cycles=1000,
integrator=femto.md.config.LangevinIntegrator(
timestep=4.0 * openmm.unit.femtosecond
),
lambda_sterics=absolv.config.DEFAULT_LAMBDA_STERICS_SOLVENT,
lambda_electrostatics=absolv.config.DEFAULT_LAMBDA_ELECTROSTATICS_SOLVENT,
),
)
lambda_sterics=absolv.config.DEFAULT_LAMBDA_STERICS_SOLVENT,
lambda_electrostatics=absolv.config.DEFAULT_LAMBDA_ELECTROSTATICS_SOLVENT,
),
)


def _download_ref_mols():
Expand Down Expand Up @@ -168,10 +162,11 @@ def run_replica(
):
output_dir.mkdir(parents=True, exist_ok=False)

config_generator = default_config_neq if method == "neq" else default_config_eq
config = config_generator(system)
config = DEFAULT_CONFIG_NEQ if method == "neq" else DEFAULT_CONFIG_EQ

prepared_system_a, prepared_system_b = absolv.runner.setup(config, system_generator)
prepared_system_a, prepared_system_b = absolv.runner.setup(
system, config, system_generator
)

femto.md.system.apply_hmr(
prepared_system_a.system,
Expand Down Expand Up @@ -232,7 +227,7 @@ def main(solutes: list[str], methods: list[str], replicas: list[str]):
DEFAULT_SYSTEMS[solute], system_generator, method, replica_dir
)
except BaseException as e:
logging.exception(f"failed to run {method} {solute} {replica}", e)
logging.error(f"failed to run {method} {solute} {replica}: {e}")


if __name__ == "__main__":
Expand Down

0 comments on commit 5dfce5b

Please sign in to comment.