diff --git a/examples/mvp/run.py b/examples/mvp/run.py index 90ede52..246245d 100644 --- a/examples/mvp/run.py +++ b/examples/mvp/run.py @@ -2,38 +2,21 @@ from pathlib import Path from datasets import Dataset -from openff.toolkit import ForceField, Molecule +from openff.toolkit import ForceField +from rich import print -import tyff -import tyff.converters -import tyff.train from tyff.compute._logging import _set_up_logger -from tyff.compute.configs import local_config -from tyff.compute.workflow import SimulationWorkflow -from tyff.converters.openff._tensors import convert_tensor_force_field +from tyff.converters import topologies_for from tyff.targets import energy, thermo +from tyff.targets.new_thermo import predict, simulate +from tyff.train import ParameterConfig, Trainable N_EPOCHS = 1 -N_REPLICATES = 1 +N_REPLICATES = 3 N_MOLECULES = 200 INITIAL_FORCE_FIELD = "openff-2.3.0.offxml" -def topologies_for( - smiles_set: set[str], - smirnoff_force_field: ForceField, -) -> tuple[dict[str, tyff.TensorTopology], tyff.TensorForceField]: - # this function is maybe meant to be surely there is a more straightforward way to do this? - interchanges = [ - smirnoff_force_field.create_interchange(Molecule.from_mapped_smiles(smiles).to_topology()) - for smiles in smiles_set - ] - - tensor_ff, topologies = tyff.converters.convert_interchange(interchanges) - - return {smiles: topology for smiles, topology in zip(smiles_set, topologies)}, tensor_ff - - def load_datasets(root: Path) -> tuple[Dataset, Dataset]: """Load datasets stored at data/*.tar.gz""" if not (root / "data/energy").exists(): @@ -47,70 +30,14 @@ 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 -def simulate( - force_field: tyff.TensorForceField, - root_dir: Path, - thermo_data: Dataset, - n_molecules: int, - n_replicates: int, -) -> None: - """ - Simulate the given force field using the provided thermo data. - - This is meant to be used inside of a training loop. - - Parameters - ---------- - force_field : tyff.TensorForceField - The force field to be used in the simulation. - root_dir : Path - The root directory where simulation results will be stored. - thermo_data : Dataset - The dataset containing thermodynamic data for the simulation. - n_molecules : int - The number of molecules to simulate. - n_replicates : int - The number of replicates for each simulation. - - Returns - ------- - None - """ - - # 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(INITIAL_FORCE_FIELD), - tensor_force_field=force_field, - ) - - logger.info("Making directory for this epoch's simulations") - root_dir.mkdir(exist_ok=True) - - logger.info("Writing out the current SMIRNOFF force field to file") - smirnoff_force_field_file = str(root_dir / "current_force_field.offxml") - current_smirnoff_force_field.to_file(smirnoff_force_field_file) - - logger.info("Setting up simulation workflow") - with SimulationWorkflow( - base_dir=str(root_dir / "simulations"), - parsl_config=local_config(max_workers=8), - ) as workflow: - logger.info("Submitting dataset") - workflow.submit_dataset( - dataset=thermo_data, - force_field=smirnoff_force_field_file, - n_molecules=n_molecules, - n_replicates=n_replicates, - ) - - workflow.wait() - - def main(root: Path) -> None: global logger @@ -118,8 +45,6 @@ def main(root: Path) -> None: thermo_data, energy_data = load_datasets(root) - # should this be convert_interchanges? - # is `topologies` list[tyff.TensorTopology] or dict[str, tyff.TensorTopology]? _topologies, tensor_ff = topologies_for( smiles_set={ *thermo.extract_smiles(thermo_data), @@ -128,10 +53,10 @@ def main(root: Path) -> None: smirnoff_force_field=ForceField(INITIAL_FORCE_FIELD), ) - trainable = tyff.train.Trainable( + trainable = Trainable( tensor_ff, parameters={ - "vdW": tyff.train.ParameterConfig( + "vdW": ParameterConfig( # type: ignore cols=["epsilon", "sigma"], scales={"epsilon": 10.0, "sigma": 1.0}, limits={"epsilon": (0.0001, None), "sigma": (0.5, None)}, @@ -144,15 +69,31 @@ def main(root: Path) -> None: for epoch in range(N_EPOCHS): current_ff = trainable.to_force_field(theta) + # no return value? simulate( force_field=current_ff, - root_dir=root / f"EPOCH_{epoch}", + epoch_dir=root / f"EPOCH_{epoch}", thermo_data=thermo_data, + smirnoff_force_field=ForceField(INITIAL_FORCE_FIELD), n_replicates=N_REPLICATES, n_molecules=N_MOLECULES, ) - return + y_ref, y_ref_std, y_pred, y_pred_std = predict( + dataset=thermo_data, + force_field=current_ff, + epoch_dir=root / f"EPOCH_{epoch}", + ) + + print(f"{y_ref=}\n{y_ref_std=}\n{y_pred=}\n{y_pred_std=}\n") + # y_ref=tensor([0.8736, 0.8060, 1.1335, 0.8579, 0.8215, 0.8375, 0.8543, 1.0950, 1.1203, + # 1.1290]) + # y_ref_std=tensor([4.0000e-05, 5.7000e-05, 5.0000e-04, 3.3500e-04, 1.1150e-04, 1.0100e-04, + # 9.1000e-05, 1.5000e-04, 3.5000e-04, 1.0000e-03]) + # y_pred=tensor([0.8859, 0.7895, 1.1050, 0.8637, 0.8140, 0.8355, 0.8586, 1.0800, 1.0986, + # 1.1028], dtype=torch.float64, grad_fn=) + # y_pred_std=tensor([0.0037, 0.0043, 0.0068, 0.0021, 0.0024, 0.0025, 0.0065, 0.0065, 0.0092, + # 0.0034], dtype=torch.float64, grad_fn=) if __name__ == "__main__": diff --git a/prototype.py b/prototype.py new file mode 100644 index 0000000..e9ad198 --- /dev/null +++ b/prototype.py @@ -0,0 +1,191 @@ +import json +import pathlib +from collections.abc import Sequence + +import torch + +import tyff +import tyff.converters +from tyff._serialization import load_tensor_force_field +from tyff.compute.jobs import make_job_id +from tyff.configs.liquid import BulkLiquid +from tyff.configs.targets.thermo import DataEntry + +# in leiu of better target-job tracking, just define known argumnets as ~globals +FORCE_FIELD = "openff-2.3.0.offxml" +N_MOLECULES = 200 +N_REPLICATES = 3 + + +def load_remote_observables( + force_field_per_epoch: tyff.TensorForceField, + base_dir: str, + job_ids: list[str], +) -> dict[str, dict]: + """Load observables and attach to graph.""" + from tyff.mm._ops import _pack_force_field + + tensors, *_ = _pack_force_field( + force_field_per_epoch + ) # force_field_per_epoch should also be the global reference FF passed to the workers + parameters: torch.Tensor = torch.cat( + [t.reshape(-1) for t in tensors if t is not None] + ) # has to be flat, and not a tuple + + # TODO: Move this class to not be inside of this function + class RemoteEnsembleAverage(torch.autograd.Function): + @staticmethod + def forward( + ctx: torch.autograd.function.FunctionCtx, + values: torch.Tensor, + jacobian: torch.Tensor, + parameters: torch.Tensor, + ) -> torch.Tensor: + ctx.save_for_backward(jacobian) + return values.clone() + + @staticmethod + def backward( + ctx: torch.autograd.function.FunctionCtx, + grad_output: torch.Tensor, + ) -> tuple[None, None, torch.Tensor]: + (jacobian,) = ctx.saved_tensors + return None, None, grad_output @ jacobian + + job_id_to_observables = {} + + for job_id in job_ids: + job_dir = pathlib.Path(base_dir) / job_id + + assert job_dir.exists() + + with open(job_dir / "ensemble_averages.json") as f: + means_dict: dict[str, float] = json.load(f) # actually str, torch.float64 + + jacobians: torch.Tensor = torch.load(job_dir / "jacobian.pt") + + names = list( + means_dict + ) # get the names of the computed properties in the means_dict so we can convert dict -> tensor -> dict again + tensor_means = parameters.new_tensor([means_dict[x] for x in names]) + attached_means = RemoteEnsembleAverage.apply( + tensor_means, + jacobians.to("cpu"), # to(device) ?, get from somewhere? + parameters, + ) + + job_id_to_observables[job_id] = dict(zip(names, attached_means)) + + return job_id_to_observables + + +def compute_properties( + dataset: Sequence[DataEntry], + observables: dict[str, dict[str, torch.Tensor]], +) -> list[tuple[dict, dict]]: + from tyff.compute.prep import _compute_configs_from_data_entry + + means_and_stds: list[tuple[dict, dict]] = list() + + for entry in dataset: + if entry["tag"] == "density": + # this includes replicates, but only jobs for this physical property target + # structure is a list (length n_replicates) of tuples, each tuple is + # 1-length and contains a compute config for (a single) liquid job + compute_configs: Sequence[tuple[BulkLiquid, ...]] = _compute_configs_from_data_entry( + data_entry=entry, + force_field=FORCE_FIELD, + n_molecules=N_MOLECULES, + n_replicates=N_REPLICATES, + ) + + # only one job per replicate in density calculatiosn + job_ids = [make_job_id(replicate[0]) for replicate in compute_configs] + + types_of_observables = [*observables[job_ids[0]].keys()] + + # could be simpler to return these as 1-D tensors - make sure to sort + # keys to avoid ordering issues + mean = { + key: torch.stack([observables[job_id][key] for job_id in job_ids]).mean() + for key in types_of_observables + } + std = { + key: torch.stack([observables[job_id][key] for job_id in job_ids]).std() + for key in types_of_observables + } + + # figure out which job id/s (if multiple replicas) in the observable corresponds to + # and calculate means and stds + elif entry["tag"] == "dielectric_constant": + 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 + elif entry["tag"] == "excess_molar_volume": + raise NotImplementedError() + elif entry["tag"] == "enthalpy_of_vaporization": + raise NotImplementedError() + else: + raise ValueError(f"I do not know what to do with {entry['tag']}") + + means_and_stds.append((mean, std)) + + return means_and_stds + + +def predict( + dataset: Sequence[DataEntry], + force_field: tyff.TensorForceField, + job_ids: list[str], + base_dir: str, +) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor]: + attached_means = load_remote_observables(force_field, base_dir, job_ids) + + means_and_stds: list[dict[str, torch.Tensor]] = compute_properties(dataset, attached_means) + + ref_values = [entry["value"] for entry in dataset] + ref_stds = [entry["std"] for entry in dataset] + predicted_values = [mean["density"] for mean, _ in means_and_stds] + predicted_stds = [std["density"] for _, std in means_and_stds] + + ref_values = torch.tensor(ref_values) + ref_stds = torch.tensor(ref_stds) + predicted_values = torch.tensor(predicted_values) + predicted_stds = torch.tensor(predicted_stds) + + return ref_values, ref_stds, predicted_values, predicted_stds + + +if __name__ == "__main__": + with open("density_example/thermo_targets/targets.json") as f: + targets = json.load(f) + + with open( + "density_example/reference_force_fields/openff-2.3.0.offxml.json", + ) as f: + global_force_field = load_tensor_force_field(json.load(f)) + + predict( + dataset=targets, + force_field=global_force_field, + job_ids=[ + "0f6ef5a4fdd601e764c9a70b185a6e27b1b91db6a8b65fb25a8a26922a3fca09", + "2ed40b18f3773e1a0eb756147e2b56e3da4c47149cdbf97c7bf6b85107df5c1c", + "5090adac8840d612600a2542c0236d47c2de3cdef6d6517961aa68e84f7ab8da", + "5a423d141a2e9d2d9ba9d30f988179ca8230bd8b93e7b854cc9c5c61f9ee9093", + "8418e9155a226c72f4d382d18e190ba023cbe55c9ad90175333d30637267059a", + "86cf94ed4a70093ce9d6d1f480eebaf806c457ca1c7347e5b2f153e59f12d59e", + "93e2b533fdec76225cd6af47938c69c40180d90a398a73b99892221471bcebc8", + "d721ef4f95b38106e81cc51e1d5638d9f09f372f4245cfe96ce2269d17e6a971", + "f08c269b4006de30978674f5ea1615258f471b5c21e045d57641d67ee86c3bd6", + "86cf94ed4a70093ce9d6d1f480eebaf806c457ca1c7347e5b2f153e59f12d59e", + "5090adac8840d612600a2542c0236d47c2de3cdef6d6517961aa68e84f7ab8da", + "2ed40b18f3773e1a0eb756147e2b56e3da4c47149cdbf97c7bf6b85107df5c1c", + ], + base_dir="density_example", + ) + + # TODO: Would be nice to dummy-check .backwards, can do this by constructing + # a minimal loss function of (predicted - expected) for each density value diff --git a/tyff/_tests/compute/test_workflow.py b/tyff/_tests/compute/test_workflow.py index b4781c8..52bf83e 100644 --- a/tyff/_tests/compute/test_workflow.py +++ b/tyff/_tests/compute/test_workflow.py @@ -15,3 +15,5 @@ def test_init(self, tmp_path): assert str(tmp_path) in lines[0] parsl.clear() + + assert (tmp_path / "target_job_map.json").exists() diff --git a/tyff/compute/_jacobian.py b/tyff/compute/_jacobian.py index 0eea661..2745a06 100644 --- a/tyff/compute/_jacobian.py +++ b/tyff/compute/_jacobian.py @@ -22,10 +22,17 @@ def _gather_from_reference(local: tyff.TensorForceField, reference: tyff.TensorF for potential in local.potentials: ref = reference_by_type[potential.handler_type] # KeyError: potential type missing from reference - assert potential.parameter_cols == ref.parameter_cols, potential.handler_type + assert potential.parameter_cols == ref.parameter_cols, ( + f"parameter_cols doesn't match in potential {potential.handler_type=}" + ) ref_idx = {key: i for i, key in enumerate(ref.parameter_keys)} - idx = torch.tensor([ref_idx[key] for key in potential.parameter_keys]) # KeyError: unseen parameter + try: + idx = torch.tensor([ref_idx[key] for key in potential.parameter_keys]) # KeyError: unseen parameter + except KeyError as error: + raise KeyError( + f"Parameter key {error} in local potential {potential.handler_type} not found in reference" + ) from error attributes = potential.attributes if potential.attribute_cols is not None: diff --git a/tyff/compute/_prepare.py b/tyff/compute/_prepare.py index 3c2cc7a..b5fe853 100644 --- a/tyff/compute/_prepare.py +++ b/tyff/compute/_prepare.py @@ -43,7 +43,7 @@ def _prepare_openmm_system( packed_topology: Topology = Topology.from_pdb( file_path=packing_files["packed_topology"].filepath, - unique_molecules=[Molecule.from_smiles(smiles) for smiles in compute_config["smiles"]], + unique_molecules=[Molecule.from_mapped_smiles(smiles) for smiles in compute_config["smiles"]], ) force_field = ForceField(compute_config["force_field"]) @@ -55,7 +55,7 @@ def _prepare_openmm_system( # # ordering is quite fragile, but hope order of smiles list never changes mid-job for unique_molecule_index, unique_molecule in enumerate( - [Molecule.from_smiles(smiles) for smiles in compute_config["smiles"]] + [Molecule.from_mapped_smiles(smiles) for smiles in compute_config["smiles"]] ): single_molecule_interchange = force_field.create_interchange(unique_molecule.to_topology()) with open(f"{job_dir}/single_molecule_interchange_{unique_molecule_index}.json", "w") as f: diff --git a/tyff/compute/workflow.py b/tyff/compute/workflow.py index 93b4674..a8990be 100644 --- a/tyff/compute/workflow.py +++ b/tyff/compute/workflow.py @@ -52,7 +52,17 @@ def __init__(self, base_dir: str, parsl_config: parsl.Config): # at scale these should become databases or something more robust self._targets = dict() - self._target_compute_mapping: dict[tuple[int, str, int, int], list[tuple[str, ...]]] = dict() + + target_job_map_path = pathlib.Path(self.base_dir) / "simulations" / "target_job_map.json" + if target_job_map_path.exists(): + with open(target_job_map_path) as f: + # don't overwrite this if it already exists + loaded_mapping = {tuple(k): v for k, v in json.load(f).items()} + if len(loaded_mapping) > 0: + 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() # futures submitted via submit_target/submit_target_batch that haven't # been waited on yet; drained by wait() @@ -123,11 +133,23 @@ def submit_target( ) # track targets (+ other arg) mapped to compute runs (as job_id) so we can look up results later - self._target_compute_mapping[(target_config["id"], force_field, n_molecules, n_replicates)] = list() - for this_targets_compute_configs in compute_configs: - self._target_compute_mapping[(target_config["id"], force_field, n_molecules, n_replicates)].append( - tuple([make_job_id(compute_config) for compute_config in this_targets_compute_configs]) + self._target_compute_mapping[ + ( + target_config["id"], # this is unmodified from input data and therefore could look funky + force_field, + n_molecules, + n_replicates, ) + ] = list() + for this_targets_compute_configs in compute_configs: + self._target_compute_mapping[ + ( + target_config["id"], + force_field, + n_molecules, + n_replicates, + ) + ].append(tuple([make_job_id(compute_config) for compute_config in this_targets_compute_configs])) # submit each compute job without blocking - can be >1 compute job per property. # NOTE: this used to call self.run(...), which calls .result() on every @@ -148,6 +170,21 @@ def submit_dataset( from tyff._serialization import dump_tensor_force_field + # save these targets since they're probably going to be needed later + (pathlib.Path(self.base_dir) / "thermo_targets").mkdir(exist_ok=True) + + if (pathlib.Path(self.base_dir) / "thermo_targets" / "targets.json").exists(): + with open(pathlib.Path(self.base_dir) / "thermo_targets" / "targets.json") as f: + targets = json.load(f) + else: + targets = list() + + # list-ify dicts, which should quack like `DataEntry`s, via iterating over datasets + targets.extend([*dataset]) + + with open(pathlib.Path(self.base_dir) / "thermo_targets" / "targets.json", "w") as f: + json.dump(targets, f, indent=4) + unique_smiles = set() for target in dataset: @@ -171,15 +208,32 @@ def submit_dataset( f.write(interchange.model_dump_json(indent=4)) self._reference_force_fields = pathlib.Path(self.base_dir) / "reference_force_fields" - pathlib.Path(self._reference_force_fields).mkdir(exist_ok=True) + self._reference_force_fields.mkdir(exist_ok=True) + + assert self._reference_force_fields.exists() + assert self._reference_force_fields.is_dir() + logger.info(f"Made a path to store reference force fields at {self._reference_force_fields}") # this is a "global"/reference force field composed from all unique molecules, # don't think we're going to use the topologies here reference_force_field, _tensor_topologies = tyff.converters.convert_interchange(interchanges) + # make sure this is not empty + assert len(reference_force_field.potentials) > 0 + + logger.info( + "Writing reference force field from path (which is probably outside of " + f"this workflow's root): {force_field=}" + f"to path {self._reference_force_fields / 'reference.json'}" + ) + # dumps a JSON representation of the tensor representation of the SMIRNOFF force field - with open(self._reference_force_fields / f"{force_field}.json", "w") as f: - json.dump(dump_tensor_force_field(reference_force_field), f) + with open(self._reference_force_fields / "reference.json", "w") as f: + json.dump( + dump_tensor_force_field(reference_force_field), + f, + indent=4, + ) return [ result @@ -255,6 +309,14 @@ def wait(self): made, so Parsl can actually run them concurrently in the meantime. Returns results/errors per job and clears the pending queue. """ + with open(pathlib.Path(self.base_dir) / "target_job_map.json", "w") as f: + json.dump( + # keys are tuples, so just serialize key-val pairs as lists + [[list(k), v] for k, v in self._target_compute_mapping.items()], + f, + indent=4, + ) + results = self._drain(self._pending) self._pending = [] return results @@ -351,7 +413,9 @@ def _common_run(self, compute_config: BaseComputeConfig) -> tuple[dict[str, torc jacobian_future = create_jacobian( production_future=production_future, job_dir=job_dir, - reference_force_field_path=self._reference_force_fields / f"{compute_config['force_field']}.json", + # pathing here is weird, but there should only be one global reference + # force field for this workflow + reference_force_field_path=self._reference_force_fields / "reference.json", ) return {"job_id": job_id, "future": jacobian_future} diff --git a/tyff/converters/__init__.py b/tyff/converters/__init__.py index 82bfd1a..a24acf4 100644 --- a/tyff/converters/__init__.py +++ b/tyff/converters/__init__.py @@ -4,6 +4,7 @@ convert_handlers, convert_interchange, smirnoff_parameter_converter, + topologies_for, ) from tyff.converters.openff._tensors import convert_tensor_force_field from tyff.converters.openmm import ( @@ -24,4 +25,5 @@ "convert_to_openmm_topology", "ffxml_converter", "smirnoff_parameter_converter", + "topologies_for", ] diff --git a/tyff/converters/openff/__init__.py b/tyff/converters/openff/__init__.py index ca44cb0..aa53e3e 100644 --- a/tyff/converters/openff/__init__.py +++ b/tyff/converters/openff/__init__.py @@ -4,6 +4,12 @@ convert_handlers, convert_interchange, smirnoff_parameter_converter, + topologies_for, ) -__all__ = ["convert_handlers", "convert_interchange", "smirnoff_parameter_converter"] +__all__ = ( + "convert_handlers", + "convert_interchange", + "smirnoff_parameter_converter", + "topologies_for", +) diff --git a/tyff/converters/openff/_openff.py b/tyff/converters/openff/_openff.py index 08e52a4..f8c9b04 100644 --- a/tyff/converters/openff/_openff.py +++ b/tyff/converters/openff/_openff.py @@ -494,3 +494,18 @@ def convert_interchange( tensor_force_field = tyff.TensorForceField(potentials, v_sites) return tensor_force_field, tensor_topologies + + +def topologies_for( + smiles_set: set[str], + smirnoff_force_field: openff.toolkit.ForceField, +) -> tuple[dict[str, tyff.TensorTopology], tyff.TensorForceField]: + # this function is maybe meant to be surely there is a more straightforward way to do this? + interchanges = [ + smirnoff_force_field.create_interchange(openff.toolkit.Molecule.from_mapped_smiles(smiles).to_topology()) + for smiles in smiles_set + ] + + tensor_ff, topologies = convert_interchange(interchanges) + + return {smiles: topology for smiles, topology in zip(smiles_set, topologies)}, tensor_ff diff --git a/tyff/mm/_ops.py b/tyff/mm/_ops.py index f652a97..c25b05a 100644 --- a/tyff/mm/_ops.py +++ b/tyff/mm/_ops.py @@ -552,3 +552,23 @@ def reweight_ensemble_averages( *avg_outputs, columns = _ReweightAverageOp.apply(kwargs, *tensors) return {column: avg for avg, column in zip(avg_outputs, columns, strict=True)} + + +class RemoteEnsembleAverage(torch.autograd.Function): + @staticmethod + def forward( + ctx: torch.autograd.function.FunctionCtx, + values: torch.Tensor, + jacobian: torch.Tensor, + parameters: torch.Tensor, + ) -> torch.Tensor: + ctx.save_for_backward(jacobian) + return values.clone() + + @staticmethod + def backward( + ctx: torch.autograd.function.FunctionCtx, + grad_output: torch.Tensor, + ) -> tuple[None, None, torch.Tensor]: + (jacobian,) = ctx.saved_tensors + return None, None, grad_output @ jacobian diff --git a/tyff/targets/new_thermo.py b/tyff/targets/new_thermo.py new file mode 100644 index 0000000..853c7bb --- /dev/null +++ b/tyff/targets/new_thermo.py @@ -0,0 +1,224 @@ +import json +from pathlib import Path + +import torch +from datasets import Dataset +from openff.toolkit import ForceField + +import tyff +import tyff.converters +import tyff.train +from tyff.compute.configs import local_config +from tyff.compute.workflow import SimulationWorkflow +from tyff.converters.openff._tensors import convert_tensor_force_field +from tyff.mm._ops import RemoteEnsembleAverage +from tyff.targets.thermo import DataEntry + + +def simulate( + force_field: tyff.TensorForceField, + epoch_dir: Path, + thermo_data: Dataset, + smirnoff_force_field: ForceField, + n_molecules: int, + n_replicates: int, +) -> None: + """ + Simulate the given force field using the provided thermo data. + + This is meant to be used inside of a training loop. + + Parameters + ---------- + force_field : tyff.TensorForceField + The 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 : ForceField + The original SMIRNOFF force field corresponding to the tensor force field. + n_molecules : int + The number of molecules to simulate. + n_replicates : int + The number of replicates for each simulation. + + Returns + ------- + None + """ + + # write out this epoch's (tensor) force field as a SMIRNOFF force field + current_smirnoff_force_field = convert_tensor_force_field( + original_force_field=smirnoff_force_field, + tensor_force_field=force_field, + ) + + epoch_dir.mkdir(exist_ok=True) + + smirnoff_force_field_file = str(epoch_dir / "current_force_field.offxml") + current_smirnoff_force_field.to_file(smirnoff_force_field_file) + + with SimulationWorkflow( + base_dir=str(epoch_dir / "simulations"), + parsl_config=local_config(max_workers=6), + ) as workflow: + workflow.submit_dataset( + dataset=thermo_data, + force_field=smirnoff_force_field_file, + n_molecules=n_molecules, + n_replicates=n_replicates, + ) + + workflow.wait() + + +def load_remote_observables( + force_field_per_epoch: tyff.TensorForceField, + base_dir: Path, + job_ids: list[str], +) -> dict[str, dict[str, torch.Tensor]]: + """ + Load observables and attach to graph. + + Returns a mapping between job IDs and associated observables. + """ + from tyff.mm._ops import _pack_force_field + + tensors, *_ = _pack_force_field( + force_field_per_epoch + ) # force_field_per_epoch should also be the global reference FF passed to the workers + parameters: torch.Tensor = torch.cat( + [t.reshape(-1) for t in tensors if t is not None] + ) # has to be flat, and not a tuple + + job_id_to_observables = {} + + for job_id in job_ids: + job_dir = base_dir / "simulations" / job_id + + assert job_dir.exists() + + with open(job_dir / "ensemble_averages.json") as f: + means_dict: dict[str, float] = json.load(f) # actually str, torch.float64 + + jacobians: torch.Tensor = torch.load(job_dir / "jacobian.pt") + + names = list( + means_dict + ) # get the names of the computed properties in the means_dict so we can convert dict -> tensor -> dict again + tensor_means = parameters.new_tensor([means_dict[x] for x in names]) + attached_means = RemoteEnsembleAverage.apply( + tensor_means, + jacobians.to("cpu"), # to(device) ?, get from somewhere? + parameters, + ) + + job_id_to_observables[job_id] = dict(zip(names, attached_means)) + + return job_id_to_observables + + +def compute_properties( + dataset: Dataset, + observables: dict[str, dict[str, torch.Tensor]], + epoch_dir: Path, +) -> list[tuple[dict[str, torch.Tensor], dict[str, torch.Tensor]]]: + + means_and_stds: list[tuple[dict[str, torch.Tensor], dict[str, torch.Tensor]]] = list() + + # think about better standardizing some of these paths + with open(epoch_dir / "simulations" / "target_job_map.json") as f: + target_compute_mapping = {tuple(k): v for k, v in json.load(f)} + + 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"]: + break + + # this is set with an iteration variable - can we assume that + # this will be the same for all job_ids within this target? + # all targets will be the same (i.e. always computing dhvap) + # but job-to-job maybe not? (i.e. gas vs. liquid jobs) + types_of_observables = [*observables[job_ids[0][0]].keys()] + + # could be simpler to return these as 1-D tensors - make sure to sort + # keys to avoid ordering issues + mean: dict[str, torch.Tensor] = { + key: torch.stack([observables[job_id[0]][key] for job_id in job_ids]).mean() + for key in types_of_observables + } + std: dict[str, torch.Tensor] = { + key: torch.stack([observables[job_id[0]][key] for job_id in job_ids]).std() + for key in types_of_observables + } + + # figure out which job id/s (if multiple replicas) in the observable corresponds to + # and calculate means and stds + elif entry["tag"] == "dielectric_constant": + 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 + elif entry["tag"] == "excess_molar_volume": + raise NotImplementedError() + elif entry["tag"] == "enthalpy_of_vaporization": + raise NotImplementedError() + else: + raise ValueError(f"I do not know what to do with {entry['tag']}") + + means_and_stds.append((mean, std)) + + return means_and_stds + + +def predict( + dataset: Dataset, + 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 + # dataset, we need the job IDs of these + + with open(epoch_dir / "simulations" / "target_job_map.json") as f: + target_compute_mapping = {tuple(k): v for k, v in json.load(f)} + + this_datasets_ids: set[str] = set(entry["id"] for entry in dataset) # type: ignore + + # this needs to be refactored to pass in _targets_, not jobs. + # synonymous with pure density, but for others (dhmix, dhvap, etc.) need to + # collect information from multiple jobs (i.e. gas and liquid jobs for dhvap) + # 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(): + 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) + + means_and_stds: list[tuple[dict[str, torch.Tensor], dict[str, torch.Tensor]]] = compute_properties( + dataset=dataset, + observables=attached_observables, + epoch_dir=epoch_dir, + ) + + ref_values = [entry["value"] for entry in dataset] # type: ignore + ref_stds = [entry["std"] for entry in dataset] # type: ignore + predicted_values = [mean["density"] for mean, _ in means_and_stds] + predicted_stds = [std["density"] for _, std in means_and_stds] + + ref_values = torch.tensor(ref_values) + ref_stds = torch.tensor(ref_stds) + predicted_values = torch.stack(predicted_values) + predicted_stds = torch.stack(predicted_stds) + + return ref_values, ref_stds, predicted_values, predicted_stds