graphchem.preprocessing.MoleculeEncoder

Bases: object

A class to encode molecular SMILES strings into numerical (integer) representations using tokenized atom and bond information.

Attributes

_atom_tokenizer : Tokenizer integer Tokenizer for atom representations. _bond_tokenizer : Tokenizer integer Tokenizer for bond representations.

Source code in graphchem/preprocessing/features.py

class MoleculeEncoder(object):
    """
    A class to encode molecular SMILES strings into numerical (integer)
    representations using tokenized atom and bond information.

    Attributes
    ----------
    _atom_tokenizer : Tokenizer
        integer Tokenizer for atom representations.
    _bond_tokenizer : Tokenizer
        integer Tokenizer for bond representations.
    """

    def __init__(self, smiles: List[str]):
        """
        Initializes the MoleculeEncoder with a list of SMILES strings and
        creates/trains integer tokenizers for atoms and bonds.

        Parameters
        ----------
        smiles : List[str]
            A list of SMILES strings representing molecules used for tokenizer
            creation/training.

        Raises
        ------
        ValueError
            If any provided SMILES string cannot be parsed by RDKit.
        """
        mols = [Chem.MolFromSmiles(smi) for smi in smiles]
        for idx, mol in enumerate(mols):
            if mol is None:
                raise ValueError(f"Unable to parse SMILES: {smiles[idx]}")

        atoms = np.concatenate([mol.GetAtoms() for mol in mols])
        atom_reprs = [atom_to_str(atom) for atom in atoms]
        bond_reprs = np.concatenate([
            [bond_to_str(bond) for bond in atom.GetBonds()]
            for atom in atoms
        ])

        self._atom_tokenizer = Tokenizer()
        for rep in atom_reprs:
            self._atom_tokenizer(rep)
        self._atom_tokenizer.train = False

        self._bond_tokenizer = Tokenizer()
        for rep in bond_reprs:
            self._bond_tokenizer(rep)
        self._bond_tokenizer.train = False

    @property
    def vocab_sizes(self) -> Tuple[int, int]:
        """
        Returns the vocabulary sizes of the atom and bond tokenizers.

        Returns
        -------
        Tuple[int, int]
            A tuple containing two integers representing the sizes of the atom
            and bond tokenizers' vocabularies respectively.
        """
        return (
            self._atom_tokenizer.vocab_size,
            self._bond_tokenizer.vocab_size
        )

    def encode(
            self,
            smiles: str
         ) -> Tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
        """
        Encodes a single SMILES string into three tensors representing atoms,
        bonds, and connectivity.

        Parameters
        ----------
        smiles : str
            A SMILES string representing the molecule to be encoded.

        Returns
        -------
        Tuple[torch.Tensor, torch.Tensor, torch.Tensor]
            A tuple containing:
                - A tensor of atom encodings, shape (n_atoms,).
                - A tensor of bond encodings, shape (n_bonds,).
                - A connectivity matrix as a tensor, shape (2, n_bonds).

        Raises
        ------
        ValueError
            If the provided SMILES string cannot be parsed by RDKit.
        """
        mol = rdkit.Chem.MolFromSmiles(smiles)
        if mol is None:
            raise ValueError(f"Unable to parse SMILES string: {smiles}")
        atoms = mol.GetAtoms()

        atom_reprs = [atom_to_str(atom) for atom in atoms]
        enc_atoms = torch.tensor([
            self._atom_tokenizer(atom) for atom in atom_reprs
        ]).type(torch.int)

        bond_reprs = np.concatenate([
            [bond_to_str(bond) for bond in atom.GetBonds()]
            for atom in atoms
        ])
        enc_bonds = torch.tensor([
            self._bond_tokenizer(bond) for bond in bond_reprs
        ]).type(torch.int)

        connectivity = np.zeros((2, 2 * mol.GetNumBonds()))
        bond_index = 0
        for atom in atoms:
            start_idx = atom.GetIdx()
            for bond in atom.GetBonds():
                if bond.GetBeginAtomIdx() == start_idx:
                    connectivity[0, bond_index] = bond.GetBeginAtomIdx()
                    connectivity[1, bond_index] = bond.GetEndAtomIdx()
                else:
                    connectivity[0, bond_index] = bond.GetEndAtomIdx()
                    connectivity[1, bond_index] = bond.GetBeginAtomIdx()
                bond_index += 1
        connectivity = torch.tensor(connectivity).type(torch.long)

        return enc_atoms, enc_bonds, connectivity

    def encode_many(
            self,
            smiles: Iterable[str]
         ) -> List[Tuple[torch.Tensor, torch.Tensor, torch.Tensor]]:
        """
        Encodes a list of SMILES strings into tensors representing atoms,
        bonds, and connectivities.

        Parameters
        ----------
        smiles : Iterable[str]
            An iterable collection of SMILES strings representing molecules to
            be encoded.

        Returns
        -------
        List[Tuple[torch.Tensor, torch.Tensor, torch.Tensor]]
            A list containing tuples with three elements:
                - A tensor of atom encodings, shape (n_atoms,).
                - A tensor of bond encodings, shape (n_bonds,).
                - A connectivity matrix as a tensor, shape (2, n_bonds).

        Raises
        ------
        ValueError
            If any provided SMILES string cannot be parsed by RDKit.
        """
        encodings = []
        for smi in smiles:
            encodings.append(self.encode(smi))
        return encodings

    def save(self, filename: str) -> None:
        """
        Save the encoder to a file.

        Parameters
        ----------
        filename : str
            filename/path to save the encoder to.
        """
        with open(filename, "wb") as outp:
            pickle.dump(self, outp, pickle.HIGHEST_PROTOCOL)

