PAFlow

🚀 Implementation of NeurIPS 2025 paper "Prior-Guided Flow Matching for Target-Aware Molecule Design with Learnable Atom Number"

🧩 Environment

The environment configuration is the same as that of KGDiff.

📊 Datasets

To run the code download data.zip and extract it to the ./data directory. The original CrossDocked protein pocket files are also included in crossdocked_v1.1_rmsd1.0_pocket10.zip.

📄 Pretrained Models

Download the pretrained_models.zip and place it in the ./pretrained_models directory.

🛠️ Training

cd PAFlow

To train PAFlow:

python ./scripts/train_vpflow_guide_ba.py

To train atom number predictor:

python ./scripts/train_num_predictor.py

⚙️ Sampling

To sample molecules for the test set in CrossDocked2020:

python ./scripts/sample_flow_VP_guide.py --config ./configs/sampling_guide.yml --result_path ./results -i {i} --pos_grad_w 350.0

Replace {i} with the index of the data. i should be between 0 and 99 for the testset. --pos_grad_w controls the guidance strength. The value is set to 350 in the paper but can be adjusted as needed.

⚖️ Evaluating

python ./scripts/evaluate_diffusion.py ./results

Our evaluated results are provided in PAFlow_results.pt.

💊 Generating for a Given Pocket

When generating molecules for a given protein pocket, you need to calculate the pocket’s volume and area first to predict the number of atoms. In this work we use pyKVFinder for these calculations but you may use any other tool you are familiar with. Follow the documentation to install.

It is strongly recommended to install pyKVFinder in a new environment since its dependencies conflict with those used in PAFlow.

Assume the environment used for PAFlow is named paflow and pyKVFinder is installed in the environment pykvfinder. Given the PDB file of the protein pocket 2Z3H and the SDF file of its reference ligand as example, the molecule generation process proceeds as follows:

• Convert the reference ligand SDF file to a PDB file:

python ./scripts/data_preparation/sdf_to_pdb.py --ligand_sdf_path ./example/2z3h_ligand.sdf

• Calculate the pocket’s volume and surface area:

conda activate pykvfinder
python ./scripts/data_preparation/calculate_volume_area.py --pocket_path ./example/2z3h_pocket10.pdb --ligand_path ./example/2z3h_ligand.pdb

Output:

Volume: 232.63
Area: 229.08

• Generate molecules for the pocket:

conda activate paflow
python ./scripts/sample_for_pocket.py --pdb_path ./example/2z3h_pocket10.pdb --volume 232.63 --area 229.08