torch-molecule 0.1.0.dev2__py3-none-any.whl

torch_molecule/generator/jtvae/modeling_jtvae.py

@@ -0,0 +1,247 @@
+from tqdm import tqdm
+from dataclasses import dataclass, field
+from typing import Optional, Dict, Any, Tuple, List, Type
+import numpy as np
+
+import torch
+
+from .jtnn_vae import JTNNVAE
+from .jtnn.mol_tree import MolTree
+from .jtnn.vocab import Vocab
+from .jtnn.datautils import MolTreeFolder
+from ...base import BaseMolecularGenerator
+
+@dataclass
+class JTVAEMolecularGenerator(BaseMolecularGenerator):
+    """ 
+    JT-VAE-based molecular generator. Implemented for unconditional moleculargeneration.     
+
+    References
+    ----------
+    - Junction Tree Variational Autoencoder for Molecular Graph Generation. ICML 2018. https://arxiv.org/pdf/1802.04364
+    - Code: https://github.com/kamikaze0923/jtvae
+
+    Parameters
+    ----------
+    hidden_size : int, default=450
+        Dimension of hidden layers in the model.
+    latent_size : int, default=56
+        Dimension of the latent space.
+    depthT : int, default=20
+        Depth of the tree encoder.
+    depthG : int, default=3
+        Depth of the graph decoder.
+    batch_size : int, default=32
+        Number of samples per batch during training.
+    epochs : int, default=20
+        Number of epochs to train the model.
+    learning_rate : float, default=0.003
+        Initial learning rate for the optimizer.
+    weight_decay : float, default=0.0
+        L2 regularization factor.
+    grad_norm_clip : Optional[float], default=None
+        Maximum norm for gradient clipping. None means no clipping.
+    beta : float, default=0.0
+        Initial KL divergence weight for VAE training.
+    step_beta : float, default=0.002
+        Step size for KL annealing.
+    max_beta : float, default=1.0
+        Maximum value for KL weight.
+    warmup : int, default=40000
+        Number of steps for KL annealing warmup.
+    use_lr_scheduler : bool, default=True
+        Whether to use learning rate scheduling.
+    anneal_rate : float, default=0.9
+        Learning rate annealing factor.
+    anneal_iter : int, default=40000
+        Number of iterations between learning rate updates.
+    kl_anneal_iter : int, default=2000
+        Number of iterations between KL weight updates.
+    verbose : bool, default=False
+        Whether to print detailed training information.
+    """
+    # Model parameters
+    hidden_size: int = 450
+    latent_size: int = 56
+    depthT: int = 20
+    depthG: int = 3
+
+    # Training parameters
+    batch_size: int = 32
+    epochs: int = 20
+    learning_rate: float = 0.003
+    weight_decay: float = 0.0
+    grad_norm_clip: Optional[float] = None
+    
+    # KL annealing parameters
+    beta: float = 0.0
+    step_beta: float = 0.002
+    max_beta: float = 1.0
+    warmup: int = 40000
+
+    # Learning rate scheduling
+    use_lr_scheduler: bool = True
+    anneal_rate: float = 0.9
+    anneal_iter: int = 40000
+    kl_anneal_iter: int = 2000
+
+    # Other parameters
+    verbose: bool = False
+
+    # attributes
+    model_name: str = "JTVAEMolecularGenerator"
+    fitting_loss: List[float] = field(default_factory=list, init=False)
+    fitting_epoch: int = field(default=0, init=False)
+    model_class: Type[JTNNVAE] = field(default=JTNNVAE, init=False)
+    
+    def __post_init__(self):
+        super().__post_init__()
+        self.vocab = None
+
+    @staticmethod
+    def _get_param_names() -> List[str]:
+        return [
+            "hidden_size", "latent_size", "depthT", "depthG",
+            "batch_size", "epochs", "learning_rate", "weight_decay",
+            "beta", "step_beta", "max_beta", "warmup",
+            "use_lr_scheduler", "anneal_rate", "anneal_iter", "kl_anneal_iter",
+            "verbose", 'model_name', 'vocab'
+        ]
+    
+    def _get_model_params(self, checkpoint: Optional[Dict] = None) -> Dict[str, Any]:
+        params = [
+            "vocab", "hidden_size", "latent_size", "depthT", "depthG"
+        ]
+        if checkpoint is not None:
+            if "hyperparameters" not in checkpoint:
+                raise ValueError("Checkpoint missing 'hyperparameters' key")
+            return {k: checkpoint["hyperparameters"][k] for k in params}
+    
+        return {k: getattr(self, k) for k in params}
+
+    def _setup_optimizers(self) -> Tuple[torch.optim.Optimizer, Optional[Any]]:
+        """Setup optimization components including optimizer and learning rate scheduler.
+        """
+        optimizer = torch.optim.Adam(
+            self.model.parameters(), lr=self.learning_rate, weight_decay=self.weight_decay
+        )
+        scheduler = None
+        if self.use_lr_scheduler:
+            scheduler = torch.optim.lr_scheduler.ExponentialLR(optimizer, self.anneal_rate)
+
+        return optimizer, scheduler
+
+    def _extract_vocab(self, X):
+        cset = set()
+        for ii, smiles in enumerate(X):
+            try:
+                mol = MolTree(smiles)
+                for c in mol.nodes:
+                    cset.add(c.smiles)
+            except Exception as e:
+                print(f'Error {e} in extracting vocab for smiles: {smiles}')
+                pass
+        vocab = list(cset)
+        return vocab
+
+    def _convert_to_tensor(self, X):
+        all_data = []
+        for smiles in X:
+            try:
+                mol_tree = MolTree(smiles)
+                mol_tree.recover()
+                mol_tree.assemble()
+                for node in mol_tree.nodes:
+                    if node.label not in node.cands:
+                        node.cands.append(node.label)
+
+                del mol_tree.mol
+                for node in mol_tree.nodes:
+                    del node.mol
+            except Exception as e:
+                print(f'Error {e} in tensorizing smiles: {smiles}')
+                mol_tree = None
+            all_data.append(mol_tree)
+        return all_data
+
+    def fit(
+        self,
+        X_train: List[str],
+    ) -> "JTVAEMolecularGenerator":
+
+        X_train, _ = self._validate_inputs(X_train, None, num_task=0, return_rdkit_mol=False)
+        vocab = self._extract_vocab(X_train)
+        vocab = Vocab(vocab)
+        self.vocab = vocab
+
+        self._initialize_model(self.model_class)
+        self.model.initialize_parameters()
+        optimizer, scheduler = self._setup_optimizers()
+
+        train_dataset = self._convert_to_tensor(X_train)
+        train_dataset = list(filter(lambda x: x is not None, train_dataset))
+        train_loader = MolTreeFolder(train_dataset, vocab, self.batch_size, num_workers=0)
+        step_len = len(X_train) // self.batch_size
+        self.fitting_loss = []
+        self.fitting_epoch = 0
+        total_step = 0
+        for epoch in range(self.epochs):
+            train_losses, total_step = self._train_epoch(train_loader, optimizer, scheduler, epoch, total_step, step_len)
+            self.fitting_loss.append(np.mean(train_losses).item())
+
+        self.fitting_epoch = epoch
+        self.is_fitted_ = True
+        return self
+
+    def _train_epoch(self, train_loader, optimizer, scheduler, epoch, total_step, step_len):
+        self.model.train()
+        losses = []
+        iterator = (
+            tqdm(train_loader, desc="Training", leave=False, total=step_len)
+            if self.verbose
+            else train_loader
+        )
+        for step, batched_data in enumerate(iterator):
+            total_step += 1
+            optimizer.zero_grad()
+
+            word_loss, topo_loss, assm_loss, kl_div = self.model.compute_loss(batched_data)
+            total_loss = word_loss + topo_loss + assm_loss + kl_div * self.beta
+            total_loss.backward()
+            if self.grad_norm_clip is not None:
+                torch.nn.utils.clip_grad_norm_(self.model.parameters(), self.grad_norm_clip)
+            optimizer.step()
+            losses.append(total_loss.item())
+
+            if total_step % self.anneal_iter == 0:
+                scheduler.step()
+                
+            if total_step % self.kl_anneal_iter == 0 and total_step >= self.warmup:
+                self.beta = min(self.max_beta, self.beta + self.step_beta)
+
+            if self.verbose:
+                iterator.set_postfix({"Epoch": epoch, "Loss": f"{total_loss.item():.4f}, word_loss: {word_loss.item():.4f}, topo_loss: {topo_loss.item():.4f}, assm_loss: {assm_loss.item():.4f}, kl_div: {kl_div.item():.4f}"})
+            
+        return losses, total_step
+
+    def generate(
+        self, 
+        batch_size: int = 32
+    ) -> List[str]:
+        """Generate molecules using JT-VAE.
+        
+        Parameters
+        ----------  
+        batch_size : int
+            Number of molecules to generate.
+            
+        Returns
+        -------
+        List[str]
+            Generated molecules as SMILES strings.
+        """
+        if not self.is_fitted_:
+            raise ValueError("Model must be fitted before generating molecules.")
+        
+        return self.model.sample_prior(self.device)
+    