vocab_sizes property

Returns the vocabulary sizes of the atom and bond tokenizers.

Returns

Tuple[int, int] A tuple containing two integers representing the sizes of the atom and bond tokenizers' vocabularies respectively.

__init__(smiles)

Initializes the MoleculeEncoder with a list of SMILES strings and creates/trains integer tokenizers for atoms and bonds.

Parameters

smiles : List[str] A list of SMILES strings representing molecules used for tokenizer creation/training.

Raises

ValueError If any provided SMILES string cannot be parsed by RDKit.

Source code in graphchem/preprocessing/features.py

def __init__(self, smiles: List[str]):
    """
    Initializes the MoleculeEncoder with a list of SMILES strings and
    creates/trains integer tokenizers for atoms and bonds.

    Parameters
    ----------
    smiles : List[str]
        A list of SMILES strings representing molecules used for tokenizer
        creation/training.

    Raises
    ------
    ValueError
        If any provided SMILES string cannot be parsed by RDKit.
    """
    mols = [Chem.MolFromSmiles(smi) for smi in smiles]
    for idx, mol in enumerate(mols):
        if mol is None:
            raise ValueError(f"Unable to parse SMILES: {smiles[idx]}")

    atoms = np.concatenate([mol.GetAtoms() for mol in mols])
    atom_reprs = [atom_to_str(atom) for atom in atoms]
    bond_reprs = np.concatenate([
        [bond_to_str(bond) for bond in atom.GetBonds()]
        for atom in atoms
    ])

    self._atom_tokenizer = Tokenizer()
    for rep in atom_reprs:
        self._atom_tokenizer(rep)
    self._atom_tokenizer.train = False

    self._bond_tokenizer = Tokenizer()
    for rep in bond_reprs:
        self._bond_tokenizer(rep)
    self._bond_tokenizer.train = False

encode(smiles)

Encodes a single SMILES string into three tensors representing atoms, bonds, and connectivity.

Parameters

smiles : str A SMILES string representing the molecule to be encoded.

Returns

Tuple[torch.Tensor, torch.Tensor, torch.Tensor] A tuple containing: - A tensor of atom encodings, shape (n_atoms,). - A tensor of bond encodings, shape (n_bonds,). - A connectivity matrix as a tensor, shape (2, n_bonds).

Raises

ValueError If the provided SMILES string cannot be parsed by RDKit.

Source code in graphchem/preprocessing/features.py

