synplan/ml/training/preprocessing.py

"""Module containing functions for preparation of the training sets for policy and value
network."""

import logging
import os
import pickle
from abc import ABC
from typing import Any, Dict, List, Optional, Tuple

import ray
import torch
from CGRtools import smiles
from CGRtools.containers import MoleculeContainer
from CGRtools.exceptions import InvalidAromaticRing
from CGRtools.reactor import Reactor
from ray.util.queue import Empty, Queue
from torch import Tensor
from torch_geometric.data import InMemoryDataset
from torch_geometric.data.data import Data
from torch_geometric.data.makedirs import makedirs
from torch_geometric.transforms import ToUndirected
from tqdm import tqdm

from synplan.chem.utils import unite_molecules
from synplan.utils.files import ReactionReader
from synplan.utils.loading import load_reaction_rules


class ValueNetworkDataset(InMemoryDataset, ABC):
    """Value network dataset."""

    def __init__(self, extracted_precursor: Dict[str, float]) -> None:
        """Initializes a value network dataset object.

        :param extracted_precursor: The dictionary with the extracted from the built
            search trees precursor and their labels.
        """
        super().__init__(None, None, None)

        if extracted_precursor:
            self.data, self.slices = self.graphs_from_extracted_precursor(
                extracted_precursor
            )

    @staticmethod
    def mol_to_graph(molecule: MoleculeContainer, label: float) -> Optional[Data]:
        """Takes a molecule as input, and converts the molecule to a PyTorch geometric
        graph, assigns the reward value (label) to the graph, and returns the graph.

        :param molecule: The input molecule.
        :param label: The label (solved/unsolved routes in the tree) of the molecule
            (precursor).
        :return: A PyTorch Geometric graph representation of a molecule.
        """
        if len(molecule) > 2:
            pyg = mol_to_pyg(molecule)
            if pyg:
                pyg.y = torch.tensor([label])
                return pyg

        return None

    def graphs_from_extracted_precursor(
        self, extracted_precursor: Dict[str, float]
    ) -> Tuple[Data, Dict]:
        """Converts the extracted from the search trees precursor to the PyTorch geometric
        graphs.

        :param extracted_precursor: The dictionary with the extracted from the built
            search trees precursor and their labels.
        :return: The PyTorch geometric graphs and slices.
        """
        processed_data = []
        for smi, label in extracted_precursor.items():
            mol = smiles(smi)
            pyg = self.mol_to_graph(mol, label)
            if pyg:
                processed_data.append(pyg)
        data, slices = self.collate(processed_data)
        return data, slices


class RankingPolicyDataset(InMemoryDataset):
    """Ranking policy network dataset."""

    def __init__(self, reactions_path: str, reaction_rules_path: str, output_path: str):
        """Initializes a policy network dataset.

        :param reactions_path: The path to the file containing the reaction data used
            for extraction of reaction rules.
        :param reaction_rules_path: The path to the file containing the reaction rules.
        :param output_path: The output path to the file where policy network dataset
            will be saved.
        """
        super().__init__(None, None, None)

        self.reactions_path = reactions_path
        self.reaction_rules_path = reaction_rules_path
        self.output_path = output_path

        if output_path and os.path.exists(output_path):
            self.data, self.slices = torch.load(self.output_path)
        else:
            self.data, self.slices = self.prepare_data()

    @property
    def num_classes(self) -> int:
        return self._infer_num_classes(self._data.y_rules)

    def prepare_data(self) -> Tuple[Data, Dict[str, Tensor]]:
        """Prepares data by loading reaction rules, preprocessing the molecules,
        collating the data, and returning the data and slices.

        :return: The PyTorch geometric graphs and slices.
        """

        with open(self.reaction_rules_path, "rb") as inp:
            reaction_rules = pickle.load(inp)
        reaction_rules = sorted(reaction_rules, key=lambda x: len(x[1]), reverse=True)

        reaction_rule_pairs = {}
        for rule_i, (_, reactions_ids) in enumerate(reaction_rules):
            for reaction_id in reactions_ids:
                reaction_rule_pairs[reaction_id] = rule_i
        reaction_rule_pairs = dict(sorted(reaction_rule_pairs.items()))

        list_of_graphs = []
        with ReactionReader(self.reactions_path) as reactions:

            for reaction_id, reaction in tqdm(
                enumerate(reactions),
                desc="Number of reactions processed: ",
                bar_format="{desc}{n} [{elapsed}]",
            ):

                rule_id = reaction_rule_pairs.get(reaction_id)
                if rule_id:
                    try:  #  MENDEL_INFO does not contain cadmium (Cd) properties
                        molecule = unite_molecules(reaction.products)
                        pyg_graph = mol_to_pyg(molecule)

                    except (
                        Exception
                    ) as e:  # TypeError: can't assign a NoneType to a torch.ByteTensor
                        logging.debug(e)
                        continue

                    if pyg_graph is not None:
                        pyg_graph.y_rules = torch.tensor([rule_id], dtype=torch.long)
                        list_of_graphs.append(pyg_graph)
                else:
                    continue

        data, slices = self.collate(list_of_graphs)
        if self.output_path:
            makedirs(os.path.dirname(self.output_path))
            torch.save((data, slices), self.output_path)

        return data, slices


