PharmacoNet: Accelerating Large-Scale Virtual Screening by Deep Pharmacophore Modeling
import argparse
from src.modeling import ModelingModule


class Modeling_ArgParser(argparse.ArgumentParser):
    def __init__(self):
        super().__init__('pharmacophore modeling script')
        self.formatter_class = argparse.ArgumentDefaultsHelpFormatter

        # config
        cfg_args = self.add_argument_group('config')
        cfg_args.add_argument('--model_path', type=str, help='path of embedded model (.tar)', default='./weights/model.tar')
        cfg_args.add_argument('-r', '--receptor', type=str, help='path of receptor pdb file (.pdb)', required=True)
        cfg_args.add_argument('-p', '--pharmacophore_model', type=str, help='path to save pharmacophore model (.pkl)', required=True)

        cfg_args.add_argument('--autobox_ligand', type=str, help='path of ligand to define the center of box (.sdf, .pdb, .mol2)')
        cfg_args.add_argument('--center', nargs='+', type=float, help='coordinate of the center')

        # system config
        env_args = self.add_argument_group('environment')
        env_args.add_argument('--cuda', action='store_true', help='use gpu acceleration with CUDA')


def main(args):
    if args.autobox_ligand is not None:
        print(f'Using center of {args.autobox_ligand} as center of box')
    else:
        assert args.center is not None, \
            'No Center!. Enter the input `--autobox_ligand <LIGAND_PATH>` or `--center x y z`'
        assert len(args.center) == 3, \
            'Wrong Center!. The arguments for center coordinates should be 3. (ex. --center 1.00 2.00 -1.50)'
        print(f'Using center {tuple(args.center)}')
    module = ModelingModule(args.model_path, 'cuda' if args.cuda else 'cpu')

    if args.autobox_ligand is not None:
        pharmacophore_model = module.run(args.receptor, ref_ligand_path=args.autobox_ligand)
    else:
        pharmacophore_model = module.run(args.receptor, center=tuple(args.center))
    pharmacophore_model.save(args.pharmacophore_model)


if __name__ == '__main__':
    parser = Modeling_ArgParser()
    args = parser.parse_args()
    main(args)

modeling.py
SeonghwanSeo/PharmacoNet
import argparse
import numpy as np

from openbabel import pybel
from rdkit import Chem
from rdkit.Chem import AllChem

from src import PharmacophoreModel


class Scoring_ArgParser(argparse.ArgumentParser):
    def __init__(self):
        super().__init__('scoring')
        self.formatter_class = argparse.ArgumentDefaultsHelpFormatter
        self.add_argument('-p', '--pharmacophore_model', type=str, help='path of pharmacophore model (.pkl)', required=True)
        self.add_argument('-s', '--smiles', type=str, help='molecule SMILES', required=True)
        self.add_argument('--num_conformers', type=int, help='number of RDKit conformer to use', default=10)


def scoring(smiles: str, model: PharmacophoreModel, num_conformers: int = 10):
    pbmol = pybel.readstring('smi', smiles)

    # NOTE: Create Conformers
    rdmol = Chem.MolFromSmiles(smiles)
    rdmol = Chem.AddHs(rdmol)
    AllChem.EmbedMultipleConfs(rdmol, num_conformers, AllChem.srETKDGv3())
    assert rdmol.GetNumConformers() > 0

    # NOTE: Scoring
    if True:
        # NOTE: Scoring with RDKit Object (Using rdmol.GetConformers())
        return model.scoring(pbmol, rdmol)
    else:
        # NOTE: Scoring with custom ligand conformers.
        atom_positions = np.stack([conformer.GetPositions() for conformer in rdmol.GetConformers()], axis=0)
        return model.scoring(pbmol, atom_positions=atom_positions)


if __name__ == '__main__':
    parser = Scoring_ArgParser()
    args = parser.parse_args()
    model = PharmacophoreModel.load(args.pharmacophore_model)
    score = scoring(args.smiles, model, args.num_conformers)
    print(score)

scoring.py
SeonghwanSeo/PharmacoNet
import argparse

from openbabel import pybel
from rdkit import Chem
from rdkit.Chem import AllChem
import multiprocessing
import functools

from src.scoring import PharmacophoreModel


class Scoring_ArgParser(argparse.ArgumentParser):
    def __init__(self):
        super().__init__('scoring')
        self.formatter_class = argparse.ArgumentDefaultsHelpFormatter
        self.add_argument('-p', '--pharmacophore_model', type=str, help='path of pharmacophore model (.pkl)', required=True)
        self.add_argument('-s', '--smiles_path', type=str, help='molecules SMILES file', required=True)
        self.add_argument('--num_conformers', type=int, help='number of RDKit conformer to use', default=10)
        self.add_argument('--num_cpus', type=int, help='number of cpu cores. default: (try to detect the number of CPUs)')


def scoring(smiles: str, model: PharmacophoreModel, num_conformers: int = 10):
    try:
        pbmol = pybel.readstring('smi', smiles)

        # NOTE: Create Conformers
        rdmol = Chem.MolFromSmiles(smiles)
        rdmol = Chem.AddHs(rdmol)
        AllChem.EmbedMultipleConfs(rdmol, num_conformers, AllChem.srETKDGv3())
        assert rdmol.GetNumConformers() > 0

        # NOTE: Scoring
        return model.scoring(pbmol, rdmol)

    except KeyboardInterrupt:
        raise KeyboardInterrupt
    except Exception:
        return None


if __name__ == '__main__':
    parser = Scoring_ArgParser()
    args = parser.parse_args()
    model = PharmacophoreModel.load(args.pharmacophore_model)

    with open(args.smiles_path) as f:
        lines = f.readlines()
    key_list = [line.strip().split(',')[0] for line in lines]
    smiles_list = [line.strip().split(',')[1] for line in lines]

    num_cpus = args.num_cpus if args.num_cpus is not None else multiprocessing.cpu_count()
    partial_func = functools.partial(scoring, model=model, num_conformers=args.num_conformers)
    with multiprocessing.Pool(num_cpus) as p:
        scores = p.map(partial_func, smiles_list)
    for key, smiles, score in zip(key_list, smiles_list, scores):
        print(key, smiles, score)

scoring_file.py
SeonghwanSeo/PharmacoNet
from .scoring import PharmacophoreModel

src/__init__.py
SeonghwanSeo/PharmacoNet
from .constant import INTERACTION_LIST
from .pointcloud import PROTEIN_CHANNEL_LIST

src/data/__init__.py
SeonghwanSeo/PharmacoNet
from typing import Sequence, Set

INTERACTION_LIST: Sequence[str] = (
    'Hydrophobic',
    'PiStacking_P',
    'PiStacking_T',
    'PiCation_lring',
    'PiCation_pring',
    'HBond_ldon',
    'HBond_pdon',
    'SaltBridge_lneg',
    'SaltBridge_pneg',
    'XBond'
)

NUM_INTERACTION_TYPES: int = 10

HYDROPHOBIC = 0
PISTACKING_P = 1
PISTACKING_T = 2
PICATION_LRING = 3
PICATION_PRING = 4
HBOND_LDON = 5
HBOND_PDON = 6
SALTBRIDGE_LNEG = 7
SALTBRIDGE_PNEG = 8
XBOND = 9

# PLIP Distance + 0.5 A
INTERACTION_DIST = {
    HYDROPHOBIC: 4.5,       # 4.0 + 0.5
    PISTACKING_P: 6.0,      # 5.5 + 0.5
    PISTACKING_T: 6.0,      # 5.5 + 0.5
    PICATION_LRING: 6.5,    # 6.0 + 0.5
    PICATION_PRING: 6.5,    # 6.0 + 0.5
    HBOND_LDON: 4.5,        # 4.1 + 0.5 - 0.1 (to be devided to 0.5)
    HBOND_PDON: 4.5,        # 4.1 + 0.5 - 0.1
    SALTBRIDGE_LNEG: 6.0,   # 5.5 + 0.5
    SALTBRIDGE_PNEG: 6.0,   # 5.5 + 0.5
    XBOND: 4.5,             # 4.0 + 0.5
}

LONG_INTERACTION: Set[int] = {
    PISTACKING_P,
    PISTACKING_T,
    PICATION_PRING,
    PICATION_LRING,
    SALTBRIDGE_LNEG,
    SALTBRIDGE_PNEG
}

SHORT_INTERACTION: Set[int] = {
    HYDROPHOBIC,
    HBOND_LDON,
    HBOND_PDON,
    XBOND,
}

src/data/constant.py
SeonghwanSeo/PharmacoNet
import os
import numpy as np

from Bio.PDB import PDBParser, PDBIO
from Bio.PDB.PDBIO import Select

from numpy.typing import ArrayLike

import warnings
warnings.filterwarnings("ignore")

AMINO_ACID = [
    'GLY', 'ALA', 'VAL', 'LEU', 'ILE', 'PRO', 'PHE', 'TYR', 'TRP', 'SER',
    'THR', 'CYS', 'MET', 'ASN', 'GLN', 'ASP', 'GLU', 'LYS', 'ARG', 'HIS',
    'HIP', 'HIE', 'TPO', 'HID', 'LEV', 'MEU', 'PTR', 'GLV', 'CYT', 'SEP',
    'HIZ', 'CYM', 'GLM', 'ASQ', 'TYS', 'CYX', 'GLZ', 'MSE', 'CSO', 'KCX',
    'CSD', 'MLY', 'PCA', 'LLP'
]


class DistSelect(Select):
    def __init__(self, center, cutoff: float = 40.0):
        self.center = np.array(center).reshape(1, 3)
        self.cutoff = cutoff

    def accept_residue(self, residue):
        if super().accept_residue(residue) == 0:
            return 0
        if residue.get_resname() not in AMINO_ACID:
            return 0
        residue_positions = np.array([
            list(atom.get_vector())
            for atom in residue.get_atoms()
            if "H" not in atom.get_id()
        ])
        if residue_positions.shape[0] == 0:
            return 0
        min_dis = np.min(np.linalg.norm(residue_positions - self.center, axis=-1))
        if min_dis < self.cutoff:
            return 1
        else:
            return 0


def extract_pocket(
    protein_pdb_path: str,
    out_pocket_pdb_path: str,
    center: ArrayLike,
    cutoff: float
):
    parser = PDBParser()
    structure = parser.get_structure("protein", protein_pdb_path)
    io = PDBIO()
    io.set_structure(structure)
    io.save(out_pocket_pdb_path, DistSelect(center, cutoff))
    command = f"obabel {out_pocket_pdb_path} -O {out_pocket_pdb_path} -d 2>/dev/null"
    os.system(command)

src/data/extract_pocket.py
SeonghwanSeo/PharmacoNet
from .objects import Protein

src/data/objects/__init__.py
SeonghwanSeo/PharmacoNet
from dataclasses import dataclass, field
from openbabel.pybel import ob
import numpy as np

from collections.abc import Sequence
from typing import Sequence, List, Tuple
from numpy.typing import NDArray
from functools import cached_property

from . import utils

Tuple3D = Tuple[float, float, float]


@dataclass
class Point3D(Sequence):
    x: float
    y: float
    z: float

    @classmethod
    def from_obatom(cls, obatom: ob.OBAtom):
        return cls(*utils.ob_coords(obatom))

    @classmethod
    def from_array(cls, array: NDArray):
        x, y, z = array
        return cls(x, y, z)

    def __array__(self):
        return np.array((self.x, self.y, self.z))

    def __iter__(self):
        yield self.x
        yield self.y
        yield self.z

    def __getitem__(self, idx: int):
        if idx == 0:
            return self.x
        elif idx == 1:
            return self.y
        elif idx == 2:
            return self.z
        raise ValueError

    def __len__(self):
        return 3


@dataclass
class BaseInteractablePart():

    @property
    def small(self):
        if self.__small is None:
            self.__small = self.to_small()
        return self.__small

    def to_small(self):
        raise NotImplemented


@dataclass
class BaseHydrophobicAtom(BaseInteractablePart):
    obatom: ob.OBAtom
    coords: Point3D = field(init=False)

    def __post_init__(self):
        self.coords = Point3D.from_obatom(self.obatom)

    @property
    def index(self) -> int:
        return self.obatom.GetIdx() - 1


@dataclass
class BaseHBondAcceptor(BaseInteractablePart):
    obatom: ob.OBAtom
    coords: Point3D = field(init=False)

    def __post_init__(self):
        self.coords = Point3D.from_obatom(self.obatom)

    @property
    def index(self) -> int:
        return self.obatom.GetIdx() - 1


@dataclass
class BaseHBondDonor(BaseInteractablePart):
    obatom: ob.OBAtom
    coords: Point3D = field(init=False)
    hydrogens: Sequence[ob.OBAtom] = field(init=False)
    hydrogen_coords_list: Sequence[Point3D] = field(init=False)

    def __post_init__(self):
        self.coords = Point3D.from_obatom(self.obatom)
        hydrogens = [neigh for neigh in ob.OBAtomAtomIter(self.obatom)
                     if neigh.GetAtomicNum() == 1
                     ]
        self.hydrogens = hydrogens
        self.hydrogen_coords_list = [Point3D.from_obatom(h) for h in hydrogens]

    @property
    def index(self) -> int:
        return self.obatom.GetIdx() - 1


@dataclass
class BaseRing(BaseInteractablePart):
    obatoms: Sequence[ob.OBAtom]
    center: Point3D = field(init=False)
    normal: NDArray = field(init=False)

    def __post_init__(self):
        coords_list = np.array(
            [utils.ob_coords(obatom) for obatom in self.obatoms]
        )
        self.center = Point3D.from_array(np.mean(coords_list, axis=0))
        p1, p2, p3 = coords_list[0], coords_list[2], coords_list[4]
        v1, v2 = utils.vector(p1, p2), utils.vector(p1, p3)
        self.normal = utils.normalize(np.cross(v1, v2))

    @property
    def indices(self) -> List[int]:
        return [obatom.GetIdx() - 1 for obatom in self.obatoms]


@dataclass
class BaseCharged(BaseInteractablePart):
    obatoms: Sequence[ob.OBAtom]
    center: Point3D = None

    def __post_init__(self):
        if self.center is None:
            if len(self.obatoms) == 1:
                self.center = Point3D.from_obatom(self.obatoms[0])
            else:
                coords_list = np.array([utils.ob_coords(obatom) for obatom in self.obatoms])
                self.center = Point3D.from_array(np.mean(coords_list, axis=0))

    @property
    def indices(self) -> List[int]:
        return [obatom.GetIdx() - 1 for obatom in self.obatoms]


@dataclass
class BasePosCharged(BaseCharged):
    pass


@dataclass
class BaseNegCharged(BaseCharged):
    pass


@dataclass
class BaseXBondDonor(BaseInteractablePart):
    X: ob.OBAtom
    C: ob.OBAtom = field(init=False)
    X_coords: Point3D = field(init=False)
    C_coords: Point3D = field(init=False)

    def __post_init__(self):
        for neigh in ob.OBAtomAtomIter(self.X):
            if neigh.GetAtomicNum() == 6:
                self.C = neigh
                break
        assert self.C is not None
        self.X_coords = Point3D.from_obatom(self.X)
        self.C_coords = Point3D.from_obatom(self.C)

    @property
    def X_index(self) -> int:
        return self.X.GetIdx() - 1

    @property
    def C_index(self) -> int:
        return self.C.GetIdx() - 1

    @property
    def indices(self) -> List[int]:
        return [self.X_index, self.C_index]


@dataclass
class BaseXBondAcceptor(BaseInteractablePart):
    O: ob.OBAtom
    Y: ob.OBAtom
    O_coords: Point3D = field(init=False)
    Y_coords: Point3D = field(init=False)

    def __post_init__(self):
        self.O_coords = Point3D.from_obatom(self.O)
        self.Y_coords = Point3D.from_obatom(self.Y)

    @property
    def O_index(self) -> int:
        return self.O.GetIdx() - 1

    @property
    def Y_index(self) -> int:
        return self.Y.GetIdx() - 1

    @property
    def indices(self) -> List[int]:
        return [self.O_index, self.Y_index]


""" PROTEIN """


class ProteinAtom:
    @cached_property
    def obresidue_index(self) -> int:
        return self.obresidue.GetIdx() - 1

    @cached_property
    def obresidue_name(self) -> str:
        return self.obresidue.GetName()

    @cached_property
    def obresidue(self) -> ob.OBResidue:
        if hasattr(self, "obatom"):
            obatom = self.obatom
        elif hasattr(self, "obatoms"):
            obatom = self.obatoms[0]
        else:
            obatom = self.Y
        return obatom.GetResidue()


@dataclass
class HydrophobicAtom_P(ProteinAtom, BaseHydrophobicAtom):
    pass


@dataclass
class HBondAcceptor_P(ProteinAtom, BaseHBondAcceptor):
    pass


@dataclass
class HBondDonor_P(ProteinAtom, BaseHBondDonor):
    pass


@dataclass
class Ring_P(ProteinAtom, BaseRing):
    pass


@dataclass
class PosCharged_P(ProteinAtom, BasePosCharged):
    pass


@dataclass
class NegCharged_P(ProteinAtom, BaseNegCharged):
    pass


@dataclass
class XBondAcceptor_P(ProteinAtom, BaseXBondAcceptor):
    pass

src/data/objects/atom_classes.py
SeonghwanSeo/PharmacoNet
from __future__ import annotations
from openbabel import pybel
from openbabel.pybel import ob

from typing import List, Tuple

from .atom_classes import (
    HydrophobicAtom_P,
    HBondAcceptor_P,
    HBondDonor_P,
    Ring_P,
    PosCharged_P,
    NegCharged_P,
    XBondAcceptor_P,
)


class Protein():
    def __init__(
            self,
            pbmol: pybel.Molecule,
            addh: bool = True,
    ):
        """
        pbmol: Pybel Mol
        addh: if True, call OBMol.AddPolarHydrogens()
        setup: if True, find interactable parts
        ligand: if ligand is not None, extract pocket
        """

        self.addh: bool = addh

        self.pbmol = pbmol.clone
        self.pbmol.removeh()
        self.obmol = self.pbmol.OBMol
        self.obatoms: List[ob.OBAtom] = list(ob.OBMolAtomIter(self.obmol))
        self.num_heavyatoms = len(self.obatoms)

        self.pbmol_hyd: pybel.Molecule
        if addh:
            self.pbmol_hyd = self.pbmol.clone
            self.pbmol_hyd.OBMol.AddPolarHydrogens()
        else:
            self.pbmol_hyd = pbmol
        self.obmol_hyd = self.pbmol_hyd.OBMol
        self.obatoms_hyd: List[ob.OBAtom] = list(ob.OBMolAtomIter(self.obmol_hyd))[:self.num_heavyatoms]
        self.obatoms_hyd_nonwater: List[ob.OBAtom] = [
            obatom for obatom in self.obatoms_hyd
            if obatom.GetResidue().GetName() != 'HOH'
            and obatom.GetAtomicNum() in [6, 7, 8, 16]
        ]
        self.obresidues_hyd: List[ob.OBResidue] = list(ob.OBResidueIter(self.obmol_hyd))

        self.hydrophobic_atoms_all: List[HydrophobicAtom_P]
        self.rings_all: List[Ring_P]
        self.pos_charged_atoms_all: List[PosCharged_P]
        self.neg_charged_atoms_all: List[NegCharged_P]
        self.hbond_donors_all: List[HBondDonor_P]
        self.hbond_acceptors_all: List[HBondAcceptor_P]
        self.xbond_acceptors_all: List[XBondAcceptor_P]

        self.hydrophobic_atoms_all = self.__find_hydrophobic_atoms()
        self.rings_all = self.__find_rings()
        self.pos_charged_atoms_all, self.neg_charged_atoms_all = self.__find_charged_atoms()
        self.hbond_donors_all = self.__find_hbond_donors()
        self.hbond_acceptors_all = self.__find_hbond_acceptors()
        self.xbond_acceptors_all = self.__find_xbond_acceptors()

    @classmethod
    def from_pdbfile(cls, path, addh=True, **kwargs):
        pbmol = next(pybel.readfile('pdb', path))
        return cls(pbmol, addh, **kwargs)

    # Search Interactable Part
    def __find_hydrophobic_atoms(self) -> List[HydrophobicAtom_P]:
        hydrophobics = [HydrophobicAtom_P(obatom) for obatom in self.obatoms_hyd_nonwater
                        if obatom.GetAtomicNum() == 6
                        and all(neigh.GetAtomicNum() in (1, 6) for neigh in ob.OBAtomAtomIter(obatom))
                        ]
        return hydrophobics

    def __find_hbond_acceptors(self) -> List[HBondAcceptor_P]:
        acceptors = [HBondAcceptor_P(obatom) for obatom in self.obatoms_hyd_nonwater
                     if obatom.IsHbondAcceptor()
                     ]
        return acceptors

    def __find_hbond_donors(self) -> List[HBondDonor_P]:
        donors = [HBondDonor_P(obatom) for obatom in self.obatoms_hyd_nonwater
                  if obatom.IsHbondDonor()
                  ]
        return donors

    def __find_rings(self) -> List[Ring_P]:
        rings = []
        ring_candidates = self.pbmol_hyd.sssr
        for ring in ring_candidates:
            if not 4 < len(ring._path) <= 6:
                continue
            obatoms = [self.obatoms_hyd[idx - 1] for idx in sorted(ring._path)]
            residue = obatoms[0].GetResidue()
            if residue.GetName() not in ["TYR", "TRP", "HIS", "PHE"]:
                continue
            rings.append(Ring_P(obatoms))
        return rings

    def __find_charged_atoms(self) -> Tuple[List[PosCharged_P], List[NegCharged_P]]:
        pos_charged = []
        neg_charged = []

        for obresidue in self.obresidues_hyd:
            obresname = obresidue.GetName()
            if obresname in ("ARG", "HIS", "LYS"):
                obatoms = [obatom for obatom in ob.OBResidueAtomIter(obresidue)
                           if obatom.GetAtomicNum() == 7
                           and obresidue.GetAtomProperty(obatom, ob.SIDECHAIN)
                           ]
                if len(obatoms) > 0:
                    pos_charged.append(PosCharged_P(obatoms))

            elif obresname in ("GLU", "ASP"):
                obatoms = [obatom for obatom in ob.OBResidueAtomIter(obresidue)
                           if obatom.GetAtomicNum() == 8
                           and obresidue.GetAtomProperty(obatom, ob.SIDECHAIN)
                           ]
                if len(obatoms) > 0:
                    neg_charged.append(NegCharged_P(obatoms))

        return pos_charged, neg_charged

    def __find_xbond_acceptors(self) -> List[XBondAcceptor_P]:
        """Look for halogen bond acceptors (Y-{O|N|S}, with Y=N,C)"""
        acceptors = []
        for obatom in self.obatoms_hyd_nonwater:
            if obatom.GetAtomicNum() not in [8, 7, 16]:
                continue
            neighbors = [neigh for neigh in ob.OBAtomAtomIter(obatom)
                         if neigh.GetAtomicNum() in [6, 7, 16]
                         ]
            if len(neighbors) == 1:
                O, Y = obatom, neighbors[0]
                acceptors.append(XBondAcceptor_P(O, Y))
        return acceptors

src/data/objects/objects.py
SeonghwanSeo/PharmacoNet
import numpy as np
import math

from openbabel.pybel import ob

from typing import Sequence, Tuple, Union
from numpy.typing import NDArray


def check_in_cutoff(coords, neighbor_coords_list, cutoff: float):
    """
    coords: (3,)
    neighbor_coords: (N, 3)
    cutoff: scalar
    """
    x1, y1, z1 = coords
    cutoff_square = cutoff**2
    for neighbor_coords in neighbor_coords_list:
        x2, y2, z2 = neighbor_coords
        distance_sq = (x1 - x2) ** 2 + (y1 - y2) ** 2 + (z1 - z2) ** 2
        if distance_sq < cutoff_square:
            return True
    return False


def angle_btw_vectors(vec1: NDArray, vec2: NDArray, degree=True, normalized=False) -> float:
    if np.array_equal(vec1, vec2):
        return 0.0
    if normalized:
        cosval = np.dot(vec1, vec2)
    else:
        cosval = np.dot(vec1, vec2) / (np.linalg.norm(vec1) * np.linalg.norm(vec2))
    angle = math.acos(np.clip(cosval, -1, 1))
    return math.degrees(angle) if degree else angle


def vector(p1: Union[Sequence[float], NDArray], p2: Union[Sequence[float], NDArray]) -> NDArray:
    return np.subtract(p2, p1)


def euclidean3d(p1: Union[Sequence[float], NDArray], p2: Union[Sequence[float], NDArray]) -> float:
    return math.sqrt(sum([(a - b) ** 2 for a, b in zip(p1, p2)]))


def normalize(vec: NDArray) -> NDArray:
    norm = np.linalg.norm(vec)
    assert norm > 0, "vector size is zero"
    return vec / norm


def projection(point: Union[Sequence[float], NDArray], origin: Union[Sequence[float], NDArray],
               normal: NDArray) -> NDArray:
    """
    point: point to be projected
    normal, orig: normal vector & origin of projection plane
    """
    c = np.dot(normal, np.subtract(point, origin))
    return np.subtract(point, c * normal)


def ob_coords(obatom: ob.OBAtom) -> Tuple[float, float, float]:
    return (obatom.x(), obatom.y(), obatom.z())

src/data/objects/utils.py
SeonghwanSeo/PharmacoNet
import numpy as np
from typing import Tuple, Sequence
from openbabel.pybel import ob
from numpy.typing import NDArray

from .objects import Protein


protein_atom_num_list = (6, 7, 8, 16, -1)
protein_atom_symbol_list = ('C', 'N', 'O', 'S', 'UNK_ATOM')
protein_aminoacid_list = (
    'GLY', 'ALA', 'VAL', 'LEU', 'ILE', 'PRO', 'PHE', 'TYR', 'TRP', 'SER',
    'THR', 'CYS', 'MET', 'ASN', 'GLN', 'ASP', 'GLU', 'LYS', 'ARG', 'HIS',
    'UNK_AA',
)
protein_interactable_list = ('HydrophobicAtom', 'Ring', 'HBondDonor', 'HBondAcceptor', 'Cation', 'Anion', 'XBondAcceptor')

