Adding iSIM as a trackable value during training

configs/toml/example.toml

@@ -0,0 +1,55 @@
+run_type = "staged_learning"
+device = "cuda:0"
+tb_logdir = "tb_RL"
+json_out_config = "_RL.json"
+
+[parameters]
+
+use_checkpoint = false
+
+prior_file = "reinvent.prior"  # change as needed
+agent_file = "reinvent.prior"  # same here
+summary_csv_prefix = "RL"
+
+batch_size = 100
+randomize_smiles = true
+
+
+[learning_strategy]
+
+type = "dap"
+sigma = 128
+rate = 0.0001
+
+
+[diversity_filter]
+
+type = "IdenticalMurckoScaffold"
+bucket_size = 10
+minscore = 0.7
+minsimilarity = 0.5
+
+
+[inception]
+
+memory_size = 50
+sample_size = 10
+
+
+[[stage]]
+termination = "simple"
+max_score = 1.0
+max_steps = 200
+
+chkpt_file = 'RL.chkpt'
+
+[stage.scoring]
+type = "geometric_mean"
+
+
+[[stage.scoring.component]]
+[stage.scoring.component.QED]
+
+[[stage.scoring.component.QED.endpoint]]
+name = "QED"
+weight = 1.0

reinvent/runmodes/RL/learning.py

@@ -17,6 +17,10 @@
 from torch.utils.tensorboard import SummaryWriter
 import numpy as np
 
+#ISIM imports
+from iSIM.comp import calculate_isim
+from iSIM.utils import binary_fps
+
 from .reports import RLTBReporter, RLCSVReporter, RLRemoteReporter, RLReportData
 from reinvent.runmodes.RL.data_classes import ModelState
 from reinvent.models.model_factory.sample_batch import SmilesState
@@ -52,6 +56,7 @@ def __init__(
         inception: Inception = None,
         responder_config: dict = None,
         tb_logdir: str = None,
+        tb_isim: bool = False,
     ):
         """Setup of the common framework"""
 
@@ -94,13 +99,14 @@ def __init__(
         self.reporters = []
         self.tb_reporter = None
         self._setup_reporters(tb_logdir)
+        self.tb_isim = tb_isim
 
         self.start_time = 0
 
     def optimize(self, converged: terminator_callable) -> bool:
         """Run the multistep optimization loop
 
-        Sample from the agent, score the SNILES, update the agent parameters.
+        Sample from the agent, score the SMILES, update the agent parameters.
         Log some key characteristics of the current step.
 
         :param converged: a callable that determines convergence
@@ -310,6 +316,12 @@ def report(
         fract_duplicate_smiles = num_duplicate_smiles / len(mask_duplicates)
 
         smilies = np.array(self.sampled.smilies)[mask_valid]
+
+        isim = None
+        if self.tb_isim:
+            fingerprints = binary_fps(smilies, fp_type='RDKIT', n_bits=None) #Use isim utilities to compute RDKIT binary fingerprints
+            isim = calculate_isim(fingerprints, n_ary ='JT') #Use isim calculator for average Tanimoto similarity
+
         if self.prior.model_type == "Libinvent":
             smilies = normalize(smilies, keep_all=True)
         mask_idx = (np.argwhere(mask_valid).flatten(),)
@@ -318,6 +330,7 @@ def report(
             step=step_no,
             stage=self.stage_no,
             smilies=smilies,
+            isim=isim, #Add isim to report_data
             scaffolds=scaffolds,
             sampled=self.sampled,
             score_results=score_results,

reinvent/runmodes/RL/reports/data.py

@@ -17,6 +17,7 @@ class RLReportData:
     step: int
     stage: int
     smilies: list
+    isim: Optional[float]
     scaffolds: list
     sampled: SampleBatch
     score_results: ScoreResults

reinvent/runmodes/RL/reports/tensorboard.py

@@ -56,6 +56,9 @@ def submit(self, data: RLReportData) -> None:
             self.reporter.add_scalar(f"{name} (raw)", np.nanmean(_scores[mask_idx]), step)
 
         self.reporter.add_scalar(f"Loss", data.loss, step)
+        #Add iSIM to board as scalar per step 
+        if data.isim:
+            self.reporter.add_scalar(f"iSIM: Average similarity", data.isim, step)
 
         # NOTE: for some reason this breaks on Windows because the necessary
         #       subdirectory cannot be created

reinvent/runmodes/RL/run_staged_learning.py

@@ -219,7 +219,7 @@ def run_staged_learning(
     )
 
     parameters = config.parameters
-
+    
     # NOTE: The model files are a dictionary with model attributes from
     #       Reinvent and a set of tensors, each with an attribute for the
     #       device (CPU or GPU) and if gradients are required
@@ -301,7 +301,7 @@ def run_staged_learning(
     distance_threshold = parameters.distance_threshold
 
     model_learning = getattr(RL, f"{model_type}Learning")
-
+    
     if callable(write_config):
         write_config(config.model_dump())
 
@@ -329,7 +329,7 @@ def run_staged_learning(
             else:
                 state = ModelState(agent, package.diversity_filter)
                 logger.debug(f"Using stage DF")
-
+            
             optimize = model_learning(
                 max_steps=package.max_steps,
                 stage_no=stage_no,
@@ -344,6 +344,7 @@ def run_staged_learning(
                 inception=inception,
                 responder_config=responder_config,
                 tb_logdir=logdir,
+                tb_isim=parameters.tb_isim,
             )
 
             if device.type == "cuda" and torch.cuda.is_available():

reinvent/runmodes/RL/validation.py

@@ -20,6 +20,7 @@ class SectionParameters(GlobalConfig):
     randomize_smiles: bool = True
     unique_sequences: bool = False
     temperature: float = 1.0
+    tb_isim: Optional[bool] = False  #Add iSIM tracking as optional parameter
 
 
 class SectionLearningStrategy(GlobalConfig):

reinvent/utils/config_parse.py

@@ -198,14 +198,12 @@ def read_config(filename: Optional[Path], fmt: str) -> dict:
     """
 
     pkg = FMT_CONVERT[fmt]
-
     if isinstance(filename, (str, Path)):
         with open(filename, "rb") as tf:
             config = pkg.load(tf)
     else:
         config_str = "\n".join(sys.stdin.readlines())
         config = pkg.loads(config_str)
-
     return config
 
 

reinvent/validation.py

@@ -15,6 +15,7 @@ class ReinventConfig(GlobalConfig):
     use_cuda: Optional[bool] = Field(True, deprecated="use 'device' instead")
     tb_logdir: Optional[str] = None
     json_out_config: Optional[str] = None
+    tb_isim: Optional[bool] = False
     parameters: dict
 
     # run mode dependent