diff --git a/absolv/config.py b/absolv/config.py index 5ae017c..a4c470d 100644 --- a/absolv/config.py +++ b/absolv/config.py @@ -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.""" @@ -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]." diff --git a/absolv/runner.py b/absolv/runner.py index a744345..ab2754f 100644 --- a/absolv/runner.py +++ b/absolv/runner.py @@ -87,6 +87,7 @@ 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, @@ -94,6 +95,7 @@ def setup( """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 @@ -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 ) diff --git a/absolv/tests/test_config.py b/absolv/tests/test_config.py index 81b0634..0b614bf 100644 --- a/absolv/tests/test_config.py +++ b/absolv/tests/test_config.py @@ -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: diff --git a/docs/user-guide/overview.md b/docs/user-guide/overview.md index f58313b..8cf0838 100644 --- a/docs/user-guide/overview.md +++ b/docs/user-guide/overview.md @@ -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, @@ -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: diff --git a/regression/run.py b/regression/run.py index e322cad..76a9960 100644 --- a/regression/run.py +++ b/regression/run.py @@ -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(): @@ -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, @@ -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__":