NUM_PROTEIN_ATOMIC_NUM = len(protein_atom_num_list)
NUM_PROTEIN_AMINOACID_NUM = len(protein_aminoacid_list)
NUM_PROTEIN_INTERACTABLE_NUM = len(protein_interactable_list)

PROTEIN_CHANNEL_LIST: Sequence[str] = protein_atom_symbol_list + protein_aminoacid_list + protein_interactable_list
NUM_PROTEIN_CHANNEL = len(PROTEIN_CHANNEL_LIST)


def get_position(obatom: ob.OBAtom) -> Tuple[float, float, float]:
    return (obatom.x(), obatom.y(), obatom.z())


def protein_atom_function(atom: ob.OBAtom, out: NDArray, **kwargs) -> NDArray[np.float32]:
    atomicnum = atom.GetAtomicNum()
    if atomicnum in protein_atom_num_list:
        out[protein_atom_num_list.index(atomicnum)] = 1
    else:
        out[NUM_PROTEIN_ATOMIC_NUM - 1] = 1
    residue_type = atom.GetResidue().GetName()
    if residue_type in protein_aminoacid_list:
        out[NUM_PROTEIN_ATOMIC_NUM + protein_aminoacid_list.index(residue_type)] = 1
    else:
        out[NUM_PROTEIN_ATOMIC_NUM + NUM_PROTEIN_AMINOACID_NUM - 1] = 1
    return out


def get_protein_pointcloud(pocket_obj: Protein) -> Tuple[NDArray[np.float32], NDArray[np.float32]]:
    positions = np.array(
        [(obatom.x(), obatom.y(), obatom.z()) for obatom in pocket_obj.obatoms],
        dtype=np.float32
    )

    channels = np.zeros((pocket_obj.num_heavyatoms, NUM_PROTEIN_CHANNEL), dtype=np.float32)
    for i, atom in enumerate(pocket_obj.obatoms):
        protein_atom_function(atom, channels[i])

    offset = NUM_PROTEIN_ATOMIC_NUM + NUM_PROTEIN_AMINOACID_NUM
    for hydrop in pocket_obj.hydrophobic_atoms_all:
        channels[hydrop.index, offset] = 1
    for ring in pocket_obj.rings_all:
        channels[ring.indices, offset + 1] = 1
    for donor in pocket_obj.hbond_donors_all:
        channels[donor.index, offset + 2] = 1
    for acceptor in pocket_obj.hbond_acceptors_all:
        channels[acceptor.index, offset + 3] = 1
    for cation in pocket_obj.pos_charged_atoms_all:
        channels[cation.indices, offset + 4] = 1
    for anion in pocket_obj.neg_charged_atoms_all:
        channels[anion.indices, offset + 5] = 1
    for acceptor in pocket_obj.xbond_acceptors_all:
        channels[acceptor.indices, offset + 6] = 1
    return positions, channels

src/data/pointcloud.py
SeonghwanSeo/PharmacoNet
import numpy as np
import math

from typing import Tuple, List
from numpy.typing import NDArray, ArrayLike

from .objects import Protein
from . import constant as C


def get_token_informations(
    protein_obj: Protein,
) -> Tuple[NDArray[np.float32], NDArray[np.int16]]:
    """get token information

    Args:
        protein_obj: Union[Protein]

    Returns:
        token_positions: [float, (N, 3)] token center positions
        token_classes: [int, (N,)] token interaction type
    """
    num_tokens = \
        len(protein_obj.hydrophobic_atoms_all) + \
        len(protein_obj.rings_all) * 3 + \
        len(protein_obj.hbond_donors_all) + \
        len(protein_obj.hbond_acceptors_all) + \
        len(protein_obj.pos_charged_atoms_all) * 2 + \
        len(protein_obj.neg_charged_atoms_all) + \
        len(protein_obj.xbond_acceptors_all)

    positions: List[Tuple[float, float, float]] = []
    classes: List[int] = []

    # NOTE: Hydrophobic
    positions.extend(tuple(hydrop.coords) for hydrop in protein_obj.hydrophobic_atoms_all)
    classes.extend([C.HYDROPHOBIC] * len(protein_obj.hydrophobic_atoms_all))

    # NOTE: PiStacking_P
    positions.extend(tuple(ring.center) for ring in protein_obj.rings_all)
    classes.extend([C.PISTACKING_P] * len(protein_obj.rings_all))

    # NOTE: PiStacking_T
    positions.extend(tuple(ring.center) for ring in protein_obj.rings_all)
    classes.extend([C.PISTACKING_T] * len(protein_obj.rings_all))

    # NOTE: PiCation_lring
    positions.extend(tuple(cation.center) for cation in protein_obj.pos_charged_atoms_all)
    classes.extend([C.PICATION_LRING] * len(protein_obj.pos_charged_atoms_all))

    # NOTE: PiCation_pring
    positions.extend(tuple(ring.center) for ring in protein_obj.rings_all)
    classes.extend([C.PICATION_PRING] * len(protein_obj.rings_all))

    # NOTE: HBond_ldon
    positions.extend(tuple(acceptor.coords) for acceptor in protein_obj.hbond_acceptors_all)
    classes.extend([C.HBOND_LDON] * len(protein_obj.hbond_acceptors_all))

    # NOTE: HBond_pdon
    positions.extend(tuple(donor.coords) for donor in protein_obj.hbond_donors_all)
    classes.extend([C.HBOND_PDON] * len(protein_obj.hbond_donors_all))

    # NOTE: Saltbridge_lneg
    positions.extend(tuple(cation.center) for cation in protein_obj.pos_charged_atoms_all)
    classes.extend([C.SALTBRIDGE_LNEG] * len(protein_obj.pos_charged_atoms_all))

    # NOTE: Saltbridge_pneg
    positions.extend(tuple(anion.center) for anion in protein_obj.neg_charged_atoms_all)
    classes.extend([C.SALTBRIDGE_PNEG] * len(protein_obj.neg_charged_atoms_all))

    # NOTE: XBond
    positions.extend(tuple(acceptor.O_coords) for acceptor in protein_obj.xbond_acceptors_all)
    classes.extend([C.XBOND] * len(protein_obj.xbond_acceptors_all))

    assert len(positions) == len(classes) == num_tokens
    return (
        np.array(positions, dtype=np.float32),
        np.array(classes, dtype=np.int16),
    )


