신약개발 AI 파이프라인

# ═══════════════════════════════════════════════
# STEP 02: 화합물 라이브러리 생성
# 벤치마크: Insilico Medicine Pharma.AI
# ═══════════════════════════════════════════════

import torch
import torch.nn as nn
from rdkit import Chem
from rdkit.Chem import Descriptors, QED, AllChem
import numpy as np

# ── 2-1. Lipinski Rule of 5 필터 ──
def lipinski_filter(smiles: str) -> bool:
    mol = Chem.MolFromSmiles(smiles)
    if not mol: return False
    mw    = Descriptors.MolWt(mol)
    logp  = Descriptors.MolLogP(mol)
    hbd   = Descriptors.NumHDonors(mol)
    hba   = Descriptors.NumHAcceptors(mol)
    return (mw < 500 and logp < 5 and hbd <= 5 and hba <= 10)

# ── 2-2. VAE 분자 생성 모델 ──
class MolecularVAE(nn.Module):
    def __init__(self, vocab_size=60, latent_dim=256):
        super().__init__()
        # 인코더: SMILES → 잠재벡터
        self.encoder = nn.Sequential(
            nn.Embedding(vocab_size, 128),
            nn.GRU(128, 512, batch_first=True)
        )
        self.fc_mu     = nn.Linear(512, latent_dim)
        self.fc_logvar = nn.Linear(512, latent_dim)
        # 디코더: 잠재벡터 → SMILES
        self.decoder = nn.Sequential(
            nn.Linear(latent_dim, 512),
            nn.GRU(512, 512, batch_first=True),
            nn.Linear(512, vocab_size)
        )
    
    def reparameterize(self, mu, logvar):
        std = torch.exp(0.5 * logvar)
        eps = torch.randn_like(std)
        return mu + eps * std
    
    def generate(self, n_samples=1000, latent_dim=256):
        """잠재공간 샘플링으로 신규 분자 생성"""
        z = torch.randn(n_samples, latent_dim)
        with torch.no_grad():
            decoded = self.decoder(z.unsqueeze(1))
        return decoded

# ── 2-3. GAN 기반 생성 모듈 ──
class MolGAN:
    def __init__(self):
        self.generator     = self._build_generator()
        self.discriminator = self._build_discriminator()
    
    def _build_generator(self) -> nn.Module:
        return nn.Sequential(
            nn.Linear(128, 512), nn.ReLU(),
            nn.Linear(512, 1024), nn.ReLU(),
            nn.Linear(1024, 60),  nn.Softmax(dim=-1)
        )
    
    def _build_discriminator(self) -> nn.Module:
        return nn.Sequential(
            nn.Linear(60, 512),  nn.LeakyReLU(0.2),
            nn.Linear(512, 256),  nn.LeakyReLU(0.2),
            nn.Linear(256, 1),    nn.Sigmoid()
        )

# ── 2-4. 통합 라이브러리 생성 파이프라인 ──
class LibraryGenerator:
    def __init__(self):
        self.vae = MolecularVAE()
        self.gan = MolGAN()
    
    def generate_library(self, n_candidates=10000) -> list:
        generated = []
        batch     = n_candidates // 100  # 배치 처리
        for _ in range(batch):
            raw = self.vae.generate(100)
            # SMILES 변환 후 Lipinski 필터 적용
            valid = [s for s in raw if lipinski_filter(s)]
            generated.extend(valid)
        print(f"✅ {len(generated):,}개 유효 후보 생성 (목표: {n_candidates:,})")
        return generated

if __name__ == "__main__":
    gen = LibraryGenerator()
    library = gen.generate_library(n_candidates=10000)

# ═══════════════════════════════════════════════
# STEP 03: 바인딩 예측 (Binding Prediction)
# 모델: Graph Convolutional Network (GCN)
# ═══════════════════════════════════════════════

