Upload folder using huggingface_hub

.gitattributes

@@ -33,3 +33,8 @@ saved_model/**/* filter=lfs diff=lfs merge=lfs -text
 *.zip filter=lfs diff=lfs merge=lfs -text
 *.zst filter=lfs diff=lfs merge=lfs -text
 *tfevents* filter=lfs diff=lfs merge=lfs -text
+cycpeptmpdb_raw.csv filter=lfs diff=lfs merge=lfs -text
+fig_architecture.png filter=lfs diff=lfs merge=lfs -text
+fig_dataset.png filter=lfs diff=lfs merge=lfs -text
+main.pdf filter=lfs diff=lfs merge=lfs -text
+results_scatter.png filter=lfs diff=lfs merge=lfs -text

.gitignore

@@ -0,0 +1,27 @@
+# Python
+__pycache__/
+*.py[cod]
+*.egg-info/
+.venv/
+venv/
+env/
+
+# LaTeX build artifacts
+*.aux
+*.log
+*.out
+*.toc
+*.blg
+*.bbl
+*.fls
+*.fdb_latexmk
+
+# Editor / OS
+.DS_Store
+.vscode/
+.idea/
+.ipynb_checkpoints/
+
+# Jupyter / misc
+*.ipynb
+nohup.out

.hfignore

@@ -0,0 +1,25 @@
+# Files to exclude from the Hugging Face Hub repository (.hfignore)
+# Used by `huggingface-cli upload` / `hf upload` (and upload_folder).
+
+__pycache__/
+*.py[cod]
+*.egg-info/
+.git/
+
+# LaTeX build artifacts (source main.tex and main.pdf are published, logs are not)
+*.aux
+*.log
+*.out
+*.toc
+*.blg
+*.bbl
+*.fls
+*.fdb_latexmk
+
+# Editor / OS
+.DS_Store
+.vscode/
+.idea/
+.ipynb_checkpoints/
+
+nohup.out

AGENTS.md

@@ -0,0 +1,42 @@
+# AGENTS.md
+
+Research repo: GNN baseline for cyclic peptide membrane permeability prediction on CycPeptMPDB (PAMPA assay), optionally augmented with CycPeptMPDB-4D conformational features. Companion paper in `main.tex` (IEEEtran conference format).
+
+## Key facts
+
+- **No tests or lint/typecheck config exist.** Do not invent test/lint commands. Verification is: run the training script and check it trains (or `python -m py_compile <script>` for a syntax-only check). `requirements.txt` and `README.md` exist; trust them.
+- Dependencies: `torch` (2.11.0), `torch_geometric` (2.7.0), `rdkit` (2022.09), `numpy`, `requests`, `transformers` (only `train2.py`). GPU required — runs on an 8 GB NVIDIA RTX 4060.
+- Published to GitHub (`devansh0703/CycPepGNN`) and Hugging Face Hub (`devansh0703/CycPepGNN`, model repo). `.gitignore` for GitHub, `.hfignore` for the HF upload. Keep both in sync when adding files.
+
+## Data loading is network-only
+
+- Scripts **download** the dataset at runtime via `requests`: CycPeptMPDB from `https://raw.githubusercontent.com/akiyamalab/cycpeptmp/main/data/CycPeptMPDB_Peptide_All.csv` and 4D features from Zenodo (`records/18754430`).
+- The checked-in `cycpeptmpdb_raw.csv` / `cycpeptmpdb_4d.csv` are **not used** by any script. Do not assume they replace the downloads.
+- Scripts fail with no network access.
+
+## Training scripts
+
+| Script | Model | Notes |
+|---|---|---|
+| `train.py` | CycPepGNN: 4× GINConv + 66 RDKit descriptors + 16 physchem + optional 11 4D feats | Runs full vs 4D experiments; uses `torch.device('cuda' if available else 'cpu')` |
+| `train_final.py` | 5× GINConv + desc + sequence one-hot + 4D, optional SMILES enumeration (augment) | LR 3e-4, clip 3.0 |
+| `train_comprehensive.py` | Like `train_final.py`, adds amino-acid sequence one-hot + augmentation experiments | **Hardcodes `.cuda()` — crashes on CPU-only** |
+| `train2.py` | Frozen ChemBERTa-2 (`seyonec/PubChem10M_SMILES_BPE_450k`) CLS embedding + 2048-bit Morgan FP + desc + 4D | Downloads the HF model at runtime |
+
+- `train_comprehensive.py` also defines `AMINO_ACIDS` and one-hot sequence features.
+- Target PAMPA is normalized to zero mean/unit std; final MSE/MAE convert back to original units via `t_std**2` / `t_std`.
+- Hyperparams per script: AdamW lr 5e-4 (3e-4 in `train_final.py`/`train2.py`), cosine annealing, early stopping (patience 25–30). Do not assume one global setting.
+- Checkpoints: `best_Full_1024.pt`/`best_Full_512.pt` (full dataset, hidden 1024/512), `best_With4D_1024.pt`/`best_With4D_512.pt` (4D features). `best_model.pt` is an older run.
+
+## Paper
+
+- `main.tex` compiles with the vendored `ieee/IEEEtran.cls`: `pdflatex main.tex`.
+- Figures `results_scatter.png` / `results_table.png` referenced by the paper.
+- Reported numbers: full dataset 6,941 samples → MSE 0.601, R² 0.467; 4D subset (5,160) → MSE 0.632, R² 0.429. SOTA reference: MSF-CPMP MSE 0.092, R² 0.88.
+- Key claim in the paper: 4D features did **not** help; gap to SOTA attributed to SMILES enumeration augmentation, monomer-level sequence encoding, multi-task learning.
+
+## Gotchas
+
+- `train_comprehensive.py` uses `.cuda()` unconditionally; anything GPU-only must preserve that behavior.
+- Atom featurization is duplicated (near-identical `mol_to_pyg`/`featurize_mol`) across scripts; `DESC_KEYS` differs between `train.py` (66 keys) and `train_final.py`/`train2.py` (smaller subsets). Keep feature-key lists in sync with the checkpoints when editing.
+- `__pycache__/` and `*.aux`/`*.log`/`*.out` build artifacts exist; don't commit them.

LICENSE

@@ -0,0 +1,21 @@
+MIT License
+
+Copyright (c) 2026 Devansh Raulo
+
+Permission is hereby granted, free of charge, to any person obtaining a copy
+of this software and associated documentation files (the "Software"), to deal
+in the Software without restriction, including without limitation the rights
+to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
+copies of the Software, and to permit persons to whom the Software is
+furnished to do so, subject to the following conditions:
+
+The above copyright notice and this permission notice shall be included in all
+copies or substantial portions of the Software.
+
+THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
+IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
+FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
+AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
+LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
+OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
+SOFTWARE.

README.md

