Adding iSIM as a trackable value during training

configs/toml/staged_learning.toml

@@ -24,10 +24,14 @@ summary_csv_prefix = "staged_learning"  # prefix for the CSV file
 use_checkpoint = false  # if true read diversity filter from agent_file
 purge_memories = false  # if true purge all diversity filter memories after each stage
 
+
 ## Reinvent
 prior_file = "priors/reinvent.prior"
 agent_file = "priors/reinvent.prior"
 
+##ISIM
+tb_isim = false # if true track isim value of smilies across training epochs in tensorboard 
+
 ## LibInvent
 #prior_file = "priors/libinvent.prior"
 #agent_file = "priors/libinvent.prior"
@@ -176,7 +180,7 @@ max_steps = 100
 
 [stage.scoring]  # the scoring components can be read from a score file
 type = "geometric_mean"  # aggregation function
-filename = "stage2_scoring.toml"  # file with scoring setup for this stage
+filename = "configs/toml/stage2_scoring.toml"  # file with scoring setup for this stage
 filetype = "toml"  # file format: TOML or JSON, no default, must be present
 
 

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,14 @@ 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
     seed: Optional[int] = None
     parameters: dict