class FilteringPolicyDataset(InMemoryDataset):
    """Filtering policy network dataset."""

    def __init__(
        self,
        molecules_path: str,
        reaction_rules_path: str,
        output_path: str,
        num_cpus: int,
    ) -> None:
        """Initializes a policy network dataset object.

        :param molecules_path: The path to the file containing the molecules for
            reaction rule appliance.
        :param reaction_rules_path: The path to the file containing the reaction rules.
        :param output_path: The output path to the file where policy network dataset
            will be stored.
        :param num_cpus: The number of CPUs to be used for the dataset preparation.
        :return: None.
        """
        super().__init__(None, None, None)

        self.molecules_path = molecules_path
        self.reaction_rules_path = reaction_rules_path
        self.output_path = output_path
        self.num_cpus = num_cpus
        self.batch_size = 100

        if output_path and os.path.exists(output_path):
            self.data, self.slices = torch.load(self.output_path)
        else:
            self.data, self.slices = self.prepare_data()

    @property
    def num_classes(self) -> int:
        return self._data.y_rules.shape[1]

    def prepare_data(self) -> Tuple[Data, Dict]:
        """Prepares data by loading reaction rules, initializing Ray, preprocessing the
        molecules, collating the data, and returning the data and slices.

        :return: The PyTorch geometric graphs and slices.
        """

        ray.init(num_cpus=self.num_cpus, ignore_reinit_error=True)
        reaction_rules = load_reaction_rules(self.reaction_rules_path)
        reaction_rules_ids = ray.put(reaction_rules)

        to_process = Queue(maxsize=self.batch_size * self.num_cpus)
        processed_data = []
        results_ids = [
            preprocess_filtering_policy_molecules.remote(to_process, reaction_rules_ids)
            for _ in range(self.num_cpus)
        ]

        with open(self.molecules_path, "r", encoding="utf-8") as inp_data:
            for molecule in tqdm(
                inp_data.read().splitlines(),
                desc="Number of molecules processed: ",
                bar_format="{desc}{n} [{elapsed}]",
            ):

                to_process.put(molecule)

        results = [graph for res in ray.get(results_ids) if res for graph in res]
        processed_data.extend(results)

        ray.shutdown()

        for pyg in processed_data:
            pyg.y_rules = pyg.y_rules.to_dense()
            pyg.y_priority = pyg.y_priority.to_dense()

        data, slices = self.collate(processed_data)
        if self.output_path:
            makedirs(os.path.dirname(self.output_path))
            torch.save((data, slices), self.output_path)

        return data, slices