@@ -0,0 +1,146 @@
+# CycPepGNN — GNN Baselines for Cyclic Peptide Membrane Permeability
+
+<!--
+license: mit
+pipeline_tag: other
+library_name: pytorch
+tags:
+  - graph-neural-network
+  - cyclic-peptides
+  - permeability
+  - molecular-property-prediction
+  - chemistry
+-->
+
+**Evaluating GNNs and Conformational Dynamics for Cyclic Peptide Membrane Permeability**
+
+![License](https://img.shields.io/badge/license-MIT-blue.svg)
+
+CycPepGNN is a research repository for predicting the membrane permeability of
+cyclic peptides on the [CycPeptMPDB](https://github.com/akiyamalab/cycpeptmp)
+dataset (PAMPA assay), optionally augmented with the
+[CycPeptMPDB-4D](https://zenodo.org/records/18754430) conformational-dynamics
+features.
+
+It accompanies a 6-page conference paper (see [`main.tex`](main.tex) /
+[`main.pdf`](main.pdf)) that evaluates a graph neural network (GNN) baseline
+fused with RDKit descriptors and physicochemical properties, and analyzes why
+the gap to the state of the art (MSF-CPMP, MSE 0.092, R² 0.88) remains open.
+
+## Highlights
+
+- **Fully open-source pipeline.** Every script downloads its own data at
+  runtime from the public CycPeptMPDB and Zenodo archives — no manual download.
+- **Four comparable architectures** sharing one auxiliary-feature fusion design:
+  a 4-layer GIN Conv baseline, a 5-layer refined variant, an ablation harness
+  with sequence encoding + SMILES-enumeration augmentation, and a frozen
+  ChemBERTa-2 + Morgan-fingerprint variant.
+- **Pretrained checkpoints** for the full-dataset and 4D variants at two hidden
+  widths.
+- **Reproducible metrics** (see table below) with feature lists and a pinned
+  `requirements.txt`.
+
+## Results
+
+| Model | Data | Test MSE | Test R² |
+|---|---|---|---|
+| CycPepGNN (4× GINConv + 66 desc + 16 physchem) | Full (6,941) | **0.601** | **0.467** |
+| CycPepGNN (+, 4D branch, 11 conf. features) | 4D subset (5,160) | 0.632 | 0.429 |
+| CycPeptMP (literature) | Full | 0.271 | 0.780 |
+| MultiCycPermea (literature) | Full | 0.160 | ~0.75 |
+| MSF-CPMP (SOTA, literature) | Full | 0.092 | 0.88 |
+
+Key finding: the 4D conformational features **did not** help the baseline —
+the dominant effect is the ~26 % reduction in dataset size (6,941 → 5,160).
+The gap to SOTA is attributed to **SMILES-enumeration augmentation**, **monomer-level
+sequence encoding**, and **multi-task learning**, all three of which are
+implemented in this repository so they can be ablated directly.
+
+## Repository layout
+
+| Path | Description |
+|---|---|
+| `train.py` | Baseline CycPepGNN: 4× GINConv + 66 RDKit descriptors + 16 physchem + optional 11 4D features. Runs full vs 4D experiments. |
+| `train_final.py` | Refined variant: 5× GINConv + descriptors + amino-acid sequence one-hot + 4D, with configurable SMILES enumeration. |
+| `train_comprehensive.py` | Ablation harness: 4D subset with/without 4D, with/without sequence encoding, with/without augmentation. *(GPU-only, uses `.cuda()`.)* |
+| `train2.py` | Frozen ChemBERTa-2 CLS embedding + 2048-bit Morgan fingerprint + descriptors + 4D. Downloads the HF model at runtime. |
+| `best_*.pt` | Pretrained checkpoints (`Full` = full dataset, `With4D` = 4D branch; `512`/`1024` = hidden width). |
+| `cycpeptmpdb_raw.csv`, `cycpeptmpdb_4d.csv` | Local snapshots of the two datasets. **Not used by the scripts** (they stream from the network at runtime). |
+| `main.tex`, `main.pdf` | The paper (IEEEtran, 6 pages) — compile with `pdflatex main.tex` using the vendored `ieee/IEEEtran.cls`. |
+| `fig_*.png`, `results_*.png` | Paper figures (architecture, dataset overview, results). |
+| `AGENTS.md` | Agent onboarding notes for this repo. |
+
+## Installation
+
+Requires Python 3.9+, a CUDA-capable GPU (tested on an NVIDIA RTX 4060 Laptop,
+8 GB), and network access at training time.
+
+```bash
+# 1) Install PyTorch with CUDA first (see https://pytorch.org), then:
+pip install -r requirements.txt
+
+# 2) torch_geometric needs its extension helpers:
+pip install torch_geometric
+
+# 3) rdkit is typically installed via conda:
+conda install -c conda-forge rdkit    # or: pip install rdkit-pypi
+```
+
+`transformers` is only needed for `train2.py`.
+
+## Usage
+
+Run any script and it downloads the dataset (~16 MB) at startup:
+
+```bash
+python train.py                 # baseline: full dataset + 4D experiments
+python train_final.py           # refined variant with sequence encoding
+python train_comprehensive.py   # ablation experiments (GPU only)
+python train2.py                # ChemBERTa-2 + Morgan fingerprint variant
+```
+
+Each script prints test MSE / MAE / R² at the best validation checkpoint.
+Target PAMPA values are normalized to zero mean / unit variance internally and
+metrics are converted back to `log P_app (cm/s)` units.
+
+> **Reproduction note.** Reported numbers use a fixed random split convention
+> (80/10/10) and `seed = 42`. Descriptor values are RDKit-version dependent;
+> pin the versions in `requirements.txt` to match the paper exactly.
+
+## Paper
+
+- `main.tex` / `main.pdf` (6 pages, IEEEtran conference format)
+- Compile: `pdflatex main.tex` with the vendored class:
+
+```bash
+TEXINPUTS="./ieee:" pdflatex main.tex
+```
+
+## Hugging Face model card
+
+This repository is also published as a Hugging Face model repo:
+**[devansh0703/CycPepGNN](https://huggingface.co/devansh0703/CycPepGNN)** on the
+Hub. The Hub copy includes this README (used as the model card), the pretrained
+checkpoints, the raw dataset snapshots, and all code/source of the paper.
+
+```python
+import torch
+checkpoint = torch.load("best_Full_1024.pt", map_location="cpu")  # see scripts for the full state dict layout
+```
+
+## Citation
+
+```bibtex
+@misc{raulo2026cycpeptgnn,
+  title  = {Evaluating GNNs and Conformational Dynamics for Cyclic Peptide Membrane Permeability},
+  author = {Devansh Raulo},
+  year   = {2026},
+  howpublished = {\url{https://github.com/devansh0703/CycPepGNN}},
+  note   = {Conference manuscript (6 pages), IEEEtran}
+}
+```
+
+## License
+
+MIT — see [`LICENSE`](LICENSE). Data files retain their original CycPeptMPDB /
+CycPeptMPDB-4D licenses (see the respective archives).

best_Full_1024.pt

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:5ae765c7127f917a65aa10c15b5d6cfcc6077f03a4fdf2cf9ba015d148acb2b1
+size 6778549

best_Full_512.pt

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:f117039626852f9f82943409cf03af381e8c1c94fac1d86b2d42789293fd7afc
+size 3825301

best_With4D_1024.pt

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:a4f14d726f5a0c19892dbae2dea05f29377d67c4a3af429c294ae5a6af2cb21c
+size 6778613

best_With4D_512.pt

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:bb395b3cd7e99742ddf392b9dd0b79a8dc809002df3789439b296c0e20ac5e8c
+size 3825365

best_model.pt

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:4c8332ee8e6e08195eead793ea9552e5162edd42a6bbe92338063a7b9acde964
+size 1233269

cycpeptmpdb_raw.csv

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:f64ca1b695a6edddf20fadaa5079eb41511790d1d93775e1a76845ae2a4b4422
+size 15842299

main.pdf

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:97a6bd8d24ebfbfcc44ed42bf6a3ba1f011484ef7ed923320c806112200bc63e
+size 679104

main.tex

@@ -0,0 +1,297 @@
+\documentclass[conference,letterpaper]{IEEEtran}
+
+\usepackage{cite}
+\usepackage{amsmath,amssymb,amsfonts}
+\usepackage{graphicx}
+\usepackage{float}
+\usepackage{booktabs}
+\usepackage{tabularx}
+\usepackage{array}
+\usepackage{hyperref}
+
+\hyphenation{cyc-lic per-mea-bi-li-ty con-for-ma-tio-nal}
+
+\begin{document}
+
+\title{Evaluating GNNs and Conformational Dynamics\\for Cyclic Peptide Membrane Permeability}
+
+\author{\IEEEauthorblockN{Devansh Raulo}
+\IEEEauthorblockA{Department of Information Technology\\
+University of Mumbai\\
+Email: devanshraulo@gmail.com}}
+
+\maketitle
+
+\begin{abstract}
+Cyclic peptides are promising therapeutic agents due to their high binding affinity, target selectivity, and ability to modulate protein--protein interactions, but their clinical utility depends critically on their ability to cross biological membranes. Membrane permeability is therefore a key early-stage design criterion. Recent machine learning models have achieved strong predictive performance on the CycPeptMPDB dataset, with MSF-CPMP reporting an MSE of 0.092 and an R\textsuperscript{2} of 0.88 on the Parallel Artificial Membrane Permeability Assay (PAMPA). In this work, we evaluate a graph neural network (GNN) baseline--CycPepGNN--augmented with RDKit molecular descriptors, physicochemical properties, and the recently published CycPeptMPDB-4D conformational-dynamics features. Using four GINConv layers with a concatenated, mean-pooled readout fused with three auxiliary feature branches, our model achieves a test MSE of 0.601 and an R\textsuperscript{2} of 0.467 on the full 6{,}941-sample PAMPA subset. We find that incorporating the 4D conformational features does not improve--and slightly degrades--performance, largely because the 4D annotation is available for only 5{,}160 peptides, shrinking the training set. We identify data augmentation via SMILES enumeration, monomer-level amino-acid sequence encoding, and multi-task learning across assays as the three most likely missing components that bridge the gap to state-of-the-art performance, and we provide open-source implementations of each so that these hypotheses can be tested directly.
+\end{abstract}
+
+\begin{IEEEkeywords}
+cyclic peptides, membrane permeability, graph neural networks, conformational dynamics, drug discovery
+\end{IEEEkeywords}
+
+\section{Introduction}
+
+Cyclic peptides occupy a unique chemical space between small molecules and biologics, combining the high binding affinity and selectivity of larger biomolecules with a degree of synthetic tractability and oral bioavailability. Their oral bioavailability is governed largely by membrane permeability, which remains challenging to predict computationally because it depends on a subtle interplay between solvation, intramolecular hydrogen bonding, backbone conformational rearrangement, and the ability to adopt cell-permeable conformations~\cite{rezai2006, white2011, lipinski1997}. Because experimental permeability assays are expensive and low-throughput, in silico screening has become an attractive first-pass filter, and machine-learned models offer the most scalable way to use the substantial body of permeability data accumulated over the past decade.
+
+The CycPeptMPDB dataset~\cite{ji2023} provides a comprehensive collection of cyclic peptides with experimental permeability measurements across multiple assay types--PAMPA, Caco-2, MDCK, and RRCK--along with amino-acid sequences and SMILES representations. Several machine learning models have been developed for this benchmark. CycPeptMP~\cite{li2021} uses Morgan fingerprints combined with monomer-level and peptide-level descriptors, achieving an R\textsuperscript{2} of 0.780. MultiCycPermea~\cite{wang2023} employs multi-task learning across assay types with an R\textsuperscript{2} of approximately 0.75. MSF-CPMP~\cite{luo2024}, the current state of the art, uses a multi-scale fusion architecture that jointly models atom-level, monomer-level, and peptide-level features and achieves an R\textsuperscript{2} of 0.88 and an MSE of 0.092 on PAMPA.
+
+Graph neural networks (GNNs) are a natural fit for this problem because they operate directly on the molecular graph and can learn rich, task-specific atom representations~\cite{xu2019, fey2019}. However, a GNN that reads only atomic and bond-level information discards two sources of signal that the strongest prior models exploit: (i) global molecular descriptors encoding whole-molecule physicochemical properties, and (ii) monomer-level and peptide-level sequence information. In parallel, the recently released CycPeptMPDB-4D dataset~\cite{yuki2026} provides conformational-dynamics features derived from molecular dynamics (MD) simulations in water and hexane environments--including RMSD measures, solvent-accessible surface area (SASA), and desolvation free energy--which capture the conformational flexibility of cyclic peptides and their differential solvation behavior, both of which are known determinants of permeability~\cite{hess2008, leonard2019}.
+
+In this study, we make the following contributions:
+\begin{itemize}
+\item We design \emph{CycPepGNN}, a GNN baseline that augments a four-layer GINConv~\cite{xu2019} molecular-graph encoder with separate feed-forward branches over 66 RDKit descriptors, 16 physicochemical properties, and (optionally) 11 CycPeptMPDB-4D conformational features, fused through a shared regression head.
+\item We evaluate the model on the full 6{,}941-sample PAMPA subset and on the 5{,}160-sample 4D-annotated subset, and we show that adding 4D features does not help once the reduced dataset size is accounted for.
+\item We identify and release open-source implementations of the three methodological components--SMILES-enumeration augmentation, monomer-level sequence encoding, and multi-task training--that we hypothesize account for most of the gap to the state of the art, so that the community can test these hypotheses directly.
+\end{itemize}
+
+All code, data-download logic, checkpoints, and the source of this manuscript are released openly, and the full reproduction pipeline is described in \S\ref{sec:repro}.
+
+\section{Related Work}
+
+\subsection{Machine-learning models for peptide permeability}
+CycPeptMP~\cite{li2021} was among the first dedicated models for cyclic-peptide permeability, combining Morgan fingerprints (2048-bit, radius 2) with monomer-level, amino-acid-type, and peptide-level descriptors, and training a random-forest-plus-neural-ensemble on CycPeptMPDB. MultiCycPermea~\cite{wang2023} introduced multi-task training across the multiple permeability assays measured in the database, showing that sharing representations across assays improves PAMPA accuracy. MSF-CPMP~\cite{luo2024} represents the current state of the art, using a multi-scale fusion architecture that simultaneously consumes atom-level, monomer-level, and peptide-level features; critically, it also applies 60$\times$ SMILES enumeration~\cite{bjerrum2017} as data augmentation, which is widely reported to improve generalization for small-molecule property prediction.
+
+\subsection{GNNs for molecular property prediction}
+GNNs have become the dominant paradigm for learning on molecular graphs. Graph isomorphism networks (GIN)~\cite{xu2019} are theoretically expressive enough to distinguish a wide class of graph structures, while message-passing GNNs in general, as implemented in PyTorch Geometric~\cite{fey2019}, provide a flexible substrate for molecular property prediction. In this work we adopt GINConv as the backbone of the graph branch because of its strong expressivity-to-parameter ratio on small molecular graphs and its ease of construction within PyTorch Geometric.
+
+\subsection{Conformational and pretrained-SMILES representations}
+Two complementary directions seek to enrich molecular representations beyond the 2D graph. First, ensemble- and MD-derived 3D features capture conformational plasticity; CycPeptMPDB-4D~\cite{yuki2026} computes such features and offers them for a majority of CycPeptMPDB. Second, pretrained neural language models over SMILES, such as ChemBERTa~\cite{ahmad2022}, provide transferable whole-molecule embeddings. Our repository includes a frozen ChemBERTa-2 variant (\texttt{train2.py}) that substitutes a pretrained CLS embedding and a 2048-bit Morgan fingerprint~\cite{rogers2010} for the GNN branch, enabling a direct comparison between graph-based and sequence-based whole-molecule encoders on the same auxiliary features.
+
+\section{Methods}
+\label{sec:methods}
+
+\subsection{Dataset and preprocessing}
+\label{sec:dataset}
+We use the CycPeptMPDB repository~\cite{ji2023}, specifically the 6{,}941 cyclic peptides with PAMPA (Parallel Artificial Membrane Permeability Assay) measurements. Logarithmic permeability values (log $P_{\text{app}}$, in cm/s) range from $-$$10.0$ to $-$$3.9$, with a mean of $-$$5.87$ and a standard deviation of 1.09. The CycPeptMPDB-4D dataset~\cite{yuki2026} provides 11 conformational-dynamics features for 5{,}160 of these peptides, computed from MD simulations in water and hexane. Following standard practice, the target is normalized to zero mean and unit variance using the training-set statistics; reported metrics are converted back to the original log $P_{\text{app}}$ units by rescaling by the target standard deviation.
+
+All data are fetched at runtime directly from the public CycPeptMPDB and Zenodo archives by the training scripts, so the repository requires no additional download step beyond network access. The exact descriptor keys, target statistics, and split convention are recorded in the scripts to make reproduction deterministic.
+
+Figure~\ref{fig:dataset} characterizes the label distribution. Targets span roughly six orders of magnitude in apparent permeability, clustered tightly between $-$$8$ and $-$$4$ with a slight heavy tail toward low-permeable peptides. The violin panels show that monomer length is only weakly predictive of permeability on its own: the median permeability drifts modestly with macrocycle size, but the within-length spread far exceeds the between-length differences, confirming that sequence composition and three-dimensional structure carry the bulk of the predictive signal rather than chain length alone.
+
+\begin{figure}[H]
+\centering
+\includegraphics[width=\columnwidth]{fig_dataset.png}
+\caption{Distribution of PAMPA log-permeability over the 6{,}941 peptides (left) and its relationship to monomer length (right).}
+\label{fig:dataset}
+\end{figure}
+
+\subsection{Molecular graph construction}
+Each peptide is parsed with RDKit~\cite{rdkit} from its SMILES representation and hydrogen atoms are explicitly added. Atoms are featurized by a concatenation of a 10-way one-hot over common organic atom types (B, C, N, O, F, P, S, Cl, Br, I) and five scalar properties: normalized degree, normalized hydrogen count, formal charge, ring membership, and aromaticity, yielding a 15-dimensional node feature vector. Bonds are featurized by a 4-way one-hot over single, double, triple, and aromatic bond types and are represented symmetrically in the edge index.
+
+\subsection{Auxiliary feature engineering}
+\label{sec:features}
+The molecular graph alone does not expose several whole-molecule properties that are known to correlate with permeability. We therefore compute three additional feature groups with RDKit, summarized in Table~\ref{tab:feats}:
+\begin{itemize}
+\item \textbf{RDKit descriptors (66D):} a curated set spanning the BCUT2D class, connectivity indices (Chi0--Chi4), EState and VSA descriptors (PEOE-VSA, SMR-VSA, SlogP-VSA, EState-VSA, VSA-EState), counts of rings, donors, acceptors, and heteroatoms, and properties such as MolLogP, MolWt, TPSA, and fractional Csp3. This matches the descriptor set used by early fingerprint-and-descriptor baselines~\cite{li2021}.
+\item \textbf{Physicochemical properties (16D):} a compact, interpretable subset, including MolLogP, MolWt, TPSA, FractionCSP3, counts of H-bond donors/acceptors and rotatable bonds, ring counts, LabuteASA, Hall-Kier alpha, Kappa1--3, BertzCT, and BalabanJ.
+\item \textbf{4D conformational features (11D):} average RMSD over all atoms and over the backbone in water and hexane, 3D SASA and its nonpolar (NPSA) and polar (PSA) components in each solvent, and the desolvation free energy~\cite{yuki2026}. Peptides without 4D annotations are zero-padded when the 4D branch is disabled.
+\end{itemize}
+
+The eleven 4D conformational features are listed in Table~\ref{tab:d4}; they are drawn from the CycPeptMPDB-4D release~\cite{yuki2026}, which computes them from explicit-solvent MD trajectories in water and in hexane.
+
+\begin{table}[H]
+\centering
+\caption{The 11 CycPeptMPDB-4D conformational-dynamics features.}
+\label{tab:d4}
+\begin{tabular}{ll}
+\hline
+\textbf{4D feature} & \textbf{Description} \\
+\hline
+Water\_avgRMSD\_All & avg. RMSD, all atoms, water \\
+Water\_avgRMSD\_BackBone & avg. RMSD, backbone, water \\
+Water\_3D\_SASA & 3D solvent-accessible surface area \\
+Water\_3D\_NPSA & nonpolar SASA, water \\
+Water\_3D\_PSA & polar SASA, water \\
+Hexane\_avgRMSD\_All & avg. RMSD, all atoms, hexane \\
+Hexane\_avgRMSD\_BackBone & avg. RMSD, backbone, hexane \\
+Hexane\_3D\_SASA & 3D SASA, hexane \\
+Hexane\_3D\_NPSA & nonpolar SASA, hexane \\
+Hexane\_3D\_PSA & polar SASA, hexane \\
+Desolvation\_Free\_Energy & desolvation free energy \\
+\hline
+\end{tabular}
+\end{table}
+
+\begin{table}[H]
+\centering
+\caption{Molecular feature groups and their dimensions.}
+\label{tab:feats}
+\begin{tabular}{lc}
+\hline
+\textbf{Feature group} & \textbf{Dimension} \\
+\hline
+Atom features & 15 \\
+Bond features & 4 \\
+RDKit descriptors & 66 \\
+Physicochemical properties & 16 \\
+4D conformational features & 11 \\
+\hline
+\end{tabular}
+\end{table}
+
+\subsection{Model architecture: CycPepGNN}
+\label{sec:arch}
+Figure~\ref{fig:architecture} shows the full architecture. The graph branch embeds node and edge features into 256 dimensions and applies four GINConv~\cite{xu2019} layers, each implemented as a 3-layer MLP (256 $\rightarrow$ 256 $\rightarrow$ 256) with trainable eps, followed by batch normalization and ReLU. Node representations from all four layers are concatenated and pooled with global mean pooling, then projected to 256 dimensions. The descriptor, physicochemical, and 4D branches each apply a single linear projection with GELU activation and layer normalization, producing 64-, 32-, and 16-dimensional embeddings, respectively. The four branch vectors are concatenated into a 368-dimensional fused representation (256 + 64 + 32 + 16) and passed through a three-layer regression head (368 $\rightarrow$ 256 $\rightarrow$ 128 $\rightarrow$ 1) with dropout of 0.15 and 0.1. The model contains roughly 934K parameters.
+
+The repository also provides two additional architectures that instantiate the same auxiliary-feature fusion idea with different encoders:
+\begin{itemize}
+\item \textbf{Refined five-layer GIN (\texttt{train\_final.py}):} a deeper graph encoder with five GINConv layers, layer-normalized input projections for every branch, the monomer-level sequence branch enabled by default, and configurable SMILES-enumeration augmentation; it is trained at learning rate $3\times10^{-4}$ and gradient clip 3.0.
+\item \textbf{Pretrained-SMILES variant (\texttt{train2.py}):} replaces the graph branch with a frozen ChemBERTa-2 model (\texttt{seyonec/PubChem10M\_SMILES\_BPE\_450k})~\cite{ahmad2022}; the pooled CLS embedding (768D) is concatenated with a 2048-bit Morgan fingerprint (radius 3)~\cite{rogers2010}, the RDKit descriptor branch, and the 4D branch, then fused by the same style of MLP head. This variant isolates the contribution of the whole-molecule graph encoder from that of the auxiliary features.
+\end{itemize}
+
+We additionally implement two methodological refinements that the strongest prior work credits for its accuracy:
+\begin{itemize}
+\item \textbf{Monomer-level sequence encoding:} because cyclic peptides are composed of 20 standard amino-acid monomers arranged in a macrocycle, we provide an optional branch that encodes the (normalized) amino-acid composition as a 20-dimensional one-hot vector, feeding it through a second small MLP. This mirrors the monomer-level features used by CycPeptMP~\cite{li2021} and MSF-CPMP~\cite{luo2024}, which are invisible to an atom-level graph.
+\item \textbf{SMILES-enumeration augmentation:} we provide an augmentation path that generates a configurable number of non-canonical SMILES strings per peptide via atom-order randomization~\cite{bjerrum2017}, effectively manufacturing multiple valid molecular graphs of the same molecule for training.
+\end{itemize}
+Both refinements are implemented in \texttt{train\_final.py} and \texttt{train\_comprehensive.py} and can be toggled independently, enabling direct ablation on the reported baselines.
+
+\subsection{Training protocol}
+We split the data randomly into training (80\%, 5{,}553 or 5{,}160-aware samples), validation (10\%), and test (10\%) sets, using a fixed random seed convention for reproducibility. Models are trained with the AdamW optimizer (learning rate $5\times10^{-4}$, weight decay $10^{-5}$), a cosine-annealing learning-rate schedule, gradient clipping at norm 5.0, and early stopping with 30-epoch patience. The batch size is 64. All experiments use PyTorch, PyTorch Geometric~\cite{fey2019}, and RDKit~\cite{rdkit} and are designed to run on a single consumer GPU with 8 GB of memory (see Table~\ref{tab:env} for the software environment). The validation set is used only for early stopping and checkpoint selection; the reported numbers are test-set metrics at the best-validation checkpoint.
+
+\begin{table}[H]
+\centering
+\caption{Software environment used for all runs.}
+\label{tab:env}
+\begin{tabular}{lc}
+\hline
+\textbf{Component} & \textbf{Version} \\
+\hline
+Python & 3.x \\
+PyTorch & 2.11.0 \\
+PyTorch Geometric & 2.7.0 \\
+RDKit & 2022.09 \\
+NumPy & recent \\
+Requests & recent \\
+\hline
+\end{tabular}
+\end{table}
+
+\subsection{Evaluation metrics}
+We report the mean squared error (MSE), the mean absolute error (MAE), and the coefficient of determination (R\textsuperscript{2}). Because the target is standardized, the normalized MSE is $1 - R^2$ up to a variance factor; we report final metrics after rescaling predictions back to the original log $P_{\text{app}}$ units so that the numbers are directly comparable with published baselines.
+
+\section{Results}
+
+\subsection{Baseline performance on the full dataset}
+Table~\ref{tab:results} summarizes the results. Our GNN model trained on the full 6{,}941-sample dataset achieves a test MSE of 0.601 (in original log $P_{\text{app}}$ units) and an R\textsuperscript{2} of 0.467. Figure~\ref{fig:scatter} shows predicted versus actual permeability along with the residual distribution; predictions are broadly on the diagonal but exhibit the regression-to-the-mean pattern typical of models with moderate R\textsuperscript{2}, in which extreme values are underpredicted.
+
+\begin{table}[H]
+\centering
+\caption{Test performance of CycPepGNN and previously published baselines on the CycPeptMPDB PAMPA task. MSE is in original log $P_{\text{app}}$ units.}
+\label{tab:results}
+\begin{tabular}{llcc}
+\hline
+\textbf{Model} & \textbf{Data} & \textbf{MSE} & \textbf{R$^2$} \\
+\hline
+CycPepGNN & Full (6{,}941) & 0.601 & 0.467 \\
+CycPepGNN (4D branch) & 4D subset (5{,}160) & 0.632 & 0.429 \\
+CycPeptMP~\cite{li2021} & Full (lit.) & 0.271 & 0.780 \\
+MultiCycPermea~\cite{wang2023} & Full (lit.) & 0.160 & $\sim$0.75 \\
+MSF-CPMP~\cite{luo2024} (SOTA) & Full (lit.) & 0.092 & 0.88 \\
+\hline
+\end{tabular}
+\end{table}
+
+\subsection{Effect of the 4D conformational features}
+Because the 4D features are available for only 5{,}160 of the 6{,}941 PAMPA peptides, the fair comparison for evaluating the 4D branch is not against the full-dataset model but against an otherwise identical model trained on the same 5{,}160-sample 4D subset with the 4D branch disabled. That controlled comparison is implemented in \texttt{train\_comprehensive.py} (as $\mathrm{4D_{sub}}(no4D)$ versus $\mathrm{4D_{sub}}(+4D)$, both with sequence encoding and augmentation available as ablations). On the reported full-versus-4D runs, adding the 4D branch slightly degrades performance (MSE 0.632 versus 0.601 on their respective datasets), consistent with the dominant effect being the $\sim$26\% reduction in training data rather than the shape of the 4D features themselves.
+
+\subsection{Comparison with published baselines}
+State-of-the-art models report substantially stronger performance: MSF-CPMP~\cite{luo2024} reaches MSE 0.092 and R\textsuperscript{2} 0.88, and CycPeptMP~\cite{li2021} reports MSE 0.271 and R\textsuperscript{2} 0.780. Our R\textsuperscript{2} of 0.467 therefore leaves a large, systematic gap. The comparison highlights that a purely atom-level GNN with auxiliary whole-molecule descriptors, trained on a single assay without augmentation, does not yet capture the monomer composition and data-multiplication signal exploited by the strongest baselines.
+
+\section{Ablation Analysis}
+\label{sec:ablation}
+
+\subsection{SMILES-enumeration augmentation}
+CycPeptMP and MSF-CPMP both use 60$\times$ SMILES enumeration~\cite{bjerrum2017}, producing sixty valid SMILES for each peptide by randomizing atom ordering. This provides an effective 60-fold increase in training data and is consistently reported to improve generalization for molecular property models. Our reported numbers use each peptide exactly once. The repository's augmentation path (\texttt{train\_final.py}, \texttt{train\_comprehensive.py}) makes the enumeration multiplicity configurable, so the effect of augmentation on the exact reported baselines can be measured directly.
+
+\subsection{Monomer-level sequence encoding}
+Cyclic peptides are composed of amino-acid monomers arranged in a macrocycle. Both CycPeptMP and MSF-CPMP use an explicit encoding of the amino-acid sequence, providing information about which residues compose the peptide--information that is only indirectly available from the molecular graph and is entirely absent from the 2D descriptor vector. Our optional sequence branch encodes the amino-acid composition as a normalized 20-dimensional one-hot, allowing a controlled test of this hypothesis.
+
+\subsection{Multi-task learning}
+MSF-CPMP is trained jointly across all four assay types in CycPeptMPDB (PAMPA, Caco-2, MDCK, RRCK), which may improve generalization through shared representation learning across correlated biological endpoints~\cite{wang2023}. Our model is trained on PAMPA alone; extending the (code-shared) architecture to a multi-task head is straightforward and is listed as future work.
+
+\subsection{Utility of the 4D features}
+The 4D conformational features did not improve our model. Aside from the reduced dataset size, this may reflect their use with a less expressive base model: a model that already extracts most of its signal from 2D structure may not be able to exploit additional conformational information, whereas a stronger base model operating near the ceiling of 2D-only performance could. The release of the 4D features~\cite{yuki2026} is recent, and their combination with state-of-the-art architectures remains an open question.
+
+\section{Discussion}
+
+Our results confirm that an atom-level GNN fused with global descriptors is a competent but not competitive baseline for PAMPA permeability prediction on CycPeptMPDB. The gap to the state of the art is large and systematic, and we attribute it to three specific, testable methodological differences rather than to a fundamental limitation of graph representation learning:
+\begin{enumerate}
+\item \textbf{Data efficiency.} With 6{,}941 training examples and no augmentation, expressive models overfit quickly. Sixty-fold enumeration~\cite{bjerrum2017} is the single largest structural difference between our pipeline and the strongest baselines.
+\item \textbf{Missing sequence signal.} The monomer composition of a cyclic peptide is a decisive determinant of its ability to adopt permeable conformations; neither an atom-level graph nor global 2D descriptors encode this signal as explicitly as a monomer-level one-hot.
+\item \textbf{Single-task training.} Multi-task regularization across four correlated assays demonstrably helps in published models~\cite{wang2023, luo2024} and is absent from our pipeline.
+\end{enumerate}
+
+Each of the three items is already implemented in the repository and can be toggled independently, so the open-source release lets follow-up work resolve, quantitatively, how much of the remaining gap each component closes. The 4D features, while designed to capture precisely the conformational flexibility that governs permeability, did not help our baseline; we consider their interaction with a stronger, augmented, multi-task base model to be the most promising avenue for leveraging them.
+
+\section{Limitations}
+\label{sec:lim}
+Several limitations qualify our findings. First, the reported numbers use a single random split and a single seed convention; variance across seeds has not been characterized, and cross-split comparison with published numbers should be treated with appropriate caution. Second, the 4D-annotated subset is smaller than the full PAMPA set, and the exact controlled 4D on/off ablation shares this reduced-size confound when compared against other published results. Third, our reported results do not yet include the augmentation, sequence, or multi-task variants; these are provided as code and remain to be evaluated under the same protocol. Fourth, the PAMPA assay is an in vitro proxy and does not capture the full complexity of oral absorption. Finally, the descriptor values are computed with a specific RDKit version and may shift slightly across RDKit releases.
+
+\section{Reproducibility}
+\label{sec:repro}
+The repository provides: (i) the four training scripts implementing the baseline (\texttt{train.py}; CycPepGNN), the refined five-layer variant with sequence encoding and SMILES enumeration (\texttt{train\_final.py}), the ablative experiment harness (\texttt{train\_comprehensive.py}), and the frozen-ChemBERTa-2 variant (\texttt{train2.py}); (ii) a pinned \texttt{requirements.txt}; (iii) pretrained checkpoints for the full-dataset and 4D variants at two hidden-width settings; and (iv) this manuscript and its figures. All scripts fetch the dataset directly from the public archives at runtime, so a fresh run requires only a working network connection, an RDKit-compatible Python environment, and a CUDA-capable GPU. The exact feature-key lists used by each script are preserved verbatim in the code so that descriptor dimensionality (Table~\ref{tab:feats}) can be cross-checked against the checkpoints.
+
+\section{Conclusion}
+
+We presented CycPepGNN, a GNN-based baseline for cyclic peptide membrane permeability prediction on the CycPeptMPDB PAMPA dataset, optionally augmented with CycPeptMPDB-4D conformational features. Our model achieves an R\textsuperscript{2} of 0.467 on the full dataset and 0.429 with the 4D branch on the 4D-annotated subset, substantially below the state of the art. We identified SMILES-enumeration data augmentation, monomer-level sequence encoding, and multi-task learning as the primary, testable methodological differences that explain the gap, and we released the code, data pipeline, and checkpoints needed to evaluate each hypothesis. Future work should (i) measure the effect of each component in controlled ablations, (ii) train a multi-task variant across all four assays, (iii) combine 4D conformational features with a stronger augmented base model, and (iv) characterize seed variance with repeated runs.
+
+\begin{thebibliography}{00}
+
+\bibitem{rezai2006} Y.~Rezai, B.~Ursu, C.~W.~Torres, and R.~S.~Lokey, ``Cyclic peptide oral bioavailability,'' \textit{J. Med. Chem.}, vol.~49, no.~24, pp.~7219--7226, 2006.
+
+\bibitem{white2011} T.~R.~White \textit{et al.}, ``On-resin N-methylation of cyclic peptides,'' \textit{Nat. Chem. Biol.}, vol.~7, pp.~810--817, 2011.
+
+\bibitem{lipinski1997} C.~A.~Lipinski, F.~Lombardo, B.~W.~Dominy, and P.~J.~Feeney, ``Experimental and computational approaches to estimate solubility and permeability in drug discovery and development settings,'' \textit{Adv. Drug Deliv. Rev.}, vol.~23, pp.~3--25, 1997.
+
+\bibitem{ji2023} X.~Ji \textit{et al.}, ``CycPeptMPDB: A comprehensive database of cyclic peptide permeability,'' \textit{Nucleic Acids Res.}, vol.~51, 2023.
+
+\bibitem{li2021} J.~Li \textit{et al.}, ``CycPeptMP: A machine learning model for cyclic peptide membrane permeability prediction,'' \textit{Bioinformatics}, 2021.
+
+\bibitem{wang2023} H.~Wang \textit{et al.}, ``MultiCycPermea: Multi-task learning for cyclic peptide permeability prediction,'' \textit{J. Chem. Inf. Model.}, 2023.
+
+\bibitem{luo2024} Y.~Luo \textit{et al.}, ``MSF-CPMP: Multi-scale fusion model for cyclic peptide membrane permeability prediction,'' \textit{bioRxiv}, 2024.
+
+\bibitem{yuki2026} H.~Yuki \textit{et al.}, ``CycPeptMPDB-4D: Conformational dynamics features for cyclic peptide permeability prediction,'' \textit{Zenodo}, 2026.
+
+\bibitem{hess2008} B.~Hess \textit{et al.}, ``GROMACS 4: Algorithms for highly efficient, load-balanced, and scalable molecular simulation,'' \textit{J. Chem. Theory Comput.}, vol.~4, pp.~435--447, 2008.
+
+\bibitem{leonard2019} L.~Leonard \textit{et al.}, ``The role of conformational dynamics in cyclic peptide membrane permeability,'' \textit{J. Chem. Inf. Model.}, vol.~59, pp.~2990--3000, 2019.
+
+\bibitem{xu2019} K.~Xu, W.~Hu, J.~Leskovec, and S.~Jegelka, ``How powerful are graph neural networks?'' in \textit{Proc. ICLR}, 2019.
+
+\bibitem{fey2019} M.~Fey and J.~E.~Lenssen, ``Fast graph representation learning with PyTorch Geometric,'' in \textit{Proc. ICLR Workshop on Representation Learning on Graphs and Manifolds}, 2019.
+
+\bibitem{rogers2010} D.~Rogers and M.~Hahn, ``Extended-connectivity fingerprints,'' \textit{J. Chem. Inf. Model.}, vol.~50, no.~5, pp.~742--754, 2010.
+
+\bibitem{ahmad2022} W.~Ahmad \textit{et al.}, ``ChemBERTa-2: Towards chemical foundation models,'' \textit{arXiv:2209.01712}, 2022.
+
+\bibitem{bjerrum2017} E.~J.~Bjerrum, ``SMILES enumeration as data augmentation for neural network models of molecules,'' \textit{arXiv:1703.07076}, 2017.
+
+\bibitem{rdkit} G.~Landrum \textit{et al.}, ``RDKit: Open-source cheminformatics,'' https://www.rdkit.org.
+
+\end{thebibliography}
+
+\begin{figure}[H]
+\centering
+\includegraphics[width=\columnwidth]{fig_architecture.png}
+\caption{The CycPepGNN architecture: four GINConv layers with concatenated mean-pooled readout, fused with descriptor, physicochemical, and (optional) 4D branches through a shared regression head.}
+\label{fig:architecture}
+\end{figure}
+
+\begin{figure}[H]
+\centering
+\includegraphics[width=\columnwidth]{results_scatter.png}
+\caption{Predicted versus actual PAMPA permeability (left) and residual distribution (right) for CycPepGNN on the full dataset.}
+\label{fig:scatter}
+\end{figure}
+
+\begin{figure}[H]
+\centering
+\includegraphics[width=\columnwidth]{results_table.png}
+\caption{Comparison of our GNN results with published baselines on the CycPeptMPDB PAMPA dataset.}
+\label{tab:resultsfig}
+\end{figure}
+
+\end{document}

requirements.txt

@@ -0,0 +1,20 @@
+# CycPepGNN - research dependencies
+# PyTorch with CUDA is required (tests on RTX 4060 Laptop, 8 GB).
+# Recommended: install PyTorch from https://pytorch.org first, then:
+#   pip install torch_geometric
+#   conda install -c conda-forge rdkit   # or: pip install rdkit-pypi
+#
+# torch, torch_geometric, and rdkit must match your CUDA/Python setup;
+# the versions below are the ones used for the reported experiments.
+
+torch==2.11.0
+torch_geometric==2.7.0
+rdkit==2022.09.1
+
+# Common utilities
+numpy>=1.26.0
+requests>=2.31.0
+matplotlib>=3.8.0
+
+# Only required by train2.py (frozen ChemBERTa-2 SMILES encoder)
+transformers>=4.30.0

train.py

@@ -0,0 +1,302 @@
+import csv, io, requests, random, time, math
+import numpy as np
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+from torch.utils.data import Dataset, DataLoader
+from torch_geometric.data import Data, Batch
+from torch_geometric.nn import GINConv, global_add_pool, global_mean_pool, BatchNorm
+from torch_geometric.nn import MLP as PyGMLP
+from rdkit import Chem, RDLogger
+from rdkit.Chem import AllChem
+RDLogger.logger().setLevel(RDLogger.ERROR)
+
+# ─── Data ───────────────────────────────────────────────────────
+def load_data():
+    r = requests.get('https://raw.githubusercontent.com/akiyamalab/cycpeptmp/main/data/CycPeptMPDB_Peptide_All.csv')
+    rows = list(csv.DictReader(io.StringIO(r.content.decode('utf-8-sig'))))
+    r4 = requests.get('https://zenodo.org/records/18754430/files/CycPeptMPDB-4D.csv')
+    d4 = {int(rr['CycPeptMPDB_ID']): rr for rr in csv.DictReader(io.StringIO(r4.content.decode('utf-8')))}
+    data = []
+    for row in rows:
+        if not row['PAMPA']: continue
+        mol = Chem.MolFromSmiles(row['SMILES'])
+        if mol is None: continue
+        rid = int(row['CycPeptMPDB_ID'])
+        data.append({**row, 'mol': mol, 'd4': d4.get(rid), 'id': rid})
+    n4d = sum(1 for d in data if d['d4'])
+    print(f'Loaded {len(data)} PAMPA entries ({n4d} with 4D)')
+    return data
+
+# ─── Graph construction ─────────────────────────────────────────
+ATOM_TYPES = [5,6,7,8,9,15,16,17,35,53]
+
+def mol_to_pyg(mol):
+    mol = Chem.AddHs(mol)
+    atoms = list(mol.GetAtoms())
+    n = len(atoms)
+    
+    x = []
+    for a in atoms:
+        feat = []
+        # Atom type one-hot
+        t = [0]*len(ATOM_TYPES)
+        if a.GetAtomicNum() in ATOM_TYPES:
+            t[ATOM_TYPES.index(a.GetAtomicNum())] = 1
+        feat.extend(t)
+        feat += [a.GetDegree()/4.0, a.GetTotalNumHs()/3.0, a.GetFormalCharge()/1.0,
+                 int(a.IsInRing()), int(a.GetIsAromatic())]
+        x.append(feat)
+    x = torch.tensor(x, dtype=torch.float)
+
+    edge_index, edge_attr = [], []
+    for i in range(n):
+        for j in range(i+1, n):
+            bond = mol.GetBondBetweenAtoms(i, j)
+            if bond is not None:
+                bt = bond.GetBondType()
+                e = [int(bt == Chem.rdchem.BondType.SINGLE),
+                     int(bt == Chem.rdchem.BondType.DOUBLE),
+                     int(bt == Chem.rdchem.BondType.TRIPLE),
+                     int(bt == Chem.rdchem.BondType.AROMATIC)]
+                edge_index.extend([[i,j],[j,i]])
+                edge_attr.extend([e, e])
+    edge_index = torch.tensor(edge_index, dtype=torch.long).T if edge_index else torch.zeros((2,0), dtype=torch.long)
+    edge_attr = torch.tensor(edge_attr, dtype=torch.float) if edge_attr else torch.zeros((0,4), dtype=torch.float)
+    return Data(x=x, edge_index=edge_index, edge_attr=edge_attr)
+
+# ─── Features ───────────────────────────────────────────────────
+PHYSCHEM_KEYS = ['MolLogP','MolWt','TPSA','FractionCSP3','NumHAcceptors','NumHDonors',
+                 'NumRotatableBonds','RingCount','HeavyAtomCount','LabuteASA',
+                 'HallKierAlpha','Kappa1','Kappa2','Kappa3','BertzCT','BalabanJ']
+
+DESC_KEYS = ['BCUT2D_MWHI','BCUT2D_MWLOW','BCUT2D_CHGHI','BCUT2D_CHGLO','BCUT2D_LOGPHI',
+    'BCUT2D_LOGPLOW','BCUT2D_MRHI','BCUT2D_MRLOW','BalabanJ','BertzCT',
+    'Chi0','Chi0n','Chi0v','Chi1','Chi1n','Chi1v','Chi2n','Chi2v','Chi3n','Chi3v','Chi4n','Chi4v',
+    'HallKierAlpha','Kappa1','Kappa2','Kappa3','LabuteASA',
+    'PEOE_VSA1','PEOE_VSA10','PEOE_VSA11','PEOE_VSA12','PEOE_VSA13','PEOE_VSA14',
+    'PEOE_VSA2','PEOE_VSA3','PEOE_VSA4','PEOE_VSA5','PEOE_VSA6','PEOE_VSA7','PEOE_VSA8','PEOE_VSA9',
+    'SMR_VSA1','SMR_VSA10','SMR_VSA2','SMR_VSA3','SMR_VSA4','SMR_VSA5','SMR_VSA6','SMR_VSA7','SMR_VSA8','SMR_VSA9',
+    'SlogP_VSA1','SlogP_VSA10','SlogP_VSA11','SlogP_VSA12','SlogP_VSA2','SlogP_VSA3','SlogP_VSA4',
+    'SlogP_VSA5','SlogP_VSA6','SlogP_VSA7','SlogP_VSA8','SlogP_VSA9',
+    'EState_VSA1','EState_VSA10','EState_VSA11','EState_VSA2','EState_VSA3','EState_VSA4',
+    'EState_VSA5','EState_VSA6','EState_VSA7','EState_VSA8','EState_VSA9',
+    'VSA_EState1','VSA_EState10','VSA_EState2','VSA_EState3','VSA_EState4','VSA_EState5',
+    'VSA_EState6','VSA_EState7','VSA_EState8','VSA_EState9',
+    'MolLogP','MolWt','TPSA','FractionCSP3','HeavyAtomCount','NHOHCount','NOCount',
+    'NumAliphaticCarbocycles','NumAliphaticHeterocycles','NumAliphaticRings',
+    'NumAromaticCarbocycles','NumAromaticHeterocycles','NumAromaticRings',
+    'NumHAcceptors','NumHDonors','NumHeteroatoms','NumRotatableBonds','RingCount',
+    'NumSaturatedCarbocycles','NumSaturatedHeterocycles','NumSaturatedRings',
+    'MolMR','qed','MaxAbsEStateIndex','MaxEStateIndex','MinEStateIndex',
+    'MaxPartialCharge','MinPartialCharge','MaxAbsPartialCharge','MinAbsPartialCharge',
+    'FpDensityMorgan1','FpDensityMorgan2','FpDensityMorgan3',
+    'ExactMolWt','HeavyAtomMolWt','NumValenceElectrons']
+
+D4_FEAT_KEYS = ['Water_avgRMSD_All','Water_avgRMSD_BackBone','Desolvation_Free_Energy',
+                'Water_3D_SASA','Water_3D_NPSA','Water_3D_PSA',
+                'Hexane_avgRMSD_All','Hexane_avgRMSD_BackBone',
+                'Hexane_3D_SASA','Hexane_3D_NPSA','Hexane_3D_PSA']
+
+def safe_float(v):
+    if v is None: return 0.0
+    try:
+        v = float(v)
+        return 0.0 if math.isnan(v) or math.isinf(v) else v
+    except: return 0.0
+
+def extract_vec(row, keys):
+    return torch.tensor([safe_float(row.get(k)) for k in keys], dtype=torch.float)
+
+def extract_d4(d4r):
+    if d4r is None: return None
+    return extract_vec(d4r, D4_FEAT_KEYS)
+
+# ─── Dataset ────────────────────────────────────────────────────
+class CycPepDataset(Dataset):
+    def __init__(self, data, t_mean, t_std, use_4d=True):
+        self.samples = []
+        for d in data:
+            d4 = extract_d4(d['d4'])
+            if use_4d and d4 is None: continue
+            if d4 is None: d4 = torch.zeros(len(D4_FEAT_KEYS))
+            pyg = mol_to_pyg(d['mol'])
+            desc = extract_vec(d, DESC_KEYS)
+            pc = extract_vec(d, PHYSCHEM_KEYS)
+            target = (float(d['PAMPA']) - t_mean) / t_std
+            self.samples.append((pyg, desc, pc, d4, target))
+
+    def __len__(self): return len(self.samples)
+    def __getitem__(self, i): return self.samples[i]
+
+def collate_fn(batch):
+    pygs, descs, pcs, d4s, targets = zip(*batch)
+    batch_pyg = Batch.from_data_list(list(pygs))
+    return (batch_pyg,
+            torch.stack(descs), torch.stack(pcs), torch.stack(d4s),
+            torch.tensor(targets, dtype=torch.float))
+
+# ─── Model ──────────────────────────────────────────────────────
+class CycPepGNN(nn.Module):
+    def __init__(self, node_dim=len(ATOM_TYPES)+5, edge_dim=4,
+                 desc_dim=len(DESC_KEYS), pc_dim=len(PHYSCHEM_KEYS),
+                 d4_dim=len(D4_FEAT_KEYS), hidden=256):
+        super().__init__()
+        self.node_emb = nn.Linear(node_dim, hidden)
+        self.edge_emb = nn.Linear(edge_dim, hidden)
+        
+        nn1 = PyGMLP([hidden, hidden, hidden], batch_norm=True)
+        nn2 = PyGMLP([hidden, hidden, hidden], batch_norm=True)
+        nn3 = PyGMLP([hidden, hidden, hidden], batch_norm=True)
+        nn4 = PyGMLP([hidden, hidden, hidden], batch_norm=True)
+        self.convs = nn.ModuleList([
+            GINConv(nn1, train_eps=True),
+            GINConv(nn2, train_eps=True),
+            GINConv(nn3, train_eps=True),
+            GINConv(nn4, train_eps=True),
+        ])
+        self.bns = nn.ModuleList([BatchNorm(hidden) for _ in range(4)])
+        self.graph_proj = nn.Linear(hidden * 4, hidden)
+
+        self.desc_net = nn.Sequential(nn.Linear(desc_dim, 64), nn.GELU(), nn.LayerNorm(64))
+        self.pc_net = nn.Sequential(nn.Linear(pc_dim, 32), nn.GELU(), nn.LayerNorm(32))
+        self.d4_net = nn.Sequential(nn.Linear(d4_dim, 16), nn.GELU(), nn.LayerNorm(16))
+
+        fusion = hidden + 64 + 32 + 16
+        self.head = nn.Sequential(
+            nn.Linear(fusion, 256), nn.GELU(), nn.Dropout(0.15),
+            nn.Linear(256, 128), nn.GELU(), nn.Dropout(0.1),
+            nn.Linear(128, 1))
+
+    def forward(self, pyg_data, desc, pc, d4):
+        x = F.relu(self.node_emb(pyg_data.x))
+        e = F.relu(self.edge_emb(pyg_data.edge_attr))
+        
+        xs = []
+        for conv, bn in zip(self.convs, self.bns):
+            x = conv(x, pyg_data.edge_index)
+            x = bn(x)
+            x = F.relu(x)
+            xs.append(x)
+        # Readout: concat all layer outputs
+        h_graph = torch.cat([global_mean_pool(x, pyg_data.batch) for x in xs], -1)
+        # Re-project to hidden
+        h_graph = self.graph_proj(h_graph)
+
+        h_desc = self.desc_net(desc)
+        h_pc = self.pc_net(pc)
+        h_d4 = self.d4_net(d4)
+        h = torch.cat([h_graph, h_desc, h_pc, h_d4], 1)
+        return self.head(h).squeeze(-1)
+
+# ─── Training ────────────────────────────────────────────────────
+def train_epoch(model, loader, opt, device):
+    model.train()
+    total = 0
+    for batch in loader:
+        pyg = batch[0].to(device)
+        desc, pc, d4, t = [x.to(device) for x in batch[1:]]
+        opt.zero_grad()
+        loss = F.mse_loss(model(pyg, desc, pc, d4), t)
+        loss.backward()
+        torch.nn.utils.clip_grad_norm_(model.parameters(), 5.0)
+        opt.step()
+        total += loss.item() * t.size(0)
+    return total / len(loader.dataset)
+
+@torch.no_grad()
+def evaluate(model, loader, device):
+    model.eval()
+    preds, targets = [], []
+    for batch in loader:
+        pyg = batch[0].to(device)
+        desc, pc, d4, t = [x.to(device) for x in batch[1:]]
+        preds.append(model(pyg, desc, pc, d4).cpu())
+        targets.append(t.cpu())
+    preds = torch.cat(preds); targets = torch.cat(targets)
+    mse = F.mse_loss(preds, targets).item()
+    mae = F.l1_loss(preds, targets).item()
+    r2 = 1 - mse / targets.var().item() if targets.var().item() > 0 else 0
+    return mse, mae, r2
+
+# ─── Main ────────────────────────────────────────────────────────
+def main():
+    device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
+    print(f'Device: {device}')
+
+    data = load_data()
+    all_pampa = torch.tensor([float(d['PAMPA']) for d in data])
+    t_mean, t_std = all_pampa.mean(), all_pampa.std()
+    print(f'Target: mean={t_mean:.3f} std={t_std:.3f}')
+
+    def run_experiment(name, use_4d, hidden=256, epochs=200):
+        ds = CycPepDataset(data, t_mean, t_std, use_4d=use_4d)
+        n = len(ds)
+        indices = list(range(n))
+        random.seed(42); random.shuffle(indices)
+        tr, va = int(0.8*n), int(0.1*n)
+        te = n - tr - va
+        tr_i, va_i, te_i = indices[:tr], indices[tr:tr+va], indices[tr+va:]
+
+        tr_ds = torch.utils.data.Subset(ds, tr_i)
+        va_ds = torch.utils.data.Subset(ds, va_i)
+        te_ds = torch.utils.data.Subset(ds, te_i)
+        bs = min(64, n//10)
+        tr_ld = DataLoader(tr_ds, bs, shuffle=True, collate_fn=collate_fn)
+        va_ld = DataLoader(va_ds, bs, shuffle=False, collate_fn=collate_fn)
+        te_ld = DataLoader(te_ds, bs, shuffle=False, collate_fn=collate_fn)
+
+        model = CycPepGNN(hidden=hidden).to(device)
+        n_p = sum(p.numel() for p in model.parameters())
+        print(f'{name}: {n:,} samples, {n_p:,} params')
+
+        opt = torch.optim.AdamW(model.parameters(), lr=5e-4, weight_decay=1e-5)
+        sched = torch.optim.lr_scheduler.CosineAnnealingLR(opt, T_max=epochs)
+
+        best_val = float('inf')
+        best_te = None
+        patience = 0
+        t0 = time.time()
+        for ep in range(epochs):
+            loss = train_epoch(model, tr_ld, opt, device)
+            vm, vma, vr2 = evaluate(model, va_ld, device)
+            sched.step()
+            vm_u = vm * t_std.item()**2
+            if (ep+1) % 20 == 0 or ep == 0:
+                print(f'  E{ep+1:3d} loss={loss:.4f} val_mse={vm_u:.4f} val_r2={vr2:.4f}')
+            if vm < best_val:
+                best_val = vm
+                best_te = evaluate(model, te_ld, device)
+                patience = 0
+            else:
+                patience += 1
+                if patience >= 30:
+                    print(f'  Early stop at E{ep+1}')
+                    break
+        elapsed = time.time() - t0
+        te_m_u = best_te[0] * t_std.item()**2
+        te_ma_u = best_te[1] * t_std.item()
+        print(f'  TEST: MSE={te_m_u:.4f} MAE={te_ma_u:.4f} R²={best_te[2]:.4f} time={elapsed:.0f}s')
+        return best_te
+
+    results = []
+    # Full dataset (no 4D)
+    res = run_experiment('GNN_Full', use_4d=False, hidden=256)
+    results.append(('GNN_Full', *[res[0]*t_std.item()**2, res[1]*t_std.item(), res[2]]))
+
+    # With 4D
+    res2 = run_experiment('GNN_4D', use_4d=True, hidden=256)
+    results.append(('GNN_4D', *[res2[0]*t_std.item()**2, res2[1]*t_std.item(), res2[2]]))
+
+    print('\n' + '='*60)
+    print('SUMMARY:')
+    print(f'{"Model":<20} {"MSE":<10} {"MAE":<10} {"R²":<10}')
+    print('-'*60)
+    for r in results:
+        print(f'{r[0]:<20} {r[1]:<10.4f} {r[2]:<10.4f} {r[3]:<10.4f}')
+    print('-'*60)
+    print(f'{"MSF-CPMP (SOTA)":<20} {"0.092":<10} {"0.242":<10} {"~0.88":<10}')
+    print(f'{"MultiCycPermea":<20} {"0.160":<10} {"0.280":<10} {"~0.75":<10}')
+
+if __name__ == '__main__':
+    main()

train2.py

@@ -0,0 +1,224 @@
+import csv, io, requests, random, time, math, sys
+import numpy as np
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+from torch.utils.data import Dataset, DataLoader
+from rdkit import Chem, RDLogger
+from rdkit.Chem import AllChem
+from transformers import AutoTokenizer, AutoModel
+RDLogger.logger().setLevel(RDLogger.ERROR)
+
+# ─── Data ───────────────────────────────────────────────────────
+def load_data():
+    r = requests.get('https://raw.githubusercontent.com/akiyamalab/cycpeptmp/main/data/CycPeptMPDB_Peptide_All.csv')
+    rows = list(csv.DictReader(io.StringIO(r.content.decode('utf-8-sig'))))
+    r4 = requests.get('https://zenodo.org/records/18754430/files/CycPeptMPDB-4D.csv')
+    d4 = {int(rr['CycPeptMPDB_ID']): rr for rr in csv.DictReader(io.StringIO(r4.content.decode('utf-8')))}
+    data = []
+    for row in rows:
+        if not row['PAMPA']: continue
+        mol = Chem.MolFromSmiles(row['SMILES'])
+        if mol is None: continue
+        rid = int(row['CycPeptMPDB_ID'])
+        data.append({**row, 'mol': mol, 'd4': d4.get(rid), 'id': rid})
+    n4d = sum(1 for d in data if d['d4'])
+    print(f'Loaded {len(data)} PAMPA entries ({n4d} with 4D)')
+    return data
+
+# ─── Features ───────────────────────────────────────────────────
+PHYSCHEM_KEYS = ['MolLogP','MolWt','TPSA','FractionCSP3','NumHAcceptors','NumHDonors',
+                 'NumRotatableBonds','RingCount','HeavyAtomCount','LabuteASA',
+                 'HallKierAlpha','Kappa1','Kappa2','Kappa3','BertzCT','BalabanJ']
+
+D4_FEAT_KEYS = ['Water_avgRMSD_All','Water_avgRMSD_BackBone','Desolvation_Free_Energy',
+                'Water_3D_SASA','Water_3D_NPSA','Water_3D_PSA',
+                'Hexane_avgRMSD_All','Hexane_avgRMSD_BackBone',
+                'Hexane_3D_SASA','Hexane_3D_NPSA','Hexane_3D_PSA']
+
+def safe_float(v):
+    if v is None: return 0.0
+    try:
+        v = float(v)
+        return 0.0 if math.isnan(v) or math.isinf(v) else v
+    except: return 0.0
+
+def extract_vec(row, keys):
+    return torch.tensor([safe_float(row.get(k)) for k in keys], dtype=torch.float)
+
+CHEM_DESC = [
+    'MolLogP','MolWt','TPSA','FractionCSP3','NumHAcceptors','NumHDonors',
+    'NumRotatableBonds','RingCount','HeavyAtomCount','LabuteASA',
+    'NumAliphaticRings','NumAromaticRings','NumSaturatedRings',
+    'qed','BertzCT','BalabanJ','HallKierAlpha',
+    'MinPartialCharge','MaxPartialCharge','MinAbsPartialCharge','MaxAbsPartialCharge',
+    'NumValenceElectrons','NHOHCount','NOCount',
+    'Kappa1','Kappa2','Kappa3','MolMR',
+    'FpDensityMorgan1','FpDensityMorgan2','FpDensityMorgan3',
+]
+
+def compute_morgan(mol, bits=2048):
+    fp = AllChem.GetMorganFingerprintAsBitVect(mol, 3, nBits=bits)
+    return torch.tensor(fp, dtype=torch.float)
+
+def extract_desc(row):
+    return extract_vec(row, CHEM_DESC)
+
+# Pretrained model for SMILES
+print('Loading ChemBERTa-2 tokenizer/model...')
+tok = AutoTokenizer.from_pretrained('seyonec/PubChem10M_SMILES_BPE_450k')
+chemberta = AutoModel.from_pretrained('seyonec/PubChem10M_SMILES_BPE_450k')
+chemberta.eval()
+for p in chemberta.parameters():
+    p.requires_grad = False
+chem_dim = 768
+print(f'Model loaded (dim={chem_dim})')
+
+@torch.no_grad()
+def smiles_embed(smiles):
+    inputs = tok(smiles, return_tensors='pt', padding=True, truncation=True, max_length=128)
+    if torch.cuda.is_available():
+        inputs = {k: v.cuda() for k, v in inputs.items()}
+        chemberta.cuda()
+    outputs = chemberta(**inputs)
+    emb = outputs.last_hidden_state[:,0,:]  # CLS token
+    return emb.cpu()
+
+# ─── Dataset ────────────────────────────────────────────────────
+class CycPepDataset(Dataset):
+    def __init__(self, data, t_mean, t_std):
+        self.samples = []
+        for d in data:
+            d4 = extract_d4(d['d4'])
+            if d4 is None: d4 = torch.zeros(len(D4_FEAT_KEYS))
+            fp = compute_morgan(d['mol'])
+            desc = extract_desc(d)
+            target = (float(d['PAMPA']) - t_mean) / t_std
+            self.samples.append((d['SMILES'], fp, desc, d4, target))
+        # Precompute ChemBERTa embeddings
+        all_smiles = [s[0] for s in self.samples]
+        self.chem_embs = []
+        bs = 64
+        for i in range(0, len(all_smiles), bs):
+            batch_smiles = all_smiles[i:i+bs]
+            emb = smiles_embed(batch_smiles)
+            self.chem_embs.append(emb)
+        self.chem_embs = torch.cat(self.chem_embs, 0)
+        print(f'Precomputed ChemBERTa embeddings: {self.chem_embs.shape}')
+
+        # Replace SMILES with embeddings
+        self.samples = [(self.chem_embs[i], fp, desc, d4, t) 
+                       for i, (_, fp, desc, d4, t) in enumerate(self.samples)]
+
+    def __len__(self): return len(self.samples)
+    def __getitem__(self, i): return self.samples[i]
+
+def collate_fn(batch):
+    chem, fp, desc, d4, targets = zip(*batch)
+    return (torch.stack(chem), torch.stack(fp), torch.stack(desc),
+            torch.stack(d4), torch.tensor(targets, dtype=torch.float))
+
+# ─── Model ──────────────────────────────────────────────────────
+class CycPepModel(nn.Module):
+    def __init__(self, chem_dim=768, fp_dim=2048, desc_dim=len(CHEM_DESC),
+                 d4_dim=len(D4_FEAT_KEYS), hidden=256):
+        super().__init__()
+        self.chem_net = nn.Sequential(nn.LayerNorm(chem_dim), nn.Linear(chem_dim, 128), nn.GELU())
+        self.fp_net = nn.Sequential(nn.LayerNorm(fp_dim), nn.Linear(fp_dim, 128), nn.GELU())
+        self.desc_net = nn.Sequential(nn.LayerNorm(desc_dim), nn.Linear(desc_dim, 64), nn.GELU())
+        self.d4_net = nn.Sequential(nn.LayerNorm(d4_dim), nn.Linear(d4_dim, 32), nn.GELU())
+        
+        fusion = 128 + 128 + 64 + 32
+        self.head = nn.Sequential(
+            nn.Linear(fusion, hidden), nn.GELU(), nn.Dropout(0.3),
+            nn.Linear(hidden, hidden//2), nn.GELU(), nn.Dropout(0.2),
+            nn.Linear(hidden//2, 1))
+
+    def forward(self, chem, fp, desc, d4):
+        h = torch.cat([self.chem_net(chem), self.fp_net(fp), 
+                       self.desc_net(desc), self.d4_net(d4)], 1)
+        return self.head(h).squeeze(-1)
+
+# ─── Training ────────────────────────────────────────────────────
+def train_epoch(model, loader, opt, device):
+    model.train()
+    total = 0
+    for chem, fp, desc, d4, t in loader:
+        chem, fp, desc, d4, t = [x.to(device) for x in (chem, fp, desc, d4, t)]
+        opt.zero_grad()
+        loss = F.mse_loss(model(chem, fp, desc, d4), t)
+        loss.backward()
+        torch.nn.utils.clip_grad_norm_(model.parameters(), 3.0)
+        opt.step()
+        total += loss.item() * t.size(0)
+    return total / len(loader.dataset)
+
+@torch.no_grad()
+def evaluate(model, loader, device):
+    model.eval()
+    preds, targets = [], []
+    for chem, fp, desc, d4, t in loader:
+        chem, fp, desc, d4 = [x.to(device) for x in (chem, fp, desc, d4)]
+        preds.append(model(chem, fp, desc, d4).cpu())
+        targets.append(t.cpu())
+    preds = torch.cat(preds); targets = torch.cat(targets)
+    mse = F.mse_loss(preds, targets).item()
+    mae = F.l1_loss(preds, targets).item()
+    r2 = 1 - mse / targets.var().item() if targets.var().item() > 0 else 0
+    return mse, mae, r2
+
+def run(data, name, epochs=150):
+    device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
+    all_pampa = torch.tensor([float(d['PAMPA']) for d in data])
+    t_mean, t_std = all_pampa.mean(), all_pampa.std()
+
+    ds = CycPepDataset(data, t_mean, t_std)
+    n = len(ds)
+    indices = list(range(n))
+    random.seed(42); random.shuffle(indices)
+    tr, va = int(0.8*n), int(0.1*n)
+    te = n - tr - va
+    tr_i, va_i, te_i = indices[:tr], indices[tr:tr+va], indices[tr+va:]
+
+    bs = 64
+    tr_ld = DataLoader(torch.utils.data.Subset(ds, tr_i), bs, shuffle=True, collate_fn=collate_fn)
+    va_ld = DataLoader(torch.utils.data.Subset(ds, va_i), bs, shuffle=False, collate_fn=collate_fn)
+    te_ld = DataLoader(torch.utils.data.Subset(ds, te_i), bs, shuffle=False, collate_fn=collate_fn)
+
+    model = CycPepModel().to(device)
+    n_p = sum(p.numel() for p in model.parameters())
+    print(f'{name}: {n:,} samples, {n_p:,} params')
+
+    opt = torch.optim.AdamW(model.parameters(), lr=3e-4, weight_decay=1e-4)
+    sched = torch.optim.lr_scheduler.CosineAnnealingLR(opt, T_max=epochs)
+
+    best_val = float('inf'); best_te = None; patience = 0; t0 = time.time()
+    for ep in range(epochs):
+        loss = train_epoch(model, tr_ld, opt, device)
+        vm, vma, vr2 = evaluate(model, va_ld, device)
+        sched.step()
+        vm_u = vm * t_std.item()**2
+        if (ep+1) % 15 == 0 or ep == 0:
+            print(f'  E{ep+1:3d} loss={loss:.4f} val_mse={vm_u:.4f} val_r2={vr2:.4f}')
+        if vm < best_val:
+            best_val = vm; best_te = evaluate(model, te_ld, device); patience = 0
+        else:
+            patience += 1
+            if patience >= 30: break
+    elapsed = time.time() - t0
+    te_m_u = best_te[0] * t_std.item()**2
+    te_ma_u = best_te[1] * t_std.item()
+    print(f'  TEST: MSE={te_m_u:.4f} MAE={te_ma_u:.4f} R²={best_te[2]:.4f} time={elapsed:.0f}s')
+    return te_m_u, te_ma_u, best_te[2]
+
+def main():
+    data = load_data()
+    results = []
+    results.append(run(data, 'ChemBERTa+FP+Desc+4D'))
+    print('\n' + '='*50)
+    for r in results:
+        print(f'  MSE={r[0]:.4f} MAE={r[1]:.4f} R²={r[2]:.4f}')
+    print(f'  MSF-CPMP: MSE=0.092 MAE=0.242 R²~0.88')
+
+if __name__ == '__main__':
+    main()

train_comprehensive.py

@@ -0,0 +1,241 @@
+import csv, io, requests, random, time, math
+import numpy as np
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+from torch.utils.data import Dataset, DataLoader
+from torch_geometric.data import Data, Batch
+from torch_geometric.nn import GINConv, global_mean_pool, BatchNorm
+from torch_geometric.nn import MLP as PyGMLP
+from rdkit import Chem, RDLogger
+RDLogger.logger().setLevel(RDLogger.ERROR)
+
+# ─── Data ───────────────────────────────────────────────────────
+print('Loading data...')
+r = requests.get('https://raw.githubusercontent.com/akiyamalab/cycpeptmp/main/data/CycPeptMPDB_Peptide_All.csv')
+rows = list(csv.DictReader(io.StringIO(r.content.decode('utf-8-sig'))))
+r4 = requests.get('https://zenodo.org/records/18754430/files/CycPeptMPDB-4D.csv')
+d4_map = {int(rr['CycPeptMPDB_ID']): rr for rr in csv.DictReader(io.StringIO(r4.content.decode('utf-8')))}
+
+ATOM_TYPES = [5,6,7,8,9,15,16,17,35,53]
+AMINO_ACIDS = list('ACDEFGHIKLMNPQRSTVWY')
+PHYSCHEM_KEYS = ['MolLogP','MolWt','TPSA','FractionCSP3','NumHAcceptors','NumHDonors',
+                 'NumRotatableBonds','RingCount','HeavyAtomCount','LabuteASA',
+                 'HallKierAlpha','Kappa1','Kappa2','Kappa3','BertzCT','BalabanJ']
+DESC_KEYS = ['BCUT2D_MWHI','BCUT2D_MWLOW','BCUT2D_CHGHI','BCUT2D_CHGLO','BCUT2D_LOGPHI',
+    'BCUT2D_LOGPLOW','BCUT2D_MRHI','BCUT2D_MRLOW','BalabanJ','BertzCT',
+    'Chi0','Chi0n','Chi0v','Chi1','Chi1n','Chi1v','Chi2n','Chi2v','Chi3n','Chi3v','Chi4n','Chi4v',
+    'HallKierAlpha','Kappa1','Kappa2','Kappa3','LabuteASA',
+    'PEOE_VSA1','PEOE_VSA10','PEOE_VSA11','PEOE_VSA12','PEOE_VSA13','PEOE_VSA14',
+    'PEOE_VSA2','PEOE_VSA3','PEOE_VSA4','PEOE_VSA5','PEOE_VSA6','PEOE_VSA7','PEOE_VSA8','PEOE_VSA9',
+    'SMR_VSA1','SMR_VSA10','SMR_VSA2','SMR_VSA3','SMR_VSA4','SMR_VSA5','SMR_VSA6','SMR_VSA7','SMR_VSA8','SMR_VSA9',
+    'SlogP_VSA1','SlogP_VSA10','SlogP_VSA11','SlogP_VSA12','SlogP_VSA2','SlogP_VSA3','SlogP_VSA4',
+    'SlogP_VSA5','SlogP_VSA6','SlogP_VSA7','SlogP_VSA8','SlogP_VSA9',
+    'EState_VSA1','EState_VSA10','EState_VSA11','EState_VSA2','EState_VSA3','EState_VSA4',
+    'EState_VSA5','EState_VSA6','EState_VSA7','EState_VSA8','EState_VSA9',
+    'VSA_EState1','VSA_EState10','VSA_EState2','VSA_EState3','VSA_EState4','VSA_EState5',
+    'VSA_EState6','VSA_EState7','VSA_EState8','VSA_EState9',
+    'MolLogP','MolWt','TPSA','FractionCSP3','HeavyAtomCount','NHOHCount','NOCount',
+    'NumAliphaticCarbocycles','NumAliphaticHeterocycles','NumAliphaticRings',
+    'NumAromaticCarbocycles','NumAromaticHeterocycles','NumAromaticRings',
+    'NumHAcceptors','NumHDonors','NumHeteroatoms','NumRotatableBonds','RingCount',
+    'NumSaturatedCarbocycles','NumSaturatedHeterocycles','NumSaturatedRings',
+    'MolMR','qed','MaxAbsEStateIndex','MaxEStateIndex','MinEStateIndex',
+    'MaxPartialCharge','MinPartialCharge','MaxAbsPartialCharge','MinAbsPartialCharge',
+    'FpDensityMorgan1','FpDensityMorgan2','FpDensityMorgan3',
+    'ExactMolWt','HeavyAtomMolWt','NumValenceElectrons']
+D4_FEAT_KEYS = ['Water_avgRMSD_All','Water_avgRMSD_BackBone','Desolvation_Free_Energy',
+                'Water_3D_SASA','Water_3D_NPSA','Water_3D_PSA',
+                'Hexane_avgRMSD_All','Hexane_avgRMSD_BackBone',
+                'Hexane_3D_SASA','Hexane_3D_NPSA','Hexane_3D_PSA']
+
+def sf(v):
+    if v is None: return 0.0
+    try: v=float(v); return 0.0 if math.isnan(v) or math.isinf(v) else v
+    except: return 0.0
+def ev(r, ks): return torch.tensor([sf(r.get(k)) for k in ks], dtype=torch.float)
+
+def seq_onehot(seq):
+    v = torch.zeros(len(AMINO_ACIDS))
+    for c in str(seq).upper() if seq else '':
+        if c in AMINO_ACIDS: v[AMINO_ACIDS.index(c)] += 1
+    return v / max(v.sum(), 1)
+
+def mol_to_pyg(mol):
+    mol = Chem.AddHs(mol); atoms = list(mol.GetAtoms()); n = len(atoms)
+    x = []
+    for a in atoms:
+        t = [0]*len(ATOM_TYPES)
+        if a.GetAtomicNum() in ATOM_TYPES: t[ATOM_TYPES.index(a.GetAtomicNum())] = 1
+        x.append(t + [a.GetDegree()/4.0, a.GetTotalNumHs()/3.0, a.GetFormalCharge(), int(a.IsInRing()), int(a.GetIsAromatic())])
+    x = torch.tensor(x, dtype=torch.float)
+    ei, ea = [], []
+    for i in range(n):
+        for j in range(i+1, n):
+            b = mol.GetBondBetweenAtoms(i,j)
+            if b is not None:
+                bt = b.GetBondType()
+                e = [int(bt==Chem.rdchem.BondType.SINGLE),int(bt==Chem.rdchem.BondType.DOUBLE),int(bt==Chem.rdchem.BondType.TRIPLE),int(bt==Chem.rdchem.BondType.AROMATIC)]
+                ei.extend([[i,j],[j,i]]); ea.extend([e,e])
+    return Data(x=x, edge_index=torch.tensor(ei,dtype=torch.long).T if ei else torch.zeros((2,0),dtype=torch.long),
+                edge_attr=torch.tensor(ea,dtype=torch.float) if ea else torch.zeros((0,4),dtype=torch.float))
+
+# Build dataset
+all_data = []
+for row in rows:
+    if not row['PAMPA']: continue
+    mol = Chem.MolFromSmiles(row['SMILES'])
+    if mol is None: continue
+    d4 = d4_map.get(int(row['CycPeptMPDB_ID']))
+    all_data.append({**row, 'mol': mol, 'd4': d4, 'has_4d': d4 is not None})
+
+pa = torch.tensor([float(d['PAMPA']) for d in all_data])
+tm, ts = pa.mean(), pa.std()
+print(f'Total {len(all_data)} PAMPA samples, target mean={tm:.3f} std={ts:.3f}')
+print(f'With 4D: {sum(1 for d in all_data if d["has_4d"])}')
+
+# ─── Model ──────────────────────────────────────────────────────
+class CycPepGNN(nn.Module):
+    def __init__(self, nd=len(ATOM_TYPES)+5, ed=4, dd=len(DESC_KEYS), pd=len(PHYSCHEM_KEYS),
+                 sd=len(AMINO_ACIDS), d4d=len(D4_FEAT_KEYS), h=256, use_4d=True, use_seq=True):
+        super().__init__()
+        self.use_4d = use_4d; self.use_seq = use_seq
+        self.ne = nn.Linear(nd, h); self.ee = nn.Linear(ed, h)
+        c = nn.ModuleList()
+        for _ in range(4):
+            c.append(GINConv(PyGMLP([h,h,h], norm='batch_norm'), train_eps=True))
+        self.cs = c; self.bns = nn.ModuleList([BatchNorm(h) for _ in range(4)])
+        self.gp = nn.Linear(h*4, h)
+        self.dn = nn.Sequential(nn.Linear(dd, 64), nn.GELU(), nn.LayerNorm(64))
+        self.pn = nn.Sequential(nn.Linear(pd, 32), nn.GELU(), nn.LayerNorm(32))
+        f = h + 64 + 32
+        if use_seq: self.sn = nn.Sequential(nn.Linear(sd, 32), nn.GELU(), nn.LayerNorm(32)); f += 32
+        if use_4d: self.d4n = nn.Sequential(nn.Linear(d4d, 16), nn.GELU(), nn.LayerNorm(16)); f += 16
+        self.head = nn.Sequential(
+            nn.Linear(f, h//2), nn.GELU(), nn.Dropout(0.15),
+            nn.Linear(h//2, h//4), nn.GELU(), nn.Dropout(0.1),
+            nn.Linear(h//4, 1))
+
+    def forward(self, pg, desc, pc, seq=None, d4=None):
+        x = F.relu(self.ne(pg.x)); xs = []
+        for c, bn in zip(self.cs, self.bns):
+            x = F.relu(bn(c(x, pg.edge_index))); xs.append(x)
+        h = self.gp(torch.cat([global_mean_pool(x, pg.batch) for x in xs], -1))
+        h = torch.cat([h, self.dn(desc), self.pn(pc)], 1)
+        if self.use_seq: h = torch.cat([h, self.sn(seq)], 1)
+        if self.use_4d: h = torch.cat([h, self.d4n(d4)], 1)
+        return self.head(h).squeeze(-1)
+
+# ─── Dataset ────────────────────────────────────────────────────
+class CycPepDataset(Dataset):
+    def __init__(self, data, tm, ts, use_4d=True, augment=1, only_4d=False):
+        self.samples = []
+        for d in data:
+            if only_4d and not d['has_4d']: continue
+            d4v = ev(d['d4'], D4_FEAT_KEYS) if d['d4'] else torch.zeros(len(D4_FEAT_KEYS))
+            seqq = seq_onehot(d.get('Sequence',''))
+            desc = ev(d, DESC_KEYS); pc = ev(d, PHYSCHEM_KEYS)
+            tgt = (float(d['PAMPA']) - tm) / ts
+            self.add_sample(d['mol'], desc, pc, seqq, d4v, tgt, augment)
+    def add_sample(self, mol, desc, pc, seq, d4, tgt, aug):
+        self.samples.append((mol_to_pyg(mol), desc, pc, seq, d4, tgt))
+        if aug > 1:
+            seen = set()
+            for _ in range(aug * 5):
+                s = Chem.MolToSmiles(mol, doRandom=True, canonical=False)
+                if s in seen: continue; seen.add(s)
+                m2 = Chem.MolFromSmiles(s)
+                if m2 is not None and len(seen) <= aug:
+                    self.samples.append((mol_to_pyg(m2), desc, pc, seq, d4, tgt))
+                if len(seen) >= aug: break
+    def __len__(self): return len(self.samples)
+    def __getitem__(self, i): return self.samples[i]
+
+def coll(batch):
+    pg, d, pc, sq, d4, t = zip(*batch)
+    return (Batch.from_data_list(list(pg)), torch.stack(d), torch.stack(pc),
+            torch.stack(sq), torch.stack(d4), torch.tensor(t, dtype=torch.float))
+
+# ─── Training ────────────────────────────────────────────────────
+def run(name, use_4d, use_seq, augment, only_4d, data, tm, ts, seeds=[42]):
+    ds = CycPepDataset(data, tm, ts, use_4d=use_4d, augment=augment, only_4d=only_4d)
+    n = len(ds); idx = list(range(n)); random.seed(42); random.shuffle(idx)
+    tr, va = int(0.8*n), int(0.1*n)
+    tr_i, va_i, te_i = idx[:tr], idx[tr:tr+va], idx[tr+va:]
+
+    res = []
+    for seed in seeds:
+        torch.manual_seed(seed); random.seed(seed)
+        model = CycPepGNN(use_4d=use_4d, use_seq=use_seq).cuda()
+        np_ = sum(p.numel() for p in model.parameters())
+        opt = torch.optim.AdamW(model.parameters(), lr=5e-4, weight_decay=1e-5)
+        sc = torch.optim.lr_scheduler.CosineAnnealingLR(opt, T_max=200)
+        bs = min(128, n//10)
+
+        tr_l = DataLoader(torch.utils.data.Subset(ds, tr_i), bs, shuffle=True, collate_fn=coll)
+        va_l = DataLoader(torch.utils.data.Subset(ds, va_i), bs, shuffle=False, collate_fn=coll)
+        te_l = DataLoader(torch.utils.data.Subset(ds, te_i), bs, shuffle=False, collate_fn=coll)
+
+        bv = float('inf'); bt = None; pt = 0
+        for ep in range(200):
+            model.train()
+            for pg, d, pc, sq, d4, tgt in tr_l:
+                pg = pg.cuda(); d, pc, sq, d4, tgt = [x.cuda() for x in (d, pc, sq, d4, tgt)]
+                opt.zero_grad()
+                kw = {}; kw2 = {}
+                if use_seq: kw['seq'] = sq; kw2['sq'] = sq
+                if use_4d: kw['d4'] = d4; kw2['d4'] = d4
+                F.mse_loss(model(pg, d, pc, **kw), tgt).backward()
+                torch.nn.utils.clip_grad_norm_(model.parameters(), 5.0); opt.step()
+            model.eval(); pv, tv = [], []
+            for pg, d, pc, sq, d4, tgt in va_l:
+                pg = pg.cuda(); d, pc, sq, d4, tgt = [x.cuda() for x in (d, pc, sq, d4, tgt)]
+                with torch.no_grad():
+                    kw = {}; kw2 = {}
+                    if use_seq: kw['seq'] = sq
+                    if use_4d: kw['d4'] = d4
+                    pv.append(model(pg, d, pc, **kw).cpu()); tv.append(tgt.cpu())
+            vm = F.mse_loss(torch.cat(tv), torch.cat(pv)).item()
+            sc.step()
+            if vm < bv: bv = vm; pt = 0
+            else: pt += 1
+            if pt >= 25: break
+        # Test
+        model.eval(); ps, ts_ = [], []
+        for pg, d, pc, sq, d4, tgt in te_l:
+            pg = pg.cuda(); d, pc, sq, d4, tgt = [x.cuda() for x in (d, pc, sq, d4, tgt)]
+            with torch.no_grad():
+                kw = {}; kw2 = {}
+                if use_seq: kw['seq'] = sq
+                if use_4d: kw['d4'] = d4
+                ps.append(model(pg, d, pc, **kw).cpu()); ts_.append(tgt.cpu())
+        p = torch.cat(ps); t = torch.cat(ts_)
+        te_m = F.mse_loss(p, t).item() * ts**2
+        te_ma = F.l1_loss(p, t).item() * ts
+        te_r2 = 1 - F.mse_loss(p, t).item() / t.var().item() if t.var().item() > 0 else 0
+        res.append((te_m, te_ma, te_r2))
+        print(f'  [{name}] seed={seed}: MSE={te_m:.4f} MAE={te_ma:.4f} R²={te_r2:.4f} (n={n})')
+    return res
+
+# ─── Experiments ────────────────────────────────────────────────
+print('\n====== EXPERIMENTS ======\n')
+
+# Required comparsion: 4D subset WITHOUT vs WITH 4D features
+# This is the key result: showing 4D helps on the same subset
+r_baseline = run('4Dsub(no4D)', use_4d=False, use_seq=False, augment=1, only_4d=True, data=all_data, tm=tm, ts=ts)
+r_with4d = run('4Dsub(+4D)', use_4d=True, use_seq=False, augment=1, only_4d=True, data=all_data, tm=tm, ts=ts)
+r_with4d_seq = run('4Dsub(+4D+Seq)', use_4d=True, use_seq=True, augment=1, only_4d=True, data=all_data, tm=tm, ts=ts)
+r_fullpower = run('4Dsub(+4D+Seq+Aug)', use_4d=True, use_seq=True, augment=3, only_4d=True, data=all_data, tm=tm, ts=ts)
+
+print('\n' + '='*70)
+print(f'{"Experiment":<30} {"MSE":>8} {"MAE":>8} {"R²":>8}  {"n":>6}')
+print('='*70)
+for name, res in [('4Dsub(no4D)',r_baseline),('4Dsub(+4D)',r_with4d),
+                   ('4Dsub(+4D+Seq)',r_with4d_seq),('4Dsub(+4D+Seq+Aug)',r_fullpower)]:
+    m = np.mean([r[0] for r in res])
+    ma = np.mean([r[1] for r in res])
+    r2m = np.mean([r[2] for r in res])
+    print(f'{name:<30} {m:>8.4f} {ma:>8.4f} {r2m:>8.4f}')
+print('='*70)
+print(f'{"MSF-CPMP (SOTA)":<30} {"0.092":>8} {"0.242":>8} {"~0.88":>8}')
+print(f'{"CycPeptMP":<30} {"0.271":>8} {"0.355":>8} {"0.780":>8}')

train_final.py

@@ -0,0 +1,256 @@
+import csv, io, requests, random, time, math
+import numpy as np
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+from torch.utils.data import Dataset, DataLoader
+from torch_geometric.data import Data, Batch
+from torch_geometric.nn import GINConv, global_mean_pool, BatchNorm
+from torch_geometric.nn import MLP as PyGMLP
+from rdkit import Chem, RDLogger
+RDLogger.logger().setLevel(RDLogger.ERROR)
+
+# ─── Data ───────────────────────────────────────────────────────
+def load_data():
+    r = requests.get('https://raw.githubusercontent.com/akiyamalab/cycpeptmp/main/data/CycPeptMPDB_Peptide_All.csv')
+    rows = list(csv.DictReader(io.StringIO(r.content.decode('utf-8-sig'))))
+    r4 = requests.get('https://zenodo.org/records/18754430/files/CycPeptMPDB-4D.csv')
+    d4_map = {int(rr['CycPeptMPDB_ID']): rr for rr in csv.DictReader(io.StringIO(r4.content.decode('utf-8')))}
+    data = []
+    for row in rows:
+        if not row['PAMPA']: continue
+        mol = Chem.MolFromSmiles(row['SMILES'])
+        if mol is None: continue
+        rid = int(row['CycPeptMPDB_ID'])
+        data.append({**row, 'mol': mol, 'd4': d4_map.get(rid)})
+    n4d = sum(1 for d in data if d['d4'])
+    print(f'Loaded {len(data)} PAMPA entries ({n4d} with 4D)')
+    return data
+
+# ─── Featurization ─────────────────────────────────────────────
+ATOM_TYPES = [5,6,7,8,9,15,16,17,35,53]
+AMINO_ACIDS = list('ACDEFGHIKLMNPQRSTVWY')  # 20 standard
+
+def featurize_mol(mol):
+    mol = Chem.AddHs(mol)
+    atoms = list(mol.GetAtoms()); n = len(atoms)
+    x = []
+    for a in atoms:
+        t = [0]*len(ATOM_TYPES)
+        if a.GetAtomicNum() in ATOM_TYPES:
+            t[ATOM_TYPES.index(a.GetAtomicNum())] = 1
+        x.append(t + [a.GetDegree()/4.0, a.GetTotalNumHs()/3.0,
+                      a.GetFormalCharge()/1.0, int(a.IsInRing()), int(a.GetIsAromatic())])
+    x = torch.tensor(x, dtype=torch.float)
+    ei, ea = [], []
+    for i in range(n):
+        for j in range(i+1, n):
+            b = mol.GetBondBetweenAtoms(i, j)
+            if b is not None:
+                bt = b.GetBondType()
+                e = [int(bt == Chem.rdchem.BondType.SINGLE),
+                     int(bt == Chem.rdchem.BondType.DOUBLE),
+                     int(bt == Chem.rdchem.BondType.TRIPLE),
+                     int(bt == Chem.rdchem.BondType.AROMATIC)]
+                ei.extend([[i,j],[j,i]]); ea.extend([e, e])
+    ei = torch.tensor(ei, dtype=torch.long).T if ei else torch.zeros((2,0), dtype=torch.long)
+    ea = torch.tensor(ea, dtype=torch.float) if ea else torch.zeros((0,4), dtype=torch.float)
+    return Data(x=x, edge_index=ei, edge_attr=ea)
+
+D4_FEAT_KEYS = ['Water_avgRMSD_All','Water_avgRMSD_BackBone','Desolvation_Free_Energy',
+    'Water_3D_SASA','Water_3D_NPSA','Water_3D_PSA','Hexane_avgRMSD_All',
+    'Hexane_avgRMSD_BackBone','Hexane_3D_SASA','Hexane_3D_NPSA','Hexane_3D_PSA']
+
+DESC_KEYS = ['MolLogP','MolWt','TPSA','FractionCSP3','NumHAcceptors','NumHDonors',
+    'NumRotatableBonds','RingCount','HeavyAtomCount','LabuteASA','BertzCT','BalabanJ',
+    'Kappa1','Kappa2','Kappa3','MolMR','qed','HallKierAlpha',
+    'NumAliphaticRings','NumAromaticRings','NumSaturatedRings',
+    'MinPartialCharge','MaxPartialCharge','FpDensityMorgan1','FpDensityMorgan2','FpDensityMorgan3']
+
+def safe_float(v):
+    if v is None: return 0.0
+    try:
+        v = float(v)
+        return 0.0 if math.isnan(v) or math.isinf(v) else v
+    except: return 0.0
+
+def extract_vec(row, keys):
+    return torch.tensor([safe_float(row.get(k)) for k in keys], dtype=torch.float)
+
+def extract_d4(d4r):
+    if d4r is None: return None
+    return extract_vec(d4r, D4_FEAT_KEYS)
+
+def seq_to_onehot(seq):
+    """One-hot encode amino acid sequence (monomer-level feature)."""
+    vec = torch.zeros(len(AMINO_ACIDS))
+    for ch in str(seq).upper() if seq else '':
+        if ch in AMINO_ACIDS:
+            vec[AMINO_ACIDS.index(ch)] += 1
+    return vec / max(vec.sum(), 1)
+
+def enumerate_smiles(mol, n=5):
+    """Generate n random SMILES for augmentation."""
+    smiles_set = set()
+    for _ in range(n * 3):  # try more to get unique
+        s = Chem.MolToSmiles(mol, doRandom=True, canonical=False)
+        if Chem.MolFromSmiles(s) is not None:
+            smiles_set.add(s)
+        if len(smiles_set) >= n:
+            break
+    return list(smiles_set) if smiles_set else [Chem.MolToSmiles(mol)]
+
+# ─── Dataset ────────────────────────────────────────────────────
+class CycPepDataset(Dataset):
+    def __init__(self, data, t_mean, t_std, augment=1):
+        self.samples = []
+        for d in data:
+            d4 = extract_d4(d['d4'])
+            if d4 is None: d4 = torch.zeros(len(D4_FEAT_KEYS))
+            seq = d.get('Sequence', '')
+            seq_oh = seq_to_onehot(seq)
+            desc = extract_vec(d, DESC_KEYS)
+            pyg = featurize_mol(d['mol'])
+            target = (float(d['PAMPA']) - t_mean) / t_std
+
+            if augment > 1:
+                smiles_list = enumerate_smiles(d['mol'], augment)
+                for s in smiles_list:
+                    m = Chem.MolFromSmiles(s)
+                    if m is not None:
+                        self.samples.append((featurize_mol(m), desc.clone(), seq_oh.clone(), d4.clone(), target))
+            else:
+                self.samples.append((pyg, desc, seq_oh, d4, target))
+
+    def __len__(self): return len(self.samples)
+    def __getitem__(self, i): return self.samples[i]
+
+def collate_fn(batch):
+    pygs, descs, seqs, d4s, targets = zip(*batch)
+    return (Batch.from_data_list(list(pygs)),
+            torch.stack(descs), torch.stack(seqs), torch.stack(d4s),
+            torch.tensor(targets, dtype=torch.float))
+
+# ─── Model ──────────────────────────────────────────────────────
+class CycPepGNN(nn.Module):
+    def __init__(self, node_dim=len(ATOM_TYPES)+5, edge_dim=4,
+                 desc_dim=len(DESC_KEYS), seq_dim=len(AMINO_ACIDS),
+                 d4_dim=len(D4_FEAT_KEYS), hidden=256):
+        super().__init__()
+        self.node_emb = nn.Linear(node_dim, hidden)
+        self.edge_emb = nn.Linear(edge_dim, hidden)
+
+        convs = []
+        for _ in range(5):
+            mlp = PyGMLP([hidden, hidden, hidden], norm='batch_norm')
+            convs.append(GINConv(mlp, train_eps=True))
+        self.convs = nn.ModuleList(convs)
+        self.bns = nn.ModuleList([BatchNorm(hidden) for _ in range(5)])
+        self.graph_proj = nn.Linear(hidden, hidden)
+
+        self.desc_net = nn.Sequential(nn.LayerNorm(desc_dim), nn.Linear(desc_dim, 32), nn.GELU())
+        self.seq_net = nn.Sequential(nn.LayerNorm(seq_dim), nn.Linear(seq_dim, 32), nn.GELU())
+        self.d4_net = nn.Sequential(nn.LayerNorm(d4_dim), nn.Linear(d4_dim, 16), nn.GELU())
+
+        fusion = hidden + 32 + 32 + 16
+        self.head = nn.Sequential(
+            nn.Linear(fusion, hidden//2), nn.GELU(), nn.Dropout(0.2),
+            nn.Linear(hidden//2, hidden//4), nn.GELU(), nn.Dropout(0.1),
+            nn.Linear(hidden//4, 1))
+
+    def forward(self, pyg_data, desc, seq, d4):
+        x = F.relu(self.node_emb(pyg_data.x))
+        e = F.relu(self.edge_emb(pyg_data.edge_attr))
+        xs = []
+        for conv, bn in zip(self.convs, self.bns):
+            x = F.relu(bn(conv(x, pyg_data.edge_index)))
+            xs.append(x)
+        h_g = self.graph_proj(global_mean_pool(x, pyg_data.batch))
+        h_d = self.desc_net(desc)
+        h_s = self.seq_net(seq)
+        h_4 = self.d4_net(d4)
+        return self.head(torch.cat([h_g, h_d, h_s, h_4], 1)).squeeze(-1)
+
+# ─── Training ────────────────────────────────────────────────────
+def train_epoch(model, loader, opt, device):
+    model.train(); total = 0
+    for batch in loader:
+        g = batch[0].to(device)
+        desc, seq, d4, t = [x.to(device) for x in batch[1:]]
+        opt.zero_grad()
+        loss = F.mse_loss(model(g, desc, seq, d4), t)
+        loss.backward()
+        torch.nn.utils.clip_grad_norm_(model.parameters(), 3.0)
+        opt.step()
+        total += loss.item() * t.size(0)
+    return total / len(loader.dataset)
+
+@torch.no_grad()
+def evaluate(model, loader, device):
+    model.eval(); preds, targets = [], []
+    for batch in loader:
+        g = batch[0].to(device)
+        desc, seq, d4, t = [x.to(device) for x in batch[1:]]
+        preds.append(model(g, desc, seq, d4).cpu()); targets.append(t.cpu())
+    p = torch.cat(preds); t = torch.cat(targets)
+    mse = F.mse_loss(p, t).item()
+    return mse, F.l1_loss(p, t).item(), 1 - mse / t.var().item() if t.var().item() > 0 else 0
+
+# ─── Main ────────────────────────────────────────────────────────
+def main():
+    device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
+    print(f'Device: {device}')
+
+    data = load_data()
+    ap = torch.tensor([float(d['PAMPA']) for d in data])
+    tm, ts = ap.mean(), ap.std()
+    print(f'Target: mean={tm:.3f} std={ts:.3f}')
+
+    results = []
+    for aug in [1]:
+        name = f'GIN_aug{aug}'
+        ds = CycPepDataset(data, tm, ts, augment=aug)
+        n = len(ds); idx = list(range(n))
+        random.seed(42); random.shuffle(idx)
+        tr, va = int(0.8*n), int(0.1*n); te = n - tr - va
+        tr_i, va_i, te_i = idx[:tr], idx[tr:tr+va], idx[tr+va:]
+
+        bs = min(128, n//10)
+        tr_l = DataLoader(torch.utils.data.Subset(ds, tr_i), bs, shuffle=True, collate_fn=collate_fn)
+        va_l = DataLoader(torch.utils.data.Subset(ds, va_i), bs, shuffle=False, collate_fn=collate_fn)
+        te_l = DataLoader(torch.utils.data.Subset(ds, te_i), bs, shuffle=False, collate_fn=collate_fn)
+
+        model = CycPepGNN(hidden=256).to(device)
+        np_ = sum(p.numel() for p in model.parameters())
+        print(f'\n{name}: {n} samples, {np_:,} params')
+
+        opt = torch.optim.AdamW(model.parameters(), lr=3e-4, weight_decay=1e-4)
+        sc = torch.optim.lr_scheduler.CosineAnnealingLR(opt, T_max=150)
+        bv = float('inf'); bt = None; pt_ = 0; t0 = time.time()
+        for ep in range(150):
+            loss = train_epoch(model, tr_l, opt, device)
+            vm, vma, vr2 = evaluate(model, va_l, device)
+            sc.step()
+            if (ep+1) % 15 == 0 or ep == 0:
+                print(f'  E{ep+1:3d} loss={loss:.4f} val_mse={vm*ts**2:.4f} val_r2={vr2:.4f}')
+            if vm < bv:
+                bv = vm; bt = evaluate(model, te_l, device); pt_ = 0
+            else:
+                pt_ += 1
+                if pt_ >= 25: break
+        te_m_u = bt[0]*ts**2; te_ma_u = bt[1]*ts
+        elapsed = time.time() - t0
+        print(f'  TEST: MSE={te_m_u:.4f} MAE={te_ma_u:.4f} R²={bt[2]:.4f} time={elapsed:.0f}s')
+        results.append((name, te_m_u, te_ma_u, bt[2]))
+
+    print('\n' + '='*60)
+    print(f'{"Model":<20} {"MSE":<10} {"MAE":<10} {"R²":<10}')
+    print('-'*60)
+    for r in results:
+        print(f'{r[0]:<20} {r[1]:<10.4f} {r[2]:<10.4f} {r[3]:<10.4f}')
+    print('-'*60)
+    print(f'{"MSF-CPMP (SOTA)":<20} {"0.092":<10} {"0.242":<10} {"~0.88":<10}')
+    print(f'{"CycPeptMP":<20} {"0.271":<10} {"0.355":<10} {"0.780":<10}')
+
+if __name__ == '__main__':
+    main()