import torch
import torch.nn as nn
import torch_geometric.nn as gnn
from torch_geometric.data import Data
from rdkit import Chem
from rdkit.Chem import AllChem
import numpy as np

# ── 3-1. 분자 → 그래프 변환 ──
def mol_to_graph(smiles: str) -> Data:
    """SMILES 문자열을 PyG 그래프로 변환"""
    mol   = Chem.MolFromSmiles(smiles)
    if not mol: return None
    
    # 원자 특성 (노드)
    atom_features = []
    for atom in mol.GetAtoms():
        features = [
            atom.GetAtomicNum(),          # 원자번호
            atom.GetDegree(),             # 결합 수
            atom.GetFormalCharge(),       # 형식 전하
            int(atom.GetIsAromatic()),    # 방향족 여부
            atom.GetHybridization().real  # 혼성화
        ]
        atom_features.append(features)
    
    # 결합 인덱스 (엣지)
    edge_index = []
    for bond in mol.GetBonds():
        i, j = bond.GetBeginAtomIdx(), bond.GetEndAtomIdx()
        edge_index.extend([[i,j],[j,i]])
    
    x  = torch.tensor(atom_features, dtype=torch.float)
    ei = torch.tensor(edge_index, dtype=torch.long).t().contiguous()
    return Data(x=x, edge_index=ei)

# ── 3-2. GCN 결합 예측 모델 ──
class BindingGCN(nn.Module):
    def __init__(self, node_feat=5, hidden=256):
        super().__init__()
        self.conv1 = gnn.GCNConv(node_feat, hidden)
        self.conv2 = gnn.GCNConv(hidden, hidden)
        self.conv3 = gnn.GCNConv(hidden, hidden // 2)
        self.pool  = gnn.global_mean_pool
        self.fc    = nn.Sequential(
            nn.Linear(hidden // 2, 128),
            nn.ReLU(), nn.Dropout(0.2),
            nn.Linear(128, 1)  # 결합 에너지 출력
        )
    
    def forward(self, data):
        x, ei, batch = data.x, data.edge_index, data.batch
        x = torch.relu(self.conv1(x, ei))
        x = torch.relu(self.conv2(x, ei))
        x = torch.relu(self.conv3(x, ei))
        x = self.pool(x, batch)
        return self.fc(x)  # ΔG 결합 자유에너지 예측

# ── 3-3. ADMET 예측 모듈 ──
class ADMETPredictor(nn.Module):
    def __init__(self):
        super().__init__()
        self.model = nn.Sequential(
            nn.Linear(2048, 512),  # Morgan 핑거프린트 입력
            nn.ReLU(), nn.Dropout(0.2),
            nn.Linear(512, 256),
            nn.ReLU(),
            nn.Linear(256, 5)   # A, D, M, E, T 5개 출력
        )
    
    def predict(self, smiles: str) -> dict:
        mol  = Chem.MolFromSmiles(smiles)
        fp   = AllChem.GetMorganFingerprintAsBitVect(mol, 2, 2048)
        x    = torch.tensor(list(fp), dtype=torch.float).unsqueeze(0)
        pred = self.model(x).squeeze().detach().numpy()
        return {
            'absorption':   pred[0],
            'distribution': pred[1],
            'metabolism':   pred[2],
            'excretion':    pred[3],
            'toxicity':     pred[4]
        }

# ── 3-4. 통합 스크리닝 파이프라인 ──
def screen_library(library: list, top_n=100) -> list:
    binding_model = BindingGCN()
    admet_model   = ADMETPredictor()
    
    results = []
    for smiles in library:
        graph      = mol_to_graph(smiles)
        binding_dg = binding_model(graph).item()
        admet      = admet_model.predict(smiles)
        if admet['toxicity'] < 0.3:  # 독성 필터
            results.append({'smiles': smiles, 'dG': binding_dg, 'admet': admet})
    
    return sorted(results, key=lambda x: x['dG'])[:top_n]