# Designing Chelators for Radiopharmaceuticals: Solution *English translation by SOTA – AI Community of the Russian original. Organisers who would like this translation removed can email sota.ai.community@gmail.com.* *Source: the 2023/24 task collection of the National Technology Olympiad (NTO), "Artificial Intelligence" profile, Final Stage, engineering tour, "Solution of the task", pp. 119–128: [original PDF](https://ntcontest.ru/upload/problembooks-2324/13_%D0%98%D1%81%D0%BA%D1%83%D1%81%D1%81%D1%82%D0%B2%D0%B5%D0%BD%D0%BD%D1%8B%D0%B9%20%D0%B8%D0%BD%D1%82%D0%B5%D0%BB%D0%BB%D0%B5%D0%BA%D1%82.pdf#page=119).* ## Solution of the task The solution of the winning team is offered as an example. It used a machine-readable representation of the molecules based on computed physicochemical descriptors obtained from the `rdkit` library, the gradient boosting method to build the predictive model, and a variational autoencoder to build the generative model. Because of the large volume of code written by the participants and the presence of additional files that cannot be converted to text format, this section presents the main elements of the solution that demonstrate the approach proposed by the participants. The complete solution in its original form is available at: . Most of the solution was run on the laptops provided; for the most resource-intensive part, the participants used publicly available computing resources. *Translator's note: the listing below is transcribed from the typeset code in the PDF (lines 1–530); only the comments are translated. Line numbers are omitted, and lines that the PDF wraps are joined. As in the original, the listing is an excerpt: it uses some names that it does not define (for example `X_train`, `sascorer`, `GetRDKitFPGenerator`, `DataStructs`, `c2i`, `string2ids`, `device`).* ```python # importing the required libraries: import pandas as pd import numpy as np from xgboost import XGBRegressor from sklearn.model_selection import train_test_split from rdkit import Chem from rdkit.Chem import PandasTools, Descriptors from rdkit.ML.Descriptors.MoleculeDescriptors import MolecularDescriptorCalculator from joblib import load, dump import seaborn as sns import matplotlib.pyplot as plt # selecting descriptors to describe the molecular data in a format suitable for training machine learning models st = [ 'MaxAbsEStateIndex', 'MinAbsEStateIndex', 'MinEStateIndex', 'qed', 'MolWt', 'NumRadicalElectrons', 'MaxPartialCharge', 'MinPartialCharge', 'MaxAbsPartialCharge', 'FpDensityMorgan1', 'FpDensityMorgan3', 'BCUT2D_MWHI', 'BCUT2D_MWLOW', 'BCUT2D_CHGHI', 'BCUT2D_LOGPLOW', 'BCUT2D_MRLOW', 'AvgIpc', 'BalabanJ', 'Chi4n', 'PEOE_VSA1', 'PEOE_VSA12', 'PEOE_VSA13', 'PEOE_VSA2', 'PEOE_VSA3', 'PEOE_VSA5', 'PEOE_VSA6', 'PEOE_VSA7', 'SMR_VSA1', 'SMR_VSA10', 'SMR_VSA3', 'SMR_VSA5', 'SMR_VSA7', 'SMR_VSA8', 'SMR_VSA9', 'SlogP_VSA1', 'SlogP_VSA11', 'SlogP_VSA2', 'SlogP_VSA7', 'SlogP_VSA8', 'SlogP_VSA9', 'TPSA', 'EState_VSA10', 'EState_VSA3', 'EState_VSA4', 'EState_VSA5', 'EState_VSA6', 'EState_VSA7', 'EState_VSA8', 'EState_VSA9', 'VSA_EState2', 'VSA_EState3', 'VSA_EState4', 'VSA_EState5', 'VSA_EState7', 'VSA_EState8', 'FractionCSP3', 'NHOHCount', 'NumAliphaticCarbocycles', 'NumAromaticHeterocycles', 'NumAromaticRings', 'NumRotatableBonds', 'RingCount', 'fr_Al_COO', 'fr_Al_OH', 'fr_ArN', 'fr_Ar_NH', 'fr_NH0', 'fr_NH1', 'fr_NH2', 'fr_alkyl_halide', 'fr_allylic_oxid', 'fr_amide', 'fr_azide', 'fr_bicyclic', 'fr_diazo', 'fr_ester', 'fr_ether', 'fr_halogen', 'fr_imidazole', 'fr_isocyan', 'fr_para_hydroxylation', 'fr_piperdine', 'fr_priamide', 'fr_prisulfonamd', 'fr_quatN', 'fr_sulfonamd', 'fr_unbrch_alkane'] def my_mol2vec(smiles): global st mol = Chem.MolFromSmiles(smiles) chosen_descriptors = st mol_descriptor_calculator = MolecularDescriptorCalculator(chosen_descriptors) list_of_descriptor_vals = list(mol_descriptor_calculator.CalcDescriptors(mol)) return list_of_descriptor_vals #training the predictive model: def eval_model(model_, X_test, y_test): pred = model_.predict(X_test) summ = 0 for i in range(len(pred)): summ += (pred[i] - y_test.to_numpy()[i])**2 return (summ/len(y_test))**0.5 def save_model(name, model): dump(model, f'{name}.joblib') def my_fit(X_train, X_test, y_train, y_test, rs=123): model = XGBRegressor(n_estimators=500, max_depth=2, learning_rate=0.03, subsample=0.5, seed=rs, num_parallel_tree=16)#, early_stopping_rounds=10) model.fit(X_train, y_train, early_stopping_rounds=100, eval_set=[(X_test, y_test)], verbose=100) return model, eval_model(model, X_test, y_test) def my_fit2(X, y): model = XGBRegressor(n_estimators=2000, max_depth=2, learning_rate=0.03, subsample=0.5, seed=123, num_parallel_tree=16)#, early_stopping_rounds=10) model.fit(X, y) return model model, eval_ = my_fit(X_train, X_test, y_train, y_test) save_model('lgK_hope_to_ans', model) # Checking that the task conditions are met def only_permitted_elements(smiles): mol = Chem.MolFromSmiles(smiles) for atom in mol.GetAtoms(): if not (atom.GetSymbol() in ['C', 'H', 'O', 'N', 'P', 'S']): return False return True def only_permitted_elements_count_min(smiles): mol = Chem.MolFromSmiles(smiles) unique_el = [] for atom in mol.GetAtoms(): if not (atom.GetSymbol() in unique_el): unique_el.append(atom.GetSymbol()) return len(unique_el) >= 3 def only_permitted_elements_count_max(smiles): mol = Chem.MolFromSmiles(smiles) d = {'O': 0, 'N': 0, 'P': 0, 'S': 0} for atom in mol.GetAtoms(): if atom.GetSymbol() in d.keys(): d[atom.GetSymbol()] += 1 sum_el = sum(d.values()) return sum_el <= 12 def MW(smiles): mol = Chem.MolFromSmiles(smiles) MW = Descriptors.MolWt(mol) return MW <= 500 def check_sascore(smiles): mol = Chem.MolFromSmiles(smiles) try: sascore = sascorer.calculateScore(mol) except: return False return sascore < 5 def average_similarity(smiles): simis =[] for i in range(len(smiles)): for j in range(i+1, len(smiles)): smi1 = smiles[i] smi2 = smiles[j] # create molecule using given smiles mol1 = Chem.MolFromSmiles(smi1) mol2 = Chem.MolFromSmiles(smi2) # generate fingerprints fpgen = GetRDKitFPGenerator() fp1 = fpgen.GetFingerprint(mol1) fp2 = fpgen.GetFingerprint(mol2) # compute Tanimoto Similarity using fingerprints simi = DataStructs.TanimotoSimilarity(fp1, fp2) simis.append(simi) return (sum(simis)/len(simis)) def all_check(smiles): flags = [only_permitted_elements(smiles), only_permitted_elements_count_min(smiles), only_permitted_elements_count_max(smiles), MW(smiles), check_sascore(smiles)] return all(flags) # creating a variational autoencoder to be used as the generative model # -*- coding: utf-8 -*- import torch import torch.nn as nn import torch.nn.functional as F q_bidir = True q_d_h = 256 q_n_layers = 1 q_dropout = 0.5 d_n_layers = 3 d_dropout = 0 d_z = 128 d_d_h = 512 class VAE(nn.Module): def __init__(self,vocab,vector): super().__init__() self.vocabulary = vocab self.vector = vector n_vocab, d_emb = len(vocab), vector.size(1) self.x_emb = nn.Embedding(n_vocab, d_emb, c2i['']) self.x_emb.weight.data.copy_(vector) #ENCODER self.encoder_rnn = nn.GRU(d_emb,q_d_h,num_layers=q_n_layers,batch_first=True,dropout=q_dropout if q_n_layers > 1 else 0,bidirectional=q_bidir) q_d_last = q_d_h * (2 if q_bidir else 1) self.q_mu = nn.Linear(q_d_last, d_z) self.q_logvar = nn.Linear(q_d_last, d_z) # Decoder self.decoder_rnn = nn.GRU(d_emb + d_z,d_d_h,num_layers=d_n_layers,batch_first=True,dropout=d_dropout if d_n_layers > 1 else 0) self.decoder_latent = nn.Linear(d_z, d_d_h) self.decoder_fullyc = nn.Linear(d_d_h, n_vocab) # Grouping the model's parameters self.encoder = nn.ModuleList([self.encoder_rnn,self.q_mu,self.q_logvar]) self.decoder = nn.ModuleList([self.decoder_rnn,self.decoder_latent,self.decoder_fullyc]) self.vae = nn.ModuleList([self.x_emb,self.encoder,self.decoder]) @property def device(self): return next(self.parameters()).device def string2tensor(self, string, device='model'): ids = string2ids(string, add_bos=True, add_eos=True) tensor = torch.tensor(ids, dtype=torch.long,device=self.device if device == 'model' else device) return tensor def tensor2string(self, tensor): ids = tensor.tolist() string = ids2string(ids, rem_bos=True, rem_eos=True) return string def forward(self,x): global count_epoch z, kl_loss = self.forward_encoder(x) recon_loss = self.forward_decoder(x, z) count_epoch += 1 if count_epoch % 100 == 0: print(count_epoch) return kl_loss, recon_loss def forward_encoder(self,x): x = [self.x_emb(i_x) for i_x in x] x = nn.utils.rnn.pack_sequence(x) _, h = self.encoder_rnn(x, None) h = h[-(1 + int(self.encoder_rnn.bidirectional)):] h = torch.cat(h.split(1), dim=-1).squeeze(0) mu, logvar = self.q_mu(h), self.q_logvar(h) eps = torch.randn_like(mu) z = mu + (logvar / 2).exp() * eps kl_loss = 0.5 * (logvar.exp() + mu ** 2 - 1 - logvar).sum(1).mean() return z, kl_loss def forward_decoder(self,x, z): lengths = [len(i_x) for i_x in x] x = nn.utils.rnn.pad_sequence(x, batch_first=True, padding_value= c2i['']) x_emb = self.x_emb(x) z_0 = z.unsqueeze(1).repeat(1, x_emb.size(1), 1) x_input = torch.cat([x_emb, z_0], dim=-1) x_input = nn.utils.rnn.pack_padded_sequence(x_input, lengths, batch_first=True) h_0 = self.decoder_latent(z) h_0 = h_0.unsqueeze(0).repeat(self.decoder_rnn.num_layers, 1, 1) output, _ = self.decoder_rnn(x_input, h_0) output, _ = nn.utils.rnn.pad_packed_sequence(output, batch_first=True) y = self.decoder_fullyc(output) recon_loss = F.cross_entropy(y[:, :-1].contiguous().view(-1, y.size(-1)),x[:, 1:].contiguous().view(-1),ignore_index= c2i['']) return recon_loss def sample_z_prior(self,n_batch): return torch.randn(n_batch,self.q_mu.out_features,device= self.x_emb.weight.device) def sample(self,n_batch, max_len=100, z=None, temp=1.0): with torch.no_grad(): if z is None: z = self.sample_z_prior(n_batch) z = z.to(self.device) z_0 = z.unsqueeze(1) h = self.decoder_latent(z) h = h.unsqueeze(0).repeat(self.decoder_rnn.num_layers, 1, 1) w = torch.tensor(c2i[''], device=self.device).repeat(n_batch) x = torch.tensor([c2i['']], device=device).repeat(n_batch, max_len) x[:, 0] = c2i[''] end_pads = torch.tensor([max_len], device=self.device).repeat(n_batch) eos_mask = torch.zeros(n_batch, dtype=torch.bool, device=self.device) for i in range(1, max_len): x_emb = self.x_emb(w).unsqueeze(1) x_input = torch.cat([x_emb, z_0], dim=-1) o, h = self.decoder_rnn(x_input, h) y = self.decoder_fullyc(o.squeeze(1)) y = F.softmax(y / temp, dim=-1) w = torch.multinomial(y, 1)[:, 0] x[~eos_mask, i] = w[~eos_mask] i_eos_mask = ~eos_mask & (w == c2i['']) end_pads[i_eos_mask] = i + 1 eos_mask = eos_mask | i_eos_mask new_x = [] for i in range(x.size(0)): new_x.append(x[i, :end_pads[i]]) return [self.tensor2string(i_x) for i_x in new_x] from torch.optim.lr_scheduler import _LRScheduler import torch.optim as optim from torch.utils.data import DataLoader from torch.nn.utils import clip_grad_norm_ import math import numpy as np from collections import UserList, defaultdict n_last = 1000 n_batch = 32 kl_start = 0 kl_w_start = 0.0 kl_w_end = 1.0 n_epoch = 200 n_workers = 0 clip_grad = 50 lr_start = 0.003 lr_n_period = 10 lr_n_mult = 1 lr_end = 3 * 1e-4 lr_n_restarts = 6 def _n_epoch(): return sum(lr_n_period * (lr_n_mult ** i) for i in range(lr_n_restarts)) def _train_epoch(model, epoch, train_loader, kl_weight, optimizer=None): if optimizer is None: model.eval() else: model.train() kl_loss_values = CircularBuffer(n_last) recon_loss_values = CircularBuffer(n_last) loss_values = CircularBuffer(n_last) for i, input_batch in enumerate(train_loader): input_batch = tuple(data.to(device) for data in input_batch) #forward kl_loss, recon_loss = model(input_batch) loss = kl_weight * kl_loss + recon_loss #backward if optimizer is not None: optimizer.zero_grad() loss.backward() clip_grad_norm_(get_optim_params(model),clip_grad) optimizer.step() kl_loss_values.add(kl_loss.item()) recon_loss_values.add(recon_loss.item()) loss_values.add(loss.item()) lr = (optimizer.param_groups[0]['lr'] if optimizer is not None else None) #update train_loader kl_loss_value = kl_loss_values.mean() recon_loss_value = recon_loss_values.mean() loss_value = loss_values.mean() postfix = [f'loss={loss_value:.5f}', f'(kl={kl_loss_value:.5f}', f'recon={recon_loss_value:.5f})', f'klw={kl_weight:.5f} lr={lr:.5f}'] postfix = {'epoch': epoch,'kl_weight': kl_weight,'lr': lr,'kl_loss': kl_loss_value,'recon_loss': recon_loss_value,'loss': loss_value,'mode': 'Eval' if optimizer is None else 'Train'} return postfix def _train(model, train_loader, val_loader=None, logger=None): optimizer = optim.Adam(get_optim_params(model),lr= lr_start) lr_annealer = CosineAnnealingLRWithRestart(optimizer) model.zero_grad() for epoch in range(n_epoch): kl_annealer = KLAnnealer(n_epoch) kl_weight = kl_annealer(epoch) postfix = _train_epoch(model, epoch,train_loader, kl_weight, optimizer) lr_annealer.step() def fit(model, train_data, val_data=None): logger = Logger() if False is not None else None train_loader = get_dataloader(model,train_data,shuffle=True) val_loader = None if val_data is None else get_dataloader(model, val_data, shuffle=False) _train(model, train_loader, val_loader, logger) return model def get_collate_device(model): return model.device def get_dataloader(model, train_data, collate_fn=None, shuffle=True): if collate_fn is None: collate_fn = get_collate_fn(model) print(collate_fn) return DataLoader(train_data, batch_size=n_batch, shuffle=shuffle, num_workers=n_workers, collate_fn=collate_fn) def get_collate_fn(model): device = get_collate_device(model) def collate(train_data): train_data.sort(key=len, reverse=True) tensors = [string2tensor(string, device=device) for string in train_data] return tensors return collate def get_optim_params(model): return (p for p in model.parameters() if p.requires_grad) class KLAnnealer: def __init__(self,n_epoch): self.i_start = kl_start self.w_start = kl_w_start self.w_max = kl_w_end self.n_epoch = n_epoch self.inc = (self.w_max - self.w_start) / (self.n_epoch - self.i_start) def __call__(self, i): k = (i - self.i_start) if i >= self.i_start else 0 return self.w_start + k * self.inc class CosineAnnealingLRWithRestart(_LRScheduler): def __init__(self , optimizer): self.n_period = lr_n_period self.n_mult = lr_n_mult self.lr_end = lr_end self.current_epoch = 0 self.t_end = self.n_period # Also calls first epoch super().__init__(optimizer, -1) def get_lr(self): return [self.lr_end + (base_lr - self.lr_end) * (1 + math.cos(math.pi * self.current_epoch / self.t_end)) / 2 for base_lr in self.base_lrs] def step(self, epoch=None): if epoch is None: epoch = self.last_epoch + 1 self.last_epoch = epoch self.current_epoch += 1 for param_group, lr in zip(self.optimizer.param_groups, self.get_lr()): param_group['lr'] = lr if self.current_epoch == self.t_end: self.current_epoch = 0 self.t_end = self.n_mult * self.t_end class CircularBuffer: def __init__(self, size): self.max_size = size self.data = np.zeros(self.max_size) self.size = 0 self.pointer = -1 def add(self, element): self.size = min(self.size + 1, self.max_size) self.pointer = (self.pointer + 1) % self.max_size self.data[self.pointer] = element return element def last(self): assert self.pointer != -1, "Can't get an element from an empty buffer!" return self.data[self.pointer] def mean(self): return self.data.mean() class Logger(UserList): def __init__(self, data=None): super().__init__() self.sdata = defaultdict(list) for step in (data or []): self.append(step) def __getitem__(self, key): if isinstance(key, int): return self.data[key] elif isinstance(key, slice): return Logger(self.data[key]) else: ldata = self.sdata[key] if isinstance(ldata[0], dict): return Logger(ldata) else: return ldata def append(self, step_dict): super().append(step_dict) for k, v in step_dict.items(): self.sdata[k].append(v) ```