def reaction_rules_appliance(
    molecule: MoleculeContainer, reaction_rules: List[Reactor]
) -> Tuple[List[int], List[int]]:
    """Applies each reaction rule from the list of reaction rules to a given molecule
    and returns the indexes of the successfully applied regular and prioritized reaction
    rules.

    :param molecule: The input molecule.
    :param reaction_rules: The list of reaction rules.
    :return: The two lists of indexes of successfully applied regular reaction rules and
        priority reaction rules.
    """

    applied_rules, priority_rules = [], []
    for i, rule in enumerate(reaction_rules):

        rule_applied = False
        rule_prioritized = False

        try:
            for reaction in rule([molecule]):
                for prod in reaction.products:
                    prod.kekule()
                    if prod.check_valence():
                        break
                    rule_applied = True

                    # check priority rules
                    if len(reaction.products) > 1:
                        # check coupling retro manual
                        if all(len(mol) > 6 for mol in reaction.products):
                            if (
                                sum(len(mol) for mol in reaction.products)
                                - len(reaction.reactants[0])
                                < 6
                            ):
                                rule_prioritized = True
                    else:
                        # check cyclization retro manual
                        if sum(len(mol.sssr) for mol in reaction.products) < sum(
                            len(mol.sssr) for mol in reaction.reactants
                        ):
                            rule_prioritized = True
            #
            if rule_applied:
                applied_rules.append(i)
                #
                if rule_prioritized:
                    priority_rules.append(i)
        except Exception as e:
            logging.debug(e)
            continue

    return applied_rules, priority_rules


@ray.remote
def preprocess_filtering_policy_molecules(
    to_process: Queue, reaction_rules: List[Reactor]
) -> List[Optional[Data]]:
    """Preprocesses a list of molecules by applying reaction rules and converting
    molecules into PyTorch geometric graphs. Successfully applied reaction rules are
    converted to binary vectors for policy network training.

    :param to_process: The queue containing SMILES of molecules to be converted to the
        training data.
    :param reaction_rules: The list of reaction rules.
    :return: The list of PyGraph objects.
    """

    pyg_graphs = []
    while True:
        try:
            molecule = smiles(to_process.get(timeout=30))
            if not isinstance(molecule, MoleculeContainer):
                continue

            # reaction reaction_rules application
            applied_rules, priority_rules = reaction_rules_appliance(
                molecule, reaction_rules
            )

            y_rules = torch.sparse_coo_tensor(
                [applied_rules],
                torch.ones(len(applied_rules)),
                (len(reaction_rules),),
                dtype=torch.uint8,
            )
            y_priority = torch.sparse_coo_tensor(
                [priority_rules],
                torch.ones(len(priority_rules)),
                (len(reaction_rules),),
                dtype=torch.uint8,
            )

            y_rules = torch.unsqueeze(y_rules, 0)
            y_priority = torch.unsqueeze(y_priority, 0)

            pyg_graph = mol_to_pyg(molecule)
            if not pyg_graph:
                continue
            pyg_graph.y_rules = y_rules
            pyg_graph.y_priority = y_priority
            pyg_graphs.append(pyg_graph)

        except Empty:
            break

    return pyg_graphs


def atom_to_vector(atom: Any) -> Tensor:
    """Given an atom, return a vector of length 8 with the following
    information:

    1. Atomic number
    2. Period
    3. Group
    4. Number of electrons + atom's charge
    5. Shell
    6. Total number of hydrogens
    7. Whether the atom is in a ring
    8. Number of neighbors

    :param atom: The atom object.

    :return: The vector of the atom.
    """
    vector = torch.zeros(8, dtype=torch.uint8)
    period, group, shell, electrons = MENDEL_INFO[atom.atomic_symbol]
    vector[0] = atom.atomic_number
    vector[1] = period
    vector[2] = group
    vector[3] = electrons + atom.charge
    vector[4] = shell
    vector[5] = atom.total_hydrogens
    vector[6] = int(atom.in_ring)
    vector[7] = atom.neighbors
    return vector


def bonds_to_vector(molecule: MoleculeContainer, atom_ind: int) -> Tensor:
    """Takes a molecule and an atom index as input, and returns a vector representing
    the bond orders of the atom's bonds.

    :param molecule: The given molecule.
    :param atom_ind: The index of the atom in the molecule to be converted to the bond
        vector.
    :return: The torch tensor of size 3, with each element representing the order of
        bonds connected to the atom with the given index in the molecule.
    """

    vector = torch.zeros(3, dtype=torch.uint8)
    for b_order in molecule._bonds[atom_ind].values():
        vector[int(b_order) - 1] += 1
    return vector


def mol_to_matrix(molecule: MoleculeContainer) -> Tensor:
    """Given a molecule, it returns a vector of shape (max_atoms, 12) where each row is
    an atom and each column is a feature.

    :param molecule: The molecule to be converted to a vector
    :return: The atoms vectors array.
    """

    atoms_vectors = torch.zeros((len(molecule), 11), dtype=torch.uint8)
    for n, atom in molecule.atoms():
        atoms_vectors[n - 1][:8] = atom_to_vector(atom)
    for n, _ in molecule.atoms():
        atoms_vectors[n - 1][8:] = bonds_to_vector(molecule, n)

    return atoms_vectors