torch_molecule/generator/lstm/__init__.py

@@ -0,0 +1,3 @@
+from .modeling_lstm import LSTMMolecularGenerator
+
+__all__ = ['LSTMMolecularGenerator']

torch_molecule/generator/lstm/action_sampler.py

@@ -0,0 +1,91 @@
+from typing import Type
+
+import torch
+import torch.nn.functional as F
+from torch.distributions import Categorical, Distribution
+
+from .lstm import LSTM
+from .utils import rnn_start_token_vector
+
+class ActionSampler:
+    """
+    Sampler for a SmilesRNN models.
+
+    Does not return SMILES strings directly, but instead the actions (i.e. which SMILES character to select).
+    Those values are more general and are for instance necessary for other RL algorithms.
+
+    The class will sample the RNN model multiple times if the number of desired samples is larger than the
+    maximal allowed batch size.
+    """
+
+    def __init__(self, max_batch_size, max_seq_length, device,
+                 distribution_cls: Type[Distribution] = Categorical) -> None:
+        """
+        Args:
+            max_batch_size: maximal batch size for the RNN model
+            max_seq_length: max length for a sampled SMILES string
+            device: cuda | cpu
+            distribution_cls: distribution type to sample from. If None, will be a multinomial distribution. Useful for testing purposes.
+        """
+        self.max_batch_size = max_batch_size
+        self.max_seq_length = max_seq_length
+        self.device = device
+        self.distribution_cls = distribution_cls
+
+    def sample(self, model: LSTM, num_samples: int, target: torch.Tensor) -> torch.Tensor:
+        """
+        Samples a specified number of actions from an RNN model based on a multinomial distribution.
+
+        Args:
+            model: Smiles RNN model to sample from
+            num_samples: Number of samples to generate
+
+        Returns:
+            tensor of actions (num_samples x max_seq_length)
+        """
+
+        # Round up division to get the number of batches that are necessary:
+        number_batches = (num_samples + self.max_batch_size - 1) // self.max_batch_size
+        remaining_samples = num_samples
+
+        actions = torch.LongTensor(num_samples, self.max_seq_length).to(self.device)
+
+        batch_start = 0
+
+        for i in range(number_batches):
+            batch_size = min(self.max_batch_size, remaining_samples)
+            batch_end = batch_start + batch_size
+
+            actions[batch_start:batch_end, :] = self._sample_batch(model, batch_size, target)
+
+            batch_start += batch_size
+            remaining_samples -= batch_size
+
+        return actions
+
+    def _sample_batch(self, model: LSTM, batch_size: int, target: torch.Tensor) -> torch.Tensor:
+        """
+        Samples a batch of actions based on a multinomial distribution.
+
+        Args:
+            model: Smiles RNN model to sample from
+            num_samples: Number of samples to generate
+
+        Returns:
+            tensor of actions (batch_size x max_seq_length)
+        """
+        hidden, cell = model.init_hidden(batch_size, target)
+        inp = rnn_start_token_vector(batch_size, self.device)
+        actions = torch.zeros((batch_size, self.max_seq_length), dtype=torch.long).to(self.device)
+
+        for char in range(self.max_seq_length):
+            output, hidden, cell = model(inp, hidden, cell)
+            prob = F.softmax(output, dim=2)
+            distribution = self.distribution_cls(probs=prob)
+            action = distribution.sample()
+
+            actions[:, char] = action.squeeze()
+
+            inp = action
+
+        return actions

