From b7324d222a49d18b08335a01649abdb0ac66a734 Mon Sep 17 00:00:00 2001 From: GuoJeff Date: Wed, 1 Jun 2022 17:17:01 +0200 Subject: [PATCH] REINVENT v3.2 --- Dockerfile | 34 +++++ README.md | 4 +- configs/example.config.json | 11 +- input.py | 39 +++-- pytest.ini | 4 - reinvent.yml | 8 +- running_modes/.directory | 3 + .../actions/__init__.py | 4 + .../actions/base_action.py | 11 ++ .../actions/base_sample_action.py | 11 ++ .../actions/lib_invent_sample_model.py | 68 +++++++++ .../actions/link_invent_sample_model.py | 72 ++++++++++ .../actions/reinvent_sample_model.py | 38 +++++ .../automated_curriculum_runner.py | 28 ++-- .../base_curriculum_strategy.py | 131 ++++++----------- .../curriculum_strategy.py | 43 ++++-- .../linkinvent_curriculum_strategy.py | 110 ++++++++++++++ .../no_curriculum_strategy.py | 24 ++++ .../reinvent_curriculum_strategy.py | 77 ++++++++++ .../user_defined_curriculum.py | 34 ----- .../dto/__init__.py | 5 + .../dto/curriculum_outcome_dto.py | 9 ++ .../dto/sampled_batch_dto.py | 12 ++ .../dto/sampled_sequences_dto.py | 8 ++ .../dto/timestep_dto.py | 15 ++ .../dto/updated_likelihoods_dto.py | 10 ++ .../inception/__init__.py | 0 .../inception/inception.py | 50 +++++++ .../learning_strategy/__init__.py | 6 + .../base_double_query_learning_strategy.py | 25 ++++ .../base_learning_strategy.py | 40 ++++++ .../learning_strategy/base_linker_strategy.py | 30 ++++ .../base_single_query_learning_strategy.py | 30 ++++ .../dap_single_query_strategy.py | 41 ++++++ .../learning_strategy/dap_strategy.py | 33 +++++ .../learning_strategy/learning_strategy.py | 23 +++ .../learning_strategy_configuration.py | 9 ++ .../learning_strategy_enum.py | 10 ++ .../learning_strategy/mascof_strategy.py | 27 ++++ .../learning_strategy/mauli_strategy.py | 29 ++++ .../learning_strategy/sdap_strategy.py | 30 ++++ .../logging/__init__.py | 1 - .../logging/auto_cl_logger.py | 22 +-- .../logging/base_auto_cl_logger.py | 99 ------------- .../logging/base_logger.py | 81 +++++++++++ .../logging/console_message.py | 41 +++--- .../logging/local_auto_cl_logger.py | 95 ------------ .../logging/local_logger.py | 90 ++++++++++++ .../base_production_strategy.py | 110 ++++---------- .../link_invent_production_strategy.py | 111 ++++++++++++++ .../production_strategy.py | 33 +++-- .../reinvent_production_strategy.py | 82 +++++++++++ .../standard_production_strategy.py | 40 ------ .../automated_curriculum_learning/__init__.py | 8 +- ...mated_curriculum_learning_configuration.py | 3 - ...curriculum_learning_input_configuration.py | 12 ++ ...culum_learning_linkinvent_configuration.py | 12 ++ .../base_configuration.py | 5 +- .../curriculum_strategy_configuration.py | 10 +- ...curriculum_strategy_input_configuration.py | 26 ++++ .../inception_configuration.py | 6 +- ...nvent_curriculum_strategy_configuration.py | 25 ++++ ...nvent_production_strategy_configuration.py | 25 ++++ ...prodcution_strategy_input_configuration.py | 27 ++++ .../production_strategy_configuration.py | 14 +- .../curriculum_learning_components.py | 3 - .../curriculum_learning_configuration.py | 5 +- .../logging/curriculum_log_configuration.py | 3 + .../reinforcement_learning_configuration.py | 2 - ...ed_curriculum_learning_mode_constructor.py | 20 --- .../curriculum_learning_mode_constructor.py | 51 +++++-- running_modes/constructors/running_mode.py | 2 - .../curriculum_learning/curriculum_runner.py | 50 ++++--- .../logging/base_curriculum_logger.py | 4 +- .../logging/local_curriculum_logger.py | 2 +- .../logging/remote_curriculum_logger.py | 19 ++- .../enums/curriculum_strategy_enum.py | 2 + running_modes/enums/merging_strategy_enum.py | 9 ++ running_modes/enums/model_type_enum.py | 2 +- .../enums/production_strategy_enum.py | 3 +- running_modes/enums/ranking_strategy_enum.py | 8 ++ running_modes/enums/scoring_table_enum.py | 10 ++ running_modes/manager.py | 24 ++-- .../base_learning_strategy.py | 11 +- .../lib_invent_reinforcement_learning.py | 2 +- .../link_invent_reinforcement_learning.py | 2 +- .../link_logging/base_reinforcement_logger.py | 2 +- .../bond_link_reinforcement_logger.py | 4 +- .../local_bond_link_reinforcement_logger.py | 3 +- .../remote_bond_link_reinforcement_logger.py | 136 ++++++++++++++++++ .../logging/reinforcement_logger.py | 5 +- .../logging/remote_reinforcement_logger.py | 17 +-- .../logging/remote_sampling_logger.py | 7 +- .../remote_transfer_learning_logger.py | 7 +- running_modes/utils/general.py | 4 +- .../fixtures/predictive_model_fixtures.py | 130 ++--------------- .../curriculum_tests/__init__.py | 2 +- .../test_automated_curriculum_learning.py | 54 ++++--- .../base_learning_strategy.py | 20 +-- .../test_learning_strategy_mascof_strategy.py | 3 +- .../test_learning_strategy_mauli_strategy.py | 3 +- .../test_learning_strategy_sdap_strategy.py | 3 +- 102 files changed, 1973 insertions(+), 835 deletions(-) create mode 100644 Dockerfile delete mode 100644 pytest.ini create mode 100644 running_modes/.directory create mode 100644 running_modes/automated_curriculum_learning/actions/__init__.py create mode 100644 running_modes/automated_curriculum_learning/actions/base_action.py create mode 100644 running_modes/automated_curriculum_learning/actions/base_sample_action.py create mode 100644 running_modes/automated_curriculum_learning/actions/lib_invent_sample_model.py create mode 100644 running_modes/automated_curriculum_learning/actions/link_invent_sample_model.py create mode 100644 running_modes/automated_curriculum_learning/actions/reinvent_sample_model.py create mode 100644 running_modes/automated_curriculum_learning/curriculum_strategy/linkinvent_curriculum_strategy.py create mode 100644 running_modes/automated_curriculum_learning/curriculum_strategy/no_curriculum_strategy.py create mode 100644 running_modes/automated_curriculum_learning/curriculum_strategy/reinvent_curriculum_strategy.py delete mode 100644 running_modes/automated_curriculum_learning/curriculum_strategy/user_defined_curriculum.py create mode 100644 running_modes/automated_curriculum_learning/dto/__init__.py create mode 100644 running_modes/automated_curriculum_learning/dto/curriculum_outcome_dto.py create mode 100644 running_modes/automated_curriculum_learning/dto/sampled_batch_dto.py create mode 100644 running_modes/automated_curriculum_learning/dto/sampled_sequences_dto.py create mode 100644 running_modes/automated_curriculum_learning/dto/timestep_dto.py create mode 100644 running_modes/automated_curriculum_learning/dto/updated_likelihoods_dto.py create mode 100644 running_modes/automated_curriculum_learning/inception/__init__.py create mode 100644 running_modes/automated_curriculum_learning/inception/inception.py create mode 100644 running_modes/automated_curriculum_learning/learning_strategy/__init__.py create mode 100644 running_modes/automated_curriculum_learning/learning_strategy/base_double_query_learning_strategy.py create mode 100644 running_modes/automated_curriculum_learning/learning_strategy/base_learning_strategy.py create mode 100644 running_modes/automated_curriculum_learning/learning_strategy/base_linker_strategy.py create mode 100644 running_modes/automated_curriculum_learning/learning_strategy/base_single_query_learning_strategy.py create mode 100644 running_modes/automated_curriculum_learning/learning_strategy/dap_single_query_strategy.py create mode 100644 running_modes/automated_curriculum_learning/learning_strategy/dap_strategy.py create mode 100644 running_modes/automated_curriculum_learning/learning_strategy/learning_strategy.py create mode 100644 running_modes/automated_curriculum_learning/learning_strategy/learning_strategy_configuration.py create mode 100644 running_modes/automated_curriculum_learning/learning_strategy/learning_strategy_enum.py create mode 100644 running_modes/automated_curriculum_learning/learning_strategy/mascof_strategy.py create mode 100644 running_modes/automated_curriculum_learning/learning_strategy/mauli_strategy.py create mode 100644 running_modes/automated_curriculum_learning/learning_strategy/sdap_strategy.py delete mode 100644 running_modes/automated_curriculum_learning/logging/base_auto_cl_logger.py create mode 100644 running_modes/automated_curriculum_learning/logging/base_logger.py delete mode 100644 running_modes/automated_curriculum_learning/logging/local_auto_cl_logger.py create mode 100644 running_modes/automated_curriculum_learning/logging/local_logger.py create mode 100644 running_modes/automated_curriculum_learning/production_strategy/link_invent_production_strategy.py create mode 100644 running_modes/automated_curriculum_learning/production_strategy/reinvent_production_strategy.py delete mode 100644 running_modes/automated_curriculum_learning/production_strategy/standard_production_strategy.py create mode 100644 running_modes/configurations/automated_curriculum_learning/automated_curriculum_learning_input_configuration.py create mode 100644 running_modes/configurations/automated_curriculum_learning/automated_curriculum_learning_linkinvent_configuration.py create mode 100644 running_modes/configurations/automated_curriculum_learning/curriculum_strategy_input_configuration.py create mode 100644 running_modes/configurations/automated_curriculum_learning/linkinvent_curriculum_strategy_configuration.py create mode 100644 running_modes/configurations/automated_curriculum_learning/linkinvent_production_strategy_configuration.py create mode 100644 running_modes/configurations/automated_curriculum_learning/prodcution_strategy_input_configuration.py delete mode 100644 running_modes/constructors/automated_curriculum_learning_mode_constructor.py create mode 100644 running_modes/enums/merging_strategy_enum.py create mode 100644 running_modes/enums/ranking_strategy_enum.py create mode 100644 running_modes/enums/scoring_table_enum.py create mode 100644 running_modes/reinforcement_learning/logging/link_logging/remote_bond_link_reinforcement_logger.py diff --git a/Dockerfile b/Dockerfile new file mode 100644 index 0000000..c4fd622 --- /dev/null +++ b/Dockerfile @@ -0,0 +1,34 @@ +FROM docker.io/continuumio/conda-ci-linux-64-python3.7:latest + +USER root + +RUN apt-get update && \ + apt-get -y install rsync procps && \ + wget https://sourceforge.net/projects/lmod/files/lua-5.1.4.9.tar.bz2 && \ + tar xf lua-5.1.4.9.tar.bz2 && \ + cd lua-5.1.4.9 && \ + ./configure --prefix=/opt/apps/lua/5.1.4.9 && \ + make; make install && \ + cd /opt/apps/lua; ln -s 5.1.4.9 lua && \ + ln -s /opt/apps/lua/lua/bin/lua /usr/local/bin && \ + ln -s /opt/apps/lua/lua/bin/luac /usr/local/bin && \ + cd; wget https://sourceforge.net/projects/lmod/files/Lmod-8.2.tar.bz2 && \ + tar xf Lmod-8.2.tar.bz2 && \ + cd Lmod-8.2; ./configure --prefix=/opt/apps --with-fastTCLInterp=no && \ + make install && \ + ln -s /opt/apps/lmod/lmod/init/profile /etc/profile.d/z00_lmod.sh + +ENV LMOD_ROOT=/opt/apps/lmod \ + LMOD_PKG=/opt/apps/lmod/lmod \ + LMOD_VERSION=8.2 \ + LMOD_CMD=/opt/apps/lmod/lmod/libexec/lmod \ + LMOD_DIR=/opt/apps/lmod/lmod/libexec \ + BASH_ENV=/opt/apps/lmod/lmod/init/bash + +COPY . /reinventcli/ + +WORKDIR /reinventcli + +RUN conda update -n base -c defaults conda && \ + conda env update --name=base --file=reinvent.yml && \ + chmod -R "a+rx" /reinventcli diff --git a/README.md b/README.md index 16384d1..fddcdab 100644 --- a/README.md +++ b/README.md @@ -1,4 +1,4 @@ -REINVENT 3.1 +REINVENT 3.2 ================================================================================================================= Installation @@ -12,7 +12,7 @@ Installation 4. Activate the environment: - $ conda activate reinvent.v3.0 + $ conda activate reinvent.v3.2 5. Use the tool. diff --git a/configs/example.config.json b/configs/example.config.json index 97a0af2..a2ffe50 100644 --- a/configs/example.config.json +++ b/configs/example.config.json @@ -18,6 +18,13 @@ "AZGARD_EXECUTOR_SCRIPT_PATH": "//executor.py", "AZGARD_ENV_PATH": "//miniconda3/envs/AZgard/bin/python", "AZGARD_DEBUG": true + }, + "ICOLOS": { + "ICOLOS_EXECUTOR_PATH": "//miniconda3/envs/icolosprod/bin/icolos", + "ICOLOS_DEBUG": true + }, + "AIZYNTH": { + "CONFIG": "/projects/mai/synthesisplanning/minimal_config.yml" } }, "ENVIRONMENTAL_VARIABLES": { @@ -28,5 +35,7 @@ "ACTIVITY_CLASSIFICATION": "", "SMILES_SET_PATH": "", "PRIOR_PATH": "", - "LIBINVENT_PRIOR_PATH": "" + "LIBINVENT_PRIOR_PATH": "", + "SMILES_SET_LINK_INVENT_PATH":"", + "LINK_INVENT_PRIOR_PATH": "" } \ No newline at end of file diff --git a/input.py b/input.py index d23d7dc..eb66a31 100644 --- a/input.py +++ b/input.py @@ -3,19 +3,38 @@ import sys import json +import argparse +from pathlib import Path from running_modes.manager import Manager -if __name__ == "__main__": +DEFAULT_BASE_CONFIG_PATH = (Path(__file__).parent / 'configs/config.json').resolve() + +parser = argparse.ArgumentParser(description='Run Reinvent.') +parser.add_argument( + '--base_config', type=str, default=DEFAULT_BASE_CONFIG_PATH, + help='Path to basic configuration for Reinvent environment.' +) +parser.add_argument( + 'run_config', type=str, + help='Path to configuration json file for this run.' +) - with open(sys.argv[1]) as f: - json_input = f.read().replace('\r', '').replace('\n', '') - configuration = {} +def read_json_file(path): + with open(path) as f: + json_input = f.read().replace('\r', '').replace('\n', '') try: - configuration = json.loads(json_input) - except (ValueError, KeyError, TypeError): - print("JSON format error") - else: - manager = Manager(configuration) - manager.run() + return json.loads(json_input) + except (ValueError, KeyError, TypeError) as e: + print(f"JSON format error in file ${path}: \n ${e}") + + +if __name__ == "__main__": + args = parser.parse_args() + + base_config = read_json_file(args.base_config) + run_config = read_json_file(args.run_config) + + manager = Manager(base_config, run_config) + manager.run() diff --git a/pytest.ini b/pytest.ini deleted file mode 100644 index 0834c57..0000000 --- a/pytest.ini +++ /dev/null @@ -1,4 +0,0 @@ -[pytest] -markers = - integration: mark a test as an integration test. - clab: mark a test as a test that requires CLAB. \ No newline at end of file diff --git a/reinvent.yml b/reinvent.yml index 96e31fe..84a095e 100644 --- a/reinvent.yml +++ b/reinvent.yml @@ -1,4 +1,4 @@ -name: reinvent.v3.0 +name: reinvent.v3.2 channels: - rdkit - pytorch @@ -211,9 +211,9 @@ dependencies: - markdown==3.2.1 - opt-einsum==3.2.0 - protobuf==3.11.3 - - reinvent-chemistry==0.0.40 - - reinvent-models==0.0.12 - - reinvent-scoring==0.0.57 + - reinvent-chemistry==0.0.50 + - reinvent-models==0.0.15rc1 + - reinvent-scoring==0.0.73 - tensorboard==1.15.0 - tensorflow==1.15.2 - tensorflow-estimator==1.15.1 diff --git a/running_modes/.directory b/running_modes/.directory new file mode 100644 index 0000000..f964202 --- /dev/null +++ b/running_modes/.directory @@ -0,0 +1,3 @@ +[Dolphin] +Timestamp=2022,4,8,15,57,33 +Version=3 diff --git a/running_modes/automated_curriculum_learning/actions/__init__.py b/running_modes/automated_curriculum_learning/actions/__init__.py new file mode 100644 index 0000000..67680ed --- /dev/null +++ b/running_modes/automated_curriculum_learning/actions/__init__.py @@ -0,0 +1,4 @@ +from running_modes.automated_curriculum_learning.actions.base_action import BaseAction +from running_modes.automated_curriculum_learning.actions.base_sample_action import BaseSampleAction +from running_modes.automated_curriculum_learning.actions.lib_invent_sample_model import LibInventSampleModel +from running_modes.automated_curriculum_learning.actions.link_invent_sample_model import LinkInventSampleModel diff --git a/running_modes/automated_curriculum_learning/actions/base_action.py b/running_modes/automated_curriculum_learning/actions/base_action.py new file mode 100644 index 0000000..a0c78ac --- /dev/null +++ b/running_modes/automated_curriculum_learning/actions/base_action.py @@ -0,0 +1,11 @@ +import abc +from running_modes.automated_curriculum_learning.logging.base_logger import BaseLogger + + +class BaseAction(abc.ABC): + def __init__(self, logger=None): + """ + (Abstract) Initializes an action. + :param logger: An optional logger instance. + """ + self.logger: BaseLogger = logger diff --git a/running_modes/automated_curriculum_learning/actions/base_sample_action.py b/running_modes/automated_curriculum_learning/actions/base_sample_action.py new file mode 100644 index 0000000..70a3972 --- /dev/null +++ b/running_modes/automated_curriculum_learning/actions/base_sample_action.py @@ -0,0 +1,11 @@ +import numpy as np +from running_modes.automated_curriculum_learning.actions import BaseAction + + +class BaseSampleAction(BaseAction): + + def _get_indices_of_unique_smiles(self, smiles: [str]) -> np.array: + """Returns an np.array of indices corresponding to the first entries in a list of smiles strings""" + _, idxs = np.unique(smiles, return_index=True) + sorted_indices = np.sort(idxs) + return sorted_indices \ No newline at end of file diff --git a/running_modes/automated_curriculum_learning/actions/lib_invent_sample_model.py b/running_modes/automated_curriculum_learning/actions/lib_invent_sample_model.py new file mode 100644 index 0000000..d9925fd --- /dev/null +++ b/running_modes/automated_curriculum_learning/actions/lib_invent_sample_model.py @@ -0,0 +1,68 @@ +from typing import List + +import numpy as np +from reinvent_chemistry import Conversions +from reinvent_chemistry.library_design import BondMaker, AttachmentPoints +from reinvent_models.lib_invent.models.dataset import Dataset +from reinvent_models.model_factory.generative_model_base import GenerativeModelBase +from torch.utils.data import DataLoader + +from running_modes.automated_curriculum_learning.actions import BaseSampleAction +from running_modes.automated_curriculum_learning.dto.sampled_sequences_dto import SampledSequencesDTO + + +class LibInventSampleModel(BaseSampleAction): + + def __init__(self, model: GenerativeModelBase, batch_size: int, logger=None, randomize=False, sample_uniquely=True): + """ + Creates an instance of SampleModel. + :params model: A model instance (better in scaffold_decorating mode). + :params batch_size: Batch size to use. + :return: + """ + super().__init__(logger) + self.model = model + self._batch_size = batch_size + self._bond_maker = BondMaker() + self._attachment_points = AttachmentPoints() + self._randomize = randomize + self._conversions = Conversions() + self._sample_uniquely = sample_uniquely + + def run(self, scaffold_list: List[str]) -> List[SampledSequencesDTO]: + """ + Samples the model for the given number of SMILES. + :params scaffold_list: A list of scaffold SMILES. + :return: A list of SampledSequencesDTO. + """ + scaffold_list = self._randomize_scaffolds(scaffold_list) if self._randomize else scaffold_list + clean_scaffolds = [self._attachment_points.remove_attachment_point_numbers(scaffold) for scaffold in scaffold_list] + dataset = Dataset(clean_scaffolds, self.model.get_vocabulary().scaffold_vocabulary, + self.model.get_vocabulary().scaffold_tokenizer) + dataloader = DataLoader(dataset, batch_size=len(dataset), shuffle=False, collate_fn=Dataset.collate_fn) + + for batch in dataloader: + sampled_sequences = [] + + for _ in range(self._batch_size): + scaffold_seqs, scaffold_seq_lengths = batch + packed = self.model.sample(scaffold_seqs, scaffold_seq_lengths) + for scaffold, decoration, nll in packed: + sampled_sequences.append(SampledSequencesDTO(scaffold, decoration, nll)) + + if self._sample_uniquely: + sampled_sequences = self._sample_unique_sequences(sampled_sequences) + + return sampled_sequences + + def _sample_unique_sequences(self, sampled_sequences: List[SampledSequencesDTO]) -> List[SampledSequencesDTO]: + strings = ["".join([ss.input, ss.output]) for index, ss in enumerate(sampled_sequences)] + unique_idxs = self._get_indices_of_unique_smiles(strings) + sampled_sequences_np = np.array(sampled_sequences) + unique_sampled_sequences = sampled_sequences_np[unique_idxs] + return unique_sampled_sequences.tolist() + + def _randomize_scaffolds(self, scaffolds: List[str]): + scaffold_mols = [self._conversions.smile_to_mol(scaffold) for scaffold in scaffolds] + randomized = [self._bond_maker.randomize_scaffold(mol) for mol in scaffold_mols] + return randomized diff --git a/running_modes/automated_curriculum_learning/actions/link_invent_sample_model.py b/running_modes/automated_curriculum_learning/actions/link_invent_sample_model.py new file mode 100644 index 0000000..55f2cf4 --- /dev/null +++ b/running_modes/automated_curriculum_learning/actions/link_invent_sample_model.py @@ -0,0 +1,72 @@ +from typing import List + +import numpy as np +from reinvent_chemistry import Conversions, TransformationTokens +from reinvent_chemistry.library_design import BondMaker, AttachmentPoints +from reinvent_models.link_invent.dataset.dataset import Dataset +from reinvent_models.model_factory.generative_model_base import GenerativeModelBase +from torch.utils.data import DataLoader + +from running_modes.automated_curriculum_learning.actions import BaseSampleAction +from running_modes.automated_curriculum_learning.dto.sampled_sequences_dto import SampledSequencesDTO + + +class LinkInventSampleModel(BaseSampleAction): + def __init__(self, model: GenerativeModelBase, batch_size: int, logger=None, randomize=False, sample_uniquely=True): + """ + Creates an instance of SampleModel. + :params model: A model instance. + :params batch_size: Batch size to use. + :return: + """ + super().__init__(logger) + self.model = model + self._batch_size = batch_size + self._bond_maker = BondMaker() + self._randomize = randomize + self._sample_uniquely = sample_uniquely + + self._conversions = Conversions() + self._attachment_points = AttachmentPoints() + self._tokens = TransformationTokens() + + def run(self, warheads_list: List[str]) -> List[SampledSequencesDTO]: + """ + Samples the model for the given number of SMILES. + :params warheads_list: A list of warhead pair SMILES. + :return: A list of SampledSequencesDTO. + """ + warheads_list = self._randomize_warheads(warheads_list) if self._randomize else warheads_list + clean_warheads = [self._attachment_points.remove_attachment_point_numbers(warheads) for warheads in warheads_list] + dataset = Dataset(clean_warheads, self.model.get_vocabulary().input) + data_loader = DataLoader(dataset, batch_size=len(dataset), shuffle=False, collate_fn=dataset.collate_fn) + + for batch in data_loader: + sampled_sequences = [] + for _ in range(self._batch_size): + sampled_sequences.extend(self.model.sample(*batch)) + + if self._sample_uniquely: + sampled_sequences = self._sample_unique_sequences(sampled_sequences) + + return sampled_sequences + + def _sample_unique_sequences(self, sampled_sequences: List[SampledSequencesDTO]) -> List[SampledSequencesDTO]: + # TODO could be part of a base sample action as it is the same for link and lib invent + strings = ["".join([ss.input, ss.output]) for index, ss in enumerate(sampled_sequences)] + unique_idxs = self._get_indices_of_unique_smiles(strings) + sampled_sequences_np = np.array(sampled_sequences) + unique_sampled_sequences = sampled_sequences_np[unique_idxs] + return unique_sampled_sequences.tolist() + + def _randomize_warheads(self, warhead_pair_list: List[str]): + randomized_warhead_pair_list = [] + for warhead_pair in warhead_pair_list: + warhead_list = warhead_pair.split(self._tokens.ATTACHMENT_SEPARATOR_TOKEN) + warhead_mol_list = [self._conversions.smile_to_mol(warhead) for warhead in warhead_list] + warhead_randomized_list = [self._conversions.mol_to_random_smiles(mol) for mol in warhead_mol_list] + # Note do not use self.self._bond_maker.randomize_scaffold, as it would add unwanted brackets to the + # attachment points (which are not part of the warhead vocabulary) + warhead_pair_randomized = self._tokens.ATTACHMENT_SEPARATOR_TOKEN.join(warhead_randomized_list) + randomized_warhead_pair_list.append(warhead_pair_randomized) + return randomized_warhead_pair_list diff --git a/running_modes/automated_curriculum_learning/actions/reinvent_sample_model.py b/running_modes/automated_curriculum_learning/actions/reinvent_sample_model.py new file mode 100644 index 0000000..1f16828 --- /dev/null +++ b/running_modes/automated_curriculum_learning/actions/reinvent_sample_model.py @@ -0,0 +1,38 @@ +from typing import Tuple, Any + +import numpy as np +from reinvent_chemistry import Conversions +from reinvent_models.model_factory.generative_model_base import GenerativeModelBase + +from running_modes.automated_curriculum_learning.actions import BaseSampleAction +from running_modes.automated_curriculum_learning.dto import SampledBatchDTO + + +class ReinventSampleModel(BaseSampleAction): + def __init__(self, model: GenerativeModelBase, batch_size: int, logger=None): + """ + Creates an instance of SampleModel. + :params model: A model instance. + :params batch_size: Batch size to use. + :return: + """ + super().__init__(logger) + self.model = model + self._batch_size = batch_size + + self._conversions = Conversions() + + def run(self) -> SampledBatchDTO: + seqs, smiles, agent_likelihood = self._sample_unique_sequences(self.model, self._batch_size) + batch = SampledBatchDTO(seqs, smiles, agent_likelihood) + + return batch + + def _sample_unique_sequences(self, agent: GenerativeModelBase, batch_size: int) -> Tuple[Any, Any, Any]: + seqs, smiles, agent_likelihood = agent.sample(batch_size) + unique_idxs = self._get_indices_of_unique_smiles(smiles) + seqs_unique = seqs[unique_idxs] + smiles_np = np.array(smiles) + smiles_unique = smiles_np[unique_idxs] + agent_likelihood_unique = agent_likelihood[unique_idxs] + return seqs_unique, smiles_unique, agent_likelihood_unique diff --git a/running_modes/automated_curriculum_learning/automated_curriculum_runner.py b/running_modes/automated_curriculum_learning/automated_curriculum_runner.py index 9e7757e..53346bb 100644 --- a/running_modes/automated_curriculum_learning/automated_curriculum_runner.py +++ b/running_modes/automated_curriculum_learning/automated_curriculum_runner.py @@ -1,17 +1,18 @@ -from reinvent_models.reinvent_core.models.model import Model +from reinvent_models.model_factory.generative_model_base import GenerativeModelBase from reinvent_scoring import ScoringFunctionFactory -from reinvent_scoring.scoring.diversity_filters.reinvent_core.diversity_filter import DiversityFilter from running_modes.automated_curriculum_learning.curriculum_strategy.curriculum_strategy import CurriculumStrategy -from running_modes.automated_curriculum_learning.logging.base_auto_cl_logger import BaseAutoCLLogger +from running_modes.automated_curriculum_learning.inception.inception import Inception +from running_modes.automated_curriculum_learning.logging.base_logger import BaseLogger from running_modes.automated_curriculum_learning.production_strategy.production_strategy import ProductionStrategy -from running_modes.configurations import AutomatedCLConfiguration +from running_modes.configurations.automated_curriculum_learning.automated_curriculum_learning_input_configuration import \ + AutomatedCurriculumLearningInputConfiguration from running_modes.constructors.base_running_mode import BaseRunningMode -from running_modes.reinforcement_learning.inception import Inception class AutomatedCurriculumRunner(BaseRunningMode): - def __init__(self, config: AutomatedCLConfiguration, logger: BaseAutoCLLogger, prior: Model, agent: Model): + def __init__(self, config: AutomatedCurriculumLearningInputConfiguration, + logger: BaseLogger, prior: GenerativeModelBase, agent: GenerativeModelBase): self._config = config self._prior = prior @@ -21,21 +22,22 @@ def __init__(self, config: AutomatedCLConfiguration, logger: BaseAutoCLLogger, p self._logger) production_sf = ScoringFunctionFactory(self._config.production_strategy.scoring_function) - production_df = DiversityFilter(self._config.production_strategy.diversity_filter) if self._config.production_strategy.retain_inception: production_inception = self._curriculum_strategy.inception else: production_inception = Inception(self._config.production_strategy.inception, production_sf, self._prior) - self.production_strategy = ProductionStrategy(prior=self._prior, - diversity_filter=production_df, inception=production_inception, + self.production_strategy = ProductionStrategy(prior=self._prior, inception=production_inception, configuration=self._config.production_strategy, logger=self._logger) def run(self): self._logger.log_message("Starting Curriculum Learning") - cl_agent, step_counter = self._curriculum_strategy.run() - # Production - self._logger.log_message("Starting Production phase") - self.production_strategy.run(cl_agent, step_counter) + outcome = self._curriculum_strategy.run() + + if outcome.successful_curriculum: + self._logger.log_message("Starting Production phase") + self.production_strategy.run(outcome.agent, outcome.step_counter) + else: + self._logger.log_message("Failing to qualify for Production phase") \ No newline at end of file diff --git a/running_modes/automated_curriculum_learning/curriculum_strategy/base_curriculum_strategy.py b/running_modes/automated_curriculum_learning/curriculum_strategy/base_curriculum_strategy.py index 618103e..2a3c8c3 100644 --- a/running_modes/automated_curriculum_learning/curriculum_strategy/base_curriculum_strategy.py +++ b/running_modes/automated_curriculum_learning/curriculum_strategy/base_curriculum_strategy.py @@ -1,82 +1,63 @@ from abc import ABC, abstractmethod -from typing import Tuple, List, Any +from typing import Union -import numpy as np import torch -from reinvent_chemistry import get_indices_of_unique_smiles -from reinvent_models.reinvent_core.models.model import Model -from reinvent_scoring import ScoringFunctionFactory, FinalSummary -from reinvent_scoring.scoring.diversity_filters.reinvent_core.base_diversity_filter import BaseDiversityFilter -from reinvent_scoring.scoring.diversity_filters.reinvent_core.diversity_filter import DiversityFilter +from reinvent_chemistry.library_design import BondMaker, AttachmentPoints +from reinvent_models.model_factory.generative_model_base import GenerativeModelBase +from reinvent_scoring import ScoringFunctionFactory, Conversions +from reinvent_scoring.scoring.diversity_filters.curriculum_learning.base_diversity_filter import BaseDiversityFilter +from reinvent_scoring.scoring.diversity_filters.curriculum_learning.diversity_filter import DiversityFilter from reinvent_scoring.scoring.function.base_scoring_function import BaseScoringFunction -from running_modes.automated_curriculum_learning.logging.base_auto_cl_logger import BaseAutoCLLogger +from running_modes.automated_curriculum_learning.dto import CurriculumOutcomeDTO +from running_modes.automated_curriculum_learning.inception.inception import Inception +from running_modes.automated_curriculum_learning.learning_strategy.learning_strategy import LearningStrategy +from running_modes.automated_curriculum_learning.logging.base_logger import BaseLogger from running_modes.configurations import CurriculumStrategyConfiguration -from running_modes.reinforcement_learning.inception import Inception -from running_modes.utils import to_tensor +from running_modes.configurations.automated_curriculum_learning.curriculum_strategy_input_configuration import \ + CurriculumStrategyInputConfiguration class BaseCurriculumStrategy(ABC): - def __init__(self, prior: Model, agent: Model, diversity_filter: BaseDiversityFilter, - inception: Inception, configuration: CurriculumStrategyConfiguration, - logger: BaseAutoCLLogger): + def __init__(self, prior: GenerativeModelBase, agent: GenerativeModelBase, + configuration: Union[CurriculumStrategyConfiguration, CurriculumStrategyInputConfiguration], + diversity_filter: BaseDiversityFilter, inception: Inception, + logger: BaseLogger): + + self._bond_maker = BondMaker() + self._attachment_points = AttachmentPoints() + self._conversion = Conversions() self._parameters = configuration self._prior = prior self._agent = agent + self._logger = logger + optimizer = torch.optim.Adam(self._agent.get_network_parameters(), lr=self._parameters.learning_rate) + self.learning_strategy = LearningStrategy(self._prior, optimizer, configuration.learning_strategy, logger) self._diversity_filter = diversity_filter self.inception = inception - self._logger = logger - self._max_num_iterations = configuration.max_num_iterations @abstractmethod - def run(self) -> Tuple[Model, int]: - raise NotImplementedError("run not implemented for the current merging strategy") - - def _log_sf_update(self, current_parameters: List[dict], step: int): - text_to_log = f" Merging \n Step: {step}, #SF parameters: {len(current_parameters)} " \ - f"\n used components: {[component.get('name') for component in current_parameters]}" - self._logger.log_message(text_to_log) - - def _sample_unique_sequences(self, agent: Model, batch_size: int) -> Tuple[Any, Any, Any]: - seqs, smiles, agent_likelihood = agent.sample_sequences_and_smiles(batch_size) - unique_idxs = get_indices_of_unique_smiles(smiles) - seqs_unique = seqs[unique_idxs] - smiles_np = np.array(smiles) - smiles_unique = smiles_np[unique_idxs] - agent_likelihood_unique = agent_likelihood[unique_idxs] - return seqs_unique, smiles_unique, agent_likelihood_unique + def run(self) -> CurriculumOutcomeDTO: + raise NotImplementedError("run() method is not implemented ") - def disable_prior_gradients(self): - # There might be a more elegant way of disabling gradients - for param in self._prior.network.parameters(): - param.requires_grad = False + @abstractmethod + def take_step(self, agent: GenerativeModelBase, scoring_function: BaseScoringFunction, + step:int, start_time: float) -> float: + raise NotImplementedError("take_step() method is not implemented ") def save_and_flush_memory(self, agent, memory_name: str): self._logger.save_merging_state(agent, self._diversity_filter, name=memory_name) self._diversity_filter = DiversityFilter(self._parameters.diversity_filter) - def _stats_and_chekpoint(self, agent: Model, score: Any, start_time: float, step: int, smiles: List, - score_summary: FinalSummary, agent_likelihood: torch.tensor, - prior_likelihood: torch.tensor, - augmented_likelihood: torch.tensor): - mean_score = np.mean(score) - self._logger.timestep_report(start_time, self._max_num_iterations, step, smiles, - mean_score, score_summary, score, - agent_likelihood, prior_likelihood, augmented_likelihood, self._diversity_filter) - self._logger.save_checkpoint(step, self._diversity_filter, agent) - - def _inception_filter(self, agent, loss, agent_likelihood, prior_likelihood, - sigma, smiles, score): - if self.inception: - exp_smiles, exp_scores, exp_prior_likelihood = self.inception.sample() - if len(exp_smiles) > 0: - exp_agent_likelihood = -agent.likelihood_smiles(exp_smiles) - exp_augmented_likelihood = exp_prior_likelihood + sigma * exp_scores - exp_loss = torch.pow((to_tensor(exp_augmented_likelihood) - exp_agent_likelihood), 2) - loss = torch.cat((loss, exp_loss), 0) - agent_likelihood = torch.cat((agent_likelihood, exp_agent_likelihood), 0) - self.inception.add(smiles, score, prior_likelihood) - return loss, agent_likelihood + def disable_prior_gradients(self): + for param in self._prior.get_network_parameters(): + param.requires_grad = False + + def _setup_scoring_function(self, item_id: int) -> BaseScoringFunction: + parameters = self._parameters.curriculum_objectives[item_id] + scoring_function_instance = ScoringFunctionFactory(parameters.scoring_function) + self._logger.log_message(f"Loading a curriculum step number {item_id}") + return scoring_function_instance def _is_ready_to_promote(self, promotion_threshold: float, score: float) -> bool: if score >= promotion_threshold: @@ -92,41 +73,11 @@ def _is_step_quota_exceeded(self, current_step: int) -> bool: return True return False - def _setup_scoring_function(self, item_id: int) -> BaseScoringFunction: - parameters = self._parameters.curriculum_objectives[item_id] - scoring_function_instance = ScoringFunctionFactory(parameters.scoring_function) - self._logger.log_message(f"Loading a curriculum step number {item_id}") - return scoring_function_instance - - def take_step(self, agent: Model, optimiser: Any, scoring_function_current: BaseScoringFunction, step: int, - start_time: float) -> float: - seqs, smiles, agent_likelihood = self._sample_unique_sequences(agent, self._parameters.batch_size) - agent_likelihood = -agent_likelihood - prior_likelihood = -self._prior.likelihood(seqs) - score_summary_current: FinalSummary = scoring_function_current.get_final_score_for_step(smiles, step) - score = self._diversity_filter.update_score(score_summary_current, step) - augmented_likelihood = prior_likelihood + self._parameters.sigma * to_tensor(score) - loss = torch.pow((augmented_likelihood - agent_likelihood), 2) - loss, agent_likelihood = self._inception_filter(agent, loss, agent_likelihood, prior_likelihood, - self._parameters.sigma, smiles, score) - - loss = loss.mean() - optimiser.zero_grad() - loss.backward() - optimiser.step() - - self._stats_and_chekpoint(agent=agent, score=score, start_time=start_time, step=step, smiles=smiles, - score_summary=score_summary_current, agent_likelihood=agent_likelihood, - prior_likelihood=prior_likelihood, augmented_likelihood=augmented_likelihood) - - score = score.mean() - return score - - def promote_agent(self, agent: Model, optimiser: Any, scoring_function: BaseScoringFunction, step_counter: int, - start_time: float, merging_threshold: float) -> int: + def promote_agent(self, agent: GenerativeModelBase, scoring_function: BaseScoringFunction, + step_counter: int, start_time: float, merging_threshold: float) -> int: while not self._is_step_quota_exceeded(step_counter): - score = self.take_step(agent=agent, optimiser=optimiser, scoring_function_current=scoring_function, + score = self.take_step(agent=agent, scoring_function=scoring_function, step=step_counter, start_time=start_time) if self._is_ready_to_promote(merging_threshold, score): break diff --git a/running_modes/automated_curriculum_learning/curriculum_strategy/curriculum_strategy.py b/running_modes/automated_curriculum_learning/curriculum_strategy/curriculum_strategy.py index 974e7e8..ad68c9a 100644 --- a/running_modes/automated_curriculum_learning/curriculum_strategy/curriculum_strategy.py +++ b/running_modes/automated_curriculum_learning/curriculum_strategy/curriculum_strategy.py @@ -1,28 +1,43 @@ -from reinvent_models.reinvent_core.models.model import Model +from reinvent_models.model_factory.generative_model_base import GenerativeModelBase from reinvent_scoring import ScoringFunctionFactory -from reinvent_scoring.scoring.diversity_filters.reinvent_core.diversity_filter import DiversityFilter +from reinvent_scoring.scoring.diversity_filters.curriculum_learning.diversity_filter import DiversityFilter from running_modes.automated_curriculum_learning.curriculum_strategy.base_curriculum_strategy import \ BaseCurriculumStrategy -from running_modes.automated_curriculum_learning.curriculum_strategy.user_defined_curriculum import \ - UserDefinedCurriculum -from running_modes.automated_curriculum_learning.logging.base_auto_cl_logger import BaseAutoCLLogger -from running_modes.configurations import CurriculumStrategyConfiguration +from running_modes.automated_curriculum_learning.curriculum_strategy.linkinvent_curriculum_strategy import \ + LinkInventCurriculumStrategy +from running_modes.automated_curriculum_learning.curriculum_strategy.no_curriculum_strategy import NoCurriculumStrategy +from running_modes.automated_curriculum_learning.curriculum_strategy.reinvent_curriculum_strategy import \ + ReinventCurriculumStrategy +from running_modes.automated_curriculum_learning.inception.inception import Inception +from running_modes.automated_curriculum_learning.logging.base_logger import BaseLogger +from running_modes.configurations.automated_curriculum_learning.curriculum_strategy_input_configuration import \ + CurriculumStrategyInputConfiguration from running_modes.enums.curriculum_strategy_enum import CurriculumStrategyEnum -from running_modes.reinforcement_learning.inception import Inception class CurriculumStrategy: - def __new__(cls, prior: Model, agent: Model, configuration: CurriculumStrategyConfiguration, - logger: BaseAutoCLLogger) -> BaseCurriculumStrategy: - curriculum_strategy_enum = CurriculumStrategyEnum() + def __new__(cls, prior: GenerativeModelBase, agent: GenerativeModelBase, + configuration: CurriculumStrategyInputConfiguration, + logger: BaseLogger) -> BaseCurriculumStrategy: - diversity_filter = DiversityFilter(configuration.diversity_filter) - first_scoring_function_instance = ScoringFunctionFactory(configuration.curriculum_objectives[0].scoring_function) + curriculum_strategy_enum = CurriculumStrategyEnum() + first_objective = configuration.curriculum_objectives[0] + first_scoring_function_instance = ScoringFunctionFactory(first_objective.scoring_function) inception = Inception(configuration.inception, first_scoring_function_instance, prior) + diversity_filter = DiversityFilter(configuration.diversity_filter) if curriculum_strategy_enum.USER_DEFINED == configuration.name: - return UserDefinedCurriculum(prior=prior, agent=agent, configuration=configuration, - diversity_filter=diversity_filter, logger=logger, inception=inception) + cl_strategy = ReinventCurriculumStrategy(prior=prior, agent=agent, configuration=configuration, + diversity_filter=diversity_filter, logger=logger, inception=inception) + return cl_strategy + elif curriculum_strategy_enum.LINK_INVENT == configuration.name: + cl_strategy = LinkInventCurriculumStrategy(prior=prior, agent=agent, configuration=configuration, + diversity_filter=diversity_filter, logger=logger, inception=inception) + return cl_strategy + elif curriculum_strategy_enum.NO_CURRICULUM == configuration.name: + cl_strategy = NoCurriculumStrategy(prior=prior, agent=agent, configuration=configuration, + diversity_filter=diversity_filter, logger=logger, inception=inception) + return cl_strategy else: raise NotImplementedError(f"Unknown curriculum strategy name {configuration.name}") diff --git a/running_modes/automated_curriculum_learning/curriculum_strategy/linkinvent_curriculum_strategy.py b/running_modes/automated_curriculum_learning/curriculum_strategy/linkinvent_curriculum_strategy.py new file mode 100644 index 0000000..091d080 --- /dev/null +++ b/running_modes/automated_curriculum_learning/curriculum_strategy/linkinvent_curriculum_strategy.py @@ -0,0 +1,110 @@ +import time +from typing import List + +from reinvent_models.model_factory.generative_model_base import GenerativeModelBase +from reinvent_scoring import FinalSummary +from reinvent_scoring.scoring.diversity_filters.curriculum_learning.loggable_data_dto import UpdateLoggableDataDTO +from reinvent_scoring.scoring.diversity_filters.curriculum_learning.update_diversity_filter_dto import \ + UpdateDiversityFilterDTO +from reinvent_scoring.scoring.function.base_scoring_function import BaseScoringFunction + +from running_modes.automated_curriculum_learning.actions import LinkInventSampleModel +from running_modes.automated_curriculum_learning.curriculum_strategy.base_curriculum_strategy import \ + BaseCurriculumStrategy +from running_modes.automated_curriculum_learning.dto import SampledSequencesDTO, CurriculumOutcomeDTO, TimestepDTO, \ + UpdatedLikelihoodsDTO + + +class LinkInventCurriculumStrategy(BaseCurriculumStrategy): + + def run(self) -> CurriculumOutcomeDTO: + step_counter = 0 + self.disable_prior_gradients() + + for item_id, sf_configuration in enumerate(self._parameters.curriculum_objectives): + start_time = time.time() + scoring_function = self._setup_scoring_function(item_id) + step_counter = self.promote_agent(agent=self._agent, scoring_function=scoring_function, + step_counter=step_counter, start_time=start_time, + merging_threshold=sf_configuration.score_threshold) + self.save_and_flush_memory(agent=self._agent, memory_name=f"_merge_{item_id}") + is_successful_curriculum = step_counter < self._parameters.max_num_iterations + outcome_dto = CurriculumOutcomeDTO(self._agent, step_counter, successful_curriculum=is_successful_curriculum) + + return outcome_dto + + def take_step(self, agent: GenerativeModelBase, scoring_function: BaseScoringFunction, + step:int, start_time: float) -> float: + # 1. Sampling + sampled_sequences = self._sampling(agent) + # 2. Scoring + score_summary = self._scoring(scoring_function, sampled_sequences, step) + # 3. Updating + dto = self._updating(sampled_sequences, score_summary.total_score) + # 4. Logging + self._logging(start_time, step, score_summary, dto) + + score = score_summary.total_score.mean() + return score + + def _sampling(self, agent) -> List[SampledSequencesDTO]: + sampling_action = LinkInventSampleModel(agent, self._parameters.batch_size, self._logger, + self._parameters.randomize_input) + sampled_sequences = sampling_action.run(self._parameters.input) + return sampled_sequences + + def _scoring(self, scoring_function, sampled_sequences, step: int) -> FinalSummary: + score_summary = self._apply_scoring_function(scoring_function, sampled_sequences, step) + score_summary = self._clean_scored_smiles(score_summary) + loggable_data = [UpdateLoggableDataDTO(dto.input, dto.output, dto.nll) for dto in sampled_sequences] + dto = UpdateDiversityFilterDTO(score_summary, loggable_data, step) + score_summary.total_score = self._diversity_filter.update_score(dto) + return score_summary + + def _updating(self, sampled_sequences, score) -> UpdatedLikelihoodsDTO: + likelihood_dto = self._agent.likelihood_smiles(sampled_sequences) + dto = self.learning_strategy.run(likelihood_dto, score) + return dto + + def _logging(self, start_time, step, score_summary, dto: UpdatedLikelihoodsDTO): + report_dto = TimestepDTO(start_time, self._parameters.max_num_iterations, step, score_summary, + dto.agent_likelihood, dto.prior_likelihood, dto.augmented_likelihood) + self._logger.timestep_report(report_dto, self._diversity_filter, self._agent) + + def _apply_scoring_function(self, scoring_function, sampled_sequences: List[SampledSequencesDTO], step) -> FinalSummary: + molecules = self._join_linker_and_warheads(sampled_sequences, keep_labels=True) + smiles = [] + for idx, molecule in enumerate(molecules): + try: + smiles_str = self._conversion.mol_to_smiles(molecule) if molecule else "INVALID" + except RuntimeError as exception: + # NOTE: Current implementation of BondMaker (reinvent_chemistry.library_design.bond_maker) results in + # impossible conversion of mol to smiles if one single atom has two attachment points and labels are + # kept. As this case is not relevant in the context of link_invent, the can be discarded as invalid. + smiles_str = "INVALID" + self._logger.log_message(exception.__str__() + f'\n\tinput: {sampled_sequences[idx].input}' + f'\n\toutput: {sampled_sequences[idx].output}\n') + finally: + smiles.append(smiles_str) + final_score: FinalSummary = scoring_function.get_final_score_for_step(smiles, step) + return final_score + + def _join_linker_and_warheads(self, sampled_sequences: List[SampledSequencesDTO], keep_labels=False): + molecules = [] + for sample in sampled_sequences: + linker = self._attachment_points.add_attachment_point_numbers(sample.output, canonicalize=False) + molecule = self._bond_maker.join_scaffolds_and_decorations(linker, sample.input, + keep_labels_on_atoms=keep_labels) + molecules.append(molecule) + return molecules + + def _clean_scored_smiles(self, score_summary: FinalSummary) -> FinalSummary: + """ + Remove attachment point numbers from scored smiles + """ + # Note: method AttachmentPoints.remove_attachment_point_numbers does not work in this context, as it searches + # for attachment point token ('*') + score_summary.scored_smiles = [self._conversion.mol_to_smiles( + self._attachment_points.remove_attachment_point_numbers_from_mol(self._conversion.smile_to_mol(smile)) + ) if idx in score_summary.valid_idxs else smile for idx, smile in enumerate(score_summary.scored_smiles)] + return score_summary diff --git a/running_modes/automated_curriculum_learning/curriculum_strategy/no_curriculum_strategy.py b/running_modes/automated_curriculum_learning/curriculum_strategy/no_curriculum_strategy.py new file mode 100644 index 0000000..1d5b52d --- /dev/null +++ b/running_modes/automated_curriculum_learning/curriculum_strategy/no_curriculum_strategy.py @@ -0,0 +1,24 @@ +from reinvent_models.model_factory.generative_model_base import GenerativeModelBase +from reinvent_scoring.scoring.diversity_filters.curriculum_learning.diversity_filter import DiversityFilter +from reinvent_scoring.scoring.function.base_scoring_function import BaseScoringFunction + +from running_modes.automated_curriculum_learning.curriculum_strategy.base_curriculum_strategy import \ + BaseCurriculumStrategy +from running_modes.automated_curriculum_learning.dto import CurriculumOutcomeDTO + + +class NoCurriculumStrategy(BaseCurriculumStrategy): + + def run(self) -> CurriculumOutcomeDTO: + step_counter = 0 + self.disable_prior_gradients() + outcome_dto = CurriculumOutcomeDTO(self._agent, step_counter, successful_curriculum=True) + return outcome_dto + + def take_step(self, agent: GenerativeModelBase, scoring_function: BaseScoringFunction, + step:int, start_time: float) -> float: + pass + + def save_and_flush_memory(self, agent, memory_name: str): + self._logger.save_merging_state(agent, self._diversity_filter, name=memory_name) + self._diversity_filter = DiversityFilter(self._parameters.diversity_filter) diff --git a/running_modes/automated_curriculum_learning/curriculum_strategy/reinvent_curriculum_strategy.py b/running_modes/automated_curriculum_learning/curriculum_strategy/reinvent_curriculum_strategy.py new file mode 100644 index 0000000..2f61fb6 --- /dev/null +++ b/running_modes/automated_curriculum_learning/curriculum_strategy/reinvent_curriculum_strategy.py @@ -0,0 +1,77 @@ +import time +from typing import List, Tuple + +import numpy as np +import torch +from reinvent_models.model_factory.generative_model_base import GenerativeModelBase +from reinvent_scoring import FinalSummary +from reinvent_scoring.scoring.diversity_filters.curriculum_learning.diversity_filter import DiversityFilter +from reinvent_scoring.scoring.diversity_filters.curriculum_learning.update_diversity_filter_dto import \ + UpdateDiversityFilterDTO +from reinvent_scoring.scoring.function.base_scoring_function import BaseScoringFunction + +from running_modes.automated_curriculum_learning.actions.reinvent_sample_model import ReinventSampleModel +from running_modes.automated_curriculum_learning.curriculum_strategy.base_curriculum_strategy import \ + BaseCurriculumStrategy +from running_modes.automated_curriculum_learning.dto import SampledBatchDTO, CurriculumOutcomeDTO, TimestepDTO + + +class ReinventCurriculumStrategy(BaseCurriculumStrategy): + + def run(self) -> CurriculumOutcomeDTO: + step_counter = 0 + self.disable_prior_gradients() + + for item_id, sf_configuration in enumerate(self._parameters.curriculum_objectives): + start_time = time.time() + scoring_function = self._setup_scoring_function(item_id) + step_counter = self.promote_agent(agent=self._agent, scoring_function=scoring_function, + step_counter=step_counter, start_time=start_time, + merging_threshold=sf_configuration.score_threshold) + self.save_and_flush_memory(agent=self._agent, memory_name=f"_merge_{item_id}") + is_successful_curriculum = step_counter < self._parameters.max_num_iterations + outcome_dto = CurriculumOutcomeDTO(self._agent, step_counter, successful_curriculum=is_successful_curriculum) + + return outcome_dto + + def take_step(self, agent: GenerativeModelBase, scoring_function: BaseScoringFunction, + step:int, start_time: float) -> float: + # 1. Sampling + sampled = self._sampling(agent) + # 2. Scoring + score, score_summary = self._scoring(scoring_function, sampled.smiles, step) + # 3. Updating + agent_likelihood, prior_likelihood, augmented_likelihood = self._updating(sampled, score, self.inception, agent) + # 4. Logging + self._logging(agent=agent, start_time=start_time, step=step, + score_summary=score_summary, agent_likelihood=agent_likelihood, + prior_likelihood=prior_likelihood, augmented_likelihood=augmented_likelihood) + + score = score.mean() + return score + + def _sampling(self, agent) -> SampledBatchDTO: + sampling_action = ReinventSampleModel(agent, self._parameters.batch_size, self._logger) + sampled_sequences = sampling_action.run() + return sampled_sequences + + def _scoring(self, scoring_function, smiles: List[str], step) -> Tuple[np.ndarray, FinalSummary] : + score_summary = scoring_function.get_final_score_for_step(smiles, step) + dto = UpdateDiversityFilterDTO(score_summary, [], step) + score = self._diversity_filter.update_score(dto) + return score, score_summary + + def _updating(self, sampled, score, inception, agent): + agent_likelihood, prior_likelihood, augmented_likelihood = \ + self.learning_strategy.run(sampled, score, inception, agent) + return agent_likelihood, prior_likelihood, augmented_likelihood + + def _logging(self, agent: GenerativeModelBase, start_time: float, step: int, score_summary: FinalSummary, + agent_likelihood: torch.tensor, prior_likelihood: torch.tensor, augmented_likelihood: torch.tensor): + report_dto = TimestepDTO(start_time, self._parameters.max_num_iterations, step, score_summary, + agent_likelihood, prior_likelihood, augmented_likelihood) + self._logger.timestep_report(report_dto, self._diversity_filter, agent) + + def save_and_flush_memory(self, agent, memory_name: str): + self._logger.save_merging_state(agent, self._diversity_filter, name=memory_name) + self._diversity_filter = DiversityFilter(self._parameters.diversity_filter) diff --git a/running_modes/automated_curriculum_learning/curriculum_strategy/user_defined_curriculum.py b/running_modes/automated_curriculum_learning/curriculum_strategy/user_defined_curriculum.py deleted file mode 100644 index fccdf7a..0000000 --- a/running_modes/automated_curriculum_learning/curriculum_strategy/user_defined_curriculum.py +++ /dev/null @@ -1,34 +0,0 @@ -import time -from typing import Tuple - -import torch -from reinvent_models.reinvent_core.models.model import Model -from reinvent_scoring.scoring.diversity_filters.reinvent_core.base_diversity_filter import BaseDiversityFilter - -from running_modes.automated_curriculum_learning.curriculum_strategy.base_curriculum_strategy import \ - BaseCurriculumStrategy -from running_modes.automated_curriculum_learning.logging.base_auto_cl_logger import BaseAutoCLLogger -from running_modes.configurations import CurriculumStrategyConfiguration -from running_modes.reinforcement_learning.inception import Inception - - -class UserDefinedCurriculum(BaseCurriculumStrategy): - def __init__(self, prior: Model, agent: Model, diversity_filter: BaseDiversityFilter, - inception: Inception, configuration: CurriculumStrategyConfiguration, logger: BaseAutoCLLogger): - super().__init__(prior=prior, agent=agent, configuration=configuration, diversity_filter=diversity_filter, - inception=inception, logger=logger) - - def run(self) -> Tuple[Model, int]: - optimiser = torch.optim.Adam(self._agent.network.parameters(), lr=self._parameters.learning_rate) - step_counter = 0 - self.disable_prior_gradients() - - for item_id, sf_configuration in enumerate(self._parameters.curriculum_objectives): - start_time = time.time() - scoring_function = self._setup_scoring_function(item_id) - step_counter = self.promote_agent(agent=self._agent, optimiser=optimiser, scoring_function=scoring_function, - step_counter=step_counter, start_time=start_time, - merging_threshold=sf_configuration.score_threshold) - self.save_and_flush_memory(agent=self._agent, memory_name=f"_merge_{item_id}") - - return self._agent, step_counter diff --git a/running_modes/automated_curriculum_learning/dto/__init__.py b/running_modes/automated_curriculum_learning/dto/__init__.py new file mode 100644 index 0000000..28fde81 --- /dev/null +++ b/running_modes/automated_curriculum_learning/dto/__init__.py @@ -0,0 +1,5 @@ +from running_modes.automated_curriculum_learning.dto.curriculum_outcome_dto import CurriculumOutcomeDTO +from running_modes.automated_curriculum_learning.dto.sampled_batch_dto import SampledBatchDTO +from running_modes.automated_curriculum_learning.dto.sampled_sequences_dto import SampledSequencesDTO +from running_modes.automated_curriculum_learning.dto.timestep_dto import TimestepDTO +from running_modes.automated_curriculum_learning.dto.updated_likelihoods_dto import UpdatedLikelihoodsDTO diff --git a/running_modes/automated_curriculum_learning/dto/curriculum_outcome_dto.py b/running_modes/automated_curriculum_learning/dto/curriculum_outcome_dto.py new file mode 100644 index 0000000..128f7fb --- /dev/null +++ b/running_modes/automated_curriculum_learning/dto/curriculum_outcome_dto.py @@ -0,0 +1,9 @@ +from dataclasses import dataclass +from reinvent_models.model_factory.generative_model_base import GenerativeModelBase + + +@dataclass +class CurriculumOutcomeDTO: + agent: GenerativeModelBase + step_counter: int + successful_curriculum: bool diff --git a/running_modes/automated_curriculum_learning/dto/sampled_batch_dto.py b/running_modes/automated_curriculum_learning/dto/sampled_batch_dto.py new file mode 100644 index 0000000..7c4b37c --- /dev/null +++ b/running_modes/automated_curriculum_learning/dto/sampled_batch_dto.py @@ -0,0 +1,12 @@ +from dataclasses import dataclass +from typing import List + +import numpy as np +import torch + + +@dataclass +class SampledBatchDTO: + sequences: torch.Tensor + smiles: List[str] + likelihoods: torch.Tensor diff --git a/running_modes/automated_curriculum_learning/dto/sampled_sequences_dto.py b/running_modes/automated_curriculum_learning/dto/sampled_sequences_dto.py new file mode 100644 index 0000000..8632b23 --- /dev/null +++ b/running_modes/automated_curriculum_learning/dto/sampled_sequences_dto.py @@ -0,0 +1,8 @@ +from dataclasses import dataclass + + +@dataclass +class SampledSequencesDTO: + input: str + output: str + nll: float diff --git a/running_modes/automated_curriculum_learning/dto/timestep_dto.py b/running_modes/automated_curriculum_learning/dto/timestep_dto.py new file mode 100644 index 0000000..55708f5 --- /dev/null +++ b/running_modes/automated_curriculum_learning/dto/timestep_dto.py @@ -0,0 +1,15 @@ +from dataclasses import dataclass + +import torch +from reinvent_scoring.scoring import FinalSummary + + +@dataclass +class TimestepDTO: + start_time: float + n_steps: int + step: int + score_summary: FinalSummary + agent_likelihood: torch.tensor + prior_likelihood: torch.tensor + augmented_likelihood: torch.tensor \ No newline at end of file diff --git a/running_modes/automated_curriculum_learning/dto/updated_likelihoods_dto.py b/running_modes/automated_curriculum_learning/dto/updated_likelihoods_dto.py new file mode 100644 index 0000000..9dbf02a --- /dev/null +++ b/running_modes/automated_curriculum_learning/dto/updated_likelihoods_dto.py @@ -0,0 +1,10 @@ +from dataclasses import dataclass +import torch + + +@dataclass +class UpdatedLikelihoodsDTO: + agent_likelihood: torch.Tensor + prior_likelihood: torch.Tensor + augmented_likelihood: torch.Tensor + loss: torch.Tensor \ No newline at end of file diff --git a/running_modes/automated_curriculum_learning/inception/__init__.py b/running_modes/automated_curriculum_learning/inception/__init__.py new file mode 100644 index 0000000..e69de29 diff --git a/running_modes/automated_curriculum_learning/inception/inception.py b/running_modes/automated_curriculum_learning/inception/inception.py new file mode 100644 index 0000000..f538939 --- /dev/null +++ b/running_modes/automated_curriculum_learning/inception/inception.py @@ -0,0 +1,50 @@ +import numpy as np +import pandas as pd +from typing import Tuple, List + +from reinvent_chemistry.conversions import Conversions + +from running_modes.configurations import InceptionConfiguration + + +class Inception: + def __init__(self, configuration: InceptionConfiguration, scoring_function, prior): + self.configuration = configuration + self._chemistry = Conversions() + self.memory: pd.DataFrame = pd.DataFrame(columns=['smiles', 'score', 'likelihood']) + self._load_to_memory(scoring_function, prior, self.configuration.smiles) + + + def _load_to_memory(self, scoring_function, prior, smiles): + if len(smiles): + standardized_and_nulls = [self._chemistry.convert_to_rdkit_smiles(smile) for smile in smiles] + standardized = [smile for smile in standardized_and_nulls if smile is not None] + self.evaluate_and_add(standardized, scoring_function, prior) + + def _purge_memory(self): + unique_df = self.memory.drop_duplicates(subset=["smiles"]) + sorted_df = unique_df.sort_values('score', ascending=False) + self.memory = sorted_df.head(self.configuration.memory_size) + + def evaluate_and_add(self, smiles, scoring_function, prior): + if len(smiles) > 0: + score = scoring_function.get_final_score(smiles) + likelihood = prior.likelihood_smiles(smiles) + df = pd.DataFrame({"smiles": smiles, "score": score.total_score, "likelihood": -likelihood.detach().cpu().numpy()}) + self.memory = self.memory.append(df) + self._purge_memory() + + def add(self, smiles, score, neg_likelihood): + df = pd.DataFrame({"smiles": smiles, "score": score, "likelihood": neg_likelihood.detach().cpu().numpy()}) + self.memory = self.memory.append(df) + self._purge_memory() + + def sample(self) -> Tuple[List[str], np.array, np.array]: + sample_size = min(len(self.memory), self.configuration.sample_size) + if sample_size > 0: + sampled = self.memory.sample(sample_size) + smiles = sampled["smiles"].values + scores = sampled["score"].values + prior_likelihood = sampled["likelihood"].values + return smiles, scores, prior_likelihood + return [], [], [] diff --git a/running_modes/automated_curriculum_learning/learning_strategy/__init__.py b/running_modes/automated_curriculum_learning/learning_strategy/__init__.py new file mode 100644 index 0000000..7031453 --- /dev/null +++ b/running_modes/automated_curriculum_learning/learning_strategy/__init__.py @@ -0,0 +1,6 @@ +from running_modes.automated_curriculum_learning.learning_strategy.dap_single_query_strategy import \ + DAPSingleQueryStrategy +from running_modes.automated_curriculum_learning.learning_strategy.dap_strategy import DAPStrategy +from running_modes.automated_curriculum_learning.learning_strategy.mascof_strategy import MASCOFStrategy +from running_modes.automated_curriculum_learning.learning_strategy.mauli_strategy import MAULIStrategy +from running_modes.automated_curriculum_learning.learning_strategy.sdap_strategy import SDAPStrategy diff --git a/running_modes/automated_curriculum_learning/learning_strategy/base_double_query_learning_strategy.py b/running_modes/automated_curriculum_learning/learning_strategy/base_double_query_learning_strategy.py new file mode 100644 index 0000000..eee8251 --- /dev/null +++ b/running_modes/automated_curriculum_learning/learning_strategy/base_double_query_learning_strategy.py @@ -0,0 +1,25 @@ +from abc import abstractmethod + +from reinvent_models.link_invent.dto.batch_likelihood_dto import BatchLikelihoodDTO + +from running_modes.automated_curriculum_learning.dto import UpdatedLikelihoodsDTO +from running_modes.automated_curriculum_learning.learning_strategy.base_learning_strategy import BaseLearningStrategy +from running_modes.automated_curriculum_learning.learning_strategy.learning_strategy_configuration import \ + LearningStrategyConfiguration + + +class BaseDoubleQueryLearningStrategy(BaseLearningStrategy): + def __init__(self, critic_model, optimizer, configuration: LearningStrategyConfiguration, logger): + super().__init__(critic_model, optimizer, configuration, logger) + + def run(self, likelihood_dto: BatchLikelihoodDTO, score) -> UpdatedLikelihoodsDTO: + dto = self._calculate_loss(likelihood_dto, score) + self.optimizer.zero_grad() + dto.loss.backward() + + self.optimizer.step() + return dto + + @abstractmethod + def _calculate_loss(self, likelihood_dto: BatchLikelihoodDTO, score) -> UpdatedLikelihoodsDTO: + raise NotImplementedError("_calculate_loss method is not implemented") diff --git a/running_modes/automated_curriculum_learning/learning_strategy/base_learning_strategy.py b/running_modes/automated_curriculum_learning/learning_strategy/base_learning_strategy.py new file mode 100644 index 0000000..f26ae52 --- /dev/null +++ b/running_modes/automated_curriculum_learning/learning_strategy/base_learning_strategy.py @@ -0,0 +1,40 @@ +from abc import ABC, abstractmethod +from typing import Tuple + +import torch +from reinvent_models.lib_invent.enums.generative_model_regime import GenerativeModelRegimeEnum + +from running_modes.automated_curriculum_learning.learning_strategy.learning_strategy_configuration import \ + LearningStrategyConfiguration + + +class BaseLearningStrategy(ABC): + def __init__(self, critic_model, optimizer, configuration: LearningStrategyConfiguration, logger): + self.critic_model = critic_model + self.optimizer = optimizer + self._configuration = configuration + self._running_mode_enum = GenerativeModelRegimeEnum() + self._logger = logger + self._disable_prior_gradients() + + def log_message(self, message: str): + self._logger.log_message(message) + + @abstractmethod + def run(self, *args, **kwargs) -> Tuple[torch.Tensor, torch.Tensor, torch.Tensor]: + raise NotImplementedError("run() method is not implemented") + + @abstractmethod + def _calculate_loss(self, *args, **kwargs): + raise NotImplementedError("_calculate_loss() method is not implemented") + + def _to_tensor(self, array, use_cuda=True): + if torch.cuda.is_available() and use_cuda: + return torch.tensor(array, device=torch.device("cuda")) + return torch.tensor(array, device=torch.device("cpu")) + + def _disable_prior_gradients(self): + # There might be a more elegant way of disabling gradients + self.critic_model.set_mode(self._running_mode_enum.INFERENCE) + for param in self.critic_model.network.parameters(): + param.requires_grad = False diff --git a/running_modes/automated_curriculum_learning/learning_strategy/base_linker_strategy.py b/running_modes/automated_curriculum_learning/learning_strategy/base_linker_strategy.py new file mode 100644 index 0000000..b291ee0 --- /dev/null +++ b/running_modes/automated_curriculum_learning/learning_strategy/base_linker_strategy.py @@ -0,0 +1,30 @@ +from abc import abstractmethod +from typing import Tuple + +import numpy as np +import torch + +from running_modes.automated_curriculum_learning.learning_strategy.base_learning_strategy import BaseLearningStrategy +from running_modes.automated_curriculum_learning.learning_strategy.learning_strategy_configuration import \ + LearningStrategyConfiguration + + +class BaseLinkerStrategy(BaseLearningStrategy): + def __init__(self, critic_model, optimizer, configuration: LearningStrategyConfiguration, logger): + super().__init__(critic_model, optimizer, configuration, logger) + + # TODO: Return the loss as well. + def run(self, scaffold_batch: np.ndarray, decorator_batch: np.ndarray, + score: torch.Tensor, actor_nlls: torch.Tensor) -> Tuple[torch.Tensor, torch.Tensor, torch.Tensor]: + loss, negative_actor_nlls, negative_critic_nlls, augmented_nlls = \ + self._calculate_loss(scaffold_batch, decorator_batch, score, actor_nlls) + + self.optimizer.zero_grad() + loss.backward() + + self.optimizer.step() + return negative_actor_nlls, negative_critic_nlls, augmented_nlls + + @abstractmethod + def _calculate_loss(self, scaffold_batch, decorator_batch, score, actor_nlls): + raise NotImplementedError("_calculate_loss method is not implemented") diff --git a/running_modes/automated_curriculum_learning/learning_strategy/base_single_query_learning_strategy.py b/running_modes/automated_curriculum_learning/learning_strategy/base_single_query_learning_strategy.py new file mode 100644 index 0000000..94a3908 --- /dev/null +++ b/running_modes/automated_curriculum_learning/learning_strategy/base_single_query_learning_strategy.py @@ -0,0 +1,30 @@ +from abc import abstractmethod +from typing import Tuple +import numpy as np +import torch + +from running_modes.automated_curriculum_learning.dto import SampledBatchDTO +from running_modes.automated_curriculum_learning.learning_strategy.base_learning_strategy import BaseLearningStrategy +from running_modes.automated_curriculum_learning.learning_strategy.learning_strategy_configuration import \ + LearningStrategyConfiguration + + +class BaseSingleQueryLearningStrategy(BaseLearningStrategy): + def __init__(self, critic_model, optimizer, configuration: LearningStrategyConfiguration, logger): + super().__init__(critic_model, optimizer, configuration, logger) + + # TODO: Return the loss as well. + def run(self, sampled: SampledBatchDTO , score: torch.Tensor, inception, agent) -> \ + Tuple[torch.Tensor, torch.Tensor, torch.Tensor]: + loss, negative_actor_nlls, negative_critic_nlls, augmented_nlls = \ + self._calculate_loss(sampled.smiles, sampled.sequences, score, sampled.likelihoods, inception, agent) + + self.optimizer.zero_grad() + loss.backward() + + self.optimizer.step() + return negative_actor_nlls, negative_critic_nlls, augmented_nlls + + @abstractmethod + def _calculate_loss(self, smiles, sampled_sequences: np.ndarray, score, actor_nlls, inception, agent): + raise NotImplementedError("_calculate_loss method is not implemented") diff --git a/running_modes/automated_curriculum_learning/learning_strategy/dap_single_query_strategy.py b/running_modes/automated_curriculum_learning/learning_strategy/dap_single_query_strategy.py new file mode 100644 index 0000000..805c3ed --- /dev/null +++ b/running_modes/automated_curriculum_learning/learning_strategy/dap_single_query_strategy.py @@ -0,0 +1,41 @@ +import numpy as np +import torch + +from running_modes.automated_curriculum_learning.learning_strategy.base_single_query_learning_strategy import \ + BaseSingleQueryLearningStrategy +from running_modes.automated_curriculum_learning.learning_strategy.learning_strategy_configuration import \ + LearningStrategyConfiguration + + +class DAPSingleQueryStrategy(BaseSingleQueryLearningStrategy): + + def __init__(self, critic_model, optimizer, configuration: LearningStrategyConfiguration, logger=None): + """ + TODO: Provide description of the current strategy + """ + super().__init__(critic_model, optimizer, configuration, logger) + + self._sigma = self._configuration.parameters.get("sigma", 120) + + def _calculate_loss(self, smiles, sampled_sequences: np.ndarray, score, actor_nlls, inception, agent): + critic_nlls = self.critic_model.likelihood(sampled_sequences) + negative_critic_nlls = -critic_nlls + negative_actor_nlls = -actor_nlls + augmented_nlls = negative_critic_nlls + self._sigma * self._to_tensor(score) + loss = torch.pow((augmented_nlls - negative_actor_nlls), 2) + loss, agent_likelihood = self._inception_filter(agent, loss, negative_actor_nlls, negative_critic_nlls, + self._sigma, smiles, score, inception) + loss = loss.mean() + return loss, negative_actor_nlls, negative_critic_nlls, augmented_nlls + + def _inception_filter(self, agent, loss, agent_likelihood, prior_likelihood, sigma, smiles, score, inception): + if inception: + exp_smiles, exp_scores, exp_prior_likelihood = inception.sample() + if len(exp_smiles) > 0: + exp_agent_likelihood = -agent.likelihood_smiles(exp_smiles) + exp_augmented_likelihood = exp_prior_likelihood + sigma * exp_scores + exp_loss = torch.pow((self._to_tensor(exp_augmented_likelihood) - exp_agent_likelihood), 2) + loss = torch.cat((loss, exp_loss), 0) + agent_likelihood = torch.cat((agent_likelihood, exp_agent_likelihood), 0) + inception.add(smiles, score, prior_likelihood) + return loss, agent_likelihood diff --git a/running_modes/automated_curriculum_learning/learning_strategy/dap_strategy.py b/running_modes/automated_curriculum_learning/learning_strategy/dap_strategy.py new file mode 100644 index 0000000..eecb3b4 --- /dev/null +++ b/running_modes/automated_curriculum_learning/learning_strategy/dap_strategy.py @@ -0,0 +1,33 @@ +import torch +from reinvent_models.link_invent.dto import BatchLikelihoodDTO + +from running_modes.automated_curriculum_learning.dto import UpdatedLikelihoodsDTO +from running_modes.automated_curriculum_learning.learning_strategy.base_double_query_learning_strategy import \ + BaseDoubleQueryLearningStrategy +from running_modes.automated_curriculum_learning.learning_strategy.learning_strategy_configuration import \ + LearningStrategyConfiguration + + +class DAPStrategy(BaseDoubleQueryLearningStrategy): + + def __init__(self, critic_model, optimizer, configuration: LearningStrategyConfiguration, logger=None): + """ + TODO: Provide description of the current strategy + """ + super().__init__(critic_model, optimizer, configuration, logger) + + self._sigma = self._configuration.parameters.get("sigma", 120) + + def _calculate_loss(self, likelihood_dto: BatchLikelihoodDTO, score) -> UpdatedLikelihoodsDTO: + batch = likelihood_dto.batch + critic_nlls = self.critic_model.likelihood(*batch.input, *batch.output) + negative_critic_nlls = -critic_nlls + negative_actor_nlls = -likelihood_dto.likelihood + augmented_nlls = negative_critic_nlls + self._sigma * self._to_tensor(score) + loss = torch.pow((augmented_nlls - negative_actor_nlls), 2) + loss = loss.mean() + dto = UpdatedLikelihoodsDTO(negative_actor_nlls, negative_critic_nlls, augmented_nlls, loss) + return dto + + + diff --git a/running_modes/automated_curriculum_learning/learning_strategy/learning_strategy.py b/running_modes/automated_curriculum_learning/learning_strategy/learning_strategy.py new file mode 100644 index 0000000..81fe517 --- /dev/null +++ b/running_modes/automated_curriculum_learning/learning_strategy/learning_strategy.py @@ -0,0 +1,23 @@ +from running_modes.automated_curriculum_learning.learning_strategy import DAPStrategy, MAULIStrategy, MASCOFStrategy, \ + SDAPStrategy, DAPSingleQueryStrategy +from running_modes.automated_curriculum_learning.learning_strategy.base_learning_strategy import BaseLearningStrategy +from running_modes.automated_curriculum_learning.learning_strategy.learning_strategy_configuration import \ + LearningStrategyConfiguration +from running_modes.automated_curriculum_learning.learning_strategy.learning_strategy_enum import LearningStrategyEnum + + +class LearningStrategy: + + def __new__(cls, critic_model, optimizer, configuration: LearningStrategyConfiguration, logger=None) \ + -> BaseLearningStrategy: + learning_strategy_enum = LearningStrategyEnum() + if learning_strategy_enum.DAP_SINGLE_QUERY == configuration.name: + return DAPSingleQueryStrategy(critic_model, optimizer, configuration, logger) + if learning_strategy_enum.DAP == configuration.name: + return DAPStrategy(critic_model, optimizer, configuration, logger) + if learning_strategy_enum.MAULI == configuration.name: + return MAULIStrategy(critic_model, optimizer, configuration, logger) + if learning_strategy_enum.MASCOF == configuration.name: + return MASCOFStrategy(critic_model, optimizer, configuration, logger) + if learning_strategy_enum.SDAP == configuration.name: + return SDAPStrategy(critic_model, optimizer, configuration, logger) diff --git a/running_modes/automated_curriculum_learning/learning_strategy/learning_strategy_configuration.py b/running_modes/automated_curriculum_learning/learning_strategy/learning_strategy_configuration.py new file mode 100644 index 0000000..3a77297 --- /dev/null +++ b/running_modes/automated_curriculum_learning/learning_strategy/learning_strategy_configuration.py @@ -0,0 +1,9 @@ +from dataclasses import dataclass + +from running_modes.automated_curriculum_learning.learning_strategy.learning_strategy_enum import LearningStrategyEnum + + +@dataclass +class LearningStrategyConfiguration: + name: str = LearningStrategyEnum().DAP_SINGLE_QUERY + parameters: dict = None diff --git a/running_modes/automated_curriculum_learning/learning_strategy/learning_strategy_enum.py b/running_modes/automated_curriculum_learning/learning_strategy/learning_strategy_enum.py new file mode 100644 index 0000000..a808074 --- /dev/null +++ b/running_modes/automated_curriculum_learning/learning_strategy/learning_strategy_enum.py @@ -0,0 +1,10 @@ +from dataclasses import dataclass + + +@dataclass(frozen=True) +class LearningStrategyEnum: + DAP = "dap" + MAULI = "mauli" + MASCOF = "mascof" + SDAP = "sdap" + DAP_SINGLE_QUERY = "dap_single_query" diff --git a/running_modes/automated_curriculum_learning/learning_strategy/mascof_strategy.py b/running_modes/automated_curriculum_learning/learning_strategy/mascof_strategy.py new file mode 100644 index 0000000..8984939 --- /dev/null +++ b/running_modes/automated_curriculum_learning/learning_strategy/mascof_strategy.py @@ -0,0 +1,27 @@ +import torch +from reinvent_models.link_invent.dto import BatchLikelihoodDTO + +from running_modes.automated_curriculum_learning.dto import UpdatedLikelihoodsDTO +from running_modes.automated_curriculum_learning.learning_strategy.base_double_query_learning_strategy import \ + BaseDoubleQueryLearningStrategy +from running_modes.automated_curriculum_learning.learning_strategy.learning_strategy_configuration import \ + LearningStrategyConfiguration + + +class MASCOFStrategy(BaseDoubleQueryLearningStrategy): + + def __init__(self, critic_model, optimizer, configuration: LearningStrategyConfiguration, logger=None): + """ + TODO: Provide description of the current strategy + """ + super().__init__(critic_model, optimizer, configuration, logger) + + def _calculate_loss(self, likelihood_dto: BatchLikelihoodDTO, score) -> UpdatedLikelihoodsDTO: + batch = likelihood_dto.batch + critic_nlls = self.critic_model.likelihood(*batch.input, *batch.output) + negative_critic_nlls = -critic_nlls + negative_actor_nlls = -likelihood_dto.likelihood + augmented_nlls = self._to_tensor(score) + loss = -torch.sum(self._to_tensor(score)) * torch.sum(negative_actor_nlls) + dto = UpdatedLikelihoodsDTO(negative_actor_nlls, negative_critic_nlls, augmented_nlls, loss) + return dto \ No newline at end of file diff --git a/running_modes/automated_curriculum_learning/learning_strategy/mauli_strategy.py b/running_modes/automated_curriculum_learning/learning_strategy/mauli_strategy.py new file mode 100644 index 0000000..000942e --- /dev/null +++ b/running_modes/automated_curriculum_learning/learning_strategy/mauli_strategy.py @@ -0,0 +1,29 @@ +from reinvent_models.link_invent.dto import BatchLikelihoodDTO + +from running_modes.automated_curriculum_learning.dto import UpdatedLikelihoodsDTO +from running_modes.automated_curriculum_learning.learning_strategy.base_double_query_learning_strategy import \ + BaseDoubleQueryLearningStrategy +from running_modes.automated_curriculum_learning.learning_strategy.learning_strategy_configuration import \ + LearningStrategyConfiguration + + +class MAULIStrategy(BaseDoubleQueryLearningStrategy): + + def __init__(self, critic_model, optimizer, configuration: LearningStrategyConfiguration, logger=None): + """ + TODO: Provide description of the current strategy + """ + super().__init__(critic_model, optimizer, configuration, logger) + + #TODO: Create a StrategySpecificEnums + self._sigma = self._configuration.parameters.get("sigma", 120) + + def _calculate_loss(self, likelihood_dto: BatchLikelihoodDTO, score) -> UpdatedLikelihoodsDTO: + batch = likelihood_dto.batch + critic_nlls = self.critic_model.likelihood(*batch.input, *batch.output) + negative_critic_nlls = -critic_nlls + negative_actor_nlls = -likelihood_dto.likelihood + augmented_nlls = negative_critic_nlls + self._sigma * self._to_tensor(score) + loss = -(sum(augmented_nlls) * sum(negative_actor_nlls)) + dto = UpdatedLikelihoodsDTO(negative_actor_nlls, negative_critic_nlls, augmented_nlls, loss) + return dto \ No newline at end of file diff --git a/running_modes/automated_curriculum_learning/learning_strategy/sdap_strategy.py b/running_modes/automated_curriculum_learning/learning_strategy/sdap_strategy.py new file mode 100644 index 0000000..118b1e2 --- /dev/null +++ b/running_modes/automated_curriculum_learning/learning_strategy/sdap_strategy.py @@ -0,0 +1,30 @@ +import torch +from reinvent_models.link_invent.dto import BatchLikelihoodDTO + +from running_modes.automated_curriculum_learning.dto import UpdatedLikelihoodsDTO +from running_modes.automated_curriculum_learning.learning_strategy.base_double_query_learning_strategy import \ + BaseDoubleQueryLearningStrategy +from running_modes.automated_curriculum_learning.learning_strategy.learning_strategy_configuration import \ + LearningStrategyConfiguration + + +class SDAPStrategy(BaseDoubleQueryLearningStrategy): + + def __init__(self, critic_model, optimizer, configuration: LearningStrategyConfiguration, logger=None): + """ + TODO: Provide description of the current strategy + """ + super().__init__(critic_model, optimizer, configuration, logger) + + self._sigma = self._configuration.parameters.get("sigma", 120) + + def _calculate_loss(self, likelihood_dto: BatchLikelihoodDTO, score) -> UpdatedLikelihoodsDTO: + batch = likelihood_dto.batch + critic_nlls = self.critic_model.likelihood(*batch.input, *batch.output) + negative_critic_nlls = -critic_nlls + negative_actor_nlls = -likelihood_dto.likelihood + augmented_nlls = negative_critic_nlls + self._sigma * self._to_tensor(score) + reward_score = torch.pow((augmented_nlls - negative_actor_nlls), 2).mean() + loss = -(reward_score) * (negative_actor_nlls).mean() + dto = UpdatedLikelihoodsDTO(negative_actor_nlls, negative_critic_nlls, augmented_nlls, loss) + return dto \ No newline at end of file diff --git a/running_modes/automated_curriculum_learning/logging/__init__.py b/running_modes/automated_curriculum_learning/logging/__init__.py index 90f8224..c7d6b89 100644 --- a/running_modes/automated_curriculum_learning/logging/__init__.py +++ b/running_modes/automated_curriculum_learning/logging/__init__.py @@ -1,3 +1,2 @@ from running_modes.automated_curriculum_learning.logging.console_message import ConsoleMessage -from running_modes.automated_curriculum_learning.logging.local_auto_cl_logger import LocalAutoCLLogger from running_modes.automated_curriculum_learning.logging.auto_cl_logger import AutoCLLogger diff --git a/running_modes/automated_curriculum_learning/logging/auto_cl_logger.py b/running_modes/automated_curriculum_learning/logging/auto_cl_logger.py index be7d6f9..2b0745b 100644 --- a/running_modes/automated_curriculum_learning/logging/auto_cl_logger.py +++ b/running_modes/automated_curriculum_learning/logging/auto_cl_logger.py @@ -1,20 +1,24 @@ -from running_modes.automated_curriculum_learning.logging import LocalAutoCLLogger -from running_modes.automated_curriculum_learning.logging.base_auto_cl_logger import BaseAutoCLLogger -from running_modes.configurations import ReinforcementLoggerConfiguration -from running_modes.configurations.general_configuration_envelope import GeneralConfigurationEnvelope +from reinvent_models.model_factory.enums.model_type_enum import ModelTypeEnum +from running_modes.automated_curriculum_learning.logging.base_logger import BaseLogger +from running_modes.automated_curriculum_learning.logging.local_logger import LocalLogger +from running_modes.configurations import CurriculumLoggerConfiguration +from running_modes.configurations.general_configuration_envelope import GeneralConfigurationEnvelope from running_modes.enums.logging_mode_enum import LoggingModeEnum class AutoCLLogger: - def __new__(cls, configuration: GeneralConfigurationEnvelope) -> BaseAutoCLLogger: + def __new__(cls, configuration: GeneralConfigurationEnvelope) -> BaseLogger: logging_mode_enum = LoggingModeEnum() - auto_cl_config = ReinforcementLoggerConfiguration.parse_obj(configuration.logging) + model_type = ModelTypeEnum() + auto_cl_config = CurriculumLoggerConfiguration.parse_obj(configuration.logging) if auto_cl_config.recipient == logging_mode_enum.LOCAL: - logger_instance = LocalAutoCLLogger(configuration) - + if model_type.DEFAULT == configuration.model_type: + return LocalLogger(configuration, auto_cl_config) + elif model_type.LINK_INVENT == configuration.model_type: + return LocalLogger(configuration, auto_cl_config) else: raise NotImplementedError("Remote Auto CL logging is not implemented.") - return logger_instance + diff --git a/running_modes/automated_curriculum_learning/logging/base_auto_cl_logger.py b/running_modes/automated_curriculum_learning/logging/base_auto_cl_logger.py deleted file mode 100644 index f49999a..0000000 --- a/running_modes/automated_curriculum_learning/logging/base_auto_cl_logger.py +++ /dev/null @@ -1,99 +0,0 @@ -import time -import json -import logging -import os -from abc import ABC, abstractmethod - -import numpy as np -import torch -from reinvent_scoring import ComponentSpecificParametersEnum -from reinvent_scoring.scoring.enums.scoring_function_component_enum import ScoringFunctionComponentNameEnum -from reinvent_scoring.scoring.score_summary import FinalSummary - -from reinvent_scoring.scoring.diversity_filters.reinvent_core.base_diversity_filter import BaseDiversityFilter -from running_modes.configurations.general_configuration_envelope import GeneralConfigurationEnvelope -from running_modes.configurations import ReinforcementLoggerConfiguration - - -class BaseAutoCLLogger(ABC): - def __init__(self, configuration: GeneralConfigurationEnvelope): - self._configuration = configuration - self._log_config = ReinforcementLoggerConfiguration(**self._configuration.logging) - self._setup_workfolder() - self._logger = self._setup_logger() - self._start_time = time.time() - self._specific_parameters_enum = ComponentSpecificParametersEnum() - - @abstractmethod - def log_message(self, message: str): - raise NotImplementedError("log_message method is not implemented") - - @abstractmethod - def timestep_report(self, start_time, n_steps, step, smiles, - mean_score: np.array, score_summary: FinalSummary, score, - agent_likelihood: torch.tensor, prior_likelihood: torch.tensor, - augmented_likelihood: torch.tensor, diversity_filter: BaseDiversityFilter): - raise NotImplementedError("timestep_report method is not implemented") - - def log_out_input_configuration(self): - file = os.path.join(self._log_config.result_folder, "input.json") - jsonstr = json.dumps(self._configuration, default=lambda x: x.__dict__, sort_keys=True, indent=4, - separators=(',', ': ')) - with open(file, 'w') as f: - f.write(jsonstr) - - def log_text_to_file(self, text: str): - file = os.path.join(self._log_config.logging_path, "curriculum_log.txt") - with open(file, "a+") as f: - time_passed = round(time.time() - self._start_time, 4) - f.write(f"Time from start: {time_passed}s, {text} \n") - - def save_checkpoint(self, step, scaffold_filter, agent): - actual_step = step + 1 - if self._log_config.logging_frequency > 0 and actual_step % self._log_config.logging_frequency == 0: - self.save_scaffold_memory(scaffold_filter) - agent.save(os.path.join(self._log_config.result_folder, f'Agent.{actual_step}.ckpt')) - - @abstractmethod - def save_final_state(self, agent, scaffold_filter): - raise NotImplementedError("save_final_state method is not implemented") - - def save_merging_state(self, agent, scaffold_filter, name): - agent.save(os.path.join(self._log_config.result_folder, f'Agent{name}.ckpt')) - self.save_scaffold_memory(scaffold_filter, memory_name=name) - - def save_scaffold_memory(self, scaffold_filter, memory_name: str = ""): - scaffold_memory = scaffold_filter.get_memory_as_dataframe() - self.save_to_csv(scaffold_memory=scaffold_memory, path=self._log_config.result_folder, memory_name=memory_name, - job_name=self._log_config.job_name) - - def _setup_workfolder(self): - if not os.path.isdir(self._log_config.logging_path): - os.makedirs(self._log_config.logging_path) - if not os.path.isdir(self._log_config.result_folder): - os.makedirs(self._log_config.result_folder) - - def _setup_logger(self): - handler = logging.StreamHandler() - formatter = logging.Formatter( - fmt="%(asctime)s: %(module)s.%(funcName)s +%(lineno)s: %(levelname)-8s %(message)s", - datefmt="%H:%M:%S" - ) - handler.setFormatter(formatter) - logger = logging.getLogger("auto_cl_logger") - if not logger.handlers: - logger.addHandler(handler) - logger.setLevel(logging.INFO) - logger.propagate = False - return logger - - def save_to_csv(self, scaffold_memory, path: str, memory_name: str = "", job_name: str = "default_job"): - sf_enum = ScoringFunctionComponentNameEnum() - if not os.path.isdir(path): - os.makedirs(path) - file_name = os.path.join(path, f"scaffold_memory{memory_name}.csv") - - if len(scaffold_memory) > 0: - sorted_df = scaffold_memory.sort_values(sf_enum.TOTAL_SCORE, ascending=False) - sorted_df["ID"] = [f"{job_name}_{e}" for e, _ in enumerate(sorted_df.index.array)] - sorted_df.to_csv(file_name, index=False) diff --git a/running_modes/automated_curriculum_learning/logging/base_logger.py b/running_modes/automated_curriculum_learning/logging/base_logger.py new file mode 100644 index 0000000..8bf5139 --- /dev/null +++ b/running_modes/automated_curriculum_learning/logging/base_logger.py @@ -0,0 +1,81 @@ +import json +import logging +import os +from abc import ABC, abstractmethod + +from reinvent_scoring.scoring.enums.scoring_function_component_enum import ScoringFunctionComponentNameEnum + +from running_modes.automated_curriculum_learning.dto.timestep_dto import TimestepDTO +from running_modes.configurations import GeneralConfigurationEnvelope, CurriculumLoggerConfiguration + + +class BaseLogger(ABC): + def __init__(self, configuration: GeneralConfigurationEnvelope, log_config: CurriculumLoggerConfiguration): + self._configuration = configuration + self._log_config = log_config + self._setup_workfolder() + self._logger = self._setup_logger() + + @abstractmethod + def log_message(self, message: str): + raise NotImplementedError("log_message method is not implemented") + + @abstractmethod + def timestep_report(self, report_dto: TimestepDTO, diversity_filter, agent): + raise NotImplementedError("timestep_report method is not implemented") + + @abstractmethod + def save_final_state(self, agent, diversity_filter): + raise NotImplementedError("save_final_state method is not implemented") + + def log_out_input_configuration(self): + file = os.path.join(self._log_config.result_folder, "input.json") + jsonstr = json.dumps(self._configuration, default=lambda x: x.__dict__, sort_keys=True, indent=4, + separators=(',', ': ')) + with open(file, 'w') as f: + f.write(jsonstr) + + def save_checkpoint(self, step: int, diversity_filter, agent): + actual_step = step + 1 + if self._log_config.logging_frequency > 0 and actual_step % self._log_config.logging_frequency == 0: + self.save_filter_memory(diversity_filter) + agent.save_to_file(os.path.join(self._log_config.result_folder, f'Agent.{actual_step}.ckpt')) + + def save_merging_state(self, agent, diversity_filter, name): + agent.save_to_file(os.path.join(self._log_config.result_folder, f'Agent{name}.ckpt')) + self.save_filter_memory(diversity_filter, memory_name=name) + + def save_filter_memory(self, diversity_filter, memory_name: str = ""): + diversity_memory = diversity_filter.get_memory_as_dataframe() + self._save_to_csv(diversity_memory, self._log_config.result_folder, memory_name, self._log_config.job_name) + + def _save_to_csv(self, diversity_memory, path: str, memory_name: str = "", job_name: str = "default_job"): + sf_enum = ScoringFunctionComponentNameEnum() + if not os.path.isdir(path): + os.makedirs(path) + file_name = os.path.join(path, f"scaffold_memory{memory_name}.csv") + + if len(diversity_memory) > 0: + sorted_df = diversity_memory.sort_values(sf_enum.TOTAL_SCORE, ascending=False) + sorted_df["ID"] = [f"{job_name}_{e}" for e, _ in enumerate(sorted_df.index.array)] + sorted_df.to_csv(file_name, index=False) + + def _setup_workfolder(self): + if not os.path.isdir(self._log_config.logging_path): + os.makedirs(self._log_config.logging_path) + if not os.path.isdir(self._log_config.result_folder): + os.makedirs(self._log_config.result_folder) + + def _setup_logger(self): + handler = logging.StreamHandler() + formatter = logging.Formatter( + fmt="%(asctime)s: %(module)s.%(funcName)s +%(lineno)s: %(levelname)-8s %(message)s", + datefmt="%H:%M:%S" + ) + handler.setFormatter(formatter) + logger = logging.getLogger("curriculum_logger") + if not logger.handlers: + logger.addHandler(handler) + logger.setLevel(logging.INFO) + logger.propagate = False + return logger diff --git a/running_modes/automated_curriculum_learning/logging/console_message.py b/running_modes/automated_curriculum_learning/logging/console_message.py index 7d2188d..805b681 100644 --- a/running_modes/automated_curriculum_learning/logging/console_message.py +++ b/running_modes/automated_curriculum_learning/logging/console_message.py @@ -1,39 +1,42 @@ import time +import numpy as np + from reinvent_scoring.scoring.score_summary import FinalSummary from reinvent_chemistry.logging import fraction_valid_smiles +from running_modes.automated_curriculum_learning.dto.timestep_dto import TimestepDTO + class ConsoleMessage: - def create(self, start_time, n_steps, step, smiles, - mean_score, score_summary: FinalSummary, score, - agent_likelihood, prior_likelihood, augmented_likelihood): - time_message = self._time_progress(start_time, n_steps, step, smiles, mean_score) - score_message = self._score_profile(score_summary.scored_smiles, agent_likelihood, prior_likelihood, - augmented_likelihood, score) - score_breakdown = self._score_summary_breakdown(score_summary) + def create(self, report_dto: TimestepDTO): + time_message = self._time_progress(report_dto) + score_message = self._score_profile(report_dto) + score_breakdown = self._score_summary_breakdown(report_dto.score_summary) message = time_message + score_message + score_breakdown return message - def _time_progress(self, start_time, n_steps, step, smiles, mean_score): - time_elapsed = int(time.time() - start_time) - time_left = (time_elapsed * ((n_steps - step) / (step + 1))) - valid_fraction = fraction_valid_smiles(smiles) - message = (f"\n Step {step} Fraction valid SMILES: {valid_fraction:4.1f} Score: {mean_score:.4f} " + def _time_progress(self, report_dto: TimestepDTO): + mean_score = np.mean(report_dto.score_summary.total_score) + time_elapsed = int(time.time() - report_dto.start_time) + time_left = (time_elapsed * ((report_dto.n_steps - report_dto.step) / (report_dto.step + 1))) + valid_fraction = fraction_valid_smiles(report_dto.score_summary.scored_smiles) + message = (f"\n Step {report_dto.step} Fraction valid SMILES: {valid_fraction:4.1f} Score: {mean_score:.4f} " + f"Sample size: {len(report_dto.score_summary.scored_smiles)} " f"Time elapsed: {time_elapsed} " f"Time left: {time_left:.1f}\n") return message - def _score_profile(self, smiles, agent_likelihood, prior_likelihood, augmented_likelihood, score): + def _score_profile(self, report_dto: TimestepDTO): # Convert to numpy arrays so that we can print them - augmented_likelihood = augmented_likelihood.data.cpu().numpy() - agent_likelihood = agent_likelihood.data.cpu().numpy() + augmented_likelihood = report_dto.augmented_likelihood.data.cpu().numpy() + agent_likelihood = report_dto.agent_likelihood.data.cpu().numpy() message = " ".join([" Agent", "Prior", "Target", "Score"] + ["SMILES\n"]) - for i in range(min(10, len(smiles))): - message += f'{agent_likelihood[i]:6.2f} {prior_likelihood[i]:6.2f} ' \ - f'{augmented_likelihood[i]:6.2f} {score[i]:6.2f} ' - message += f" {smiles[i]}\n" + for i in range(min(10, len(report_dto.score_summary.scored_smiles))): + message += f'{agent_likelihood[i]:6.2f} {report_dto.prior_likelihood[i]:6.2f} ' \ + f'{augmented_likelihood[i]:6.2f} {report_dto.score_summary.total_score[i]:6.2f} ' + message += f" {report_dto.score_summary.scored_smiles[i]}\n" return message def _score_summary_breakdown(self, score_summary: FinalSummary): diff --git a/running_modes/automated_curriculum_learning/logging/local_auto_cl_logger.py b/running_modes/automated_curriculum_learning/logging/local_auto_cl_logger.py deleted file mode 100644 index 762caa3..0000000 --- a/running_modes/automated_curriculum_learning/logging/local_auto_cl_logger.py +++ /dev/null @@ -1,95 +0,0 @@ -import os - -import numpy as np -import torch -from reinvent_scoring.scoring.enums.scoring_function_component_enum import ScoringFunctionComponentNameEnum -from reinvent_scoring.scoring.score_summary import FinalSummary -from torch.utils.tensorboard import SummaryWriter - -import reinvent_chemistry.logging as ul_rl -from reinvent_chemistry.logging import add_mols -from reinvent_chemistry.logging import fraction_valid_smiles -from reinvent_scoring.scoring.diversity_filters.reinvent_core.base_diversity_filter import BaseDiversityFilter -from running_modes.automated_curriculum_learning.logging.base_auto_cl_logger import BaseAutoCLLogger -from running_modes.configurations.general_configuration_envelope import GeneralConfigurationEnvelope -from running_modes.automated_curriculum_learning.logging import ConsoleMessage - - -class LocalAutoCLLogger(BaseAutoCLLogger): - def __init__(self, configuration: GeneralConfigurationEnvelope): - super().__init__(configuration) - self._summary_writer = SummaryWriter(log_dir=self._log_config.logging_path) - self._summary_writer.add_text('Legends', - 'The values under each compound are read as: [Agent; Prior; Target; Score]') - # _rows and _columns define the shape of the output grid of molecule images in tensorboard. - self._rows = 4 - self._columns = 4 - self._sample_size = self._rows * self._columns - self._sf_component_enum = ScoringFunctionComponentNameEnum() - self._console_message_formatter = ConsoleMessage() - - def log_message(self, message: str): - self._logger.info(message) - - def timestep_report(self, start_time, n_steps, step, smiles: np.array, - mean_score: np.float32, score_summary: FinalSummary, score: np.array, - agent_likelihood: torch.tensor, prior_likelihood: torch.tensor, - augmented_likelihood: torch.tensor, diversity_filter: BaseDiversityFilter): - message = self._console_message_formatter.create(start_time, n_steps, step, smiles, mean_score, score_summary, - score, agent_likelihood, prior_likelihood, - augmented_likelihood) - self._logger.info(message) - self._tensorboard_report(step, smiles, score, score_summary, agent_likelihood, prior_likelihood, - augmented_likelihood, diversity_filter) - - def save_final_state(self, agent, scaffold_filter): - agent.save(os.path.join(self._log_config.result_folder, 'Agent.ckpt')) - self.save_scaffold_memory(scaffold_filter) - self._summary_writer.close() - self.log_out_input_configuration() - - def _tensorboard_report(self, step, smiles, score, score_summary: FinalSummary, agent_likelihood, prior_likelihood, - augmented_likelihood, diversity_filter: BaseDiversityFilter): - self._summary_writer.add_scalars("nll/avg", { - "prior": prior_likelihood.mean(), - "augmented": augmented_likelihood.mean(), - "agent": agent_likelihood.mean() - }, step) - mean_score = np.mean(score) - for i, log in enumerate(score_summary.profile): - self._summary_writer.add_scalar(score_summary.profile[i].name, np.mean(score_summary.profile[i].score), - step) - self._summary_writer.add_scalar("Valid SMILES", fraction_valid_smiles(smiles), step) - self._summary_writer.add_scalar("Number of SMILES found", diversity_filter.number_of_smiles_in_memory(), step) - self._summary_writer.add_scalar("Average score", mean_score, step) - if step % 10 == 0: - self._log_out_smiles_sample(smiles, score, step, score_summary) - - def _log_out_smiles_sample(self, smiles, score, step, score_summary: FinalSummary): - self._visualize_structures(smiles, score, step, score_summary) - - def _visualize_structures(self, smiles, score, step, score_summary: FinalSummary): - - list_of_mols, legends, pattern = self._check_for_invalid_mols_and_create_legends(smiles, score, score_summary) - try: - add_mols(self._summary_writer, "Molecules from epoch", list_of_mols[:self._sample_size], self._rows, - [x for x in legends], global_step=step, size_per_mol=(320, 320), pattern=pattern) - except: - self.log_message(f"Error in RDKit has occurred, skipping report for step {step}.") - - def _check_for_invalid_mols_and_create_legends(self, smiles, score, score_summary: FinalSummary): - smiles = ul_rl.padding_with_invalid_smiles(smiles, self._sample_size) - list_of_mols, legend = ul_rl.check_for_invalid_mols_and_create_legend(smiles, score, self._sample_size) - smarts_pattern = self._get_matching_substructure_from_config(score_summary) - pattern = ul_rl.find_matching_pattern_in_smiles(list_of_mols=list_of_mols, smarts_pattern=smarts_pattern) - - return list_of_mols, legend, pattern - - def _get_matching_substructure_from_config(self, score_summary: FinalSummary): - smarts_pattern = "" - for summary_component in score_summary.scaffold_log: - if summary_component.parameters.component_type == self._sf_component_enum.MATCHING_SUBSTRUCTURE: - smarts = summary_component.parameters.specific_parameters.get(self._specific_parameters_enum.SMILES, []) - if len(smarts) > 0: - smarts_pattern = smarts[0] - return smarts_pattern diff --git a/running_modes/automated_curriculum_learning/logging/local_logger.py b/running_modes/automated_curriculum_learning/logging/local_logger.py new file mode 100644 index 0000000..0612d44 --- /dev/null +++ b/running_modes/automated_curriculum_learning/logging/local_logger.py @@ -0,0 +1,90 @@ +import os + +import numpy as np +from reinvent_chemistry.logging import padding_with_invalid_smiles, \ + check_for_invalid_mols_and_create_legend, find_matching_pattern_in_smiles, add_mols, fraction_valid_smiles +from reinvent_scoring.scoring.diversity_filters.lib_invent.base_diversity_filter import BaseDiversityFilter +from reinvent_scoring.scoring.enums import ComponentSpecificParametersEnum, ScoringFunctionComponentNameEnum +from reinvent_scoring.scoring.score_summary import FinalSummary +from torch.utils.tensorboard import SummaryWriter + +from running_modes.automated_curriculum_learning.dto.timestep_dto import TimestepDTO +from running_modes.automated_curriculum_learning.logging.base_logger import BaseLogger +from running_modes.automated_curriculum_learning.logging.console_message import ConsoleMessage +from running_modes.configurations import GeneralConfigurationEnvelope, CurriculumLoggerConfiguration + + +class LocalLogger(BaseLogger): + def __init__(self, configuration: GeneralConfigurationEnvelope, log_config: CurriculumLoggerConfiguration): + super().__init__(configuration, log_config) + self._summary_writer = SummaryWriter(log_dir=self._log_config.logging_path) + self._sample_size = self._log_config.rows * self._log_config.columns + self._sf_component_enum = ScoringFunctionComponentNameEnum() + self._specific_parameters_enum = ComponentSpecificParametersEnum() + self._console_message_formatter = ConsoleMessage() + + def log_message(self, message: str): + self._logger.info(message) + + def timestep_report(self, report_dto: TimestepDTO, diversity_filter: BaseDiversityFilter, agent): + message = self._console_message_formatter.create(report_dto) + self._logger.info(message) + self._tensorboard_report(report_dto, diversity_filter) + self.save_checkpoint(report_dto.step, diversity_filter, agent) + + def save_final_state(self, agent, diversity_filter): + agent.save_to_file(os.path.join(self._log_config.result_folder, 'Agent.ckpt')) + self.save_filter_memory(diversity_filter) + self._summary_writer.close() + self.log_out_input_configuration() + + def _tensorboard_report(self, report_dto: TimestepDTO, diversity_filter: BaseDiversityFilter): + self._summary_writer.add_scalars("nll/avg", { + "prior": report_dto.prior_likelihood.mean(), + "augmented": report_dto.augmented_likelihood.mean(), + "agent": report_dto.agent_likelihood.mean() + }, report_dto.step) + self._summary_writer.add_scalars("nll/variance", { + "prior": report_dto.prior_likelihood.var(), + "augmented": report_dto.augmented_likelihood.var(), + "agent": report_dto.agent_likelihood.var() + }, report_dto.step) + mean_score = np.mean(report_dto.score_summary.total_score) + for i, log in enumerate(report_dto.score_summary.profile): + self._summary_writer.add_scalar(report_dto.score_summary.profile[i].name, + np.mean(report_dto.score_summary.profile[i].score), report_dto.step) + self._summary_writer.add_scalar("Valid SMILES", fraction_valid_smiles(report_dto.score_summary.scored_smiles), + report_dto.step) + self._summary_writer.add_scalar("Number of SMILES found", diversity_filter.number_of_smiles_in_memory(), + report_dto.step) + self._summary_writer.add_scalar("Average score", mean_score, report_dto.step) + self._log_out_smiles_sample(report_dto) + + def _log_out_smiles_sample(self, report_dto: TimestepDTO): + self._visualize_structures(report_dto.score_summary.scored_smiles, report_dto.score_summary.total_score, + report_dto.step, report_dto.score_summary) + + def _visualize_structures(self, smiles, score, step, score_summary: FinalSummary): + list_of_mols, legends, pattern = self._check_for_invalid_mols_and_create_legends(smiles, score, score_summary) + try: + add_mols(self._summary_writer, "Molecules from epoch", list_of_mols[:self._sample_size], self._log_config.rows, + [x for x in legends], global_step=step, size_per_mol=(320, 320), pattern=pattern) + except: + self.log_message(f"Error in RDKit has occurred, skipping report for step {step}.") + + def _check_for_invalid_mols_and_create_legends(self, smiles, score, score_summary: FinalSummary): + smiles = padding_with_invalid_smiles(smiles, self._sample_size) + list_of_mols, legend = check_for_invalid_mols_and_create_legend(smiles, score, self._sample_size) + smarts_pattern = self._get_matching_substructure_from_config(score_summary) + pattern = find_matching_pattern_in_smiles(list_of_mols=list_of_mols, smarts_pattern=smarts_pattern) + + return list_of_mols, legend, pattern + + def _get_matching_substructure_from_config(self, score_summary: FinalSummary): + smarts_pattern = "" + for summary_component in score_summary.scaffold_log: + if summary_component.parameters.component_type == self._sf_component_enum.MATCHING_SUBSTRUCTURE: + smarts = summary_component.parameters.specific_parameters.get(self._specific_parameters_enum.SMILES, []) + if len(smarts) > 0: + smarts_pattern = smarts[0] + return smarts_pattern diff --git a/running_modes/automated_curriculum_learning/production_strategy/base_production_strategy.py b/running_modes/automated_curriculum_learning/production_strategy/base_production_strategy.py index 5d6e7fd..c63b650 100644 --- a/running_modes/automated_curriculum_learning/production_strategy/base_production_strategy.py +++ b/running_modes/automated_curriculum_learning/production_strategy/base_production_strategy.py @@ -1,103 +1,49 @@ from abc import ABC, abstractmethod -from typing import List, Dict, Tuple, Any +from typing import List -import numpy as np -import torch -from reinvent_chemistry import get_indices_of_unique_smiles -from reinvent_models.reinvent_core.models.model import Model -from reinvent_scoring import ScoringFunctionFactory, FinalSummary, ScoringFunctionParameters -from reinvent_scoring.scoring.diversity_filters.reinvent_core.base_diversity_filter import BaseDiversityFilter +from reinvent_chemistry import Conversions +from reinvent_chemistry.library_design import BondMaker, AttachmentPoints +from reinvent_models.model_factory.generative_model_base import GenerativeModelBase +from reinvent_scoring.scoring.diversity_filters.curriculum_learning.base_diversity_filter import BaseDiversityFilter from reinvent_scoring.scoring.function.base_scoring_function import BaseScoringFunction -from running_modes.automated_curriculum_learning.logging.base_auto_cl_logger import BaseAutoCLLogger -from running_modes.configurations.automated_curriculum_learning.production_strategy_configuration import \ - ProductionStrategyConfiguration -from running_modes.reinforcement_learning.inception import Inception -from running_modes.utils import to_tensor +from running_modes.automated_curriculum_learning.dto import CurriculumOutcomeDTO +from running_modes.automated_curriculum_learning.inception.inception import Inception +from running_modes.automated_curriculum_learning.logging.base_logger import BaseLogger +from running_modes.configurations.automated_curriculum_learning.prodcution_strategy_input_configuration import \ + ProductionStrategyInputConfiguration class BaseProductionStrategy(ABC): - def __init__(self, prior: Model, diversity_filter: BaseDiversityFilter, inception: Inception, - scoring_function: BaseScoringFunction, configuration: ProductionStrategyConfiguration, - logger: BaseAutoCLLogger): + def __init__(self, prior: GenerativeModelBase, diversity_filter: BaseDiversityFilter, inception: Inception, + scoring_function: BaseScoringFunction, configuration: ProductionStrategyInputConfiguration, + logger: BaseLogger): + + self._bond_maker = BondMaker() + self._attachment_points = AttachmentPoints() + self._conversion = Conversions() self._parameters = configuration self._prior = prior + self._logger = logger self._diversity_filter = diversity_filter - self._inception = inception + self.inception = inception self._scoring_function = scoring_function - self._logger = logger @abstractmethod - def run(self, cl_agent: Model, steps_so_far: int): - raise NotImplementedError("run not implemented.") + def run(self, cl_agent: GenerativeModelBase, steps_so_far: int) -> CurriculumOutcomeDTO: + raise NotImplementedError("run() method is not implemented ") - def setup_scoring_function(self, name: str, parameter_list: List[Dict]) -> BaseScoringFunction: - scoring_function_parameters = ScoringFunctionParameters(name=name, parameters=parameter_list, parallel=False) - scoring_function_instance = ScoringFunctionFactory(scoring_function_parameters) + @abstractmethod + def take_step(self, agent: GenerativeModelBase, learning_strategy, scoring_function: BaseScoringFunction, + step:int, start_time: float) -> float: + raise NotImplementedError("take_step() method is not implemented ") - self._log_sf_update(current_parameters=parameter_list) - return scoring_function_instance + def disable_prior_gradients(self): + for param in self._prior.get_network_parameters(): + param.requires_grad = False def _log_sf_update(self, current_parameters: List[dict]): text_to_log = f"** Production setup **\n scoring_function: " \ f"{[component.get('name') for component in current_parameters]}" - # log in console self._logger.log_message(text_to_log) - def _disable_prior_gradients(self): - # There might be a more elegant way of disabling gradients - for param in self._prior.network.parameters(): - param.requires_grad = False - - def _sample_unique_sequences(self, agent: Model, batch_size: int) -> Tuple[Any, Any, Any]: - seqs, smiles, agent_likelihood = agent.sample_sequences_and_smiles(batch_size) - unique_idxs = get_indices_of_unique_smiles(smiles) - seqs_unique = seqs[unique_idxs] - smiles_np = np.array(smiles) - smiles_unique = smiles_np[unique_idxs] - agent_likelihood_unique = agent_likelihood[unique_idxs] - return seqs_unique, smiles_unique, agent_likelihood_unique - - def _inception_filter(self, agent, loss, agent_likelihood, prior_likelihood, - sigma, smiles, score): - if self._inception is not None: - exp_smiles, exp_scores, exp_prior_likelihood = self._inception.sample() - if len(exp_smiles) > 0: - exp_agent_likelihood = -agent.likelihood_smiles(exp_smiles) - exp_augmented_likelihood = exp_prior_likelihood + sigma * exp_scores - exp_loss = torch.pow((to_tensor(exp_augmented_likelihood) - exp_agent_likelihood), 2) - loss = torch.cat((loss, exp_loss), 0) - agent_likelihood = torch.cat((agent_likelihood, exp_agent_likelihood), 0) - self._inception.add(smiles, score, prior_likelihood) - return loss, agent_likelihood - - def _stats_and_chekpoint(self, agent: Model, score: Any, start_time: float, step: int, smiles: List, - score_summary: FinalSummary, agent_likelihood: torch.tensor, prior_likelihood: torch.tensor, - augmented_likelihood: torch.tensor): - mean_score = np.mean(score) - self._logger.timestep_report(start_time, self._parameters.n_steps, step, smiles, - mean_score, score_summary, score, - agent_likelihood, prior_likelihood, augmented_likelihood, self._diversity_filter) - self._logger.save_checkpoint(step, self._diversity_filter, agent) - - def _take_step(self, agent: Model, optimizer: Any, scoring_function_current: BaseScoringFunction, step: int, - start_time: float): - seqs, smiles, agent_likelihood = self._sample_unique_sequences(agent, self._parameters.batch_size) - agent_likelihood = -agent_likelihood - prior_likelihood = -self._prior.likelihood(seqs) - score_summary_current: FinalSummary = scoring_function_current.get_final_score_for_step(smiles, step) - - score = self._diversity_filter.update_score(score_summary_current, step) - augmented_likelihood = prior_likelihood + self._parameters.sigma * to_tensor(score) - loss = torch.pow((augmented_likelihood - agent_likelihood), 2) - loss, agent_likelihood = self._inception_filter(agent, loss, agent_likelihood, prior_likelihood, - self._parameters.sigma, smiles, score) - - loss = loss.mean() - optimizer.zero_grad() - loss.backward() - optimizer.step() - - self._stats_and_chekpoint(agent=agent, score=score, start_time=start_time, step=step, smiles=smiles, - score_summary=score_summary_current, agent_likelihood=agent_likelihood, - prior_likelihood=prior_likelihood, augmented_likelihood=augmented_likelihood) \ No newline at end of file diff --git a/running_modes/automated_curriculum_learning/production_strategy/link_invent_production_strategy.py b/running_modes/automated_curriculum_learning/production_strategy/link_invent_production_strategy.py new file mode 100644 index 0000000..e397f03 --- /dev/null +++ b/running_modes/automated_curriculum_learning/production_strategy/link_invent_production_strategy.py @@ -0,0 +1,111 @@ +import time +from abc import abstractmethod +from typing import List + +import torch +from reinvent_models.model_factory.generative_model_base import GenerativeModelBase +from reinvent_scoring import FinalSummary +from reinvent_scoring.scoring.diversity_filters.curriculum_learning.loggable_data_dto import UpdateLoggableDataDTO +from reinvent_scoring.scoring.diversity_filters.curriculum_learning.update_diversity_filter_dto import \ + UpdateDiversityFilterDTO +from reinvent_scoring.scoring.function.base_scoring_function import BaseScoringFunction + +from running_modes.automated_curriculum_learning.actions import LinkInventSampleModel +from running_modes.automated_curriculum_learning.dto import SampledSequencesDTO, TimestepDTO, UpdatedLikelihoodsDTO +from running_modes.automated_curriculum_learning.learning_strategy.learning_strategy import LearningStrategy +from running_modes.automated_curriculum_learning.production_strategy.base_production_strategy import \ + BaseProductionStrategy + + +class LinkInventProductionStrategy(BaseProductionStrategy): + + def run(self, cl_agent: GenerativeModelBase, steps_so_far: int): + self.disable_prior_gradients() + step_limit = steps_so_far + self._parameters.number_of_steps + optimizer = torch.optim.Adam(cl_agent.get_network_parameters(), lr=self._parameters.learning_rate) + learning_strategy = LearningStrategy(self._prior, optimizer, self._parameters.learning_strategy, self._logger) + + for step in range(steps_so_far, step_limit): + start_time = time.time() + self.take_step(agent=cl_agent, scoring_function=self._scoring_function, step=step, start_time=start_time, + learning_strategy=learning_strategy) + + self._logger.log_message(f"Production finished at step {step_limit}") + self._logger.save_final_state(cl_agent, self._diversity_filter) + + def take_step(self, agent: GenerativeModelBase, scoring_function: BaseScoringFunction, + step: int, start_time: float, learning_strategy) -> float: + # 1. Sampling + sampled_sequences = self._sampling(agent) + # 2. Scoring + score_summary = self._scoring(scoring_function, sampled_sequences, step) + # 3. Updating + dto = self._updating(sampled_sequences, score_summary.total_score, learning_strategy, agent) + # 4. Logging + self._logging(start_time, step, score_summary, dto, agent) + + score = score_summary.total_score.mean() + return score + + def _sampling(self, agent) -> List[SampledSequencesDTO]: + sampling_action = LinkInventSampleModel(agent, self._parameters.batch_size, self._logger, + self._parameters.randomize_input) + sampled_sequences = sampling_action.run(self._parameters.input) + return sampled_sequences + + def _scoring(self, scoring_function, sampled_sequences, step: int) -> FinalSummary: + score_summary = self._apply_scoring_function(scoring_function, sampled_sequences, step) + score_summary = self._clean_scored_smiles(score_summary) + loggable_data = [UpdateLoggableDataDTO(dto.input, dto.output, dto.nll) for dto in sampled_sequences] + dto = UpdateDiversityFilterDTO(score_summary, loggable_data, step) + score_summary.total_score = self._diversity_filter.update_score(dto) + return score_summary + + def _updating(self, sampled_sequences, score, learning_strategy, agent) -> UpdatedLikelihoodsDTO: + likelihood_dto = agent.likelihood_smiles(sampled_sequences) + dto = learning_strategy.run(likelihood_dto, score) + return dto + + def _logging(self, start_time, step, score_summary, dto: UpdatedLikelihoodsDTO, agent): + report_dto = TimestepDTO(start_time, self._parameters.number_of_steps, step, score_summary, + dto.agent_likelihood, dto.prior_likelihood, dto.augmented_likelihood) + self._logger.timestep_report(report_dto, self._diversity_filter, agent) + + def _apply_scoring_function(self, scoring_function, sampled_sequences: List[SampledSequencesDTO], + step) -> FinalSummary: + molecules = self._join_linker_and_warheads(sampled_sequences, keep_labels=True) + smiles = [] + for idx, molecule in enumerate(molecules): + try: + smiles_str = self._conversion.mol_to_smiles(molecule) if molecule else "INVALID" + except RuntimeError as exception: + # NOTE: Current implementation of BondMaker (reinvent_chemistry.library_design.bond_maker) results in + # impossible conversion of mol to smiles if one single atom has two attachment points and labels are + # kept. As this case is not relevant in the context of link_invent, then can be discarded as invalid. + smiles_str = "INVALID" + self._logger.log_message(exception.__str__() + f'\n\tinput: {sampled_sequences[idx].input}' + f'\n\toutput: {sampled_sequences[idx].output}\n') + finally: + smiles.append(smiles_str) + final_score: FinalSummary = scoring_function.get_final_score_for_step(smiles, step) + return final_score + + def _join_linker_and_warheads(self, sampled_sequences: List[SampledSequencesDTO], keep_labels=False): + molecules = [] + for sample in sampled_sequences: + linker = self._attachment_points.add_attachment_point_numbers(sample.output, canonicalize=False) + molecule = self._bond_maker.join_scaffolds_and_decorations(linker, sample.input, + keep_labels_on_atoms=keep_labels) + molecules.append(molecule) + return molecules + + def _clean_scored_smiles(self, score_summary: FinalSummary) -> FinalSummary: + """ + Remove attachment point numbers from scored smiles + """ + # Note: method AttachmentPoints.remove_attachment_point_numbers does not work in this context, as it searches + # for attachment point token ('*') + score_summary.scored_smiles = [self._conversion.mol_to_smiles( + self._attachment_points.remove_attachment_point_numbers_from_mol(self._conversion.smile_to_mol(smile)) + ) if idx in score_summary.valid_idxs else smile for idx, smile in enumerate(score_summary.scored_smiles)] + return score_summary diff --git a/running_modes/automated_curriculum_learning/production_strategy/production_strategy.py b/running_modes/automated_curriculum_learning/production_strategy/production_strategy.py index 22bd031..f071d2b 100644 --- a/running_modes/automated_curriculum_learning/production_strategy/production_strategy.py +++ b/running_modes/automated_curriculum_learning/production_strategy/production_strategy.py @@ -1,30 +1,39 @@ -from reinvent_models.reinvent_core.models.model import Model +from reinvent_models.model_factory.generative_model_base import GenerativeModelBase from reinvent_scoring import ScoringFunctionFactory -from reinvent_scoring.scoring.diversity_filters.reinvent_core.base_diversity_filter import BaseDiversityFilter +from reinvent_scoring.scoring.diversity_filters.curriculum_learning.diversity_filter import DiversityFilter -from running_modes.automated_curriculum_learning.logging.base_auto_cl_logger import BaseAutoCLLogger +from running_modes.automated_curriculum_learning.inception.inception import Inception +from running_modes.automated_curriculum_learning.logging.base_logger import BaseLogger +from running_modes.automated_curriculum_learning.production_strategy.link_invent_production_strategy import \ + LinkInventProductionStrategy from running_modes.automated_curriculum_learning.production_strategy.base_production_strategy import \ BaseProductionStrategy -from running_modes.automated_curriculum_learning.production_strategy.standard_production_strategy import \ - StandardProductionStrategy -from running_modes.configurations.automated_curriculum_learning.production_strategy_configuration import \ - ProductionStrategyConfiguration +from running_modes.automated_curriculum_learning.production_strategy.reinvent_production_strategy import \ + ReinventProductionStrategy +from running_modes.configurations import ProductionStrategyInputConfiguration from running_modes.enums.production_strategy_enum import ProductionStrategyEnum -from running_modes.reinforcement_learning.inception import Inception class ProductionStrategy: - def __new__(cls, prior: Model, diversity_filter: BaseDiversityFilter, inception: Inception, - configuration: ProductionStrategyConfiguration, logger: BaseAutoCLLogger) -> BaseProductionStrategy: + def __new__(cls, prior: GenerativeModelBase, inception: Inception, + configuration: ProductionStrategyInputConfiguration, + logger: BaseLogger) -> BaseProductionStrategy: production_strategy_enum = ProductionStrategyEnum() scoring_function_instance = ScoringFunctionFactory(configuration.scoring_function) + diversity_filter = DiversityFilter(configuration.diversity_filter) if production_strategy_enum.STANDARD == configuration.name: - return StandardProductionStrategy(prior=prior, + production = ReinventProductionStrategy(prior=prior, diversity_filter=diversity_filter, inception=inception, scoring_function=scoring_function_instance, configuration=configuration, logger=logger) - + return production + elif production_strategy_enum.LINK_INVENT == configuration.name: + production = LinkInventProductionStrategy(prior=prior, diversity_filter=diversity_filter, + inception=inception, + scoring_function=scoring_function_instance, + configuration=configuration, logger=logger) + return production else: raise NotImplementedError(f"Unknown production strategy {configuration.name}") diff --git a/running_modes/automated_curriculum_learning/production_strategy/reinvent_production_strategy.py b/running_modes/automated_curriculum_learning/production_strategy/reinvent_production_strategy.py new file mode 100644 index 0000000..076143d --- /dev/null +++ b/running_modes/automated_curriculum_learning/production_strategy/reinvent_production_strategy.py @@ -0,0 +1,82 @@ +import time +from abc import abstractmethod +from typing import List, Dict, Tuple + +import numpy as np +import torch +from reinvent_models.model_factory.generative_model_base import GenerativeModelBase +from reinvent_scoring import ScoringFunctionFactory, FinalSummary, ScoringFunctionParameters +from reinvent_scoring.scoring.diversity_filters.curriculum_learning.update_diversity_filter_dto import \ + UpdateDiversityFilterDTO +from reinvent_scoring.scoring.function.base_scoring_function import BaseScoringFunction + +from running_modes.automated_curriculum_learning.actions.reinvent_sample_model import ReinventSampleModel +from running_modes.automated_curriculum_learning.dto import SampledBatchDTO, TimestepDTO +from running_modes.automated_curriculum_learning.learning_strategy.learning_strategy import LearningStrategy +from running_modes.automated_curriculum_learning.production_strategy.base_production_strategy import \ + BaseProductionStrategy + + +class ReinventProductionStrategy(BaseProductionStrategy): + + def run(self, cl_agent: GenerativeModelBase, steps_so_far: int): + start_time = time.time() + self._disable_prior_gradients() + optimizer = torch.optim.Adam(cl_agent.get_network_parameters(), lr=self._parameters.learning_rate) + learning_strategy = LearningStrategy(self._prior, optimizer, self._parameters.learning_strategy, self._logger) + + step_limit = steps_so_far + self._parameters.number_of_steps + + for step in range(steps_so_far, step_limit): + self.take_step(agent=cl_agent, learning_strategy=learning_strategy, scoring_function=self._scoring_function, + step=step, start_time=start_time) + + self._logger.log_message(f"Production finished at step {step_limit}") + self._logger.save_final_state(cl_agent, self._diversity_filter) + + def setup_scoring_function(self, name: str, parameter_list: List[Dict]) -> BaseScoringFunction: + scoring_function_parameters = ScoringFunctionParameters(name=name, parameters=parameter_list, parallel=False) + scoring_function_instance = ScoringFunctionFactory(scoring_function_parameters) + self._log_sf_update(current_parameters=parameter_list) + return scoring_function_instance + + def _disable_prior_gradients(self): + for param in self._prior.get_network_parameters(): + param.requires_grad = False + + def take_step(self, agent: GenerativeModelBase, learning_strategy, scoring_function: BaseScoringFunction, + step:int, start_time: float) -> float: + # 1. Sampling + sampled = self._sampling(agent) + # 2. Scoring + score, score_summary = self._scoring(scoring_function, sampled.smiles, step) + # 3. Updating + agent_likelihood, prior_likelihood, augmented_likelihood = self._updating(sampled, score, self.inception, agent, learning_strategy) + # 4. Logging + self._logging(agent=agent, start_time=start_time, step=step, + score_summary=score_summary, agent_likelihood=agent_likelihood, + prior_likelihood=prior_likelihood, augmented_likelihood=augmented_likelihood) + + score = score.mean() + return score + + def _sampling(self, agent) -> SampledBatchDTO: + sampling_action = ReinventSampleModel(agent, self._parameters.batch_size, self._logger) + sampled_sequences = sampling_action.run() + return sampled_sequences + + def _scoring(self, scoring_function, smiles: List[str], step) -> Tuple[np.ndarray, FinalSummary] : + score_summary = scoring_function.get_final_score_for_step(smiles, step) + dto = UpdateDiversityFilterDTO(score_summary, [], step) + score = self._diversity_filter.update_score(dto) + return score, score_summary + + def _updating(self, sampled, score, inception, agent, learning_strategy): + agent_likelihood, prior_likelihood, augmented_likelihood = learning_strategy.run(sampled, score, inception, agent) + return agent_likelihood, prior_likelihood, augmented_likelihood + + def _logging(self, agent: GenerativeModelBase, start_time: float, step: int, score_summary: FinalSummary, + agent_likelihood: torch.tensor, prior_likelihood: torch.tensor, augmented_likelihood: torch.tensor): + report_dto = TimestepDTO(start_time, self._parameters.number_of_steps, step, score_summary, + agent_likelihood, prior_likelihood, augmented_likelihood) + self._logger.timestep_report(report_dto, self._diversity_filter, agent) \ No newline at end of file diff --git a/running_modes/automated_curriculum_learning/production_strategy/standard_production_strategy.py b/running_modes/automated_curriculum_learning/production_strategy/standard_production_strategy.py deleted file mode 100644 index c04fda7..0000000 --- a/running_modes/automated_curriculum_learning/production_strategy/standard_production_strategy.py +++ /dev/null @@ -1,40 +0,0 @@ -import time - -import torch -from reinvent_models.reinvent_core.models.model import Model -from reinvent_scoring.scoring.diversity_filters.reinvent_core.base_diversity_filter import BaseDiversityFilter -from reinvent_scoring.scoring.function.base_scoring_function import BaseScoringFunction - -from running_modes.automated_curriculum_learning.logging.base_auto_cl_logger import BaseAutoCLLogger -from running_modes.automated_curriculum_learning.production_strategy.base_production_strategy import \ - BaseProductionStrategy -from running_modes.configurations.automated_curriculum_learning.production_strategy_configuration import \ - ProductionStrategyConfiguration -from running_modes.reinforcement_learning.inception import Inception - - -class StandardProductionStrategy(BaseProductionStrategy): - def __init__(self, prior: Model, - diversity_filter: BaseDiversityFilter, inception: Inception, scoring_function: BaseScoringFunction, - configuration: ProductionStrategyConfiguration, - logger: BaseAutoCLLogger): - super().__init__(prior=prior, - diversity_filter=diversity_filter, inception=inception, scoring_function=scoring_function, - configuration=configuration, - logger=logger) - - def run(self, cl_agent: Model, steps_so_far: int): - # self._diversity_filter.flush_memory() - start_time = time.time() - self._disable_prior_gradients() - optimizer = torch.optim.Adam(cl_agent.network.parameters(), lr=self._parameters.learning_rate) - - step_limit = steps_so_far + self._parameters.n_steps - - for step in range(steps_so_far, step_limit): - self._take_step(agent=cl_agent, optimizer=optimizer, scoring_function_current=self._scoring_function, - step=step, start_time=start_time) - - self._logger.log_message(f"Production finished at step {step_limit}") - self._logger.save_final_state(cl_agent, self._diversity_filter) - self._logger.log_out_input_configuration() diff --git a/running_modes/configurations/automated_curriculum_learning/__init__.py b/running_modes/configurations/automated_curriculum_learning/__init__.py index 7f62319..4606a99 100644 --- a/running_modes/configurations/automated_curriculum_learning/__init__.py +++ b/running_modes/configurations/automated_curriculum_learning/__init__.py @@ -1,6 +1,12 @@ +from running_modes.configurations.automated_curriculum_learning.automated_curriculum_learning_input_configuration import \ + AutomatedCurriculumLearningInputConfiguration from running_modes.configurations.automated_curriculum_learning.curriculum_strategy_configuration import \ CurriculumStrategyConfiguration +from running_modes.configurations.automated_curriculum_learning.curriculum_strategy_input_configuration import \ + CurriculumStrategyInputConfiguration from running_modes.configurations.automated_curriculum_learning.inception_configuration import InceptionConfiguration from running_modes.configurations.automated_curriculum_learning.automated_curriculum_learning_configuration import \ AutomatedCLConfiguration -from running_modes.configurations.automated_curriculum_learning.production_strategy_configuration import ProductionStrategyConfiguration +from running_modes.configurations.automated_curriculum_learning.prodcution_strategy_input_configuration import \ + ProductionStrategyInputConfiguration + diff --git a/running_modes/configurations/automated_curriculum_learning/automated_curriculum_learning_configuration.py b/running_modes/configurations/automated_curriculum_learning/automated_curriculum_learning_configuration.py index 7aa19b1..4224251 100644 --- a/running_modes/configurations/automated_curriculum_learning/automated_curriculum_learning_configuration.py +++ b/running_modes/configurations/automated_curriculum_learning/automated_curriculum_learning_configuration.py @@ -1,5 +1,3 @@ -from dataclasses import dataclass - from running_modes.configurations.automated_curriculum_learning.base_configuration import BaseConfiguration from running_modes.configurations.automated_curriculum_learning.curriculum_strategy_configuration import \ CurriculumStrategyConfiguration @@ -7,7 +5,6 @@ ProductionStrategyConfiguration -@dataclass class AutomatedCLConfiguration(BaseConfiguration): prior: str agent: str diff --git a/running_modes/configurations/automated_curriculum_learning/automated_curriculum_learning_input_configuration.py b/running_modes/configurations/automated_curriculum_learning/automated_curriculum_learning_input_configuration.py new file mode 100644 index 0000000..8257e88 --- /dev/null +++ b/running_modes/configurations/automated_curriculum_learning/automated_curriculum_learning_input_configuration.py @@ -0,0 +1,12 @@ +from running_modes.configurations.automated_curriculum_learning.base_configuration import BaseConfiguration +from running_modes.configurations.automated_curriculum_learning.curriculum_strategy_input_configuration import \ + CurriculumStrategyInputConfiguration +from running_modes.configurations.automated_curriculum_learning.prodcution_strategy_input_configuration import \ + ProductionStrategyInputConfiguration + + +class AutomatedCurriculumLearningInputConfiguration(BaseConfiguration): + agent: str + prior: str + curriculum_strategy: CurriculumStrategyInputConfiguration + production_strategy: ProductionStrategyInputConfiguration diff --git a/running_modes/configurations/automated_curriculum_learning/automated_curriculum_learning_linkinvent_configuration.py b/running_modes/configurations/automated_curriculum_learning/automated_curriculum_learning_linkinvent_configuration.py new file mode 100644 index 0000000..851d5a5 --- /dev/null +++ b/running_modes/configurations/automated_curriculum_learning/automated_curriculum_learning_linkinvent_configuration.py @@ -0,0 +1,12 @@ +from running_modes.configurations.automated_curriculum_learning.base_configuration import BaseConfiguration +from running_modes.configurations.automated_curriculum_learning.linkinvent_curriculum_strategy_configuration import \ + LinkInventCurriculumStrategyConfiguration +from running_modes.configurations.automated_curriculum_learning.linkinvent_production_strategy_configuration import \ + LinkInventProductionStrategyConfiguration + + +class AutomatedCurriculumLearningLinkInventConfiguration(BaseConfiguration): + agent: str + prior: str + curriculum_strategy: LinkInventCurriculumStrategyConfiguration + production_strategy: LinkInventProductionStrategyConfiguration \ No newline at end of file diff --git a/running_modes/configurations/automated_curriculum_learning/base_configuration.py b/running_modes/configurations/automated_curriculum_learning/base_configuration.py index 43a470b..3e734a4 100644 --- a/running_modes/configurations/automated_curriculum_learning/base_configuration.py +++ b/running_modes/configurations/automated_curriculum_learning/base_configuration.py @@ -1,6 +1,5 @@ -from dataclasses import dataclass +from pydantic import BaseModel -@dataclass -class BaseConfiguration: +class BaseConfiguration(BaseModel): curriculum_type: str \ No newline at end of file diff --git a/running_modes/configurations/automated_curriculum_learning/curriculum_strategy_configuration.py b/running_modes/configurations/automated_curriculum_learning/curriculum_strategy_configuration.py index e015b71..16d12f1 100644 --- a/running_modes/configurations/automated_curriculum_learning/curriculum_strategy_configuration.py +++ b/running_modes/configurations/automated_curriculum_learning/curriculum_strategy_configuration.py @@ -1,9 +1,10 @@ -from dataclasses import dataclass from typing import List -from reinvent_scoring.scoring.diversity_filters.reinvent_core.diversity_filter_parameters import \ - DiversityFilterParameters +from pydantic.dataclasses import dataclass +from reinvent_scoring.scoring.diversity_filters.curriculum_learning import DiversityFilterParameters +from running_modes.automated_curriculum_learning.learning_strategy.learning_strategy_configuration import \ + LearningStrategyConfiguration from running_modes.configurations import InceptionConfiguration from running_modes.configurations.automated_curriculum_learning.curriculum_objective import CurriculumObjective @@ -11,10 +12,11 @@ @dataclass class CurriculumStrategyConfiguration: name: str + learning_strategy: LearningStrategyConfiguration curriculum_objectives: List[CurriculumObjective] diversity_filter: DiversityFilterParameters inception: InceptionConfiguration max_num_iterations: int batch_size: int = 64 learning_rate: float = 0.0001 - sigma: float = 120 + sigma: float = 120. diff --git a/running_modes/configurations/automated_curriculum_learning/curriculum_strategy_input_configuration.py b/running_modes/configurations/automated_curriculum_learning/curriculum_strategy_input_configuration.py new file mode 100644 index 0000000..5bd5e42 --- /dev/null +++ b/running_modes/configurations/automated_curriculum_learning/curriculum_strategy_input_configuration.py @@ -0,0 +1,26 @@ +from dataclasses import dataclass +from typing import List + +from pydantic import Field +from reinvent_scoring.scoring.diversity_filters.curriculum_learning import DiversityFilterParameters + +from running_modes.automated_curriculum_learning.learning_strategy.learning_strategy_configuration import \ + LearningStrategyConfiguration +from running_modes.configurations.automated_curriculum_learning.curriculum_objective import CurriculumObjective +from running_modes.configurations.automated_curriculum_learning.inception_configuration import InceptionConfiguration + + +@dataclass +class CurriculumStrategyInputConfiguration: + name: str + learning_strategy: LearningStrategyConfiguration + curriculum_objectives: List[CurriculumObjective] + diversity_filter: DiversityFilterParameters + inception: InceptionConfiguration + max_num_iterations: int + input: List[str] = Field(default_factory=list) + randomize_input: bool = False + batch_size: int = 64 + learning_rate: float = 0.0001 + sigma: float = 120. + distance_threshold: float = 100. \ No newline at end of file diff --git a/running_modes/configurations/automated_curriculum_learning/inception_configuration.py b/running_modes/configurations/automated_curriculum_learning/inception_configuration.py index 436cebe..1e8bf5d 100644 --- a/running_modes/configurations/automated_curriculum_learning/inception_configuration.py +++ b/running_modes/configurations/automated_curriculum_learning/inception_configuration.py @@ -1,10 +1,12 @@ from dataclasses import dataclass - from typing import List +from pydantic import Field @dataclass class InceptionConfiguration: - smiles: List[str] memory_size: int sample_size: int + smiles: List[str] = Field(default_factory=list) + # inputs: List[str] = Field(default_factory=list) + diff --git a/running_modes/configurations/automated_curriculum_learning/linkinvent_curriculum_strategy_configuration.py b/running_modes/configurations/automated_curriculum_learning/linkinvent_curriculum_strategy_configuration.py new file mode 100644 index 0000000..ac3c721 --- /dev/null +++ b/running_modes/configurations/automated_curriculum_learning/linkinvent_curriculum_strategy_configuration.py @@ -0,0 +1,25 @@ +from dataclasses import dataclass +from typing import List + +from reinvent_scoring.scoring.diversity_filters.curriculum_learning import DiversityFilterParameters + +from running_modes.automated_curriculum_learning.learning_strategy.learning_strategy_configuration import \ + LearningStrategyConfiguration +from running_modes.configurations import InceptionConfiguration +from running_modes.configurations.automated_curriculum_learning.curriculum_objective import CurriculumObjective + + + +@dataclass +class LinkInventCurriculumStrategyConfiguration: + name: str + input: List[str] + learning_strategy: LearningStrategyConfiguration + curriculum_objectives: List[CurriculumObjective] + diversity_filter: DiversityFilterParameters + inception: InceptionConfiguration + max_num_iterations: int + batch_size: int = 64 + learning_rate: float = 0.0001 + sigma: float = 120 + randomize_input: bool = False diff --git a/running_modes/configurations/automated_curriculum_learning/linkinvent_production_strategy_configuration.py b/running_modes/configurations/automated_curriculum_learning/linkinvent_production_strategy_configuration.py new file mode 100644 index 0000000..cbc6ed5 --- /dev/null +++ b/running_modes/configurations/automated_curriculum_learning/linkinvent_production_strategy_configuration.py @@ -0,0 +1,25 @@ +from dataclasses import dataclass +from typing import List + +from reinvent_scoring import ScoringFunctionParameters +from reinvent_scoring.scoring.diversity_filters.curriculum_learning import DiversityFilterParameters + +from running_modes.configurations import InceptionConfiguration +from running_modes.reinforcement_learning.configurations.learning_strategy_configuration import \ + LearningStrategyConfiguration + + +@dataclass +class LinkInventProductionStrategyConfiguration: + name: str + input: List[str] + learning_strategy: LearningStrategyConfiguration + scoring_function: ScoringFunctionParameters + diversity_filter: DiversityFilterParameters + inception: InceptionConfiguration + retain_inception: bool + batch_size: int = 64 + learning_rate: float = 0.0001 + sigma: float = 120 + number_of_steps: int = 100 + randomize_input: bool = False diff --git a/running_modes/configurations/automated_curriculum_learning/prodcution_strategy_input_configuration.py b/running_modes/configurations/automated_curriculum_learning/prodcution_strategy_input_configuration.py new file mode 100644 index 0000000..92a9d5c --- /dev/null +++ b/running_modes/configurations/automated_curriculum_learning/prodcution_strategy_input_configuration.py @@ -0,0 +1,27 @@ +from dataclasses import dataclass +from typing import List + +from pydantic import Field +from reinvent_scoring import ScoringFunctionParameters +from reinvent_scoring.scoring.diversity_filters.curriculum_learning import DiversityFilterParameters + +from running_modes.automated_curriculum_learning.learning_strategy.learning_strategy_configuration import \ + LearningStrategyConfiguration +from running_modes.configurations import InceptionConfiguration + + +@dataclass +class ProductionStrategyInputConfiguration: + name: str + learning_strategy: LearningStrategyConfiguration + scoring_function: ScoringFunctionParameters + diversity_filter: DiversityFilterParameters + inception: InceptionConfiguration + retain_inception: bool + input: List[str] = Field(default_factory=list) + randomize_input: bool = False + batch_size: int = 64 + learning_rate: float = 0.0001 + sigma: float = 120 + number_of_steps: int = 100 + distance_threshold: float = 100. \ No newline at end of file diff --git a/running_modes/configurations/automated_curriculum_learning/production_strategy_configuration.py b/running_modes/configurations/automated_curriculum_learning/production_strategy_configuration.py index 2b59d7c..0ed6c5e 100644 --- a/running_modes/configurations/automated_curriculum_learning/production_strategy_configuration.py +++ b/running_modes/configurations/automated_curriculum_learning/production_strategy_configuration.py @@ -1,15 +1,15 @@ -from dataclasses import dataclass - +from pydantic import BaseModel from reinvent_scoring import ScoringFunctionParameters -from reinvent_scoring.scoring.diversity_filters.reinvent_core.diversity_filter_parameters import \ - DiversityFilterParameters +from reinvent_scoring.scoring.diversity_filters.curriculum_learning import DiversityFilterParameters +from running_modes.automated_curriculum_learning.learning_strategy.learning_strategy_configuration import \ + LearningStrategyConfiguration from running_modes.configurations import InceptionConfiguration -@dataclass -class ProductionStrategyConfiguration: +class ProductionStrategyConfiguration(BaseModel): name: str + learning_strategy: LearningStrategyConfiguration scoring_function: ScoringFunctionParameters diversity_filter: DiversityFilterParameters inception: InceptionConfiguration @@ -17,4 +17,4 @@ class ProductionStrategyConfiguration: batch_size: int = 64 learning_rate: float = 0.0001 sigma: float = 120 - n_steps: int = 100 + number_of_steps: int = 100 diff --git a/running_modes/configurations/curriculum_learning/curriculum_learning_components.py b/running_modes/configurations/curriculum_learning/curriculum_learning_components.py index 7bcfd3a..d9d9216 100644 --- a/running_modes/configurations/curriculum_learning/curriculum_learning_components.py +++ b/running_modes/configurations/curriculum_learning/curriculum_learning_components.py @@ -1,9 +1,6 @@ -from dataclasses import dataclass - from running_modes.configurations.automated_curriculum_learning.base_configuration import BaseConfiguration -@dataclass class CurriculumLearningComponents(BaseConfiguration): """This class holds the necessary configuration components to run CL""" curriculum_learning: dict diff --git a/running_modes/configurations/curriculum_learning/curriculum_learning_configuration.py b/running_modes/configurations/curriculum_learning/curriculum_learning_configuration.py index 72fb667..b8d6435 100644 --- a/running_modes/configurations/curriculum_learning/curriculum_learning_configuration.py +++ b/running_modes/configurations/curriculum_learning/curriculum_learning_configuration.py @@ -1,8 +1,7 @@ -from dataclasses import dataclass +from pydantic import BaseModel -@dataclass -class CurriculumLearningConfiguration: +class CurriculumLearningConfiguration(BaseModel): prior: str agent: str update_lock: str diff --git a/running_modes/configurations/logging/curriculum_log_configuration.py b/running_modes/configurations/logging/curriculum_log_configuration.py index 302129e..17c5351 100644 --- a/running_modes/configurations/logging/curriculum_log_configuration.py +++ b/running_modes/configurations/logging/curriculum_log_configuration.py @@ -4,3 +4,6 @@ class CurriculumLoggerConfiguration(BaseLoggerConfiguration): result_folder: str logging_frequency: int = 0 + # rows and columns define the shape of the output grid of molecule images in tensorboard. + rows: int = 4 + columns: int = 4 diff --git a/running_modes/configurations/reinforcement_learning/reinforcement_learning_configuration.py b/running_modes/configurations/reinforcement_learning/reinforcement_learning_configuration.py index dc53607..945f846 100644 --- a/running_modes/configurations/reinforcement_learning/reinforcement_learning_configuration.py +++ b/running_modes/configurations/reinforcement_learning/reinforcement_learning_configuration.py @@ -9,6 +9,4 @@ class ReinforcementLearningConfiguration: sigma: int = 120 learning_rate: float = 0.0001 batch_size: int = 128 - reset: int = 0 - reset_score_cutoff: float = 0.5 margin_threshold: int = 50 diff --git a/running_modes/constructors/automated_curriculum_learning_mode_constructor.py b/running_modes/constructors/automated_curriculum_learning_mode_constructor.py deleted file mode 100644 index d46549d..0000000 --- a/running_modes/constructors/automated_curriculum_learning_mode_constructor.py +++ /dev/null @@ -1,20 +0,0 @@ -from dacite import from_dict -from reinvent_models.reinvent_core.models.model import Model - -from running_modes.automated_curriculum_learning.automated_curriculum_runner import AutomatedCurriculumRunner -from running_modes.automated_curriculum_learning.logging import AutoCLLogger -from running_modes.configurations import GeneralConfigurationEnvelope, AutomatedCLConfiguration -from running_modes.constructors.base_running_mode import BaseRunningMode -from running_modes.utils.general import set_default_device_cuda - - -class AutomatedCurriculumLearningModeConstructor: - def __new__(self, configuration: GeneralConfigurationEnvelope) -> BaseRunningMode: - self._configuration = configuration - config = from_dict(data_class=AutomatedCLConfiguration, data=self._configuration.parameters) - set_default_device_cuda() - _logger = AutoCLLogger(self._configuration) - _prior = Model.load_from_file(config.prior) - _agent = Model.load_from_file(config.agent) - runner = AutomatedCurriculumRunner(config, _logger, _prior, _agent) - return runner diff --git a/running_modes/constructors/curriculum_learning_mode_constructor.py b/running_modes/constructors/curriculum_learning_mode_constructor.py index cdd1d51..f222b25 100644 --- a/running_modes/constructors/curriculum_learning_mode_constructor.py +++ b/running_modes/constructors/curriculum_learning_mode_constructor.py @@ -1,13 +1,17 @@ -from dacite import from_dict -from reinvent_models.reinvent_core.models.model import Model +from reinvent_models.lib_invent.enums.generative_model_regime import GenerativeModelRegimeEnum +from reinvent_models.model_factory.configurations.model_configuration import ModelConfiguration +from reinvent_models.model_factory.generative_model import GenerativeModel from running_modes.automated_curriculum_learning.automated_curriculum_runner import AutomatedCurriculumRunner from running_modes.automated_curriculum_learning.logging import AutoCLLogger -from running_modes.configurations import GeneralConfigurationEnvelope, AutomatedCLConfiguration +from running_modes.configurations import GeneralConfigurationEnvelope +from running_modes.configurations.automated_curriculum_learning.automated_curriculum_learning_input_configuration import \ + AutomatedCurriculumLearningInputConfiguration from running_modes.configurations.automated_curriculum_learning.base_configuration import BaseConfiguration from running_modes.constructors.base_running_mode import BaseRunningMode from running_modes.curriculum_learning.curriculum_runner import CurriculumRunner from running_modes.enums.curriculum_type_enum import CurriculumTypeEnum +from running_modes.enums.model_type_enum import ModelTypeEnum from running_modes.utils.general import set_default_device_cuda @@ -15,17 +19,38 @@ class CurriculumLearningModeConstructor: def __new__(self, configuration: GeneralConfigurationEnvelope) -> BaseRunningMode: self._configuration = configuration cl_enum = CurriculumTypeEnum - base_config = from_dict(data_class=BaseConfiguration, data=self._configuration.parameters) - set_default_device_cuda() + + base_config = BaseConfiguration.parse_obj(self._configuration.parameters) if base_config.curriculum_type == cl_enum.MANUAL: + set_default_device_cuda() runner = CurriculumRunner(self._configuration) + elif base_config.curriculum_type == cl_enum.AUTOMATED: + runner = self._create_automated_curriculum(self._configuration) else: - self._configuration = configuration - config = from_dict(data_class=AutomatedCLConfiguration, data=self._configuration.parameters) - _logger = AutoCLLogger(self._configuration) - _prior = Model.load_from_file(config.prior) - _agent = Model.load_from_file(config.agent) - runner = AutomatedCurriculumRunner(config, _logger, _prior, _agent) - - return runner \ No newline at end of file + raise KeyError(f"Incorrect curriculum type: `{base_config.curriculum_type}` provided") + + return runner + + @staticmethod + def _create_automated_curriculum(configuration): + model_type = ModelTypeEnum() + model_regime = GenerativeModelRegimeEnum() + + if model_type.DEFAULT == configuration.model_type: + set_default_device_cuda() + config = AutomatedCurriculumLearningInputConfiguration.parse_obj(configuration.parameters) + elif model_type.LINK_INVENT == configuration.model_type: + set_default_device_cuda() + config = AutomatedCurriculumLearningInputConfiguration.parse_obj(configuration.parameters) + else: + raise KeyError(f"Incorrect model type: `{configuration.model_type}` provided") + + _logger = AutoCLLogger(configuration) + prior_config = ModelConfiguration(configuration.model_type, model_regime.INFERENCE, config.prior) + agent_config = ModelConfiguration(configuration.model_type, model_regime.TRAINING, config.agent) + _prior = GenerativeModel(prior_config) + _agent = GenerativeModel(agent_config) + + runner = AutomatedCurriculumRunner(config, _logger, _prior, _agent) + return runner diff --git a/running_modes/constructors/running_mode.py b/running_modes/constructors/running_mode.py index 13ba18a..a8f985e 100644 --- a/running_modes/constructors/running_mode.py +++ b/running_modes/constructors/running_mode.py @@ -28,7 +28,5 @@ def __new__(cls, configuration: GeneralConfigurationEnvelope) -> BaseRunningMode return CreateModelModeConstructor(configuration) if configuration.run_type == running_mode_enum.VALIDATION: return ValidationModeConstructor(configuration) - # if configuration.run_type == running_mode_enum.AUTOMATED_CURRICULUM_LEARNING: - # return AutomatedCurriculumLearningModeConstructor(configuration) else: raise TypeError(f"Requested run type: '{configuration.run_type}' is not implemented.") \ No newline at end of file diff --git a/running_modes/curriculum_learning/curriculum_runner.py b/running_modes/curriculum_learning/curriculum_runner.py index 7861ff9..af34312 100644 --- a/running_modes/curriculum_learning/curriculum_runner.py +++ b/running_modes/curriculum_learning/curriculum_runner.py @@ -2,44 +2,51 @@ import numpy as np import torch - -from reinvent_models.reinvent_core.models.model import Model - +from reinvent_chemistry.utils import get_indices_of_unique_smiles +from reinvent_models.lib_invent.enums.generative_model_regime import GenerativeModelRegimeEnum +from reinvent_models.model_factory.configurations.model_configuration import ModelConfiguration +from reinvent_models.model_factory.generative_model import GenerativeModel from reinvent_scoring.scoring.diversity_filters.reinvent_core.diversity_filter import DiversityFilter -from reinvent_scoring.scoring.diversity_filters.reinvent_core.diversity_filter_parameters import DiversityFilterParameters -from running_modes.constructors.base_running_mode import BaseRunningMode +from reinvent_scoring.scoring.diversity_filters.reinvent_core.diversity_filter_parameters import \ + DiversityFilterParameters +from reinvent_scoring.scoring.score_summary import FinalSummary +from reinvent_scoring.scoring.scoring_function_factory import ScoringFunctionFactory +from reinvent_scoring.scoring.scoring_function_parameters import ScoringFunctionParameters + +from running_modes.automated_curriculum_learning.inception.inception import Inception from running_modes.configurations import GeneralConfigurationEnvelope, InceptionConfiguration from running_modes.configurations.curriculum_learning.curriculum_learning_components import CurriculumLearningComponents from running_modes.configurations.curriculum_learning.curriculum_learning_configuration import \ CurriculumLearningConfiguration +from running_modes.constructors.base_running_mode import BaseRunningMode from running_modes.curriculum_learning.logging import CurriculumLogger from running_modes.curriculum_learning.update_watcher import UpdateWatcher -from running_modes.reinforcement_learning.inception import Inception +from running_modes.enums.model_type_enum import ModelTypeEnum from running_modes.reinforcement_learning.margin_guard import MarginGuard from running_modes.utils import to_tensor -from reinvent_chemistry.utils import get_indices_of_unique_smiles - -from reinvent_scoring.scoring.score_summary import FinalSummary -from reinvent_scoring.scoring.scoring_function_factory import ScoringFunctionFactory -from reinvent_scoring.scoring.scoring_function_parameters import ScoringFunctionParameters - class CurriculumRunner(BaseRunningMode): def __init__(self, envelope: GeneralConfigurationEnvelope): self.envelope = envelope config_components = CurriculumLearningComponents(**self.envelope.parameters) self.config = CurriculumLearningConfiguration(**config_components.curriculum_learning) - self._prior = Model.load_from_file(self.config.prior) - self._agent = Model.load_from_file(self.config.agent) + + model_regime = GenerativeModelRegimeEnum() + prior_config = ModelConfiguration(ModelTypeEnum().DEFAULT, model_regime.INFERENCE, self.config.prior) + agent_config = ModelConfiguration(ModelTypeEnum().DEFAULT, model_regime.TRAINING, self.config.agent) + _prior = GenerativeModel(prior_config) + _agent = GenerativeModel(agent_config) + + self._prior = _prior + self._agent = _agent self.logger = CurriculumLogger(self.envelope) self.scoring_function = self.setup_scoring_function(config_components.scoring_function) self.diversity_filter = self._setup_diversity_filter(config_components.diversity_filter) self.inception = self.setup_inception(config_components.inception) self._margin_guard = MarginGuard(self) - self._optimizer = torch.optim.Adam(self._agent.network.parameters(), lr=self.config.learning_rate) + self._optimizer = torch.optim.Adam(self._agent.get_network_parameters(), lr=self.config.learning_rate) - assert self._prior.vocabulary == self._agent.vocabulary, "The agent and the prior must have the same vocabulary" self._update_watcher = UpdateWatcher(self) def run(self): @@ -84,7 +91,7 @@ def run(self): def _disable_prior_gradients(self): # There might be a more elegant way of disabling gradients - for param in self._prior.network.parameters(): + for param in self._prior.get_network_parameters(): param.requires_grad = False def _stats_and_chekpoint(self, score, start_time, step, smiles, score_summary: FinalSummary, @@ -98,7 +105,7 @@ def _stats_and_chekpoint(self, score, start_time, step, smiles, score_summary: F return self._update_reset_countdown(reset_countdown, mean_score) def _sample_unique_sequences(self, agent, batch_size): - seqs, smiles, agent_likelihood = agent.sample_sequences_and_smiles(batch_size) + seqs, smiles, agent_likelihood = agent.sample(batch_size) unique_idxs = get_indices_of_unique_smiles(smiles) seqs_unique = seqs[unique_idxs] smiles_np = np.array(smiles) @@ -139,8 +146,11 @@ def _inception_filter(self, agent, loss, agent_likelihood, prior_likelihood, return loss, agent_likelihood def reset(self, reset_countdown=0): - self._agent = Model.load_from_file(self.config.agent) - self._optimizer = torch.optim.Adam(self._agent.network.parameters(), lr=self.config.learning_rate) + model_regime = GenerativeModelRegimeEnum() + agent_config = ModelConfiguration(self.envelope.model_type, model_regime.TRAINING, self.config.agent) + + self._agent = GenerativeModel(agent_config) + self._optimizer = torch.optim.Adam(self._agent.get_network_parameters(), lr=self.config.learning_rate) self.logger.log_message("Resetting Agent") self.logger.log_message(f"Adjusting sigma to: {self.config.sigma}") return reset_countdown diff --git a/running_modes/curriculum_learning/logging/base_curriculum_logger.py b/running_modes/curriculum_learning/logging/base_curriculum_logger.py index e8ece8b..1a9d013 100644 --- a/running_modes/curriculum_learning/logging/base_curriculum_logger.py +++ b/running_modes/curriculum_learning/logging/base_curriculum_logger.py @@ -6,9 +6,9 @@ import torch from reinvent_scoring import ComponentSpecificParametersEnum +from running_modes.automated_curriculum_learning.inception.inception import Inception from running_modes.configurations.general_configuration_envelope import GeneralConfigurationEnvelope from running_modes.configurations.logging.reinforcement_log_configuration import ReinforcementLoggerConfiguration -from running_modes.reinforcement_learning.inception import Inception from reinvent_scoring.scoring.enums.scoring_function_component_enum import ScoringFunctionComponentNameEnum from reinvent_scoring.scoring.score_summary import FinalSummary @@ -46,7 +46,7 @@ def save_checkpoint(self, step, scaffold_filter, agent): actual_step = step + 1 if self._log_config.logging_frequency > 0 and actual_step % self._log_config.logging_frequency == 0: self.save_diversity_memory(scaffold_filter) - agent.save(os.path.join(self._log_config.result_folder, f'Agent.{actual_step}.ckpt')) + agent.save_to_file(os.path.join(self._log_config.result_folder, f'Agent.{actual_step}.ckpt')) @abstractmethod def save_final_state(self, agent, scaffold_filter): diff --git a/running_modes/curriculum_learning/logging/local_curriculum_logger.py b/running_modes/curriculum_learning/logging/local_curriculum_logger.py index 67c143a..2ab421f 100644 --- a/running_modes/curriculum_learning/logging/local_curriculum_logger.py +++ b/running_modes/curriculum_learning/logging/local_curriculum_logger.py @@ -41,7 +41,7 @@ def timestep_report(self, start_time, n_steps, step, smiles: np.array, augmented_likelihood, diversity_filter) def save_final_state(self, agent, diversity_filter): - agent.save(os.path.join(self._log_config.result_folder, 'Agent.ckpt')) + agent.save_to_file(os.path.join(self._log_config.result_folder, 'Agent.ckpt')) self.save_diversity_memory(diversity_filter) self._summary_writer.close() diff --git a/running_modes/curriculum_learning/logging/remote_curriculum_logger.py b/running_modes/curriculum_learning/logging/remote_curriculum_logger.py index 0828529..9227e8c 100644 --- a/running_modes/curriculum_learning/logging/remote_curriculum_logger.py +++ b/running_modes/curriculum_learning/logging/remote_curriculum_logger.py @@ -31,18 +31,17 @@ def timestep_report(self, start_time, n_steps, step, smiles, mean_score: np.array, score_summary: FinalSummary, score, agent_likelihood: torch.tensor, prior_likelihood: torch.tensor, augmented_likelihood: torch.tensor, diversity_filter: BaseDiversityFilter): - score_components = self._score_summary_breakdown(score_summary, mean_score, diversity_filter) + score_components = self._score_summary_breakdown(score_summary, mean_score) learning_curves = self._learning_curve_profile(agent_likelihood, prior_likelihood, augmented_likelihood) - structures_table = self._visualize_structures(smiles, score, score_summary) smiles_report = self._create_sample_report(smiles, score, score_summary) time_estimation = running_modes.utils.general.estimate_run_time(start_time, n_steps, step) - data = self._assemble_timestep_report(step, score_components, structures_table, learning_curves, + data = self._assemble_timestep_report(step, score_components, diversity_filter, learning_curves, time_estimation, ul_rl.fraction_valid_smiles(smiles), smiles_report) self._notify_server(data, self._log_config.recipient) def save_final_state(self, agent, scaffold_filter): - agent.save(os.path.join(self._log_config.result_folder, 'Agent.ckpt')) + agent.save_to_file(os.path.join(self._log_config.result_folder, 'Agent.ckpt')) self.save_diversity_memory(scaffold_filter) def _notify_server(self, data, to_address): @@ -109,25 +108,23 @@ def _learning_curve_profile(self, agent_likelihood, prior_likelihood, augmented_ } return learning_curves - def _score_summary_breakdown(self, score_summary: FinalSummary, mean_score: np.array, - diversity_filter: BaseDiversityFilter): + def _score_summary_breakdown(self, score_summary: FinalSummary, mean_score: np.array): score_components = {} for i, log in enumerate(score_summary.profile): score_components[f"{score_summary.profile[i].component_type}:{score_summary.profile[i].name}"] = \ float(np.mean(score_summary.profile[i].score)) score_components["total_score:total_score"] = float(mean_score) - score_components["collected smiles in memory"] = diversity_filter.number_of_smiles_in_memory() return score_components - def _assemble_timestep_report(self, step, score_components, structures_table, learning_curves, time_estimation, - fraction_valid_smiles, smiles_report): + def _assemble_timestep_report(self, step, score_components, diversity_filter: BaseDiversityFilter, + learning_curves, time_estimation, fraction_valid_smiles, smiles_report): actual_step = step + 1 timestep_report = {"step": actual_step, "components": score_components, - # "structures": structures_table, "learning": learning_curves, "time_estimation": time_estimation, "fraction_valid_smiles": fraction_valid_smiles, - "smiles_report": smiles_report + "smiles_report": smiles_report, + "collected smiles in memory": diversity_filter.number_of_smiles_in_memory() } return timestep_report diff --git a/running_modes/enums/curriculum_strategy_enum.py b/running_modes/enums/curriculum_strategy_enum.py index 04c5427..bea827c 100644 --- a/running_modes/enums/curriculum_strategy_enum.py +++ b/running_modes/enums/curriculum_strategy_enum.py @@ -4,3 +4,5 @@ @dataclass(frozen=True) class CurriculumStrategyEnum: USER_DEFINED = "user_defined" + LINK_INVENT = "link_invent" + NO_CURRICULUM = "no_curriculum" diff --git a/running_modes/enums/merging_strategy_enum.py b/running_modes/enums/merging_strategy_enum.py new file mode 100644 index 0000000..4ad7382 --- /dev/null +++ b/running_modes/enums/merging_strategy_enum.py @@ -0,0 +1,9 @@ +from dataclasses import dataclass + + +@dataclass(frozen=True) +class MergingStrategyEnum: + LINEAR_MERGE = "linear" + LINEAR_MERGE_WITH_THRESHOLD = "linear_with_threshold" + LINEAR_MERGE_WITH_BULK = "linear_with_bulk" + LINEAR_MERGE_WITH_REMOVAL = "linear_with_removal" diff --git a/running_modes/enums/model_type_enum.py b/running_modes/enums/model_type_enum.py index 0494136..6266e81 100644 --- a/running_modes/enums/model_type_enum.py +++ b/running_modes/enums/model_type_enum.py @@ -5,4 +5,4 @@ class ModelTypeEnum: DEFAULT = "default" LIB_INVENT = "lib_invent" - LINK_INVENT = "link_invent" + LINK_INVENT = "link_invent" \ No newline at end of file diff --git a/running_modes/enums/production_strategy_enum.py b/running_modes/enums/production_strategy_enum.py index dba5c72..17dfc92 100644 --- a/running_modes/enums/production_strategy_enum.py +++ b/running_modes/enums/production_strategy_enum.py @@ -4,4 +4,5 @@ @dataclass(frozen=True) class ProductionStrategyEnum: STANDARD = "standard" - SPECIFIC_COMPONENTS = "specific_components" \ No newline at end of file + SPECIFIC_COMPONENTS = "specific_components" + LINK_INVENT = "link_invent" \ No newline at end of file diff --git a/running_modes/enums/ranking_strategy_enum.py b/running_modes/enums/ranking_strategy_enum.py new file mode 100644 index 0000000..ab68539 --- /dev/null +++ b/running_modes/enums/ranking_strategy_enum.py @@ -0,0 +1,8 @@ +from dataclasses import dataclass + + +@dataclass(frozen=True) +class RankingStrategyEnum: + SEQUENTIAL_FIXED_TIME = "sequential_fixed_time" + SEQUENTIAL_THRESHOLD = "sequential_threshold" + USER_DEFINED_ORDER = "user_defined_order" \ No newline at end of file diff --git a/running_modes/enums/scoring_table_enum.py b/running_modes/enums/scoring_table_enum.py new file mode 100644 index 0000000..f14ddb1 --- /dev/null +++ b/running_modes/enums/scoring_table_enum.py @@ -0,0 +1,10 @@ +from dataclasses import dataclass + + +@dataclass(frozen=True) +class ScoringTableEnum: + + AGENTS = "agents" + SCORES = "scores" + SCORING_FUNCTIONS = "scoring_functions" + COMPONENT_NAMES = "component_names" diff --git a/running_modes/manager.py b/running_modes/manager.py index dfededb..02523f4 100644 --- a/running_modes/manager.py +++ b/running_modes/manager.py @@ -1,29 +1,23 @@ -import json import os from running_modes.configurations import GeneralConfigurationEnvelope -from running_modes.enums.running_mode_enum import RunningModeEnum from running_modes.constructors.running_mode import RunningMode +from running_modes.enums.running_mode_enum import RunningModeEnum + class Manager: - def __init__(self, configuration): + def __init__(self, base_configuration, run_configuration): self.running_mode_enum = RunningModeEnum() - self.configuration = GeneralConfigurationEnvelope(**configuration) + self.base_configuration = base_configuration + self.run_configuration = GeneralConfigurationEnvelope(**run_configuration) self._load_environmental_variables() def run(self): - runner = RunningMode(self.configuration) + runner = RunningMode(self.run_configuration) runner.run() def _load_environmental_variables(self): - try: - project_root = os.path.dirname(__file__) - with open(os.path.join(project_root, '../configs/config.json'), 'r') as f: - config = json.load(f) - environmental_variables = config["ENVIRONMENTAL_VARIABLES"] - for key, value in environmental_variables.items(): - os.environ[key] = value - - except KeyError as ex: - raise ex + environmental_variables = self.base_configuration["ENVIRONMENTAL_VARIABLES"] + for key, value in environmental_variables.items(): + os.environ[key] = value diff --git a/running_modes/reinforcement_learning/learning_strategy/base_learning_strategy.py b/running_modes/reinforcement_learning/learning_strategy/base_learning_strategy.py index 2841787..32db329 100644 --- a/running_modes/reinforcement_learning/learning_strategy/base_learning_strategy.py +++ b/running_modes/reinforcement_learning/learning_strategy/base_learning_strategy.py @@ -35,13 +35,10 @@ def run(self, scaffold_batch: np.ndarray, decorator_batch: np.ndarray, def _calculate_loss(self, scaffold_batch, decorator_batch, score, actor_nlls): raise NotImplementedError("_calculate_loss method is not implemented") - # TODO: Don't use CUDA. - def _to_tensor(self, tensor): - if isinstance(tensor, np.ndarray): - tensor = torch.from_numpy(tensor) - if torch.cuda.is_available(): - return torch.autograd.Variable(tensor).cuda() - return torch.autograd.Variable(tensor) + def _to_tensor(self, array, use_cuda=True): + if torch.cuda.is_available() and use_cuda: + return torch.tensor(array, device=torch.device("cuda")) + return torch.tensor(array, device=torch.device("cpu")) def _disable_prior_gradients(self): # There might be a more elegant way of disabling gradients diff --git a/running_modes/reinforcement_learning/lib_invent_reinforcement_learning.py b/running_modes/reinforcement_learning/lib_invent_reinforcement_learning.py index a3520db..b1d7574 100644 --- a/running_modes/reinforcement_learning/lib_invent_reinforcement_learning.py +++ b/running_modes/reinforcement_learning/lib_invent_reinforcement_learning.py @@ -61,7 +61,7 @@ def _updating(self, sampled_sequences, score): def _logging(self, start_time, step, score_summary, actor_nlls, critic_nlls, augmented_nlls): self.logger.timestep_report(start_time, self.configuration.n_steps, step, score_summary, actor_nlls, - critic_nlls, augmented_nlls, self.scoring_strategy.diversity_filter) + critic_nlls, augmented_nlls, self.scoring_strategy.diversity_filter, self.actor) def _calculate_likelihood(self, sampled_sequences: List[SampledSequencesDTO]): nll_calculation_action = LikelihoodEvaluation(self.actor, self.configuration.batch_size, self.logger) diff --git a/running_modes/reinforcement_learning/link_invent_reinforcement_learning.py b/running_modes/reinforcement_learning/link_invent_reinforcement_learning.py index b6536ef..f454bf3 100644 --- a/running_modes/reinforcement_learning/link_invent_reinforcement_learning.py +++ b/running_modes/reinforcement_learning/link_invent_reinforcement_learning.py @@ -55,7 +55,7 @@ def _updating(self, sampled_sequences, score): def _logging(self, start_time, step, score_summary, actor_nlls, critic_nlls, augmented_nlls): self.logger.timestep_report(start_time, self.configuration.n_steps, step, score_summary, actor_nlls, - critic_nlls, augmented_nlls, self.scoring_strategy.diversity_filter) + critic_nlls, augmented_nlls, self.scoring_strategy.diversity_filter, self.actor) def _calculate_likelihood(self, sampled_sequences: List[SampledSequencesDTO]): nll_calculation_action = LinkInventLikelihoodEvaluation(self.actor, self.logger) diff --git a/running_modes/reinforcement_learning/logging/link_logging/base_reinforcement_logger.py b/running_modes/reinforcement_learning/logging/link_logging/base_reinforcement_logger.py index 9fe4263..94683c1 100644 --- a/running_modes/reinforcement_learning/logging/link_logging/base_reinforcement_logger.py +++ b/running_modes/reinforcement_learning/logging/link_logging/base_reinforcement_logger.py @@ -24,7 +24,7 @@ def log_message(self, message: str): @abstractmethod def timestep_report(self, start_time, n_steps, step, score_summary: FinalSummary, agent_likelihood: torch.tensor, prior_likelihood: torch.tensor, - augmented_likelihood: torch.tensor, diversity_filter): + augmented_likelihood: torch.tensor, diversity_filter, agent): raise NotImplementedError("timestep_report method is not implemented") def log_out_input_configuration(self): diff --git a/running_modes/reinforcement_learning/logging/link_logging/bond_link_reinforcement_logger.py b/running_modes/reinforcement_learning/logging/link_logging/bond_link_reinforcement_logger.py index 0ba243f..64f1861 100644 --- a/running_modes/reinforcement_learning/logging/link_logging/bond_link_reinforcement_logger.py +++ b/running_modes/reinforcement_learning/logging/link_logging/bond_link_reinforcement_logger.py @@ -3,6 +3,8 @@ from running_modes.reinforcement_learning.logging.link_logging.base_reinforcement_logger import BaseReinforcementLogger from running_modes.reinforcement_learning.logging.link_logging.local_bond_link_reinforcement_logger import \ LocalBondLinkReinforcementLogger +from running_modes.reinforcement_learning.logging.link_logging.remote_bond_link_reinforcement_logger import \ + RemoteLinkReinforcementLogger class BondLinkReinforcementLogger: @@ -14,6 +16,6 @@ def __new__(cls, configuration: GeneralConfigurationEnvelope, log_config: Reinfo if log_config.recipient == logging_mode_enum.LOCAL: logger_instance = LocalBondLinkReinforcementLogger(configuration, log_config) else: - raise NotImplemented("Remote logging mode is not implemented yet !") + logger_instance = RemoteLinkReinforcementLogger(configuration, log_config) return logger_instance \ No newline at end of file diff --git a/running_modes/reinforcement_learning/logging/link_logging/local_bond_link_reinforcement_logger.py b/running_modes/reinforcement_learning/logging/link_logging/local_bond_link_reinforcement_logger.py index 57f56b3..d538300 100644 --- a/running_modes/reinforcement_learning/logging/link_logging/local_bond_link_reinforcement_logger.py +++ b/running_modes/reinforcement_learning/logging/link_logging/local_bond_link_reinforcement_logger.py @@ -28,13 +28,14 @@ def log_message(self, message: str): def timestep_report(self, start_time, n_steps, step, score_summary: FinalSummary, agent_likelihood: torch.tensor, prior_likelihood: torch.tensor, - augmented_likelihood: torch.tensor, diversity_filter): + augmented_likelihood: torch.tensor, diversity_filter, actor): message = self._console_message_formatter.create(start_time, n_steps, step, score_summary, agent_likelihood, prior_likelihood, augmented_likelihood) self._logger.info(message) self._tensorboard_report(step, score_summary, agent_likelihood, prior_likelihood, augmented_likelihood, diversity_filter) + self.save_checkpoint(step, diversity_filter, actor) def _tensorboard_report(self, step, score_summary: FinalSummary, agent_likelihood, prior_likelihood, augmented_likelihood, diversity_filter: BaseDiversityFilter): diff --git a/running_modes/reinforcement_learning/logging/link_logging/remote_bond_link_reinforcement_logger.py b/running_modes/reinforcement_learning/logging/link_logging/remote_bond_link_reinforcement_logger.py new file mode 100644 index 0000000..09fd787 --- /dev/null +++ b/running_modes/reinforcement_learning/logging/link_logging/remote_bond_link_reinforcement_logger.py @@ -0,0 +1,136 @@ +import os +import numpy as np +import requests +import torch +from reinvent_scoring import ScoringFunctionComponentNameEnum, FinalSummary, ComponentSpecificParametersEnum +from reinvent_scoring.scoring.diversity_filters.lib_invent.base_diversity_filter import BaseDiversityFilter + +import running_modes.utils.configuration as ull +import reinvent_chemistry.logging as ul_rl +import running_modes.utils.general as ul_gen + +from running_modes.configurations import GeneralConfigurationEnvelope, ReinforcementLoggerConfiguration, \ + get_remote_logging_auth_token +from running_modes.reinforcement_learning.logging.link_logging.base_reinforcement_logger import BaseReinforcementLogger + + +class RemoteLinkReinforcementLogger(BaseReinforcementLogger): + def __init__(self, configuration: GeneralConfigurationEnvelope, rl_config: ReinforcementLoggerConfiguration): + super().__init__(configuration, rl_config) + self._rows = 2 + self._columns = 5 + self._sample_size = self._rows * self._columns + self._sf_component_enum = ScoringFunctionComponentNameEnum() + self._specific_parameters_enum = ComponentSpecificParametersEnum() + self._is_dev = ull._is_development_environment() + + def log_message(self, message: str): + self._logger.info(message) + + def timestep_report(self, start_time, n_steps, step, score_summary: FinalSummary, + agent_likelihood: torch.tensor, prior_likelihood: torch.tensor, + augmented_likelihood: torch.tensor, diversity_filter, actor): + + mean_score = np.mean(score_summary.total_score) + smiles = score_summary.scored_smiles + score = score_summary.total_score + + score_components = self._score_summary_breakdown(score_summary, mean_score) + learning_curves = self._learning_curve_profile(agent_likelihood, prior_likelihood, augmented_likelihood) + smiles_report = self._create_sample_report(smiles, score, score_summary) + + time_estimation = ul_gen.estimate_run_time(start_time, n_steps, step) + data = self._assemble_timestep_report(step, score_components, diversity_filter, learning_curves, + time_estimation, ul_rl.fraction_valid_smiles(smiles), smiles_report) + self._notify_server(data, self._log_config.recipient) + self.save_checkpoint(step, diversity_filter, actor) + + def save_final_state(self, agent, scaffold_filter): + agent.save_to_file(os.path.join(self._log_config.result_folder, 'Agent.ckpt')) + self.save_filter_memory(scaffold_filter) + self.log_out_input_configuration() + + def _notify_server(self, data, to_address): + """This is called every time we are posting data to server""" + try: + self._logger.warning(f"posting to {to_address}") + headers = { + 'Accept': 'application/json', 'Content-Type': 'application/json', + 'Authorization': get_remote_logging_auth_token() + } + response = requests.post(to_address, json=data, headers=headers) + + if self._is_dev: + """logs out the response content only when running a test instance""" + if response.status_code == requests.codes.ok: + self._logger.info(f"SUCCESS: {response.status_code}") + self._logger.info(response.content) + else: + self._logger.info(f"PROBLEM: {response.status_code}") + self._logger.exception(data, exc_info=False) + except Exception as t_ex: + self._logger.exception("Exception occurred", exc_info=True) + self._logger.exception(f"Attempted posting the following data:") + self._logger.exception(data, exc_info=False) + + def _get_matching_substructure_from_config(self, score_summary: FinalSummary): + smarts_pattern = "" + for summary_component in score_summary.scaffold_log: + if summary_component.parameters.component_type == self._sf_component_enum.MATCHING_SUBSTRUCTURE: + smarts = summary_component.parameters.specific_parameters.get(self._specific_parameters_enum.SMILES, []) + if len(smarts) > 0: + smarts_pattern = smarts[0] + return smarts_pattern + + def _visualize_structures(self, smiles, score, score_summary: FinalSummary): + score, smiles = ul_rl.sort_smiles_by_score(score, smiles) + smiles = ul_rl.padding_with_invalid_smiles(smiles, self._sample_size) + list_of_mols, legend = ul_rl.check_for_invalid_mols_and_create_legend(smiles, score, self._sample_size) + smarts_pattern = self._get_matching_substructure_from_config(score_summary) + pattern = ul_rl.find_matching_pattern_in_smiles(list_of_mols=list_of_mols, smarts_pattern=smarts_pattern) + mol_in_base64_string = ul_rl.mol_to_png_string(list_of_mols, molsPerRow=self._columns, subImgSize=(300, 300), + legend=legend, matches=pattern) + return mol_in_base64_string + + def _create_sample_report(self, smiles, score, score_summary: FinalSummary): + score, smiles = ul_rl.sort_smiles_by_score(score, smiles) + smiles = ul_rl.padding_with_invalid_smiles(smiles, self._sample_size) + _, legend = ul_rl.check_for_invalid_mols_and_create_legend(smiles, score, self._sample_size) + smarts_pattern = self._get_matching_substructure_from_config(score_summary) + + smiles_legend_pairs = [{"smiles": smiles[indx], "legend": legend[indx]} for indx in range(self._sample_size)] + + report = { + "smarts_pattern": smarts_pattern, + "smiles_legend_pairs": smiles_legend_pairs + } + return report + + def _learning_curve_profile(self, agent_likelihood, prior_likelihood, augmented_likelihood): + learning_curves = { + "prior": float(np.float(prior_likelihood.detach().mean().cpu())), + "augmented": float(np.float(augmented_likelihood.detach().mean().cpu())), + "agent": float(np.float(agent_likelihood.detach().mean().cpu())) + } + return learning_curves + + def _score_summary_breakdown(self, score_summary: FinalSummary, mean_score: np.array): + score_components = {} + for i, log in enumerate(score_summary.profile): + score_components[f"{score_summary.profile[i].component_type}:{score_summary.profile[i].name}"] = \ + float(np.mean(score_summary.profile[i].score)) + score_components["total_score:total_score"] = float(mean_score) + return score_components + + def _assemble_timestep_report(self, step, score_components, diversity_filter: BaseDiversityFilter, + learning_curves, time_estimation, fraction_valid_smiles, smiles_report): + actual_step = step + 1 + timestep_report = {"step": actual_step, + "components": score_components, + "learning": learning_curves, + "time_estimation": time_estimation, + "fraction_valid_smiles": fraction_valid_smiles, + "smiles_report": smiles_report, + "collected smiles in memory": diversity_filter.number_of_smiles_in_memory() + } + return timestep_report diff --git a/running_modes/reinforcement_learning/logging/reinforcement_logger.py b/running_modes/reinforcement_learning/logging/reinforcement_logger.py index 301e4d7..720d561 100644 --- a/running_modes/reinforcement_learning/logging/reinforcement_logger.py +++ b/running_modes/reinforcement_learning/logging/reinforcement_logger.py @@ -2,11 +2,11 @@ from running_modes.configurations.logging.reinforcement_log_configuration import ReinforcementLoggerConfiguration from running_modes.enums.logging_mode_enum import LoggingModeEnum from running_modes.enums.model_type_enum import ModelTypeEnum +from running_modes.reinforcement_learning.logging.base_reinforcement_logger import BaseReinforcementLogger from running_modes.reinforcement_learning.logging.link_logging.bond_link_reinforcement_logger import \ BondLinkReinforcementLogger -from running_modes.reinforcement_learning.logging.remote_reinforcement_logger import RemoteReinforcementLogger from running_modes.reinforcement_learning.logging.local_reinforcement_logger import LocalReinforcementLogger -from running_modes.reinforcement_learning.logging.base_reinforcement_logger import BaseReinforcementLogger +from running_modes.reinforcement_learning.logging.remote_reinforcement_logger import RemoteReinforcementLogger class ReinforcementLogger: @@ -23,4 +23,5 @@ def __new__(cls, configuration: GeneralConfigurationEnvelope, log_config: Reinfo logger_instance = RemoteReinforcementLogger(configuration, log_config) else: logger_instance = BondLinkReinforcementLogger(configuration, log_config) + return logger_instance diff --git a/running_modes/reinforcement_learning/logging/remote_reinforcement_logger.py b/running_modes/reinforcement_learning/logging/remote_reinforcement_logger.py index f10d26a..6ad05e3 100644 --- a/running_modes/reinforcement_learning/logging/remote_reinforcement_logger.py +++ b/running_modes/reinforcement_learning/logging/remote_reinforcement_logger.py @@ -32,13 +32,12 @@ def timestep_report(self, start_time, n_steps, step, smiles, mean_score: np.array, score_summary: FinalSummary, score, agent_likelihood: torch.tensor, prior_likelihood: torch.tensor, augmented_likelihood: torch.tensor, diversity_filter: BaseDiversityFilter): - score_components = self._score_summary_breakdown(score_summary, mean_score, diversity_filter) + score_components = self._score_summary_breakdown(score_summary, mean_score) learning_curves = self._learning_curve_profile(agent_likelihood, prior_likelihood, augmented_likelihood) - structures_table = self._visualize_structures(smiles, score, score_summary) smiles_report = self._create_sample_report(smiles, score, score_summary) time_estimation = running_modes.utils.general.estimate_run_time(start_time, n_steps, step) - data = self._assemble_timestep_report(step, score_components, structures_table, learning_curves, + data = self._assemble_timestep_report(step, score_components, diversity_filter, learning_curves, time_estimation, ul_rl.fraction_valid_smiles(smiles), smiles_report) self._notify_server(data, self._log_config.recipient) @@ -111,25 +110,23 @@ def _learning_curve_profile(self, agent_likelihood, prior_likelihood, augmented_ } return learning_curves - def _score_summary_breakdown(self, score_summary: FinalSummary, mean_score: np.array, - diversity_filter: BaseDiversityFilter): + def _score_summary_breakdown(self, score_summary: FinalSummary, mean_score: np.array): score_components = {} for i, log in enumerate(score_summary.profile): score_components[f"{score_summary.profile[i].component_type}:{score_summary.profile[i].name}"] = \ float(np.mean(score_summary.profile[i].score)) score_components["total_score:total_score"] = float(mean_score) - score_components["collected smiles in memory"] = diversity_filter.number_of_smiles_in_memory() return score_components - def _assemble_timestep_report(self, step, score_components, structures_table, learning_curves, time_estimation, - fraction_valid_smiles, smiles_report): + def _assemble_timestep_report(self, step, score_components, diversity_filter:BaseDiversityFilter, + learning_curves, time_estimation, fraction_valid_smiles, smiles_report): actual_step = step + 1 timestep_report = {"step": actual_step, "components": score_components, - # "structures": structures_table, "learning": learning_curves, "time_estimation": time_estimation, "fraction_valid_smiles": fraction_valid_smiles, - "smiles_report": smiles_report + "smiles_report": smiles_report, + "collected smiles in memory": diversity_filter.number_of_smiles_in_memory() } return timestep_report diff --git a/running_modes/sampling/logging/remote_sampling_logger.py b/running_modes/sampling/logging/remote_sampling_logger.py index 077d605..5b1d3d7 100644 --- a/running_modes/sampling/logging/remote_sampling_logger.py +++ b/running_modes/sampling/logging/remote_sampling_logger.py @@ -49,13 +49,10 @@ def _notify_server(self, data, to_address): self._logger.exception(data, exc_info=False) def _create_sample_report(self, smiles): - legend, list_of_mols = self._count_unique_inchi_keys(smiles) + legends, list_of_mols = self._count_unique_inchi_keys(smiles) list_of_smiles = [self._conversions.mol_to_smiles(mol) if mol is not None else "INVALID" for mol in list_of_mols] + report = [{"smiles": smiles, "legend": legend} for smiles, legend in zip(list_of_smiles, legends)] - report = { - "smiles": list_of_smiles, - "legend": legend - } return report def _assemble_timestep_report(self, fraction_valid_smiles, fraction_unique_smiles, smiles_report): diff --git a/running_modes/transfer_learning/logging/remote_transfer_learning_logger.py b/running_modes/transfer_learning/logging/remote_transfer_learning_logger.py index e82e661..647fcf6 100644 --- a/running_modes/transfer_learning/logging/remote_transfer_learning_logger.py +++ b/running_modes/transfer_learning/logging/remote_transfer_learning_logger.py @@ -56,13 +56,10 @@ def log_timestep(self, lr, epoch, sampled_smiles, sampled_nlls, self._notify_server(data, self._log_config.recipient) def _create_sample_report(self, smiles): - legend, list_of_mols = self._count_compound_frequency(smiles) + legends, list_of_mols = self._count_compound_frequency(smiles) list_of_smiles = [self._conversions.mol_to_smiles(mol) if mol is not None else "INVALID" for mol in list_of_mols] + report = [{"smiles": smiles, "legend": legend} for smiles, legend in zip(list_of_smiles, legends)] - report = { - "smiles": list_of_smiles, - "legend": legend - } return report def _mean_learning_curve_profile(self, sampled_nlls: np.array, training_nlls: np.array): diff --git a/running_modes/utils/general.py b/running_modes/utils/general.py index f7302e0..44049f4 100644 --- a/running_modes/utils/general.py +++ b/running_modes/utils/general.py @@ -12,9 +12,9 @@ def to_tensor(tensor): return torch.autograd.Variable(tensor) -def set_default_device_cuda(): +def set_default_device_cuda(dont_use_cuda=False): """Sets the default device (cpu or cuda) used for all tensors.""" - if torch.cuda.is_available() == False: + if torch.cuda.is_available() == False or dont_use_cuda: tensor = torch.FloatTensor torch.set_default_tensor_type(tensor) return False diff --git a/unittest_reinvent/fixtures/predictive_model_fixtures.py b/unittest_reinvent/fixtures/predictive_model_fixtures.py index bc59abe..f54e42f 100644 --- a/unittest_reinvent/fixtures/predictive_model_fixtures.py +++ b/unittest_reinvent/fixtures/predictive_model_fixtures.py @@ -1,11 +1,11 @@ -from unittest_reinvent.fixtures.paths import ACTIVITY_REGRESSION, ACTIVITY_CLASSIFICATION, MAIN_TEST_PATH - - +from reinvent_scoring import TransformationParametersEnum +from reinvent_scoring.scoring.component_parameters import ComponentParameters +from reinvent_scoring.scoring.enums.component_specific_parameters_enum import ComponentSpecificParametersEnum +from reinvent_scoring.scoring.enums.descriptor_types_enum import DescriptorTypesEnum from reinvent_scoring.scoring.enums.scoring_function_component_enum import ScoringFunctionComponentNameEnum from reinvent_scoring.scoring.enums.transformation_type_enum import TransformationTypeEnum -from reinvent_scoring.scoring.enums.descriptor_types_enum import DescriptorTypesEnum -from reinvent_scoring.scoring.enums.component_specific_parameters_enum import ComponentSpecificParametersEnum -from reinvent_scoring.scoring.component_parameters import ComponentParameters + +from unittest_reinvent.fixtures.paths import ACTIVITY_REGRESSION def create_activity_component_regression(): @@ -18,19 +18,6 @@ def create_activity_component_regression(): return parameters -def create_offtarget_activity_component_regression(): - csp_enum = ComponentSpecificParametersEnum() - sf_enum = ScoringFunctionComponentNameEnum() - specific_parameters = _specific_parameters_regression_predictive_model(ACTIVITY_REGRESSION) - specific_parameters[csp_enum.HIGH] = 3 - specific_parameters[csp_enum.LOW] = 0 - specific_parameters[csp_enum.TRANSFORMATION] = False - parameters = ComponentParameters(component_type=sf_enum.PREDICTIVE_PROPERTY, - name="offtarget_activity", - weight=1., - specific_parameters=specific_parameters) - return parameters - def create_predictive_property_component_regression(): sf_enum = ScoringFunctionComponentNameEnum() specific_parameters = _specific_parameters_regression_predictive_model(ACTIVITY_REGRESSION) @@ -40,108 +27,17 @@ def create_predictive_property_component_regression(): specific_parameters=specific_parameters) return parameters -def create_activity_component_classification(): - sf_enum = ScoringFunctionComponentNameEnum() - specific_parameters = _specific_parameters_classifiaction_predictive_model(ACTIVITY_CLASSIFICATION) - parameters = ComponentParameters(component_type=sf_enum.PREDICTIVE_PROPERTY, - name="activity_classification", - weight=1., - specific_parameters=specific_parameters) - return parameters - -def create_offtarget_activity_component_classification(): - csp_enum = ComponentSpecificParametersEnum() - sf_enum = ScoringFunctionComponentNameEnum() - specific_parameters = _specific_parameters_classifiaction_predictive_model(ACTIVITY_CLASSIFICATION) - specific_parameters[csp_enum.HIGH] = 3 - specific_parameters[csp_enum.LOW] = 0 - specific_parameters[csp_enum.TRANSFORMATION] = False - parameters = ComponentParameters(component_type=sf_enum.PREDICTIVE_PROPERTY, - name="predictive_property_classification", - weight=1., - specific_parameters=specific_parameters) - return parameters - -def create_predictive_property_component_classification(): - sf_enum = ScoringFunctionComponentNameEnum() - specific_parameters = _specific_parameters_classifiaction_predictive_model(ACTIVITY_CLASSIFICATION) - parameters = ComponentParameters(component_type=sf_enum.PREDICTIVE_PROPERTY, - name="predictive_property_classification", - weight=1., - specific_parameters=specific_parameters) - return parameters - -def create_c_lab_component(somponent_type): - csp_enum = ComponentSpecificParametersEnum() - transf_type = TransformationTypeEnum() - specific_parameters = {csp_enum.CLAB_INPUT_FILE: f"{MAIN_TEST_PATH}/clab_input.json", - csp_enum.HIGH: 9, - csp_enum.LOW: 4, - csp_enum.K: 0.25, - csp_enum.TRANSFORMATION: True, - csp_enum.TRANSFORMATION_TYPE: transf_type.SIGMOID} - parameters = ComponentParameters(component_type=somponent_type, - name="c_lab", - weight=1., - specific_parameters=specific_parameters) - return parameters - - def _specific_parameters_regression_predictive_model(path=None): csp_enum = ComponentSpecificParametersEnum() transf_type = TransformationTypeEnum() descriptor_types = DescriptorTypesEnum() - specific_parameters = {csp_enum.HIGH: 9, - csp_enum.LOW: 4, - csp_enum.K: 0.25, - csp_enum.TRANSFORMATION: True, - csp_enum.TRANSFORMATION_TYPE: transf_type.SIGMOID, + transform_params = TransformationParametersEnum + specific_parameters = {csp_enum.TRANSFORMATION: {transform_params.HIGH: 9, + transform_params.LOW: 4, + transform_params.K: 0.25, + transform_params.TRANSFORMATION_TYPE: transf_type.SIGMOID}, csp_enum.SCIKIT: "regression", "model_path": path, - csp_enum.DESCRIPTOR_TYPE: descriptor_types.ECFP_COUNTS} - return specific_parameters - -def _specific_parameters_classifiaction_predictive_model(path=None): - csp_enum = ComponentSpecificParametersEnum() - descriptor_types = DescriptorTypesEnum() - specific_parameters = {csp_enum.HIGH: 9, - csp_enum.LOW: 4, - csp_enum.K: 0.25, - csp_enum.TRANSFORMATION: False, - csp_enum.SCIKIT: "classification", - "model_path": path, - csp_enum.DESCRIPTOR_TYPE: descriptor_types.ECFP_COUNTS} + csp_enum.DESCRIPTOR_TYPE: descriptor_types.ECFP_COUNTS + } return specific_parameters - -def create_custom_alerts_configuration(): - custom_alerts_list = [ - '[*;r7]', - '[*;r8]', - '[*;r9]', - '[*;r10]', - '[*;r11]', - '[*;r12]', - '[*;r13]', - '[*;r14]', - '[*;r15]', - '[*;r16]', - '[*;r17]', - '[#8][#8]', - '[#6;+]', - '[#16][#16]', - '[#7;!n][S;!$(S(=O)=O)]', - '[#7;!n][#7;!n]', - 'C#C', - 'C(=[O,S])[O,S]', - '[#7;!n][C;!$(C(=[O,N])[N,O])][#16;!s]', - '[#7;!n][C;!$(C(=[O,N])[N,O])][#7;!n]', - '[#7;!n][C;!$(C(=[O,N])[N,O])][#8;!o]', - '[#8;!o][C;!$(C(=[O,N])[N,O])][#16;!s]', - '[#8;!o][C;!$(C(=[O,N])[N,O])][#8;!o]', - '[#16;!s][C;!$(C(=[O,N])[N,O])][#16;!s]'] - sf_enum = ScoringFunctionComponentNameEnum() - parameters = ComponentParameters(component_type=sf_enum.CUSTOM_ALERTS, - name="custom_alerts", - weight=1., - specific_parameters={"smiles":custom_alerts_list}) - return parameters diff --git a/unittest_reinvent/running_modes/curriculum_tests/__init__.py b/unittest_reinvent/running_modes/curriculum_tests/__init__.py index 08f8583..0109178 100644 --- a/unittest_reinvent/running_modes/curriculum_tests/__init__.py +++ b/unittest_reinvent/running_modes/curriculum_tests/__init__.py @@ -2,4 +2,4 @@ TestManualCurriculumLearning from unittest_reinvent.running_modes.curriculum_tests.test_automated_curriculum_learning import \ - TestAutomatedCurriculumLearning + TestAutomatedCurriculumLearning \ No newline at end of file diff --git a/unittest_reinvent/running_modes/curriculum_tests/test_automated_curriculum_learning.py b/unittest_reinvent/running_modes/curriculum_tests/test_automated_curriculum_learning.py index 9dd85e6..9d2da54 100644 --- a/unittest_reinvent/running_modes/curriculum_tests/test_automated_curriculum_learning.py +++ b/unittest_reinvent/running_modes/curriculum_tests/test_automated_curriculum_learning.py @@ -2,30 +2,31 @@ import shutil import unittest -from running_modes.automated_curriculum_learning.automated_curriculum_runner import AutomatedCurriculumRunner -from reinvent_models.reinvent_core.models.model import Model -from running_modes.automated_curriculum_learning.logging import AutoCLLogger -from running_modes.configurations.automated_curriculum_learning.automated_curriculum_learning_configuration import \ - AutomatedCLConfiguration -from running_modes.configurations.automated_curriculum_learning.curriculum_strategy_configuration import CurriculumStrategyConfiguration -from running_modes.configurations.automated_curriculum_learning.curriculum_objective import CurriculumObjective -from running_modes.configurations.automated_curriculum_learning.production_strategy_configuration import ProductionStrategyConfiguration - +from reinvent_models.lib_invent.enums.generative_model_regime import GenerativeModelRegimeEnum +from reinvent_models.model_factory.configurations.model_configuration import ModelConfiguration +from reinvent_models.model_factory.generative_model import GenerativeModel from reinvent_scoring.scoring.component_parameters import ComponentParameters -from reinvent_scoring.scoring.diversity_filters.reinvent_core.diversity_filter_parameters import \ - DiversityFilterParameters +from reinvent_scoring.scoring.diversity_filters.curriculum_learning import DiversityFilterParameters from reinvent_scoring.scoring.enums.diversity_filter_enum import DiversityFilterEnum from reinvent_scoring.scoring.enums.scoring_function_component_enum import ScoringFunctionComponentNameEnum from reinvent_scoring.scoring.enums.scoring_function_enum import ScoringFunctionNameEnum from reinvent_scoring.scoring.scoring_function_parameters import ScoringFunctionParameters +from running_modes.automated_curriculum_learning.automated_curriculum_runner import AutomatedCurriculumRunner +from running_modes.automated_curriculum_learning.learning_strategy.learning_strategy_configuration import \ + LearningStrategyConfiguration +from running_modes.automated_curriculum_learning.logging import AutoCLLogger +from running_modes.configurations import AutomatedCurriculumLearningInputConfiguration, \ + CurriculumStrategyInputConfiguration, ProductionStrategyInputConfiguration +from running_modes.configurations.automated_curriculum_learning.curriculum_objective import CurriculumObjective from running_modes.configurations.general_configuration_envelope import GeneralConfigurationEnvelope from running_modes.configurations.logging.reinforcement_log_configuration import ReinforcementLoggerConfiguration from running_modes.configurations.reinforcement_learning.inception_configuration import InceptionConfiguration -from running_modes.enums.curriculum_type_enum import CurriculumTypeEnum from running_modes.enums.curriculum_strategy_enum import CurriculumStrategyEnum -from running_modes.enums.production_strategy_enum import ProductionStrategyEnum +from running_modes.enums.curriculum_type_enum import CurriculumTypeEnum from running_modes.enums.logging_mode_enum import LoggingModeEnum +from running_modes.enums.model_type_enum import ModelTypeEnum +from running_modes.enums.production_strategy_enum import ProductionStrategyEnum from running_modes.enums.running_mode_enum import RunningModeEnum from running_modes.utils import set_default_device_cuda from unittest_reinvent.fixtures.paths import MAIN_TEST_PATH, PRIOR_PATH @@ -46,12 +47,17 @@ def setUp(self): self.workfolder = MAIN_TEST_PATH self.logging_path = f"{self.workfolder}/log" smiles = [ASPIRIN, CELECOXIB] + model_type = ModelTypeEnum() + model_regime = GenerativeModelRegimeEnum() automated_cl_parameters, parameters = self._create_configuration(smiles) + prior_config = ModelConfiguration(model_type.DEFAULT, model_regime.INFERENCE, PRIOR_PATH) + agent_config = ModelConfiguration(model_type.DEFAULT, model_regime.TRAINING, PRIOR_PATH) + prior = GenerativeModel(prior_config) + agent = GenerativeModel(agent_config) - self.runner = AutomatedCurriculumRunner(automated_cl_parameters, AutoCLLogger(parameters), - prior=Model.load_from_file(PRIOR_PATH), - agent=Model.load_from_file(PRIOR_PATH)) + self.runner = AutomatedCurriculumRunner(automated_cl_parameters, AutoCLLogger(parameters), prior=prior, + agent=agent) def tearDown(self): if os.path.isdir(self.workfolder): @@ -68,11 +74,13 @@ def _create_configuration(self, smiles): curriculum_df = DiversityFilterParameters(self.filter_enum.NO_FILTER, 0.05, 25, 0.4) curriculum_inception = InceptionConfiguration(smiles, 100, 10) + learning_strategy = LearningStrategyConfiguration(parameters={"sigma": 120}) - curriculum_config = CurriculumStrategyConfiguration(name=self.cs_enum.USER_DEFINED, + curriculum_config = CurriculumStrategyInputConfiguration(name=self.cs_enum.USER_DEFINED, curriculum_objectives=curriculum_objectives, diversity_filter=curriculum_df, - inception=curriculum_inception, max_num_iterations=3) + inception=curriculum_inception, max_num_iterations=3, + learning_strategy=learning_strategy) # Production Phase Configuration production_sf_parameters = vars(ComponentParameters(name="tanimoto similarity", weight=1, @@ -82,15 +90,17 @@ def _create_configuration(self, smiles): production_df = DiversityFilterParameters(self.filter_enum.IDENTICAL_MURCKO_SCAFFOLD, 0.05, 25, 0.4) production_inception = InceptionConfiguration(smiles, 100, 10) - production_config = ProductionStrategyConfiguration(name=self.ps_enum.STANDARD, + production_config = ProductionStrategyInputConfiguration(name=self.ps_enum.STANDARD, scoring_function= ScoringFunctionParameters( name=self.sf_enum.CUSTOM_PRODUCT, parameters=[production_sf_parameters]), - diversity_filter=production_df, inception=production_inception, - retain_inception=False, n_steps=3) + diversity_filter=production_df, + inception=production_inception, + retain_inception=False, number_of_steps=3, + learning_strategy=learning_strategy) - automated_cl_parameters = AutomatedCLConfiguration(prior=PRIOR_PATH, agent=PRIOR_PATH, + automated_cl_parameters = AutomatedCurriculumLearningInputConfiguration(prior=PRIOR_PATH, agent=PRIOR_PATH, curriculum_strategy=curriculum_config, production_strategy=production_config, curriculum_type=CurriculumTypeEnum.AUTOMATED) diff --git a/unittest_reinvent/running_modes/lib_invent_tests/learning_strategy_tests/base_learning_strategy.py b/unittest_reinvent/running_modes/lib_invent_tests/learning_strategy_tests/base_learning_strategy.py index c6ff24e..b961769 100644 --- a/unittest_reinvent/running_modes/lib_invent_tests/learning_strategy_tests/base_learning_strategy.py +++ b/unittest_reinvent/running_modes/lib_invent_tests/learning_strategy_tests/base_learning_strategy.py @@ -23,13 +23,10 @@ DECORATION_SUZUKI -def dummy_func(a, b) -> float: - return 0.3 - - class BaseTestLearningStrategy(unittest.TestCase): - def setUp(self): + def arrange(self, learning_strategy): + self.learning_strategy = learning_strategy self.sf_component_enum = ScoringFunctionComponentNameEnum() self.sf_enum = ScoringFunctionNameEnum() self.rf_enum = ReactionFiltersEnum() @@ -51,7 +48,7 @@ def setUp(self): critic_model = Mock() critic_model.set_mode = lambda _: None critic_model.network.parameters = lambda: [] - critic_model.likelihood = dummy_func + critic_model.likelihood = self.dummy_func reaction_filter_config = ReactionFilterConfiguration(type=self.rf_enum.SELECTIVE, reactions=reactions) @@ -73,14 +70,17 @@ def setUp(self): n_steps=2, batch_size=4) self.runner = LearningStrategy(critic_model, optimizer, config.learning_strategy) - self.runner._to_tensor = self._to_tensor - + # self.runner._to_tensor = self._to_tensor + # self.scaffold_batch = np.array([CELECOXIB_SCAFFOLD]) self.decorator_batch = np.array([DECORATION_SUZUKI]) - self.score = torch.tensor([0.9], device=torch.device("cpu")) - self.actor_nlls = torch.tensor([0.2], device=torch.device("cpu"), requires_grad=True) + self.score = np.array([0.9]) # torch.tensor([0.9], device=torch.device("cpu")) + self.actor_nlls = torch.tensor([0.2], device=torch.device("cuda"), requires_grad=True) @staticmethod + def dummy_func(a, b) -> float: + return 0.3 + @staticmethod def _to_tensor(tensor): if isinstance(tensor, np.ndarray): tensor = torch.from_numpy(tensor) diff --git a/unittest_reinvent/running_modes/lib_invent_tests/learning_strategy_tests/test_learning_strategy_mascof_strategy.py b/unittest_reinvent/running_modes/lib_invent_tests/learning_strategy_tests/test_learning_strategy_mascof_strategy.py index 7c26f04..fe3c08d 100644 --- a/unittest_reinvent/running_modes/lib_invent_tests/learning_strategy_tests/test_learning_strategy_mascof_strategy.py +++ b/unittest_reinvent/running_modes/lib_invent_tests/learning_strategy_tests/test_learning_strategy_mascof_strategy.py @@ -8,8 +8,7 @@ class TestLearningStrategyMascofStrategy(BaseTestLearningStrategy): def setUp(self): - self.learning_strategy = LearningStrategyEnum().MASCOF - super().setUp() + super().arrange(LearningStrategyEnum().MASCOF) def test_mascof_strategy(self): actor_nlls, critic_nlls, augmented_nlls = \ diff --git a/unittest_reinvent/running_modes/lib_invent_tests/learning_strategy_tests/test_learning_strategy_mauli_strategy.py b/unittest_reinvent/running_modes/lib_invent_tests/learning_strategy_tests/test_learning_strategy_mauli_strategy.py index b35a70d..725a3b3 100644 --- a/unittest_reinvent/running_modes/lib_invent_tests/learning_strategy_tests/test_learning_strategy_mauli_strategy.py +++ b/unittest_reinvent/running_modes/lib_invent_tests/learning_strategy_tests/test_learning_strategy_mauli_strategy.py @@ -8,8 +8,7 @@ class TestLearningStrategyMauliStrategy(BaseTestLearningStrategy): def setUp(self): - self.learning_strategy = LearningStrategyEnum().MAULI - super().setUp() + super().arrange(LearningStrategyEnum().MAULI) def test_mauli_strategy(self): actor_nlls, critic_nlls, augmented_nlls = \ diff --git a/unittest_reinvent/running_modes/lib_invent_tests/learning_strategy_tests/test_learning_strategy_sdap_strategy.py b/unittest_reinvent/running_modes/lib_invent_tests/learning_strategy_tests/test_learning_strategy_sdap_strategy.py index 7025dee..68158ee 100644 --- a/unittest_reinvent/running_modes/lib_invent_tests/learning_strategy_tests/test_learning_strategy_sdap_strategy.py +++ b/unittest_reinvent/running_modes/lib_invent_tests/learning_strategy_tests/test_learning_strategy_sdap_strategy.py @@ -8,8 +8,7 @@ class TestLearningStrategySdapStrategy(BaseTestLearningStrategy): def setUp(self): - self.learning_strategy = LearningStrategyEnum().SDAP - super().setUp() + super().arrange(LearningStrategyEnum().SDAP) def test_sdap_strategy(self): actor_nlls, critic_nlls, augmented_nlls = \