def encode(
        self,
        smiles: str
     ) -> Tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
    """
    Encodes a single SMILES string into three tensors representing atoms,
    bonds, and connectivity.

    Parameters
    ----------
    smiles : str
        A SMILES string representing the molecule to be encoded.

    Returns
    -------
    Tuple[torch.Tensor, torch.Tensor, torch.Tensor]
        A tuple containing:
            - A tensor of atom encodings, shape (n_atoms,).
            - A tensor of bond encodings, shape (n_bonds,).
            - A connectivity matrix as a tensor, shape (2, n_bonds).

    Raises
    ------
    ValueError
        If the provided SMILES string cannot be parsed by RDKit.
    """
    mol = rdkit.Chem.MolFromSmiles(smiles)
    if mol is None:
        raise ValueError(f"Unable to parse SMILES string: {smiles}")
    atoms = mol.GetAtoms()

    atom_reprs = [atom_to_str(atom) for atom in atoms]
    enc_atoms = torch.tensor([
        self._atom_tokenizer(atom) for atom in atom_reprs
    ]).type(torch.int)

    bond_reprs = np.concatenate([
        [bond_to_str(bond) for bond in atom.GetBonds()]
        for atom in atoms
    ])
    enc_bonds = torch.tensor([
        self._bond_tokenizer(bond) for bond in bond_reprs
    ]).type(torch.int)

    connectivity = np.zeros((2, 2 * mol.GetNumBonds()))
    bond_index = 0
    for atom in atoms:
        start_idx = atom.GetIdx()
        for bond in atom.GetBonds():
            if bond.GetBeginAtomIdx() == start_idx:
                connectivity[0, bond_index] = bond.GetBeginAtomIdx()
                connectivity[1, bond_index] = bond.GetEndAtomIdx()
            else:
                connectivity[0, bond_index] = bond.GetEndAtomIdx()
                connectivity[1, bond_index] = bond.GetBeginAtomIdx()
            bond_index += 1
    connectivity = torch.tensor(connectivity).type(torch.long)

    return enc_atoms, enc_bonds, connectivity

encode_many(smiles)

Encodes a list of SMILES strings into tensors representing atoms, bonds, and connectivities.

Parameters

smiles : Iterable[str] An iterable collection of SMILES strings representing molecules to be encoded.

Returns

List[Tuple[torch.Tensor, torch.Tensor, torch.Tensor]] A list containing tuples with three elements: - A tensor of atom encodings, shape (n_atoms,). - A tensor of bond encodings, shape (n_bonds,). - A connectivity matrix as a tensor, shape (2, n_bonds).

Raises

ValueError If any provided SMILES string cannot be parsed by RDKit.

Source code in graphchem/preprocessing/features.py

def encode_many(
        self,
        smiles: Iterable[str]
     ) -> List[Tuple[torch.Tensor, torch.Tensor, torch.Tensor]]:
    """
    Encodes a list of SMILES strings into tensors representing atoms,
    bonds, and connectivities.

    Parameters
    ----------
    smiles : Iterable[str]
        An iterable collection of SMILES strings representing molecules to
        be encoded.

    Returns
    -------
    List[Tuple[torch.Tensor, torch.Tensor, torch.Tensor]]
        A list containing tuples with three elements:
            - A tensor of atom encodings, shape (n_atoms,).
            - A tensor of bond encodings, shape (n_bonds,).
            - A connectivity matrix as a tensor, shape (2, n_bonds).

    Raises
    ------
    ValueError
        If any provided SMILES string cannot be parsed by RDKit.
    """
    encodings = []
    for smi in smiles:
        encodings.append(self.encode(smi))
    return encodings

save(filename)

Save the encoder to a file.

Parameters

filename : str filename/path to save the encoder to.

Source code in graphchem/preprocessing/features.py

def save(self, filename: str) -> None:
    """
    Save the encoder to a file.

    Parameters
    ----------
    filename : str
        filename/path to save the encoder to.
    """
    with open(filename, "wb") as outp:
        pickle.dump(self, outp, pickle.HIGHEST_PROTOCOL)