torch_molecule/generator/lstm/lstm.py

@@ -0,0 +1,59 @@
+import torch
+import torch.nn as nn
+
+class LSTM(nn.Module):
+    def __init__(self, num_task, input_size, hidden_size, output_size, num_layer, dropout):
+        super().__init__()
+        self.input_size = input_size
+        self.hidden_size = hidden_size
+        self.output_size = output_size
+        self.num_layer = num_layer
+        self.dropout = dropout
+        if num_task == 0:
+            self.input_dim = 1
+        else:
+            self.input_dim = num_task
+        self.hidden_transform = nn.Linear(self.input_dim, num_layer * hidden_size)
+        self.cell_transform = nn.Linear(self.input_dim, num_layer * hidden_size)
+        self.encoder = nn.Embedding(input_size, hidden_size)
+        self.decoder = nn.Linear(hidden_size, output_size)
+
+        self.rnn = nn.LSTM(hidden_size, hidden_size, batch_first=True, num_layers=num_layer, dropout=dropout)
+        self.initialize_parameters()
+
+    def initialize_parameters(self):
+        # encoder / decoder
+        nn.init.xavier_uniform_(self.encoder.weight)
+        nn.init.xavier_uniform_(self.decoder.weight)
+        nn.init.constant_(self.decoder.bias, 0)
+        # RNN
+        for name, param in self.rnn.named_parameters():
+            if 'weight' in name:
+                nn.init.orthogonal_(param)
+            elif 'bias' in name:
+                nn.init.constant_(param, 0)
+                # LSTM remember gate bias should be initialised to 1
+                # https://github.com/pytorch/pytorch/issues/750
+                r_gate = param[int(0.25 * len(param)):int(0.5 * len(param))]
+                nn.init.constant_(r_gate, 1)
+
+    def forward(self, input, hidden, cell):
+        embeds = self.encoder(input)
+        output, (hidden, cell) = self.rnn(embeds, (hidden, cell))
+        output = self.decoder(output)
+        return output, hidden, cell
+
+    def init_hidden(self, bsz, target):
+        hidden = self.hidden_transform(target)
+        cell = self.cell_transform(target)
+        hidden = hidden.view(self.num_layer, bsz, self.hidden_size)
+        cell = cell.view(self.num_layer, bsz, self.hidden_size)
+        return hidden, cell
+    
+    def compute_loss(self, batch_data, criterion):
+        ipt, tgt, y = batch_data
+        hidden, cell = self.init_hidden(ipt.size(0), y)
+        output, hidden, cell = self.forward(ipt, hidden, cell)
+        output = output.view(output.size(0) * output.size(1), -1)
+        loss = criterion(output, tgt.view(-1))
+        return loss