def mol_to_pyg(
    molecule: MoleculeContainer, canonicalize: bool = True
) -> Optional[Data]:
    """Takes a list of molecules and returns a list of PyTorch Geometric graphs, a one-
    hot encoded vectors of the atoms, and a matrices of the bonds.

    :param molecule: The molecule to be converted to PyTorch Geometric graph.
    :param canonicalize: If True, the input molecule is canonicalized.
    :return: The list of PyGraph objects.
    """

    if len(molecule) == 1:  # to avoid a precursor to be a single atom
        return None

    tmp_molecule = molecule.copy()
    try:
        if canonicalize:
            tmp_molecule.canonicalize()
        tmp_molecule.kekule()
        if tmp_molecule.check_valence():
            return None
    except InvalidAromaticRing:
        return None

    # remapping target for torch_geometric because
    # it is necessary that the elements in edge_index only hold nodes_idx in the range { 0, ..., num_nodes - 1}
    new_mappings = {n: i for i, (n, _) in enumerate(tmp_molecule.atoms(), 1)}
    tmp_molecule.remap(new_mappings)

    # get edge indexes from target mapping
    edge_index = []
    for atom, neighbour, bond in tmp_molecule.bonds():
        edge_index.append([atom - 1, neighbour - 1])
    edge_index = torch.tensor(edge_index, dtype=torch.long)

    #
    x = mol_to_matrix(tmp_molecule)

    mol_pyg_graph = Data(x=x, edge_index=edge_index.t().contiguous())
    mol_pyg_graph = ToUndirected()(mol_pyg_graph)

    assert mol_pyg_graph.is_undirected()

    return mol_pyg_graph


MENDEL_INFO = {
    "Ag": (5, 11, 1, 1),
    "Al": (3, 13, 2, 1),
    "Ar": (3, 18, 2, 6),
    "As": (4, 15, 2, 3),
    "B": (2, 13, 2, 1),
    "Ba": (6, 2, 1, 2),
    "Bi": (6, 15, 2, 3),
    "Br": (4, 17, 2, 5),
    "C": (2, 14, 2, 2),
    "Ca": (4, 2, 1, 2),
    "Ce": (6, None, 1, 2),
    "Cl": (3, 17, 2, 5),
    "Cr": (4, 6, 1, 1),
    "Cs": (6, 1, 1, 1),
    "Cu": (4, 11, 1, 1),
    "Dy": (6, None, 1, 2),
    "Er": (6, None, 1, 2),
    "F": (2, 17, 2, 5),
    "Fe": (4, 8, 1, 2),
    "Ga": (4, 13, 2, 1),
    "Gd": (6, None, 1, 2),
    "Ge": (4, 14, 2, 2),
    "Hg": (6, 12, 1, 2),
    "I": (5, 17, 2, 5),
    "In": (5, 13, 2, 1),
    "K": (4, 1, 1, 1),
    "La": (6, 3, 1, 2),
    "Li": (2, 1, 1, 1),
    "Mg": (3, 2, 1, 2),
    "Mn": (4, 7, 1, 2),
    "N": (2, 15, 2, 3),
    "Na": (3, 1, 1, 1),
    "Nd": (6, None, 1, 2),
    "O": (2, 16, 2, 4),
    "P": (3, 15, 2, 3),
    "Pb": (6, 14, 2, 2),
    "Pd": (5, 10, 3, 10),
    "Pr": (6, None, 1, 2),
    "Rb": (5, 1, 1, 1),
    "S": (3, 16, 2, 4),
    "Sb": (5, 15, 2, 3),
    "Se": (4, 16, 2, 4),
    "Si": (3, 14, 2, 2),
    "Sm": (6, None, 1, 2),
    "Sn": (5, 14, 2, 2),
    "Sr": (5, 2, 1, 2),
    "Te": (5, 16, 2, 4),
    "Ti": (4, 4, 1, 2),
    "Tl": (6, 13, 2, 1),
    "Yb": (6, None, 1, 2),
    "Zn": (4, 12, 1, 2),
}