FlowMatching-SBDD

This repository contains the code for FlowMatching-SBDD.

Installation

To install the required dependencies, run the following command with python version 3.10:

pip install -r requirements.txt

Notebook train BAPNet: https://www.kaggle.com/code/minhtrancolab1/bapnet-train-pdbbindv2016

Put data into ./dataset from this link: https://drive.google.com/drive/folders/1tCegA_jvoQbc3_yKF2aQKoDBs9JAxQKZ?usp=sharing

Put pretrained model in ./pretrained_models: https://drive.google.com/drive/folders/1OBEkUYoSN-QY1g5eL-UxW72RMpVkTJwW?usp=sharing. This folder include:

• checkpoint for BAPNet
• checkpoint for backbone

Train model from scratch

To train the model from scratch, run the train.py script. You can specify the configuration file using the --config argument. Setting parameter fm_variant between "pv-shift" for shift version and "pv" for no shift version. For other argument please refer in train.py.

Remember to change the root_dir

python train.py --config configs/training.yml --logdir ...

Sample the data

To sample the data, run the sample_split.py script. You can specify the sampling configuration and training configuration files especially the model path.

Remember to change the root_dir

python sample_split.py --batch_size 20 --result_path ...

Evaluate generated samples

To evaluate the generated samples, run the eval_split.py script. You can specify the path to the sampled results using the --sample_path argument.

Remember to change the root_dir

python eval_split.py --sample_path sampled_results/

Hyper-parameters Configuration

1. Backbone (uni02transformer)

Hyperparameter	Value	Description
num_layers	9	The number of layers inside the transformer network. Determines the depth of the message passing process on the graph.
num_blocks	1	The number of main blocks. Setting it to 1 optimizes computational cost while maintaining model capacity.
hidden_dim	128	The dimensionality of the hidden feature vector for each atom (node).
n_heads	16	The number of attention heads. Allows the model to capture spatial structure and chemical properties from multiple perspectives simultaneously.
cutoff_mode	knn	Graph connection mode using $K$-Nearest Neighbors instead of a fixed radius.
knn	32	The number of maximum neighbors connected to each atom, which keeps the graph structure stable.
num_r_gaussian	20	The number of Gaussian basis functions used to embed the continuous distance between atoms.
r_max	10.0	The maximum distance (in Ångström) for spatial interactions to be considered valid.
num_x2h / num_h2x	1 / 1	The number of update steps from Coordinates to Hidden features (x2h) and vice versa (h2x). Ensures 3D O(2)/Equivariant properties.
ew_net_type	global	Edge weight prediction network using local distance learning to discount less important edges/bonds.

2. BAPNet & Flow Matching

Hyperparameter	Value	Description
use_flow_matching	True	Enables Flow Matching generative mechanism instead of standard Diffusion (DDPM).
fm_variant	pv-shift	The Flow Matching variant. Two main values are pv and pv-shift. The pv-shift version uses the prior from BAPNet as a shift function to guide the generation trajectory directly toward a chemically stable state.
cond_dim	128	The dimensionality of the condition vector extracted from BAPNet, synchronized with the hidden_dim of the backbone.
fm_pred_x0	True	Indicates that we optimize (calculate loss) directly on the predicted $x_0$, rather than calculating the velocity at time $t$ from $x_0$ and optimizing on that velocity.
model_mean_type	C0	The optimization target of the model is C0 (recovering the original coordinate $x_0$).
num_diffusion_timesteps	1000	The number of integration steps splitting the time $t \in [0,1]$ for ODE solving, leading to smoother generated molecular geometries.
beta_schedule	sigmoid	The noise schedule for continuous coordinates. A sigmoid function better preserves the global structure in early generation steps.
v_beta_schedule	cosine	The optimal noise schedule designed specifically for categorical variables (atom types).
loss_v_weight	100.0	A very large penalty weight for the categorical loss (atom type prediction). It forces the model to strictly pick the correct chemical elements, even if coordinates are accurate.

3. Sampling Generation

Hyperparameter	Value	Description
num_samples	100	The number of ligand molecules generated per protein pocket.
num_steps	600	The number of integration steps during the sampling/generation phase. Can differ from training num_diffusion_timesteps.
sample_num_atoms	prior	Defines how the number of atoms for the generated ligand is chosen. prior samples it based on the protein pocket size.
center_pos_mode	protein	Centers the coordinate system relative to the protein center of mass during generation.
guidance_rho	0.0	The strength of affinity guidance during sampling. Default 0.0 means OFF. Setting it to >0 (e.g., 0.3) enables guided generation toward higher binding affinity.
guidance_t_thresh	0.2	The time threshold for applying guidance. 0.2 means guidance is only applied during the last 20% of the integration steps.
pos_only	False	When False, the model generates both the 3D coordinates and categorical atom types. If True, it only generates coordinates.