diff --git a/examples/mvp/run.py b/examples/mvp/run.py index c10a478..7bbc15b 100644 --- a/examples/mvp/run.py +++ b/examples/mvp/run.py @@ -15,8 +15,8 @@ from tyff.train import ParameterConfig, Trainable N_EPOCHS = 10 -N_REPLICATES = 2 -N_MOLECULES = 200 +N_REPLICATES = 1 +N_MOLECULES = 480 INITIAL_FORCE_FIELD = "openff-2.3.0.offxml" # these weights should be exposed to the user in some sort of training config @@ -41,9 +41,6 @@ def load_datasets(root: Path) -> tuple[Dataset, Dataset]: thermo_data = Dataset.load_from_disk(root / "data/thermo") energy_data = Dataset.load_from_disk(root / "data/energy") - # temporary - need to go back and add dhmix, dhvap, etc. - thermo_data = thermo_data.filter(lambda x: x["tag"] == "density") - logger.info(f"Loaded {len(thermo_data)} thermo entries and {len(energy_data)} energy entries") return thermo_data, energy_data @@ -82,17 +79,17 @@ def main(root: Path) -> None: theta = trainable.to_values().requires_grad_(True) - # TODO: Possible drop in L-M optimizer + # TODO: Possibly drop in L-M optimizer optimizer = Adam([theta], lr=0.01) # TODO: inner/outer loop, or any way of stepping through QM optimization # without running MD jobs each time for epoch in range(N_EPOCHS): - # this is linked under the hood to theta - needed to gradients together - current_ff = trainable.to_force_field(theta) # 287 + # this is linked under the hood to theta - needed to hook gradients together + current_ff = trainable.to_force_field(theta) simulate( - force_field=current_ff, + tensor_force_field=current_ff, epoch_dir=root / f"EPOCH_{epoch}", thermo_data=thermo_data, smirnoff_force_field=INITIAL_FORCE_FIELD, @@ -102,7 +99,7 @@ def main(root: Path) -> None: y_ref, _y_ref_std, y_pred, _y_pred_std = predict( dataset=thermo_data, - force_field=current_ff, + tensor_force_field=current_ff, epoch_dir=root / f"EPOCH_{epoch}", ) diff --git a/pixi.lock b/pixi.lock index 7e958ff..0f9e447 100644 --- a/pixi.lock +++ b/pixi.lock @@ -5646,7 +5646,7 @@ packages: license: BSD-3-Clause license_family: BSD purls: - - pkg:pypi/cytoolz?source=hash-mapping + - pkg:pypi/cytoolz?source=compressed-mapping run_exports: {} size: 760733 timestamp: 1791511265780 @@ -9141,7 +9141,7 @@ packages: - python_abi 3.12.* *_cp312 license: MPL-2.0 AND Apache-2.0 purls: - - pkg:pypi/ml-dtypes?source=compressed-mapping + - pkg:pypi/ml-dtypes?source=hash-mapping run_exports: {} size: 369809 timestamp: 1792786261608 @@ -10497,7 +10497,7 @@ packages: - python_abi 3.12.* *_cp312 license: LGPL-2.1-only OR MPL-1.1 purls: - - pkg:pypi/pycairo?source=compressed-mapping + - pkg:pypi/pycairo?source=hash-mapping run_exports: {} size: 122097 timestamp: 1791296118521 @@ -12324,7 +12324,8 @@ packages: - openssl !=1.1.1e license: BSD-3-Clause license_family: BSD - purls: [] + purls: + - pkg:pypi/dask?source=compressed-mapping run_exports: {} size: 10415 timestamp: 1787917206356 @@ -12345,7 +12346,7 @@ packages: license: BSD-3-Clause license_family: BSD purls: - - pkg:pypi/dask?source=hash-mapping + - pkg:pypi/dask?source=compressed-mapping run_exports: {} size: 1071072 timestamp: 1787845029550 @@ -12466,7 +12467,7 @@ packages: license: BSD-3-Clause license_family: BSD purls: - - pkg:pypi/flexcache?source=hash-mapping + - pkg:pypi/flexcache?source=compressed-mapping run_exports: {} size: 18768 timestamp: 1788605834431 @@ -13995,7 +13996,7 @@ packages: license: BSD-3-Clause license_family: BSD purls: - - pkg:pypi/xyzservices?source=hash-mapping + - pkg:pypi/xyzservices?source=compressed-mapping run_exports: {} size: 53963 timestamp: 1788530754725 @@ -14997,7 +14998,7 @@ packages: license: BSD-3-Clause license_family: BSD purls: - - pkg:pypi/cytoolz?source=hash-mapping + - pkg:pypi/cytoolz?source=compressed-mapping run_exports: {} size: 695392 timestamp: 1791511337371 @@ -17627,7 +17628,7 @@ packages: - numpy >=1.25,<3 license: MPL-2.0 AND Apache-2.0 purls: - - pkg:pypi/ml-dtypes?source=hash-mapping + - pkg:pypi/ml-dtypes?source=compressed-mapping run_exports: {} size: 246811 timestamp: 1792786317299 diff --git a/tyff/_tests/conftest.py b/tyff/_tests/conftest.py index d959323..a60930c 100644 --- a/tyff/_tests/conftest.py +++ b/tyff/_tests/conftest.py @@ -24,6 +24,28 @@ _E = openff.units.unit.elementary_charge +def pytest_addoption(parser): + parser.addoption( + "--run-gpu", + action="store_true", + default=False, + help="run tests marked with @pytest.mark.gpu", + ) + + +def pytest_configure(config): + config.addinivalue_line("markers", "gpu: tests that require a GPU (skipped unless --run-gpu is given)") + + +def pytest_collection_modifyitems(config, items): + if config.getoption("--run-gpu"): + return + skip_gpu = pytest.mark.skip(reason="need --run-gpu option to run") + for item in items: + if "gpu" in item.keywords: + item.add_marker(skip_gpu) + + @pytest.fixture def density_entry() -> DataEntry: return { diff --git a/tyff/_tests/integration_tests/compute/test_workflow.py b/tyff/_tests/integration_tests/compute/test_workflow.py new file mode 100644 index 0000000..e607dae --- /dev/null +++ b/tyff/_tests/integration_tests/compute/test_workflow.py @@ -0,0 +1,110 @@ +from pathlib import Path + +import openff.toolkit +import pytest +import torch + +from tyff.converters.openff import topologies_for +from tyff.targets.new_thermo import predict, simulate +from tyff.targets.thermo import create_dataset + +# this variable is magically detected by pytest, important that it is defined at module level +pytestmark = pytest.mark.gpu + +# turn on to store simulation results in local path for debugging/iteration, otherwise (as in CI) +# let pytest start everything fresh in a temporary directory +USE_LOCAL_FILES = False + + +class TestEnthalpyOfVaporization: + def test_enthalpy_of_vaporization(self, tmp_path): + + dataset = create_dataset( + *[ + { # type: ignore + "id": "methanol", + "tag": "enthalpy_of_vaporization", + "x": [1.0], + "smiles": ["CO"], + "temperature": 300.0, + "pressure": 101.325, + "value": 38.3, + "std": 0.01, # arbitrary/unused + "units": "kJ/mol", + "source": "test_source", + } + ] + ) + + # I could be forgetting a cleaner way of going from SMIRNOFF + set[SMILES] -> TensorForceField + _topologies, tensor_force_field = topologies_for( + smiles_set={smiles for smiles_list in dataset["smiles"] for smiles in smiles_list}, + smirnoff_force_field=openff.toolkit.ForceField("openff-2.3.0.offxml"), + ) + + simulate( + tensor_force_field=tensor_force_field, + epoch_dir=Path("vaporization") if USE_LOCAL_FILES else tmp_path, + thermo_data=dataset, + smirnoff_force_field="openff-2.3.0.offxml", + n_replicates=2, + n_molecules=240, + ) + + y_ref, y_ref_std, y_pred, y_pred_std = predict( + dataset=dataset, + tensor_force_field=tensor_force_field, + epoch_dir=Path("vaporization") if USE_LOCAL_FILES else tmp_path, + ) + + assert not torch.isnan(torch.stack([y_ref, y_ref_std, y_pred, y_pred_std])).any() + + # experimental value is around 38 kJ/mol, could get away with much smaller range + assert 30 < y_pred < 50, f"{y_pred=} {y_ref=}, {y_pred_std=}, {y_ref_std=}" + + +class TestEnthalpyOfMixing: + def test_enthalpy_of_mixing(self, tmp_path): + + dataset = create_dataset( + *[ + { # type: ignore + "id": "foo", + "tag": "enthalpy_of_mixing", + "x": [0.8, 0.2], + "smiles": ["O", "CCO"], + "temperature": 298.15, + "pressure": 101.325, + "value": -0.765, + "std": 0.01, # arbitrary/unused + "units": "kJ/mol", + "source": "test_source", + } + ] + ) + + # I could be forgetting a cleaner way of going from SMIRNOFF + set[SMILES] -> TensorForceField + _topologies, tensor_force_field = topologies_for( + smiles_set={smiles for smiles_list in dataset["smiles"] for smiles in smiles_list}, + smirnoff_force_field=openff.toolkit.ForceField("openff-2.3.0.offxml"), + ) + + simulate( + tensor_force_field=tensor_force_field, + epoch_dir=Path("mixing") if USE_LOCAL_FILES else tmp_path, + thermo_data=dataset, + smirnoff_force_field="openff-2.3.0.offxml", + n_replicates=3, + n_molecules=480, + ) + + y_ref, y_ref_std, y_pred, y_pred_std = predict( + dataset=dataset, + tensor_force_field=tensor_force_field, + epoch_dir=Path("mixing") if USE_LOCAL_FILES else tmp_path, + ) + + assert not torch.isnan(torch.stack([y_ref, y_ref_std, y_pred, y_pred_std])).any() + + # experimental value is around -0.7 kJ/mol, but it should be negative (for all concentrations) + assert -1 < y_pred < 0, f"{y_pred=} {y_ref=}, {y_pred_std=}, {y_ref_std=}" diff --git a/tyff/compute/workflow.py b/tyff/compute/workflow.py index 26c58e7..ab18495 100644 --- a/tyff/compute/workflow.py +++ b/tyff/compute/workflow.py @@ -62,7 +62,19 @@ def __init__(self, base_dir: str, parsl_config: parsl.Config): self._target_compute_mapping = loaded_mapping else: # initialize to empty dict otherwise - self._target_compute_mapping: dict[tuple[str, str, int, int], list[tuple[str, ...]]] = dict() + self._target_compute_mapping: dict[ + tuple[ + str, # target id + str, # target tag + str, # force field path + int, # n_molecules + int, # n_replicates + ], + list[ # len n_replicates + # each job_dir in this replicate, i.e. 1 for density, 2 for dhvap, 3 for dhmix, etc. + tuple[str, ...] + ], + ] = dict() # futures submitted via submit_target/submit_target_batch that haven't # been waited on yet; drained by wait() @@ -136,6 +148,7 @@ def submit_target( self._target_compute_mapping[ ( target_config["id"], # this is unmodified from input data and therefore could look funky + target_config["tag"], # this is unmodified from input data and therefore could look funky force_field, n_molecules, n_replicates, @@ -145,6 +158,7 @@ def submit_target( self._target_compute_mapping[ ( target_config["id"], + target_config["tag"], force_field, n_molecules, n_replicates, @@ -237,7 +251,15 @@ def estimate_target( if target_config["tag"] == "density": # pull job ids for all replicates of this target/force field/n_molecules combination # TODO: other querying into jobs i.e. given a target and force field but any n_molecules - job_ids = self._target_compute_mapping[(target_config["id"], force_field, n_molecules, n_replicates)] + job_ids = self._target_compute_mapping[ + ( + target_config["id"], + target_config["tag"], + force_field, + n_molecules, + n_replicates, + ) + ] density_futures = [ run_density_analysis( @@ -258,7 +280,15 @@ def estimate_target( elif target_config["tag"] == "enthalpy_of_vaporization": # pull job ids for all replicates of this target/force field/n_molecules combination # TODO: other querying into jobs i.e. given a target and force field but any n_molecules - job_ids = self._target_compute_mapping[(target_config["id"], force_field, n_molecules, n_replicates)] + job_ids = self._target_compute_mapping[ + ( + target_config["id"], + target_config["tag"], + force_field, + n_molecules, + n_replicates, + ) + ] dhvap_futures = [ # this function should separate the two job ids into gas and liquid, @@ -281,6 +311,16 @@ def estimate_target( else: raise ValueError(f"Unknown target tag {target_config['tag']}") + def estimate_dataset( + self, + dataset: datasets.Dataset, # should be something that unpacks to DataEntry on iteration + tensor_force_field: tyff.TensorForceField, + force_field: str, + n_molecules: int, + n_replicates: int = 3, + ): + raise NotImplementedError() + def run(self, compute_configs: Sequence[BaseComputeConfig]): """Submit a batch and block until all complete.""" pending = self._submit_compute_batch(compute_configs) diff --git a/tyff/targets/new_thermo.py b/tyff/targets/new_thermo.py index 6326859..255bdc4 100644 --- a/tyff/targets/new_thermo.py +++ b/tyff/targets/new_thermo.py @@ -4,6 +4,7 @@ import torch from datasets import Dataset from openff.toolkit import ForceField +from openff.units import unit import tyff import tyff.converters @@ -14,9 +15,11 @@ from tyff.mm._ops import RemoteEnsembleAverage from tyff.targets.thermo import DataEntry +KCAL_TO_KJ = 4.184 + def simulate( - force_field: tyff.TensorForceField, + tensor_force_field: tyff.TensorForceField, epoch_dir: Path, thermo_data: Dataset, smirnoff_force_field: str, @@ -30,14 +33,14 @@ def simulate( Parameters ---------- - force_field : tyff.TensorForceField - The force field to be used in the simulation. + tensor_force_field : tyff.TensorForceField + The (tensor) force field to be used in the simulation. epoch_dir : Path The directory for this epoch where simulation results will be stored. thermo_data : Dataset The dataset containing thermodynamic data for the simulation. smirnoff_force_field : str - The path to the original SMIRNOFF force field corresponding to the tensor force field. + The path to the original (SMIRNOFF) force field corresponding to the tensor force field. n_molecules : int The number of molecules to simulate. n_replicates : int @@ -51,7 +54,7 @@ def simulate( # write out this epoch's (tensor) force field as a SMIRNOFF force field current_smirnoff_force_field = convert_tensor_force_field( original_force_field=ForceField(smirnoff_force_field), - tensor_force_field=force_field, + tensor_force_field=tensor_force_field, ) epoch_dir.mkdir(exist_ok=True) @@ -65,7 +68,7 @@ def simulate( ) as workflow: workflow.submit_dataset( dataset=thermo_data, - tensor_force_field=force_field, + tensor_force_field=tensor_force_field, force_field=smirnoff_force_field_file, # possibly smoother to pass through openff.toolkit.ForceField n_molecules=n_molecules, n_replicates=n_replicates, @@ -86,6 +89,15 @@ def load_remote_observables( Load observables and attach to graph. Returns a mapping between job IDs and associated observables. + + Parameters + ---------- + force_field_per_epoch : tyff.TensorForceField + The tensor force field used for the simulations. + base_dir : Path + The base directory where simulation results are stored. + job_ids : list[str] + A 1-D list of job IDs corresponding to the simulations. """ from tyff.mm._ops import _pack_force_field @@ -135,13 +147,14 @@ def compute_properties( with open(epoch_dir / "simulations" / "target_job_map.json") as f: target_compute_mapping = {tuple(k): v for k, v in json.load(f)} + # TODO: Wrap this all up into SimulationWorkflow.estimate_dataset for entry in dataset: # type: ignore entry: DataEntry if entry["tag"] == "density": # this lookup can be improved, but INITIAL_FORCE_FIELD # gets mutated by this point so it's not so straightforward - for as_submitted, job_ids in target_compute_mapping.items(): - if as_submitted[0] == entry["id"]: + for (target_id, *_), job_ids in target_compute_mapping.items(): + if target_id == entry["id"]: break # this is set with an iteration variable - can we assume that @@ -167,12 +180,168 @@ def compute_properties( raise NotImplementedError() # likely identical to density case elif entry["tag"] == "enthalpy_of_mixing": - raise NotImplementedError() - # figure out the three job IDs necessary from observables, calculate, etc + for (target_id, *_), job_ids in target_compute_mapping.items(): + if target_id == entry["id"]: + break + + values = list() + + for job_ids_in_this_replicate in job_ids: + assert len(job_ids_in_this_replicate) >= 3, ( + "Enthalpy of vaporization requires three or more jobs (pure and each component)" + ) + with ( + open( + epoch_dir / "simulations" / job_ids_in_this_replicate[0] / "compute_config.json", + ) as f0, + open( + epoch_dir / "simulations" / job_ids_in_this_replicate[1] / "compute_config.json", + ) as f1, + open( + epoch_dir / "simulations" / job_ids_in_this_replicate[2] / "compute_config.json", + ) as f2, + ): + compute_configs = [ + json.load(f0), + json.load(f1), + json.load(f2), + ] + + mixture_config = next(d for d in compute_configs if len(d["smiles"]) == 2) + + mixture_job_id = job_ids_in_this_replicate[compute_configs.index(mixture_config)] + + # try to keep this flexible to ternary+ mixtures + components = mixture_config["smiles"] + x = mixture_config["x"] + + pure_configs = [ + next(d for d in compute_configs if d["smiles"] == [component_i]) for component_i in components + ] + + component_job_ids = list() + for component_i in components: + component_job_id = next( + job_ids_in_this_replicate[compute_configs.index(d)] + for d in compute_configs + if d["smiles"] == [component_i] + ) + component_job_ids.append(component_job_id) + + # the provided job IDs should be represented by these variables, add several checks to make sure + # parsing+bookkeepping is all lined up + assert sorted([*component_job_ids, mixture_job_id]) == sorted(job_ids_in_this_replicate) + + assert len(component_job_ids) == len(pure_configs) == len(mixture_config["smiles"]) + + for component_config in pure_configs: + assert len(component_config["smiles"]) == len(component_config["x"]) == 1 + assert component_config["smiles"][0] in mixture_config["smiles"] + + mixture_energy = observables[mixture_job_id]["potential_energy"] + + component_energies = [ + observables[component_job_id]["potential_energy"] for component_job_id in component_job_ids + ] + component_n_molecules = [config["n_molecules"] for config in pure_configs] + + # this is kcal/mol because of what observables is computed as in tyff/mm/_ops.py + dhmix = mixture_energy / mixture_config["n_molecules"] - sum( + x_i * component_energy_i / n_i + for x_i, component_energy_i, n_i in zip(x, component_energies, component_n_molecules) + ) + + values.append(dhmix * KCAL_TO_KJ) # convert to kJ/mol + + types_of_observables = [*observables[job_ids[0][0]].keys()] + + mean: dict[str, torch.Tensor] = {key: torch.stack(values).mean() for key in types_of_observables} + std: dict[str, torch.Tensor] = {key: torch.stack(values).std() for key in types_of_observables} + elif entry["tag"] == "excess_molar_volume": raise NotImplementedError() elif entry["tag"] == "enthalpy_of_vaporization": - raise NotImplementedError() + for (target_id, *_), job_ids in target_compute_mapping.items(): + if target_id == entry["id"]: + break + + values = list() + + for job_ids_in_this_replicate in job_ids: + assert len(job_ids_in_this_replicate) == 2, ( + "Enthalpy of vaporization requires two jobs (gas and liquid)" + ) + with ( + open( + epoch_dir / "simulations" / job_ids_in_this_replicate[0] / "compute_config.json", + ) as f0, + open(epoch_dir / "simulations" / job_ids_in_this_replicate[1] / "compute_config.json") as f1, + ): + compute_config0 = json.load(f0) + compute_config1 = json.load(f1) + + if compute_config0["tag"] == "liquid" and compute_config1["tag"] == "gas": + liquid_job_id = job_ids_in_this_replicate[0] + liquid_config = compute_config0 + + gas_job_id = job_ids_in_this_replicate[1] + gas_config = compute_config1 + + elif compute_config0["tag"] == "gas" and compute_config1["tag"] == "liquid": + gas_job_id = job_ids_in_this_replicate[0] + gas_config = compute_config0 + + liquid_job_id = job_ids_in_this_replicate[1] + liquid_config = compute_config1 + + else: + raise ValueError( + f"Unexpected conmpute configs with tags {compute_config0['tag']} and " + f"{compute_config1['tag']}" + ) + + assert gas_config["temperature"] == liquid_config["temperature"], ( + "Gas and liquid temperatures must be the same for dHvap calculation" + ) + + # these are stored in kcal/mol, so need to convert + gas_energy = observables[gas_job_id]["potential_energy"] * KCAL_TO_KJ + liquid_energy = observables[liquid_job_id]["potential_energy"] * KCAL_TO_KJ + + n_molecules = liquid_config["n_molecules"] + + # 0.00831446261815324 + R = (1 * unit.avogadro_number * unit.boltzmann_constant).m_as("kJ/K") + + print( + f"Enthalpy of vaporization calculation: {gas_energy=}, {liquid_energy=}, " + f"{n_molecules=}, {R=}, {gas_config['temperature']=}" + ) + + # Enthalpy of vaporization calculation: + # gas_energy=tensor(1308.0884, dtype=torch.float64), + # liquid_energy=tensor(386691.8175, dtype=torch.float64), + # n_molecules=240, + # R=0.00831446261815324, + # gas_config['temperature']=300.0 + # + # Enthalpy of vaporization calculation: + # gas_energy=tensor(1300.6276, dtype=torch.float64), + # liquid_energy=tensor(385886.6263, dtype=torch.float64), + # n_molecules=240, + # R=0.00831446261815324, + # gas_config['temperature']=300.0 + + dhvap = gas_energy - liquid_energy / n_molecules + R * gas_config["temperature"] # kJ/mol + + values.append(dhvap) + + # assume the types of observables are uniform across all these job IDs + types_of_observables = [*observables[job_ids[0][0]].keys()] + + mean: dict[str, torch.Tensor] = {key: torch.stack(values).mean() for key in types_of_observables} + std: dict[str, torch.Tensor] = {key: torch.stack(values).std() for key in types_of_observables} + else: raise ValueError(f"I do not know what to do with {entry['tag']}") @@ -183,7 +352,7 @@ def compute_properties( def predict( dataset: Dataset, - force_field: tyff.TensorForceField, + tensor_force_field: tyff.TensorForceField, epoch_dir: Path, ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor]: # in order to get all of the (remote) observables for all simulations run from this @@ -200,14 +369,15 @@ def predict( # and condense down to one value (dhvap) to be reported job_ids: list[str] = list() - for (target_id, _, _, _), these_job_ids in target_compute_mapping.items(): + for (target_id, _target_tag, _, _, _), these_job_ids in target_compute_mapping.items(): if target_id in this_datasets_ids: job_ids.extend(these_job_ids) - # see comment above - job_ids = [x[0] for x in job_ids] - - attached_observables = load_remote_observables(force_field, epoch_dir, job_ids) + attached_observables = load_remote_observables( + force_field_per_epoch=tensor_force_field, + base_dir=epoch_dir, + job_ids=[x for sublist in job_ids for x in sublist], + ) means_and_stds: list[tuple[dict[str, torch.Tensor], dict[str, torch.Tensor]]] = compute_properties( dataset=dataset,