def get_token_and_filter(
    positions: NDArray[np.float32],
    classes: NDArray[np.int16],
    center: NDArray[np.float32],
    resolution: float,
    dimension: int,
) -> Tuple[NDArray[np.int16], NDArray[np.int16]]:
    """Create token and Filtering valid instances

    Args:
        positions: [float, (N, 3)] token center positions
        classes: [int, (N,)] token interaction type
        center: [float, (3,)] voxel image center
        resolution: voxel image resolution
        dimension: voxel imzge dimension (size)

    Returns:
        token: [int, (N_token, 4)]
        filter: [int, (N_token,)]
    """
    filter = []
    tokens = []
    x_center, y_center, z_center = center
    x_start = x_center - (dimension / 2) * resolution
    y_start = y_center - (dimension / 2) * resolution
    z_start = z_center - (dimension / 2) * resolution
    for i, ((x, y, z), c) in enumerate(zip(positions, classes)):
        _x = int((x - x_start) // resolution)
        _y = int((y - y_start) // resolution)
        _z = int((z - z_start) // resolution)
        if (0 <= _x < dimension) and (0 <= _y < dimension) and (0 <= _z < dimension):
            filter.append(i)
            tokens.append((_x, _y, _z, c))

    return np.array(tokens, dtype=np.int16), np.array(filter, dtype=np.int16)


def get_box_area(
    tokens: ArrayLike,
    pharmacophore_size: float,
    resolution: float,
    dimension: int,
) -> NDArray[np.bool_]:
    """Create Box Area

    Args:
        tokens: [Ntoken, 4]
        resolution: float, default = 0.5
        dimension: int, default = 64,

    Returns:
        box_areas: BoolArray [Ntoken, D, H, W] D=H=W=dimension
    """
    num_tokens = len(tokens)
    box_areas = np.zeros((num_tokens, dimension, dimension, dimension), dtype=np.bool_)
    grids = np.stack(np.meshgrid(np.arange(dimension), np.arange(dimension), np.arange(dimension), indexing='ij'), 3)
    for i, (x, y, z, t) in enumerate(tokens):
        x, y, z, t = int(x), int(y), int(z), int(t)
        distances = np.linalg.norm(grids - np.array([[x, y, z]]), axis=-1)
        threshold = math.ceil((C.INTERACTION_DIST[int(t)] + pharmacophore_size) / resolution)
        box_areas[i] = distances < threshold
    return box_areas

src/data/token_inference.py
SeonghwanSeo/PharmacoNet
import os
import tempfile
import math

import torch
import numpy as np

from typing import Optional, Dict, Union, Tuple
from omegaconf import OmegaConf
from torch import Tensor
from numpy.typing import NDArray, ArrayLike

from molvoxel import create_voxelizer, BaseVoxelizer

from .network import build_model
from .network.detector import PharmacoFormer
from .data import token_inference, pointcloud
from .data import constant as C
from .data import INTERACTION_LIST
from .data.objects import Protein
from .data.extract_pocket import extract_pocket

from .scoring import PharmacophoreModel
from . import utils
from .utils.smoothing import GaussianSmoothing


DEFAULT_FOCUS_THRESHOLD = 0.5
DEFAULT_BOX_THRESHOLD = 0.5
DEFAULT_SCORE_THRESHOLD = {
    'PiStacking_P': 0.6,    # Top 40%
    'PiStacking_T': 0.6,
    'SaltBridge_lneg': 0.6,
    'SaltBridge_pneg': 0.6,
    'PiCation_lring': 0.6,
    'PiCation_pring': 0.6,
    'XBond': 0.8,           # Top 20%
    'HBond_ldon': 0.8,
    'HBond_pdon': 0.8,
    'Hydrophobic': 0.8,
}


class ModelingModule():
    def __init__(
        self,
        model_path: str,
        device: str = 'cpu',
        molvoxel_library: str = 'numpy',
        focus_threshold: float = DEFAULT_FOCUS_THRESHOLD,
        box_threshold: float = DEFAULT_BOX_THRESHOLD,
        score_threshold: Union[float, Dict[str, float]] = DEFAULT_SCORE_THRESHOLD,
    ):
        checkpoint = torch.load(model_path, map_location='cpu')
        self.config = config = OmegaConf.create(checkpoint['config'])
        self.device = device
        model = build_model(config.MODEL)
        model.load_state_dict(checkpoint['model'])
        model.eval()
        self.model: PharmacoFormer = model.to(device)
        self.smoothing = GaussianSmoothing(kernel_size=5, sigma=0.5).to(device)
        self.focus_threshold = focus_threshold
        self.box_threshold = box_threshold
        self.score_distributions = {typ: np.array(distribution['focus']) for typ, distribution in checkpoint['score_distributions'].items()}

        self.score_threshold: Dict[str, float]
        if isinstance(score_threshold, float):
            self.score_threshold = {typ: score_threshold for typ in INTERACTION_LIST}
        else:
            self.score_threshold = score_threshold
        # self.score_threshold = utils.get_score_threshold(checkpoint['score_distributions'], score_threshold, INTERACTION_LIST)
        del checkpoint

        in_resolution = config.VOXEL.IN.RESOLUTION
        in_size = config.VOXEL.IN.SIZE
        self.in_voxelizer: BaseVoxelizer = create_voxelizer(in_resolution, in_size, sigma=(1 / 3), library=molvoxel_library)
        self.pocket_cutoff = (in_resolution * in_size * math.sqrt(3) / 2) + 5.0
        self.out_resolution = config.VOXEL.OUT.RESOLUTION
        self.out_size = config.VOXEL.OUT.SIZE

    def run(
        self,
        protein_pdb_path: str,
        ref_ligand_path: Optional[str] = None,
        center: Optional[ArrayLike] = None,
    ) -> PharmacophoreModel:
        assert (ref_ligand_path is not None) or (center is not None)
        if ref_ligand_path is not None:
            ref_ligand = utils.load_ligand(ref_ligand_path)
            center_array = np.mean([atom.coords for atom in ref_ligand.atoms], axis=0, dtype=np.float32)
        else:
            center_array = np.array(center, dtype=np.float32)
        assert center_array is not None
        assert center_array.shape == (3,)

        return self._run(protein_pdb_path, center_array)

    @torch.no_grad()
    def _run(
        self,
        protein_pdb_path: str,
        center: NDArray[np.float32],
    ):
        protein_image, non_protein_area, token_positions, tokens = self.__parse_protein(protein_pdb_path, center)
        density_maps = self.__create_density_maps(
            torch.from_numpy(protein_image),
            torch.from_numpy(non_protein_area) if non_protein_area is not None else None,
            torch.from_numpy(token_positions),
            torch.from_numpy(tokens),
        )
        x, y, z = center.tolist()
        pharmacophore_model = PharmacophoreModel.create((x, y, z), self.out_resolution, self.out_size, density_maps)
        return pharmacophore_model

    def __parse_protein(
        self,
        protein_pdb_path: str,
        center: NDArray[np.float32],
    ) -> Tuple[NDArray, Optional[NDArray], NDArray, NDArray]:

        with tempfile.TemporaryDirectory() as dirname:
            pocket_path = os.path.join(dirname, 'pocket.pdb')
            extract_pocket(protein_pdb_path, pocket_path, center, self.pocket_cutoff)   # root(3)
            protein_obj: Protein = Protein.from_pdbfile(pocket_path)

        token_positions, token_classes = token_inference.get_token_informations(protein_obj)
        tokens, filter = token_inference.get_token_and_filter(
            token_positions, token_classes, center, self.out_resolution, self.out_size
        )
        token_positions = token_positions[filter]

        protein_positions, protein_features = pointcloud.get_protein_pointcloud(protein_obj)

        protein_image = np.asarray(
            self.in_voxelizer.forward_features(protein_positions, center, protein_features, radii=self.config.VOXEL.RADII.PROTEIN),
            np.float32
        )
        if self.config.VOXEL.RADII.PROTEIN_MASKING > 0:
            non_protein_area = np.logical_not(
                np.asarray(
                    self.in_voxelizer.forward_single(protein_positions, center, radii=self.config.VOXEL.RADII.PROTEIN_MASKING),
                    np.bool_
                )
            )
        else:
            non_protein_area = None

        return protein_image, non_protein_area, token_positions, tokens,

    def __create_density_maps(
        self,
        protein_image: Tensor,
        non_protein_area: Optional[Tensor],
        token_positions: Tensor,
        tokens: Tensor,
    ):
        protein_image = protein_image.to(device=self.device, dtype=torch.float)
        token_positions = token_positions.to(device=self.device, dtype=torch.float)
        tokens = tokens.to(device=self.device, dtype=torch.long)
        non_protein_area = non_protein_area.to(device=self.device, dtype=torch.bool) if non_protein_area is not None else None

        with torch.amp.autocast(self.device, enabled=self.config.AMP_ENABLE):
            protein_image = protein_image.unsqueeze(0)
            multi_scale_features = self.model.forward_feature(protein_image)                    # List[[1, D, H, W, F]]
            bottom_features = multi_scale_features[-1]

            token_scores, token_features = self.model.forward_token_prediction(bottom_features, [tokens])   # [[Ntoken,]], [[Ntoken, F]]
            token_scores = token_scores[0].sigmoid()                                            # [Ntoken,]
            token_features = token_features[0]                                                  # [Ntoken, F]

            cavity_narrow, cavity_wide = self.model.forward_cavity_extraction(bottom_features)  # [1, 1, D, H, W], [1, 1, D, H, W]
            cavity_narrow = cavity_narrow[0].sigmoid() > self.focus_threshold                   # [1, D, H, W]
            cavity_wide = cavity_wide[0].sigmoid() > self.focus_threshold                       # [1, D, H, W]

            num_tokens = tokens.shape[0]
            indices = []
            relative_scores = []
            for i in range(num_tokens):
                x, y, z, typ = tokens[i].tolist()
                # NOTE: Check the token score
                absolute_score = token_scores[i].item()
                relative_score = float((self.score_distributions[INTERACTION_LIST[int(typ)]] < absolute_score).mean())
                if relative_score < self.score_threshold[INTERACTION_LIST[int(typ)]]:
                    continue
                # NOTE: Check the token exists in cavity
                if typ in C.LONG_INTERACTION:
                    if not cavity_wide[0, x, y, z]:
                        continue
                else:
                    if not cavity_narrow[0, x, y, z]:
                        continue
                indices.append(i)
                relative_scores.append(relative_score)
            selected_indices = torch.tensor(indices, device=self.device, dtype=torch.long)  # [Ntoken',]

            tokens = tokens[selected_indices]                                               # [Ntoken',]
            # token_scores = token_scores[selected_indices]                                   # [Ntoken',]
            token_positions = token_positions[selected_indices]                             # [Ntoken', 3]
            token_features = token_features[selected_indices]                               # [Ntoken', F]

            density_maps_list = []
            step = 8
            for idx in range(0, tokens.size(0), step):
                _tokens, _token_features = tokens[idx:idx + step], token_features[idx:idx + step]
                density_maps = self.model.forward_segmentation(multi_scale_features, [_tokens], [_token_features])[0]   # [[4, D, H, W]]
                density_maps = density_maps[0].sigmoid()                                    # [4, D, H, W]
                density_maps_list.append(density_maps)
            density_maps = torch.cat(density_maps_list, dim=0)                              # [Ntoken', D, H, W]

            box_area = token_inference.get_box_area(
                tokens, self.config.VOXEL.RADII.PHARMACOPHORE, self.out_resolution, self.out_size,
            )
            box_area = torch.from_numpy(box_area).to(device=self.device, dtype=torch.bool)  # [Ntoken', D, H, W]
            unavailable_area = ~ (box_area & non_protein_area & cavity_narrow)              # [Ntoken', D, H, W]

            # NOTE: masking should be performed before smoothing - masked area is not trained.
            density_maps.masked_fill_(unavailable_area, 0.)
            density_maps = self.smoothing(density_maps)
            density_maps.masked_fill_(unavailable_area, 0.)
            density_maps[density_maps < self.box_threshold] = 0.

        out = []
        assert len(tokens) == len(relative_scores)
        for token, score, position, map in zip(tokens, relative_scores, token_positions, density_maps):
            if torch.all(map < 1e-6):
                continue
            out.append({
                'coords': tuple(token[:3].tolist()),
                'type': INTERACTION_LIST[int(token[3])],
                'position': tuple(position.tolist()),
                'score': float(score),
                'map': map.cpu().numpy(),
            })
        return out

src/modeling.py
SeonghwanSeo/PharmacoNet
from .builder import build_model
from . import detector, backbones, decoders, necks, feature_embedding, mask_head, token_head

src/network/__init__.py
SeonghwanSeo/PharmacoNet
from .swin import SwinTransformer
from .swinv2 import SwinTransformerV2

src/network/backbones/__init__.py
SeonghwanSeo/PharmacoNet
#  --------------------------------------------------------
# Swin Transformer
# Copyright (c) 2021 Microsoft
# Licensed under The MIT License [see LICENSE for details]
# Written by Ze Liu, Yutong Lin, Yixuan Wei
# --------------------------------------------------------

import torch
import torch.nn as nn
import torch.nn.functional as F
import numpy as np
from timm.models.layers import DropPath, to_3tuple, trunc_normal_

from typing import Sequence, Tuple, Optional, Type, Union, List
from torch import Tensor

from ..builder import BACKBONE

class Mlp(nn.Module):
    """ Multilayer perceptron."""
    def __init__(
            self,
            in_features: int,
            hidden_features: Optional[int] = None,
            out_features: Optional[int] = None,
            act_layer: Type[nn.Module] = nn.GELU,
            drop: float = 0.
        ):
        super().__init__()
        out_features = out_features or in_features
        hidden_features = hidden_features or in_features
        self.fc1 = nn.Linear(in_features, hidden_features)
        self.act = act_layer()
        self.fc2 = nn.Linear(hidden_features, out_features)
        self.drop = nn.Dropout(drop)

    def forward(self, x):
        x = self.fc1(x)
        x = self.act(x)
        x = self.drop(x)
        x = self.fc2(x)
        x = self.drop(x)
        return x


def window_partition(x: Tensor, window_size: int) -> Tensor:
    """
    Args:
        x: (B, D, H, W, C)
        window_size (int): window size

    Returns:
        windows: (num_windows*B, window_size, window_size, window_size, C)
    """
    B, D, H, W, C = x.shape
    x = x.view(B, 
               D // window_size, window_size,
               H // window_size, window_size,
               W // window_size, window_size,
               C)
    windows = x.permute(0, 1, 3, 5, 2, 4, 6, 7).contiguous().view(-1, window_size, window_size, window_size, C)
    return windows


def window_reverse(windows, window_size, D, H, W):
    """
    Args:
        windows: (num_windows*B, window_size, window_size, window_size, C)
        window_size (int): Window size
        D (int): Depth of image
        H (int): Height of image
        W (int): Width of image

    Returns:
        x: (B, D, H, W, C)
    """
    B = int(windows.shape[0] / (D * H * W / window_size / window_size / window_size))
    x = windows.view(
            B,
            D // window_size, H // window_size, W // window_size,
            window_size, window_size, window_size, -1)
    x = x.permute(0, 1, 4, 2, 5, 3, 6, 7).contiguous().view(B, D, H, W, -1)
    return x


class WindowAttention(nn.Module):
    """ Window based multi-head self attention (W-MSA) module with relative position bias.
    It supports both of shifted and non-shifted window.

    Args:
        dim (int): Number of input channels.
        window_size (tuple[int]): The height and width of the window.
        num_heads (int): Number of attention heads.
        qkv_bias (bool, optional):  If True, add a learnable bias to query, key, value. Default: True
        qk_scale (float | None, optional): Override default qk scale of head_dim ** -0.5 if set
        attn_drop (float, optional): Dropout ratio of attention weight. Default: 0.0
        proj_drop (float, optional): Dropout ratio of output. Default: 0.0
    """

    def __init__(
            self,
            dim: int,
            window_size: Tuple[int, int, int],
            num_heads: int,
            qkv_bias: bool = True,
            qk_scale: Optional[float] = None,
            attn_drop: float = 0.,
            proj_drop: float = 0.
        ):

        super().__init__()
        self.dim = dim
        self.window_size = window_size  # Wd, Wh, Ww
        self.num_heads = num_heads
        head_dim = dim // num_heads
        self.scale = qk_scale or head_dim ** -0.5

        # define a parameter table of relative position bias
        self.relative_position_bias_table = nn.Parameter(
            torch.zeros(
                (2 * window_size[0] - 1) *
                (2 * window_size[1] - 1) * 
                (2 * window_size[2] - 1) 
                , num_heads))  # 2*Wd-1 * 2*Wh-1 * 2*Ww-1, nH

        # get pair-wise relative position index for each token inside the window
        coords_d = torch.arange(self.window_size[0])
        coords_h = torch.arange(self.window_size[1])
        coords_w = torch.arange(self.window_size[2])
        coords = torch.stack(torch.meshgrid([coords_d, coords_h, coords_w], indexing='ij'))     # 3, Wd, Wh, Ww
        coords_flatten = torch.flatten(coords, 1)                                               # 3, Wd*Wh*Ww
        relative_coords = coords_flatten[:, :, None] - coords_flatten[:, None, :]               # 3, Wd*Wh*Ww, Wd*Wh*Ww
        relative_coords = relative_coords.permute(1, 2, 0).contiguous()                         # Wd*Wh*Ww, Wd*Wh*Ww, 3
        relative_coords[:, :, 0] += self.window_size[0] - 1
        relative_coords[:, :, 1] += self.window_size[1] - 1
        relative_coords[:, :, 2] += self.window_size[2] - 1
        relative_coords[:, :, 0] *= (2 * self.window_size[1] - 1) * (2 * self.window_size[2] - 1)
        relative_coords[:, :, 1] *= 2 * self.window_size[2] - 1
        relative_position_index = relative_coords.sum(-1)                           # Wd*Wh*Ww, Wd*Wh*Ww
        self.register_buffer("relative_position_index", relative_position_index)

        self.qkv = nn.Linear(dim, dim * 3, bias=qkv_bias)
        self.attn_drop = nn.Dropout(attn_drop)
        self.proj = nn.Linear(dim, dim)
        self.proj_drop = nn.Dropout(proj_drop)

        trunc_normal_(self.relative_position_bias_table, std=.02)
        self.softmax = nn.Softmax(dim=-1)

    def forward(self, x, mask=None):
        """ Forward function.

        Args:
            x: input features with shape of (num_windows*B, N, C)
            mask: (0/-inf) mask with shape of (num_windows, Wd*Wh*Ww, Wd*Wh*Ww) or None
        """
        B_, N, C = x.shape
        qkv = self.qkv(x).reshape(B_, N, 3, self.num_heads, C // self.num_heads).permute(2, 0, 3, 1, 4)
        q, k, v = qkv[0], qkv[1], qkv[2]  # make torchscript happy (cannot use tensor as tuple)

        q = q * self.scale
        attn = (q @ k.transpose(-2, -1))

        num_windows = self.window_size[0] * self.window_size[1] * self.window_size[2]   # Wd*Wh*Ww
        relative_position_bias = self.relative_position_bias_table[self.relative_position_index.view(-1)].view(
            num_windows, num_windows, -1)   # Wd*Wh*Ww, Wd*Wh*Ww, nH
        relative_position_bias = relative_position_bias.permute(2, 0, 1).contiguous()                   # nH, Wd*Wh*Ww, Wd*Wh*Ww
        attn = attn + relative_position_bias.unsqueeze(0)

        if mask is not None:
            nW = mask.shape[0]
            attn = attn.view(B_ // nW, nW, self.num_heads, N, N) + mask.unsqueeze(1).unsqueeze(0)
            attn = attn.view(-1, self.num_heads, N, N)
            attn = self.softmax(attn)
        else:
            attn = self.softmax(attn)

        attn = self.attn_drop(attn)

        x = (attn @ v).transpose(1, 2).reshape(B_, N, C)
        x = self.proj(x)
        x = self.proj_drop(x)
        return x


class SwinTransformerBlock(nn.Module):
    """ Swin Transformer Block.

    Args:
        dim (int): Number of input channels.
        num_heads (int): Number of attention heads.
        window_size (int): Window size.
        shift_size (int): Shift size for SW-MSA.
        mlp_ratio (float): Ratio of mlp hidden dim to embedding dim.
        qkv_bias (bool, optional): If True, add a learnable bias to query, key, value. Default: True
        qk_scale (float | None, optional): Override default qk scale of head_dim ** -0.5 if set.
        drop (float, optional): Dropout rate. Default: 0.0
        attn_drop (float, optional): Attention dropout rate. Default: 0.0
        drop_path (float, optional): Stochastic depth rate. Default: 0.0
        act_layer (nn.Module, optional): Activation layer. Default: nn.GELU
        norm_layer (nn.Module, optional): Normalization layer.  Default: nn.LayerNorm
    """

    def __init__(
            self,
            dim: int,
            num_heads: int,
            window_size: int = 7,
            shift_size: int = 0,
            mlp_ratio: float = 4.,
            qkv_bias: bool = True,
            qk_scale: Optional[float] = None,
            drop: float = 0.,
            attn_drop: float = 0.,
            drop_path: float = 0.,
            act_layer: Type[nn.Module] = nn.GELU,
            norm_layer: Type[nn.Module] = nn.LayerNorm,
    ):
        super().__init__()
        self.dim = dim
        self.num_heads = num_heads
        self.window_size = window_size
        self.shift_size = shift_size
        self.mlp_ratio = mlp_ratio
        assert 0 <= self.shift_size < self.window_size, "shift_size must in 0-window_size"

        self.norm1 = norm_layer(dim)
        self.attn = WindowAttention(
            dim, window_size=to_3tuple(self.window_size), num_heads=num_heads,
            qkv_bias=qkv_bias, qk_scale=qk_scale, attn_drop=attn_drop, proj_drop=drop)

        self.drop_path = DropPath(drop_path) if drop_path > 0. else nn.Identity()
        self.norm2 = norm_layer(dim)
        mlp_hidden_dim = int(dim * mlp_ratio)
        self.mlp = Mlp(in_features=dim, hidden_features=mlp_hidden_dim, act_layer=act_layer, drop=drop)

    def forward(self, x, mask_matrix):
        """ Forward function.

        Args:
            x: Input feature, tensor size (B, D, H, W, C).
            D, H, W: Spatial resolution of the input feature.
            mask_matrix: Attention mask for cyclic shift.

        Return:
            x: Output feature, tensor size (B, D, H, W, C).
        """
        B, D, H, W, C = x.shape
        shortcut = x
        x = self.norm1(x)

        # pad feature maps to multiples of window size
        pad_w0 = pad_h0 = pad_d0 = 0
        pad_w1 = (self.window_size - H % self.window_size) % self.window_size
        pad_h1 = (self.window_size - W % self.window_size) % self.window_size
        pad_d1 = (self.window_size - D % self.window_size) % self.window_size
        x = F.pad(x, (0, 0, pad_w0, pad_w1, pad_h0, pad_h1, pad_d0, pad_d1))
        _, Dp, Hp, Wp, _ = x.shape

        # cyclic shift
        if self.shift_size > 0:
            shifted_x = torch.roll(x, shifts=(-self.shift_size, -self.shift_size, -self.shift_size), dims=(1, 2, 3))
            attn_mask = mask_matrix
        else:
            shifted_x = x
            attn_mask = None

        # partition windows
        x_windows = window_partition(shifted_x, self.window_size)               # nW*B, window_size, window_size, window_size, C (nW: num_windows)
        x_windows = x_windows.view(-1, self.window_size * self.window_size * self.window_size, C)
                                                                                # nW*B, window_size*window_size*window_size, C

        # W-MSA/SW-MSA
        attn_windows = self.attn(x_windows, mask=attn_mask)                     # nW*B, window_size*window_size*window_size, C

        # merge windows
        attn_windows = attn_windows.view(-1, self.window_size, self.window_size, self.window_size, C)
        shifted_x = window_reverse(attn_windows, self.window_size, Dp, Hp, Wp)  # B, D, H, W, C

        # reverse cyclic shift
        if self.shift_size > 0:
            x = torch.roll(shifted_x, shifts=(self.shift_size, self.shift_size, self.shift_size), dims=(1, 2, 3))
        else:
            x = shifted_x

        if pad_d1 > 0 or pad_h1 > 0 or pad_w1 > 0:
            x = x[:, :D, :H, :W, :].contiguous()

        # FFN
        x = shortcut + self.drop_path(x)
        x = x + self.drop_path(self.mlp(self.norm2(x)))

        return x    # B, D, H, W, C


class PatchMerging(nn.Module):
    """ Patch Merging Layer

    Args:
        dim (int): Number of input channels.
        norm_layer (nn.Module, optional): Normalization layer.  Default: nn.LayerNorm
    """
    def __init__(self, dim, norm_layer=nn.LayerNorm):
        super().__init__()
        self.dim = dim
        self.reduction = nn.Linear(8 * dim, 2 * dim, bias=False)
        self.norm = norm_layer(8 * dim)

    def forward(self, x):
        """ Forward function.

        Args:
            x: Input feature, tensor size (B, D, H, W, C).

        Return:
            x: Input feature, tensor size (B, D/2, H/2, W/2, 2*C).
        """
        B, D, H, W, C = x.shape

        # padding
        pad_input = (D % 2 == 1) or (H % 2 == 1) or (W % 2 == 1)
        if pad_input:
            x = F.pad(x, (0, 0, 0, W % 2, 0, H % 2, 0, D % 2))

        x0 = x[:, 0::2, 0::2, 0::2, :]  # B, D/2, H/2, W/2, C
        x1 = x[:, 1::2, 0::2, 0::2, :]  # B, D/2, H/2, W/2, C
        x2 = x[:, 0::2, 1::2, 0::2, :]  # B, D/2, H/2, W/2, C
        x3 = x[:, 1::2, 1::2, 0::2, :]  # B, D/2, H/2, W/2, C
        x4 = x[:, 0::2, 0::2, 1::2, :]  # B, D/2, H/2, W/2, C
        x5 = x[:, 1::2, 0::2, 1::2, :]  # B, D/2, H/2, W/2, C
        x6 = x[:, 0::2, 1::2, 1::2, :]  # B, D/2, H/2, W/2, C
        x7 = x[:, 1::2, 1::2, 1::2, :]  # B, D/2, H/2, W/2, C
        x = torch.cat([x0, x1, x2, x3, x4, x5, x6, x7], -1)  # B, D/2, H/2, W/2, 8*C

        x = self.norm(x)
        x = self.reduction(x)           # B, D/2, H/2, W/2, 2*C

        return x


class BasicLayer(nn.Module):
    """ A basic Swin Transformer layer for one stage.
Args:
        dim (int): Number of feature channels
        depth (int): Depths of this stage.
        num_heads (int): Number of attention head.
        window_size (int): Local window size. Default: 7.
        mlp_ratio (float): Ratio of mlp hidden dim to embedding dim. Default: 4.
        qkv_bias (bool, optional): If True, add a learnable bias to query, key, value. Default: True
        qk_scale (float | None, optional): Override default qk scale of head_dim ** -0.5 if set.
        drop (float, optional): Dropout rate. Default: 0.0
        attn_drop (float, optional): Attention dropout rate. Default: 0.0
        drop_path (float | tuple[float], optional): Stochastic depth rate. Default: 0.0
        norm_layer (nn.Module, optional): Normalization layer. Default: nn.LayerNorm
        downsample (nn.Module | None, optional): Downsample layer at the end of the layer. Default: None
    """

    def __init__(
            self,
            dim: int,
            depth: int,
            num_heads: int,
            window_size: int = 7,
            mlp_ratio: float = 4.,
            qkv_bias: bool = True,
            qk_scale: Optional[float] = None,
            drop: float = 0.,
            attn_drop: float = 0.,
            drop_path: Union[float, Sequence[float]] = 0.,
            norm_layer: Type[nn.Module] = nn.LayerNorm,
            downsample: Optional[Type[nn.Module]] = None,
        ):
        super().__init__()
        self.window_size = window_size
        self.shift_size = window_size // 2
        self.depth = depth

        # build blocks
        self.blocks = nn.ModuleList([
            SwinTransformerBlock(
                dim=dim,
                num_heads=num_heads,
                window_size=window_size,
                shift_size=0 if (i % 2 == 0) else window_size // 2,
                mlp_ratio=mlp_ratio,
                qkv_bias=qkv_bias,
                qk_scale=qk_scale,
                drop=drop,
                attn_drop=attn_drop,
                drop_path=drop_path if isinstance(drop_path, float) else drop_path[i],
                norm_layer=norm_layer
            )
            for i in range(depth)])

        # patch merging layer
        if downsample is not None:
            self.downsample = downsample(dim=dim, norm_layer=norm_layer)
        else:
            self.downsample = None

    def forward(self, x: Tensor) -> Tuple[Tensor, int, int, int, Tensor, int, int, int]:
        """ Forward function.

        Args:
            x: Input feature, tensor size (B, D, H, W, C).
        """
        B, D, H, W, C = x.shape

        # calculate attention mask for SW-MSA
        Dp = int(np.ceil(D / self.window_size)) * self.window_size
        Hp = int(np.ceil(H / self.window_size)) * self.window_size
        Wp = int(np.ceil(W / self.window_size)) * self.window_size
        img_mask = torch.zeros((1, Dp, Hp, Wp, 1), device=x.device)  # 1 Dp Hp Wp 1
        d_slices = (slice(0, -self.window_size),
                    slice(-self.window_size, -self.shift_size),
                    slice(-self.shift_size, None))
        h_slices = (slice(0, -self.window_size),
                    slice(-self.window_size, -self.shift_size),
                    slice(-self.shift_size, None))
        w_slices = (slice(0, -self.window_size),
                    slice(-self.window_size, -self.shift_size),
                    slice(-self.shift_size, None))
        cnt = 0
        for d in d_slices:
            for h in h_slices:
                for w in w_slices:
                    img_mask[:, d, h, w, :] = cnt
                    cnt += 1

        mask_windows = window_partition(img_mask, self.window_size)  # nW, window_size, window_size, window_size, 1
        mask_windows = mask_windows.view(-1, self.window_size * self.window_size * self.window_size)
        attn_mask = mask_windows.unsqueeze(1) - mask_windows.unsqueeze(2)
        attn_mask = attn_mask.masked_fill(attn_mask != 0, float(-100.0)).masked_fill(attn_mask == 0, float(0.0))

        for blk in self.blocks:
            x = blk(x, attn_mask)
        if self.downsample is not None:
            x_down = self.downsample(x)
            Wd, Wh, Ww = (D + 1) // 2, (H + 1) // 2, (W + 1) // 2
            return x, D, H, W, x_down, Wd, Wh, Ww
        else:
            return x, D, H, W, x, D, H, W


class PatchEmbed(nn.Module):
    """ Image to Patch Embedding

    Args:
        patch_size (int): Patch token size. Default: 4.
        in_chans (int): Number of input image channels. Default: 3.
        embed_dim (int): Number of linear projection output channels. Default: 96.
        norm_layer (nn.Module, optional): Normalization layer. Default: None
    """

    def __init__(
            self,
            in_chans: int = 3,
            patch_size: Union[int, Tuple[int, int, int]] = 4,
            embed_dim: int = 96,
            norm_layer: Optional[Type[nn.Module]] = None):
        super().__init__()
        self.patch_size: Tuple[int, int, int] = to_3tuple(patch_size)

        self.in_chans = in_chans
        self.embed_dim = embed_dim

        self.proj = nn.Conv3d(in_chans, embed_dim, kernel_size=patch_size, stride=patch_size,
                              bias = (norm_layer is None))
        if norm_layer is not None:
            self.norm = norm_layer(embed_dim)
        else:
            self.norm = None

    def forward(self, x):
        """Forward function.

        Args:
            x: Input feature, tensor size (B, C, D, H, W).
        Return:
            tensor size (B, D, H, W, C').
        """
        # padding
        B, C, D, H, W = x.size()
        if W % self.patch_size[2] != 0:
            x = F.pad(x, (0, self.patch_size[2] - W % self.patch_size[2]))
        if H % self.patch_size[1] != 0:
            x = F.pad(x, (0, 0, 0, self.patch_size[1] - H % self.patch_size[1]))
        if D % self.patch_size[0] != 0:
            x = F.pad(x, (0, 0, 0, 0, 0, self.patch_size[0] - D % self.patch_size[0]))

        x = self.proj(x)                        # B, C, Wd, Wh, Ww
        if self.norm is not None:
            _, _, Wd, Wh, Ww = x.shape
            x = x.flatten(2).transpose(1, 2)    # B, Wd*Wh*Ww, C'
            x = self.norm(x)
            x = x.transpose(1, 2).view(B, self.embed_dim, Wd, Wh, Ww)  # B, C', Wd, Wh, Ww

        return x


@BACKBONE.register()
class SwinTransformer(nn.Module):
    """ Swin Transformer backbone.
        A PyTorch impl of : `Swin Transformer: Hierarchical Vision Transformer using Shifted Windows`  -
          https://arxiv.org/pdf/2103.14030

    Args:
        in_channels (int): Number of input image channels.
        image_size (int): Input image size,
            used in absolute postion embedding. Default 48.
        patch_size (int | tuple(int)): Patch size. Default: 2.
        embed_dim (int): Number of linear projection output channels. Default: 96.
        depths (tuple[int]): Depths of each Swin Transformer stage.
        num_heads (tuple[int]): Number of attention head of each stage.
        window_size (int): Window size. Default: 7.
        mlp_ratio (float): Ratio of mlp hidden dim to embedding dim. Default: 4.
        qkv_bias (bool): If True, add a learnable bias to query, key, value. Default: True
        qk_scale (float): Override default qk scale of head_dim ** -0.5 if set.
        drop_rate (float): Dropout rate.
        attn_drop_rate (float): Attention dropout rate. Default: 0.
        drop_path_rate (float): Stochastic depth rate. Default: 0.2.
        norm_layer (nn.Module): Normalization layer. Default: nn.LayerNorm.
        ape (bool): If True, add absolute position embedding to the patch embedding. Default: False.
        patch_norm (bool): If True, add normalization after patch embedding. Default: True.
        out_indices (Sequence[int]): Output from which stages.
        frozen_stages (int): Stages to be frozen (stop grad and set eval mode).
            -1 means not freezing any parameters.
    """

    def __init__(
            self,
            in_channels: int,
            image_size: int = 48,
            patch_size: Union[int, Tuple[int, int, int]] = 2,
            embed_dim: int = 64,
            depths: Sequence[int] = [2, 2, 6, 2],
            num_heads: Sequence[int] = [3, 6, 12, 12],
            window_size: int = 6,
            mlp_ratio: float = 4.,
            qkv_bias: bool = True,
            qk_scale: Optional[float] = None,
            drop_rate: float = 0.,
            attn_drop_rate: float = 0.,
            drop_path_rate: float = 0.2,
            norm_layer: Type[nn.Module] = nn.LayerNorm,
            ape: bool = False,
            patch_norm: bool = True,
            out_indices: Union[int, Sequence[int]] = [0, 1, 2, 3],
        ):
        super().__init__()

        self.image_size = image_size
        self.num_layers = len(depths)
        self.embed_dim = embed_dim
        self.ape = ape
        self.patch_norm = patch_norm
        out_indices = out_indices if not isinstance(out_indices, int) else [out_indices]
        self.out_indices = [(index if index >= 0 else self.num_layers + index) for index in out_indices]

        # split image into non-overlapping patches
        self.patch_embed = PatchEmbed(
            in_chans=in_channels, patch_size=patch_size, embed_dim=embed_dim,
            norm_layer=norm_layer if self.patch_norm else None)

        # absolute position embedding
        if self.ape:
            patches_resolution = [imgsz // patchsz for imgsz, patchsz in zip(to_3tuple(image_size), to_3tuple(patch_size))]
            self.absolute_pos_embed = nn.Parameter(torch.zeros(1, embed_dim, *patches_resolution))  # (1, C, D, H, W)
            trunc_normal_(self.absolute_pos_embed, std=.02)

        self.pos_drop = nn.Dropout(p=drop_rate)

        # stochastic dept
        dpr = [x.item() for x in torch.linspace(0, drop_path_rate, sum(depths))]  # stochastic depth decay rule

        # build layers
        self.layers = nn.ModuleList()
        for i_layer in range(self.num_layers):
            layer = BasicLayer(
                dim=int(embed_dim * 2 ** i_layer),
                depth=depths[i_layer],
                num_heads=num_heads[i_layer],
                window_size=window_size,
                mlp_ratio=mlp_ratio,
                qkv_bias=qkv_bias,
                qk_scale=qk_scale,
                drop=drop_rate,
                attn_drop=attn_drop_rate,
                drop_path=dpr[sum(depths[:i_layer]):sum(depths[:i_layer + 1])],
                norm_layer=norm_layer,
                downsample=PatchMerging if (i_layer < self.num_layers - 1) else None)
            self.layers.append(layer)

        num_features = [int(embed_dim * 2 ** i) for i in range(self.num_layers)]
        self.num_features = num_features

        # add a norm layer for each output
        for i_layer in out_indices:
            layer = norm_layer(num_features[i_layer])
            layer_name = f'norm{i_layer}'
            self.add_module(layer_name, layer)

    def initialize_weights(self):
        """Initialize the weights in backbone.

        Args:
            pretrained (str, optional): Path to pre-trained weights.
                Defaults to None.
        """

        def _init_weights(m):
            if isinstance(m, (nn.Linear, nn.Conv3d)):
                trunc_normal_(m.weight, std=.02)
                if isinstance(m, nn.Linear) and m.bias is not None:
                    nn.init.constant_(m.bias, 0)
            elif isinstance(m, (nn.BatchNorm3d, nn.GroupNorm, nn.LayerNorm)):
                nn.init.constant_(m.bias, 0)
                nn.init.constant_(m.weight, 1.0)
        self.apply(_init_weights)

    def forward(self, x: Tensor) -> Sequence[Tensor]:
        """Forward function."""
        x = self.patch_embed(x)             # B, C, D, H, W -> B, C', Wd, Wh, Ww

        B, _, Wd, Wh, Ww = x.shape
        if self.ape:
            # interpolate the position embedding to the corresponding size
            absolute_pos_embed = F.interpolate(self.absolute_pos_embed, size=(Wd, Wh, Ww), mode='trilinear')
            x = (x + absolute_pos_embed).flatten(2).transpose(1, 2)  # B, Wd*Wh*Ww, C'
        else:
            x = x.flatten(2).transpose(1, 2)
        x = self.pos_drop(x)
        x = x.view(B, Wd, Wh, Ww, -1)       # B, Wd, Wh, Ww, C'

        outs = []
        for i in range(self.num_layers):
            layer = self.layers[i]
            x_out, _, _, _, x, _, _, _ = layer(x)             # B, Wd, Wh, Ww, C' | B, Wd', Wh', Ww', C''
            if i in self.out_indices:
                norm_layer = getattr(self, f'norm{i}')
                x_out = norm_layer(x_out)
                x_out = x_out.permute(0, 4, 1, 2, 3).contiguous()     # B, C', Wd, Wh, Ww
                outs.append(x_out)

        return outs

src/network/backbones/swin.py
SeonghwanSeo/PharmacoNet
# --------------------------------------------------------
# Swin Transformer V2
# Copyright (c) 2022 Microsoft
# Licensed under The MIT License [see LICENSE for details]
# Written by Ze Liu
# --------------------------------------------------------

import torch
import torch.nn as nn
import torch.nn.functional as F
import numpy as np
from timm.models.layers import DropPath, to_3tuple, trunc_normal_

from typing import Sequence, Tuple, Optional, Type, Union, List
from torch import Tensor

from .swin import Mlp, window_partition, window_reverse
from ..builder import BACKBONE

class WindowAttention(nn.Module):
    r""" Window based multi-head self attention (W-MSA) module with relative position bias.
    It supports both of shifted and non-shifted window.

    Args:
        dim (int): Number of input channels.
        window_size (tuple[int]): The height and width of the window.
        num_heads (int): Number of attention heads.
        qkv_bias (bool, optional):  If True, add a learnable bias to query, key, value. Default: True
        attn_drop (float, optional): Dropout ratio of attention weight. Default: 0.0
        proj_drop (float, optional): Dropout ratio of output. Default: 0.0
    """

    def __init__(
            self,
            dim: int,
            window_size: Tuple[int, int, int],
            num_heads: int,
            qkv_bias: bool = True,
            attn_drop: float = 0.,
            proj_drop: float = 0.,
        ):

        super().__init__()
        self.dim = dim
        self.window_size = window_size  # Wd, Wh, Ww
        self.num_heads = num_heads

        self.logit_scale = nn.Parameter(torch.log(10 * torch.ones((num_heads, 1, 1))), requires_grad=True)

        # mlp to generate continuous relative position bias
        self.cpb_mlp = nn.Sequential(nn.Linear(3, 512, bias=True),
                                     nn.ReLU(inplace=True),
                                     nn.Linear(512, num_heads, bias=False))

        # get relative_coords_table
        relative_coords_d = torch.arange(-(self.window_size[0] - 1), self.window_size[0], dtype=torch.float32)
        relative_coords_h = torch.arange(-(self.window_size[1] - 1), self.window_size[1], dtype=torch.float32)
        relative_coords_w = torch.arange(-(self.window_size[2] - 1), self.window_size[2], dtype=torch.float32)
        relative_coords_table = torch.stack(
            torch.meshgrid([relative_coords_d,
                            relative_coords_h,
                            relative_coords_w], indexing='ij')).permute(1, 2, 3, 0).contiguous().unsqueeze(0)  # 1, 2*Wd-1, 2*Wh-1, 2*Ww-1, 3
        relative_coords_table[:, :, :, 0] /= (self.window_size[0] - 1)
        relative_coords_table[:, :, :, 1] /= (self.window_size[1] - 1)
        relative_coords_table[:, :, :, 2] /= (self.window_size[2] - 1)
        relative_coords_table *= 8  # normalize to -8, 8
        relative_coords_table = torch.sign(relative_coords_table) * torch.log2(
            torch.abs(relative_coords_table) + 1.0) / np.log2(8)

        self.register_buffer("relative_coords_table", relative_coords_table)

        # get pair-wise relative position index for each token inside the window
        coords_d = torch.arange(self.window_size[0])
        coords_h = torch.arange(self.window_size[1])
        coords_w = torch.arange(self.window_size[2])
        coords = torch.stack(torch.meshgrid([coords_d, coords_h, coords_w], indexing='ij')) # 3, Wd, Wh, Ww
        coords_flatten = torch.flatten(coords, 1)                                           # 3, Wd*Wh*Ww
        relative_coords = coords_flatten[:, :, None] - coords_flatten[:, None, :]           # 3, Wd*Wh*Ww, Wd*Wh*Ww
        relative_coords = relative_coords.permute(1, 2, 0).contiguous()                     # Wd*Wh*Ww, Wd*Wh*Ww, 3
        relative_coords[:, :, 0] += self.window_size[0] - 1
        relative_coords[:, :, 1] += self.window_size[1] - 1
        relative_coords[:, :, 2] += self.window_size[2] - 1
        relative_coords[:, :, 0] *= (2 * self.window_size[1] - 1) * (2 * self.window_size[2] - 1)
        relative_coords[:, :, 1] *= 2 * self.window_size[2] - 1
        relative_position_index = relative_coords.sum(-1)                           # Wd*Wh*Ww, Wd*Wh*Ww
        self.register_buffer("relative_position_index", relative_position_index)

        self.qkv = nn.Linear(dim, dim * 3, bias=False)
        if qkv_bias:
            self.q_bias = nn.Parameter(torch.zeros(dim))
            self.v_bias = nn.Parameter(torch.zeros(dim))
        else:
            self.q_bias = None
            self.v_bias = None
        self.attn_drop = nn.Dropout(attn_drop)
        self.proj = nn.Linear(dim, dim)
        self.proj_drop = nn.Dropout(proj_drop)
        self.softmax = nn.Softmax(dim=-1)

    def forward(self, x, mask=None):
        """
        Args:
            x: input features with shape of (num_windows*B, N, C)
            mask: (0/-inf) mask with shape of (num_windows, Wh*Ww, Wh*Ww) or None
        """
        B_, N, C = x.shape
        qkv_bias = None
        if self.q_bias is not None:
            qkv_bias = torch.cat((self.q_bias, torch.zeros_like(self.v_bias, requires_grad=False), self.v_bias))
        qkv = F.linear(input=x, weight=self.qkv.weight, bias=qkv_bias)
        qkv = qkv.reshape(B_, N, 3, self.num_heads, -1).permute(2, 0, 3, 1, 4)
        q, k, v = qkv[0], qkv[1], qkv[2]  # make torchscript happy (cannot use tensor as tuple)

        # cosine attention
        attn = (F.normalize(q, dim=-1) @ F.normalize(k, dim=-1).transpose(-2, -1))
        logit_scale = torch.clamp(self.logit_scale, max=torch.log(torch.tensor(1. / 0.01, device=x.device))).exp()
        attn = attn * logit_scale

        num_windows = self.window_size[0] * self.window_size[1] * self.window_size[2]   # Wd*Wh*Ww
        relative_position_bias_table = self.cpb_mlp(self.relative_coords_table).view(-1, self.num_heads)
        relative_position_bias = relative_position_bias_table[self.relative_position_index.view(-1)].view(
            num_windows, num_windows, -1)  # Wd*Wh*Ww, Wd*Wh*Ww, nH
        relative_position_bias = relative_position_bias.permute(2, 0, 1).contiguous()  # nH, Wd*Wh*Ww, Wd*Wh*Ww
        relative_position_bias = 16 * torch.sigmoid(relative_position_bias)
        attn = attn + relative_position_bias.unsqueeze(0)

        if mask is not None:
            nW = mask.shape[0]
            attn = attn.view(B_ // nW, nW, self.num_heads, N, N) + mask.unsqueeze(1).unsqueeze(0)
            attn = attn.view(-1, self.num_heads, N, N)
            attn = self.softmax(attn)
        else:
            attn = self.softmax(attn)

        attn = self.attn_drop(attn)

        x = (attn @ v).transpose(1, 2).reshape(B_, N, C)
        x = self.proj(x)
        x = self.proj_drop(x)
        return x

    def extra_repr(self) -> str:
        return f'dim={self.dim}, window_size={self.window_size}, num_heads={self.num_heads} '

class SwinTransformerBlock(nn.Module):
    r""" Swin Transformer Block.

    Args:
        dim (int): Number of input channels.
        input_resolution (tuple[int]): Input resulotion.
        num_heads (int): Number of attention heads.
        window_size (int): Window size.
        shift_size (int): Shift size for SW-MSA.
        mlp_ratio (float): Ratio of mlp hidden dim to embedding dim.
        qkv_bias (bool, optional): If True, add a learnable bias to query, key, value. Default: True
        drop (float, optional): Dropout rate. Default: 0.0
        attn_drop (float, optional): Attention dropout rate. Default: 0.0
        drop_path (float, optional): Stochastic depth rate. Default: 0.0
        act_layer (nn.Module, optional): Activation layer. Default: nn.GELU
        norm_layer (nn.Module, optional): Normalization layer.  Default: nn.LayerNorm
    """

    def __init__(
            self,
            dim: int,
            input_resolution: Tuple[int, int, int],
            num_heads: int,
            window_size: int = 7,
            shift_size:int = 0,
            mlp_ratio:float = 4.,
            qkv_bias: bool = True,
            drop:float = 0.,
            attn_drop:float = 0.,
            drop_path:float = 0.,
            act_layer: Type[nn.Module] = nn.GELU,
            norm_layer: Type[nn.Module] = nn.LayerNorm,
    ):
        super().__init__()
        self.dim = dim
        self.input_resolution = input_resolution
        self.num_heads = num_heads
        self.window_size = window_size
        self.shift_size = shift_size
        self.mlp_ratio = mlp_ratio
        if min(self.input_resolution) <= self.window_size:
            # if window size is larger than input resolution, we don't partition windows
            self.shift_size = 0
            self.window_size = min(self.input_resolution)
        assert 0 <= self.shift_size < self.window_size, "shift_size must in 0-window_size"

        self.norm1 = norm_layer(dim)
        self.attn = WindowAttention(
            dim, window_size=to_3tuple(self.window_size), num_heads=num_heads,
            qkv_bias=qkv_bias, attn_drop=attn_drop, proj_drop=drop
        )

        self.drop_path = DropPath(drop_path) if drop_path > 0. else nn.Identity()
        self.norm2 = norm_layer(dim)
        mlp_hidden_dim = int(dim * mlp_ratio)
        self.mlp = Mlp(in_features=dim, hidden_features=mlp_hidden_dim, act_layer=act_layer, drop=drop)

        if self.shift_size > 0:
            # calculate attention mask for SW-MSA
            D, H, W = self.input_resolution
            img_mask = torch.zeros((1, D, H, W, 1))  # 1 D H W 1
            d_slices = (slice(0, -self.window_size),
                        slice(-self.window_size, -self.shift_size),
                        slice(-self.shift_size, None))
            h_slices = (slice(0, -self.window_size),
                        slice(-self.window_size, -self.shift_size),
                        slice(-self.shift_size, None))
            w_slices = (slice(0, -self.window_size),
                        slice(-self.window_size, -self.shift_size),
                        slice(-self.shift_size, None))
            cnt = 0
            for d in d_slices :
                for h in h_slices:
                    for w in w_slices:
                        img_mask[:, d, h, w, :] = cnt
                        cnt += 1

            mask_windows = window_partition(img_mask, self.window_size)  # nW, window_size, window_size, window_size, 1
            mask_windows = mask_windows.view(-1, self.window_size * self.window_size * self.window_size)
            attn_mask = mask_windows.unsqueeze(1) - mask_windows.unsqueeze(2)
            attn_mask = attn_mask.masked_fill(attn_mask != 0, float(-100.0)).masked_fill(attn_mask == 0, float(0.0))
        else:
            attn_mask = None

        self.register_buffer("attn_mask", attn_mask)

    def forward(self, x):
        D, H, W = self.input_resolution
        B, L, C = x.shape
        assert L == D * H * W, "input feature has wrong size"

        shortcut = x
        x = x.view(B, D, H, W, C)

        # cyclic shift
        if self.shift_size > 0:
            shifted_x = torch.roll(x, shifts=(-self.shift_size, -self.shift_size), dims=(1, 2))
        else:
            shifted_x = x

        # partition windows
        x_windows = window_partition(shifted_x, self.window_size)  # nW*B, window_size, window_size, window_size, C
        x_windows = x_windows.view(-1, self.window_size * self.window_size * self.window_size, C)  # nW*B, window_size*window_size*window_size, C

        # W-MSA/SW-MSA
        attn_windows = self.attn(x_windows, mask=self.attn_mask)  # nW*B, window_size*window_size*window_size, C

        # merge windows
        attn_windows = attn_windows.view(-1, self.window_size, self.window_size, self.window_size, C)
        shifted_x = window_reverse(attn_windows, self.window_size, D, H, W)  # B D' H' W' C

        # reverse cyclic shift
        if self.shift_size > 0:
            x = torch.roll(shifted_x, shifts=(self.shift_size, self.shift_size), dims=(1, 2))
        else:
            x = shifted_x
        x = x.view(B, D * H * W, C)
        x = shortcut + self.drop_path(self.norm1(x))

        # FFN
        x = x + self.drop_path(self.norm2(self.mlp(x)))

        return x

    def extra_repr(self) -> str:
        return f"dim={self.dim}, input_resolution={self.input_resolution}, num_heads={self.num_heads}, " \
               f"window_size={self.window_size}, shift_size={self.shift_size}, mlp_ratio={self.mlp_ratio}"

class PatchMerging(nn.Module):
    r""" Patch Merging Layer.

    Args:
        input_resolution (tuple[int]): Resolution of input feature.
        dim (int): Number of input channels.
        norm_layer (nn.Module, optional): Normalization layer.  Default: nn.LayerNorm
    """

    def __init__(
            self,
            input_resolution: Tuple[int, int, int],
            dim: int,
            norm_layer: Type[nn.Module] = nn.LayerNorm
        ):
        super().__init__()
        self.input_resolution = input_resolution
        self.dim = dim
        self.reduction = nn.Linear(8 * dim, 2 * dim, bias=False)
        self.norm = norm_layer(2 * dim)

    def forward(self, x):
        """
        x: B, D*H*W, C
        """
        D, H, W = self.input_resolution
        B, L, C = x.shape
        assert L == D * H * W, "input feature has wrong size"
        assert D % 2 == 0 and H % 2 == 0 and W % 2 == 0, f"x size ({D}*{H}*{W}) are not even."

        x = x.view(B, D, H, W, C)

        x0 = x[:, 0::2, 0::2, 0::2, :]  # B, D/2, H/2, W/2, C
        x1 = x[:, 1::2, 0::2, 0::2, :]  # B, D/2, H/2, W/2, C
        x2 = x[:, 0::2, 1::2, 0::2, :]  # B, D/2, H/2, W/2, C
        x3 = x[:, 1::2, 1::2, 0::2, :]  # B, D/2, H/2, W/2, C
        x4 = x[:, 0::2, 0::2, 1::2, :]  # B, D/2, H/2, W/2, C
        x5 = x[:, 1::2, 0::2, 1::2, :]  # B, D/2, H/2, W/2, C
        x6 = x[:, 0::2, 1::2, 1::2, :]  # B, D/2, H/2, W/2, C
        x7 = x[:, 1::2, 1::2, 1::2, :]  # B, D/2, H/2, W/2, C
        x = torch.cat([x0, x1, x2, x3, x4, x5, x6, x7], -1)  # B, D/2, H/2, W/2, 8*C
        x = x.view(B, -1, 8 * C)        # B, D/2*H/2*W/2, 8*C

        x = self.reduction(x)           # B, D/2, H/2, W/2, 2*C
        x = self.norm(x)

        return x

    def extra_repr(self) -> str:
        return f"input_resolution={self.input_resolution}, dim={self.dim}"

class BasicLayer(nn.Module):
    """ A basic Swin Transformer layer for one stage.

    Args:
        dim (int): Number of input channels.
        input_resolution (tuple[int]): Input resolution.
        depth (int): Number of blocks.
        num_heads (int): Number of attention heads.
        window_size (int): Local window size.
        mlp_ratio (float): Ratio of mlp hidden dim to embedding dim.
        qkv_bias (bool, optional): If True, add a learnable bias to query, key, value. Default: True
        drop (float, optional): Dropout rate. Default: 0.0
        attn_drop (float, optional): Attention dropout rate. Default: 0.0
        drop_path (float | tuple[float], optional): Stochastic depth rate. Default: 0.0
        norm_layer (nn.Module, optional): Normalization layer. Default: nn.LayerNorm
        downsample (nn.Module | None, optional): Downsample layer at the end of the layer. Default: None
    """

    def __init__(
            self,
            dim: int,
            input_resolution: Tuple[int, int, int],
            depth: int,
            num_heads: int,
            window_size: int,
            mlp_ratio: float = 4.,
            qkv_bias: bool = True,
            drop: float = 0.,
            attn_drop: float = 0.,
            drop_path:float = 0.,
            norm_layer: Type[nn.Module] = nn.LayerNorm,
            downsample: Optional[Type[nn.Module]] = None
        ):

        super().__init__()
        self.dim = dim
        self.input_resolution = input_resolution
        self.depth = depth

        # build blocks
        self.blocks = nn.ModuleList([
            SwinTransformerBlock(dim=dim, input_resolution=input_resolution,
                                 num_heads=num_heads, window_size=window_size,
                                 shift_size=0 if (i % 2 == 0) else window_size // 2,
                                 mlp_ratio=mlp_ratio,
                                 qkv_bias=qkv_bias,
                                 drop=drop, attn_drop=attn_drop,
                                 drop_path=drop_path[i] if isinstance(drop_path, list) else drop_path,
                                 norm_layer=norm_layer)
            for i in range(depth)])

        # patch merging layer
        if downsample is not None:
            self.downsample = downsample(input_resolution, dim=dim, norm_layer=norm_layer)
        else:
            self.downsample = None

    def forward(self, x: Tensor) -> Tuple[Tensor, int, int, int, Tensor, int, int, int]:
        for blk in self.blocks :
            x = blk(x)
        if self.downsample is not None:
            x_down = self.downsample(x)
            return x, x_down
        else:
            return x, x

    def extra_repr(self) -> str:
        return f"dim={self.dim}, input_resolution={self.input_resolution}, depth={self.depth}"

    def _init_respostnorm(self):
        for blk in self.blocks:
            nn.init.constant_(blk.norm1.bias, 0)
            nn.init.constant_(blk.norm1.weight, 0)
            nn.init.constant_(blk.norm2.bias, 0)
            nn.init.constant_(blk.norm2.weight, 0)


class PatchEmbed(nn.Module):
    r""" Image to Patch Embedding

    Args:
        img_size (int): Image size.  Default: 48.
        patch_size (int): Patch token size. Default: 4.
        in_chans (int): Number of input image channels. Default: 3.
        embed_dim (int): Number of linear projection output channels. Default: 96.
        norm_layer (nn.Module, optional): Normalization layer. Default: None
    """

    def __init__(
            self,
            img_size: int = 48,
            patch_size: int = 4,
            in_chans: int = 3,
            embed_dim: int = 96,
            norm_layer: Optional[Type[nn.Module]] = None
    ):
        super().__init__()
        img_size = to_3tuple(img_size)
        patch_size = to_3tuple(patch_size)
        patches_resolution = [img_size[0] // patch_size[0], img_size[1] // patch_size[1], img_size[2] // patch_size[2]]
        self.img_size = img_size
        self.patch_size = patch_size
        self.patches_resolution = patches_resolution
        self.num_patches = patches_resolution[0] * patches_resolution[1] * patches_resolution[2]

        self.in_chans = in_chans
        self.embed_dim = embed_dim

        self.proj = nn.Conv3d(in_chans, embed_dim, kernel_size=patch_size, stride=patch_size)
        if norm_layer is not None:
            self.norm = norm_layer(embed_dim)
        else:
            self.norm = None

    def forward(self, x):
        B, C, D, H, W = x.shape
        # FIXME look at relaxing size constraints
        assert D == self.img_size[0] and H == self.img_size[1] and W == self.img_size[2], \
            f"Input image size ({D}*{H}*{W}) doesn't match model ({self.img_size[0]}*{self.img_size[1]}*{self.img_size[2]})."
        x = self.proj(x).flatten(2).transpose(1, 2)  # B Pd*Ph*Pw C
        if self.norm is not None:
            x = self.norm(x)
        return x


@BACKBONE.register()
class SwinTransformerV2(nn.Module):
    r""" Swin Transformer
        A PyTorch impl of : `Swin Transformer: Hierarchical Vision Transformer using Shifted Windows`  -
          https://arxiv.org/pdf/2103.14030

    Args:
        in_channels (int): Number of input image channels.
        image_size (int): Input image size,
            used in absolute postion embedding. Default 48.
        patch_size (int | tuple(int)): Patch size. Default: 2
        embed_dim (int): Patch embedding dimension. Default: 96
        depths (tuple(int)): Depth of each Swin Transformer layer.
        num_heads (tuple(int)): Number of attention heads in different layers.
        window_size (int): Window size. Default: 7
        mlp_ratio (float): Ratio of mlp hidden dim to embedding dim. Default: 4
        qkv_bias (bool): If True, add a learnable bias to query, key, value. Default: True
        drop_rate (float): Dropout rate. Default: 0
        attn_drop_rate (float): Attention dropout rate. Default: 0
        drop_path_rate (float): Stochastic depth rate. Default: 0.1
        norm_layer (nn.Module): Normalization layer. Default: nn.LayerNorm.
        ape (bool): If True, add absolute position embedding to the patch embedding. Default: False
        patch_norm (bool): If True, add normalization after patch embedding. Default: True
        use_checkpoint (bool): Whether to use checkpointing to save memory. Default: False
        pretrained_window_sizes (tuple(int)): Pretrained window sizes of each layer.
    """

    def __init__(
            self,
            in_channels: int,
            image_size: int = 48,
            patch_size: int = 4,
            embed_dim: int = 96,
            depths: Sequence[int] = [2, 2, 6, 2],
            num_heads: Sequence[int] = [3, 6, 12, 24],
            window_size: int = 7,
            mlp_ratio: float = 4.,
            qkv_bias: bool = True,
            drop_rate: float = 0.,
            attn_drop_rate: float = 0.,
            drop_path_rate: float = 0.1,
            norm_layer: Type[nn.Module] = nn.LayerNorm,
            ape: bool = False,
            patch_norm: bool = True,
            out_indices: Union[int, Sequence[int]] = [0, 1, 2, 3],
        ) :
        super().__init__()

        self.image_size = image_size
        self.num_layers = len(depths)
        self.embed_dim = embed_dim
        self.ape = ape
        self.patch_norm = patch_norm
        self.mlp_ratio = mlp_ratio
        out_indices = out_indices if not isinstance(out_indices, int) else [out_indices]
        self.out_indices = [(index if index >= 0 else self.num_layers + index) for index in out_indices]

        # split image into non-overlapping patches
        self.patch_embed = PatchEmbed(
            img_size=image_size, patch_size=patch_size, in_chans=in_channels, embed_dim=embed_dim,
            norm_layer=norm_layer if self.patch_norm else None)
        num_patches = self.patch_embed.num_patches
        patches_resolution = self.patch_embed.patches_resolution
        self.patches_resolution = patches_resolution

        # absolute position embedding
        if self.ape:
            self.absolute_pos_embed = nn.Parameter(torch.zeros(1, num_patches, embed_dim))
            trunc_normal_(self.absolute_pos_embed, std=.02)

        self.pos_drop = nn.Dropout(p=drop_rate)

        # stochastic depth
        dpr = [x.item() for x in torch.linspace(0, drop_path_rate, sum(depths))]  # stochastic depth decay rule

        # build layers
        self.layers = nn.ModuleList()
        for i_layer in range(self.num_layers):
            layer = BasicLayer(dim=int(embed_dim * 2 ** i_layer),
                               input_resolution=(patches_resolution[0] // (2 ** i_layer),
                                                 patches_resolution[1] // (2 ** i_layer),
                                                 patches_resolution[2] // (2 ** i_layer)),
                               depth=depths[i_layer],
                               num_heads=num_heads[i_layer],
                               window_size=window_size,
                               mlp_ratio=self.mlp_ratio,
                               qkv_bias=qkv_bias,
                               drop=drop_rate, attn_drop=attn_drop_rate,
                               drop_path=dpr[sum(depths[:i_layer]):sum(depths[:i_layer + 1])],
                               norm_layer=norm_layer,
                               downsample=PatchMerging if (i_layer < self.num_layers - 1) else None)
            self.layers.append(layer)

        # add a norm layer for each output
        num_features = [int(embed_dim * 2 ** i) for i in range(self.num_layers)]

        for i_layer in out_indices:
            layer = norm_layer(num_features[i_layer])
            layer_name = f'norm{i_layer}'
            self.add_module(layer_name, layer)

    def initialize_weights(self) :
        def _init_weights(m):
            if isinstance(m, nn.Linear):
                trunc_normal_(m.weight, std=.02)
                if m.bias is not None:
                    nn.init.constant_(m.bias, 0)
            elif isinstance(m, nn.LayerNorm):
                nn.init.constant_(m.bias, 0)
                nn.init.constant_(m.weight, 1.0)

        self.apply(_init_weights)
        for bly in self.layers:
            bly._init_respostnorm()

    def forward(self, x: Tensor) -> List[Tensor]:
        x = self.patch_embed(x)
        if self.ape:
            x = x + self.absolute_pos_embed
        x = self.pos_drop(x)

        outs = []
        for i in range(self.num_layers):
            layer = self.layers[i]
            x_out, x = layer(x)                 # B, Wd, Wh, Ww, C' | B, Wd', Wh', Ww', C''
            if i in self.out_indices:
                B, L, C = x_out.shape
                Wd = Wh = Ww = round(L**(1/3))
                norm_layer = getattr(self, f'norm{i}')
                x_out = x_out.view(B, Wd, Wh, Ww, C)
                x_out = norm_layer(x_out)                           # B, Wd, Wh, Ww, C
                x_out = x_out.permute(0, 4, 1, 2, 3).contiguous()   # B, C, Wd, Wh, Ww
                outs.append(x_out)
        return outs

src/network/backbones/swinv2.py
SeonghwanSeo/PharmacoNet
from torch import nn
from typing import Dict

from .utils.registry import Registry

BACKBONE = Registry('Backbone')
NECK = Registry('Neck')
DECODER = Registry('Decoder')

EMBEDDING = Registry('Embedding')
HEAD = Registry('Head')

MODEL = Registry('Model')


def build_model(config: Dict) -> nn.Module:
    registry_key = 'registry'
    module_key = 'name'
    return Registry.build_from_config(config, registry_key, module_key, convert_key_to_lower_case=True, safe_build=True)

src/network/builder.py
SeonghwanSeo/PharmacoNet
from torch import nn

from functools import partial

from typing import Tuple, Type, Optional
from torch import Tensor

from .builder import HEAD
from .nn import BaseConv3d


@HEAD.register()
class CavityHead(nn.Module):
    def __init__(
        self,
        feature_dim: int,
        hidden_dim: int,
        norm_layer: Optional[Type[nn.Module]] = nn.BatchNorm3d,
        act_layer: Optional[Type[nn.Module]] = partial(nn.ReLU, inplace=True),
    ):
        super(CavityHead, self).__init__()

        self.short_head = nn.Sequential(
            BaseConv3d(feature_dim, hidden_dim, kernel_size=3, norm_layer=norm_layer, act_layer=act_layer),
            BaseConv3d(hidden_dim, 1, kernel_size=1, norm_layer=None, act_layer=None),
        )
        self.long_head = nn.Sequential(
            BaseConv3d(feature_dim, hidden_dim, kernel_size=3, norm_layer=norm_layer, act_layer=act_layer),
            BaseConv3d(hidden_dim, 1, kernel_size=1, norm_layer=None, act_layer=None),
        )

    def initialize_weights(self):
        for m in self.short_head.children():
            m.initialize_weights()
        for m in self.long_head.children():
            m.initialize_weights()

    def forward(
        self,
        features: Tensor,
    ) -> Tuple[Tensor, Tensor]:
        """Pocket Extraction Function

        Args:
            features: FloatTensor [N, F, D, H, W]

        Returns:
            focus_area_short: FloatTensor [N, 1, D, H, W]
            focus_area_long: FloatTensor [N, 1, D, H, W]
        """
        focus_area_short = self.short_head(features)
        focus_area_long = self.long_head(features)
        return focus_area_short, focus_area_long

src/network/cavity_head.py
SeonghwanSeo/PharmacoNet
from .fpn_decoder import FPNDecoder

src/network/decoders/__init__.py
SeonghwanSeo/PharmacoNet
import torch.nn as nn
import torch.nn.functional as F

from typing import Sequence, Optional, Type, List
from torch import Tensor

from ..nn.layers import BaseConv3d
from ..builder import DECODER


@DECODER.register()
class FPNDecoder(nn.Module):
    """
    Modified FPN Structure [https://arxiv.org/abs/1807.10221]
    feature_channels: Bottom-Up Manner.
    """

    def __init__(
        self,
        feature_channels: Sequence[int],
        num_convs: Sequence[int],
        channels: int = 64,
        interpolate_mode: str = 'nearest',
        align_corners: bool = False,
        norm_layer: Optional[Type[nn.Module]] = nn.BatchNorm3d,
        act_layer: Optional[Type[nn.Module]] = nn.ReLU,
    ):
        super(FPNDecoder, self).__init__()
        self.feature_channels = feature_channels
        self.interpolate_mode = interpolate_mode
        if interpolate_mode == 'trilinear':
            self.align_corners = align_corners
        else:
            self.align_corners = None
        self.channels = channels

        lateral_conv_list = []
        fpn_convs_list = []
        for level, (channels, num_conv) in enumerate(zip(self.feature_channels, num_convs)):
            if level == (len(self.feature_channels) - 1):  # Lowest-Resolution Channels (Top)
                lateral_conv = nn.Identity()
                fpn_convs = nn.Sequential(*[
                    BaseConv3d(
                        channels if i == 0 else self.channels, self.channels,
                        kernel_size=3, norm_layer=norm_layer, act_layer=act_layer,
                    ) for i in range(num_conv)
                ])
            else:
                lateral_conv = BaseConv3d(channels, self.channels, kernel_size=1, norm_layer=norm_layer, act_layer=act_layer)
                fpn_convs = nn.Sequential(*[
                    BaseConv3d(
                        self.channels, self.channels,
                        kernel_size=3, norm_layer=norm_layer, act_layer=act_layer,
                    ) for _ in range(num_conv)
                ])
            lateral_conv_list.append(lateral_conv)
            fpn_convs_list.append(fpn_convs)

        self.lateral_conv_list = nn.ModuleList(lateral_conv_list)
        self.fpn_convs_list = nn.ModuleList(fpn_convs_list)

    def initialize_weights(self):
        for m in self.lateral_conv_list:
            if isinstance(m, BaseConv3d):
                m.initialize_weights()
        for seqm in self.fpn_convs_list:
            for m in seqm.children():
                m.initialize_weights()

    def forward(self, features: Sequence[Tensor]) -> List[Tensor]:
        """Forward function.
        Args:
            features: Bottom-Up, [Highest-Resolution Feature Map, ..., Lowest-Resolution Feature Map]
        Returns:
            features: Top-Down, [Lowest-Resolution Feature Map, ..., Highest-Resolution Feature Map]
        """
        num_levels = len(features)
        assert num_levels == len(self.feature_channels)
        fpn = None
        multi_scale_features = []
        for level in range(num_levels - 1, -1, -1):
            feature = features[level]
            lateral_conv = self.lateral_conv_list[level]
            fpn_convs = self.fpn_convs_list[level]
            current_fpn = lateral_conv(feature)
            if level == (num_levels - 1):    # Top
                assert fpn is None
                fpn = current_fpn
            else:
                assert fpn is not None
                fpn = current_fpn + F.interpolate(fpn, size=current_fpn.size()[-3:], mode=self.interpolate_mode, align_corners=self.align_corners)
            fpn = fpn_convs(fpn)
            multi_scale_features.append(fpn)
        return multi_scale_features

src/network/decoders/fpn_decoder.py
SeonghwanSeo/PharmacoNet
import torch.nn as nn

from typing import Sequence, List, Optional, Tuple, Optional
from torch import Tensor, IntTensor

from .feature_embedding import FeaturePyramidNetwork
from .token_head import TokenHead
from .cavity_head import CavityHead
from .mask_head import MaskHead
from .builder import MODEL


@MODEL.register()
class PharmacoFormer(nn.Module):
    def __init__(
        self,
        embedding: FeaturePyramidNetwork,
        cavity_head: CavityHead,
        token_head: TokenHead,
        mask_head: MaskHead,
        num_interactions: int,
    ):
        super(PharmacoFormer, self).__init__()
        self.num_interactions = num_interactions
        self.embedding = embedding
        self.cavity_head = cavity_head
        self.token_head = token_head
        self.mask_head = mask_head

    def initialize_weights(self):
        self.embedding.initialize_weights()
        self.token_head.initialize_weights()
        self.mask_head.initialize_weights()

    def setup_train(self, criterion: nn.Module):
        self.criterion = criterion

    def forward_feature(self, in_image: Tensor) -> Tuple[Tensor, ...]:
        """Feature Embedding
        Args:
            in_image: FloatTensor [N, C, Din, Hin, Win]
        Returns:
            multi-scale features: [FloatTensor [N, F, D, H, W]]
        """
        return tuple(self.embedding.forward(in_image))

    def forward_cavity_extraction(self, features: Tensor) -> Tuple[Tensor, Tensor]:
        """Cavity Extraction
        Args:
            features: FloatTensor [N, F, Dout, Hout, Wout]
        Returns:
            cavity_narrow: FloatTensor [N, 1, Dout, Hout, Wout]
            cavity_wide: FloatTensor [N, 1, Dout, Hout, Wout]
        """
        return self.cavity_head.forward(features)

    def forward_token_prediction(
        self,
        features: Tensor,
        tokens_list: Sequence[IntTensor],
    ) -> Tuple[List[Tensor], List[Tensor]]:
        """token Selection Network

        Args:
            features: FloatTensor [N, F, Dout, Hout, Wout]
            tokens_list: List[IntTensor [Ntoken, 4] - (x, y, z, i)]

        Returns:
            token_scores_list: List[FloatTensor [Ntoken,] $\\in$ [0, 1]]
            token_features_list: List[FloatTensor [Ntoken, F]]
        """
        token_scores_list, token_features_list = self.token_head.forward(features, tokens_list)
        return token_scores_list, token_features_list

    def forward_segmentation(
        self,
        multi_scale_features: Tuple[Tensor, ...],
        box_tokens_list: Sequence[IntTensor],
        box_token_features_list: Sequence[Tensor],
        return_aux: bool = False,
    ) -> Tuple[List[Tensor], Optional[List[List[Tensor]]]]:
        """Mask Prediction

        Args:
            multi_scales_features: List[FloatTensor [N, F, D_scale, H_scale, W_scale]]
            box_tokens_list: List[IntTensor [Nbox, 4] - (x, y, z, i)]
            box_token_features_list: List[FloatTensor [Nbox, F]]

        Returns:
            box_masks_list: List[FloatTensor [Nbox, D, H, W]]
            aux_box_masks_list: List[List[FloatTensor [Nbox, D_scale, H_scale, W_scale]]]
        """
        return self.mask_head.forward(multi_scale_features, box_tokens_list, box_token_features_list, return_aux)

src/network/detector.py
SeonghwanSeo/PharmacoNet
from torch import nn

from typing import Sequence, Optional
from torch import Tensor

from .builder import EMBEDDING


@EMBEDDING.register()
class FeaturePyramidNetwork(nn.Module):
    def __init__(
        self,
        backbone: nn.Module,
        decoder: nn.Module,
        neck: Optional[nn.Module] = None,
        feature_indices: Optional[Sequence[int]] = None,
        set_input_to_bottom: bool = True,
    ):
        super(FeaturePyramidNetwork, self).__init__()
        self.backbone = backbone
        self.decoder = decoder
        self.feature_indices = feature_indices
        self.input_is_bottom = set_input_to_bottom

        if neck is not None:
            self.with_neck = True
            self.neck = neck
        else:
            self.with_neck = False

    def initialize_weights(self):
        self.backbone.initialize_weights()
        self.decoder.initialize_weights()
        if self.with_neck:
            self.neck.initialize_weights()

    def forward(self, in_image: Tensor) -> Sequence[Tensor]:
        """Feature Pyramid Network -> return multi-scale feature maps
        Args:
            in_image: (N, C, D, H, W)
        Returns:
            multi-scale features (top_down): [(N, F, D, H, W)]
        """
        bottom_up_features: Sequence[Tensor] = self.backbone(in_image)
        if self.feature_indices is not None:
            bottom_up_features = [bottom_up_features[index] for index in self.feature_indices]
        if self.input_is_bottom:
            bottom_up_features = [in_image, *bottom_up_features]
        if self.with_neck:
            bottom_up_features = self.neck(bottom_up_features)
        top_down_features = self.decoder(bottom_up_features)
        return top_down_features

src/network/feature_embedding.py
SeonghwanSeo/PharmacoNet
import torch
from torch import nn

from typing import Sequence, List, Optional, Tuple
from torch import Tensor

from .builder import HEAD


@HEAD.register()
class MaskHead(nn.Module):
    def __init__(
        self,
        decoder: nn.Module,
        token_feature_dim: int,
    ):
        super(MaskHead, self).__init__()
        feature_channels_list: List[int] = decoder.feature_channels
        self.point_mlp_list = nn.ModuleList(
            [nn.Linear(token_feature_dim, channels) for channels in feature_channels_list]
        )
        self.background_mlp_list = nn.ModuleList(
            [nn.Linear(token_feature_dim, channels) for channels in feature_channels_list]
        )
        self.decoder = decoder
        self.conv_logits = nn.Conv3d(decoder.channels, 1, kernel_size=1)

    def initialize_weights(self):
        def _init_weight(m):
            if isinstance(m, (nn.Linear, nn.Conv3d)):
                nn.init.normal_(m.weight, std=0.01)
                if m.bias is not None:
                    nn.init.constant_(m.bias, 0)
        for m in [self.conv_logits, self.point_mlp_list, self.background_mlp_list]:
            m.apply(_init_weight)

        self.decoder.initialize_weights()

    def forward(
        self,
        multi_scale_features: Tensor,
        tokens_list: Sequence[Tensor],
        token_features_list: Sequence[Tensor],
        return_aux: bool = False,
    ) -> Tuple[List[Tensor], Optional[List[List[Tensor]]]]:
        """Box Predicting Function

        Args:
            multi_scale_features: Top-Down, List[FloatTensor [N, F_scale, D_scale, H_scale, W_scale]]
            tokens_list: List[IntTensor [Nbox, 4] - (x, y, z, i)]
            token_features_list: List[FloatTensor[Nbox, Ftoken]

        Returns:
            masks_list: List[FloatTensor [Nbox, D, H, W]]
            aux_masks_list(optional): List[List[FloatTensor [Nbox, D, H, W]]]
        """
        num_images = len(tokens_list)
        assert len(multi_scale_features[0]) == num_images
        multi_scale_features = multi_scale_features[::-1]   # Top-Down -> Bottom-Up
        if return_aux:
            out_masks_list = []
            aux_masks_list = []
            for image_idx in range(num_images):
                out = self.do_predict_w_aux(
                    [multi_scale_features[level][image_idx] for level in range(len(multi_scale_features))],
                    tokens_list[image_idx],
                    token_features_list[image_idx],
                )
                out_masks_list.append(out[-1])
                aux_masks_list.append(out[:-1])
        else:
            aux_masks_list = None
            out_masks_list = [
                self.do_predict_single(
                    [multi_scale_features[level][image_idx] for level in range(len(multi_scale_features))],
                    tokens_list[image_idx],
                    token_features_list[image_idx],
                ) for image_idx in range(num_images)
            ]
        return out_masks_list, aux_masks_list

    def do_predict_w_aux(
        self,
        multi_scale_features: Sequence[Tensor],
        tokens: Tensor,
        token_features: Tensor,
    ) -> List[Tensor]:
        """Box Predicting Function

        Args:
            multi_scale_features: Bottom-Up, List[FloatTensor [F_scale, D_scale, H_scale, W_scale]]
            token_features: FloatTensor [Nbox, Ftoken]
            tokens: IntTensor [Nbox, 4] - (x, y, z, i)

        Returns:
            multi_scale_box_masks: List[FloatTensor [Nbox, D_scale, H_scale, W_scale]]
        """
        Nbox = tokens.size(0)
        multi_scale_size = [features.size()[1:] for features in multi_scale_features]
        if Nbox > 0:
            Dout, Hout, Wout = multi_scale_size[0]
            token_indices = torch.split(tokens, 1, dim=1)                           # (x_list, y_list, z_list, i_list)
            xs, ys, zs, _ = token_indices

            bottom_up_box_features = []
            for level in range(len(multi_scale_features)):
                features = multi_scale_features[level]
                _, D, H, W = features.shape
                _xs = torch.div(xs, Dout // D, rounding_mode='trunc')
                _ys = torch.div(ys, Hout // H, rounding_mode='trunc')
                _zs = torch.div(zs, Wout // W, rounding_mode='trunc')
                box_features = self.get_box_features(features, (_xs, _ys, _zs), token_features, level)
                bottom_up_box_features.append(box_features)

            top_down_features = self.decoder(bottom_up_box_features)
            top_down_box_masks = [self.conv_logits(features).squeeze(1) for features in top_down_features]
            return top_down_box_masks
        else:
            return [torch.empty((0, *size), dtype=multi_scale_features[0].dtype, device=tokens.device) for size in multi_scale_size[::-1]]

    def do_predict_single(
        self,
        multi_scale_features: Sequence[Tensor],
        tokens: Tensor,
        token_features: Tensor,
    ) -> Tensor:
        """Box Predicting Function

        Args:
            multi_scale_features: Bottom-Up, List[FloatTensor [F_scale, D_scale, H_scale, W_scale]]
            token_features: FloatTensor [Nbox, Ftoken]
            tokens: IntTensor [Nbox, 4] - (x, y, z, i)

        Returns:
            box_masks: FloatTensor [Nbox, D_out, H_out, W_out]
        """
        Nbox = tokens.size(0)
        multi_scale_size = [features.size()[1:] for features in multi_scale_features]
        Dout, Hout, Wout = multi_scale_size[0]
        if Nbox > 0:
            token_indices = torch.split(tokens, 1, dim=1)                           # (x_list, y_list, z_list, i_list)
            xs, ys, zs, _ = token_indices

            bottom_up_box_features = []
            for level in range(len(multi_scale_features)):
                features = multi_scale_features[level]
                _, D, H, W = features.shape
                _xs = torch.div(xs, Dout // D, rounding_mode='trunc')
                _ys = torch.div(ys, Hout // H, rounding_mode='trunc')
                _zs = torch.div(zs, Wout // W, rounding_mode='trunc')
                box_features = self.get_box_features(features, (_xs, _ys, _zs), token_features, level)
                bottom_up_box_features.append(box_features)

            top_down_features = self.decoder(bottom_up_box_features)
            return self.conv_logits(top_down_features[-1]).squeeze(1)
        else:
            return torch.empty((0, Dout, Hout, Wout), dtype=multi_scale_features[0].dtype, device=tokens.device)

    def get_box_features(
        self,
        features: Tensor,
        token_indices: Tuple[Tensor, Tensor, Tensor],
        token_features: Tensor,
        level: int
    ) -> Tensor:
        """Extract token features

        Args:
            features: FloatTensor [F_scale, D_scale, H_scale, W_scale]
            token_indices: Tuple[IntTensor [Nbox,], IntTensor [Nbox,], IntTensor[Nbox,]] - (xs, ys, zs)
            token_features: FloatTensor [Nbox, Ftoken]

        Returns:
            box_features: FloatTensor [Nbox, F_scale, D_scale, H_scale, W_scale]
        """
        F, D, H, W = features.shape
        xs, ys, zs = token_indices
        Nbox = token_features.size(0)
        Nboxs = torch.arange(Nbox, dtype=xs.dtype, device=xs.device)
        background_features = self.background_mlp_list[level](token_features)               # [Nbox, F]
        point_features = self.point_mlp_list[level](token_features)                         # [Nbox, F]
        box_features = background_features.view(Nbox, F, 1, 1, 1).repeat(1, 1, D, H, W)     # [Nbox, F, D, H, W]
        box_features[Nboxs, :, xs, ys, zs] += point_features
        features = features.unsqueeze(0) + box_features
        return features

src/network/mask_head.py
SeonghwanSeo/PharmacoNet
from .center_crop import CenterCrop, MultipleCenterCrop

src/network/necks/__init__.py
SeonghwanSeo/PharmacoNet
from torch import nn, Tensor

from typing import List, Tuple
from timm.models.layers import to_3tuple

from ..builder import NECK


@NECK.register()
class MultipleCenterCrop(nn.Module):
    def __init__(self, crop_sizes: List[int]):
        super().__init__()
        self.crop_sizes: List[Tuple[int, int, int]] = [to_3tuple(size) for size in crop_sizes]

    def forward(self, inputs: List[Tensor]) -> List[Tensor]:
        outputs: List[Tensor] = []
        for size, tensor in zip(self.crop_sizes, inputs):
            Dc, Hc, Wc = size
            _, _, D, H, W = tensor.size()
            d, h, w = (D - Dc) // 2, (H - Hc) // 2, (W - Wc) // 2
            assert d >= 0 and h >= 0 and w >= 0
            if d == 0 and h == 0 and w == 0:
                outputs.append(tensor)
            else:
                outputs.append(tensor[:, :, d:D - d, h:H - h, w:W - w].contiguous())
        return outputs

    def initialize_weights(self):
        pass


@NECK.register()
class CenterCrop(nn.Module):
    def __init__(self, crop_size: int):
        super().__init__()
        self.crop_size: Tuple[int, int, int] = to_3tuple(crop_size)

    def forward(self, input: Tensor) -> Tensor:
        Dc, Hc, Wc = self.crop_size
        _, _, D, H, W = input.size()
        d, h, w = (D - Dc) // 2, (H - Hc) // 2, (W - Wc) // 2
        assert d >= 0 and h >= 0 and w >= 0
        if d != 0 or h != 0 or w != 0:
            return input[:, :, d:D - d, h:H - h, w:W - w].contiguous()
        else:
            return input

    def initialize_weights(self):
        pass

src/network/necks/center_crop.py
SeonghwanSeo/PharmacoNet
from torch import nn

from typing import Optional, Type


class BaseConv3d(nn.Module):
    def __init__(
            self,
            in_channels: int,
            out_channels: int,
            kernel_size: int,
            stride: int = 1,
            padding: int = None,
            dilation: int = 1,
            groups: int = 1,
            norm_layer: Optional[Type[nn.Module]] = nn.BatchNorm3d,
            act_layer: Optional[Type[nn.Module]] = nn.ReLU,
    ):
        super().__init__()
        if padding is None:
            padding = (kernel_size - 1) // 2
        bias = (norm_layer is None)
        self._conv = nn.Conv3d(in_channels, out_channels, kernel_size, stride, padding, dilation, groups, bias)
        self._norm = norm_layer(out_channels) if norm_layer is not None else nn.Identity()
        self._act = act_layer() if act_layer is not None else nn.Identity()

    def initialize_weights(self):
        if isinstance(self._act, nn.LeakyReLU):
            a = self._act.negative_slope
            nn.init.kaiming_normal_(self._conv.weight, a, mode='fan_out', nonlinearity='leaky_relu')
        else:
            nn.init.kaiming_normal_(self._conv.weight, mode='fan_out', nonlinearity='relu')
        if self._conv.bias is not None:
            nn.init.constant_(self._conv.bias, 0.)
        if not isinstance(self._norm, nn.Identity):
            nn.init.constant_(self._norm.weight, 1.0)

    def forward(self, x):
        return self._act(self._norm(self._conv(x)))

src/network/nn/layers.py
SeonghwanSeo/PharmacoNet
import torch
from torch import nn

from typing import Sequence, List, Tuple
from torch import Tensor

from .builder import HEAD


@HEAD.register()
class TokenHead(nn.Module):
    def __init__(
        self,
        feature_dim: int,
        num_interactions: int,
        token_feature_dim: int,
        num_feature_mlp_layers: int,
        num_score_mlp_layers: int,
    ):
        super(TokenHead, self).__init__()
        self.interaction_embedding = nn.Embedding(num_interactions, feature_dim)
        self.token_feature_dim = token_feature_dim

        feature_mlp = []
        dim = 2 * feature_dim
        for _ in range(num_feature_mlp_layers):
            feature_mlp.append(nn.Linear(dim, token_feature_dim))
            feature_mlp.append(nn.SiLU(inplace=True))
            dim = token_feature_dim
        self.feature_mlp = nn.Sequential(*feature_mlp)
        if 2 * feature_dim != token_feature_dim:
            self.skip = nn.Linear(2 * feature_dim, token_feature_dim)
        else:
            self.skip = nn.Identity()

        score_mlp = []
        for _ in range(num_score_mlp_layers - 1):
            score_mlp.append(nn.Linear(token_feature_dim, token_feature_dim))
            score_mlp.append(nn.ReLU(inplace=True))
        score_mlp.append(nn.Linear(token_feature_dim, 1))
        self.score_mlp = nn.Sequential(*score_mlp)

    def initialize_weights(self):
        def _init_weight(m):
            if isinstance(m, nn.Linear):
                nn.init.normal_(m.weight, std=.01)
                if m.bias is not None:
                    nn.init.constant_(m.bias, 0)

        nn.init.uniform_(self.interaction_embedding.weight, -1.0, 1.0)
        for m in [self.feature_mlp, self.score_mlp, self.skip]:
            m.apply(_init_weight)

    def forward(self, features: Tensor, tokens_list: Sequence[Tensor]) -> Tuple[List[Tensor], List[Tensor]]:
        """Token Scoring Function

        Args:
            features: FloatTensor [N, F, D, H, W]
            tokens_list: List[IntTensor [Ntoken, 4] - (x, y, z, i)]

        Returns:
            token_scores_list: List[FloatTensor [Ntoken,]]
            token_features_list: List[FloatTensor [Ntoken, F]]
        """
        num_images = len(tokens_list)
        token_features_list = [self.extract_token_features(features[idx], tokens_list[idx]) for idx in range(num_images)]
        token_scores_list = [self.score_mlp(token_features).squeeze(-1) for token_features in token_features_list]
        return token_scores_list, token_features_list

    def extract_token_features(self, features: Tensor, tokens: Tensor) -> Tensor:
        """Extract token features

        Args:
            features: FloatTensor [D, H, W, F]
            tokens: IntTensor [Ntoken, 4] - (x, y, z, i)

        Returns:
            token_features: FloatTensor [Ntoken, Fh]
        """
        if tokens.size(0) == 0:
            return torch.empty([0, self.token_feature_dim], dtype=torch.float, device=features.device)
        else:
            features = features.permute(1, 2, 3, 0).contiguous()                    # [D, H, W, F]
            x_list, y_list, z_list, i_list = torch.split(tokens, 1, dim=1)          # (x_list, y_list, z_list, i_list)
            token_features = features[x_list, y_list, z_list].squeeze(1)            # [Ntoken, F]
            embeddings = self.interaction_embedding(i_list).squeeze(1)              # [Ntoken, F]
            token_features = torch.cat([token_features, embeddings], dim=1)         # [Ntoken, 2F]
            return self.skip(token_features) + self.feature_mlp(token_features)     # [Ntoken, Fh]

src/network/token_head.py
SeonghwanSeo/PharmacoNet
from typing import Any, Dict, Optional, Callable
from collections.abc import MutableMapping
import copy


def convert_dict_key_to_lower_case(dictionary: MutableMapping) -> Dict:
    out = {}
    for key, value in dictionary.items():
        if isinstance(value, MutableMapping):
            value = convert_dict_key_to_lower_case(value)
        out[key.lower()] = value
    return out


class Registry():
    __OBJ_DICT__ = {}

    def __init__(self, name: str):
        self._name = name
        self._module_dict: Dict[str, Any] = dict()
        self.__OBJ_DICT__[name] = self

    def __len__(self):
        return len(self._module_dict)

    def __contains__(self, name: str) -> bool:
        return name in self._module_dict

    @property
    def name(self):
        return self._name

    def get(self, name: str) -> Any:
        module = self._module_dict.get(name, None)
        if module is None:
            raise KeyError(
                f"No object named '{name}' found in '{self._name}' registry"
                + f"Registry: {set(self._module_dict.keys())}"
            )
        return module

    def _do_register(self, name: str, obj: Any) -> None:
        assert (
            name not in self._module_dict
        ), "An object named '{}' was already registered in '{}' registry!".format(
            name, self._name
        )
        self._module_dict[name] = obj

    def register(self, module: Optional[Any] = None) -> Any:
        if module is None:
            def deco(_module: Any) -> Any:
                name = _module.__name__
                self._do_register(name, _module)
                return _module
            return deco
        else:
            name = module.__name__
            self._do_register(name, module)
            return module

    @classmethod
    def build_from_config(
        cls,
        config: MutableMapping,
        registry_key: str,
        module_key: str,
        convert_key_to_lower_case: bool = True,
        safe_build: bool = True
    ):
        _config = config
        _registry_key = registry_key
        _module_key = module_key
        if convert_key_to_lower_case:
            config = convert_dict_key_to_lower_case(_config)
            registry_key = _registry_key.lower()
            module_key = _module_key.lower()
        elif safe_build:
            config = copy.deepcopy(_config)
        else:
            config = _config
        assert registry_key in config, f'registry key ({_registry_key}) not in config (keys: {set(_config.keys())})'
        assert module_key in config, f'module key ({_module_key}) not in config ({set(_config.keys())})'
        return cls._do_build_from_config(config, registry_key, module_key)

    @classmethod
    def _do_build_from_config(
        cls,
        config: MutableMapping,
        registry_key: str,
        module_key: str,
    ):
        registry_name = config.pop(registry_key)
        module_name = config.pop(module_key)
        registry = cls.__OBJ_DICT__[registry_name]
        module = registry.get(module_name)

        for key, value in config.items():
            if isinstance(value, MutableMapping) and registry_key in value and module_key in value:
                config[key] = cls._do_build_from_config(value, registry_key, module_key)
        return module(**config)

src/network/utils/registry.py
SeonghwanSeo/PharmacoNet
from .pharmacophore_model import PharmacophoreModel

src/scoring/__init__.py
SeonghwanSeo/PharmacoNet
from __future__ import annotations

import numpy as np
import math
import itertools

from typing import Dict, List, Tuple, Set, Iterator
from numpy.typing import NDArray

from ..utils.constant import INTERACTION_LIST


OVERLAP_DISTANCE = 1.5
CLUSTER_DISTANCE = 3.0


def coords_to_position(coords, center, resolution, size) -> Tuple[float, float, float]:
    x_center, y_center, z_center = center
    x, y, z = coords
    x_origin = x_center - (resolution * (size - 1) / 2)
    y_origin = y_center - (resolution * (size - 1) / 2)
    z_origin = z_center - (resolution * (size - 1) / 2)
    x_pos = float(x_origin + x * resolution)
    y_pos = float(y_origin + y * resolution)
    z_pos = float(z_origin + z * resolution)
    return (x_pos, y_pos, z_pos)


class DensityMapGraph():
    def __init__(self, center: Tuple[float, float, float], resolution: float, size: int):
        self.center: Tuple[float, float, float] = center
        self.resolution: float = resolution
        self.size: int = size

        self.nodes: List[DensityMapNode] = []
        self.edges: List[DensityMapEdge] = []
        self.node_dict: Dict[str, List[DensityMapNode]] = {typ: [] for typ in INTERACTION_LIST}
        self.edge_dict_nodes: Dict[Tuple[DensityMapNode, DensityMapNode], DensityMapEdge] = {}
        self.edge_dict_indices: Dict[Tuple[int, int], DensityMapEdge] = {}

        self.node_clusters: List[DensityMapNodeCluster] = []
        self.node_cluster_dict: Dict[str, List[DensityMapNodeCluster]] = dict(
            Cation=[],
            Anion=[],
            HBond=[],
            Aromatic=[],
            Hydrophobic=[],
            Halogen=[]
        )

    def add_node(
        self,
        node_type: str,
        hotspot_coords: Tuple[int, int, int],
        score: float,
        mask: NDArray[np.float_],
    ):
        for grids, grid_scores in self.__extract_pharmacophores(mask):
            grids, grid_scores = np.array(grids), np.array(grid_scores)
            if len(grids) < 8:
                continue
            new_node = DensityMapNode(self, hotspot_coords, node_type, score, grids, grid_scores)
            self.nodes.append(new_node)
            self.node_dict[node_type].append(new_node)
            for node in self.nodes:     # Add self loop
                edge = new_node.add_neighbors(node)
                self.edges.append(edge)
                self.edge_dict_nodes[(node, new_node)] = edge
                self.edge_dict_nodes[(new_node, node)] = edge
                self.edge_dict_indices[(node.index, new_node.index)] = edge
                self.edge_dict_indices[(new_node.index, node.index)] = edge

    def setup(self):
        self.__clustering()

    @staticmethod
    def __extract_pharmacophores(mask) -> Iterator[Tuple[List[Tuple[int, int, int]], List[float]]]:
        """Get Node From Mask by Clustering Algorithm

        Args:
            mask (NDArray): FloatArray[D, H, W]

        Yields:
            coordinates: List[Tuple[int, int, int]] - [(x, y, z)]
            grid_scores: List[float] - [score]
        """
        x_indices, y_indices, z_indices = np.where(mask > 0.)
        points = {(int(x), int(y), int(z)) for x, y, z in zip(x_indices, y_indices, z_indices)}
        while (len(points) > 0):
            point = x, y, z = points.pop()
            cluster = [point]
            scores = [float(mask[x, y, z])]
            search_center = cluster
            for (x, y, z) in search_center:
                new_center = []
                for dx, dy, dz in itertools.product((-1, 0, 1), repeat=3):
                    if dx == 0 and dy == 0 and dz == 0:
                        continue
                    new_point = _x, _y, _z = (x + dx, y + dy, z + dz)
                    if new_point in points:
                        new_center.append(new_point)
                        scores.append(float(mask[_x, _y, _z]))
                        points.remove(new_point)
                cluster += new_center
                search_center = new_center
            yield cluster, scores

    def __clustering(self):
        def are_nodes_close(node1: DensityMapNode, node2: DensityMapNode):
            edge = self.edge_dict_nodes[node1, node2]
            return edge.distance_mean < CLUSTER_DISTANCE

        # NOTE: Cation, Anion, Aromatic Aromatic
        GROUP_CLUSTER_CONFIGS = [
            {'name': 'Cation', 'major_type': ('SaltBridge_pneg', 'PiCation_pring'), 'minor_type': 'HBond'},
            {'name': 'Anion', 'major_type': 'SaltBridge_lneg', 'minor_type': 'HBond'},
            {'name': 'Aromatic', 'major_type': ('PiStacking', 'PiCation_lring'), 'minor_type': 'Hydrophobic'},
        ]
        used_nodes = set()
        for node in self.nodes:
            if node in used_nodes:
                continue
            for CONFIG in GROUP_CLUSTER_CONFIGS:
                if node.type.startswith(CONFIG['major_type']):
                    cluster_nodes = {node}
                    # NOTE: Add Overlapped Nodes
                    cluster_nodes.update(
                        overlapped_node for overlapped_node in node.overlapped_nodes
                        if overlapped_node.type.startswith(CONFIG['major_type'])
                    )
                    # NOTE: Add Dependency Nodes
                    cluster_nodes.update(
                        _node for _node in self.nodes
                        if _node.type.startswith(CONFIG['minor_type'])
                        and any(are_nodes_close(_node, cluster_node) for cluster_node in cluster_nodes)
                    )
                    used_nodes.update(cluster_nodes)
                    # NOTE: Add New Cluster
                    cluster = DensityMapNodeCluster(self, cluster_nodes, CONFIG['name'])
                    self.node_cluster_dict[CONFIG['name']].append(cluster)
                    break

        # NOTE: HBond, Hydrophobic, Halogen
        SINGLE_CLUSTER_CONFIGS = [
            {'name': 'HBond', 'type': 'HBond'},
            {'name': 'Hydrophobic', 'type': 'Hydrophobic'},
            {'name': 'Halogen', 'type': 'XBond'},
        ]
        for node in self.nodes:
            if node in used_nodes:
                continue
            for CONFIG in SINGLE_CLUSTER_CONFIGS:
                if node.type.startswith(CONFIG['type']):
                    # NOTE: Create New Cluster - [Node, Close Nodes]
                    cluster_nodes = {
                        _node for _node in self.nodes
                        if _node.type.startswith(CONFIG['type'])
                        and are_nodes_close(node, _node)
                    }
                    cluster_nodes.add(node)
                    cluster = DensityMapNodeCluster(self, cluster_nodes, CONFIG['name'])
                    used_nodes.update(cluster_nodes)
                    self.node_cluster_dict[CONFIG['name']].append(cluster)
                    break

        for clusters in self.node_cluster_dict.values():
            self.node_clusters.extend(clusters)


class DensityMapNodeCluster():
    def __init__(
        self,
        graph: DensityMapGraph,
        nodes: Set[DensityMapNode],
        cluster_type: str,
    ):
        self.graph = graph
        self.type: str = cluster_type
        self.nodes: Set[DensityMapNode] = nodes
        positions = np.array([node.center for node in self.nodes])
        radii = np.array([node.radius * 2 for node in self.nodes])
        center = np.mean(positions, axis=0)
        distances = np.linalg.norm(positions - center.reshape(1, 3), axis=-1) + radii
        x, y, z = center.tolist()
        self.center: Tuple[float, float, float] = (x, y, z)
        self.size: float = np.max(distances).item()

    def __repr__(self):
        return f'DensityMapNodeCluster[{self.type}] [ {self.nodes} ]'


class DensityMapNode():
    def __init__(
        self,
        graph: DensityMapGraph,
        hotspot_coords: Tuple[int, int, int],
        node_type: str,
        score: float,
        grids: NDArray[np.int_],
        grid_scores: NDArray[np.float_],
    ):
        self.graph: DensityMapGraph = graph
        self.index: int = len(self.graph.nodes)
        self.type: str = node_type
        self.grids = grids

        self.hotspot_position: Tuple[float, float, float] = coords_to_position(hotspot_coords, self.graph.center, self.graph.resolution, self.graph.size)

        self.score: float = score
        center_coords = np.average(grids, axis=0, weights=grid_scores)
        self.center: NDArray[np.float32] = np.array(coords_to_position(center_coords, self.graph.center, self.graph.resolution, self.graph.size), dtype=np.float32)
        self.radius = (grids.shape[0] / (4 * math.pi / 3)) ** (1 / 3) * self.graph.resolution

        self.neighbor_edge_dict: Dict[DensityMapNode, DensityMapEdge] = {}
        self.overlapped_nodes: List[DensityMapNode] = []

    def __hash__(self):
        return self.index

    def __repr__(self):
        return f'DensityMapNode({self.index})[{self.type}]'

    def add_neighbors(self, neighbor: DensityMapNode) -> DensityMapEdge:
        assert neighbor not in self.neighbor_edge_dict
        edge = DensityMapEdge(self.graph, self, neighbor)
        self.neighbor_edge_dict[neighbor] = edge
        neighbor.neighbor_edge_dict[self] = edge

        if edge.overlapped:
            self.overlapped_nodes.append(neighbor)
            neighbor.overlapped_nodes.append(self)

        return edge


class DensityMapEdge():
    """Density Map Edge

    Attributes:
        graph: Parent Graph
        index: Edge Index
        node_indices: (Node1.index, Node2.index)
        nodes: (Node1, Node2)
        type: (Node1.type, Node2.type)
        center_distance: distance between map center
        overlapped: end nodes are overlapped
    """

    def __init__(self, graph, node1, node2):
        self.graph = graph
        self.index = len(self.graph.edges)
        if node2.index < node1.index:
            node1, node2 = node2, node1
        self.node_indices: Tuple[int, int] = (node1.index, node2.index)
        self.nodes: Tuple[DensityMapNode, DensityMapNode] = (node1, node2)
        type1, type2 = node1.type, node2.type

        self.type: Tuple[str, str] = (min(type1, type2), max(type1, type2))
        self.distance_mean: float = np.linalg.norm(node1.center - node2.center).item()
        self.distance_std: float = math.sqrt(node1.radius ** 2 + node2.radius ** 2)
        self.overlapped: bool = (self.distance_mean < OVERLAP_DISTANCE)

src/scoring/density_map.py
SeonghwanSeo/PharmacoNet
from __future__ import annotations

import itertools

import numpy as np
import math

from typing import Dict, List, Tuple, Optional
from numpy.typing import NDArray

from .ligand import Ligand, LigandGraph, LigandNode, LigandNodeCluster
from .pharmacophore_model import PharmacophoreModel, ModelNodeCluster, ModelNode
from .tree import ClusterMatchTreeRoot
try:
    from .match_utils_numba import scoring_matching_pair, scoring_matching_self
except:
    from .match_utils import scoring_matching_pair, scoring_matching_self


# NOTE: TYPE
LigandClusterPair = Tuple[LigandNodeCluster, LigandNodeCluster]
ModelClusterPair = Tuple[ModelNodeCluster, ModelNodeCluster]


# NOTE: Constant
WEIGHTS = dict(
    Cation=8,
    Anion=8,
    Aromatic=8,
    HBond_donor=4,
    HBond_acceptor=4,
    Halogen=4,
    Hydrophobic=1,
)
MAX_DEPTH = 20


def priority_fn(cluster: LigandNodeCluster):
    cluster_size_priority = -len(cluster.nodes)
    cluster_type = cluster.type
    atom_index = min(cluster.nodes[0].atom_indices)
    if cluster_type.startswith('Aromatic'):
        return (0, cluster_size_priority, 0, atom_index)
    elif cluster_type.startswith('Cation'):
        return (0, cluster_size_priority, 1, atom_index)
    elif cluster_type.startswith('Anion'):
        return (0, cluster_size_priority, 2, atom_index)
    elif cluster_type.startswith('HBond'):
        return (1, cluster_size_priority, 0, atom_index)
    elif cluster_type.startswith('Halogen'):
        return (1, cluster_size_priority, 1, atom_index)
    elif cluster_type.startswith('Hydrophobic'):
        return (1, cluster_size_priority, 2, atom_index)
    else:
        raise NotImplementedError


class GraphMatcher():
    def __init__(
        self,
        model: PharmacophoreModel,
        ligand: Ligand,
        max_depth: Optional[int] = MAX_DEPTH,
    ):
        self.model_graph: PharmacophoreModel = model
        self.ligand_graph: LigandGraph = ligand.graph
        self.num_conformers = self.ligand_graph.num_conformers
        self.max_depth = max_depth
        self.cluster_match_dict: Dict[LigandNodeCluster, List[ModelNodeCluster]]
        self.ligand_cluster_list: List[LigandNodeCluster]
        self.node_match_dict: Dict[Tuple[LigandNodeCluster, ModelNodeCluster], List[Tuple[LigandNode, List[ModelNode], NDArray[np.float32]]]]

    def setup(self):
        self.cluster_match_dict = self._get_cluster_match_dict()
        self.ligand_cluster_list = sorted(self.cluster_match_dict.keys(), key=priority_fn)
        if self.max_depth is not None:
            self.ligand_cluster_list = self.ligand_cluster_list[:self.max_depth]
        self.node_match_dict = self._get_node_match_dict()
        self.matching_pair_scores_dict: Dict[LigandClusterPair, Dict[ModelClusterPair, Tuple[float, ...]]] = self._get_pair_scores()

    def evaluate(self):
        self.setup()
        root_tree = self._run()
        return list(root_tree.iteration())

    def scoring(self) -> float:
        self.setup()
        root_tree = self._run()
        return max(leaf.score for leaf in root_tree.iteration())

    def _run(self) -> ClusterMatchTreeRoot:
        root_tree = ClusterMatchTreeRoot(
            self.ligand_cluster_list,
            self.cluster_match_dict,
            self.matching_pair_scores_dict,
            self.ligand_graph.num_conformers
        )
        root_tree.run()
        return root_tree

    def _get_cluster_match_dict(self) -> Dict[LigandNodeCluster, List[ModelNodeCluster]]:
        cluster_match_dict: Dict[LigandNodeCluster, List[ModelNodeCluster]] = {}
        ligand_graph, model_graph = self.ligand_graph, self.model_graph
        for ligand_cluster in ligand_graph.node_clusters:
            match_model_clusters = [
                model_cluster for model_cluster in model_graph.node_clusters
                if len(ligand_cluster.node_types.intersection(model_cluster.node_types)) > 0
            ]
            if len(match_model_clusters) > 0:
                cluster_match_dict[ligand_cluster] = match_model_clusters
        return cluster_match_dict

    def _get_node_match_dict(self) -> Dict[Tuple[LigandNodeCluster, ModelNodeCluster], List[Tuple[LigandNode, List[ModelNode], NDArray[np.float32]]]]:
        def __get_node_match(ligand_node: LigandNode, model_cluster: ModelNodeCluster) -> Tuple[LigandNode, List[ModelNode], NDArray[np.float32]]:
            match_model_nodes = [model_node for model_node in model_cluster.nodes if model_node.type in ligand_node.types]
            weights = np.array([WEIGHTS[model_node.type] * model_node.score for model_node in match_model_nodes], dtype=np.float32)
            return (ligand_node, match_model_nodes, weights)

        node_match_dict = {
            (ligand_cluster, model_cluster): [__get_node_match(ligand_node, model_cluster) for ligand_node in ligand_cluster.nodes]
            for ligand_cluster, model_cluster_list in self.cluster_match_dict.items()
            for model_cluster in model_cluster_list
        }
        node_match_dict = {
            key: [(ligand_node, model_node_list, weights) for ligand_node, model_node_list, weights in node_matches if len(model_node_list) > 0]
            for key, node_matches in node_match_dict.items()
        }
        return node_match_dict

    # NOTE: (Not Use) Code with Readability - same to _get_pair_scores().
    def _get_pair_scores_readability(self) -> Dict[LigandClusterPair, Dict[ModelClusterPair, Tuple[float, ...]]]:
        matching_pair_scores_dict: Dict[LigandClusterPair, Dict[ModelClusterPair, Tuple[float, ...]]] = {
            (ligand_cluster1, ligand_cluster2): {} for ligand_cluster1, ligand_cluster2 in itertools.combinations(self.ligand_cluster_list, 2)
        }

        NO_MATCH_SCORE = (-1,) * self.num_conformers
        for ligand_cluster1, ligand_cluster2 in itertools.combinations(self.ligand_cluster_list, 2):
            ligand_cluster_distance = np.linalg.norm(ligand_cluster1.center - ligand_cluster2.center, axis=-1)
            ligand_cluster_size = ligand_cluster1.size + ligand_cluster2.size

            model_cluster_list1, model_cluster_list2 = self.cluster_match_dict[ligand_cluster1], self.cluster_match_dict[ligand_cluster2]
            for model_cluster1, model_cluster2 in itertools.product(model_cluster_list1, model_cluster_list2):
                (x1, y1, z1), (x2, y2, z2) = model_cluster1.center, model_cluster2.center
                model_cluster_distance = math.sqrt((x1 - x2) ** 2 + (y1 - y2) ** 2 + (z1 - z2) ** 2)
                model_cluster_size = model_cluster1.size + model_cluster2.size

                if min(np.abs(ligand_cluster_distance - model_cluster_distance) - ligand_cluster_size) > model_cluster_size:
                    pair_score = NO_MATCH_SCORE
                else:
                    node_match_list1 = self.node_match_dict[ligand_cluster1, model_cluster1]
                    node_match_list2 = self.node_match_dict[ligand_cluster2, model_cluster2]
                    pair_score = scoring_matching_pair(node_match_list1, node_match_list2, self.num_conformers)
                matching_pair_scores_dict[ligand_cluster1, ligand_cluster2][model_cluster1, model_cluster2] = pair_score

        for ligand_cluster in self.ligand_cluster_list:
            for model_cluster in self.cluster_match_dict[ligand_cluster]:
                node_match_list = self.node_match_dict[ligand_cluster, model_cluster]
                self_pair_score = scoring_matching_self(node_match_list, self.num_conformers)
                matching_pair_scores_dict[ligand_cluster, ligand_cluster][model_cluster, model_cluster] = self_pair_score

        return matching_pair_scores_dict

    # NOTE: Efficient but bad-readable code
    def _get_pair_scores(self) -> Dict[LigandClusterPair, Dict[ModelClusterPair, Tuple[float, ...]]]:
        NO_MATCH_SCORE = (-1,) * self.num_conformers

        def __get_score_dict_outer(ligand_cluster_pair: LigandClusterPair) -> Dict[ModelClusterPair, Tuple[float, ...]]:
            ligand_cluster1, ligand_cluster2 = ligand_cluster_pair
            if ligand_cluster1 is ligand_cluster2:
                return {
                    (model_cluster, model_cluster): scoring_matching_self(self.node_match_dict[ligand_cluster1, model_cluster], self.num_conformers)
                    for model_cluster in self.cluster_match_dict[ligand_cluster1]
                }
            else:
                ligand_cluster_distance = np.linalg.norm(ligand_cluster1.center - ligand_cluster2.center, axis=-1)
                ligand_cluster_size = ligand_cluster1.size + ligand_cluster2.size
                return {
                    model_cluster_pair: __get_score_dict_inner(ligand_cluster_pair, model_cluster_pair, ligand_cluster_distance, ligand_cluster_size)
                    for model_cluster_pair in itertools.product(self.cluster_match_dict[ligand_cluster1], self.cluster_match_dict[ligand_cluster2])
                }

        def __get_score_dict_inner(
            ligand_cluster_pair: LigandClusterPair,
            model_cluster_pair: ModelClusterPair,
            ligand_cluster_distance: NDArray[np.float32],
            ligand_cluster_size: NDArray[np.float32],
        ) -> Tuple[float, ...]:
            ligand_cluster1, ligand_cluster2 = ligand_cluster_pair
            model_cluster1, model_cluster2 = model_cluster_pair
            (x1, y1, z1), (x2, y2, z2) = model_cluster1.center, model_cluster2.center
            model_cluster_distance = math.sqrt((x1 - x2) ** 2 + (y1 - y2) ** 2 + (z1 - z2) ** 2)
            model_cluster_size = model_cluster1.size + model_cluster2.size

            if min(np.abs(ligand_cluster_distance - model_cluster_distance) - ligand_cluster_size) > model_cluster_size:
                pair_score = NO_MATCH_SCORE
            else:
                node_match_list1 = self.node_match_dict[ligand_cluster1, model_cluster1]
                node_match_list2 = self.node_match_dict[ligand_cluster2, model_cluster2]
                pair_score = scoring_matching_pair(node_match_list1, node_match_list2, self.num_conformers)
            return pair_score

        matching_pair_scores_dict: Dict[LigandClusterPair, Dict[ModelClusterPair, Tuple[float, ...]]] = {
            ligand_cluster_pair: __get_score_dict_outer(ligand_cluster_pair)
            for ligand_cluster_pair in itertools.combinations_with_replacement(self.ligand_cluster_list, 2)
        }
        return matching_pair_scores_dict

src/scoring/graph_match.py
SeonghwanSeo/PharmacoNet
from __future__ import annotations
import os
import numpy as np

from openbabel import pybel
from openbabel.pybel import ob
from rdkit import Chem

from typing import Set, Dict, List, Sequence, Union, Tuple, Optional, Iterator
from numpy.typing import NDArray

from .ligand_utils import get_pharmacophore_nodes, PharmacophoreNode


def order(a, b):
    return min(a, b), max(a, b)


class Ligand():
    def __init__(
        self,
        pbmol: pybel.Molecule,
        rdmol: Optional[Chem.Mol] = None,
        atom_positions: Optional[Union[List[NDArray[np.float32]], NDArray[np.float32]]] = None,
        conformer_axis: Optional[int] = None
    ):
        """Ligand Object

        Args:
            pbmol: pybel.Molecule Object
            rdmol: Chem.Mol Object
            atom_positions: List[NDArray[np.float32]] | NDArray[np.float32] | None
                case: NDArray[np.float32]
                    i) conformer_axis is 0 or None
                        atom_positions: (N_conformers, N_atoms, 3)
                    ii) conformer_axis is 1
                        atom_positions: (N_atoms, N_conformers, 3)
                case: None
                    Using RDKit Conformer informations
        """
        self.pbmol: pybel.Molecule = pbmol.clone
        self.pbmol.removeh()
        self.obmol: ob.OBMol = self.pbmol.OBMol
        self.obatoms: List[ob.OBAtom] = [self.obmol.GetAtom(i + 1) for i in range(self.obmol.NumAtoms())]

        self.num_atoms: int = len(self.obatoms)
        if rdmol is not None:
            rdmol = Chem.RemoveHs(rdmol)
            assert self.num_atoms == rdmol.GetNumAtoms(), f'Atom Number ERROR - openbabel: {self.num_atoms}, rdkit: {rdmol.GetNumAtoms()}'
            for obatom, rdatom in zip(self.obatoms, rdmol.GetAtoms()):
                assert obatom.GetAtomicNum() == rdatom.GetAtomicNum(), f'Atomic Number ERROR - openbabel: {obatom.GetAtomicNum()}, rdkit: {rdatom.GetAtomicNum()}'

        self.atom_positions: NDArray[np.float32]    # [N_atoms, N_conformers, 3]
        if isinstance(atom_positions, list):
            self.atom_positions = np.stack(atom_positions, axis=1, dtype=np.float32)
        elif isinstance(atom_positions, np.ndarray):
            self.atom_positions = np.asarray(atom_positions, dtype=np.float32)
            if conformer_axis in [0, None]:
                self.atom_positions = np.ascontiguousarray(np.moveaxis(self.atom_positions, 0, 1))
        elif rdmol is not None:
            self.atom_positions = np.stack([conformer.GetPositions() for conformer in rdmol.GetConformers()], axis=1, dtype=np.float32)
        else:
            raise ValueError

        assert self.num_atoms == self.atom_positions.shape[0]
        self.num_conformers: int = self.atom_positions.shape[1]

        self.pharmacophore_nodes: Dict[str, List[PharmacophoreNode]] = get_pharmacophore_nodes(self.pbmol)
        self.pharmacophore_list: List[Tuple[str, PharmacophoreNode]] = []
        for typ, node_list in self.pharmacophore_nodes.items():
            self.pharmacophore_list.extend((typ, node) for node in node_list)

        self.graph = LigandGraph(self)

    @classmethod
    def create(
        cls,
        pbmol: Optional[pybel.Molecule] = None,
        smiles: Optional[str] = None,
        filename: Optional[str] = None,
        atom_positions: Optional[NDArray[np.float32]] = None,
        conformer_axis: Optional[int] = None,
    ) -> Ligand:
        assert pbmol is not None or smiles is not None or filename is not None
        if pbmol is not None:
            pass
        elif smiles is not None:
            pbmol = pybel.readstring('smi', smiles)
            assert atom_positions is not None
        else:
            assert filename is not None
            extension = os.path.splitext(filename)[1]
            assert extension in ['.sdf', '.pdb', '.mol2']
            pbmol = next(pybel.readfile(extension[1:], filename))
        if atom_positions is None:
            rdmol = Chem.MolFromMol2Block(pbmol.write('mol2'))
            if rdmol is None:
                rdmol = Chem.MolFromPDBBlock(pbmol.write('pdb'))
            assert rdmol is not None
        else:
            rdmol = None
        out = cls(pbmol, rdmol, atom_positions, conformer_axis)
        return out


class LigandGraph():
    def __init__(self, ligand_obj: Ligand):
        self.nodes: List[LigandNode] = []
        self.edges: List[LigandEdge] = []
        self.node_dict: Dict[str, List[LigandNode]] = {}
        self.edge_dict: Dict[Tuple[LigandNode, LigandNode], LigandEdge] = {}
        self.node_clusters: List[LigandNodeCluster]
        self.node_cluster_dict: Dict[str, List[LigandNodeCluster]] = dict(
            Cation=[],
            Anion=[],
            HBond=[],
            Aromatic=[],
            Hydrophobic=[],
            Halogen=[]
        )
        self.__add_nodes(ligand_obj)
        self.__setup_conformer(ligand_obj)
        self.__group_nodes(ligand_obj)
        self.__setup_cluster()

    def __setup_conformer(self, obj):
        assert obj.num_conformers > 0
        self.atom_positions = obj.atom_positions
        self.num_conformers = obj.num_conformers
        for node in self.nodes:
            node.set_positions()
        for edge in self.edges:
            edge.set_distances()

    def __add_nodes(self, obj):
        atom_indices_dict: Dict[Union[int, Sequence[int]], LigandNode] = dict()
        for pharmacophore_type, pharmacophore_node in obj.pharmacophore_list:
            node = atom_indices_dict.get(pharmacophore_node.atom_indices, None)
            if node is not None:
                node.types.append(pharmacophore_type)
                self.node_dict.setdefault(pharmacophore_type, []).append(node)
            else:
                new_node = LigandNode(self, len(self.nodes), pharmacophore_node.atom_indices, pharmacophore_node.center_indices, pharmacophore_type)
                self.nodes.append(new_node)
                self.node_dict.setdefault(pharmacophore_type, []).append(new_node)
                for node in self.nodes[:-1]:
                    edge = node.add_neighbors(new_node)
                    self.edges.append(edge)
                    self.edge_dict[(node, new_node)] = edge
                    self.edge_dict[(new_node, node)] = edge
                atom_indices_dict[pharmacophore_node.atom_indices] = new_node

    def __group_nodes(self, obj):
        hbond_pharmacophores: Dict[int, List[LigandNode]] = {}
        hydrop_pharmacophores: Dict[int, List[LigandNode]] = {}

        for node in self.nodes:
            pharmacophore_types = node.types
            if 'HBond_acceptor' in pharmacophore_types or 'HBond_donor' in pharmacophore_types:
                # NOTE: Clustering Nodes in Same Functional Groups (such as -PO3)
                assert len(node.atom_indices) == 1
                atom_index = next(iter(node.atom_indices))
                obatom = obj.obatoms[atom_index]
                neighbors = [neighbor.GetIdx() - 1 for neighbor in ob.OBAtomAtomIter(obatom) if neighbor.GetAtomicNum() != 1]
                if len(neighbors) == 1:
                    group_nodes = hbond_pharmacophores.setdefault(neighbors[0], [])
                    for _node in group_nodes:
                        node.group_nodes.add(_node)
                        _node.group_nodes.add(node)
                    group_nodes.append(node)

            elif 'Hydrophobic' in pharmacophore_types:
                # NOTE: Clustering Nodes in Same Functional Groups (such as -C(CH3)3)
                assert len(node.atom_indices) == 1
                atom_index = next(iter(node.atom_indices))
                obatom = obj.obatoms[atom_index]
                neighbors = [neighbor.GetIdx() - 1 for neighbor in ob.OBAtomAtomIter(obatom) if neighbor.GetAtomicNum() != 1]
                if len(neighbors) == 1:
                    group_nodes = hydrop_pharmacophores.setdefault(neighbors[0], [])
                    for _node in group_nodes:
                        node.group_nodes.add(_node)
                        _node.group_nodes.add(node)
                    group_nodes.append(node)

        hydrophobic_nodes = self.node_dict.get('Hydrophobic', [])
        index_to_node = {next(iter(node.atom_indices)): node for node in hydrophobic_nodes}
        while len(index_to_node) > 0:
            index, node = index_to_node.popitem()
            group_nodes = [node] + list(node.group_nodes)
            group_index = [next(iter(node.atom_indices)) for node in group_nodes]
            for atom_index in group_index:
                obatom = obj.obatoms[atom_index]
                neighbors = [neighbor.GetIdx() - 1 for neighbor in ob.OBAtomAtomIter(obatom) if neighbor.GetAtomicNum() == 6]
                for neighbor_index in neighbors:
                    neighbor_node = index_to_node.pop(neighbor_index, None)
                    if neighbor_node is None or neighbor_index in group_nodes:
                        continue
                    group_index.append(neighbor_index)
                    for node in group_nodes:
                        node.group_nodes.add(neighbor_node)
                        neighbor_node.group_nodes.add(node)
                    group_nodes.append(neighbor_node)

    def __setup_cluster(self):
        # NOTE: Set Cluster
        # Cluster: List[LigandNode]
        # The order inside the cluster follows priority.
        in_cluster: Set[LigandNode] = set()
        node_cluster_dict: Dict[LigandNode, LigandNodeCluster] = {}
        for pharmacophore_type in ['Aromatic', 'Cation', 'Anion', 'Halogen']:
            for node in self.node_dict.get(pharmacophore_type, []):
                if node in in_cluster:
                    continue
                in_cluster.add(node)
                cluster = LigandNodeCluster(pharmacophore_type)
                cluster.add_new_node(node, 'high')
                node_cluster_dict[node] = cluster

        for pharmacophore_type in ['Hydrophobic', 'HBond_donor', 'HBond_acceptor']:
            for node in self.node_dict.get(pharmacophore_type, []):
                if node in in_cluster:
                    continue
                in_cluster.add(node)
                add_new_cluster: bool = True
                if len(node.dependence_nodes) > 0:
                    # NOTE: Added to Dependence Nodes (Aromatic Ring, Charged Atoms)
                    cluster = node_cluster_dict[min(node.dependence_nodes)]
                    cluster.add_new_node(node, 'low')
                    add_new_cluster = False
                elif len(node.group_nodes) > 0:
                    # NOTE: Clustering Nodes in Same Functional Groups
                    for group_node in node.group_nodes:
                        if group_node in node_cluster_dict:
                            cluster = node_cluster_dict[group_node]
                            cluster.add_new_node(node, 'low')
                            add_new_cluster = False
                            break
                if add_new_cluster:
                    if pharmacophore_type.startswith('HBond'):
                        cluster = LigandNodeCluster('HBond')
                    else:
                        cluster = LigandNodeCluster('Hydrophobic')
                    cluster.add_new_node(node, 'low')
                    node_cluster_dict[node] = cluster

        self.node_clusters = list(node_cluster_dict.values())
        for cluster in self.node_clusters:
            self.node_cluster_dict[cluster.type].append(cluster)


class LigandNode():
    def __init__(
        self,
        graph: LigandGraph,
        node_index: int,
        atom_indices: Union[int, Sequence[int]],
        center_indices: Union[int, Sequence[int]],
        node_type: str,
    ):
        """Ligand Pharmacophore Node

        Args:
            graph: root graph
            index: node index
            node_type: node pharmacophoretypes
        """
        self.graph: LigandGraph = graph
        self.index: int = node_index
        self.types: List[str] = [node_type]
        self.atom_indices: Set[int] = {atom_indices} if isinstance(atom_indices, int) else set(atom_indices)
        self.center_indices: Union[int, Sequence[int]] = center_indices

        self.neighbor_edge_dict: Dict[LigandNode, LigandEdge] = {}
        self.group_nodes: Set[LigandNode] = set()
        self.dependence_nodes: Set[LigandNode] = set()

        self.positions: NDArray[np.float32]  # (N, 3)

    def __repr__(self):
        return f'LigandNode({self.index}){self.types}'

    def set_positions(self):
        if isinstance(self.center_indices, int):
            self.positions = np.asarray(self.graph.atom_positions[self.center_indices], dtype=np.float32)
        else:
            self.positions = np.mean(self.graph.atom_positions[self.center_indices, :], axis=0, dtype=np.float32)

    def add_neighbors(self, neighbor: LigandNode) -> LigandEdge:
        edge = self.neighbor_edge_dict.get(neighbor, None)
        if edge is not None:
            return edge
        edge = LigandEdge(self.graph, self, neighbor)
        self.neighbor_edge_dict[neighbor] = edge
        neighbor.neighbor_edge_dict[self] = edge

        node_types = self.types
        neighbor_node_types = neighbor.types

        def check_type(node_types, *keys):
            return any(typ.startswith(keys) for typ in node_types)

        if check_type(node_types, 'Hydrophobic') and check_type(neighbor_node_types, 'Aromatic'):
            if self.atom_indices.issubset(neighbor.atom_indices):
                self.dependence_nodes.add(neighbor)
        elif check_type(node_types, 'Aromatic') and check_type(neighbor_node_types, 'Hydrophobic'):
            if neighbor.atom_indices.issubset(self.atom_indices):
                neighbor.dependence_nodes.add(self)
        elif check_type(node_types, 'HBond') and check_type(neighbor_node_types, 'Cation', 'Anion'):
            if self.atom_indices.issubset(neighbor.atom_indices):
                self.dependence_nodes.add(neighbor)
        elif check_type(node_types, 'Cation', 'Anion') and check_type(neighbor_node_types, 'HBond'):
            if neighbor.atom_indices.issubset(self.atom_indices):
                neighbor.dependence_nodes.add(self)
        return edge

    def __gt__(self, other):
        return self.index > other.index

    def __le__(self, other):
        return self.index < other.index


class LigandEdge():
    def __init__(self, graph, node1, node2):
        self.graph = graph
        self.index = len(self.graph.edges)
        if node2.index < node1.index:
            node1, node2 = node2, node1
        self.indices: Tuple[int, int] = (node1.index, node2.index)
        self.nodes: Tuple[LigandNode, LigandNode] = (node1, node2)

        self.distances: NDArray[np.float32]  # (N_conformer,)

    def set_distances(self):
        node1, node2 = self.nodes
        self.distances = np.linalg.norm(node1.positions - node2.positions, axis=-1)


class LigandNodeCluster():
    """ Ligand Node Cluster
    Property
    - center_node
    - cluster_nodes

    Case
    - one high-priority node:
        center_node = high-priority node
        cluster_nodes = []
    - one high-priority node + multi low-priority nodes:
        center_node = high-priority node
        cluster_nodes = low-priority nodes
    - one low-priority node
        center_node = low-priority node
        cluster_nodes = []
    - multi low-priority nodes:
        center_node = None
        cluster_nodes = low-priority nodes
    """

    def __init__(self, cluster_type: str):
        self.type: str = cluster_type
        self._high_priority_node: Optional[LigandNode] = None
        self._low_priority_nodes: List[LigandNode] = []

        self._node_types: Optional[Set[str]] = None
        self._node_indices: Optional[List[int]] = None

        self._positions: Optional[NDArray[np.float32]] = None   # [N_conf, N_node, 3]
        self._center: Optional[NDArray[np.float32]] = None      # [N_conf, 3]
        self._size: Optional[NDArray[np.float32]] = None        # [N_conf,]

    def __iter__(self) -> Iterator[LigandNode]:
        if self._high_priority_node is not None:
            if len(self._low_priority_nodes) > 0:
                yield self._high_priority_node
                yield from self._low_priority_nodes
            else:
                yield self._high_priority_node
        else:
            yield from iter(self._low_priority_nodes)

    def add_new_node(self, node: LigandNode, priority: str):
        assert priority in ('low', 'high')
        if priority.startswith('high'):
            self._high_priority_node = node
        else:
            self._low_priority_nodes.append(node)
        self._node_types = None
        self._node_indices = None

    def __repr__(self):
        if self.center_node is not None and len(self.cluster_nodes) > 0:
            return f'LigandNodeCluster({self.type})[ Center: {self.center_node} Cluster: {self.cluster_nodes} ]'
        elif self.center_node is not None:
            return f'LigandNodeCluster({self.type})[ Center: {self.center_node} ]'
        else:
            return f'LigandNodeCluster({self.type})[ Cluster: {self.cluster_nodes} ]'

    @property
    def nodes(self) -> List[LigandNode]:
        return list(iter(self))

    @property
    def node_types(self) -> Set[str]:
        if self._node_types is None:
            self._node_types = set()
            for node in self:
                self._node_types.update(node.types)
        return self._node_types

    @property
    def node_indices(self) -> Set[str]:
        if self._node_types is None:
            self._node_types = set()
            for node in self:
                self._node_types.update(node.types)
        return self._node_types

    @property
    def center_node(self) -> Optional[LigandNode]:
        if self._high_priority_node is not None:
            return self._high_priority_node
        elif len(self._low_priority_nodes) == 1:
            return self._low_priority_nodes[0]
        else:
            return None

    @property
    def cluster_nodes(self) -> List[LigandNode]:
        if self._high_priority_node is not None:
            return self._low_priority_nodes
        elif len(self._low_priority_nodes) > 1:
            return self._low_priority_nodes
        else:
            return []

    @property
    def positions(self) -> NDArray[np.float32]:
        if self._positions is None:
            self._positions = np.stack([node.positions for node in self.nodes], axis=1)
        return self._positions

    @property
    def center(self) -> NDArray[np.float32]:
        if self._center is None:
            self._center = np.mean(self.positions, axis=1)
        return self._center

    @property
    def size(self) -> NDArray[np.float32]:
        if self._size is None:
            self._size = np.max(np.linalg.norm(self.positions - self.center.reshape(-1, 1, 3), axis=-1), axis=-1)
        return self._size

src/scoring/ligand.py
SeonghwanSeo/PharmacoNet
import numpy as np
from openbabel import pybel
from openbabel.pybel import ob

from typing import List, Union, Tuple, Dict
from numpy.typing import NDArray


class PharmacophoreNode():
    def __init__(self, atom_indices: Union[int, Tuple[int, ...]], center_indices: Union[None, int, Tuple[int, ...]] = None):
        if center_indices is None:
            center_indices = atom_indices
        self.atom_indices: Union[int, Tuple[int, ...]] = atom_indices
        self.center_indices: Union[int, Tuple[int, ...]] = center_indices

    def get_center(self, atom_positions: NDArray) -> NDArray:
        if isinstance(self.center_indices, int):
            return atom_positions[self.center_indices]
        else:
            return np.mean(atom_positions[self.center_indices, :], axis=0)


def get_pharmacophore_nodes(pbmol: pybel.Molecule) -> Dict[str, List[PharmacophoreNode]]:
    obmol = pbmol.OBMol
    obatoms: List[ob.OBAtom] = list(ob.OBMolAtomIter(obmol))
    pbmol_hyd = pbmol.clone
    pbmol_hyd.OBMol.AddPolarHydrogens()
    obmol_hyd = pbmol_hyd.OBMol
    num_heavy_atoms = len(obatoms)
    obatoms_hyd: List[ob.OBAtom] = list(ob.OBMolAtomIter(obmol_hyd))[:num_heavy_atoms]

    hydrophobics = [
        PharmacophoreNode(idx) for idx, obatom in enumerate(obatoms)
        if obatom.GetAtomicNum() == 6
        and all(neigh.GetAtomicNum() in (1, 6) for neigh in ob.OBAtomAtomIter(obatom))
    ]
    hbond_acceptors = [
        PharmacophoreNode(idx) for idx, obatom in enumerate(obatoms)
        if obatom.GetAtomicNum() not in [9, 17, 35, 53]
        and obatom.IsHbondAcceptor()
    ]
    hbond_donors = [
        PharmacophoreNode(idx) for idx, obatom in enumerate(obatoms_hyd)
        if obatom.IsHbondDonor()
    ]
    rings = [
        PharmacophoreNode(tuple(sorted(idx - 1 for idx in ring._path)))     # start from 1 -> minus
        for ring in pbmol.sssr if ring.IsAromatic()
    ]
    rings.sort(key=lambda ring: ring.atom_indices)

    pos_charged = [
        PharmacophoreNode(idx) for idx, obatom in enumerate(obatoms)
        if is_quartamine_N(obatom)
        or is_tertamine_N(obatom)
        or is_sulfonium_S(obatom)
    ]
    neg_charged = []

    for idx, obatom in enumerate(obatoms):
        if is_guanidine_C(obatom):
            nitrogens = tuple(
                neigh.GetIdx() - 1 for neigh in ob.OBAtomAtomIter(obatom)
                if neigh.GetAtomicNum() == 7
            )
            pos_charged.append(PharmacophoreNode((idx,) + nitrogens, idx))

        elif is_phosphate_P(obatom) or is_sulfate_S(obatom):
            neighbors = tuple(neigh.GetIdx() - 1 for neigh in ob.OBAtomAtomIter(obatom))
            neg_charged.append(PharmacophoreNode((idx,) + neighbors, idx))

        elif is_sulfonicacid_S(obatom):
            oxygens = tuple(
                neigh.GetIdx() - 1 for neigh in ob.OBAtomAtomIter(obatom)
                if neigh.GetAtomicNum() == 8
            )
            neg_charged.append(PharmacophoreNode((idx,) + oxygens, idx))

        elif is_carboxylate_C(obatom):
            oxygens = tuple(
                neigh.GetIdx() - 1 for neigh in ob.OBAtomAtomIter(obatom)
                if neigh.GetAtomicNum() == 8
            )
            neg_charged.append(PharmacophoreNode((idx,) + oxygens, oxygens))

    xbond_donors = [
        PharmacophoreNode(idx) for idx, obatom in enumerate(obatoms)
        if is_halocarbon_X(obatom)
    ]

    return {
        'Hydrophobic': hydrophobics,
        'Aromatic': rings,
        'Cation': pos_charged,
        'Anion': neg_charged,
        'HBond_donor': hbond_donors,
        'HBond_acceptor': hbond_acceptors,
        'Halogen': xbond_donors
    }


""" FUNCTIONAL GROUP """


def is_quartamine_N(obatom: ob.OBAtom):
    # It's a nitrogen, so could be a protonated amine or quaternary ammonium
    if obatom.GetAtomicNum() != 7:  # Nitrogen
        return False
    if obatom.GetExplicitDegree() != 4:
        return False
    for neigh in ob.OBAtomAtomIter(obatom):
        if neigh.GetAtomicNum() == 1:  # It's a quat. ammonium (N with 4 residues != H)
            return False
    return True


def is_tertamine_N(obatom: ob.OBAtom):  # Nitrogen
    return (
        obatom.GetAtomicNum() == 7
        and obatom.GetHyb() == 3
        and obatom.GetHvyDegree() == 3
    )


def is_sulfonium_S(obatom: ob.OBAtom):
    if obatom.GetAtomicNum() != 16:  # Sulfur
        return False
    if obatom.GetExplicitDegree() != 3:
        return False
    for neigh in ob.OBAtomAtomIter(obatom):
        if neigh.GetAtomicNum() == 1:  # It's a sulfonium (S with 3 residues != H)
            return False
    return True


def is_guanidine_C(obatom: ob.OBAtom):
    if obatom.GetAtomicNum() != 6:  # It's a carbon atom
        return False
    numNs = 0
    numN_with_only_C = 0
    for neigh in ob.OBAtomAtomIter(obatom):
        if neigh.GetAtomicNum() == 7:
            numNs += 1
            if neigh.GetHvyDegree() == 1:
                numN_with_only_C += 1
        else:
            return False
    return numNs == 3 and numN_with_only_C > 0


def is_sulfonicacid_S(obatom: ob.OBAtom):
    if obatom.GetAtomicNum() != 16:  # Sulfur
        return False
    numOs = 0
    for neigh in ob.OBAtomAtomIter(obatom):
        if neigh.GetAtomicNum() == 8:
            numOs += 1
    return numOs == 3


def is_sulfate_S(obatom: ob.OBAtom):
    if obatom.GetAtomicNum() != 16:  # Sulfur
        return False
    numOs = 0
    for neigh in ob.OBAtomAtomIter(obatom):
        if neigh.GetAtomicNum() == 8:
            numOs += 1
    return numOs == 4


def is_phosphate_P(obatom: ob.OBAtom):
    if obatom.GetAtomicNum() != 15:  # Phosphor
        return False
    for neigh in ob.OBAtomAtomIter(obatom):
        if neigh.GetAtomicNum() != 8:  # It's a phosphate, only O
            return False
    return True


def is_carboxylate_C(obatom: ob.OBAtom):
    if obatom.GetAtomicNum() != 6:  # It's a carbon atom
        return False
    numOs = numCs = 0
    for neigh in ob.OBAtomAtomIter(obatom):
        neigh_z = neigh.GetAtomicNum()
        if neigh_z == 8:
            numOs += 1
        elif neigh_z == 6:
            numCs += 1
    return numOs == 2 and numCs == 1


def is_halocarbon_X(obatom: ob.OBAtom) -> bool:
    if obatom.GetAtomicNum() not in [9, 17, 35, 53]:  # Halogen atoms
        return False
    for neigh in ob.OBAtomAtomIter(obatom):
        if neigh.GetAtomicNum() == 6:
            return True
    return False

src/scoring/ligand_utils.py
SeonghwanSeo/PharmacoNet
import itertools
import numpy as np
import math

from typing import Tuple, List, Tuple
from numpy.typing import NDArray

from .ligand import LigandNode
from .pharmacophore_model import ModelNode


DISTANCE_SIGMA_THRESHOLD = 2.
PASS_THRESHOLD = 0.5


def scoring_matching_pair(
    cluster_node_match_list1: List[Tuple[LigandNode, List[ModelNode], NDArray[np.float32]]],
    cluster_node_match_list2: List[Tuple[LigandNode, List[ModelNode], NDArray[np.float32]]],
    num_conformers: int,
) -> Tuple[float, ...]:
    match_scores = np.zeros((num_conformers,), dtype=np.float32)
    num_fails = np.zeros((num_conformers,), dtype=np.int16)

    total_match = len(cluster_node_match_list1) * len(cluster_node_match_list2)
    match_threshold = total_match * (1 - PASS_THRESHOLD)

    num_pass = np.empty((num_conformers,), dtype=np.int16)
    likelihood = np.empty((num_conformers,), dtype=np.float32)
    for cluster_node_match1, cluster_node_match2 in itertools.product(cluster_node_match_list1, cluster_node_match_list2):
        ligand_node1, model_node_list1, weights1 = cluster_node_match1
        ligand_node2, model_node_list2, weights2 = cluster_node_match2
        ligand_edge = ligand_node1.neighbor_edge_dict[ligand_node2]
        distances = ligand_edge.distances

        num_match = len(model_node_list1) * len(model_node_list2)
        means = np.array(
            [
                [model_node1.neighbor_edge_dict[model_node2].distance_mean]
                for model_node1, model_node2 in itertools.product(model_node_list1, model_node_list2)
            ],
            dtype=np.float32
        )   # [M*N]
        stds = np.array(
            [
                [model_node1.neighbor_edge_dict[model_node2].distance_std]
                for model_node1, model_node2 in itertools.product(model_node_list1, model_node_list2)
            ],
            dtype=np.float32
        )   # [M*N]
        weights = (weights1.reshape(-1, 1) * weights2.reshape(1, -1)).reshape(-1)   # [M * N]

        weights_sum = sum(weights)
        normalize_coeff = 1 / (math.sqrt(2 * math.pi) * weights_sum)
        score_coeff = weights_sum / num_match

        distance_sigma_array = (distances.reshape(1, num_conformers) - means) / stds
        np.sum(distance_sigma_array < DISTANCE_SIGMA_THRESHOLD, axis=0, out=num_pass)
        num_fails += (num_pass < (num_match * PASS_THRESHOLD))
        if min(num_fails) > match_threshold:
            return (-1,) * num_conformers

        np.dot(weights / stds, np.exp(-0.5 * distance_sigma_array ** 2), out=likelihood)
        match_scores += likelihood * normalize_coeff * score_coeff

    return tuple(float(score) if num_fail <= match_threshold else -1 for score, num_fail in zip(match_scores, num_fails))


def scoring_matching_self(
    cluster_node_match_list: List[Tuple[LigandNode, List[ModelNode], NDArray[np.float32]]],
    num_conformers: int,
) -> Tuple[float, ...]:
    match_scores = np.zeros((num_conformers,), dtype=np.float32)
    likelihood = np.empty((num_conformers,), dtype=np.float32)
    for cluster_node_match1, cluster_node_match2 in itertools.combinations(cluster_node_match_list, 2):
        ligand_node1, model_node_list1, weights1 = cluster_node_match1
        ligand_node2, model_node_list2, weights2 = cluster_node_match2
        ligand_edge = ligand_node1.neighbor_edge_dict[ligand_node2]
        distances = ligand_edge.distances

        num_match = len(model_node_list1) * len(model_node_list2)
        means = np.array(
            [
                [model_node1.neighbor_edge_dict[model_node2].distance_mean]
                for model_node1, model_node2 in itertools.product(model_node_list1, model_node_list2)
            ],
            dtype=np.float32
        )   # [M*N]
        stds = np.array(
            [
                [model_node1.neighbor_edge_dict[model_node2].distance_std]
                for model_node1, model_node2 in itertools.product(model_node_list1, model_node_list2)
            ],
            dtype=np.float32
        )   # [M*N]
        weights = (weights1.reshape(-1, 1) * weights2.reshape(1, -1)).reshape(-1)   # [M*N]
        weights_sum = sum(weights)
        normalize_coeff = 1 / (math.sqrt(2 * math.pi) * weights_sum)
        score_coeff = weights_sum / num_match

        distance_sigma_array = (distances.reshape(1, num_conformers) - means) / stds
        np.dot(weights / stds, np.exp(-0.5 * distance_sigma_array ** 2), out=likelihood)
        match_scores += likelihood * normalize_coeff * score_coeff

    return tuple(match_scores.tolist())

src/scoring/match_utils.py
SeonghwanSeo/PharmacoNet
import math
import numpy as np
import numba as nb
import itertools

from typing import Tuple, List, Tuple
from numpy.typing import NDArray

from .ligand import LigandNode
from .pharmacophore_model import ModelNode


DISTANCE_SIGMA_THRESHOLD = 2.
PASS_THRESHOLD = 0.5


@nb.njit('void(float32[::1],float32[:, :, ::1],float32[::1],float32[::1],float32[::1],int16[::1])', fastmath=True, cache=True)
def __numba_run(
    distances: NDArray[np.float32],
    mean_stds: NDArray[np.float32],
    weights1: NDArray[np.float32],
    weights2: NDArray[np.float32],
    score_array: NDArray[np.float32],
    fail_array: NDArray[np.int16],
):
    """Scoring Function

    Args:
        distances: [C,]
        mean_stds: [M, N, 2]
        weights1: [M,]
        weights2: [N,]
        score_array: [C,]
        fail_array: [C,]
    """
    assert mean_stds.shape[0] == weights1.shape[0]
    assert mean_stds.shape[1] == weights2.shape[0]
    assert distances.shape[0] == score_array.shape[0] == fail_array.shape[0]

    num_match: int
    pass_threshold: float

    W1: float
    W1: float
    normalize_coeff: float
    score_coeff: float

    num_pass: int
    sigma_sq: float
    likelihood: float
    _likelihood: float

    M = weights1.shape[0]
    N = weights2.shape[0]
    C = distances.shape[0]

    num_match = M * N
    pass_threshold = num_match * 0.5    # PASS_THRESHOLD

    # NOTE: Coefficient Calculation
    W1 = sum(weights1)
    W2 = sum(weights2)
    normalize_coeff = 1 / (W1 * W2)
    score_coeff = (W1 * W2) / num_match

    for c in range(C):
        d = distances[c]
        num_pass = 0
        likelihood = 0.
        for m in range(M):
            w1 = weights1[m]
            _likelihood = 0.
            for n in range(N):
                w2 = weights2[n]
                mu_std = mean_stds[m, n]
                mu = mu_std[0]
                std = mu_std[1]
                sigma_sq = ((d - mu) / std) ** 2
                num_pass += (sigma_sq < 4.)  # abs(sigma) < 2.0 A (DISTANCE_SIGMA_THRESHOLD)
                _likelihood += w2 / std * math.exp(-0.5 * sigma_sq)
            likelihood += w1 * _likelihood

        score_array[c] += likelihood * normalize_coeff * score_coeff
        fail_array[c] += num_pass < pass_threshold


def __get_distance_mean_std(model_node1: ModelNode, model_node2: ModelNode) -> Tuple[float, float]:
    edge = model_node1.neighbor_edge_dict[model_node2]
    return edge.distance_mean, edge.distance_std


def scoring_matching_pair(
    cluster_node_match_list1: List[Tuple[LigandNode, List[ModelNode], NDArray[np.float32]]],
    cluster_node_match_list2: List[Tuple[LigandNode, List[ModelNode], NDArray[np.float32]]],
    num_conformers: int,
) -> Tuple[float, ...]:

    match_threshold = len(cluster_node_match_list1) * len(cluster_node_match_list2) * (1 - PASS_THRESHOLD)

    match_scores = np.zeros((num_conformers,), dtype=np.float32)
    num_fails = np.zeros((num_conformers,), dtype=np.int16)
    for ligand_node1, model_node_list1, weights1 in cluster_node_match_list1:
        for ligand_node2, model_node_list2, weights2 in cluster_node_match_list2:
            ligand_edge = ligand_node1.neighbor_edge_dict[ligand_node2]
            distances = ligand_edge.distances

            mean_stds = np.array([
                [__get_distance_mean_std(model_node1, model_node2) for model_node2 in model_node_list2]
                for model_node1 in model_node_list1
            ], dtype=np.float32)   # [M, N, 2]
            __numba_run(
                distances,
                mean_stds,
                weights1,
                weights2,
                match_scores,
                num_fails
            )
            if min(num_fails) > match_threshold:
                return (-1,) * num_conformers

    return tuple(float(score) if num_fail <= match_threshold else -1 for score, num_fail in zip(match_scores, num_fails))


def scoring_matching_self(
    cluster_node_match_list: List[Tuple[LigandNode, List[ModelNode], NDArray[np.float32]]],
    num_conformers: int,
) -> Tuple[float, ...]:
    match_scores = np.zeros((num_conformers,), dtype=np.float32)
    num_fails = np.zeros((num_conformers,), dtype=np.int16)
    for match1, match2 in itertools.combinations(cluster_node_match_list, 2):
        ligand_node1, model_node_list1, weights1 = match1
        ligand_node2, model_node_list2, weights2 = match2

        ligand_edge = ligand_node1.neighbor_edge_dict[ligand_node2]
        distances = ligand_edge.distances

        mean_stds = np.array([
            [__get_distance_mean_std(model_node1, model_node2) for model_node2 in model_node_list2]
            for model_node1 in model_node_list1
        ], dtype=np.float32)   # [M, N, 2]
        __numba_run(
            distances,
            mean_stds,
            weights1,
            weights2,
            match_scores,
            num_fails
        )

    return tuple(match_scores.tolist())

src/scoring/match_utils_numba.py
SeonghwanSeo/PharmacoNet
from __future__ import annotations

import pickle
from openbabel import pybel
from rdkit.Chem import Mol

from typing import Dict, List, Tuple, Set, Optional
from numpy.typing import NDArray

from .density_map import DensityMapGraph, DensityMapNode, DensityMapNodeCluster, DensityMapEdge


INTERACTION_TO_PHARMACOPHORE = {
    'Hydrophobic': 'Hydrophobic',
    'PiStacking_P': 'Aromatic',
    'PiStacking_T': 'Aromatic',
    'PiCation_lring': 'Aromatic',
    'PiCation_pring': 'Cation',
    'HBond_pdon': 'HBond_acceptor',
    'HBond_ldon': 'HBond_donor',
    'SaltBridge_pneg': 'Cation',
    'SaltBridge_lneg': 'Anion',
    'XBond': 'Halogen',
}


# NOTE: Pickle-Friendly Object
class PharmacophoreModel():
    def __init__(self):
        self.nodes: List[ModelNode]
        self.edges: List[ModelEdge]
        self.node_dict: Dict[str, List[ModelNode]]
        self.node_cluster_dict: Dict[str, List[ModelNodeCluster]]
        self.node_clusters: List[ModelNodeCluster]

    @classmethod
    def create(
        cls,
        center: Tuple[float, float, float],
        resolution: float,
        size: int,
        density_maps: List[dict],
    ):
        graph = DensityMapGraph(center, resolution, size)
        for node in density_maps:
            graph.add_node(node['type'], node['coords'], node['score'], node['map'])
        graph.setup()

        model = cls()
        model.nodes = [ModelNode.create(model, node) for node in graph.nodes]
        model.edges = [ModelEdge.create(model, edge) for edge in graph.edges]
        for node in model.nodes:
            node.setup()
        model.node_dict = {
            typ: [model.nodes[node.index] for node in node_list]
            for typ, node_list in graph.node_dict.items()
        }
        model.node_cluster_dict = {
            typ: [ModelNodeCluster.create(model, cluster) for cluster in cluster_list]
            for typ, cluster_list in graph.node_cluster_dict.items()
        }
        model.node_clusters = []
        for node_cluster_list in model.node_cluster_dict.values():
            model.node_clusters.extend(node_cluster_list)
        del graph
        return model

    def scoring(
        self,
        ligand_pbmol: pybel.Molecule,
        ligand_rbmol: Optional[Mol] = None,
        atom_positions: Optional[NDArray] = None,
        conformer_axis: Optional[int] = None,
    ) -> float:
        """Scoring Function

        Args:
            ligand_pbmol: pybel.Molecule
            ligand_rdmol: Chem.Mol | None
            atom_positions: List[NDArray[np.float32]] | NDArray[np.float32] | None
            conformer_axis: Optional[int]

            case: atom_positions: NDArray[np.float32]
                i) conformer_axis is 0 or None
                    atom_positions: (N_conformers, N_atoms, 3)
                ii) conformer_axis is 1
                    atom_positions: (N_atoms, N_conformers, 3)
            case: atom_positions: None
                Using RDKit Conformer informations
        """
        from .ligand import Ligand
        from .graph_match import GraphMatcher
        ligand = Ligand(ligand_pbmol, ligand_rbmol, atom_positions, conformer_axis)
        matcher = GraphMatcher(self, ligand)
        return matcher.scoring()

    def save(self, save_path: str):
        state = self.__getstate__()
        with open(save_path, 'wb') as w:
            pickle.dump(state, w)

    @classmethod
    def load(cls, save_path: str):
        with open(save_path, 'rb') as f:
            state = pickle.load(f)
        model = cls()
        model.__setstate__(state)
        return model

    def __getstate__(self):
        state = dict(
            nodes=[node.get_kwargs() for node in self.nodes],
            edges=[edge.get_kwargs() for edge in self.edges],
            node_cluster_dict={typ: [cluster.get_kwargs() for cluster in cluster_list] for typ, cluster_list in self.node_cluster_dict.items()},
            node_dict={typ: [node.index for node in nodes] for typ, nodes in self.node_dict.items()},
        )
        return state

    def __setstate__(self, state):
        self.nodes = [ModelNode(self, **kwargs) for kwargs in state['nodes']]
        self.edges = [ModelEdge(self, **kwargs) for kwargs in state['edges']]
        for node in self.nodes:
            node.setup()
        self.node_dict = {
            typ: [self.nodes[index] for index in indices] for typ, indices in state['node_dict'].items()
        }
        self.node_cluster_dict = {
            typ: [ModelNodeCluster(self, **kwargs) for kwargs in cluster_list]
            for typ, cluster_list in state['node_cluster_dict'].items()
        }
        self.node_clusters: List[ModelNodeCluster] = []
        for node_cluster_list in self.node_cluster_dict.values():
            self.node_clusters.extend(node_cluster_list)


class ModelNodeCluster():
    def __init__(
        self,
        graph: PharmacophoreModel,
        cluster_type: str,
        node_indices: Set[int],
        node_types: Set[str],
        center: Tuple[float, float, float],
        size: float,
    ):
        self.type: str = cluster_type
        self.nodes: Set[ModelNode] = {graph.nodes[index] for index in node_indices}
        self.node_indices: Set[int] = node_indices
        self.node_types: Set[str] = node_types

        self.center: Tuple[float, float, float] = center
        self.size: float = size

    @classmethod
    def create(cls, graph: PharmacophoreModel, cluster: DensityMapNodeCluster) -> ModelNodeCluster:
        return cls(
            graph,
            cluster.type,
            {node.index for node in cluster.nodes},
            {INTERACTION_TO_PHARMACOPHORE[node.type] for node in cluster.nodes},
            cluster.center,
            cluster.size
        )

    def __repr__(self):
        return f'ModelCluster({self.type})[{self.nodes}]'

    def get_kwargs(self):
        return dict(
            cluster_type=self.type,
            node_indices=self.node_indices,
            node_types=self.node_types,
            center=self.center,
            size=self.size,
        )


class ModelNode():
    def __init__(
        self,
        graph: PharmacophoreModel,
        index: int,
        interaction_type: str,
        hotspot_position: Tuple[float, float, float],
        score: float,
        center: Tuple[float, float, float],
        radius: float,
        neighbor_edge_dict: Dict[int, int],
        overlapped_nodes: List[int],
    ):
        self.graph: PharmacophoreModel = graph
        self.index: int = index
        self.interaction_type: str = interaction_type
        self.hotspot_position: Tuple[float, float, float] = hotspot_position
        self.type: str = INTERACTION_TO_PHARMACOPHORE[interaction_type]
        self.score: float = score
        self.center: Tuple[float, float, float] = center
        self.radius: float = radius

        self._neighbor_edge_dict: Dict[int, int] = neighbor_edge_dict
        self._overlapped_nodes: List[int] = overlapped_nodes
        self.neighbor_edge_dict: Dict[ModelNode, ModelEdge]
        self.overlapped_nodes: List[ModelNode]

    def setup(self):
        self.neighbor_edge_dict = {
            self.graph.nodes[node_index]: self.graph.edges[edge_index] for node_index, edge_index in self._neighbor_edge_dict.items()
        }
        self.overlapped_nodes = [self.graph.nodes[node_index] for node_index in self._overlapped_nodes]

    @classmethod
    def create(cls, graph: PharmacophoreModel, node: DensityMapNode) -> ModelNode:
        x, y, z = node.center.tolist()
        center = (x, y, z)
        return cls(
            graph,
            node.index,
            node.type,
            node.hotspot_position,
            node.score,
            center,
            node.radius,
            {neighbor.index: edge.index for neighbor, edge in node.neighbor_edge_dict.items()},
            [node.index for node in node.overlapped_nodes],
        )

    def __hash__(self):
        return self.index

    def get_kwargs(self):
        return dict(
            index=self.index,
            interaction_type=self.interaction_type,
            hotspot_position=self.hotspot_position,
            score=self.score,
            center=self.center,
            radius=self.radius,
            neighbor_edge_dict=self._neighbor_edge_dict,
            overlapped_nodes=self._overlapped_nodes
        )

    def __repr__(self):
        return f'ModelNode({self.index})[{self.type}]'


class ModelEdge():
    def __init__(
        self,
        graph: PharmacophoreModel,
        index: int,
        node_indices: Tuple[int, int],
        edge_type: Tuple[str, str],
        distance_mean: float,
        distance_std: float,
    ):
        self.graph: PharmacophoreModel = graph
        self.index: int = index
        self.nodes: Tuple[ModelNode, ModelNode] = (self.graph.nodes[node_indices[0]], self.graph.nodes[node_indices[1]])
        self.node_indices: Tuple[int, int] = node_indices
        self.type: Tuple[str, str] = edge_type
        self.distance_mean: float = distance_mean
        self.distance_std: float = distance_std

    @classmethod
    def create(cls, graph: PharmacophoreModel, edge: DensityMapEdge) -> ModelEdge:
        return cls(
            graph,
            edge.index,
            edge.node_indices,
            edge.type,
            edge.distance_mean,
            edge.distance_std,
        )

    def __hash__(self):
        return self.index

    def get_kwargs(self):
        return dict(
            index=self.index,
            node_indices=self.node_indices,
            edge_type=self.type,
            distance_mean=self.distance_mean,
            distance_std=self.distance_std,
        )

src/scoring/pharmacophore_model.py
SeonghwanSeo/PharmacoNet
from __future__ import annotations

from typing import Dict, List, Tuple, Optional, Iterator

from .ligand import LigandNodeCluster
from .pharmacophore_model import ModelNodeCluster


# NOTE: TYPE
LigandClusterPair = Tuple[LigandNodeCluster, LigandNodeCluster]
ModelClusterPair = Tuple[ModelNodeCluster, ModelNodeCluster]


class ClusterMatchTree():
    def __init__(
        self,
        model_cluster: Optional[ModelNodeCluster],
        pair_scores: Optional[Dict[int, float]],
        parent: ClusterMatchTree,
    ):
        self.level: int = parent.level + 1
        self.num_matches: int = parent.num_matches + (model_cluster is not None)
        self.parent: ClusterMatchTree = parent
        self.root: ClusterMatchTreeRoot = parent.root
        self.children: List[ClusterMatchTree] = []

        ligand_cluster: LigandNodeCluster = self.root.ligand_cluster_list[self.level]
        self.ligand_cluster: LigandNodeCluster = ligand_cluster
        self.model_cluster: Optional[ModelNodeCluster] = model_cluster

        self.pair_scores: Dict[int, float]
        if model_cluster is not None:
            assert pair_scores is not None
            self_pair_scores = self.root.matching_pair_scores_dict[ligand_cluster, ligand_cluster][model_cluster, model_cluster]
            self.pair_scores = {
                conformer_id: parent.pair_scores[conformer_id] + self_pair_scores[conformer_id] + score
                for conformer_id, score in pair_scores.items()
            }
        else:
            self.pair_scores = parent.pair_scores

    @property
    def score(self) -> float:
        return self.edge_score

    @property
    def edge_score(self) -> float:
        if self.num_matches == 0:
            return 0.
        return max(self.pair_scores.values())

    @property
    def conformer_ids(self):
        return self.pair_scores.keys()

    def dfs_run(
        self,
        match_dict: Dict[LigandNodeCluster, Dict[ModelNodeCluster, Dict[int, float]]]
    ) -> int:
        """recursive function

        Args:
            level: level of new node
            ligand_cluster: ligand cluster according to the level
            model_cluster_dict: candidate model cluster
                ModelCluster: {conformer_id: accumulate_score}
        """
        upd_match_dict: Dict[LigandNodeCluster, Dict[ModelNodeCluster, Dict[int, float]]] = {}
        if self.model_cluster is not None:
            for ligand_cluster, model_cluster_dict in match_dict.items():
                upd_model_cluster_dict = {}
                matching_pair_scores_dict = self.root.matching_pair_scores_dict[self.ligand_cluster, ligand_cluster]
                for model_cluster, conformer_pair_score_dict in model_cluster_dict.items():
                    pair_score_list: Tuple[float, ...] = matching_pair_scores_dict[self.model_cluster, model_cluster]
                    # NOTE: Update Model Cluster List accoring to Validity of Pair (Use only Valid Conformer)
                    upd_conformer_pair_score_dict: Dict[int, float] = {
                        conformer_id: total_score + pair_score_list[conformer_id]
                        for conformer_id, total_score in conformer_pair_score_dict.items()
                        if conformer_id in self.conformer_ids and pair_score_list[conformer_id] > 0
                    }
                    if len(upd_conformer_pair_score_dict) > 0:
                        upd_model_cluster_dict[model_cluster] = upd_conformer_pair_score_dict
                upd_match_dict[ligand_cluster] = upd_model_cluster_dict
        else:
            upd_match_dict = match_dict.copy()

        # NOTE: Add Child
        if self.level < len(self.root.ligand_cluster_list) - 1:
            child_ligand_cluster = self.root.ligand_cluster_list[self.level + 1]
            model_cluster_dict = upd_match_dict.pop(child_ligand_cluster)
            max_num_matches = 0
            for model_cluster, conformer_pair_score_dict in model_cluster_dict.items():
                child = self.add_child(model_cluster, conformer_pair_score_dict)
                num_matches = child.dfs_run(upd_match_dict)
                max_num_matches = max(num_matches, max_num_matches)
            if len(model_cluster_dict) == 0 or (self.num_matches + max_num_matches) < 5:
                child = self.add_child(None, None)
                num_matches = child.dfs_run(upd_match_dict)
                max_num_matches = max(num_matches, max_num_matches)
            return max_num_matches + int(self.model_cluster is not None)
        else:
            return int(self.model_cluster is not None)

    def add_child(self, model_cluster: Optional[ModelNodeCluster], pair_score_dict: Optional[Dict[int, float]]):
        child = ClusterMatchTree(model_cluster, pair_score_dict, self)
        self.children.append(child)
        return child

    def delete(self):
        assert self.level >= 0
        self.parent.children.remove(self)
        del self

    @property
    def size(self) -> int:
        if len(self.children) == 0:
            return 1
        size = 0
        for node in self.children:
            size += node.size
        return size

    @property
    def key(self) -> List[Optional[ModelNodeCluster]]:
        key = []
        node: ClusterMatchTree = self
        while node is not self.root:
            key.append(node.model_cluster)
            node = node.parent
        key.reverse()
        return key

    @property
    def item(self) -> Dict[LigandNodeCluster, Optional[ModelNodeCluster]]:
        node: ClusterMatchTree = self
        graph_match: Dict[LigandNodeCluster, Optional[ModelNodeCluster]] = {}
        while node is not self.root:
            graph_match[node.ligand_cluster] = node.model_cluster
            node = node.parent
        return graph_match

    def iteration(self, level: Optional[int] = None) -> Iterator[ClusterMatchTree]:
        if level is not None:
            yield from self.iteration_level(level)
        else:
            yield from self.iteration_leaf()

    def iteration_level(self, level: int) -> Iterator[ClusterMatchTree]:
        assert self.level <= level
        if self.level < level:
            for node in self.children:
                yield from node.iteration_level(level)
        elif self.level == level:
            yield self

    def iteration_leaf(self) -> Iterator[ClusterMatchTree]:
        if len(self.children) > 0:
            for node in self.children:
                yield from node.iteration_leaf()
        else:
            yield self

    def __tree_repr__(self):
        repr = '  ' * (self.level + 1) + f'- {self.model_cluster}'
        if len(self.children) > 0:
            repr += '  ' * (self.level + 2)
            repr += f'(level {self.level + 1}) {self.children[0].ligand_cluster}\n'
            for child in self.children:
                repr += child.__tree_repr__()
            repr += '\n'
        return repr

    def __repr__(self):
        repr = ''
        tree = self
        while tree is not self.root:
            repr = f'({tree.ligand_cluster}, {tree.model_cluster})\n' + repr
            tree = tree.parent
        return repr


class ClusterMatchTreeRoot(ClusterMatchTree):
    def __init__(
        self,
        ligand_cluster_list: List[LigandNodeCluster],
        cluster_match_dict: Dict[LigandNodeCluster, List[ModelNodeCluster]],
        matching_pair_score_dict: Dict[LigandClusterPair, Dict[ModelClusterPair, Tuple[float, ...]]],
        num_conformers: int,
    ):
        self.root = self
        self.level: int = -1
        self.num_matches: int = 0
        self.num_conformers: int = num_conformers
        self.children: List[ClusterMatchTree] = []
        self.ligand_cluster_list: List[LigandNodeCluster] = ligand_cluster_list
        self.cluster_match_dict: Dict[LigandNodeCluster, List[ModelNodeCluster]] = cluster_match_dict
        self.matching_pair_scores_dict: Dict[Tuple[LigandNodeCluster, LigandNodeCluster], Dict[Tuple[ModelNodeCluster, ModelNodeCluster], Tuple[float, ...]]] = matching_pair_score_dict

        self.model_cluster = None
        self.pair_scores: Dict[int, float] = {conformer_id: 0. for conformer_id in range(num_conformers)}

    def __repr__(self):
        repr = 'Root\n'
        if len(self.children) > 0:
            repr += f'(level {self.level + 1}) {self.children[0].ligand_cluster}\n'
        for child in self.children:
            repr += child.__tree_repr__()
        return repr

    def run(self):
        match_dict = {
            ligand_cluster: {
                model_cluster: {conformer_id: 0. for conformer_id in range(self.num_conformers)}
                for model_cluster in self.cluster_match_dict[ligand_cluster]
            }
            for ligand_cluster in self.ligand_cluster_list
        }
        self.dfs_run(match_dict)

src/scoring/tree.py
SeonghwanSeo/PharmacoNet
import os
from openbabel import pybel
from typing import Dict, Union, Sequence

import numpy as np
from numpy.typing import NDArray


def load_ligand(ligand_path: str) -> pybel.Molecule:
    extension = os.path.splitext(ligand_path)[1]
    assert extension in ['.sdf', '.pdb', '.mol2']
    return next(pybel.readfile(extension[1:], ligand_path))


def get_score_threshold(
    score_distributions: Dict[str, Dict[str, NDArray[np.float32]]],
    relative_score_threshold: Union[float, Dict[str, float]],
    interaction_list: Sequence[str],
) -> Dict[str, float]:
    abs_score_threshold = {}
    for interaction_type in interaction_list:
        focus_scores = score_distributions[interaction_type]['focus']
        if isinstance(relative_score_threshold, float):
            rel_threshold = relative_score_threshold
        else:
            rel_threshold = relative_score_threshold[interaction_type]
        if rel_threshold == 0.0:
            abs_score_threshold[interaction_type] = float('inf')
        else:
            abs_score_threshold[interaction_type] = focus_scores[int(len(focus_scores) * (1 - rel_threshold))]
    return abs_score_threshold

src/utils/__init__.py
SeonghwanSeo/PharmacoNet
from typing import Sequence, Set

INTERACTION_LIST: Sequence[str] = (
    'Hydrophobic',
    'PiStacking_P',
    'PiStacking_T',
    'PiCation_lring',
    'PiCation_pring',
    'HBond_ldon',
    'HBond_pdon',
    'SaltBridge_lneg',
    'SaltBridge_pneg',
    'XBond'
)

NUM_INTERACTION_TYPES: int = 10

HYDROPHOBIC = 0
PISTACKING_P = 1
PISTACKING_T = 2
PICATION_LRING = 3
PICATION_PRING = 4
HBOND_LDON = 5
HBOND_PDON = 6
SALTBRIDGE_LNEG = 7
SALTBRIDGE_PNEG = 8
XBOND = 9

# PLIP Distance + 0.5 A
INTERACTION_DIST = {
    HYDROPHOBIC: 4.5,       # 4.0 + 0.5
    PISTACKING_P: 6.0,      # 5.5 + 0.5
    PISTACKING_T: 6.0,      # 5.5 + 0.5
    PICATION_LRING: 6.5,    # 6.0 + 0.5
    PICATION_PRING: 6.5,    # 6.0 + 0.5
    HBOND_LDON: 4.5,        # 4.1 + 0.5 - 0.1 (to be devided to 0.5)
    HBOND_PDON: 4.5,        # 4.1 + 0.5 - 0.1
    SALTBRIDGE_LNEG: 6.0,   # 5.5 + 0.5
    SALTBRIDGE_PNEG: 6.0,   # 5.5 + 0.5
    XBOND: 4.5,             # 4.0 + 0.5
}

LONG_INTERACTION: Set[int] = {
    PISTACKING_P,
    PISTACKING_T,
    PICATION_PRING,
    PICATION_LRING,
    SALTBRIDGE_LNEG,
    SALTBRIDGE_PNEG
}

SHORT_INTERACTION: Set[int] = {
    HYDROPHOBIC,
    HBOND_LDON,
    HBOND_PDON,
    XBOND,
}


src/utils/constant.py
SeonghwanSeo/PharmacoNet
import torch
from torch import nn
from torch.nn import functional as F

from typing import Tuple, Union, List
from torch import Tensor


def clip(value, lower, upper):
    return lower if value < lower else upper if value > upper else value


def to_3tuple(value):
    if not isinstance(value, tuple):
        return (value, value, value)
    else:
        return value


class GaussianSmoothing(nn.Module):
    def __init__(
        self,
        kernel_size: Union[int, Tuple[int, int, int]],
        sigma: Union[float, Tuple[float, float, float]],
    ):
        super(GaussianSmoothing, self).__init__()
        kernel_size = to_3tuple(kernel_size)
        sigma = to_3tuple(sigma)

        # The gaussian kernel is the product of the
        # gaussian function of each dimension.
        meshgrids = torch.meshgrid(
            [torch.arange(size, dtype=torch.float) for size in kernel_size], indexing='ij'
        )
        kernel: torch.Tensor = None
        for size, std, mgrid in zip(kernel_size, sigma, meshgrids):
            mean = (size - 1) / 2
            # _kernel = 1 / (std * math.sqrt(2 * math.pi)) * torch.exp(-((mgrid - mean) / (2 * std)) ** 2)
            _kernel = torch.exp(-((mgrid - mean) / (std)) ** 2 / 2)     # omit constant part
            if kernel is None:
                kernel = _kernel
            else:
                kernel *= _kernel

        # Make sure sum of values in gaussian kernel equals 1.
        kernel /= torch.sum(kernel)                     # (Kd, Kh, Kw), Kd, Kh, Kw: kernel_size

        # Reshape to depthwise convolutional weight
        kernel = kernel.view(1, 1, *kernel.size())      # (1, 1, Kd, Kh, Kw)
        kernel = kernel.repeat(1, 1, 1, 1, 1)           # (1, 1, Kd, Kh, Kw)

        self.register_buffer('weight', kernel)
        self.pad: Tuple[int, int, int, int, int, int] = \
            (kernel_size[0] // 2, kernel_size[0] // 2,
             kernel_size[1] // 2, kernel_size[1] // 2,
             kernel_size[2] // 2, kernel_size[2] // 2,)

    @torch.no_grad()
    def forward(self, x):
        """
        Apply gaussian filter to input.
        Arguments:
            input (torch.Tensor): Input to apply gaussian filter on.
        Returns:
            filtered (torch.Tensor): Filtered output.
        """
        x = F.pad(x, self.pad, mode='constant', value=0.)
        weight = self.weight.repeat(x.shape[-4], 1, 1, 1, 1)
        return F.conv3d(x, weight=weight, groups=x.shape[-4])

src/utils/smoothing.py
SeonghwanSeo/PharmacoNet
import os
import argparse

import pymol
from pymol import cmd

import sys
sys.path.append(".")
sys.path.append("..")
from src.graph.pharmacophore_model import PharmacophoreModel


class Visualize_ArgParser(argparse.ArgumentParser):
    def __init__(self):
        super().__init__('scoring')
        self.formatter_class = argparse.ArgumentDefaultsHelpFormatter
        self.add_argument('-p', '--pharmacophore_model', type=str, help='path to save pharmacophore model (.pkl)', required=True)
        self.add_argument('-r', '--receptor', type=str, help='path of receptor file')
        self.add_argument('-l', '--ligand', type=str, help='path of ligand file')
        self.add_argument('-o', '--out', type=str, help='path of pymol session file (.pse)', required=True)


PHARMACOPHORE_COLOR_DICT = {
    'Hydrophobic': 'orange',
    'Aromatic': 'deeppurple',
    'Cation': 'blue',
    'Anion': 'red',
    'HBond_acceptor': 'magenta',
    'HBond_donor': 'cyan',
    'Halogen': 'yellow',
}

INTERACTION_COLOR_DICT = {
    'Hydrophobic': 'orange',
    'PiStacking_P': 'deeppurple',
    'PiStacking_T': 'deeppurple',
    'PiCation_lring': 'blue',
    'PiCation_pring': 'deeppurple',
    'HBond_ldon': 'magenta',
    'HBond_pdon': 'cyan',
    'SaltBridge_lneg': 'blue',
    'SaltBridge_pneg': 'red',
    'XBond': 'yellow',
}


if __name__ == '__main__':
    parser = Visualize_ArgParser()
    args = parser.parse_args()
    RECEPTOR_PATH = args.receptor
    LIGAND_PATH = args.ligand
    PHARMACOPHORE_MODEL_PATH = args.pharmacophore_model
    SAVE_PATH = args.out

    pymol.pymol_argv = ['pymol', '-pcq']
    pymol.finish_launching(args=['pymol', '-pcq', '-K'])
    cmd.reinitialize()
    cmd.feedback('disable', 'all', 'everything')

    # NOTE: Draw Molecule
    if RECEPTOR_PATH:
        cmd.load(RECEPTOR_PATH)
        cmd.set_name(os.path.splitext(os.path.basename(RECEPTOR_PATH))[0], 'Protein')
        cmd.color('gray90', 'Protein')
    if LIGAND_PATH:
        cmd.load(LIGAND_PATH)
        cmd.set_name(os.path.splitext(os.path.basename(LIGAND_PATH))[0], 'Ligand')
        cmd.color('green', 'Ligand')

    # NOTE: Pharmacophore Model
    model = PharmacophoreModel.load(PHARMACOPHORE_MODEL_PATH)
    for node in model.nodes:
        protein_color = INTERACTION_COLOR_DICT[node.interaction_type]
        pharmacophore_color = PHARMACOPHORE_COLOR_DICT[node.type]
        protein_id = f'spot{node.index}'
        cmd.pseudoatom(protein_id, pos=node.hotspot_position, color=protein_color)
        cmd.set('sphere_color', protein_color, protein_id)
        cmd.set('sphere_scale', 0.5, protein_id)

        pharmacophore_id = f'pharmacophore{node.index}_{node.type}'
        cmd.pseudoatom(pharmacophore_id, pos=node.center, color=protein_color)
        cmd.set('sphere_color', pharmacophore_color, pharmacophore_id)
        cmd.set('sphere_scale', node.radius, pharmacophore_id)

        interaction_id = f'interaction{node.index}_{node.interaction_type}'
        cmd.distance(interaction_id, protein_id, pharmacophore_id)
        cmd.set('dash_color', pharmacophore_color, interaction_id)
        cmd.set('dash_gap', 0.2, interaction_id)
        cmd.set('dash_length', 0.4, interaction_id)

        result_id = f'result{node.index}_{node.interaction_type}'
        cmd.group(result_id, protein_id + ' ' + pharmacophore_id + ' ' + interaction_id)

    cmd.group('Result', 'result*')

    cmd.enable('all')
    cmd.hide('everything', 'all')
    if RECEPTOR_PATH:
        cmd.show('sticks', 'Protein')
    if LIGAND_PATH:
        cmd.show('sticks', 'Ligand')
    cmd.show('sphere', 'Result')
    cmd.show('dash', 'Result')

    cmd.bg_color('white')
    if RECEPTOR_PATH:
        cmd.set('stick_transparency', 0.5, 'Protein')
    if LIGAND_PATH:
        cmd.set('stick_transparency', 0.5, 'Ligand')

    cmd.set('sphere_transparency', 0.5, 'pharmacophore*')
    cmd.save(f'{SAVE_PATH}')

visualize.py
SeonghwanSeo/PharmacoNet