torch_molecule/generator/lstm/modeling_lstm.py

@@ -0,0 +1,263 @@
+from tqdm import tqdm
+from typing import Optional, Union, Dict, Any, Tuple, List, Callable, Type
+from dataclasses import dataclass, field
+import numpy as np
+
+import torch
+from torch.utils.data import DataLoader, TensorDataset
+
+from .lstm import LSTM
+from .utils import canonicalize
+from .action_sampler import ActionSampler
+from .smiles_char_dict import SmilesCharDictionary
+from ...base import BaseMolecularGenerator
+
+@dataclass
+class LSTMMolecularGenerator(BaseMolecularGenerator):
+    """LSTM-based molecular generator.
+    
+    This generator implements an LSTM architecture for molecular generation. 
+    When conditions (y values) are provided during training, they are used to initialize the hidden 
+    and cell states of the LSTM.
+    
+    Parameters
+    ----------
+    num_task : int, default=0
+        Number of tasks for conditional generation (0 means unconditional generation).
+    max_len : int, default=100
+        Maximum length of the SMILES strings.
+    num_layer : int, default=3
+        Number of LSTM layers.
+    hidden_size : int, default=512
+        Dimension of hidden states in LSTM.
+    dropout : float, default=0.2
+        Dropout probability for regularization.
+    batch_size : int, default=128
+        Batch size for training.
+    epochs : int, default=10000
+        Maximum number of training epochs.
+    learning_rate : float, default=0.0002
+        Learning rate for optimizer.
+    weight_decay : float, default=0.0
+        L2 regularization factor.
+    use_lr_scheduler : bool, default=False
+        Whether to use learning rate scheduler.
+    scheduler_factor : float, default=0.5
+        Factor by which to reduce learning rate when using scheduler (if True).
+    scheduler_patience : int, default=5
+        Number of epochs with no improvement after which learning rate will be reduced (if True).
+    grad_norm_clip : Optional[float], default=None
+        Maximum norm for gradient clipping. None means no clipping.
+    verbose : bool, default=False
+        Whether to print progress during training.
+    """
+    num_task: int = 0
+    # LSTM parameters
+    max_len: int = 100
+    num_layer: int = 3
+    hidden_size: int = 512
+    dropout: float = 0.2
+
+    # Training parameters
+    batch_size: int = 128
+    epochs: int = 10000
+    learning_rate: float = 0.0002
+    weight_decay: float = 0.0
+    # Scheduler parameters
+    use_lr_scheduler: bool = False
+    scheduler_factor: float = 0.5
+    scheduler_patience: int = 5
+    grad_norm_clip: Optional[float] = None
+    # Other parameters
+    verbose: bool = False
+
+    # attributes
+    model_name: str = "LSTMMolecularGenerator"
+    fitting_loss: List[float] = field(default_factory=list, init=False)
+    fitting_epoch: int = field(default=0, init=False)
+    input_size: int = field(init=False, default=None)
+    output_size: int = field(init=False, default=None)
+    model_class: Type[LSTM] = field(default=LSTM, init=False)
+    
+    def __post_init__(self):
+        super().__post_init__()
+        self.tokenizer = SmilesCharDictionary()
+        self.input_size = len(self.tokenizer.char_idx)
+        self.output_size = len(self.tokenizer.char_idx)
+
+    @staticmethod
+    def _get_param_names() -> List[str]:
+        return [
+            "num_task", "max_len", "num_layer", "hidden_size", "dropout",
+            "batch_size", "epochs", "learning_rate", "weight_decay",
+            "use_lr_scheduler", "scheduler_factor", "scheduler_patience",
+            "grad_norm_clip", "verbose", "input_size", "output_size", 'model_name'
+        ]
+    
+    def _get_model_params(self, checkpoint: Optional[Dict] = None) -> Dict[str, Any]:
+        params = [
+            "num_task", "input_size", "hidden_size", "output_size", "num_layer", "dropout",
+        ]
+        if checkpoint is not None:
+            if "hyperparameters" not in checkpoint:
+                raise ValueError("Checkpoint missing 'hyperparameters' key")
+            return {k: checkpoint["hyperparameters"][k] for k in params}
+    
+        return {k: getattr(self, k) for k in params}
+
+    def _convert_to_pytorch_data(self, X, y=None):
+        # filter valid smiles strings
+        valid_smiles = []
+        for s in X:
+            s = s.strip()
+            if self.tokenizer.allowed(s) and len(s) <= self.max_len:
+                valid_smiles.append(s)
+            else:
+                valid_smiles.append('C')  # default placeholder
+        
+        # max len + two chars for start token 'Q' and stop token '\n'
+        max_seq_len = self.max_len + 2
+        # allocate the zero matrix to be filled
+        seqs = np.zeros((len(valid_smiles), max_seq_len), dtype=np.int32)
+        
+        for i, mol in enumerate(valid_smiles):
+            enc_smi = self.tokenizer.BEGIN + self.tokenizer.encode(mol) + self.tokenizer.END
+            for c in range(len(enc_smi)):
+                seqs[i, c] = self.tokenizer.char_idx[enc_smi[c]]
+        
+        seqs_tensor = torch.from_numpy(seqs).long()
+        inp = seqs_tensor[:, :-1]
+        target = seqs_tensor[:, 1:]
+        bsz = inp.size(0)
+        if y is not None:
+            assert len(y) == bsz
+            y = torch.tensor(y)
+            return TensorDataset(inp, target, y)
+        else:
+            target_y = torch.zeros(bsz, 1)
+            return TensorDataset(inp, target, target_y)
+
+    def _setup_optimizers(self) -> Tuple[torch.optim.Optimizer, Optional[Any]]:
+        """Setup optimization components including optimizer and learning rate scheduler.
+        """
+        optimizer = torch.optim.Adam(
+            self.model.parameters(), lr=self.learning_rate, weight_decay=self.weight_decay
+        )
+        scheduler = None
+        if self.use_lr_scheduler:
+            scheduler = torch.optim.lr_scheduler.ReduceLROnPlateau(
+                optimizer,
+                mode="min",
+                factor=self.scheduler_factor,
+                patience=self.scheduler_patience,
+                min_lr=1e-6,
+                cooldown=0,
+                eps=1e-8,
+            )
+
+        return optimizer, scheduler
+    
+    def fit(
+        self,
+        X_train: List[str],
+        y_train: Optional[Union[List, np.ndarray]] = None,
+    ) -> "LSTMMolecularGenerator":
+        X_train, y_train = self._validate_inputs(X_train, y_train, num_task=self.num_task, return_rdkit_mol=False)
+        self._initialize_model(self.model_class)
+        self.model.initialize_parameters()
+        optimizer, scheduler = self._setup_optimizers()
+
+        train_dataset = self._convert_to_pytorch_data(X_train, y_train)
+        train_loader = DataLoader(
+            train_dataset,
+            batch_size=self.batch_size,
+            shuffle=True,
+            num_workers=0
+        )
+
+        self.fitting_loss = []
+        self.fitting_epoch = 0
+        criterion = torch.nn.CrossEntropyLoss()
+        for epoch in range(self.epochs):
+            train_losses = self._train_epoch(train_loader, optimizer, epoch, criterion)
+            self.fitting_loss.append(np.mean(train_losses).item())
+            if scheduler:
+                scheduler.step(np.mean(train_losses).item())
+
+        self.fitting_epoch = epoch
+        self.is_fitted_ = True
+        return self
+
+    def _train_epoch(self, train_loader, optimizer, epoch, criterion):
+        self.model.train()
+        losses = []
+        iterator = (
+            tqdm(train_loader, desc="Training", leave=False)
+            if self.verbose
+            else train_loader
+        )
+        for step, batched_data in enumerate(iterator):
+            for i in range(len(batched_data)):
+                batched_data[i] = batched_data[i].to(self.device)
+            optimizer.zero_grad()
+
+            loss = self.model.compute_loss(batched_data, criterion)
+            loss.backward()
+            if self.grad_norm_clip is not None:
+                torch.nn.utils.clip_grad_norm_(self.model.parameters(), self.grad_norm_clip)
+            optimizer.step()
+            losses.append(loss.item())
+            
+            if self.verbose:
+                iterator.set_postfix({"Epoch": epoch, "Loss": f"{loss.item():.4f}"})
+            
+        return losses
+
+    def generate(
+        self, 
+        labels: Optional[Union[List[List], np.ndarray]] = None,
+        batch_size: int = 32
+    ) -> List[str]:
+        """Generate molecules using LSTM.
+        
+        Parameters
+        ----------
+        labels : Optional[Union[List[List], np.ndarray]]
+            Target property values for conditional generation.
+        batch_size : int
+            Number of molecules to generate.
+            
+        Returns
+        -------
+        List[str]
+            Generated molecules as SMILES strings.
+        """
+        if not self.is_fitted_:
+            raise ValueError("Model must be fitted before generating molecules.")
+        
+        if labels is not None and batch_size != len(labels):
+            print(f"batch_size is not equal to the length of labels, using the length of labels: {len(labels)}")
+            batch_size = len(labels)
+
+        # Convert properties to 2D tensor if needed
+        if isinstance(labels, list):
+            labels = torch.tensor(labels)
+        elif isinstance(labels, np.ndarray):
+            labels = torch.from_numpy(labels)
+        if labels is not None and labels.dim() == 1:
+            labels = labels.unsqueeze(-1)
+        elif labels is None:
+            labels = torch.zeros(batch_size, 1)
+        labels = labels.to(self.device)
+        
+        all_generated_mols = []
+        sampler = ActionSampler(max_batch_size=batch_size, max_seq_length=self.max_len, device=self.device)
+        self.model.eval()
+        with torch.no_grad():
+            indices = sampler.sample(self.model, num_samples=batch_size, target=labels)
+            samples = self.tokenizer.matrix_to_smiles(indices)
+            canonicalized_smiles = [canonicalize(smiles, include_stereocenters=True) for smiles in samples]
+            all_generated_mols.extend(canonicalized_smiles)
+
+        return all_generated_mols
+